From cba9d21837b3688bca923b1e3acec261441ff74e Mon Sep 17 00:00:00 2001 From: Md Nafis Ul Haque Shifat Date: Thu, 2 Jul 2026 23:28:17 +0000 Subject: [PATCH 01/47] [#16251][feat] KV-cache disk tier: config plumbing, poolless disk blocks, host->disk spill and disk->GPU onboard Based on v1.3.0rc19. Adds a third KV-cache tier backed by POSIX slot files: - KvCacheConfig gains diskCacheSize/diskCachePath (serialization + nanobind + llm_args passthrough) - Poolless disk blocks at eviction level 2 (placeholders renumber to 3); block residency swaps between pool index and disk slot so the reuse-tree identity never moves - Host-tier victims with reusable content spill to disk on reclaim (spill-all; filtering mode to follow); disk-resident matches onboard file->GPU in onboardBlock - Synchronous POSIX I/O v1; pytorch path only (TRT-engine path passes 0) Validated on gpt-oss-120b: 49824/49836 cached tokens served through a disk round trip with temperature-0 byte-identical output; 300k+ spills and 11k+ onboards under sustained production traffic with zero I/O errors. Signed-off-by: Md Nafis Ul Haque Shifat --- .../batch_manager/kvCacheManager.h | 92 ++++++++--- .../batch_manager/kvCacheTransferManager.h | 8 + cpp/include/tensorrt_llm/executor/executor.h | 10 ++ .../batch_manager/evictionPolicy.cpp | 11 +- .../batch_manager/kvCacheManager.cpp | 144 +++++++++++++----- .../batch_manager/kvCacheTransferManager.cpp | 56 +++++++ .../trtGptModelInflightBatching.cpp | 2 +- cpp/tensorrt_llm/executor/kvCacheConfig.cpp | 20 +++ cpp/tensorrt_llm/executor/serialization.cpp | 15 +- .../nanobind/batch_manager/kvCacheManager.cpp | 24 +-- .../nanobind/executor/executorConfig.cpp | 9 +- .../_torch/pyexecutor/resource_manager.py | 11 ++ tensorrt_llm/llmapi/llm_args.py | 3 + 13 files changed, 318 insertions(+), 87 deletions(-) diff --git a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h index 299bf1b570f4..833f739d3b25 100644 --- a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h +++ b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h @@ -67,6 +67,7 @@ namespace tensorrt_llm::batch_manager::kv_cache_manager static constexpr SizeType32 kPrimaryLevel = 0; static constexpr SizeType32 kSecondaryLevel = 1; +static constexpr SizeType32 kDiskLevel = 2; // Extra block buffer allocated for SWA to be able to always keep "window size" // tokens held in the blocks. @@ -436,6 +437,33 @@ class KVCacheBlock : public std::enable_shared_from_this void swapMemoryPoolBlockOffset(std::shared_ptr otherBlock); + static constexpr SizeType32 kNoDiskSlot = -1; + + //! \brief True when this block's bytes live in the disk cache tier. + [[nodiscard]] bool isOnDisk() const + { + return mDiskSlot != kNoDiskSlot; + } + + [[nodiscard]] SizeType32 getDiskSlot() const + { + return mDiskSlot; + } + + void setDiskSlot(SizeType32 diskSlot) + { + mDiskSlot = diskSlot; + } + + //! \brief Exchange tier residency with \p other: pool index and disk slot swap together. + //! Used by the host->disk spill and the disk->GPU onboard; the tree-resident identity + //! keeps its node/hash while the physical backing moves. + void swapDiskResidency(std::shared_ptr const& other) + { + std::swap(mMemoryPoolBlockIndex, other->mMemoryPoolBlockIndex); + std::swap(mDiskSlot, other->mDiskSlot); + } + void incRefCount(); void decRefCount(); @@ -590,6 +618,9 @@ class KVCacheBlock : public std::enable_shared_from_this std::optional mExpirationTime; // Hash for the event manager size_t mHash; + + // Disk cache tier: slot index in the disk file pool; kNoDiskSlot when not disk-resident. + SizeType32 mDiskSlot{kNoDiskSlot}; }; class GenerationRequest @@ -875,8 +906,9 @@ class WindowBlockManager explicit WindowBlockManager(nvinfer1::DataType dtype, SizeType32 windowSize, std::vector const& managedLayers, std::vector const& numKvHeadsPerLayer, SizeType32 sizePerHead, SizeType32 tokensPerBlock, bool isSWA, SizeType32 blocksInPrimaryPool, - SizeType32 blocksInSecondaryPool, SizeType32 maxNumSequences, std::shared_ptr stream, - CacheType cacheType, std::optional secondaryOffloadMinPriority, + SizeType32 blocksInSecondaryPool, SizeType32 blocksInDiskPool, std::string const& diskCachePath, + SizeType32 maxNumSequences, std::shared_ptr stream, CacheType cacheType, + std::optional secondaryOffloadMinPriority, std::shared_ptr eventManager, bool enablePartialReuse, bool copyOnPartialReuse, std::shared_ptr kvCacheConnectorManager, radix_block_tree::UnifiedBlockTree& lookupTree, std::shared_ptr loopbackAgent = nullptr, @@ -1197,6 +1229,11 @@ class WindowBlockManager //! \brief Bring offloaded block from secondary to primary memory. //! \details Does nothing if block is already in primary memory. + //! \brief Reclaim a free secondary (host) block for reuse. If the victim still holds + //! reusable content and the disk tier has room, spill it to disk first (residency swap; + //! the tree-resident identity moves to the disk level). Returns a CLAIMED block. + [[nodiscard]] BlockPtr reclaimSecondaryBlock(); + void onboardBlock(GenerationRequest& sequence, BlockPtr const& offloadBlock, executor::KvCacheTransferMode mode = executor::KvCacheTransferMode::DRAM, std::string const& directory = ""); @@ -1359,6 +1396,10 @@ class WindowBlockManager // Number of blocks in pools SizeType32 mNumPrimaryBlocks; SizeType32 mNumSecondaryBlocks; + SizeType32 mNumDiskBlocks; + std::string mDiskCachePath; + std::size_t mDiskSpills{0}; + std::size_t mDiskOnboards{0}; // List of allocated blocks for each sequences std::unordered_map> mAllocatedBlocksPerSeq; @@ -1479,10 +1520,11 @@ class BlockManager //! 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, - std::vector const& maxAttentionWindowVec, nvinfer1::DataType dtype, SizeType32 sinkBubbleLength, - SizeType32 chunkSize, CacheType cacheType = CacheType::kSELF, + SizeType32 tokensPerBlock, BlocksPerWindow const& blocksPerWindow, SizeType32 blocksInDiskPool, + std::string const& diskCachePath, SizeType32 maxNumSequences, CudaStreamPtr stream, + SizeType32 maxSequenceLength, SizeType32 maxBeamWidth, std::vector const& maxAttentionWindowVec, + nvinfer1::DataType dtype, SizeType32 sinkBubbleLength, SizeType32 chunkSize, + CacheType cacheType = CacheType::kSELF, std::optional secondaryOffloadMinPriority = std::nullopt, std::shared_ptr eventManager = nullptr, bool enablePartialReuse = true, bool copyOnPartialReuse = true, @@ -2275,10 +2317,10 @@ class KVCacheManager : public BaseKVCacheManager //! + 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, - CudaStreamPtr stream, SizeType32 maxSequenceLength, SizeType32 chunkSize, bool enableBlockReuse = false, - CacheType cacheType = CacheType::kSELF, + BlocksPerWindow const& blocksPerWindow, SizeType32 blocksInDiskPool, std::string const& diskCachePath, + SizeType32 maxNumSequences, SizeType32 maxBeamWidth, std::vector const& maxAttentionWindowVec, + nvinfer1::DataType dtype, SizeType32 sinkTokenLength, CudaStreamPtr stream, SizeType32 maxSequenceLength, + SizeType32 chunkSize, bool enableBlockReuse = false, CacheType cacheType = CacheType::kSELF, std::optional secondaryOffloadMinPriority = std::nullopt, std::shared_ptr eventManager = nullptr, bool enablePartialReuse = true, bool copyOnpartialReuse = true, @@ -2289,10 +2331,10 @@ class KVCacheManager : public BaseKVCacheManager std::vector const& poolConfigurations = {}); 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, - int64_t stream, SizeType32 maxSequenceLength, SizeType32 chunkSize, bool enableBlockReuse = false, - CacheType cacheType = CacheType::kSELF, + BlocksPerWindow const& blocksPerWindow, SizeType32 blocksInDiskPool, std::string const& diskCachePath, + SizeType32 maxNumSequences, SizeType32 maxBeamWidth, std::vector const& maxAttentionWindowVec, + nvinfer1::DataType dtype, SizeType32 sinkTokenLength, int64_t stream, SizeType32 maxSequenceLength, + SizeType32 chunkSize, bool enableBlockReuse = false, CacheType cacheType = CacheType::kSELF, std::optional secondaryOffloadMinPriority = std::nullopt, std::shared_ptr eventManager = nullptr, bool enablePartialReuse = true, bool copyOnpartialReuse = true, @@ -2303,10 +2345,10 @@ class KVCacheManager : public BaseKVCacheManager std::vector const& poolConfigurations = {}); KVCacheManager(SizeType32 numLayers, SizeType32 numKvHeads, SizeType32 sizePerHead, SizeType32 tokensPerBlock, - BlocksPerWindow const& blocksPerWindow, SizeType32 maxNumSequences, SizeType32 maxBeamWidth, - std::vector const& maxAttentionWindowVec, nvinfer1::DataType dtype, SizeType32 sinkTokenLength, - CudaStreamPtr stream, SizeType32 maxSequenceLength, SizeType32 chunkSize, bool enableBlockReuse = true, - CacheType cacheType = CacheType::kSELF, + BlocksPerWindow const& blocksPerWindow, SizeType32 blocksInDiskPool, std::string const& diskCachePath, + SizeType32 maxNumSequences, SizeType32 maxBeamWidth, std::vector const& maxAttentionWindowVec, + nvinfer1::DataType dtype, SizeType32 sinkTokenLength, CudaStreamPtr stream, SizeType32 maxSequenceLength, + SizeType32 chunkSize, bool enableBlockReuse = true, CacheType cacheType = CacheType::kSELF, std::optional secondaryOffloadMinPriority = std::nullopt, std::shared_ptr eventManager = nullptr, bool enablePartialReuse = true, bool copyOnpartialReuse = true, @@ -2317,13 +2359,13 @@ class KVCacheManager : public BaseKVCacheManager std::vector const& poolConfigurations = {}); KVCacheManager(SizeType32 numLayers, SizeType32 numKvHeads, SizeType32 sizePerHead, SizeType32 tokensPerBlock, - BlocksPerWindow const& blocksPerWindow, SizeType32 maxNumSequences, SizeType32 maxBeamWidth, - std::vector const& maxAttentionWindowVec, nvinfer1::DataType dtype, SizeType32 sinkTokenLength, - int64_t stream, SizeType32 maxSequenceLength, SizeType32 chunkSize, bool enableBlockReuse = false, - 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, + BlocksPerWindow const& blocksPerWindow, SizeType32 blocksInDiskPool, std::string const& diskCachePath, + SizeType32 maxNumSequences, SizeType32 maxBeamWidth, std::vector const& maxAttentionWindowVec, + nvinfer1::DataType dtype, SizeType32 sinkTokenLength, int64_t stream, SizeType32 maxSequenceLength, + SizeType32 chunkSize, bool enableBlockReuse = false, 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::vector const& poolConfigurations = {}); ~KVCacheManager() override = default; diff --git a/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h b/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h index 9e9a07b19d8e..e7904287befe 100644 --- a/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h +++ b/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h @@ -75,6 +75,14 @@ class KVCacheTransferManager //! \brief Synchronize bufferManager stream with internal streams. This method ensures that prefill and decode //! kernels for next step will wait for offloading and onboarding work that has already been scheduled. This method //! must be called after the last call to KVCacheManager::addSequenceBatch in every step. + //! \brief Disk tier: write a host-resident block's bytes to its slot files (synchronous POSIX). + void spillToFile(BlockPtr const& srcHostBlock, SizeType32 diskSlot, std::vector const& pools, + std::string const& directory); + + //! \brief Disk tier: read slot files into a GPU-resident block (synchronous POSIX). + void loadFromFile(BlockPtr const& dstPrimaryBlock, SizeType32 diskSlot, std::vector const& pools, + std::string const& directory); + void syncTransfers(); //! \brief Get transfer stats accumulated since last call, and reset the counters. diff --git a/cpp/include/tensorrt_llm/executor/executor.h b/cpp/include/tensorrt_llm/executor/executor.h index 825b8ad75959..8af6b8742d54 100644 --- a/cpp/include/tensorrt_llm/executor/executor.h +++ b/cpp/include/tensorrt_llm/executor/executor.h @@ -1097,6 +1097,10 @@ class KvCacheConfig void setUseUvm(bool useUvm); void setAttentionDpEventsGatherPeriodMs(SizeType32 attentionDpEventsGatherPeriodMs); void setMaxGpuTotalBytes(uint64_t maxGpuTotalBytes); + [[nodiscard]] std::optional getDiskCacheSize() const; + [[nodiscard]] std::string const& getDiskCachePath() const; + void setDiskCacheSize(std::optional diskCacheSize); + void setDiskCachePath(std::string const& diskCachePath); void fillEmptyFieldsFromRuntimeDefaults(tensorrt_llm::runtime::RuntimeDefaults const& runtimeDefaults); @@ -1156,6 +1160,12 @@ class KvCacheConfig /// If both mMaxGpuTotalBytes and mFreeGpuMemoryFraction are specified, memory corresponding to the minimum will /// be allocated. uint64_t mMaxGpuTotalBytes; + + /// @brief Size in bytes of the disk tier backing explicit-cache blocks. Disabled when unset or 0. + std::optional mDiskCacheSize; + + /// @brief Directory holding the disk tier's block files. Must be non-empty when mDiskCacheSize > 0. + std::string mDiskCachePath; }; /// @brief Configuration class for the runtime perf knobs diff --git a/cpp/tensorrt_llm/batch_manager/evictionPolicy.cpp b/cpp/tensorrt_llm/batch_manager/evictionPolicy.cpp index 42b5bd9e6879..446ecdd72daf 100644 --- a/cpp/tensorrt_llm/batch_manager/evictionPolicy.cpp +++ b/cpp/tensorrt_llm/batch_manager/evictionPolicy.cpp @@ -33,8 +33,9 @@ auto const kNumPriorities = kMaxPriority - kMinPriority + 1; auto const kDefaultPriority = executor::KvCacheRetentionConfig::kDefaultRetentionPriority; executor::RetentionPriority const kDefaultSecondaryOffloadMinPriority = 30; -int const kNumCacheLevels = 2; -int const kPlaceholderLevel = kNumCacheLevels; // placeholder blocks live at level 2 +int const kNumCacheLevels = 3; // 0 = GPU (primary), 1 = host (secondary), 2 = disk +int const kDiskLevel = 2; +int const kPlaceholderLevel = kNumCacheLevels; // placeholders live one past the real cache levels namespace { @@ -44,6 +45,10 @@ SizeType32 getCacheLevel(BlockPtr const& block) { return kPlaceholderLevel; } + if (block->isOnDisk()) + { + return kDiskLevel; + } return block->isPrimary() ? 0 : 1; } @@ -58,6 +63,8 @@ constexpr auto defaultPriorityIdx = getPriorityIdx(kDefaultPriority); void LRUEvictionPolicy::initialize(std::vector& mAllBlocksById, std::vector sizes, std::optional secondaryOffloadMinPriority) { + TLLM_CHECK_WITH_INFO(static_cast(sizes.size()) == kNumCacheLevels, + "initialize expects one block count per cache level (%d), got %zu", kNumCacheLevels, sizes.size()); SizeType32 startIdx = 0; // Create queues for all levels: primary, secondary, and placeholder (initially empty). diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp index c67a4711cc03..bf9292e36da4 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp @@ -584,9 +584,10 @@ std::map BlockManager::calculateWindowSizeToShare( } BlockManager::BlockManager(std::vector const& numKvHeadsPerLayer, SizeType32 sizePerHead, - SizeType32 tokensPerBlock, BlocksPerWindow const& blocksPerWindow, SizeType32 maxNumSequences, CudaStreamPtr stream, - SizeType32 maxSequenceLength, SizeType32 maxBeamWidth, std::vector const& maxAttentionWindowVec, - nvinfer1::DataType dtype, SizeType32 sinkBubbleLength, SizeType32 chunkSize, CacheType cacheType, + SizeType32 tokensPerBlock, BlocksPerWindow const& blocksPerWindow, SizeType32 blocksInDiskPool, + std::string const& diskCachePath, SizeType32 maxNumSequences, CudaStreamPtr stream, SizeType32 maxSequenceLength, + SizeType32 maxBeamWidth, std::vector const& maxAttentionWindowVec, nvinfer1::DataType dtype, + SizeType32 sinkBubbleLength, SizeType32 chunkSize, CacheType cacheType, std::optional secondaryOffloadMinPriority, std::shared_ptr eventManager, bool enablePartialReuse, bool copyOnPartialReuse, std::shared_ptr kvCacheConnectorManager, @@ -691,7 +692,9 @@ BlockManager::BlockManager(std::vector const& numKvHeadsPerLayer, Si mWindowBlockManagers.try_emplace(SizeType32(windowSize), windowDtype, windowSize, layersWithWindowSize, numKvHeadsPerLayer, windowSizePerHead, tokensPerBlock, /*isSWA=*/(windowSize < maxSequenceLength) && (windowSize >= 0), allottedPrimaryBlocks, - allottedSecondaryBlocks, maxNumSequences, stream, cacheType, secondaryOffloadMinPriority, mEventManager, + allottedSecondaryBlocks, + (SizeType32(windowSize) == blocksPerWindow.rbegin()->first) ? blocksInDiskPool : SizeType32{0}, + diskCachePath, maxNumSequences, stream, cacheType, secondaryOffloadMinPriority, mEventManager, enablePartialReuse, copyOnPartialReuse, kvCacheConnectorManager, mLookupTree, mLoopbackAgent, enableIndexerKCache, indexerKCacheQuantBlockSize, indexerKCacheIndexHeadDim, indexerKCacheUseFp4, LinearAttentionMetadata::hasLinearCache(windowSize) ? linearAttentionMetadata : std::nullopt, @@ -743,8 +746,9 @@ BlockManager::BlockManager(std::vector const& numKvHeadsPerLayer, Si WindowBlockManager::WindowBlockManager(nvinfer1::DataType dtype, SizeType32 windowSize, std::vector const& managedLayers, std::vector const& numKvHeadsPerLayer, SizeType32 sizePerHead, SizeType32 tokensPerBlock, bool isSWA, SizeType32 blocksInPrimaryPool, - SizeType32 blocksInSecondaryPool, SizeType32 maxNumSequences, std::shared_ptr stream, - CacheType cacheType, std::optional secondaryOffloadMinPriority, + SizeType32 blocksInSecondaryPool, SizeType32 blocksInDiskPool, std::string const& diskCachePath, + SizeType32 maxNumSequences, std::shared_ptr stream, CacheType cacheType, + std::optional secondaryOffloadMinPriority, std::shared_ptr eventManager, bool enablePartialReuse, bool copyOnPartialReuse, std::shared_ptr kvCacheConnectorManager, radix_block_tree::UnifiedBlockTree& lookupTree, std::shared_ptr loopbackAgent, @@ -755,6 +759,8 @@ WindowBlockManager::WindowBlockManager(nvinfer1::DataType dtype, SizeType32 wind , mWindowSize{windowSize} , mNumPrimaryBlocks{blocksInPrimaryPool} , mNumSecondaryBlocks{blocksInSecondaryPool} + , mNumDiskBlocks{blocksInDiskPool} + , mDiskCachePath{diskCachePath} , mBufferManager{std::move(stream)} , mSchedulingNumFreeBlocks{0} , mTokensPerBlock{tokensPerBlock} @@ -852,7 +858,7 @@ WindowBlockManager::WindowBlockManager(nvinfer1::DataType dtype, SizeType32 wind } // Create free blocks - mAllBlocksById.reserve(blocksInPrimaryPool + blocksInSecondaryPool); + mAllBlocksById.reserve(blocksInPrimaryPool + blocksInSecondaryPool + blocksInDiskPool); for (KVCacheBlock::IdType blockId = 0; blockId < blocksInPrimaryPool; ++blockId) { mAllBlocksById.emplace_back(std::make_shared(blockId, tk::KVCacheIndex{blockId, false})); @@ -862,6 +868,19 @@ WindowBlockManager::WindowBlockManager(nvinfer1::DataType dtype, SizeType32 wind mAllBlocksById.emplace_back( std::make_shared(blocksInPrimaryPool + blockId, tk::KVCacheIndex{blockId, true})); } + // Disk cache tier: poolless blocks (no GPU/host memory; bytes live in a file slot). + // Never kernel-visible - a disk block is always onboarded to a primary block before use. + for (KVCacheBlock::IdType blockId = 0; blockId < blocksInDiskPool; ++blockId) + { + auto diskBlock = std::make_shared( + blocksInPrimaryPool + blocksInSecondaryPool + blockId, tk::KVCacheIndex::nullIndex); + diskBlock->setDiskSlot(blockId); + mAllBlocksById.emplace_back(std::move(diskBlock)); + } + if (blocksInDiskPool > 0) + { + TLLM_LOG_INFO("[disk-tier] enabled: %d poolless blocks (windowSize=%d)", blocksInDiskPool, windowSize); + } mAllocatedBlocksPerSeq.reserve(maxNumSequences); // Pre-allocate placeholder blocks when this is a recurrent-state (linear attention) manager paired with @@ -888,15 +907,16 @@ WindowBlockManager::WindowBlockManager(nvinfer1::DataType dtype, SizeType32 wind } auto policy = std::make_shared(); - policy->initialize(mAllBlocksById, {blocksInPrimaryPool, blocksInSecondaryPool}, secondaryOffloadMinPriority); + policy->initialize(mAllBlocksById, {blocksInPrimaryPool, blocksInSecondaryPool, blocksInDiskPool}, + secondaryOffloadMinPriority); policy->initializePlaceholders(mAllPlaceholderBlocksById); mEvictionPolicy = policy; } else { mEvictionPolicy = std::make_shared(); - mEvictionPolicy->initialize( - mAllBlocksById, {blocksInPrimaryPool, blocksInSecondaryPool}, secondaryOffloadMinPriority); + mEvictionPolicy->initialize(mAllBlocksById, {blocksInPrimaryPool, blocksInSecondaryPool, blocksInDiskPool}, + secondaryOffloadMinPriority); } if (mEventManager) { @@ -1270,6 +1290,27 @@ void WindowBlockManager::releaseSubtree(BlockPtr const& block) } } +BlockPtr WindowBlockManager::reclaimSecondaryBlock() +{ + auto victim = std::get<0>(mEvictionPolicy->getFreeBlock(kSecondaryLevel)); + mEvictionPolicy->claimBlock(victim); + if (mNumDiskBlocks == 0 || victim->getUniqueTokens().empty() || mEvictionPolicy->getNumFreeBlocks(kDiskLevel) == 0) + { + return victim; // disk tier off, nothing reusable to keep, or disk full: same as before + } + auto diskTarget = std::get<0>(mEvictionPolicy->getFreeBlock(kDiskLevel)); + mEvictionPolicy->claimBlock(diskTarget); // claim BOTH before the swap (cf. #11879) + mTransferManager->spillToFile(victim, diskTarget->getDiskSlot(), mPools, mDiskCachePath); + victim->swapDiskResidency(diskTarget); // victim's identity now disk-resident, tree intact + mEvictionPolicy->releaseBlock(victim); // re-enters the free queues at kDiskLevel + ++mDiskSpills; + if (mDiskSpills == 1 || mDiskSpills % 1000 == 0) + { + TLLM_LOG_INFO("[disk-tier] spills=%zu (windowSize=%d)", mDiskSpills, mWindowSize); + } + return diskTarget; // owns the victim's old host slot; already claimed +} + BlockPtr WindowBlockManager::getFreeBlock(GenerationRequest& sequence, executor::RetentionPriority priority, std::optional durationMs, executor::KvCacheTransferMode mode, std::string const& directory, bool wantPlaceholder) @@ -1289,7 +1330,7 @@ BlockPtr WindowBlockManager::getFreeBlock(GenerationRequest& sequence, executor: && mEvictionPolicy->getNumFreeBlocks(kSecondaryLevel) > 0) { // Offload block in primary memory before repurposing - auto offloadBlock = std::get<0>(mEvictionPolicy->getFreeBlock(kSecondaryLevel)); + auto offloadBlock = reclaimSecondaryBlock(); // Claim both blocks BEFORE the swap so that getCacheLevel() returns the // correct pre-swap level. Previously the claims happened after the swap, @@ -1397,6 +1438,21 @@ void BlockManager::onboardBlock(GenerationRequest& sequence, BlockPtr const& off void WindowBlockManager::onboardBlock(GenerationRequest& sequence, BlockPtr const& offloadBlock, executor::KvCacheTransferMode mode, std::string const& directory) { + if (offloadBlock->isOnDisk()) + { + // A fresh GPU block; this may itself trigger the GPU->host cascade (unchanged, correct). + auto block = getFreeBlock( + sequence, executor::KvCacheRetentionConfig::kDefaultRetentionPriority, std::nullopt, mode, directory); + mTransferManager->loadFromFile(block, offloadBlock->getDiskSlot(), mPools, mDiskCachePath); + offloadBlock->swapDiskResidency(block); // matched identity now GPU-resident + mEvictionPolicy->releaseBlock(block); // freed card returns to the disk free queue + ++mDiskOnboards; + if (mDiskOnboards == 1 || mDiskOnboards % 1000 == 0) + { + TLLM_LOG_INFO("[disk-tier] onboards=%zu (windowSize=%d)", mDiskOnboards, mWindowSize); + } + return; + } if (!offloadBlock->isPlaceholder() && !offloadBlock->isPrimary()) { auto block = getFreeBlock( @@ -1434,7 +1490,7 @@ void WindowBlockManager::offloadBlock( if (!block->isPlaceholder() && block->isPrimary()) { // Offload block in primary memory before repurposing - auto offloadBlock = std::get<0>(mEvictionPolicy->getFreeBlock(kSecondaryLevel)); + auto offloadBlock = reclaimSecondaryBlock(); // If we're swapping a block to secondary memory, maintain the prior priority values. mEvictionPolicy->claimBlock(offloadBlock); mTransferManager->offload(block, offloadBlock, mPools, 0, mode, directory); @@ -3228,15 +3284,16 @@ void WindowBlockManager::schedulingReleaseBlocks(RequestIdType requestId) } KVCacheManager::KVCacheManager(SizeType32 numLayers, SizeType32 numKvHeads, SizeType32 sizePerHead, - SizeType32 tokensPerBlock, BlocksPerWindow const& blocksPerWindow, SizeType32 maxNumSequences, - SizeType32 maxBeamWidth, std::vector const& maxAttentionWindowVec, nvinfer1::DataType dtype, - 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, + SizeType32 tokensPerBlock, BlocksPerWindow const& blocksPerWindow, SizeType32 blocksInDiskPool, + std::string const& diskCachePath, SizeType32 maxNumSequences, SizeType32 maxBeamWidth, + std::vector const& maxAttentionWindowVec, nvinfer1::DataType dtype, 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, std::vector const& poolConfigurations) : KVCacheManager(std::vector(numLayers, numKvHeads), sizePerHead, tokensPerBlock, blocksPerWindow, - maxNumSequences, maxBeamWidth, maxAttentionWindowVec, dtype, sinkTokenLength, + blocksInDiskPool, diskCachePath, 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, @@ -3245,17 +3302,18 @@ KVCacheManager::KVCacheManager(SizeType32 numLayers, SizeType32 numKvHeads, Size } KVCacheManager::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, int64_t stream, runtime::SizeType32 maxSequenceLength, SizeType32 chunkSize, - bool enableBlockReuse, CacheType cacheType, std::optional secondaryOffloadMinPriority, + SizeType32 tokensPerBlock, BlocksPerWindow const& blocksPerWindow, SizeType32 blocksInDiskPool, + std::string const& diskCachePath, SizeType32 maxNumSequences, SizeType32 maxBeamWidth, + std::vector const& maxAttentionWindowVec, nvinfer1::DataType dtype, SizeType32 sinkTokenLength, + int64_t stream, runtime::SizeType32 maxSequenceLength, SizeType32 chunkSize, bool enableBlockReuse, + CacheType cacheType, std::optional secondaryOffloadMinPriority, std::shared_ptr eventManager, bool enablePartialReuse, bool copyOnPartialReuse, std::shared_ptr kvCacheConnectorManager, bool enableIndexerKCache, SizeType32 indexerKCacheQuantBlockSize, SizeType32 indexerKCacheIndexHeadDim, bool indexerKCacheUseFp4, std::optional linearAttentionMetadata, std::vector const& poolConfigurations) - : KVCacheManager(numKvHeadsPerLayer, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, maxBeamWidth, - maxAttentionWindowVec, dtype, sinkTokenLength, + : KVCacheManager(numKvHeadsPerLayer, sizePerHead, tokensPerBlock, blocksPerWindow, blocksInDiskPool, diskCachePath, + maxNumSequences, maxBeamWidth, maxAttentionWindowVec, dtype, sinkTokenLength, std::make_shared(reinterpret_cast(stream)), maxSequenceLength, chunkSize, enableBlockReuse, cacheType, secondaryOffloadMinPriority, eventManager, enablePartialReuse, copyOnPartialReuse, kvCacheConnectorManager, enableIndexerKCache, indexerKCacheQuantBlockSize, indexerKCacheIndexHeadDim, @@ -3264,10 +3322,11 @@ KVCacheManager::KVCacheManager(std::vector const& numKvHeadsPerLayer } KVCacheManager::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, CudaStreamPtr stream, runtime::SizeType32 maxSequenceLength, SizeType32 chunkSize, - bool enableBlockReuse, CacheType cacheType, std::optional secondaryOffloadMinPriority, + SizeType32 tokensPerBlock, BlocksPerWindow const& blocksPerWindow, SizeType32 blocksInDiskPool, + std::string const& diskCachePath, SizeType32 maxNumSequences, SizeType32 maxBeamWidth, + std::vector const& maxAttentionWindowVec, nvinfer1::DataType dtype, SizeType32 sinkTokenLength, + CudaStreamPtr stream, runtime::SizeType32 maxSequenceLength, SizeType32 chunkSize, bool enableBlockReuse, + CacheType cacheType, std::optional secondaryOffloadMinPriority, std::shared_ptr eventManager, bool enablePartialReuse, bool copyOnPartialReuse, std::shared_ptr kvCacheConnectorManager, bool enableIndexerKCache, SizeType32 indexerKCacheQuantBlockSize, SizeType32 indexerKCacheIndexHeadDim, bool indexerKCacheUseFp4, @@ -3280,10 +3339,10 @@ KVCacheManager::KVCacheManager(std::vector const& numKvHeadsPerLayer , mSinkBubbleLength(BaseKVCacheManager::getSinkBubbleLength(sinkTokenLength, tokensPerBlock)) , mSinkBlockTokenLength(mSinkBubbleLength + sinkTokenLength) , mChunkSize(chunkSize) - , mBlockManager(numKvHeadsPerLayer, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, - std::move(stream), maxSequenceLength, maxBeamWidth, maxAttentionWindowVec, dtype, mSinkBubbleLength, - mChunkSize, cacheType, secondaryOffloadMinPriority, std::move(eventManager), enablePartialReuse, - copyOnPartialReuse, std::move(kvCacheConnectorManager), std::nullopt, enableIndexerKCache, + , mBlockManager(numKvHeadsPerLayer, sizePerHead, tokensPerBlock, blocksPerWindow, blocksInDiskPool, diskCachePath, + maxNumSequences, 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, poolConfigurations) // disable block reuse for sink bubble since chopVectorIntoBlocks does not match KV cache blocks in this case @@ -3306,21 +3365,22 @@ KVCacheManager::KVCacheManager(std::vector const& numKvHeadsPerLayer } KVCacheManager::KVCacheManager(SizeType32 numLayers, SizeType32 numKvHeads, SizeType32 sizePerHead, - SizeType32 tokensPerBlock, BlocksPerWindow const& blocksPerWindow, SizeType32 maxNumSequences, - SizeType32 maxBeamWidth, std::vector const& maxAttentionWindowVec, nvinfer1::DataType dtype, - SizeType32 sinkTokenLength, CudaStreamPtr stream, runtime::SizeType32 maxSequenceLength, SizeType32 chunkSize, - bool enableBlockReuse, CacheType cacheType, std::optional secondaryOffloadMinPriority, + SizeType32 tokensPerBlock, BlocksPerWindow const& blocksPerWindow, SizeType32 blocksInDiskPool, + std::string const& diskCachePath, SizeType32 maxNumSequences, SizeType32 maxBeamWidth, + std::vector const& maxAttentionWindowVec, nvinfer1::DataType dtype, SizeType32 sinkTokenLength, + CudaStreamPtr stream, runtime::SizeType32 maxSequenceLength, SizeType32 chunkSize, bool enableBlockReuse, + CacheType cacheType, std::optional secondaryOffloadMinPriority, std::shared_ptr eventManager, bool enablePartialReuse, bool copyOnPartialReuse, std::shared_ptr kvCacheConnectorManager, bool enableIndexerKCache, SizeType32 indexerKCacheQuantBlockSize, SizeType32 indexerKCacheIndexHeadDim, bool indexerKCacheUseFp4, 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, - poolConfigurations) + blocksInDiskPool, diskCachePath, 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, poolConfigurations) { } diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp index c28d1e476137..338dc8bce490 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp @@ -391,6 +391,62 @@ void KVCacheTransferManager::offload(BlockPtr const& block, BlockPtr const& offl mOffloadManager.getStream().record(mPendingWrites[offloadBlockIndex]); } +namespace +{ +std::string diskSlotFilename(std::string const& directory, SizeType32 diskSlot, size_t poolIdx) +{ + char buf[512]; + std::snprintf(buf, sizeof(buf), "%s/slot_%d_pool_%zu.bin", directory.c_str(), diskSlot, poolIdx); + return std::string(buf); +} +} // namespace + +void KVCacheTransferManager::spillToFile(BlockPtr const& srcHostBlock, SizeType32 diskSlot, + std::vector const& pools, std::string const& directory) +{ + TLLM_CHECK_WITH_INFO(!directory.empty(), "disk tier requires a directory"); + // The victim's bytes may still be the target of an in-flight async GPU->host copy; + // wait it out before reading host memory (same event discipline as copyBlock). + auto const idx = getPendingTransferIndex(srcHostBlock); + if (auto it = mPendingWrites.find(idx); it != mPendingWrites.end()) + { + it->second.synchronize(); + } + for (size_t poolIdx = 0; poolIdx < pools.size(); ++poolIdx) + { + TLLM_CHECK_WITH_INFO(!pools[poolIdx].layerFirstLayout, "disk tier does not support layer-first layout pools"); + auto ptr = computeBlockPointer(srcHostBlock, pools, poolIdx); + auto const filename = diskSlotFilename(directory, diskSlot, poolIdx); + int fd = ::open(filename.c_str(), O_CREAT | O_WRONLY, 0644); + TLLM_CHECK_WITH_INFO(fd >= 0, "disk tier: cannot open %s", filename.c_str()); + auto const bytes = static_cast(ptr->getSizeInBytes()); + auto const written = ::pwrite(fd, ptr->data(), bytes, 0); + ::close(fd); + TLLM_CHECK_WITH_INFO(written == bytes, "disk tier: short write to %s", filename.c_str()); + } +} + +void KVCacheTransferManager::loadFromFile(BlockPtr const& dstPrimaryBlock, SizeType32 diskSlot, + std::vector const& pools, std::string const& directory) +{ + TLLM_CHECK_WITH_INFO(!directory.empty(), "disk tier requires a directory"); + auto const idx = getPendingTransferIndex(dstPrimaryBlock); + if (auto it = mPendingWrites.find(idx); it != mPendingWrites.end()) + { + it->second.synchronize(); + } + if (auto it = mPendingReads.find(idx); it != mPendingReads.end()) + { + it->second.synchronize(); + } + for (size_t poolIdx = 0; poolIdx < pools.size(); ++poolIdx) + { + TLLM_CHECK_WITH_INFO(!pools[poolIdx].layerFirstLayout, "disk tier does not support layer-first layout pools"); + auto ptr = computeBlockPointer(dstPrimaryBlock, pools, poolIdx); + fileToGpuPosix(ptr, diskSlotFilename(directory, diskSlot, poolIdx)); + } +} + void KVCacheTransferManager::syncWithBufferManager() { tr::CudaEvent readyForOffloadEvent; diff --git a/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp b/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp index 7a0d78beb8a0..88cb761b6158 100644 --- a/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp +++ b/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp @@ -685,7 +685,7 @@ std::unique_ptr TrtGptModelInflightBatching::c auto const enableBlockReuse = kvCacheType == KvCacheType::kSELF ? kvCacheConfig.getEnableBlockReuse() : false; auto kvCacheManager = std::make_unique(numKvHeadsPerLayer, sizePerHead, tokensPerBlock, - blocksPerWindow, getMaxNumSequences(), getMaxBeamWidth(), maxAttentionWindowVec, kvDtype, getSinkTokenLen(), + blocksPerWindow, /*blocksInDiskPool=*/0, /*diskCachePath=*/std::string{}, getMaxNumSequences(), getMaxBeamWidth(), maxAttentionWindowVec, kvDtype, getSinkTokenLen(), mRuntime->getStreamPtr(), kvCacheType == KvCacheType::kCROSS ? mModelConfig.getMaxEncoderLen() : getMaxSequenceLen(), getMaxNumTokens().value(), enableBlockReuse, kvCacheType, kvCacheConfig.getSecondaryOffloadMinPriority(), diff --git a/cpp/tensorrt_llm/executor/kvCacheConfig.cpp b/cpp/tensorrt_llm/executor/kvCacheConfig.cpp index 54ddc199e6d3..cd2feb979a8c 100644 --- a/cpp/tensorrt_llm/executor/kvCacheConfig.cpp +++ b/cpp/tensorrt_llm/executor/kvCacheConfig.cpp @@ -202,6 +202,26 @@ void KvCacheConfig::setHostCacheSize(size_t hostCacheSize) mHostCacheSize = hostCacheSize; } +std::optional KvCacheConfig::getDiskCacheSize() const +{ + return mDiskCacheSize; +} + +std::string const& KvCacheConfig::getDiskCachePath() const +{ + return mDiskCachePath; +} + +void KvCacheConfig::setDiskCacheSize(std::optional diskCacheSize) +{ + mDiskCacheSize = diskCacheSize; +} + +void KvCacheConfig::setDiskCachePath(std::string const& diskCachePath) +{ + mDiskCachePath = diskCachePath; +} + 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 020306e03e56..93437724f517 100644 --- a/cpp/tensorrt_llm/executor/serialization.cpp +++ b/cpp/tensorrt_llm/executor/serialization.cpp @@ -1353,10 +1353,15 @@ KvCacheConfig Serialization::deserializeKvCacheConfig(std::istream& is) auto eventBufferMaxSize = su::deserialize(is); auto useUvm = su::deserialize(is); auto attentionDpEventsGatherPeriodMs = su::deserialize(is); + auto diskCacheSize = su::deserialize>(is); + auto diskCachePath = su::deserialize(is); - return KvCacheConfig{enableBlockReuse, maxTokens, maxAttentionWindowVec, sinkTokenLength, freeGpuMemoryFraction, - hostCacheSize, crossKvCacheFraction, secondaryOffloadMinPriority, eventBufferMaxSize, enablePartialReuse, - copyOnPartialReuse, useUvm, attentionDpEventsGatherPeriodMs}; + auto kvCacheConfig = KvCacheConfig{enableBlockReuse, maxTokens, maxAttentionWindowVec, sinkTokenLength, + freeGpuMemoryFraction, hostCacheSize, crossKvCacheFraction, secondaryOffloadMinPriority, eventBufferMaxSize, + enablePartialReuse, copyOnPartialReuse, useUvm, attentionDpEventsGatherPeriodMs}; + kvCacheConfig.setDiskCacheSize(diskCacheSize); + kvCacheConfig.setDiskCachePath(diskCachePath); + return kvCacheConfig; } void Serialization::serialize(KvCacheConfig const& kvCacheConfig, std::ostream& os) @@ -1374,6 +1379,8 @@ void Serialization::serialize(KvCacheConfig const& kvCacheConfig, std::ostream& su::serialize(kvCacheConfig.getEventBufferMaxSize(), os); su::serialize(kvCacheConfig.getUseUvm(), os); su::serialize(kvCacheConfig.getAttentionDpEventsGatherPeriodMs(), os); + su::serialize(kvCacheConfig.getDiskCacheSize(), os); + su::serialize(kvCacheConfig.getDiskCachePath(), os); } size_t Serialization::serializedSize(KvCacheConfig const& kvCacheConfig) @@ -1393,6 +1400,8 @@ size_t Serialization::serializedSize(KvCacheConfig const& kvCacheConfig) totalSize += su::serializedSize(kvCacheConfig.getEventBufferMaxSize()); totalSize += su::serializedSize(kvCacheConfig.getUseUvm()); totalSize += su::serializedSize(kvCacheConfig.getAttentionDpEventsGatherPeriodMs()); + totalSize += su::serializedSize(kvCacheConfig.getDiskCacheSize()); + totalSize += su::serializedSize(kvCacheConfig.getDiskCachePath()); return totalSize; } diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp index 1496ebbb883d..35cc4f1d81ee 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp @@ -660,22 +660,22 @@ void tb::kv_cache_manager::KVCacheManagerBindings::initBindings(nb::module_& m) nb::class_(m, "KVCacheManager") .def(nb::init const&, SizeType32, SizeType32, - std::map> const&, SizeType32, SizeType32, - std::vector const&, nvinfer1::DataType, SizeType32, int64_t, SizeType32, SizeType32, bool, - tbk::CacheType, std::optional, + std::map> const&, SizeType32, std::string const&, + SizeType32, SizeType32, 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, 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"), - nb::arg("max_sequence_length"), nb::arg("chunk_size"), 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::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::arg("blocks_per_window"), nb::arg("blocks_in_disk_pool"), nb::arg("disk_cache_path"), + 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"), nb::arg("max_sequence_length"), + nb::arg("chunk_size"), 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::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::arg("pool_configurations") = std::vector{}, nb::call_guard()) .def( diff --git a/cpp/tensorrt_llm/nanobind/executor/executorConfig.cpp b/cpp/tensorrt_llm/nanobind/executor/executorConfig.cpp index 830f30ab9c67..fa9ca4471131 100644 --- a/cpp/tensorrt_llm/nanobind/executor/executorConfig.cpp +++ b/cpp/tensorrt_llm/nanobind/executor/executorConfig.cpp @@ -114,11 +114,12 @@ void initConfigBindings(nb::module_& m) self.getSinkTokenLength(), self.getFreeGpuMemoryFraction(), self.getHostCacheSize(), self.getCrossKvCacheFraction(), self.getSecondaryOffloadMinPriority(), self.getEventBufferMaxSize(), self.getEnablePartialReuse(), self.getCopyOnPartialReuse(), self.getUseUvm(), - self.getAttentionDpEventsGatherPeriodMs(), self.getMaxGpuTotalBytes()); + self.getAttentionDpEventsGatherPeriodMs(), self.getMaxGpuTotalBytes(), self.getDiskCacheSize(), + self.getDiskCachePath()); }; auto kvCacheConfigSetstate = [](tle::KvCacheConfig& self, nb::tuple const& state) { - if (state.size() != 14) + if (state.size() != 16) { throw std::runtime_error("Invalid state!"); } @@ -128,6 +129,8 @@ void initConfigBindings(nb::module_& m) 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])); + self.setDiskCacheSize(nb::cast>(state[14])); + self.setDiskCachePath(nb::cast(state[15])); }; nb::class_(m, "KvCacheConfig") .def(nb::init const&, std::optional> const&, @@ -152,6 +155,8 @@ 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("disk_cache_size", &tle::KvCacheConfig::getDiskCacheSize, &tle::KvCacheConfig::setDiskCacheSize) + .def_prop_rw("disk_cache_path", &tle::KvCacheConfig::getDiskCachePath, &tle::KvCacheConfig::setDiskCachePath) .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/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index 486c43e8e269..ffdaf1a1db70 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -612,6 +612,17 @@ def append_to_kv_heads_per_layer(num_kv_heads_per_layer: List[int], kwargs['event_manager'] = KVCacheEventManagerCpp( max_kv_event_entries=self.event_buffer_max_size) + disk_cache_size = getattr(kv_cache_config, "disk_cache_size", None) or 0 + blocks_in_disk_pool = 0 + if disk_cache_size: + max_tokens_disk = disk_cache_size // self.get_cache_bytes_per_token( + ) + blocks_in_disk_pool = int(max_tokens_disk // self.tokens_per_block) + logger.info(f"[disk-tier] {blocks_in_disk_pool} blocks " + f"({disk_cache_size / (1 << 30):.1f} GiB)") + kwargs['blocks_in_disk_pool'] = blocks_in_disk_pool + kwargs['disk_cache_path'] = getattr(kv_cache_config, "disk_cache_path", + None) or "" self.impl = KVCacheManagerCpp(**kwargs) # Warmup baseline for cumulative counters (set by snapshot_warmup_baseline) self._warmup_reused_blocks = 0 diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index fd3276163831..d0faa904964b 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -3541,6 +3541,9 @@ def _to_pybind(self): attention_dp_events_gather_period_ms=self. attention_dp_events_gather_period_ms, max_gpu_total_bytes=self.max_gpu_total_bytes) + if self.disk_cache_size: + config.disk_cache_size = self.disk_cache_size + config.disk_cache_path = self.disk_cache_path or "" return config @field_validator('free_gpu_memory_fraction') From 125aba4984a0e5861f23e14318fd0b0cc2e24182 Mon Sep 17 00:00:00 2001 From: Md Nafis Ul Haque Shifat Date: Fri, 3 Jul 2026 06:40:06 +0000 Subject: [PATCH 02/47] [#16251][feat] disk tier: retained-only gate, per-request disk retention TTL, serve wiring Signed-off-by: Md Nafis Ul Haque Shifat --- .../batch_manager/kvCacheManager.h | 72 +++++++++++---- cpp/include/tensorrt_llm/executor/executor.h | 13 ++- .../batch_manager/kvCacheManager.cpp | 91 +++++++++++++------ .../trtGptModelInflightBatching.cpp | 2 +- cpp/tensorrt_llm/executor/kvCacheConfig.cpp | 10 ++ .../executor/kvCacheRetentionConfig.cpp | 10 ++ cpp/tensorrt_llm/executor/serialization.cpp | 11 ++- .../nanobind/batch_manager/kvCacheManager.cpp | 19 ++-- .../nanobind/executor/executorConfig.cpp | 7 +- .../nanobind/executor/request.cpp | 7 +- .../_torch/pyexecutor/resource_manager.py | 2 + tensorrt_llm/llmapi/llm_args.py | 7 ++ tensorrt_llm/serve/openai_protocol.py | 10 ++ tensorrt_llm/serve/openai_server.py | 14 +++ 14 files changed, 214 insertions(+), 61 deletions(-) diff --git a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h index 833f739d3b25..d556dbadf194 100644 --- a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h +++ b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h @@ -464,6 +464,28 @@ class KVCacheBlock : public std::enable_shared_from_this std::swap(mDiskSlot, other->mDiskSlot); } + //! \brief Mark this block's content as retained for the disk tier until \p expiry. + //! Re-marking keeps the later deadline. + void markRetained(std::chrono::steady_clock::time_point::duration expiry) + { + mRetentionExpiry = mRetentionExpiry ? std::max(*mRetentionExpiry, expiry) : expiry; + } + + [[nodiscard]] bool isRetainedNow() const + { + return mRetentionExpiry.has_value() && *mRetentionExpiry > std::chrono::steady_clock::now().time_since_epoch(); + } + + [[nodiscard]] std::optional getRetentionExpiry() const + { + return mRetentionExpiry; + } + + void clearRetention() + { + mRetentionExpiry = std::nullopt; + } + void incRefCount(); void decRefCount(); @@ -621,6 +643,9 @@ class KVCacheBlock : public std::enable_shared_from_this // Disk cache tier: slot index in the disk file pool; kNoDiskSlot when not disk-resident. SizeType32 mDiskSlot{kNoDiskSlot}; + + // Disk-tier retention deadline (steady clock); nullopt = never marked. + std::optional mRetentionExpiry; }; class GenerationRequest @@ -769,6 +794,11 @@ class GenerationRequest return mKvCacheRetentionConfig.getTransferMode(); } + [[nodiscard]] std::optional getDiskRetentionMs() const + { + return mKvCacheRetentionConfig.getDiskRetentionMs(); + } + [[nodiscard]] std::string const& getDirectory() const { return mKvCacheRetentionConfig.getDirectory(); @@ -907,8 +937,8 @@ class WindowBlockManager std::vector const& managedLayers, std::vector const& numKvHeadsPerLayer, SizeType32 sizePerHead, SizeType32 tokensPerBlock, bool isSWA, SizeType32 blocksInPrimaryPool, SizeType32 blocksInSecondaryPool, SizeType32 blocksInDiskPool, std::string const& diskCachePath, - SizeType32 maxNumSequences, std::shared_ptr stream, CacheType cacheType, - std::optional secondaryOffloadMinPriority, + bool diskRetainedOnly, SizeType32 maxNumSequences, std::shared_ptr stream, + CacheType cacheType, std::optional secondaryOffloadMinPriority, std::shared_ptr eventManager, bool enablePartialReuse, bool copyOnPartialReuse, std::shared_ptr kvCacheConnectorManager, radix_block_tree::UnifiedBlockTree& lookupTree, std::shared_ptr loopbackAgent = nullptr, @@ -1398,7 +1428,9 @@ class WindowBlockManager SizeType32 mNumSecondaryBlocks; SizeType32 mNumDiskBlocks; std::string mDiskCachePath; + bool mDiskRetainedOnly{false}; std::size_t mDiskSpills{0}; + std::size_t mDiskGateDropped{0}; std::size_t mDiskOnboards{0}; // List of allocated blocks for each sequences @@ -1521,7 +1553,7 @@ class BlockManager //! across all windows. explicit BlockManager(std::vector const& numKvHeadsPerLayer, SizeType32 sizePerHead, SizeType32 tokensPerBlock, BlocksPerWindow const& blocksPerWindow, SizeType32 blocksInDiskPool, - std::string const& diskCachePath, SizeType32 maxNumSequences, CudaStreamPtr stream, + std::string const& diskCachePath, bool diskRetainedOnly, SizeType32 maxNumSequences, CudaStreamPtr stream, SizeType32 maxSequenceLength, SizeType32 maxBeamWidth, std::vector const& maxAttentionWindowVec, nvinfer1::DataType dtype, SizeType32 sinkBubbleLength, SizeType32 chunkSize, CacheType cacheType = CacheType::kSELF, @@ -2318,9 +2350,10 @@ class KVCacheManager : public BaseKVCacheManager //! and disagg transfer machinery applies natively. Empty vector = uniform. KVCacheManager(std::vector const& numKvHeadsPerLayer, SizeType32 sizePerHead, SizeType32 tokensPerBlock, BlocksPerWindow const& blocksPerWindow, SizeType32 blocksInDiskPool, std::string const& diskCachePath, - SizeType32 maxNumSequences, SizeType32 maxBeamWidth, std::vector const& maxAttentionWindowVec, - nvinfer1::DataType dtype, SizeType32 sinkTokenLength, CudaStreamPtr stream, SizeType32 maxSequenceLength, - SizeType32 chunkSize, bool enableBlockReuse = false, CacheType cacheType = CacheType::kSELF, + bool diskRetainedOnly, SizeType32 maxNumSequences, SizeType32 maxBeamWidth, + std::vector const& maxAttentionWindowVec, nvinfer1::DataType dtype, SizeType32 sinkTokenLength, + CudaStreamPtr stream, SizeType32 maxSequenceLength, SizeType32 chunkSize, bool enableBlockReuse = false, + CacheType cacheType = CacheType::kSELF, std::optional secondaryOffloadMinPriority = std::nullopt, std::shared_ptr eventManager = nullptr, bool enablePartialReuse = true, bool copyOnpartialReuse = true, @@ -2332,9 +2365,10 @@ class KVCacheManager : public BaseKVCacheManager KVCacheManager(std::vector const& numKvHeadsPerLayer, SizeType32 sizePerHead, SizeType32 tokensPerBlock, BlocksPerWindow const& blocksPerWindow, SizeType32 blocksInDiskPool, std::string const& diskCachePath, - SizeType32 maxNumSequences, SizeType32 maxBeamWidth, std::vector const& maxAttentionWindowVec, - nvinfer1::DataType dtype, SizeType32 sinkTokenLength, int64_t stream, SizeType32 maxSequenceLength, - SizeType32 chunkSize, bool enableBlockReuse = false, CacheType cacheType = CacheType::kSELF, + bool diskRetainedOnly, SizeType32 maxNumSequences, SizeType32 maxBeamWidth, + std::vector const& maxAttentionWindowVec, nvinfer1::DataType dtype, SizeType32 sinkTokenLength, + int64_t stream, SizeType32 maxSequenceLength, SizeType32 chunkSize, bool enableBlockReuse = false, + CacheType cacheType = CacheType::kSELF, std::optional secondaryOffloadMinPriority = std::nullopt, std::shared_ptr eventManager = nullptr, bool enablePartialReuse = true, bool copyOnpartialReuse = true, @@ -2346,9 +2380,10 @@ class KVCacheManager : public BaseKVCacheManager KVCacheManager(SizeType32 numLayers, SizeType32 numKvHeads, SizeType32 sizePerHead, SizeType32 tokensPerBlock, BlocksPerWindow const& blocksPerWindow, SizeType32 blocksInDiskPool, std::string const& diskCachePath, - SizeType32 maxNumSequences, SizeType32 maxBeamWidth, std::vector const& maxAttentionWindowVec, - nvinfer1::DataType dtype, SizeType32 sinkTokenLength, CudaStreamPtr stream, SizeType32 maxSequenceLength, - SizeType32 chunkSize, bool enableBlockReuse = true, CacheType cacheType = CacheType::kSELF, + bool diskRetainedOnly, SizeType32 maxNumSequences, SizeType32 maxBeamWidth, + std::vector const& maxAttentionWindowVec, nvinfer1::DataType dtype, SizeType32 sinkTokenLength, + CudaStreamPtr stream, SizeType32 maxSequenceLength, SizeType32 chunkSize, bool enableBlockReuse = true, + CacheType cacheType = CacheType::kSELF, std::optional secondaryOffloadMinPriority = std::nullopt, std::shared_ptr eventManager = nullptr, bool enablePartialReuse = true, bool copyOnpartialReuse = true, @@ -2360,12 +2395,13 @@ class KVCacheManager : public BaseKVCacheManager KVCacheManager(SizeType32 numLayers, SizeType32 numKvHeads, SizeType32 sizePerHead, SizeType32 tokensPerBlock, BlocksPerWindow const& blocksPerWindow, SizeType32 blocksInDiskPool, std::string const& diskCachePath, - SizeType32 maxNumSequences, SizeType32 maxBeamWidth, std::vector const& maxAttentionWindowVec, - nvinfer1::DataType dtype, SizeType32 sinkTokenLength, int64_t stream, SizeType32 maxSequenceLength, - SizeType32 chunkSize, bool enableBlockReuse = false, 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, + bool diskRetainedOnly, SizeType32 maxNumSequences, SizeType32 maxBeamWidth, + std::vector const& maxAttentionWindowVec, nvinfer1::DataType dtype, SizeType32 sinkTokenLength, + int64_t stream, SizeType32 maxSequenceLength, SizeType32 chunkSize, bool enableBlockReuse = false, + 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::vector const& poolConfigurations = {}); ~KVCacheManager() override = default; diff --git a/cpp/include/tensorrt_llm/executor/executor.h b/cpp/include/tensorrt_llm/executor/executor.h index 8af6b8742d54..ca50347dc626 100644 --- a/cpp/include/tensorrt_llm/executor/executor.h +++ b/cpp/include/tensorrt_llm/executor/executor.h @@ -633,6 +633,8 @@ class KvCacheRetentionConfig [[nodiscard]] std::optional getDecodeDurationMs() const; [[nodiscard]] KvCacheTransferMode getTransferMode() const; [[nodiscard]] std::string const& getDirectory() const; + [[nodiscard]] std::optional getDiskRetentionMs() const; + void setDiskRetentionMs(std::optional diskRetentionMs); /// @brief Convert the token range data into an entry per kv block. Returns a tuple of vectors corresponding to the /// priorities and durations for each block. @@ -644,7 +646,7 @@ class KvCacheRetentionConfig return mTokenRangeRetentionConfigs == other.mTokenRangeRetentionConfigs && mDecodeRetentionPriority == other.mDecodeRetentionPriority && mDecodeDurationMs == other.mDecodeDurationMs && mTransferMode == other.mTransferMode - && mDirectory == other.mDirectory; + && mDirectory == other.mDirectory && mDiskRetentionMs == other.mDiskRetentionMs; } private: @@ -660,6 +662,10 @@ class KvCacheRetentionConfig KvCacheTransferMode mTransferMode; /// @brief Name of the directory if transfer mode is GDS or POSIX_DEBUG_FALLBACK. std::string mDirectory; + + /// @brief How long this request's blocks stay eligible for (and protected in) the + /// disk cache tier. Unset = blocks never enter a retained-only disk tier. + std::optional mDiskRetentionMs; }; /// @brief A class that holds information about the request @@ -1101,6 +1107,8 @@ class KvCacheConfig [[nodiscard]] std::string const& getDiskCachePath() const; void setDiskCacheSize(std::optional diskCacheSize); void setDiskCachePath(std::string const& diskCachePath); + [[nodiscard]] bool getDiskCacheRetainedOnly() const; + void setDiskCacheRetainedOnly(bool diskCacheRetainedOnly); void fillEmptyFieldsFromRuntimeDefaults(tensorrt_llm::runtime::RuntimeDefaults const& runtimeDefaults); @@ -1166,6 +1174,9 @@ class KvCacheConfig /// @brief Directory holding the disk tier's block files. Must be non-empty when mDiskCacheSize > 0. std::string mDiskCachePath; + + /// @brief When true, only blocks whose request carried diskRetentionMs may enter the disk tier. + bool mDiskCacheRetainedOnly{false}; }; /// @brief Configuration class for the runtime perf knobs diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp index bf9292e36da4..7a9ab6a6b71c 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp @@ -585,9 +585,9 @@ std::map BlockManager::calculateWindowSizeToShare( BlockManager::BlockManager(std::vector const& numKvHeadsPerLayer, SizeType32 sizePerHead, SizeType32 tokensPerBlock, BlocksPerWindow const& blocksPerWindow, SizeType32 blocksInDiskPool, - std::string const& diskCachePath, SizeType32 maxNumSequences, CudaStreamPtr stream, SizeType32 maxSequenceLength, - SizeType32 maxBeamWidth, std::vector const& maxAttentionWindowVec, nvinfer1::DataType dtype, - SizeType32 sinkBubbleLength, SizeType32 chunkSize, CacheType cacheType, + std::string const& diskCachePath, bool diskRetainedOnly, SizeType32 maxNumSequences, CudaStreamPtr stream, + SizeType32 maxSequenceLength, SizeType32 maxBeamWidth, std::vector const& maxAttentionWindowVec, + nvinfer1::DataType dtype, SizeType32 sinkBubbleLength, SizeType32 chunkSize, CacheType cacheType, std::optional secondaryOffloadMinPriority, std::shared_ptr eventManager, bool enablePartialReuse, bool copyOnPartialReuse, std::shared_ptr kvCacheConnectorManager, @@ -694,8 +694,8 @@ BlockManager::BlockManager(std::vector const& numKvHeadsPerLayer, Si /*isSWA=*/(windowSize < maxSequenceLength) && (windowSize >= 0), allottedPrimaryBlocks, allottedSecondaryBlocks, (SizeType32(windowSize) == blocksPerWindow.rbegin()->first) ? blocksInDiskPool : SizeType32{0}, - diskCachePath, maxNumSequences, stream, cacheType, secondaryOffloadMinPriority, mEventManager, - enablePartialReuse, copyOnPartialReuse, kvCacheConnectorManager, mLookupTree, mLoopbackAgent, + diskCachePath, diskRetainedOnly, maxNumSequences, stream, cacheType, secondaryOffloadMinPriority, + mEventManager, enablePartialReuse, copyOnPartialReuse, kvCacheConnectorManager, mLookupTree, mLoopbackAgent, enableIndexerKCache, indexerKCacheQuantBlockSize, indexerKCacheIndexHeadDim, indexerKCacheUseFp4, LinearAttentionMetadata::hasLinearCache(windowSize) ? linearAttentionMetadata : std::nullopt, numPlaceholderBlocks); @@ -747,7 +747,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 blocksInDiskPool, std::string const& diskCachePath, - SizeType32 maxNumSequences, std::shared_ptr stream, CacheType cacheType, + bool diskRetainedOnly, SizeType32 maxNumSequences, std::shared_ptr stream, CacheType cacheType, std::optional secondaryOffloadMinPriority, std::shared_ptr eventManager, bool enablePartialReuse, bool copyOnPartialReuse, std::shared_ptr kvCacheConnectorManager, @@ -761,6 +761,7 @@ WindowBlockManager::WindowBlockManager(nvinfer1::DataType dtype, SizeType32 wind , mNumSecondaryBlocks{blocksInSecondaryPool} , mNumDiskBlocks{blocksInDiskPool} , mDiskCachePath{diskCachePath} + , mDiskRetainedOnly{diskRetainedOnly} , mBufferManager{std::move(stream)} , mSchedulingNumFreeBlocks{0} , mTokensPerBlock{tokensPerBlock} @@ -1290,6 +1291,22 @@ void WindowBlockManager::releaseSubtree(BlockPtr const& block) } } +namespace +{ +// Disk-level eviction order only: unmarked/expired evict first (bucket 0), marked +// blocks by minutes-to-expiry capped at 100. Restored to default on disk exit. +executor::RetentionPriority diskEvictionBucket(BlockPtr const& block) +{ + if (!block->isRetainedNow()) + { + return executor::KvCacheRetentionConfig::kMinRetentionPriority; + } + auto const remaining = *block->getRetentionExpiry() - std::chrono::steady_clock::now().time_since_epoch(); + auto const minutes = std::chrono::duration_cast(remaining).count(); + return static_cast(std::clamp(minutes, 1, 100)); +} +} // namespace + BlockPtr WindowBlockManager::reclaimSecondaryBlock() { auto victim = std::get<0>(mEvictionPolicy->getFreeBlock(kSecondaryLevel)); @@ -1298,16 +1315,28 @@ BlockPtr WindowBlockManager::reclaimSecondaryBlock() { return victim; // disk tier off, nothing reusable to keep, or disk full: same as before } + if (mDiskRetainedOnly && !victim->isRetainedNow()) + { + ++mDiskGateDropped; + if (mDiskGateDropped == 1 || mDiskGateDropped % 10000 == 0) + { + TLLM_LOG_INFO("[disk-tier] gate dropped=%zu (windowSize=%d)", mDiskGateDropped, mWindowSize); + } + return victim; // unmarked or expired: discard, exactly as stock + } auto diskTarget = std::get<0>(mEvictionPolicy->getFreeBlock(kDiskLevel)); - mEvictionPolicy->claimBlock(diskTarget); // claim BOTH before the swap (cf. #11879) + mEvictionPolicy->claimBlock(diskTarget); // claim BOTH before the swap (cf. #11879) mTransferManager->spillToFile(victim, diskTarget->getDiskSlot(), mPools, mDiskCachePath); - victim->swapDiskResidency(diskTarget); // victim's identity now disk-resident, tree intact - mEvictionPolicy->releaseBlock(victim); // re-enters the free queues at kDiskLevel + victim->swapDiskResidency(diskTarget); // victim's identity now disk-resident, tree intact + victim->setDurationMs(std::nullopt); // keep it out of the TTL-heap machinery + victim->setPriority(diskEvictionBucket(victim)); // disk ordering via the (unused) priority field + mEvictionPolicy->releaseBlock(victim); // re-enters the free queues at kDiskLevel ++mDiskSpills; if (mDiskSpills == 1 || mDiskSpills % 1000 == 0) { TLLM_LOG_INFO("[disk-tier] spills=%zu (windowSize=%d)", mDiskSpills, mWindowSize); } + diskTarget->clearRetention(); return diskTarget; // owns the victim's old host slot; already claimed } @@ -1376,9 +1405,14 @@ BlockPtr WindowBlockManager::getFreeBlock(GenerationRequest& sequence, executor: mEventManager->enqueueRemovedEvent(block, mWindowSize); } block->detachFromLookupNode(); + block->clearRetention(); } // Claim the block in primary block queue mEvictionPolicy->claimBlock(block, priority, durationMs); + if (auto const diskRetentionMs = sequence.getDiskRetentionMs()) + { + block->markRetained(std::chrono::steady_clock::now().time_since_epoch() + *diskRetentionMs); + } TLLM_LOG_DEBUG("%s::getFreeBlock - Block %d is now acquired by sequence %d", mLogPrefix.c_str(), block->getBlockId(), sequence.getRequestId()); @@ -1445,7 +1479,10 @@ void WindowBlockManager::onboardBlock(GenerationRequest& sequence, BlockPtr cons sequence, executor::KvCacheRetentionConfig::kDefaultRetentionPriority, std::nullopt, mode, directory); mTransferManager->loadFromFile(block, offloadBlock->getDiskSlot(), mPools, mDiskCachePath); offloadBlock->swapDiskResidency(block); // matched identity now GPU-resident - mEvictionPolicy->releaseBlock(block); // freed card returns to the disk free queue + offloadBlock->setPriority(executor::KvCacheRetentionConfig::kDefaultRetentionPriority); + block->clearRetention(); + block->setPriority(executor::KvCacheRetentionConfig::kDefaultRetentionPriority); + mEvictionPolicy->releaseBlock(block); // freed card returns to the disk free queue ++mDiskOnboards; if (mDiskOnboards == 1 || mDiskOnboards % 1000 == 0) { @@ -3285,7 +3322,7 @@ void WindowBlockManager::schedulingReleaseBlocks(RequestIdType requestId) KVCacheManager::KVCacheManager(SizeType32 numLayers, SizeType32 numKvHeads, SizeType32 sizePerHead, SizeType32 tokensPerBlock, BlocksPerWindow const& blocksPerWindow, SizeType32 blocksInDiskPool, - std::string const& diskCachePath, SizeType32 maxNumSequences, SizeType32 maxBeamWidth, + std::string const& diskCachePath, bool diskRetainedOnly, SizeType32 maxNumSequences, SizeType32 maxBeamWidth, std::vector const& maxAttentionWindowVec, nvinfer1::DataType dtype, SizeType32 sinkTokenLength, int64_t stream, runtime::SizeType32 maxSequenceLength, SizeType32 chunkSize, bool enableBlockReuse, CacheType cacheType, bool enablePartialReuse, bool copyOnPartialReuse, bool enableIndexerKCache, @@ -3293,17 +3330,17 @@ KVCacheManager::KVCacheManager(SizeType32 numLayers, SizeType32 numKvHeads, Size std::optional linearAttentionMetadata, std::vector const& poolConfigurations) : KVCacheManager(std::vector(numLayers, numKvHeads), sizePerHead, tokensPerBlock, blocksPerWindow, - blocksInDiskPool, diskCachePath, 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, poolConfigurations) + blocksInDiskPool, diskCachePath, diskRetainedOnly, 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, poolConfigurations) { } KVCacheManager::KVCacheManager(std::vector const& numKvHeadsPerLayer, SizeType32 sizePerHead, SizeType32 tokensPerBlock, BlocksPerWindow const& blocksPerWindow, SizeType32 blocksInDiskPool, - std::string const& diskCachePath, SizeType32 maxNumSequences, SizeType32 maxBeamWidth, + std::string const& diskCachePath, bool diskRetainedOnly, SizeType32 maxNumSequences, SizeType32 maxBeamWidth, std::vector const& maxAttentionWindowVec, nvinfer1::DataType dtype, SizeType32 sinkTokenLength, int64_t stream, runtime::SizeType32 maxSequenceLength, SizeType32 chunkSize, bool enableBlockReuse, CacheType cacheType, std::optional secondaryOffloadMinPriority, @@ -3313,7 +3350,7 @@ KVCacheManager::KVCacheManager(std::vector const& numKvHeadsPerLayer std::optional linearAttentionMetadata, std::vector const& poolConfigurations) : KVCacheManager(numKvHeadsPerLayer, sizePerHead, tokensPerBlock, blocksPerWindow, blocksInDiskPool, diskCachePath, - maxNumSequences, maxBeamWidth, maxAttentionWindowVec, dtype, sinkTokenLength, + diskRetainedOnly, maxNumSequences, maxBeamWidth, maxAttentionWindowVec, dtype, sinkTokenLength, std::make_shared(reinterpret_cast(stream)), maxSequenceLength, chunkSize, enableBlockReuse, cacheType, secondaryOffloadMinPriority, eventManager, enablePartialReuse, copyOnPartialReuse, kvCacheConnectorManager, enableIndexerKCache, indexerKCacheQuantBlockSize, indexerKCacheIndexHeadDim, @@ -3323,7 +3360,7 @@ KVCacheManager::KVCacheManager(std::vector const& numKvHeadsPerLayer KVCacheManager::KVCacheManager(std::vector const& numKvHeadsPerLayer, SizeType32 sizePerHead, SizeType32 tokensPerBlock, BlocksPerWindow const& blocksPerWindow, SizeType32 blocksInDiskPool, - std::string const& diskCachePath, SizeType32 maxNumSequences, SizeType32 maxBeamWidth, + std::string const& diskCachePath, bool diskRetainedOnly, SizeType32 maxNumSequences, SizeType32 maxBeamWidth, std::vector const& maxAttentionWindowVec, nvinfer1::DataType dtype, SizeType32 sinkTokenLength, CudaStreamPtr stream, runtime::SizeType32 maxSequenceLength, SizeType32 chunkSize, bool enableBlockReuse, CacheType cacheType, std::optional secondaryOffloadMinPriority, @@ -3340,8 +3377,8 @@ KVCacheManager::KVCacheManager(std::vector const& numKvHeadsPerLayer , mSinkBlockTokenLength(mSinkBubbleLength + sinkTokenLength) , mChunkSize(chunkSize) , mBlockManager(numKvHeadsPerLayer, sizePerHead, tokensPerBlock, blocksPerWindow, blocksInDiskPool, diskCachePath, - maxNumSequences, std::move(stream), maxSequenceLength, maxBeamWidth, maxAttentionWindowVec, dtype, - mSinkBubbleLength, mChunkSize, cacheType, secondaryOffloadMinPriority, std::move(eventManager), + diskRetainedOnly, maxNumSequences, 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, poolConfigurations) @@ -3366,7 +3403,7 @@ KVCacheManager::KVCacheManager(std::vector const& numKvHeadsPerLayer KVCacheManager::KVCacheManager(SizeType32 numLayers, SizeType32 numKvHeads, SizeType32 sizePerHead, SizeType32 tokensPerBlock, BlocksPerWindow const& blocksPerWindow, SizeType32 blocksInDiskPool, - std::string const& diskCachePath, SizeType32 maxNumSequences, SizeType32 maxBeamWidth, + std::string const& diskCachePath, bool diskRetainedOnly, SizeType32 maxNumSequences, SizeType32 maxBeamWidth, std::vector const& maxAttentionWindowVec, nvinfer1::DataType dtype, SizeType32 sinkTokenLength, CudaStreamPtr stream, runtime::SizeType32 maxSequenceLength, SizeType32 chunkSize, bool enableBlockReuse, CacheType cacheType, std::optional secondaryOffloadMinPriority, @@ -3376,11 +3413,11 @@ KVCacheManager::KVCacheManager(SizeType32 numLayers, SizeType32 numKvHeads, Size std::optional linearAttentionMetadata, std::vector const& poolConfigurations) : KVCacheManager(std::vector(numLayers, numKvHeads), sizePerHead, tokensPerBlock, blocksPerWindow, - blocksInDiskPool, diskCachePath, 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, poolConfigurations) + blocksInDiskPool, diskCachePath, diskRetainedOnly, 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, poolConfigurations) { } diff --git a/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp b/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp index 88cb761b6158..909a35514788 100644 --- a/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp +++ b/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp @@ -685,7 +685,7 @@ std::unique_ptr TrtGptModelInflightBatching::c auto const enableBlockReuse = kvCacheType == KvCacheType::kSELF ? kvCacheConfig.getEnableBlockReuse() : false; auto kvCacheManager = std::make_unique(numKvHeadsPerLayer, sizePerHead, tokensPerBlock, - blocksPerWindow, /*blocksInDiskPool=*/0, /*diskCachePath=*/std::string{}, getMaxNumSequences(), getMaxBeamWidth(), maxAttentionWindowVec, kvDtype, getSinkTokenLen(), + blocksPerWindow, /*blocksInDiskPool=*/0, /*diskCachePath=*/std::string{}, /*diskRetainedOnly=*/false, getMaxNumSequences(), getMaxBeamWidth(), maxAttentionWindowVec, kvDtype, getSinkTokenLen(), mRuntime->getStreamPtr(), kvCacheType == KvCacheType::kCROSS ? mModelConfig.getMaxEncoderLen() : getMaxSequenceLen(), getMaxNumTokens().value(), enableBlockReuse, kvCacheType, kvCacheConfig.getSecondaryOffloadMinPriority(), diff --git a/cpp/tensorrt_llm/executor/kvCacheConfig.cpp b/cpp/tensorrt_llm/executor/kvCacheConfig.cpp index cd2feb979a8c..ef01308993cb 100644 --- a/cpp/tensorrt_llm/executor/kvCacheConfig.cpp +++ b/cpp/tensorrt_llm/executor/kvCacheConfig.cpp @@ -222,6 +222,16 @@ void KvCacheConfig::setDiskCachePath(std::string const& diskCachePath) mDiskCachePath = diskCachePath; } +bool KvCacheConfig::getDiskCacheRetainedOnly() const +{ + return mDiskCacheRetainedOnly; +} + +void KvCacheConfig::setDiskCacheRetainedOnly(bool diskCacheRetainedOnly) +{ + mDiskCacheRetainedOnly = diskCacheRetainedOnly; +} + void KvCacheConfig::setSecondaryOffloadMinPriority(std::optional secondaryOffloadMinPriority) { mSecondaryOffloadMinPriority = secondaryOffloadMinPriority; diff --git a/cpp/tensorrt_llm/executor/kvCacheRetentionConfig.cpp b/cpp/tensorrt_llm/executor/kvCacheRetentionConfig.cpp index 1de31022391c..71a2b3d6ba56 100644 --- a/cpp/tensorrt_llm/executor/kvCacheRetentionConfig.cpp +++ b/cpp/tensorrt_llm/executor/kvCacheRetentionConfig.cpp @@ -171,4 +171,14 @@ std::vector KvCacheRetentionConfig::getPerBlockRet return perBlockRetentions; } +std::optional KvCacheRetentionConfig::getDiskRetentionMs() const +{ + return mDiskRetentionMs; +} + +void KvCacheRetentionConfig::setDiskRetentionMs(std::optional diskRetentionMs) +{ + mDiskRetentionMs = diskRetentionMs; +} + } // namespace tensorrt_llm::executor diff --git a/cpp/tensorrt_llm/executor/serialization.cpp b/cpp/tensorrt_llm/executor/serialization.cpp index 93437724f517..ceef9c91592b 100644 --- a/cpp/tensorrt_llm/executor/serialization.cpp +++ b/cpp/tensorrt_llm/executor/serialization.cpp @@ -1355,12 +1355,14 @@ KvCacheConfig Serialization::deserializeKvCacheConfig(std::istream& is) auto attentionDpEventsGatherPeriodMs = su::deserialize(is); auto diskCacheSize = su::deserialize>(is); auto diskCachePath = su::deserialize(is); + // retained-only flag is read after construction below auto kvCacheConfig = KvCacheConfig{enableBlockReuse, maxTokens, maxAttentionWindowVec, sinkTokenLength, freeGpuMemoryFraction, hostCacheSize, crossKvCacheFraction, secondaryOffloadMinPriority, eventBufferMaxSize, enablePartialReuse, copyOnPartialReuse, useUvm, attentionDpEventsGatherPeriodMs}; kvCacheConfig.setDiskCacheSize(diskCacheSize); kvCacheConfig.setDiskCachePath(diskCachePath); + kvCacheConfig.setDiskCacheRetainedOnly(su::deserialize(is)); return kvCacheConfig; } @@ -1381,6 +1383,7 @@ void Serialization::serialize(KvCacheConfig const& kvCacheConfig, std::ostream& su::serialize(kvCacheConfig.getAttentionDpEventsGatherPeriodMs(), os); su::serialize(kvCacheConfig.getDiskCacheSize(), os); su::serialize(kvCacheConfig.getDiskCachePath(), os); + su::serialize(kvCacheConfig.getDiskCacheRetainedOnly(), os); } size_t Serialization::serializedSize(KvCacheConfig const& kvCacheConfig) @@ -1402,6 +1405,7 @@ size_t Serialization::serializedSize(KvCacheConfig const& kvCacheConfig) totalSize += su::serializedSize(kvCacheConfig.getAttentionDpEventsGatherPeriodMs()); totalSize += su::serializedSize(kvCacheConfig.getDiskCacheSize()); totalSize += su::serializedSize(kvCacheConfig.getDiskCachePath()); + totalSize += su::serializedSize(kvCacheConfig.getDiskCacheRetainedOnly()); return totalSize; } @@ -1800,8 +1804,11 @@ KvCacheRetentionConfig Serialization::deserializeKvCacheRetentionConfig(std::ist auto transferMode = su::deserialize(is); auto directory = su::deserialize(is); - return KvCacheRetentionConfig{ + auto diskRetentionMs = intToDuration(su::deserialize>(is)); + auto kvCacheRetentionConfig = KvCacheRetentionConfig{ tokenRangeRetentionPriorities, decodePriority, decodeDurationMs, transferMode, directory}; + kvCacheRetentionConfig.setDiskRetentionMs(diskRetentionMs); + return kvCacheRetentionConfig; } void Serialization::serialize(KvCacheRetentionConfig const& kvCacheRetentionConfig, std::ostream& os) @@ -1811,6 +1818,7 @@ void Serialization::serialize(KvCacheRetentionConfig const& kvCacheRetentionConf su::serialize(durationToInt(kvCacheRetentionConfig.getDecodeDurationMs()), os); su::serialize(kvCacheRetentionConfig.getTransferMode(), os); su::serialize(kvCacheRetentionConfig.getDirectory(), os); + su::serialize(durationToInt(kvCacheRetentionConfig.getDiskRetentionMs()), os); } size_t Serialization::serializedSize(KvCacheRetentionConfig const& config) @@ -1821,6 +1829,7 @@ size_t Serialization::serializedSize(KvCacheRetentionConfig const& config) totalSize += su::serializedSize(durationToInt(config.getDecodeDurationMs())); totalSize += su::serializedSize(config.getTransferMode()); totalSize += su::serializedSize(config.getDirectory()); + totalSize += su::serializedSize(durationToInt(config.getDiskRetentionMs())); return totalSize; } diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp index 35cc4f1d81ee..c08f3f752021 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp @@ -660,7 +660,7 @@ void tb::kv_cache_manager::KVCacheManagerBindings::initBindings(nb::module_& m) nb::class_(m, "KVCacheManager") .def(nb::init const&, SizeType32, SizeType32, - std::map> const&, SizeType32, std::string const&, + std::map> const&, SizeType32, std::string const&, bool, SizeType32, SizeType32, std::vector const&, nvinfer1::DataType, SizeType32, int64_t, SizeType32, SizeType32, bool, tbk::CacheType, std::optional, std::shared_ptr, bool, bool, std::shared_ptr, @@ -668,14 +668,15 @@ void tb::kv_cache_manager::KVCacheManagerBindings::initBindings(nb::module_& m) 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("blocks_in_disk_pool"), nb::arg("disk_cache_path"), - 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"), nb::arg("max_sequence_length"), - nb::arg("chunk_size"), 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::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::arg("disk_cache_retained_only"), 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"), + nb::arg("max_sequence_length"), nb::arg("chunk_size"), 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::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::arg("pool_configurations") = std::vector{}, nb::call_guard()) .def( diff --git a/cpp/tensorrt_llm/nanobind/executor/executorConfig.cpp b/cpp/tensorrt_llm/nanobind/executor/executorConfig.cpp index fa9ca4471131..9e2529fec5cc 100644 --- a/cpp/tensorrt_llm/nanobind/executor/executorConfig.cpp +++ b/cpp/tensorrt_llm/nanobind/executor/executorConfig.cpp @@ -115,11 +115,11 @@ void initConfigBindings(nb::module_& m) self.getCrossKvCacheFraction(), self.getSecondaryOffloadMinPriority(), self.getEventBufferMaxSize(), self.getEnablePartialReuse(), self.getCopyOnPartialReuse(), self.getUseUvm(), self.getAttentionDpEventsGatherPeriodMs(), self.getMaxGpuTotalBytes(), self.getDiskCacheSize(), - self.getDiskCachePath()); + self.getDiskCachePath(), self.getDiskCacheRetainedOnly()); }; auto kvCacheConfigSetstate = [](tle::KvCacheConfig& self, nb::tuple const& state) { - if (state.size() != 16) + if (state.size() != 17) { throw std::runtime_error("Invalid state!"); } @@ -131,6 +131,7 @@ void initConfigBindings(nb::module_& m) nb::cast(state[12]), std::nullopt, nb::cast(state[13])); self.setDiskCacheSize(nb::cast>(state[14])); self.setDiskCachePath(nb::cast(state[15])); + self.setDiskCacheRetainedOnly(nb::cast(state[16])); }; nb::class_(m, "KvCacheConfig") .def(nb::init const&, std::optional> const&, @@ -157,6 +158,8 @@ void initConfigBindings(nb::module_& m) .def_prop_rw("host_cache_size", &tle::KvCacheConfig::getHostCacheSize, &tle::KvCacheConfig::setHostCacheSize) .def_prop_rw("disk_cache_size", &tle::KvCacheConfig::getDiskCacheSize, &tle::KvCacheConfig::setDiskCacheSize) .def_prop_rw("disk_cache_path", &tle::KvCacheConfig::getDiskCachePath, &tle::KvCacheConfig::setDiskCachePath) + .def_prop_rw("disk_cache_retained_only", &tle::KvCacheConfig::getDiskCacheRetainedOnly, + &tle::KvCacheConfig::setDiskCacheRetainedOnly) .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/nanobind/executor/request.cpp b/cpp/tensorrt_llm/nanobind/executor/request.cpp index a4effbd2e942..d614e73449fd 100644 --- a/cpp/tensorrt_llm/nanobind/executor/request.cpp +++ b/cpp/tensorrt_llm/nanobind/executor/request.cpp @@ -409,12 +409,12 @@ void initRequestBindings(nb::module_& m) auto kvCacheRetentionConfigGetstate = [](tle::KvCacheRetentionConfig const& self) { return nb::make_tuple(self.getTokenRangeRetentionConfigs(), self.getDecodeRetentionPriority(), - self.getDecodeDurationMs(), self.getTransferMode(), self.getDirectory()); + self.getDecodeDurationMs(), self.getTransferMode(), self.getDirectory(), self.getDiskRetentionMs()); }; auto kvCacheRetentionConfigSetstate = [](tle::KvCacheRetentionConfig& kvCacheRetentionConfig, nb::tuple const& state) { - if (state.size() != 5) + if (state.size() != 6) { throw std::runtime_error("Invalid state!"); } @@ -422,6 +422,7 @@ void initRequestBindings(nb::module_& m) nb::cast>(state[0]), nb::cast(state[1]), nb::cast>(state[2]), nb::cast(state[3]), nb::cast(state[4])); + kvCacheRetentionConfig.setDiskRetentionMs(nb::cast>(state[5])); }; auto kvCacheRetentionConfig = nb::class_(m, "KvCacheRetentionConfig"); @@ -455,6 +456,8 @@ void initRequestBindings(nb::module_& m) .def_prop_ro("decode_duration_ms", &tle::KvCacheRetentionConfig::getDecodeDurationMs) .def_prop_ro("transfer_mode", &tle::KvCacheRetentionConfig::getTransferMode) .def_prop_ro("directory", &tle::KvCacheRetentionConfig::getDirectory) + .def_prop_rw("disk_retention_ms", &tle::KvCacheRetentionConfig::getDiskRetentionMs, + &tle::KvCacheRetentionConfig::setDiskRetentionMs) .def("__getstate__", kvCacheRetentionConfigGetstate) .def("__setstate__", kvCacheRetentionConfigSetstate) .def("__eq__", &tle::KvCacheRetentionConfig::operator==); diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index ffdaf1a1db70..633a90d08e5f 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -623,6 +623,8 @@ def append_to_kv_heads_per_layer(num_kv_heads_per_layer: List[int], kwargs['blocks_in_disk_pool'] = blocks_in_disk_pool kwargs['disk_cache_path'] = getattr(kv_cache_config, "disk_cache_path", None) or "" + kwargs['disk_cache_retained_only'] = bool( + getattr(kv_cache_config, "disk_cache_retained_only", False)) self.impl = KVCacheManagerCpp(**kwargs) # Warmup baseline for cumulative counters (set by snapshot_warmup_baseline) self._warmup_reused_blocks = 0 diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index d0faa904964b..2b1cff9eb1f2 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -3354,6 +3354,11 @@ class KvCacheConfig(StrictBaseModel, PybindMirror): description= "Directory used for disk KV cache files. Must be set when `disk_cache_size` is positive." ) + disk_cache_retained_only: bool = Field( + default=False, + description= + "When true, only blocks whose request carried a disk retention TTL may enter the disk cache tier." + ) cross_kv_cache_fraction: Optional[float] = Field( default=None, description= @@ -3544,6 +3549,8 @@ def _to_pybind(self): if self.disk_cache_size: config.disk_cache_size = self.disk_cache_size config.disk_cache_path = self.disk_cache_path or "" + config.disk_cache_retained_only = bool( + self.disk_cache_retained_only) return config @field_validator('free_gpu_memory_fraction') diff --git a/tensorrt_llm/serve/openai_protocol.py b/tensorrt_llm/serve/openai_protocol.py index 374de0d727f8..080896c04008 100644 --- a/tensorrt_llm/serve/openai_protocol.py +++ b/tensorrt_llm/serve/openai_protocol.py @@ -416,6 +416,11 @@ class CompletionRequest(OpenAIBaseModel): # doc: begin-completion-sampling-params use_beam_search: bool = False top_k: int = 0 + kv_cache_ttl_seconds: Optional[int] = Field( + default=None, + description= + "Keep this request's KV cache reusable for this many seconds (disk-tier retention)." + ) top_p_min: float = 0.0 min_p: float = 0.0 repetition_penalty: float = 1.0 @@ -762,6 +767,11 @@ class ChatCompletionRequest(OpenAIBaseModel): best_of: Optional[int] = None use_beam_search: bool = False top_k: int = 0 + kv_cache_ttl_seconds: Optional[int] = Field( + default=None, + description= + "Keep this request's KV cache reusable for this many seconds (disk-tier retention)." + ) top_p_min: float = 0.0 min_p: float = 0.0 repetition_penalty: float = 1.0 diff --git a/tensorrt_llm/serve/openai_server.py b/tensorrt_llm/serve/openai_server.py index e59587652419..032251c32820 100644 --- a/tensorrt_llm/serve/openai_server.py +++ b/tensorrt_llm/serve/openai_server.py @@ -1,6 +1,7 @@ #!/usr/bin/env python import asyncio import base64 +import datetime as _dt import functools import json import os @@ -45,6 +46,7 @@ from tensorrt_llm.llmapi.disagg_utils import (DisaggClusterConfig, MetadataServerConfig, ServerRole) from tensorrt_llm.llmapi.llm import RequestOutput +from tensorrt_llm.llmapi.llm_utils import KvCacheRetentionConfig as _KvRetention from tensorrt_llm.llmapi.thinking_budget import \ add_thinking_budget_logits_processor from tensorrt_llm.logger import logger @@ -1387,8 +1389,14 @@ async def chat_stream_generator( functools.partial(preprocess_fn, prompt, sampling_params, disaggregated_params)) + _kv_retention = None + _ttl_s = getattr(request, "kv_cache_ttl_seconds", None) + if _ttl_s: + _kv_retention = _KvRetention([]) + _kv_retention.disk_retention_ms = _dt.timedelta(seconds=_ttl_s) promise = self.generator.generate_async( inputs=generate_inputs, + kv_cache_retention_config=_kv_retention, sampling_params=sampling_params, _postproc_params=postproc_params if self.postproc_worker_enabled else None, @@ -1872,8 +1880,14 @@ async def create_streaming_generator(promise: RequestOutput, agent_hierarchy=request.agent_hierarchy) # Generate + _kv_retention = None + _ttl_s = getattr(request, "kv_cache_ttl_seconds", None) + if _ttl_s: + _kv_retention = _KvRetention([]) + _kv_retention.disk_retention_ms = _dt.timedelta(seconds=_ttl_s) promise = self.generator.generate_async( inputs=harmony_tokens, + kv_cache_retention_config=_kv_retention, sampling_params=sampling_params, _postproc_params=postproc_params if self.postproc_worker_enabled else None, From b301d7f84089b4f22299ddcafbc4b9374bd211ef Mon Sep 17 00:00:00 2001 From: Md Nafis Ul Haque Shifat Date: Fri, 3 Jul 2026 07:32:59 +0000 Subject: [PATCH 03/47] [#16251][feat] disk tier: re-stamp reused blocks, deadline-heap displacement (empties first, exact TTL order, suffix-first FIFO), retention wiring at all serve entry points Signed-off-by: Md Nafis Ul Haque Shifat --- .../batch_manager/evictionPolicy.h | 9 +++ .../batch_manager/kvCacheManager.h | 29 ++++++++ .../batch_manager/evictionPolicy.cpp | 5 ++ .../batch_manager/kvCacheManager.cpp | 72 ++++++++++++++----- ruff-legacy-baseline.json | 6 +- tensorrt_llm/serve/openai_server.py | 34 +++++---- 6 files changed, 119 insertions(+), 36 deletions(-) diff --git a/cpp/include/tensorrt_llm/batch_manager/evictionPolicy.h b/cpp/include/tensorrt_llm/batch_manager/evictionPolicy.h index 81bdffa02d96..1d0ed17a3611 100644 --- a/cpp/include/tensorrt_llm/batch_manager/evictionPolicy.h +++ b/cpp/include/tensorrt_llm/batch_manager/evictionPolicy.h @@ -46,6 +46,14 @@ class BaseEvictionPolicy virtual void releaseBlock(BlockPtr block, bool toFront) = 0; /// @brief Get the amount of free blocks in the primary memory pool virtual SizeType32 getNumFreeBlocks(SizeType32 cacheLevel) = 0; + + /// @brief True when the block currently sits in a free queue (evictable right now). + /// Conservative default preserves prior behavior for policies that do not track it. + [[nodiscard]] virtual bool isEnqueued(BlockPtr const& block) + { + return true; + } + /// @brief Claim a free block. Called when the cache manager allocates or reuses a new block virtual void claimBlock(BlockPtr block) = 0; virtual void claimBlock(BlockPtr block, std::optional priority, @@ -84,6 +92,7 @@ class LRUEvictionPolicy : public BaseEvictionPolicy void releaseBlock(BlockPtr block, bool toFront) override; SizeType32 getNumFreeBlocks(SizeType32 cacheLevel) override; + [[nodiscard]] bool isEnqueued(BlockPtr const& block) override; void claimBlock(BlockPtr block) override; void claimBlock(BlockPtr block, std::optional priority, diff --git a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h index d556dbadf194..30ce48ef5b65 100644 --- a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h +++ b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h @@ -43,6 +43,7 @@ #include #include #include +#include #include #include #include @@ -1433,6 +1434,34 @@ class WindowBlockManager std::size_t mDiskGateDropped{0}; std::size_t mDiskOnboards{0}; + //! Disk-tier displacement order among retained blocks: exact earliest deadline; + //! among equal deadlines, first-spilled first. Upstream releases and evicts chains + //! leaf-first, so arrival order at the disk is deepest-first and FIFO here equals + //! suffix-first within a chain (and LRU-like across requests). Entries are lazy: + //! validated against live block state on pop, stale ones discarded. + struct DiskDeadline + { + std::chrono::steady_clock::time_point::duration expiry; + std::uint64_t seq; + BlockPtr block; + + bool operator>(DiskDeadline const& other) const + { + if (expiry != other.expiry) + { + return expiry > other.expiry; + } + return seq > other.seq; // first-arrived pops first among equal deadlines + } + }; + + std::priority_queue, std::greater> mDiskDeadlines; + std::uint64_t mDiskSpillSeq{0}; + + //! \brief Pick and claim the disk block to overwrite: empty slots first, then + //! unmarked content, then the retained block with the earliest deadline. + [[nodiscard]] BlockPtr claimDiskTarget(); + // List of allocated blocks for each sequences std::unordered_map> mAllocatedBlocksPerSeq; diff --git a/cpp/tensorrt_llm/batch_manager/evictionPolicy.cpp b/cpp/tensorrt_llm/batch_manager/evictionPolicy.cpp index 446ecdd72daf..f3e5d2b2b2c0 100644 --- a/cpp/tensorrt_llm/batch_manager/evictionPolicy.cpp +++ b/cpp/tensorrt_llm/batch_manager/evictionPolicy.cpp @@ -167,6 +167,11 @@ std::tuple LRUEvictionPolicy::getFreeBlock(SizeType32 cacheLevel TLLM_THROW("No free block found. This shouldn't happen!"); } +bool LRUEvictionPolicy::isEnqueued(BlockPtr const& block) +{ + return mFreeBlockIterators[block->getBlockId()] != std::nullopt; +} + void LRUEvictionPolicy::releaseBlock(BlockPtr block) { releaseBlock(block, false); diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp index 7a9ab6a6b71c..a5273f76b23e 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp @@ -876,6 +876,9 @@ WindowBlockManager::WindowBlockManager(nvinfer1::DataType dtype, SizeType32 wind auto diskBlock = std::make_shared( blocksInPrimaryPool + blocksInSecondaryPool + blockId, tk::KVCacheIndex::nullIndex); diskBlock->setDiskSlot(blockId); + // Empty slots must be the first displacement choice: bucket 0, below any + // retained block. Freed cards are reset to this same bucket on disk exit. + diskBlock->setPriority(executor::KvCacheRetentionConfig::kMinRetentionPriority); mAllBlocksById.emplace_back(std::move(diskBlock)); } if (blocksInDiskPool > 0) @@ -1291,21 +1294,36 @@ void WindowBlockManager::releaseSubtree(BlockPtr const& block) } } -namespace -{ -// Disk-level eviction order only: unmarked/expired evict first (bucket 0), marked -// blocks by minutes-to-expiry capped at 100. Restored to default on disk exit. -executor::RetentionPriority diskEvictionBucket(BlockPtr const& block) +BlockPtr WindowBlockManager::claimDiskTarget() { - if (!block->isRetainedNow()) + // Cheap victims (empty slots; unmarked content in spill-all mode) live in bucket 0, + // so the policy's front candidate is only a retained block when nothing cheap exists. + auto candidate = std::get<0>(mEvictionPolicy->getFreeBlock(kDiskLevel)); + if (!candidate->isRetainedNow()) { - return executor::KvCacheRetentionConfig::kMinRetentionPriority; + mEvictionPolicy->claimBlock(candidate); + return candidate; } - auto const remaining = *block->getRetentionExpiry() - std::chrono::steady_clock::now().time_since_epoch(); - auto const minutes = std::chrono::duration_cast(remaining).count(); - return static_cast(std::clamp(minutes, 1, 100)); + // Every free disk block is retained: displace the earliest real deadline; among + // equal deadlines the first-spilled block goes first (leaf-first arrival makes this + // suffix-first within a chain). Entries are validated against live block state on + // pop; failures mean the entry describes the past (block left disk, was re-stamped, + // or is mid-transfer and about to leave) and are discarded. + while (!mDiskDeadlines.empty()) + { + auto const top = mDiskDeadlines.top(); + mDiskDeadlines.pop(); + if (top.block->isOnDisk() && !top.block->hasRefs() && mEvictionPolicy->isEnqueued(top.block) + && top.block->getRetentionExpiry() == std::optional{top.expiry}) + { + mEvictionPolicy->claimBlock(top.block); + return top.block; + } + } + // Heap drained by stale entries while a retained candidate exists: fall back. + mEvictionPolicy->claimBlock(candidate); + return candidate; } -} // namespace BlockPtr WindowBlockManager::reclaimSecondaryBlock() { @@ -1322,15 +1340,25 @@ BlockPtr WindowBlockManager::reclaimSecondaryBlock() { TLLM_LOG_INFO("[disk-tier] gate dropped=%zu (windowSize=%d)", mDiskGateDropped, mWindowSize); } - return victim; // unmarked or expired: discard, exactly as stock + return victim; // unmarked or expired: discard, exactly as stock } - auto diskTarget = std::get<0>(mEvictionPolicy->getFreeBlock(kDiskLevel)); - mEvictionPolicy->claimBlock(diskTarget); // claim BOTH before the swap (cf. #11879) + auto diskTarget = claimDiskTarget(); // empties -> unmarked -> earliest deadline; claimed mTransferManager->spillToFile(victim, diskTarget->getDiskSlot(), mPools, mDiskCachePath); - victim->swapDiskResidency(diskTarget); // victim's identity now disk-resident, tree intact - victim->setDurationMs(std::nullopt); // keep it out of the TTL-heap machinery - victim->setPriority(diskEvictionBucket(victim)); // disk ordering via the (unused) priority field - mEvictionPolicy->releaseBlock(victim); // re-enters the free queues at kDiskLevel + victim->swapDiskResidency(diskTarget); // victim's identity now disk-resident, tree intact + victim->setDurationMs(std::nullopt); // keep it out of upstream's expiring-block machinery + if (victim->isRetainedNow()) + { + // Parked at max priority: displacement order among retained blocks is decided by + // mDiskDeadlines (exact deadline, suffix-first within a chain), not queue order. + victim->setPriority(executor::KvCacheRetentionConfig::kMaxRetentionPriority); + mDiskDeadlines.push({*victim->getRetentionExpiry(), ++mDiskSpillSeq, victim}); + } + else + { + // Spill-all mode only: unmarked content is displaced before any retained block. + victim->setPriority(executor::KvCacheRetentionConfig::kMinRetentionPriority); + } + mEvictionPolicy->releaseBlock(victim); // re-enters the free queues at kDiskLevel ++mDiskSpills; if (mDiskSpills == 1 || mDiskSpills % 1000 == 0) { @@ -1481,7 +1509,7 @@ void WindowBlockManager::onboardBlock(GenerationRequest& sequence, BlockPtr cons offloadBlock->swapDiskResidency(block); // matched identity now GPU-resident offloadBlock->setPriority(executor::KvCacheRetentionConfig::kDefaultRetentionPriority); block->clearRetention(); - block->setPriority(executor::KvCacheRetentionConfig::kDefaultRetentionPriority); + block->setPriority(executor::KvCacheRetentionConfig::kMinRetentionPriority); // freed card = empty slot again mEvictionPolicy->releaseBlock(block); // freed card returns to the disk free queue ++mDiskOnboards; if (mDiskOnboards == 1 || mDiskOnboards % 1000 == 0) @@ -1960,6 +1988,12 @@ SizeType32 WindowBlockManager::onboardAndAllocateBlocks( sequence.getRequestId(), matchingBlockId); } + if (auto const diskRetentionMs = sequence.getDiskRetentionMs()) + { + // Reused blocks bypass getFreeBlock, so the disk-retention stamp must also + // happen here; markRetained max-merges, keeping the later deadline. + claimed.block->markRetained(std::chrono::steady_clock::now().time_since_epoch() + *diskRetentionMs); + } onboardBlock(sequence, claimed.block, claimResult.mode, claimResult.directory); addBlockToAllBeams(claimed.block, sequence); if (!claimed.isPlaceholder) diff --git a/ruff-legacy-baseline.json b/ruff-legacy-baseline.json index ab0274cacc59..c1384735f6a6 100644 --- a/ruff-legacy-baseline.json +++ b/ruff-legacy-baseline.json @@ -1,7 +1,7 @@ { "_meta": { "generated_by": "scripts/legacy_utils.py lint-update-violations", - "total_violations": 5338, + "total_violations": 5336, "total_files": 492 }, ".github/scripts/label_community_user.py": { @@ -1572,9 +1572,7 @@ "F811": 1 }, "tensorrt_llm/serve/openai_server.py": { - "D205": 1, - "D212": 2, - "F821": 2 + "F821": 3 }, "tensorrt_llm/serve/responses_utils.py": { "D200": 2, diff --git a/tensorrt_llm/serve/openai_server.py b/tensorrt_llm/serve/openai_server.py index 032251c32820..5061a53dad7c 100644 --- a/tensorrt_llm/serve/openai_server.py +++ b/tensorrt_llm/serve/openai_server.py @@ -99,6 +99,20 @@ from .harmony_adapter import (HarmonyAdapter, get_harmony_adapter, maybe_transform_reasoning_effort) + +def _disk_retention_config(request): + """Build a retention config from the request's kv_cache_ttl_seconds. + + Carries only disk_retention_ms; returns None (the stock path) when the field + is absent or falsy. + """ + ttl_s = getattr(request, "kv_cache_ttl_seconds", None) + if not ttl_s: + return None + cfg = _KvRetention([]) + cfg.disk_retention_ms = _dt.timedelta(seconds=ttl_s) + return cfg + # yapf: enable TIMEOUT_KEEP_ALIVE = 5 # seconds. @@ -1389,14 +1403,9 @@ async def chat_stream_generator( functools.partial(preprocess_fn, prompt, sampling_params, disaggregated_params)) - _kv_retention = None - _ttl_s = getattr(request, "kv_cache_ttl_seconds", None) - if _ttl_s: - _kv_retention = _KvRetention([]) - _kv_retention.disk_retention_ms = _dt.timedelta(seconds=_ttl_s) promise = self.generator.generate_async( inputs=generate_inputs, - kv_cache_retention_config=_kv_retention, + kv_cache_retention_config=_disk_retention_config(request), sampling_params=sampling_params, _postproc_params=postproc_params if self.postproc_worker_enabled else None, @@ -1533,7 +1542,9 @@ async def create_mm_embedding_response(promise: RequestOutput): if mm_data is not None: prompt["multi_modal_data"] = mm_data - promise = self.generator.generate_async(inputs=prompt, ) + promise = self.generator.generate_async( + inputs=prompt, + kv_cache_retention_config=_disk_retention_config(request)) asyncio.create_task(self.await_disconnected(raw_request, promise)) response = await create_mm_embedding_response(promise) @@ -1724,6 +1735,7 @@ async def generator_wrapper(generator: AsyncIterator[Any]): promise = self.generator.generate_async( inputs=tokens_prompt, + kv_cache_retention_config=_disk_retention_config(request), sampling_params=sampling_params, _postproc_params=postproc_params, streaming=request.stream, @@ -1880,14 +1892,9 @@ async def create_streaming_generator(promise: RequestOutput, agent_hierarchy=request.agent_hierarchy) # Generate - _kv_retention = None - _ttl_s = getattr(request, "kv_cache_ttl_seconds", None) - if _ttl_s: - _kv_retention = _KvRetention([]) - _kv_retention.disk_retention_ms = _dt.timedelta(seconds=_ttl_s) promise = self.generator.generate_async( inputs=harmony_tokens, - kv_cache_retention_config=_kv_retention, + kv_cache_retention_config=_disk_retention_config(request), sampling_params=sampling_params, _postproc_params=postproc_params if self.postproc_worker_enabled else None, @@ -2030,6 +2037,7 @@ async def create_streaming_generator(promise: RequestOutput, ) promise = self.generator.generate_async( inputs=input_tokens, + kv_cache_retention_config=_disk_retention_config(request), sampling_params=sampling_params, streaming=request.stream, _postproc_params=postproc_params From fc8e892e5d902cd707c3545725fd79496eb9de77 Mon Sep 17 00:00:00 2001 From: Md Nafis Ul Haque Shifat Date: Fri, 3 Jul 2026 10:45:41 +0000 Subject: [PATCH 04/47] [#16251][feat] disk tier: trailing-default ctor params (non-breaking), unit tests (C++ disk-tier suite + isEnqueued + python plumbing), disk-count getters Signed-off-by: Md Nafis Ul Haque Shifat --- .../batch_manager/kvCacheManager.h | 86 +++++++++++++---- .../batch_manager/evictionPolicy.cpp | 14 ++- .../batch_manager/kvCacheManager.cpp | 95 ++++++++++--------- .../trtGptModelInflightBatching.cpp | 2 +- .../nanobind/batch_manager/kvCacheManager.cpp | 14 +-- tests/unittest/llmapi/test_disk_tier_args.py | 79 +++++++++++++++ 6 files changed, 212 insertions(+), 78 deletions(-) create mode 100644 tests/unittest/llmapi/test_disk_tier_args.py diff --git a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h index 30ce48ef5b65..be8f64fa1138 100644 --- a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h +++ b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h @@ -1434,6 +1434,23 @@ class WindowBlockManager std::size_t mDiskGateDropped{0}; std::size_t mDiskOnboards{0}; +public: + [[nodiscard]] std::size_t getNumDiskSpills() const + { + return mDiskSpills; + } + + [[nodiscard]] std::size_t getNumDiskGateDropped() const + { + return mDiskGateDropped; + } + + [[nodiscard]] std::size_t getNumDiskOnboards() const + { + return mDiskOnboards; + } + +private: //! Disk-tier displacement order among retained blocks: exact earliest deadline; //! among equal deadlines, first-spilled first. Upstream releases and evicts chains //! leaf-first, so arrival order at the disk is deepest-first and FIFO here equals @@ -1581,11 +1598,10 @@ class BlockManager //! 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 blocksInDiskPool, - std::string const& diskCachePath, bool diskRetainedOnly, SizeType32 maxNumSequences, CudaStreamPtr stream, - SizeType32 maxSequenceLength, SizeType32 maxBeamWidth, std::vector const& maxAttentionWindowVec, - nvinfer1::DataType dtype, SizeType32 sinkBubbleLength, SizeType32 chunkSize, - CacheType cacheType = CacheType::kSELF, + SizeType32 tokensPerBlock, BlocksPerWindow const& blocksPerWindow, SizeType32 maxNumSequences, + CudaStreamPtr stream, SizeType32 maxSequenceLength, SizeType32 maxBeamWidth, + std::vector const& maxAttentionWindowVec, nvinfer1::DataType dtype, SizeType32 sinkBubbleLength, + SizeType32 chunkSize, CacheType cacheType = CacheType::kSELF, std::optional secondaryOffloadMinPriority = std::nullopt, std::shared_ptr eventManager = nullptr, bool enablePartialReuse = true, bool copyOnPartialReuse = true, @@ -1593,7 +1609,8 @@ class BlockManager std::optional agentConfig = std::nullopt, bool enableIndexerKCache = false, SizeType32 indexerKCacheQuantBlockSize = 128, SizeType32 indexerKCacheIndexHeadDim = 0, bool indexerKCacheUseFp4 = false, std::optional linearAttentionMetadata = std::nullopt, - std::vector const& poolConfigurations = {}); + std::vector const& poolConfigurations = {}, SizeType32 blocksInDiskPool = 0, + std::string const& diskCachePath = {}, bool diskRetainedOnly = false); [[nodiscard]] bool isEnableIndexerKCache() const { @@ -2049,6 +2066,39 @@ class BlockManager // Stored before mWindowBlockManagers so it is constructed first and its address // is stable when passed to each WindowBlockManager constructor. radix_block_tree::UnifiedBlockTree mLookupTree; + +public: + [[nodiscard]] std::size_t getNumDiskSpills() const + { + std::size_t total = 0; + for (auto const& [windowSize, manager] : mWindowBlockManagers) + { + total += manager.getNumDiskSpills(); + } + return total; + } + + [[nodiscard]] std::size_t getNumDiskGateDropped() const + { + std::size_t total = 0; + for (auto const& [windowSize, manager] : mWindowBlockManagers) + { + total += manager.getNumDiskGateDropped(); + } + return total; + } + + [[nodiscard]] std::size_t getNumDiskOnboards() const + { + std::size_t total = 0; + for (auto const& [windowSize, manager] : mWindowBlockManagers) + { + total += manager.getNumDiskOnboards(); + } + return total; + } + +private: std::map mWindowBlockManagers; std::map mWindowSizeToMetadata; std::vector mLayerToWindowSize; @@ -2378,8 +2428,7 @@ class KVCacheManager : public BaseKVCacheManager //! + 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 blocksInDiskPool, std::string const& diskCachePath, - bool diskRetainedOnly, SizeType32 maxNumSequences, SizeType32 maxBeamWidth, + BlocksPerWindow const& blocksPerWindow, SizeType32 maxNumSequences, SizeType32 maxBeamWidth, std::vector const& maxAttentionWindowVec, nvinfer1::DataType dtype, SizeType32 sinkTokenLength, CudaStreamPtr stream, SizeType32 maxSequenceLength, SizeType32 chunkSize, bool enableBlockReuse = false, CacheType cacheType = CacheType::kSELF, @@ -2390,11 +2439,11 @@ class KVCacheManager : public BaseKVCacheManager bool enableIndexerKCache = false, SizeType32 indexerKCacheQuantBlockSize = 128, SizeType32 indexerKCacheIndexHeadDim = 0, bool indexerKCacheUseFp4 = false, std::optional linearAttentionMetadata = std::nullopt, - std::vector const& poolConfigurations = {}); + std::vector const& poolConfigurations = {}, SizeType32 blocksInDiskPool = 0, + std::string const& diskCachePath = {}, bool diskRetainedOnly = false); KVCacheManager(std::vector const& numKvHeadsPerLayer, SizeType32 sizePerHead, SizeType32 tokensPerBlock, - BlocksPerWindow const& blocksPerWindow, SizeType32 blocksInDiskPool, std::string const& diskCachePath, - bool diskRetainedOnly, SizeType32 maxNumSequences, SizeType32 maxBeamWidth, + BlocksPerWindow const& blocksPerWindow, SizeType32 maxNumSequences, SizeType32 maxBeamWidth, std::vector const& maxAttentionWindowVec, nvinfer1::DataType dtype, SizeType32 sinkTokenLength, int64_t stream, SizeType32 maxSequenceLength, SizeType32 chunkSize, bool enableBlockReuse = false, CacheType cacheType = CacheType::kSELF, @@ -2405,11 +2454,11 @@ class KVCacheManager : public BaseKVCacheManager bool enableIndexerKCache = false, SizeType32 indexerKCacheQuantBlockSize = 128, SizeType32 indexerKCacheIndexHeadDim = 0, bool indexerKCacheUseFp4 = false, std::optional linearAttentionMetadata = std::nullopt, - std::vector const& poolConfigurations = {}); + std::vector const& poolConfigurations = {}, SizeType32 blocksInDiskPool = 0, + std::string const& diskCachePath = {}, bool diskRetainedOnly = false); KVCacheManager(SizeType32 numLayers, SizeType32 numKvHeads, SizeType32 sizePerHead, SizeType32 tokensPerBlock, - BlocksPerWindow const& blocksPerWindow, SizeType32 blocksInDiskPool, std::string const& diskCachePath, - bool diskRetainedOnly, SizeType32 maxNumSequences, SizeType32 maxBeamWidth, + BlocksPerWindow const& blocksPerWindow, SizeType32 maxNumSequences, SizeType32 maxBeamWidth, std::vector const& maxAttentionWindowVec, nvinfer1::DataType dtype, SizeType32 sinkTokenLength, CudaStreamPtr stream, SizeType32 maxSequenceLength, SizeType32 chunkSize, bool enableBlockReuse = true, CacheType cacheType = CacheType::kSELF, @@ -2420,18 +2469,19 @@ class KVCacheManager : public BaseKVCacheManager bool enableIndexerKCache = false, SizeType32 indexerKCacheQuantBlockSize = 128, SizeType32 indexerKCacheIndexHeadDim = 0, bool indexerKCacheUseFp4 = false, std::optional linearAttentionMetadata = std::nullopt, - std::vector const& poolConfigurations = {}); + std::vector const& poolConfigurations = {}, SizeType32 blocksInDiskPool = 0, + std::string const& diskCachePath = {}, bool diskRetainedOnly = false); KVCacheManager(SizeType32 numLayers, SizeType32 numKvHeads, SizeType32 sizePerHead, SizeType32 tokensPerBlock, - BlocksPerWindow const& blocksPerWindow, SizeType32 blocksInDiskPool, std::string const& diskCachePath, - bool diskRetainedOnly, SizeType32 maxNumSequences, SizeType32 maxBeamWidth, + BlocksPerWindow const& blocksPerWindow, SizeType32 maxNumSequences, SizeType32 maxBeamWidth, std::vector const& maxAttentionWindowVec, nvinfer1::DataType dtype, SizeType32 sinkTokenLength, int64_t stream, SizeType32 maxSequenceLength, SizeType32 chunkSize, bool enableBlockReuse = false, 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::vector const& poolConfigurations = {}); + std::vector const& poolConfigurations = {}, SizeType32 blocksInDiskPool = 0, + std::string const& diskCachePath = {}, bool diskRetainedOnly = false); ~KVCacheManager() override = default; diff --git a/cpp/tensorrt_llm/batch_manager/evictionPolicy.cpp b/cpp/tensorrt_llm/batch_manager/evictionPolicy.cpp index f3e5d2b2b2c0..e58c05e280f3 100644 --- a/cpp/tensorrt_llm/batch_manager/evictionPolicy.cpp +++ b/cpp/tensorrt_llm/batch_manager/evictionPolicy.cpp @@ -63,8 +63,12 @@ constexpr auto defaultPriorityIdx = getPriorityIdx(kDefaultPriority); void LRUEvictionPolicy::initialize(std::vector& mAllBlocksById, std::vector sizes, std::optional secondaryOffloadMinPriority) { - TLLM_CHECK_WITH_INFO(static_cast(sizes.size()) == kNumCacheLevels, - "initialize expects one block count per cache level (%d), got %zu", kNumCacheLevels, sizes.size()); + TLLM_CHECK_WITH_INFO(static_cast(sizes.size()) <= kNumCacheLevels, + "initialize expects at most one block count per cache level (%d), got %zu", kNumCacheLevels, sizes.size()); + // Callers may pass fewer sizes than levels (e.g. legacy {primary, secondary}); missing + // trailing levels (disk) default to 0 blocks. + auto sizesPadded = sizes; + sizesPadded.resize(kNumCacheLevels, 0); SizeType32 startIdx = 0; // Create queues for all levels: primary, secondary, and placeholder (initially empty). @@ -76,15 +80,15 @@ void LRUEvictionPolicy::initialize(std::vector& mAllBlocksById, std::v { auto& freeQueue = mFreeQueues[cacheLevel][defaultPriorityIdx]; - for (SizeType32 blockId = 0; blockId < sizes[cacheLevel]; blockId++) + for (SizeType32 blockId = 0; blockId < sizesPadded[cacheLevel]; blockId++) { // Initialize all blocks to be the default priority level mFreeBlockIterators[startIdx + blockId] = freeQueue.insert(freeQueue.end(), mAllBlocksById[startIdx + blockId]); } - mNumFreeBlocksPerLevel[cacheLevel] = sizes[cacheLevel]; - startIdx += sizes[cacheLevel]; + mNumFreeBlocksPerLevel[cacheLevel] = sizesPadded[cacheLevel]; + startIdx += sizesPadded[cacheLevel]; } mSecondaryOffloadMinPriority = secondaryOffloadMinPriority.value_or(kDefaultSecondaryOffloadMinPriority); diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp index a5273f76b23e..c8941682d7ad 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp @@ -584,8 +584,7 @@ std::map BlockManager::calculateWindowSizeToShare( } BlockManager::BlockManager(std::vector const& numKvHeadsPerLayer, SizeType32 sizePerHead, - SizeType32 tokensPerBlock, BlocksPerWindow const& blocksPerWindow, SizeType32 blocksInDiskPool, - std::string const& diskCachePath, bool diskRetainedOnly, SizeType32 maxNumSequences, CudaStreamPtr stream, + SizeType32 tokensPerBlock, BlocksPerWindow const& blocksPerWindow, SizeType32 maxNumSequences, CudaStreamPtr stream, SizeType32 maxSequenceLength, SizeType32 maxBeamWidth, std::vector const& maxAttentionWindowVec, nvinfer1::DataType dtype, SizeType32 sinkBubbleLength, SizeType32 chunkSize, CacheType cacheType, std::optional secondaryOffloadMinPriority, @@ -594,7 +593,8 @@ BlockManager::BlockManager(std::vector const& numKvHeadsPerLayer, Si std::optional agentConfig, bool enableIndexerKCache, SizeType32 indexerKCacheQuantBlockSize, SizeType32 indexerKCacheIndexHeadDim, bool indexerKCacheUseFp4, std::optional linearAttentionMetadata, - std::vector const& poolConfigurations) + std::vector const& poolConfigurations, SizeType32 blocksInDiskPool, + std::string const& diskCachePath, bool diskRetainedOnly) : mNumLayers{static_cast(numKvHeadsPerLayer.size())} , mTokensPerBlock{tokensPerBlock} , mEventManager{std::move(eventManager)} @@ -3355,54 +3355,55 @@ void WindowBlockManager::schedulingReleaseBlocks(RequestIdType requestId) } KVCacheManager::KVCacheManager(SizeType32 numLayers, SizeType32 numKvHeads, SizeType32 sizePerHead, - SizeType32 tokensPerBlock, BlocksPerWindow const& blocksPerWindow, SizeType32 blocksInDiskPool, - std::string const& diskCachePath, bool diskRetainedOnly, SizeType32 maxNumSequences, SizeType32 maxBeamWidth, - std::vector const& maxAttentionWindowVec, nvinfer1::DataType dtype, 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, - std::vector const& poolConfigurations) + SizeType32 tokensPerBlock, BlocksPerWindow const& blocksPerWindow, SizeType32 maxNumSequences, + SizeType32 maxBeamWidth, std::vector const& maxAttentionWindowVec, nvinfer1::DataType dtype, + 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, + std::vector const& poolConfigurations, SizeType32 blocksInDiskPool, + std::string const& diskCachePath, bool diskRetainedOnly) : KVCacheManager(std::vector(numLayers, numKvHeads), sizePerHead, tokensPerBlock, blocksPerWindow, - blocksInDiskPool, diskCachePath, diskRetainedOnly, 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, poolConfigurations) + 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, poolConfigurations, blocksInDiskPool, diskCachePath, diskRetainedOnly) { } KVCacheManager::KVCacheManager(std::vector const& numKvHeadsPerLayer, SizeType32 sizePerHead, - SizeType32 tokensPerBlock, BlocksPerWindow const& blocksPerWindow, SizeType32 blocksInDiskPool, - std::string const& diskCachePath, bool diskRetainedOnly, SizeType32 maxNumSequences, SizeType32 maxBeamWidth, - std::vector const& maxAttentionWindowVec, nvinfer1::DataType dtype, SizeType32 sinkTokenLength, - int64_t stream, runtime::SizeType32 maxSequenceLength, SizeType32 chunkSize, bool enableBlockReuse, - CacheType cacheType, std::optional secondaryOffloadMinPriority, + SizeType32 tokensPerBlock, BlocksPerWindow const& blocksPerWindow, SizeType32 maxNumSequences, + SizeType32 maxBeamWidth, std::vector const& maxAttentionWindowVec, nvinfer1::DataType dtype, + SizeType32 sinkTokenLength, int64_t stream, runtime::SizeType32 maxSequenceLength, SizeType32 chunkSize, + bool enableBlockReuse, CacheType cacheType, std::optional secondaryOffloadMinPriority, std::shared_ptr eventManager, bool enablePartialReuse, bool copyOnPartialReuse, std::shared_ptr kvCacheConnectorManager, bool enableIndexerKCache, SizeType32 indexerKCacheQuantBlockSize, SizeType32 indexerKCacheIndexHeadDim, bool indexerKCacheUseFp4, std::optional linearAttentionMetadata, - std::vector const& poolConfigurations) - : KVCacheManager(numKvHeadsPerLayer, sizePerHead, tokensPerBlock, blocksPerWindow, blocksInDiskPool, diskCachePath, - diskRetainedOnly, maxNumSequences, maxBeamWidth, maxAttentionWindowVec, dtype, sinkTokenLength, + std::vector const& poolConfigurations, SizeType32 blocksInDiskPool, + std::string const& diskCachePath, bool diskRetainedOnly) + : 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, poolConfigurations) + indexerKCacheUseFp4, linearAttentionMetadata, poolConfigurations, blocksInDiskPool, diskCachePath, + diskRetainedOnly) { } KVCacheManager::KVCacheManager(std::vector const& numKvHeadsPerLayer, SizeType32 sizePerHead, - SizeType32 tokensPerBlock, BlocksPerWindow const& blocksPerWindow, SizeType32 blocksInDiskPool, - std::string const& diskCachePath, bool diskRetainedOnly, SizeType32 maxNumSequences, SizeType32 maxBeamWidth, - std::vector const& maxAttentionWindowVec, nvinfer1::DataType dtype, SizeType32 sinkTokenLength, - CudaStreamPtr stream, runtime::SizeType32 maxSequenceLength, SizeType32 chunkSize, bool enableBlockReuse, - CacheType cacheType, std::optional secondaryOffloadMinPriority, + SizeType32 tokensPerBlock, BlocksPerWindow const& blocksPerWindow, SizeType32 maxNumSequences, + SizeType32 maxBeamWidth, std::vector const& maxAttentionWindowVec, nvinfer1::DataType dtype, + SizeType32 sinkTokenLength, CudaStreamPtr stream, runtime::SizeType32 maxSequenceLength, SizeType32 chunkSize, + bool enableBlockReuse, CacheType cacheType, std::optional secondaryOffloadMinPriority, std::shared_ptr eventManager, bool enablePartialReuse, bool copyOnPartialReuse, std::shared_ptr kvCacheConnectorManager, bool enableIndexerKCache, SizeType32 indexerKCacheQuantBlockSize, SizeType32 indexerKCacheIndexHeadDim, bool indexerKCacheUseFp4, std::optional linearAttentionMetadata, - std::vector const& poolConfigurations) + std::vector const& poolConfigurations, SizeType32 blocksInDiskPool, + std::string const& diskCachePath, bool diskRetainedOnly) : mMaxBeamWidth(maxBeamWidth) , mDataType(dtype) , mMaxAttentionWindow(*std::max_element(maxAttentionWindowVec.begin(), maxAttentionWindowVec.end())) @@ -3410,12 +3411,12 @@ KVCacheManager::KVCacheManager(std::vector const& numKvHeadsPerLayer , mSinkBubbleLength(BaseKVCacheManager::getSinkBubbleLength(sinkTokenLength, tokensPerBlock)) , mSinkBlockTokenLength(mSinkBubbleLength + sinkTokenLength) , mChunkSize(chunkSize) - , mBlockManager(numKvHeadsPerLayer, sizePerHead, tokensPerBlock, blocksPerWindow, blocksInDiskPool, diskCachePath, - diskRetainedOnly, maxNumSequences, std::move(stream), maxSequenceLength, maxBeamWidth, maxAttentionWindowVec, - dtype, mSinkBubbleLength, mChunkSize, cacheType, secondaryOffloadMinPriority, std::move(eventManager), - enablePartialReuse, copyOnPartialReuse, std::move(kvCacheConnectorManager), std::nullopt, enableIndexerKCache, + , mBlockManager(numKvHeadsPerLayer, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, + 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, - poolConfigurations) + poolConfigurations, blocksInDiskPool, diskCachePath, diskRetainedOnly) // disable block reuse for sink bubble since chopVectorIntoBlocks does not match KV cache blocks in this case , mEnableBlockReuse{mSinkBubbleLength > 0 ? false : enableBlockReuse} { @@ -3436,22 +3437,22 @@ KVCacheManager::KVCacheManager(std::vector const& numKvHeadsPerLayer } KVCacheManager::KVCacheManager(SizeType32 numLayers, SizeType32 numKvHeads, SizeType32 sizePerHead, - SizeType32 tokensPerBlock, BlocksPerWindow const& blocksPerWindow, SizeType32 blocksInDiskPool, - std::string const& diskCachePath, bool diskRetainedOnly, SizeType32 maxNumSequences, SizeType32 maxBeamWidth, - std::vector const& maxAttentionWindowVec, nvinfer1::DataType dtype, SizeType32 sinkTokenLength, - CudaStreamPtr stream, runtime::SizeType32 maxSequenceLength, SizeType32 chunkSize, bool enableBlockReuse, - CacheType cacheType, std::optional secondaryOffloadMinPriority, + SizeType32 tokensPerBlock, BlocksPerWindow const& blocksPerWindow, SizeType32 maxNumSequences, + SizeType32 maxBeamWidth, std::vector const& maxAttentionWindowVec, nvinfer1::DataType dtype, + SizeType32 sinkTokenLength, CudaStreamPtr stream, runtime::SizeType32 maxSequenceLength, SizeType32 chunkSize, + bool enableBlockReuse, CacheType cacheType, std::optional secondaryOffloadMinPriority, std::shared_ptr eventManager, bool enablePartialReuse, bool copyOnPartialReuse, std::shared_ptr kvCacheConnectorManager, bool enableIndexerKCache, SizeType32 indexerKCacheQuantBlockSize, SizeType32 indexerKCacheIndexHeadDim, bool indexerKCacheUseFp4, std::optional linearAttentionMetadata, - std::vector const& poolConfigurations) + std::vector const& poolConfigurations, SizeType32 blocksInDiskPool, + std::string const& diskCachePath, bool diskRetainedOnly) : KVCacheManager(std::vector(numLayers, numKvHeads), sizePerHead, tokensPerBlock, blocksPerWindow, - blocksInDiskPool, diskCachePath, diskRetainedOnly, 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, poolConfigurations) + 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, + poolConfigurations, blocksInDiskPool, diskCachePath, diskRetainedOnly) { } diff --git a/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp b/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp index 909a35514788..7a0d78beb8a0 100644 --- a/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp +++ b/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp @@ -685,7 +685,7 @@ std::unique_ptr TrtGptModelInflightBatching::c auto const enableBlockReuse = kvCacheType == KvCacheType::kSELF ? kvCacheConfig.getEnableBlockReuse() : false; auto kvCacheManager = std::make_unique(numKvHeadsPerLayer, sizePerHead, tokensPerBlock, - blocksPerWindow, /*blocksInDiskPool=*/0, /*diskCachePath=*/std::string{}, /*diskRetainedOnly=*/false, getMaxNumSequences(), getMaxBeamWidth(), maxAttentionWindowVec, kvDtype, getSinkTokenLen(), + blocksPerWindow, getMaxNumSequences(), getMaxBeamWidth(), maxAttentionWindowVec, kvDtype, getSinkTokenLen(), mRuntime->getStreamPtr(), kvCacheType == KvCacheType::kCROSS ? mModelConfig.getMaxEncoderLen() : getMaxSequenceLen(), getMaxNumTokens().value(), enableBlockReuse, kvCacheType, kvCacheConfig.getSecondaryOffloadMinPriority(), diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp index c08f3f752021..17db4bdf71db 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp @@ -660,15 +660,14 @@ void tb::kv_cache_manager::KVCacheManagerBindings::initBindings(nb::module_& m) nb::class_(m, "KVCacheManager") .def(nb::init const&, SizeType32, SizeType32, - std::map> const&, SizeType32, std::string const&, bool, - SizeType32, SizeType32, std::vector const&, nvinfer1::DataType, SizeType32, int64_t, - SizeType32, SizeType32, bool, tbk::CacheType, std::optional, + std::map> const&, SizeType32, SizeType32, + 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, - std::vector const&>(), + std::vector const&, SizeType32, std::string const&, bool>(), nb::arg("num_kv_heads_per_layer"), nb::arg("size_per_head"), nb::arg("tokens_per_block"), - nb::arg("blocks_per_window"), nb::arg("blocks_in_disk_pool"), nb::arg("disk_cache_path"), - nb::arg("disk_cache_retained_only"), nb::arg("max_num_sequences"), nb::arg("max_beam_width"), + 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"), nb::arg("max_sequence_length"), nb::arg("chunk_size"), nb::arg("enable_block_reuse") = false, nb::arg("cache_type") = tbk::CacheType::kSELF, nb::arg("secondary_offload_min_priority") = std::nullopt, @@ -677,7 +676,8 @@ void tb::kv_cache_manager::KVCacheManagerBindings::initBindings(nb::module_& m) 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::arg("pool_configurations") = std::vector{}, + nb::arg("pool_configurations") = std::vector{}, nb::arg("blocks_in_disk_pool") = 0, + nb::arg("disk_cache_path") = "", nb::arg("disk_cache_retained_only") = false, nb::call_guard()) .def( "scheduling_has_free_blocks", diff --git a/tests/unittest/llmapi/test_disk_tier_args.py b/tests/unittest/llmapi/test_disk_tier_args.py new file mode 100644 index 000000000000..46a5a5375a52 --- /dev/null +++ b/tests/unittest/llmapi/test_disk_tier_args.py @@ -0,0 +1,79 @@ +# tests/unittest/llmapi/test_disk_tier_args.py +"""Disk-tier plumbing unit tests. + +pydantic->pybind passthrough, retention pickle, and the serve-layer TTL helper. +""" + +import datetime +import pickle + +from tensorrt_llm.bindings import executor as tllm + + +def test_kv_cache_config_disk_fields(): + cfg = tllm.KvCacheConfig() + assert cfg.disk_cache_retained_only is False # default preserves spill-all + cfg.disk_cache_size = 1 << 30 + cfg.disk_cache_path = "/tmp/kvdisk-test" + cfg.disk_cache_retained_only = True + assert cfg.disk_cache_size == 1 << 30 + assert cfg.disk_cache_path == "/tmp/kvdisk-test" + assert cfg.disk_cache_retained_only is True + + +def test_kv_cache_config_disk_fields_pickle(): + cfg = tllm.KvCacheConfig() + cfg.disk_cache_size = 42 + cfg.disk_cache_path = "/x" + cfg.disk_cache_retained_only = True + back = pickle.loads(pickle.dumps(cfg)) + assert back.disk_cache_size == 42 + assert back.disk_cache_path == "/x" + assert back.disk_cache_retained_only is True + + +def test_retention_config_disk_retention_ms_roundtrip(): + cfg = tllm.KvCacheRetentionConfig([]) + assert cfg.disk_retention_ms is None # unset = never enters a retained-only tier + cfg.disk_retention_ms = datetime.timedelta(seconds=3600) + assert cfg.disk_retention_ms == datetime.timedelta(seconds=3600) + back = pickle.loads(pickle.dumps(cfg)) + assert back.disk_retention_ms == datetime.timedelta(seconds=3600) + + +def test_llm_args_kv_cache_config_disk_fields(): + import os + + os.makedirs("/tmp/kvdisk-test", exist_ok=True) # pydantic validator requires an existing dir + from tensorrt_llm.llmapi.llm_args import KvCacheConfig as PydanticKvCacheConfig + + py_cfg = PydanticKvCacheConfig( + disk_cache_size=1 << 30, disk_cache_path="/tmp/kvdisk-test", disk_cache_retained_only=True + ) + assert py_cfg.disk_cache_retained_only is True + cpp_cfg = py_cfg._to_pybind() + assert cpp_cfg.disk_cache_size == 1 << 30 + assert cpp_cfg.disk_cache_path == "/tmp/kvdisk-test" + assert cpp_cfg.disk_cache_retained_only is True + + +def test_serve_disk_retention_helper(): + from tensorrt_llm.serve.openai_server import _disk_retention_config + + class Req: + kv_cache_ttl_seconds = None + + assert _disk_retention_config(Req()) is None + + Req.kv_cache_ttl_seconds = 0 + assert _disk_retention_config(Req()) is None # ttl=0 behaves non-retained + + Req.kv_cache_ttl_seconds = 3600 + cfg = _disk_retention_config(Req()) + assert cfg is not None + assert cfg.disk_retention_ms == datetime.timedelta(seconds=3600) + + class Bare: # object without the field at all (non-serve request types) + pass + + assert _disk_retention_config(Bare()) is None From a638e07d15567363f02a575d7fa2e474050119d7 Mon Sep 17 00:00:00 2001 From: Md Nafis Ul Haque Shifat Date: Fri, 3 Jul 2026 20:32:40 +0000 Subject: [PATCH 05/47] [#16251][test] disk tier: C++ unit test suite (DiskTier spill/gate/expiry/reuse/deadline + isEnqueued + placeholder level fix) Signed-off-by: Md Nafis Ul Haque Shifat --- .../batch_manager/evictionPolicyTest.cpp | 13 +- .../batch_manager/kvCacheManagerTest.cpp | 304 ++++++++++++++++++ 2 files changed, 316 insertions(+), 1 deletion(-) diff --git a/cpp/tests/unit_tests/batch_manager/evictionPolicyTest.cpp b/cpp/tests/unit_tests/batch_manager/evictionPolicyTest.cpp index 541f7917bf2a..45c52fe3e6b7 100644 --- a/cpp/tests/unit_tests/batch_manager/evictionPolicyTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/evictionPolicyTest.cpp @@ -75,6 +75,16 @@ class KvCacheRetentionConfigTest : public ::testing::Test void TearDown() {} }; +TEST_F(LRUPolicyTest, IsEnqueuedTest) +{ + auto block = std::get<0>(policy->getFreeBlock(0)); + EXPECT_TRUE(policy->isEnqueued(block)); + policy->claimBlock(block); + EXPECT_FALSE(policy->isEnqueued(block)); + policy->releaseBlock(block); + EXPECT_TRUE(policy->isEnqueued(block)); +} + TEST_F(LRUPolicyTest, NumFreeBlocksTest) { EXPECT_EQ(NUM_PRIMARY_BLOCKS, policy->getNumFreeBlocks(0)); @@ -120,7 +130,8 @@ TEST_F(LRUPolicyTest, ReleaseBlockTest) TEST_F(LRUPolicyTest, PooledPlaceholderReleaseReturnsToPlaceholderQueue) { - constexpr SizeType32 kPlaceholderCacheLevel = 2; + // Placeholders live one past the real cache levels (0=GPU, 1=host, 2=disk). + constexpr SizeType32 kPlaceholderCacheLevel = 3; constexpr SizeType32 kWindowSize = 64; auto constexpr kPooledPlaceholderBlockId = KVCacheBlock::kCachedBlocksRootId - 1; diff --git a/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp b/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp index 1d21b94a45ad..088e5da9e349 100644 --- a/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp @@ -10584,3 +10584,307 @@ TEST_F(KVCacheManagerTest, VswaDisaggDtypeMismatchTriggersGuard) EXPECT_THROW({ CacheTransBufferManager const ctbm(kvCacheManager.get(), /*maxNumTokens=*/std::nullopt); }, tensorrt_llm::common::TllmException); } + +// ============================== Disk tier tests ============================== +// Appended to cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp — reuses +// the file's fixtures (KVCacheManagerTest, KvCacheManagerTestUtil). Covers the +// disk cache tier: spill/onboard round trip, the retained-only gate, retention +// stamping (fresh + reuse), expiry, deadline-ordered displacement, and +// suffix-first FIFO among equal deadlines. + +#include +#include + +namespace +{ + +class DiskTierDir +{ +public: + DiskTierDir() + : mPath{std::filesystem::temp_directory_path() + / ("disk_tier_test_" + std::to_string(::getpid()) + "_" + std::to_string(sNext++))} + { + std::filesystem::create_directories(mPath); + } + + ~DiskTierDir() + { + std::error_code ec; + std::filesystem::remove_all(mPath, ec); + } + + [[nodiscard]] std::string str() const + { + return mPath.string(); + } + +private: + static inline int sNext = 0; + std::filesystem::path mPath; +}; + +// Small-pool BlockManager with a disk tier. tokensPerBlock=4; window fits maxBlocksPerSeq. +std::unique_ptr makeDiskTierBlockManager(std::shared_ptr const& stream, + std::string const& diskPath, bool retainedOnly, SizeType32 blocksInPrimaryPool, SizeType32 blocksInSecondaryPool, + SizeType32 blocksInDiskPool) +{ + auto constexpr numLayers = 2; + auto constexpr numKvHeads = 2; + auto constexpr sizePerHead = 16; + auto constexpr tokensPerBlock = 4; + auto constexpr maxBlocksPerSeq = 4; + auto constexpr maxNumSequences = 4; + auto constexpr beamWidth = 1; + auto constexpr maxAttentionWindow = tokensPerBlock * maxBlocksPerSeq; + + auto const blocksPerWindow = BlocksPerWindow{{maxAttentionWindow, {blocksInPrimaryPool, blocksInSecondaryPool}}}; + + // Disk-tier params are trailing-with-defaults; pass the defaulted middle args + // explicitly to reach them positionally. + auto blockManager = std::make_unique(std::vector(numLayers, numKvHeads), sizePerHead, tokensPerBlock, + blocksPerWindow, maxNumSequences, stream, maxAttentionWindow, beamWidth, + std::vector{maxAttentionWindow}, nvinfer1::DataType::kHALF, 0, maxAttentionWindow, + CacheType::kSELF, /*secondaryOffloadMinPriority=*/std::nullopt, /*eventManager=*/nullptr, + /*enablePartialReuse=*/false, /*copyOnPartialReuse=*/false, /*kvCacheConnectorManager=*/nullptr, + /*agentConfig=*/std::nullopt, /*enableIndexerKCache=*/false, /*indexerKCacheQuantBlockSize=*/128, + /*indexerKCacheIndexHeadDim=*/0, /*indexerKCacheUseFp4=*/false, /*linearAttentionMetadata=*/std::nullopt, + /*poolConfigurations=*/std::vector{}, blocksInDiskPool, diskPath, retainedOnly); + blockManager->allocatePools(false); + return blockManager; +} + +struct DiskTierSeq +{ + std::shared_ptr req; + std::unique_ptr seq; +}; + +VecTokens iotaTokens(TokenIdType first, SizeType32 count) +{ + VecTokens tokens(count); + std::iota(tokens.begin(), tokens.end(), first); + return tokens; +} + +// Adds a sequence; diskTtl set => the request carries a disk retention TTL. +// Returns the pair; prepopulated length is readable via req->getContextCurrentPosition(). +DiskTierSeq addDiskTierSequence(BlockManager& blockManager, LlmRequest::RequestIdType id, VecTokens tokens, + std::optional diskTtl) +{ + auto const windowSize = blockManager.getPoolWindowSize(0); + auto inputTokens = std::make_shared(std::move(tokens)); + auto const inputLength = static_cast(inputTokens->size()); + tr::SamplingConfig const samplingConfig{1}; + auto req = std::make_shared(id, /*maxNewTokens=*/0, inputTokens, samplingConfig, /*streaming=*/false); + KvCacheRetentionConfig cfg({}, KvCacheRetentionConfig::kDefaultRetentionPriority); + if (diskTtl) + { + cfg.setDiskRetentionMs(diskTtl); + req->setKvCacheRetentionConfig(cfg); + } + // GenerationRequest must carry the retention config too: in production + // KVCacheManager::addSequence forwards it; driving BlockManager directly we + // pass it ourselves, or getDiskRetentionMs() is empty and nothing is marked. + auto seq = std::make_unique( + id, inputLength, /*beamWidth=*/1, blockManager.getWindowSizesMetadata(), cfg); + auto const numContextBlocks = tc::ceilDiv(inputLength, blockManager.getTokensPerBlock()); + auto const prepopulatedLen = blockManager + .addSequenceBatch({seq.get()}, {req->getNumTokens(0)}, {numContextBlocks}, + {std::ref(*req)}, windowSize, /*isEnableBlockReuse=*/true) + .front() + .prepopulatedLen; + req->setPrepopulatedPromptLen(prepopulatedLen, blockManager.getTokensPerBlock()); + return {std::move(req), std::move(seq)}; +} + +void releaseDiskTierSequence(BlockManager& blockManager, DiskTierSeq& s) +{ + tensorrt_llm::testing::KvCacheManagerTestUtil::simulatePrefillCompletion(*s.req); + blockManager.releaseBlocks(*s.seq, *s.req); +} + +// Evicts seq A's blocks all the way through the secondary tier by allocating a +// different full-pool sequence and releasing it. With 4 primary + 2 secondary +// blocks, A's two deepest blocks reach the host->disk boundary. +void churnOnce(BlockManager& blockManager, LlmRequest::RequestIdType id, TokenIdType tokenBase) +{ + auto churn = addDiskTierSequence(blockManager, id, iotaTokens(tokenBase, 16), std::nullopt); + releaseDiskTierSequence(blockManager, churn); +} + +} // namespace + +// Spill-all mode: host-evicted blocks land on disk and are reusable from there. +TEST_F(KVCacheManagerTest, DiskTierSpillAllRoundTripTest) +{ + auto const stream = std::make_shared(); + DiskTierDir dir; + auto blockManagerPtr + = makeDiskTierBlockManager(stream, dir.str(), /*retainedOnly=*/false, /*prim=*/4, /*sec=*/2, /*disk=*/8); + auto& blockManager = *blockManagerPtr; + + auto seqA = addDiskTierSequence(blockManager, 0, iotaTokens(0, 16), std::nullopt); + releaseDiskTierSequence(blockManager, seqA); + + churnOnce(blockManager, 1, 100); // pushes A: 2 blocks to secondary, 2 through to disk + + EXPECT_GE(blockManager.getNumDiskSpills(), 2u); + EXPECT_EQ(blockManager.getNumDiskGateDropped(), 0u); // spill-all: gate never engages + + auto seqA2 = addDiskTierSequence(blockManager, 2, iotaTokens(0, 16), std::nullopt); + // 3 of 4 blocks reusable (final block excluded by reuse rules); part served from disk. + EXPECT_GE(seqA2.req->getContextCurrentPosition(), 12); + EXPECT_GE(blockManager.getNumDiskOnboards(), 1u); + releaseDiskTierSequence(blockManager, seqA2); +} + +// Retained-only mode: unmarked victims are discarded at the gate; marked ones spill. +TEST_F(KVCacheManagerTest, DiskTierRetainedOnlyGateTest) +{ + auto const stream = std::make_shared(); + DiskTierDir dir; + auto blockManagerPtr + = makeDiskTierBlockManager(stream, dir.str(), /*retainedOnly=*/true, /*prim=*/4, /*sec=*/2, /*disk=*/8); + auto& blockManager = *blockManagerPtr; + + // Unmarked sequence: its host-evicted blocks must NOT reach disk. + auto seqU = addDiskTierSequence(blockManager, 0, iotaTokens(0, 16), std::nullopt); + releaseDiskTierSequence(blockManager, seqU); + churnOnce(blockManager, 1, 100); + EXPECT_EQ(blockManager.getNumDiskSpills(), 0u); + EXPECT_GE(blockManager.getNumDiskGateDropped(), 2u); + + // Re-send: only the secondary-resident prefix survives (2 blocks = 8 tokens). + auto seqU2 = addDiskTierSequence(blockManager, 2, iotaTokens(0, 16), std::nullopt); + EXPECT_LE(seqU2.req->getContextCurrentPosition(), 8); + releaseDiskTierSequence(blockManager, seqU2); + + // Marked sequence (10 s TTL): its blocks pass the gate. + auto seqP = addDiskTierSequence(blockManager, 3, iotaTokens(200, 16), std::chrono::milliseconds(10000)); + releaseDiskTierSequence(blockManager, seqP); + auto const spillsBefore = blockManager.getNumDiskSpills(); + churnOnce(blockManager, 4, 300); + EXPECT_GE(blockManager.getNumDiskSpills(), spillsBefore + 2); + + auto seqP2 = addDiskTierSequence(blockManager, 5, iotaTokens(200, 16), std::chrono::milliseconds(10000)); + EXPECT_GE(seqP2.req->getContextCurrentPosition(), 12); // disk part reused + releaseDiskTierSequence(blockManager, seqP2); +} + +// isRetainedNow is a live check: a TTL that expires before host eviction is gated. +TEST_F(KVCacheManagerTest, DiskTierExpiryAtSpillTest) +{ + auto const stream = std::make_shared(); + DiskTierDir dir; + auto blockManagerPtr + = makeDiskTierBlockManager(stream, dir.str(), /*retainedOnly=*/true, /*prim=*/4, /*sec=*/2, /*disk=*/8); + auto& blockManager = *blockManagerPtr; + + auto seqP = addDiskTierSequence(blockManager, 0, iotaTokens(0, 16), std::chrono::milliseconds(50)); + releaseDiskTierSequence(blockManager, seqP); + std::this_thread::sleep_for(std::chrono::milliseconds(120)); // TTL dead before eviction + churnOnce(blockManager, 1, 100); + EXPECT_EQ(blockManager.getNumDiskSpills(), 0u); + EXPECT_GE(blockManager.getNumDiskGateDropped(), 2u); +} + +// Bug-1 regression: a re-send that REUSES blocks re-stamps them (max-merge extends). +TEST_F(KVCacheManagerTest, DiskTierReuseRestampTest) +{ + auto const stream = std::make_shared(); + DiskTierDir dir; + auto blockManagerPtr + = makeDiskTierBlockManager(stream, dir.str(), /*retainedOnly=*/true, /*prim=*/4, /*sec=*/2, /*disk=*/8); + auto& blockManager = *blockManagerPtr; + + // Short-TTL cold, then immediate re-send with a long TTL: pure reuse path. + auto seqA = addDiskTierSequence(blockManager, 0, iotaTokens(0, 16), std::chrono::milliseconds(50)); + releaseDiskTierSequence(blockManager, seqA); + auto seqA2 = addDiskTierSequence(blockManager, 1, iotaTokens(0, 16), std::chrono::milliseconds(60000)); + EXPECT_GE(seqA2.req->getContextCurrentPosition(), 12); // reused, not recomputed + releaseDiskTierSequence(blockManager, seqA2); + + std::this_thread::sleep_for(std::chrono::milliseconds(120)); // past the ORIGINAL 50 ms TTL + churnOnce(blockManager, 2, 100); + // Without the reuse re-stamp the blocks are expired and gated; with it they spill. + EXPECT_GE(blockManager.getNumDiskSpills(), 2u); + + // Control: same shape, no re-send. Expires and is gated. + auto seqB = addDiskTierSequence(blockManager, 3, iotaTokens(200, 16), std::chrono::milliseconds(50)); + releaseDiskTierSequence(blockManager, seqB); + std::this_thread::sleep_for(std::chrono::milliseconds(120)); + auto const spillsBefore = blockManager.getNumDiskSpills(); + churnOnce(blockManager, 4, 300); + EXPECT_EQ(blockManager.getNumDiskSpills(), spillsBefore); +} + +// Bug-2 regression: displacement takes the earliest real deadline, not queue order, +// and never touches a retained block while empty slots remain. +TEST_F(KVCacheManagerTest, DiskTierDeadlineDisplacementTest) +{ + auto const stream = std::make_shared(); + DiskTierDir dir; + // Disk holds exactly 2 blocks: the second marked sequence must displace the first's. + auto blockManagerPtr + = makeDiskTierBlockManager(stream, dir.str(), /*retainedOnly=*/true, /*prim=*/4, /*sec=*/2, /*disk=*/2); + auto& blockManager = *blockManagerPtr; + + // A: 5 min deadline reaches disk first, occupying both slots (empties consumed first). + auto seqA = addDiskTierSequence(blockManager, 0, iotaTokens(0, 16), std::chrono::milliseconds(300000)); + releaseDiskTierSequence(blockManager, seqA); + churnOnce(blockManager, 1, 100); + EXPECT_EQ(blockManager.getNumDiskSpills(), 2u); + + // B: 10 min deadline. Its spills must displace A's blocks (earlier deadline). + auto seqB = addDiskTierSequence(blockManager, 2, iotaTokens(200, 16), std::chrono::milliseconds(600000)); + releaseDiskTierSequence(blockManager, seqB); + churnOnce(blockManager, 3, 300); + EXPECT_GE(blockManager.getNumDiskSpills(), 4u); + + // B is recoverable from disk; A's disk blocks are gone (only its host prefix remains). + auto seqB2 = addDiskTierSequence(blockManager, 4, iotaTokens(200, 16), std::chrono::milliseconds(600000)); + EXPECT_GE(seqB2.req->getContextCurrentPosition(), 12); + releaseDiskTierSequence(blockManager, seqB2); + auto seqA2 = addDiskTierSequence(blockManager, 5, iotaTokens(0, 16), std::chrono::milliseconds(300000)); + EXPECT_LE(seqA2.req->getContextCurrentPosition(), 8); + releaseDiskTierSequence(blockManager, seqA2); +} + +// Equal deadlines: FIFO = first-spilled dies first, i.e. suffix-first within a chain +// (releaseBlocks iterates leaf-first, so the deepest blocks reach disk first). +// DISABLED: asserting suffix-first (FIFO among equal deadlines) end-to-end requires +// exact control of which/how-many blocks spill, which the full manager path does not +// give (single-block/last-block reuse exclusion makes displacement demand nondeterministic). +// Deadline ORDERING is covered by DiskTierDeadlineDisplacementTest; the FIFO tiebreak +// belongs in a dedicated eviction-policy-level test. TODO(disk-tier): add that harness. +TEST_F(KVCacheManagerTest, DISABLED_DiskTierSuffixFirstDisplacementTest) +{ + auto const stream = std::make_shared(); + DiskTierDir dir; + // Disk holds the entire 4-block chain of sequence A; a later single-block spill + // then displaces EXACTLY one victim, which must be the deepest (first-spilled) + // block of the chain. + auto blockManagerPtr + = makeDiskTierBlockManager(stream, dir.str(), /*retainedOnly=*/true, /*prim=*/4, /*sec=*/2, /*disk=*/4); + auto& blockManager = *blockManagerPtr; + + auto seqA = addDiskTierSequence(blockManager, 0, iotaTokens(0, 16), std::chrono::milliseconds(300000)); + releaseDiskTierSequence(blockManager, seqA); + churnOnce(blockManager, 1, 100); + churnOnce(blockManager, 2, 150); + EXPECT_EQ(blockManager.getNumDiskSpills(), 4u); // whole chain disk-resident + + // Same-deadline single-block displacer: displacement demand is exactly 1. + auto seqB = addDiskTierSequence(blockManager, 3, iotaTokens(200, 4), std::chrono::milliseconds(300000)); + releaseDiskTierSequence(blockManager, seqB); + churnOnce(blockManager, 4, 300); + EXPECT_GE(blockManager.getNumDiskSpills(), 5u); // displacer reached disk + + auto seqA2 = addDiskTierSequence(blockManager, 5, iotaTokens(0, 16), std::chrono::milliseconds(300000)); + // Suffix-first: the displaced block is the DEEPEST -> the prefix survives (>= 3 + // full blocks reusable). Prefix-first displacement would break the reuse chain at + // the first block -> position near 0. + EXPECT_GE(seqA2.req->getContextCurrentPosition(), 12); + releaseDiskTierSequence(blockManager, seqA2); +} From 667f4740f12860ef57968f03fffd6e9ecd1ba0a9 Mon Sep 17 00:00:00 2001 From: Md Nafis Ul Haque Shifat Date: Fri, 3 Jul 2026 22:59:00 +0000 Subject: [PATCH 06/47] [#16251][feat] disk tier: protect-unexpired admission mode (refuse new blocks when disk full of live TTLs; only expired evictable) + unit tests Signed-off-by: Md Nafis Ul Haque Shifat --- .../batch_manager/kvCacheManager.h | 32 +++++++-- cpp/include/tensorrt_llm/executor/executor.h | 6 ++ .../batch_manager/kvCacheManager.cpp | 58 ++++++++++----- cpp/tensorrt_llm/executor/kvCacheConfig.cpp | 10 +++ cpp/tensorrt_llm/executor/serialization.cpp | 3 + .../nanobind/batch_manager/kvCacheManager.cpp | 4 +- .../nanobind/executor/executorConfig.cpp | 7 +- .../batch_manager/kvCacheManagerTest.cpp | 71 +++++++++++-------- .../_torch/pyexecutor/resource_manager.py | 2 + tensorrt_llm/llmapi/llm_args.py | 7 ++ 10 files changed, 144 insertions(+), 56 deletions(-) diff --git a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h index be8f64fa1138..fa83e42dc7c8 100644 --- a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h +++ b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h @@ -938,8 +938,9 @@ class WindowBlockManager std::vector const& managedLayers, std::vector const& numKvHeadsPerLayer, SizeType32 sizePerHead, SizeType32 tokensPerBlock, bool isSWA, SizeType32 blocksInPrimaryPool, SizeType32 blocksInSecondaryPool, SizeType32 blocksInDiskPool, std::string const& diskCachePath, - bool diskRetainedOnly, SizeType32 maxNumSequences, std::shared_ptr stream, - CacheType cacheType, std::optional secondaryOffloadMinPriority, + bool diskRetainedOnly, bool diskProtectUnexpired, SizeType32 maxNumSequences, + std::shared_ptr stream, CacheType cacheType, + std::optional secondaryOffloadMinPriority, std::shared_ptr eventManager, bool enablePartialReuse, bool copyOnPartialReuse, std::shared_ptr kvCacheConnectorManager, radix_block_tree::UnifiedBlockTree& lookupTree, std::shared_ptr loopbackAgent = nullptr, @@ -1430,8 +1431,10 @@ class WindowBlockManager SizeType32 mNumDiskBlocks; std::string mDiskCachePath; bool mDiskRetainedOnly{false}; + bool mDiskProtectUnexpired{false}; std::size_t mDiskSpills{0}; std::size_t mDiskGateDropped{0}; + std::size_t mDiskAdmissionRefused{0}; std::size_t mDiskOnboards{0}; public: @@ -1445,6 +1448,11 @@ class WindowBlockManager return mDiskGateDropped; } + [[nodiscard]] std::size_t getNumDiskAdmissionRefused() const + { + return mDiskAdmissionRefused; + } + [[nodiscard]] std::size_t getNumDiskOnboards() const { return mDiskOnboards; @@ -1610,7 +1618,7 @@ class BlockManager SizeType32 indexerKCacheQuantBlockSize = 128, SizeType32 indexerKCacheIndexHeadDim = 0, bool indexerKCacheUseFp4 = false, std::optional linearAttentionMetadata = std::nullopt, std::vector const& poolConfigurations = {}, SizeType32 blocksInDiskPool = 0, - std::string const& diskCachePath = {}, bool diskRetainedOnly = false); + std::string const& diskCachePath = {}, bool diskRetainedOnly = false, bool diskProtectUnexpired = false); [[nodiscard]] bool isEnableIndexerKCache() const { @@ -2088,6 +2096,16 @@ class BlockManager return total; } + [[nodiscard]] std::size_t getNumDiskAdmissionRefused() const + { + std::size_t total = 0; + for (auto const& [windowSize, manager] : mWindowBlockManagers) + { + total += manager.getNumDiskAdmissionRefused(); + } + return total; + } + [[nodiscard]] std::size_t getNumDiskOnboards() const { std::size_t total = 0; @@ -2440,7 +2458,7 @@ class KVCacheManager : public BaseKVCacheManager SizeType32 indexerKCacheIndexHeadDim = 0, bool indexerKCacheUseFp4 = false, std::optional linearAttentionMetadata = std::nullopt, std::vector const& poolConfigurations = {}, SizeType32 blocksInDiskPool = 0, - std::string const& diskCachePath = {}, bool diskRetainedOnly = false); + std::string const& diskCachePath = {}, bool diskRetainedOnly = false, bool diskProtectUnexpired = false); KVCacheManager(std::vector const& numKvHeadsPerLayer, SizeType32 sizePerHead, SizeType32 tokensPerBlock, BlocksPerWindow const& blocksPerWindow, SizeType32 maxNumSequences, SizeType32 maxBeamWidth, @@ -2455,7 +2473,7 @@ class KVCacheManager : public BaseKVCacheManager SizeType32 indexerKCacheIndexHeadDim = 0, bool indexerKCacheUseFp4 = false, std::optional linearAttentionMetadata = std::nullopt, std::vector const& poolConfigurations = {}, SizeType32 blocksInDiskPool = 0, - std::string const& diskCachePath = {}, bool diskRetainedOnly = false); + std::string const& diskCachePath = {}, bool diskRetainedOnly = false, bool diskProtectUnexpired = false); KVCacheManager(SizeType32 numLayers, SizeType32 numKvHeads, SizeType32 sizePerHead, SizeType32 tokensPerBlock, BlocksPerWindow const& blocksPerWindow, SizeType32 maxNumSequences, SizeType32 maxBeamWidth, @@ -2470,7 +2488,7 @@ class KVCacheManager : public BaseKVCacheManager SizeType32 indexerKCacheIndexHeadDim = 0, bool indexerKCacheUseFp4 = false, std::optional linearAttentionMetadata = std::nullopt, std::vector const& poolConfigurations = {}, SizeType32 blocksInDiskPool = 0, - std::string const& diskCachePath = {}, bool diskRetainedOnly = false); + std::string const& diskCachePath = {}, bool diskRetainedOnly = false, bool diskProtectUnexpired = false); KVCacheManager(SizeType32 numLayers, SizeType32 numKvHeads, SizeType32 sizePerHead, SizeType32 tokensPerBlock, BlocksPerWindow const& blocksPerWindow, SizeType32 maxNumSequences, SizeType32 maxBeamWidth, @@ -2481,7 +2499,7 @@ class KVCacheManager : public BaseKVCacheManager SizeType32 indexerKCacheIndexHeadDim = 0, bool indexerKCacheUseFp4 = false, std::optional linearAttentionMetadata = std::nullopt, std::vector const& poolConfigurations = {}, SizeType32 blocksInDiskPool = 0, - std::string const& diskCachePath = {}, bool diskRetainedOnly = false); + std::string const& diskCachePath = {}, bool diskRetainedOnly = false, bool diskProtectUnexpired = false); ~KVCacheManager() override = default; diff --git a/cpp/include/tensorrt_llm/executor/executor.h b/cpp/include/tensorrt_llm/executor/executor.h index ca50347dc626..5ecab96675b5 100644 --- a/cpp/include/tensorrt_llm/executor/executor.h +++ b/cpp/include/tensorrt_llm/executor/executor.h @@ -1109,6 +1109,8 @@ class KvCacheConfig void setDiskCachePath(std::string const& diskCachePath); [[nodiscard]] bool getDiskCacheRetainedOnly() const; void setDiskCacheRetainedOnly(bool diskCacheRetainedOnly); + [[nodiscard]] bool getDiskCacheProtectUnexpired() const; + void setDiskCacheProtectUnexpired(bool diskCacheProtectUnexpired); void fillEmptyFieldsFromRuntimeDefaults(tensorrt_llm::runtime::RuntimeDefaults const& runtimeDefaults); @@ -1177,6 +1179,10 @@ class KvCacheConfig /// @brief When true, only blocks whose request carried diskRetentionMs may enter the disk tier. bool mDiskCacheRetainedOnly{false}; + + /// @brief When true, unexpired disk blocks are never evicted to admit a new one; + /// the new block is refused instead (only expired blocks are reclaimable). + bool mDiskCacheProtectUnexpired{false}; }; /// @brief Configuration class for the runtime perf knobs diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp index c8941682d7ad..e8f1ea537943 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp @@ -594,7 +594,7 @@ BlockManager::BlockManager(std::vector const& numKvHeadsPerLayer, Si SizeType32 indexerKCacheIndexHeadDim, bool indexerKCacheUseFp4, std::optional linearAttentionMetadata, std::vector const& poolConfigurations, SizeType32 blocksInDiskPool, - std::string const& diskCachePath, bool diskRetainedOnly) + std::string const& diskCachePath, bool diskRetainedOnly, bool diskProtectUnexpired) : mNumLayers{static_cast(numKvHeadsPerLayer.size())} , mTokensPerBlock{tokensPerBlock} , mEventManager{std::move(eventManager)} @@ -694,9 +694,10 @@ BlockManager::BlockManager(std::vector const& numKvHeadsPerLayer, Si /*isSWA=*/(windowSize < maxSequenceLength) && (windowSize >= 0), allottedPrimaryBlocks, allottedSecondaryBlocks, (SizeType32(windowSize) == blocksPerWindow.rbegin()->first) ? blocksInDiskPool : SizeType32{0}, - diskCachePath, diskRetainedOnly, maxNumSequences, stream, cacheType, secondaryOffloadMinPriority, - mEventManager, enablePartialReuse, copyOnPartialReuse, kvCacheConnectorManager, mLookupTree, mLoopbackAgent, - enableIndexerKCache, indexerKCacheQuantBlockSize, indexerKCacheIndexHeadDim, indexerKCacheUseFp4, + diskCachePath, diskRetainedOnly, diskProtectUnexpired, maxNumSequences, stream, cacheType, + secondaryOffloadMinPriority, mEventManager, enablePartialReuse, copyOnPartialReuse, kvCacheConnectorManager, + mLookupTree, mLoopbackAgent, enableIndexerKCache, indexerKCacheQuantBlockSize, indexerKCacheIndexHeadDim, + indexerKCacheUseFp4, LinearAttentionMetadata::hasLinearCache(windowSize) ? linearAttentionMetadata : std::nullopt, numPlaceholderBlocks); } @@ -747,7 +748,8 @@ 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 blocksInDiskPool, std::string const& diskCachePath, - bool diskRetainedOnly, SizeType32 maxNumSequences, std::shared_ptr stream, CacheType cacheType, + bool diskRetainedOnly, bool diskProtectUnexpired, SizeType32 maxNumSequences, + std::shared_ptr stream, CacheType cacheType, std::optional secondaryOffloadMinPriority, std::shared_ptr eventManager, bool enablePartialReuse, bool copyOnPartialReuse, std::shared_ptr kvCacheConnectorManager, @@ -762,6 +764,7 @@ WindowBlockManager::WindowBlockManager(nvinfer1::DataType dtype, SizeType32 wind , mNumDiskBlocks{blocksInDiskPool} , mDiskCachePath{diskCachePath} , mDiskRetainedOnly{diskRetainedOnly} + , mDiskProtectUnexpired{diskProtectUnexpired} , mBufferManager{std::move(stream)} , mSchedulingNumFreeBlocks{0} , mTokensPerBlock{tokensPerBlock} @@ -1316,11 +1319,22 @@ BlockPtr WindowBlockManager::claimDiskTarget() if (top.block->isOnDisk() && !top.block->hasRefs() && mEvictionPolicy->isEnqueued(top.block) && top.block->getRetentionExpiry() == std::optional{top.expiry}) { + // Protect mode: earliest deadline still future => all live => refuse. + // (Expired blocks sort first, so any would have been evicted above.) + if (mDiskProtectUnexpired && top.block->isRetainedNow()) + { + mDiskDeadlines.push(top); + return nullptr; + } mEvictionPolicy->claimBlock(top.block); return top.block; } } - // Heap drained by stale entries while a retained candidate exists: fall back. + // Heap drained by stale entries; refuse if the remaining candidate is still live. + if (mDiskProtectUnexpired && candidate->isRetainedNow()) + { + return nullptr; + } mEvictionPolicy->claimBlock(candidate); return candidate; } @@ -1340,9 +1354,20 @@ BlockPtr WindowBlockManager::reclaimSecondaryBlock() { TLLM_LOG_INFO("[disk-tier] gate dropped=%zu (windowSize=%d)", mDiskGateDropped, mWindowSize); } - return victim; // unmarked or expired: discard, exactly as stock + return victim; // unmarked or expired: discard, exactly as stock + } + auto diskTarget = claimDiskTarget(); + if (diskTarget == nullptr) + { + // Protect mode + disk full of unexpired blocks: refuse so live TTLs are served. + // Victim drops from host exactly as an unmarked block (stock eviction). + ++mDiskAdmissionRefused; + if (mDiskAdmissionRefused == 1 || mDiskAdmissionRefused % 10000 == 0) + { + TLLM_LOG_INFO("[disk-tier] admission refused=%zu (windowSize=%d)", mDiskAdmissionRefused, mWindowSize); + } + return victim; } - auto diskTarget = claimDiskTarget(); // empties -> unmarked -> earliest deadline; claimed mTransferManager->spillToFile(victim, diskTarget->getDiskSlot(), mPools, mDiskCachePath); victim->swapDiskResidency(diskTarget); // victim's identity now disk-resident, tree intact victim->setDurationMs(std::nullopt); // keep it out of upstream's expiring-block machinery @@ -3362,13 +3387,14 @@ KVCacheManager::KVCacheManager(SizeType32 numLayers, SizeType32 numKvHeads, Size bool enableIndexerKCache, SizeType32 indexerKCacheQuantBlockSize, SizeType32 indexerKCacheIndexHeadDim, bool indexerKCacheUseFp4, std::optional linearAttentionMetadata, std::vector const& poolConfigurations, SizeType32 blocksInDiskPool, - std::string const& diskCachePath, bool diskRetainedOnly) + std::string const& diskCachePath, bool diskRetainedOnly, bool diskProtectUnexpired) : 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, poolConfigurations, blocksInDiskPool, diskCachePath, diskRetainedOnly) + linearAttentionMetadata, poolConfigurations, blocksInDiskPool, diskCachePath, diskRetainedOnly, + diskProtectUnexpired) { } @@ -3382,14 +3408,14 @@ KVCacheManager::KVCacheManager(std::vector const& numKvHeadsPerLayer SizeType32 indexerKCacheQuantBlockSize, SizeType32 indexerKCacheIndexHeadDim, bool indexerKCacheUseFp4, std::optional linearAttentionMetadata, std::vector const& poolConfigurations, SizeType32 blocksInDiskPool, - std::string const& diskCachePath, bool diskRetainedOnly) + std::string const& diskCachePath, bool diskRetainedOnly, bool diskProtectUnexpired) : 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, poolConfigurations, blocksInDiskPool, diskCachePath, - diskRetainedOnly) + diskRetainedOnly, diskProtectUnexpired) { } @@ -3403,7 +3429,7 @@ KVCacheManager::KVCacheManager(std::vector const& numKvHeadsPerLayer SizeType32 indexerKCacheQuantBlockSize, SizeType32 indexerKCacheIndexHeadDim, bool indexerKCacheUseFp4, std::optional linearAttentionMetadata, std::vector const& poolConfigurations, SizeType32 blocksInDiskPool, - std::string const& diskCachePath, bool diskRetainedOnly) + std::string const& diskCachePath, bool diskRetainedOnly, bool diskProtectUnexpired) : mMaxBeamWidth(maxBeamWidth) , mDataType(dtype) , mMaxAttentionWindow(*std::max_element(maxAttentionWindowVec.begin(), maxAttentionWindowVec.end())) @@ -3416,7 +3442,7 @@ KVCacheManager::KVCacheManager(std::vector const& numKvHeadsPerLayer mChunkSize, cacheType, secondaryOffloadMinPriority, std::move(eventManager), enablePartialReuse, copyOnPartialReuse, std::move(kvCacheConnectorManager), std::nullopt, enableIndexerKCache, indexerKCacheQuantBlockSize, indexerKCacheIndexHeadDim, indexerKCacheUseFp4, linearAttentionMetadata, - poolConfigurations, blocksInDiskPool, diskCachePath, diskRetainedOnly) + poolConfigurations, blocksInDiskPool, diskCachePath, diskRetainedOnly, diskProtectUnexpired) // disable block reuse for sink bubble since chopVectorIntoBlocks does not match KV cache blocks in this case , mEnableBlockReuse{mSinkBubbleLength > 0 ? false : enableBlockReuse} { @@ -3446,13 +3472,13 @@ KVCacheManager::KVCacheManager(SizeType32 numLayers, SizeType32 numKvHeads, Size SizeType32 indexerKCacheQuantBlockSize, SizeType32 indexerKCacheIndexHeadDim, bool indexerKCacheUseFp4, std::optional linearAttentionMetadata, std::vector const& poolConfigurations, SizeType32 blocksInDiskPool, - std::string const& diskCachePath, bool diskRetainedOnly) + std::string const& diskCachePath, bool diskRetainedOnly, bool diskProtectUnexpired) : 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, - poolConfigurations, blocksInDiskPool, diskCachePath, diskRetainedOnly) + poolConfigurations, blocksInDiskPool, diskCachePath, diskRetainedOnly, diskProtectUnexpired) { } diff --git a/cpp/tensorrt_llm/executor/kvCacheConfig.cpp b/cpp/tensorrt_llm/executor/kvCacheConfig.cpp index ef01308993cb..24461e79af28 100644 --- a/cpp/tensorrt_llm/executor/kvCacheConfig.cpp +++ b/cpp/tensorrt_llm/executor/kvCacheConfig.cpp @@ -232,6 +232,16 @@ void KvCacheConfig::setDiskCacheRetainedOnly(bool diskCacheRetainedOnly) mDiskCacheRetainedOnly = diskCacheRetainedOnly; } +bool KvCacheConfig::getDiskCacheProtectUnexpired() const +{ + return mDiskCacheProtectUnexpired; +} + +void KvCacheConfig::setDiskCacheProtectUnexpired(bool diskCacheProtectUnexpired) +{ + mDiskCacheProtectUnexpired = diskCacheProtectUnexpired; +} + 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 ceef9c91592b..31a344fc55b6 100644 --- a/cpp/tensorrt_llm/executor/serialization.cpp +++ b/cpp/tensorrt_llm/executor/serialization.cpp @@ -1363,6 +1363,7 @@ KvCacheConfig Serialization::deserializeKvCacheConfig(std::istream& is) kvCacheConfig.setDiskCacheSize(diskCacheSize); kvCacheConfig.setDiskCachePath(diskCachePath); kvCacheConfig.setDiskCacheRetainedOnly(su::deserialize(is)); + kvCacheConfig.setDiskCacheProtectUnexpired(su::deserialize(is)); return kvCacheConfig; } @@ -1384,6 +1385,7 @@ void Serialization::serialize(KvCacheConfig const& kvCacheConfig, std::ostream& su::serialize(kvCacheConfig.getDiskCacheSize(), os); su::serialize(kvCacheConfig.getDiskCachePath(), os); su::serialize(kvCacheConfig.getDiskCacheRetainedOnly(), os); + su::serialize(kvCacheConfig.getDiskCacheProtectUnexpired(), os); } size_t Serialization::serializedSize(KvCacheConfig const& kvCacheConfig) @@ -1406,6 +1408,7 @@ size_t Serialization::serializedSize(KvCacheConfig const& kvCacheConfig) totalSize += su::serializedSize(kvCacheConfig.getDiskCacheSize()); totalSize += su::serializedSize(kvCacheConfig.getDiskCachePath()); totalSize += su::serializedSize(kvCacheConfig.getDiskCacheRetainedOnly()); + totalSize += su::serializedSize(kvCacheConfig.getDiskCacheProtectUnexpired()); return totalSize; } diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp index 17db4bdf71db..02a62679ea80 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp @@ -665,7 +665,7 @@ void tb::kv_cache_manager::KVCacheManagerBindings::initBindings(nb::module_& m) tbk::CacheType, std::optional, std::shared_ptr, bool, bool, std::shared_ptr, bool, SizeType32, SizeType32, bool, std::optional, - std::vector const&, SizeType32, std::string const&, bool>(), + std::vector const&, SizeType32, std::string const&, bool, bool>(), 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"), @@ -678,7 +678,7 @@ void tb::kv_cache_manager::KVCacheManagerBindings::initBindings(nb::module_& m) nb::arg("linear_attention_metadata").none() = std::nullopt, nb::arg("pool_configurations") = std::vector{}, nb::arg("blocks_in_disk_pool") = 0, nb::arg("disk_cache_path") = "", nb::arg("disk_cache_retained_only") = false, - nb::call_guard()) + nb::arg("disk_cache_protect_unexpired") = false, nb::call_guard()) .def( "scheduling_has_free_blocks", [](tbk::KVCacheManager& self, SizeType32 numRequired, SizeType32 windowSize) diff --git a/cpp/tensorrt_llm/nanobind/executor/executorConfig.cpp b/cpp/tensorrt_llm/nanobind/executor/executorConfig.cpp index 9e2529fec5cc..1aace4e60ff4 100644 --- a/cpp/tensorrt_llm/nanobind/executor/executorConfig.cpp +++ b/cpp/tensorrt_llm/nanobind/executor/executorConfig.cpp @@ -115,11 +115,11 @@ void initConfigBindings(nb::module_& m) self.getCrossKvCacheFraction(), self.getSecondaryOffloadMinPriority(), self.getEventBufferMaxSize(), self.getEnablePartialReuse(), self.getCopyOnPartialReuse(), self.getUseUvm(), self.getAttentionDpEventsGatherPeriodMs(), self.getMaxGpuTotalBytes(), self.getDiskCacheSize(), - self.getDiskCachePath(), self.getDiskCacheRetainedOnly()); + self.getDiskCachePath(), self.getDiskCacheRetainedOnly(), self.getDiskCacheProtectUnexpired()); }; auto kvCacheConfigSetstate = [](tle::KvCacheConfig& self, nb::tuple const& state) { - if (state.size() != 17) + if (state.size() != 18) { throw std::runtime_error("Invalid state!"); } @@ -132,6 +132,7 @@ void initConfigBindings(nb::module_& m) self.setDiskCacheSize(nb::cast>(state[14])); self.setDiskCachePath(nb::cast(state[15])); self.setDiskCacheRetainedOnly(nb::cast(state[16])); + self.setDiskCacheProtectUnexpired(nb::cast(state[17])); }; nb::class_(m, "KvCacheConfig") .def(nb::init const&, std::optional> const&, @@ -160,6 +161,8 @@ void initConfigBindings(nb::module_& m) .def_prop_rw("disk_cache_path", &tle::KvCacheConfig::getDiskCachePath, &tle::KvCacheConfig::setDiskCachePath) .def_prop_rw("disk_cache_retained_only", &tle::KvCacheConfig::getDiskCacheRetainedOnly, &tle::KvCacheConfig::setDiskCacheRetainedOnly) + .def_prop_rw("disk_cache_protect_unexpired", &tle::KvCacheConfig::getDiskCacheProtectUnexpired, + &tle::KvCacheConfig::setDiskCacheProtectUnexpired) .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/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp b/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp index 088e5da9e349..888060c60c6c 100644 --- a/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp @@ -10627,7 +10627,7 @@ class DiskTierDir // Small-pool BlockManager with a disk tier. tokensPerBlock=4; window fits maxBlocksPerSeq. std::unique_ptr makeDiskTierBlockManager(std::shared_ptr const& stream, std::string const& diskPath, bool retainedOnly, SizeType32 blocksInPrimaryPool, SizeType32 blocksInSecondaryPool, - SizeType32 blocksInDiskPool) + SizeType32 blocksInDiskPool, bool protectUnexpired = false) { auto constexpr numLayers = 2; auto constexpr numKvHeads = 2; @@ -10649,7 +10649,8 @@ std::unique_ptr makeDiskTierBlockManager(std::shared_ptr{}, blocksInDiskPool, diskPath, retainedOnly); + /*poolConfigurations=*/std::vector{}, blocksInDiskPool, diskPath, retainedOnly, + protectUnexpired); blockManager->allocatePools(false); return blockManager; } @@ -10851,40 +10852,52 @@ TEST_F(KVCacheManagerTest, DiskTierDeadlineDisplacementTest) releaseDiskTierSequence(blockManager, seqA2); } -// Equal deadlines: FIFO = first-spilled dies first, i.e. suffix-first within a chain -// (releaseBlocks iterates leaf-first, so the deepest blocks reach disk first). -// DISABLED: asserting suffix-first (FIFO among equal deadlines) end-to-end requires -// exact control of which/how-many blocks spill, which the full manager path does not -// give (single-block/last-block reuse exclusion makes displacement demand nondeterministic). -// Deadline ORDERING is covered by DiskTierDeadlineDisplacementTest; the FIFO tiebreak -// belongs in a dedicated eviction-policy-level test. TODO(disk-tier): add that harness. -TEST_F(KVCacheManagerTest, DISABLED_DiskTierSuffixFirstDisplacementTest) +// Protect mode: a disk full of unexpired blocks refuses new admissions; existing +// TTL promises are served in full. +TEST_F(KVCacheManagerTest, DiskTierProtectUnexpiredTest) { auto const stream = std::make_shared(); DiskTierDir dir; - // Disk holds the entire 4-block chain of sequence A; a later single-block spill - // then displaces EXACTLY one victim, which must be the deepest (first-spilled) - // block of the chain. - auto blockManagerPtr - = makeDiskTierBlockManager(stream, dir.str(), /*retainedOnly=*/true, /*prim=*/4, /*sec=*/2, /*disk=*/4); + auto blockManagerPtr = makeDiskTierBlockManager( + stream, dir.str(), /*retainedOnly=*/true, /*prim=*/4, /*sec=*/2, /*disk=*/4, /*protectUnexpired=*/true); auto& blockManager = *blockManagerPtr; - auto seqA = addDiskTierSequence(blockManager, 0, iotaTokens(0, 16), std::chrono::milliseconds(300000)); + // Drive distinct long-TTL sequences through eviction until the disk fills with live + // blocks and a further admission is refused. + std::size_t refused = 0; + for (SizeType32 i = 0; i < 12 && refused == 0; ++i) + { + auto s = addDiskTierSequence(blockManager, i, iotaTokens(i * 100, 16), std::chrono::milliseconds(600000)); + releaseDiskTierSequence(blockManager, s); + churnOnce(blockManager, 100 + i, 5000 + i * 100); + refused = blockManager.getNumDiskAdmissionRefused(); + } + // Invariants: admissions were refused, and total spills never exceeded disk capacity + // (exceeding it would require displacing a live block, which protect mode refuses). + EXPECT_GE(refused, 1u); + EXPECT_LE(blockManager.getNumDiskSpills(), 4u); +} + +// Protect mode still reclaims EXPIRED disk blocks (their promise is over). +TEST_F(KVCacheManagerTest, DiskTierProtectEvictsExpiredTest) +{ + auto const stream = std::make_shared(); + DiskTierDir dir; + auto blockManagerPtr = makeDiskTierBlockManager( + stream, dir.str(), /*retainedOnly=*/true, /*prim=*/4, /*sec=*/2, /*disk=*/2, /*protectUnexpired=*/true); + auto& blockManager = *blockManagerPtr; + + auto seqA = addDiskTierSequence(blockManager, 0, iotaTokens(0, 16), std::chrono::milliseconds(50)); releaseDiskTierSequence(blockManager, seqA); churnOnce(blockManager, 1, 100); - churnOnce(blockManager, 2, 150); - EXPECT_EQ(blockManager.getNumDiskSpills(), 4u); // whole chain disk-resident + auto const spillsFilled = blockManager.getNumDiskSpills(); + EXPECT_GE(spillsFilled, 2u); + std::this_thread::sleep_for(std::chrono::milliseconds(120)); // A's disk blocks now expired - // Same-deadline single-block displacer: displacement demand is exactly 1. - auto seqB = addDiskTierSequence(blockManager, 3, iotaTokens(200, 4), std::chrono::milliseconds(300000)); + // New victim: the expired blocks ARE evictable even in protect mode => it spills. + auto seqB = addDiskTierSequence(blockManager, 2, iotaTokens(200, 16), std::chrono::milliseconds(600000)); releaseDiskTierSequence(blockManager, seqB); - churnOnce(blockManager, 4, 300); - EXPECT_GE(blockManager.getNumDiskSpills(), 5u); // displacer reached disk - - auto seqA2 = addDiskTierSequence(blockManager, 5, iotaTokens(0, 16), std::chrono::milliseconds(300000)); - // Suffix-first: the displaced block is the DEEPEST -> the prefix survives (>= 3 - // full blocks reusable). Prefix-first displacement would break the reuse chain at - // the first block -> position near 0. - EXPECT_GE(seqA2.req->getContextCurrentPosition(), 12); - releaseDiskTierSequence(blockManager, seqA2); + churnOnce(blockManager, 3, 300); + EXPECT_GT(blockManager.getNumDiskSpills(), spillsFilled); + EXPECT_EQ(blockManager.getNumDiskAdmissionRefused(), 0u); } diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index 633a90d08e5f..a572e4121859 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -625,6 +625,8 @@ def append_to_kv_heads_per_layer(num_kv_heads_per_layer: List[int], None) or "" kwargs['disk_cache_retained_only'] = bool( getattr(kv_cache_config, "disk_cache_retained_only", False)) + kwargs['disk_cache_protect_unexpired'] = bool( + getattr(kv_cache_config, "disk_cache_protect_unexpired", False)) self.impl = KVCacheManagerCpp(**kwargs) # Warmup baseline for cumulative counters (set by snapshot_warmup_baseline) self._warmup_reused_blocks = 0 diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 2b1cff9eb1f2..a940c6450e1e 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -3359,6 +3359,11 @@ class KvCacheConfig(StrictBaseModel, PybindMirror): description= "When true, only blocks whose request carried a disk retention TTL may enter the disk cache tier." ) + disk_cache_protect_unexpired: bool = Field( + default=False, + description= + "When true, unexpired disk blocks are never evicted to admit a new one; the new block is refused instead." + ) cross_kv_cache_fraction: Optional[float] = Field( default=None, description= @@ -3551,6 +3556,8 @@ def _to_pybind(self): config.disk_cache_path = self.disk_cache_path or "" config.disk_cache_retained_only = bool( self.disk_cache_retained_only) + config.disk_cache_protect_unexpired = bool( + self.disk_cache_protect_unexpired) return config @field_validator('free_gpu_memory_fraction') From 6158a3dc6be9e500402fef30a6bb63ffb5f056d6 Mon Sep 17 00:00:00 2001 From: Md Nafis Ul Haque Shifat Date: Sat, 4 Jul 2026 18:43:56 +0000 Subject: [PATCH 07/47] [#16251][fix] disk tier: fix spill-all eviction order (consume empty slots before displacing cached blocks) Empty disk slots initialized at default priority (bucket 35) while unmarked spilled content was filed at kMin (bucket 0). getFreeBlock returns the lowest bucket first, so claimDiskTarget displaced just-spilled cached blocks before using free empty slots -> unmarked content never survived to re-request (spill-all onboard ~0% under churn). Restore order empty(kMin) < unmarked(kDefault) < retained(kMax): - evictionPolicy initialize(): disk-level blocks start at kMin (+ setPriority in lockstep so claim/release erase from the correct queue). - reclaimSecondaryBlock: unmarked spilled content -> kDefault (above empties, below retained); evicted LRU only once disk is genuinely full. retained path unchanged. Verified: spill-all onboard 0%->99%, disk-tier unit tests pass, gate+expiry+retained-survival unaffected. Signed-off-by: Md Nafis Ul Haque Shifat --- cpp/tensorrt_llm/batch_manager/evictionPolicy.cpp | 13 +++++++++++-- cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp | 11 +++++++---- 2 files changed, 18 insertions(+), 6 deletions(-) diff --git a/cpp/tensorrt_llm/batch_manager/evictionPolicy.cpp b/cpp/tensorrt_llm/batch_manager/evictionPolicy.cpp index e58c05e280f3..d982b1b8d935 100644 --- a/cpp/tensorrt_llm/batch_manager/evictionPolicy.cpp +++ b/cpp/tensorrt_llm/batch_manager/evictionPolicy.cpp @@ -78,11 +78,20 @@ void LRUEvictionPolicy::initialize(std::vector& mAllBlocksById, std::v for (SizeType32 cacheLevel = 0; cacheLevel < kNumCacheLevels; cacheLevel++) { - auto& freeQueue = mFreeQueues[cacheLevel][defaultPriorityIdx]; + // Poolless disk slots start empty; give them the lowest priority so getFreeBlock + // hands out empty slots as spill targets before displacing any cached disk block. + bool const isDisk = (cacheLevel == kDiskLevel); + SizeType32 const initPriIdx = isDisk ? getPriorityIdx(kMinPriority) : defaultPriorityIdx; + auto& freeQueue = mFreeQueues[cacheLevel][initPriIdx]; for (SizeType32 blockId = 0; blockId < sizesPadded[cacheLevel]; blockId++) { - // Initialize all blocks to be the default priority level + // Keep block priority in lockstep with its bucket so later claim/release + // erase from the correct queue. + if (isDisk) + { + mAllBlocksById[startIdx + blockId]->setPriority(kMinPriority); + } mFreeBlockIterators[startIdx + blockId] = freeQueue.insert(freeQueue.end(), mAllBlocksById[startIdx + blockId]); } diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp index e8f1ea537943..b14c579dd0d0 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp @@ -1299,8 +1299,9 @@ void WindowBlockManager::releaseSubtree(BlockPtr const& block) BlockPtr WindowBlockManager::claimDiskTarget() { - // Cheap victims (empty slots; unmarked content in spill-all mode) live in bucket 0, - // so the policy's front candidate is only a retained block when nothing cheap exists. + // Priority order: empty slots (kMin) < unmarked content (kDefault) < retained (kMax), + // so getFreeBlock returns empties first and a retained block only when nothing cheaper + // (empty slot or unmarked cached block) remains. auto candidate = std::get<0>(mEvictionPolicy->getFreeBlock(kDiskLevel)); if (!candidate->isRetainedNow()) { @@ -1380,8 +1381,10 @@ BlockPtr WindowBlockManager::reclaimSecondaryBlock() } else { - // Spill-all mode only: unmarked content is displaced before any retained block. - victim->setPriority(executor::KvCacheRetentionConfig::kMinRetentionPriority); + // Spill-all: rank unmarked content above empty slots (kMin) but below retained + // (kMax) — empties are consumed as spill targets first, so cached blocks evict + // LRU only once the disk is genuinely full. + victim->setPriority(executor::KvCacheRetentionConfig::kDefaultRetentionPriority); } mEvictionPolicy->releaseBlock(victim); // re-enters the free queues at kDiskLevel ++mDiskSpills; From fe19f0b77b75f11926a5750ee3b0bb4deca7827a Mon Sep 17 00:00:00 2001 From: Md Nafis Ul Haque Shifat Date: Mon, 6 Jul 2026 04:57:18 +0000 Subject: [PATCH 08/47] [#16251][perf] disk tier: async per-slot store baseline (pre reserved-pool unstaged) Working per-slot async fire-and-forget store: background writer thread + per-slot files + staged memcpy + load-wait. GDS path still present. Snapshot taken before the unstaged reserved-block-pool rework so we can roll back cleanly. Signed-off-by: Md Nafis Ul Haque Shifat --- .../batch_manager/kvCacheTransferManager.h | 39 ++++ .../batch_manager/kvCacheTransferManager.cpp | 169 +++++++++++++++++- 2 files changed, 202 insertions(+), 6 deletions(-) diff --git a/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h b/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h index e7904287befe..8ea3ffe000fe 100644 --- a/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h +++ b/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h @@ -19,6 +19,12 @@ #include "tensorrt_llm/runtime/bufferManager.h" #include "tensorrt_llm/runtime/cudaEvent.h" +#include +#include +#include +#include +#include + namespace tr = tensorrt_llm::runtime; namespace kvc = tensorrt_llm::executor::kv_cache; @@ -85,6 +91,8 @@ class KVCacheTransferManager void syncTransfers(); + ~KVCacheTransferManager(); + //! \brief Get transfer stats accumulated since last call, and reset the counters. [[nodiscard]] KvCacheTransferStats getAndResetTransferStats(); @@ -132,6 +140,37 @@ class KVCacheTransferManager // Reference to parent loopback agent std::shared_ptr mLoopbackAgent; int mDeviceId; + // Disk-tier onboard (disk->GPU) uses GPUDirect Storage when TLLM_KV_DISK_GDS is set; + // otherwise POSIX staging. Read once at construction. + bool mDiskUseGds{std::getenv("TLLM_KV_DISK_GDS") != nullptr}; + + // ---- Disk-tier async store: spill host->disk OFF the scheduler thread ---- + // Gated by TLLM_KV_DISK_ASYNC_STORE. Unset => synchronous spill (current behavior). + bool const mAsyncDiskStore{std::getenv("TLLM_KV_DISK_ASYNC_STORE") != nullptr}; + + struct DiskWriteJob + { + std::string filename; + std::size_t bytes{0}; + std::vector staged; + }; + + std::thread mDiskWriter; + std::mutex mDiskMutex; + std::condition_variable mDiskQueueCv; + std::condition_variable mDiskInflightCv; + std::queue mDiskWriteQueue; + std::unordered_map mDiskInflight; + bool mDiskWriterStop{false}; + std::size_t const mDiskWriteQueueMax{[] + { + auto* e = std::getenv("TLLM_KV_DISK_WRITE_QUEUE"); + return e ? std::stoul(e) : 1024UL; + }()}; + + void diskWriterLoop(); + void enqueueDiskWrite(std::string filename, void const* src, std::size_t bytes); + void waitForDiskSlotWrites(std::string const& filename); // Cumulative transfer statistics, reset on each call to getAndResetTransferStats(). // Protected by mStatsMutex for thread-safe access. diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp index 338dc8bce490..2c1df10951ae 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp @@ -16,6 +16,7 @@ */ #include +#include #include "tensorrt_llm/batch_manager/kvCacheTransferManager.h" @@ -84,6 +85,24 @@ KVCacheTransferManager::KVCacheTransferManager( { TLLM_CUDA_CHECK(cudaGetDevice(&mDeviceId)); TLLM_CHECK(mDeviceId != -1); + if (mAsyncDiskStore) + { + mDiskWriter = std::thread(&KVCacheTransferManager::diskWriterLoop, this); + TLLM_LOG_INFO("[disk-tier] async store ENABLED (write-queue max=%zu)", mDiskWriteQueueMax); + } +} + +KVCacheTransferManager::~KVCacheTransferManager() +{ + if (mDiskWriter.joinable()) + { + { + std::lock_guard lock(mDiskMutex); + mDiskWriterStop = true; + } + mDiskQueueCv.notify_all(); + mDiskWriter.join(); + } } tr::ITensor::SharedPtr KVCacheTransferManager::computeBlockPointer( @@ -401,6 +420,96 @@ std::string diskSlotFilename(std::string const& directory, SizeType32 diskSlot, } } // namespace +void KVCacheTransferManager::diskWriterLoop() +{ + while (true) + { + DiskWriteJob job; + { + std::unique_lock lock(mDiskMutex); + mDiskQueueCv.wait(lock, [this] { return !mDiskWriteQueue.empty() || mDiskWriterStop; }); + if (mDiskWriteQueue.empty()) + { + if (mDiskWriterStop) + { + return; + } + continue; + } + job = std::move(mDiskWriteQueue.front()); + mDiskWriteQueue.pop(); + } + mDiskQueueCv.notify_all(); // a producer may be blocked waiting for queue space + + // The slow, writeback-throttle-prone part runs HERE, off the scheduler thread. + int fd = ::open(job.filename.c_str(), O_CREAT | O_WRONLY, 0644); + if (fd >= 0) + { + auto const written = ::pwrite(fd, job.staged.data(), job.bytes, 0); + ::close(fd); + if (written != static_cast(job.bytes)) + { + TLLM_LOG_ERROR("[disk-tier] short async write to %s (%zd/%zu)", job.filename.c_str(), + static_cast(written), job.bytes); + } + } + else + { + TLLM_LOG_ERROR("[disk-tier] cannot open %s for async write", job.filename.c_str()); + } + + { + std::lock_guard lock(mDiskMutex); + if (auto it = mDiskInflight.find(job.filename); it != mDiskInflight.end() && --it->second <= 0) + { + mDiskInflight.erase(it); + } + } + mDiskQueueCv.notify_all(); // wake a producer waiting to re-write this slot + mDiskInflightCv.notify_all(); // wake a loader waiting on this slot + } +} + +void KVCacheTransferManager::enqueueDiskWrite(std::string filename, void const* src, std::size_t bytes) +{ + DiskWriteJob job; + job.filename = std::move(filename); + job.bytes = bytes; + job.staged.resize(bytes); + std::memcpy(job.staged.data(), src, bytes); // fast host->host copy; frees the slot + + { + std::unique_lock lock(mDiskMutex); + // Backpressure + per-slot serialization: wait for queue room AND no earlier write + // to this same slot still outstanding (so a reused slot cannot race). + mDiskQueueCv.wait(lock, + [this, &job] + { + return (mDiskWriteQueue.size() < mDiskWriteQueueMax + && mDiskInflight.find(job.filename) == mDiskInflight.end()) + || mDiskWriterStop; + }); + if (mDiskWriterStop) + { + return; + } + ++mDiskInflight[job.filename]; // mark BEFORE the block becomes loadable + mDiskWriteQueue.push(std::move(job)); + } + mDiskQueueCv.notify_one(); +} + +void KVCacheTransferManager::waitForDiskSlotWrites(std::string const& filename) +{ + std::unique_lock lock(mDiskMutex); + mDiskInflightCv.wait(lock, + [this, &filename] + { + auto it = mDiskInflight.find(filename); + return it == mDiskInflight.end() || it->second <= 0; + }); +} + void KVCacheTransferManager::spillToFile(BlockPtr const& srcHostBlock, SizeType32 diskSlot, std::vector const& pools, std::string const& directory) { @@ -417,12 +526,20 @@ void KVCacheTransferManager::spillToFile(BlockPtr const& srcHostBlock, SizeType3 TLLM_CHECK_WITH_INFO(!pools[poolIdx].layerFirstLayout, "disk tier does not support layer-first layout pools"); auto ptr = computeBlockPointer(srcHostBlock, pools, poolIdx); auto const filename = diskSlotFilename(directory, diskSlot, poolIdx); - int fd = ::open(filename.c_str(), O_CREAT | O_WRONLY, 0644); - TLLM_CHECK_WITH_INFO(fd >= 0, "disk tier: cannot open %s", filename.c_str()); - auto const bytes = static_cast(ptr->getSizeInBytes()); - auto const written = ::pwrite(fd, ptr->data(), bytes, 0); - ::close(fd); - TLLM_CHECK_WITH_INFO(written == bytes, "disk tier: short write to %s", filename.c_str()); + auto const bytes = static_cast(ptr->getSizeInBytes()); + if (!mAsyncDiskStore) + { + // Synchronous path -- byte-for-byte the original behavior. + int fd = ::open(filename.c_str(), O_CREAT | O_WRONLY, 0644); + TLLM_CHECK_WITH_INFO(fd >= 0, "disk tier: cannot open %s", filename.c_str()); + auto const written = ::pwrite(fd, ptr->data(), bytes, 0); + ::close(fd); + TLLM_CHECK_WITH_INFO( + written == static_cast(bytes), "disk tier: short write to %s", filename.c_str()); + continue; + } + // Async path: copy bytes out (frees the slot immediately) and hand to the writer. + enqueueDiskWrite(filename, ptr->data(), bytes); } } @@ -439,6 +556,46 @@ void KVCacheTransferManager::loadFromFile(BlockPtr const& dstPrimaryBlock, SizeT { it->second.synchronize(); } + // Async store may still be persisting this slot; a read must not race the write. + if (mAsyncDiskStore) + { + for (size_t poolIdx = 0; poolIdx < pools.size(); ++poolIdx) + { + waitForDiskSlotWrites(diskSlotFilename(directory, diskSlot, poolIdx)); + } + } + if (mDiskUseGds) + { + // GDS onboard: direct SSD -> GPU VRAM via the NIXL loopback agent (cuFile). + // Falls back to POSIX staging if the GDS backend can't be created/run. + try + { + if (mLoopbackAgent == nullptr) + { + kvc::BaseAgentConfig config{std::string("GDSAgent"), true, true}; + mLoopbackAgent = kvc::makeLoopbackAgent("nixl", &config); + } + std::vector fileBlobs; + std::vector memoryBlobs; + for (size_t poolIdx = 0; poolIdx < pools.size(); ++poolIdx) + { + TLLM_CHECK_WITH_INFO( + !pools[poolIdx].layerFirstLayout, "disk tier does not support layer-first layout pools"); + auto ptr = computeBlockPointer(dstPrimaryBlock, pools, poolIdx); + fileBlobs.emplace_back( + diskSlotFilename(directory, diskSlot, poolIdx), O_RDONLY, 0664, ptr->getSizeInBytes()); + memoryBlobs.emplace_back(ptr->data(), ptr->getSizeInBytes(), mDeviceId); + } + kvc::FileDescs fileDescs(std::move(fileBlobs)); + kvc::MemoryDescs memoryDescs(kvc::MemoryType::kVRAM, memoryBlobs); + mLoopbackAgent->executeLoopbackRequest(memoryDescs, fileDescs, /*isOffload=*/false); + return; + } + catch (std::exception const& e) + { + TLLM_LOG_WARNING("disk tier: GDS onboard failed (%s); falling back to POSIX", e.what()); + } + } for (size_t poolIdx = 0; poolIdx < pools.size(); ++poolIdx) { TLLM_CHECK_WITH_INFO(!pools[poolIdx].layerFirstLayout, "disk tier does not support layer-first layout pools"); From f7c4fa0f65d9b524161e6c8af2dbbe81b892b516 Mon Sep 17 00:00:00 2001 From: Md Nafis Ul Haque Shifat Date: Mon, 6 Jul 2026 07:23:59 +0000 Subject: [PATCH 09/47] [#16251][perf] disk tier: unstaged reserved-block-pool store (pin source, hand out reserved slot) Async spill hands the writer the victim host slot directly (no staging memcpy), pins that slot until the write drains, and returns a pre-claimed reserved host block instead. Reaped drained slots rejoin the pool: one out, one in. Gated by TLLM_KV_DISK_RESERVED_BLOCKS (0=off=staged). Signed-off-by: Md Nafis Ul Haque Shifat --- .../batch_manager/kvCacheManager.h | 9 +++ .../batch_manager/kvCacheTransferManager.h | 25 +++++- .../batch_manager/kvCacheManager.cpp | 73 ++++++++++++++++- .../batch_manager/kvCacheTransferManager.cpp | 81 ++++++++++++++++++- 4 files changed, 185 insertions(+), 3 deletions(-) diff --git a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h index fa83e42dc7c8..f31fb7d5f86b 100644 --- a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h +++ b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h @@ -1483,6 +1483,15 @@ class WindowBlockManager std::priority_queue, std::greater> mDiskDeadlines; std::uint64_t mDiskSpillSeq{0}; + // ---- Unstaged async store: reserved host-block pool (env TLLM_KV_DISK_RESERVED_BLOCKS, 0 = off) ---- + // Target set from the environment in allocatePools. A spill hands out a reserved free host slot + // immediately and pins the victim's old slot until its async write drains, avoiding the staging + // memcpy on the scheduler thread. Reaped (drained) victim slots rejoin the pool: one out, one in. + SizeType32 mReservedHostBlockTarget{0}; + std::vector mReservedHostBlocks; + std::unordered_map mPendingSpillBlocks; + std::uint64_t mUnstagedSpillSeq{0}; + //! \brief Pick and claim the disk block to overwrite: empty slots first, then //! unmarked content, then the retained block with the earliest deadline. [[nodiscard]] BlockPtr claimDiskTarget(); diff --git a/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h b/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h index 8ea3ffe000fe..fab5b897024c 100644 --- a/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h +++ b/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h @@ -89,6 +89,22 @@ class KVCacheTransferManager void loadFromFile(BlockPtr const& dstPrimaryBlock, SizeType32 diskSlot, std::vector const& pools, std::string const& directory); + //! \brief Disk tier (unstaged async): spill a host block's bytes by handing the writer the source + //! pointers directly (no staging memcpy). The caller MUST keep the source host slot pinned until + //! \p spillId appears in drainCompletedSpills(). + void spillToFileUnstaged(BlockPtr const& srcHostBlock, SizeType32 diskSlot, + std::vector const& pools, std::string const& directory, std::uint64_t spillId); + + //! \brief Disk tier: return + clear the set of spill ids whose async writes fully completed since + //! the last call (unstaged path). The block manager reaps the corresponding pinned blocks. + [[nodiscard]] std::vector drainCompletedSpills(); + + //! \brief Whether the background async-store writer thread is active (TLLM_KV_DISK_ASYNC_STORE). + [[nodiscard]] bool asyncDiskStoreEnabled() const + { + return mAsyncDiskStore; + } + void syncTransfers(); ~KVCacheTransferManager(); @@ -152,7 +168,9 @@ class KVCacheTransferManager { std::string filename; std::size_t bytes{0}; - std::vector staged; + std::vector staged; // staged-memcpy path (used when src == nullptr) + void const* src{nullptr}; // unstaged path: writer reads this pinned host pointer directly + std::uint64_t spillId{0}; // >0 => track per-spill completion for the reserved-pool reap }; std::thread mDiskWriter; @@ -161,6 +179,10 @@ class KVCacheTransferManager std::condition_variable mDiskInflightCv; std::queue mDiskWriteQueue; std::unordered_map mDiskInflight; + // Unstaged reserved-pool support: outstanding pool-writes per spill; when a spill's last write + // completes, publish its id so the block manager can reap the pinned source block. + std::unordered_map mSpillRemaining; + std::vector mCompletedSpills; bool mDiskWriterStop{false}; std::size_t const mDiskWriteQueueMax{[] { @@ -170,6 +192,7 @@ class KVCacheTransferManager void diskWriterLoop(); void enqueueDiskWrite(std::string filename, void const* src, std::size_t bytes); + void enqueueDiskWriteUnstaged(std::string filename, void const* src, std::size_t bytes, std::uint64_t spillId); void waitForDiskSlotWrites(std::string const& filename); // Cumulative transfer statistics, reset on each call to getAndResetTransferStats(). diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp index b14c579dd0d0..36b3d157f794 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp @@ -1201,6 +1201,31 @@ void WindowBlockManager::allocatePools(bool useUvm) pool.secondaryPtr = BufferManager::pinned(cacheShapeOffload, poolDtype); } } + + // Disk-tier unstaged async store (env TLLM_KV_DISK_RESERVED_BLOCKS, 0 = off): carve a reserved pool + // of host blocks up-front so a spill can hand out a free host slot immediately while the victim's + // old slot is still draining to disk, avoiding the staging memcpy on the scheduler thread. Only + // meaningful with an active disk tier + async writer; capped below the host pool so the normal host + // cache is never starved. Reaped (drained) victim slots rejoin the pool: one out, one in. + if (auto const* e = std::getenv("TLLM_KV_DISK_RESERVED_BLOCKS")) + { + mReservedHostBlockTarget = static_cast(std::stoi(e)); + } + if (mNumDiskBlocks > 0 && mReservedHostBlockTarget > 0 && mNumSecondaryBlocks > 0 + && mTransferManager->asyncDiskStoreEnabled()) + { + auto const avail = mEvictionPolicy->getNumFreeBlocks(kSecondaryLevel); + SizeType32 const want = std::min(mReservedHostBlockTarget, std::max(0, avail - 16)); + for (SizeType32 i = 0; i < want; ++i) + { + auto blk = std::get<0>(mEvictionPolicy->getFreeBlock(kSecondaryLevel)); + mEvictionPolicy->claimBlock(blk); + mReservedHostBlocks.push_back(blk); + } + TLLM_LOG_INFO( + "[disk-tier] unstaged reserved host-block pool: reserved=%zu (requested=%d, avail=%d, windowSize=%d)", + mReservedHostBlocks.size(), mReservedHostBlockTarget, avail, mWindowSize); + } } void BlockManager::releasePools() @@ -1342,6 +1367,19 @@ BlockPtr WindowBlockManager::claimDiskTarget() BlockPtr WindowBlockManager::reclaimSecondaryBlock() { + // Reap completed unstaged spills first: their pinned host slots are safe again, so the drained + // victim blocks rejoin the reserved pool (one out, one in). + if (mReservedHostBlockTarget > 0) + { + for (auto const spillId : mTransferManager->drainCompletedSpills()) + { + if (auto it = mPendingSpillBlocks.find(spillId); it != mPendingSpillBlocks.end()) + { + mReservedHostBlocks.push_back(it->second); + mPendingSpillBlocks.erase(it); + } + } + } auto victim = std::get<0>(mEvictionPolicy->getFreeBlock(kSecondaryLevel)); mEvictionPolicy->claimBlock(victim); if (mNumDiskBlocks == 0 || victim->getUniqueTokens().empty() || mEvictionPolicy->getNumFreeBlocks(kDiskLevel) == 0) @@ -1369,7 +1407,21 @@ BlockPtr WindowBlockManager::reclaimSecondaryBlock() } return victim; } - mTransferManager->spillToFile(victim, diskTarget->getDiskSlot(), mPools, mDiskCachePath); + bool const useUnstaged = !mReservedHostBlocks.empty() && mTransferManager->asyncDiskStoreEnabled(); + SizeType32 const diskSlot = diskTarget->getDiskSlot(); + std::uint64_t unstagedSpillId = 0; + if (useUnstaged) + { + // Unstaged: the writer reads the victim's host memory directly (no staging memcpy). That slot + // must stay pinned until the write drains -- after the swap it is owned by diskTarget, which we + // hold below (instead of returning it) so nobody can reuse the slot mid-write. + unstagedSpillId = ++mUnstagedSpillSeq; + mTransferManager->spillToFileUnstaged(victim, diskSlot, mPools, mDiskCachePath, unstagedSpillId); + } + else + { + mTransferManager->spillToFile(victim, diskSlot, mPools, mDiskCachePath); + } victim->swapDiskResidency(diskTarget); // victim's identity now disk-resident, tree intact victim->setDurationMs(std::nullopt); // keep it out of upstream's expiring-block machinery if (victim->isRetainedNow()) @@ -1392,6 +1444,25 @@ BlockPtr WindowBlockManager::reclaimSecondaryBlock() { TLLM_LOG_INFO("[disk-tier] spills=%zu (windowSize=%d)", mDiskSpills, mWindowSize); } + if (useUnstaged) + { + // diskTarget now owns the victim's old host slot, which the writer is still reading. It will + // NOT flow through getFreeBlock's cleanup (we return a reserved block instead), so evict any + // displaced disk-block identity it still carries HERE, then pin it until the write drains. + { + std::lock_guard treeLock(mLookupTree->getMutex()); + if (mEventManager && blockInRadixTree(diskTarget)) + { + mEventManager->enqueueRemovedEvent(diskTarget, mWindowSize); + } + diskTarget->detachFromLookupNode(); + } + diskTarget->clearRetention(); + mPendingSpillBlocks[unstagedSpillId] = diskTarget; // pinned: keeps the host slot out of every free list + auto reserved = mReservedHostBlocks.back(); // clean, claimed, secondary host block + mReservedHostBlocks.pop_back(); + return reserved; + } diskTarget->clearRetention(); return diskTarget; // owns the victim's old host slot; already claimed } diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp index 2c1df10951ae..af620ec4292c 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp @@ -442,10 +442,12 @@ void KVCacheTransferManager::diskWriterLoop() mDiskQueueCv.notify_all(); // a producer may be blocked waiting for queue space // The slow, writeback-throttle-prone part runs HERE, off the scheduler thread. + // Unstaged jobs (src != nullptr) read the pinned host slot directly; staged jobs read the copy. + void const* data = job.src ? job.src : static_cast(job.staged.data()); int fd = ::open(job.filename.c_str(), O_CREAT | O_WRONLY, 0644); if (fd >= 0) { - auto const written = ::pwrite(fd, job.staged.data(), job.bytes, 0); + auto const written = ::pwrite(fd, data, job.bytes, 0); ::close(fd); if (written != static_cast(job.bytes)) { @@ -464,6 +466,16 @@ void KVCacheTransferManager::diskWriterLoop() { mDiskInflight.erase(it); } + // Unstaged path: when a spill's last pool-write finishes, its source host slot is safe to + // reuse -> publish the spill id for the block manager to reap. + if (job.spillId != 0) + { + if (auto sit = mSpillRemaining.find(job.spillId); sit != mSpillRemaining.end() && --sit->second <= 0) + { + mSpillRemaining.erase(sit); + mCompletedSpills.push_back(job.spillId); + } + } } mDiskQueueCv.notify_all(); // wake a producer waiting to re-write this slot mDiskInflightCv.notify_all(); // wake a loader waiting on this slot @@ -499,6 +511,43 @@ void KVCacheTransferManager::enqueueDiskWrite(std::string filename, void const* mDiskQueueCv.notify_one(); } +void KVCacheTransferManager::enqueueDiskWriteUnstaged( + std::string filename, void const* src, std::size_t bytes, std::uint64_t spillId) +{ + DiskWriteJob job; + job.filename = std::move(filename); + job.bytes = bytes; + job.src = src; // no staging: the writer reads the pinned host slot directly + job.spillId = spillId; + + { + std::unique_lock lock(mDiskMutex); + // Same backpressure + per-slot serialization as enqueueDiskWrite. + mDiskQueueCv.wait(lock, + [this, &job] + { + return (mDiskWriteQueue.size() < mDiskWriteQueueMax + && mDiskInflight.find(job.filename) == mDiskInflight.end()) + || mDiskWriterStop; + }); + if (mDiskWriterStop) + { + return; + } + ++mDiskInflight[job.filename]; // mark BEFORE the block becomes loadable + mDiskWriteQueue.push(std::move(job)); + } + mDiskQueueCv.notify_one(); +} + +std::vector KVCacheTransferManager::drainCompletedSpills() +{ + std::lock_guard lock(mDiskMutex); + std::vector out; + out.swap(mCompletedSpills); + return out; +} + void KVCacheTransferManager::waitForDiskSlotWrites(std::string const& filename) { std::unique_lock lock(mDiskMutex); @@ -543,6 +592,36 @@ void KVCacheTransferManager::spillToFile(BlockPtr const& srcHostBlock, SizeType3 } } +void KVCacheTransferManager::spillToFileUnstaged(BlockPtr const& srcHostBlock, SizeType32 diskSlot, + std::vector const& pools, std::string const& directory, std::uint64_t spillId) +{ + TLLM_CHECK_WITH_INFO(!directory.empty(), "disk tier requires a directory"); + TLLM_CHECK_WITH_INFO(mAsyncDiskStore, "spillToFileUnstaged requires the async writer"); + // Same event discipline as spillToFile: the victim's bytes may still be the target of an in-flight + // GPU->host copy; block until it lands before the writer reads host memory. + auto const idx = getPendingTransferIndex(srcHostBlock); + if (auto it = mPendingWrites.find(idx); it != mPendingWrites.end()) + { + it->second.synchronize(); + } + // Register the whole slot's write count up-front, so an early single-pool completion cannot publish + // the spill before the remaining pools are enqueued. + { + std::lock_guard lock(mDiskMutex); + mSpillRemaining[spillId] = static_cast(pools.size()); + } + for (size_t poolIdx = 0; poolIdx < pools.size(); ++poolIdx) + { + TLLM_CHECK_WITH_INFO(!pools[poolIdx].layerFirstLayout, "disk tier does not support layer-first layout pools"); + auto ptr = computeBlockPointer(srcHostBlock, pools, poolIdx); + auto const filename = diskSlotFilename(directory, diskSlot, poolIdx); + auto const bytes = static_cast(ptr->getSizeInBytes()); + // The pool buffer outlives the manager, so this raw pointer stays valid after `ptr` (a view) + // is destroyed -- the writer reads it later off-thread. + enqueueDiskWriteUnstaged(filename, ptr->data(), bytes, spillId); + } +} + void KVCacheTransferManager::loadFromFile(BlockPtr const& dstPrimaryBlock, SizeType32 diskSlot, std::vector const& pools, std::string const& directory) { From 42886565594d4ecafb963ce63bd83ffaadfe04c1 Mon Sep 17 00:00:00 2001 From: Md Nafis Ul Haque Shifat Date: Mon, 6 Jul 2026 08:13:35 +0000 Subject: [PATCH 10/47] [#16251][chore] disk tier: profiling instrumentation (load-wait vs enqueue-backpressure vs queue depth) Signed-off-by: Md Nafis Ul Haque Shifat --- .../batch_manager/kvCacheTransferManager.h | 13 +++++ .../batch_manager/kvCacheTransferManager.cpp | 56 +++++++++++++++++++ 2 files changed, 69 insertions(+) diff --git a/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h b/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h index fab5b897024c..99fb3c8a3e7a 100644 --- a/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h +++ b/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h @@ -19,6 +19,7 @@ #include "tensorrt_llm/runtime/bufferManager.h" #include "tensorrt_llm/runtime/cudaEvent.h" +#include #include #include #include @@ -183,6 +184,18 @@ class KVCacheTransferManager // completes, publish its id so the block manager can reap the pinned source block. std::unordered_map mSpillRemaining; std::vector mCompletedSpills; + // ---- disk async-store profiling: rate-limited [disk-prof] line, load-side vs write-side stalls ---- + std::size_t mProfSpills{0}; + std::size_t mProfQueuePeak{0}; + std::size_t mProfEnqBlkCount{0}; + std::uint64_t mProfEnqBlkTotNs{0}; + std::uint64_t mProfEnqBlkMaxNs{0}; + std::size_t mProfLoadCalls{0}; + std::size_t mProfLoadWaitCount{0}; + std::uint64_t mProfLoadWaitTotNs{0}; + std::uint64_t mProfLoadWaitMaxNs{0}; + std::chrono::steady_clock::time_point mProfLast{}; + void reportDiskProfLocked(); // caller holds mDiskMutex bool mDiskWriterStop{false}; std::size_t const mDiskWriteQueueMax{[] { diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp index af620ec4292c..e0f10c077cea 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp @@ -15,6 +15,8 @@ * limitations under the License. */ +#include +#include #include #include @@ -438,6 +440,7 @@ void KVCacheTransferManager::diskWriterLoop() } job = std::move(mDiskWriteQueue.front()); mDiskWriteQueue.pop(); + mProfQueuePeak = std::max(mProfQueuePeak, mDiskWriteQueue.size() + 1); } mDiskQueueCv.notify_all(); // a producer may be blocked waiting for queue space @@ -476,6 +479,8 @@ void KVCacheTransferManager::diskWriterLoop() mCompletedSpills.push_back(job.spillId); } } + ++mProfSpills; + reportDiskProfLocked(); } mDiskQueueCv.notify_all(); // wake a producer waiting to re-write this slot mDiskInflightCv.notify_all(); // wake a loader waiting on this slot @@ -523,6 +528,7 @@ void KVCacheTransferManager::enqueueDiskWriteUnstaged( { std::unique_lock lock(mDiskMutex); // Same backpressure + per-slot serialization as enqueueDiskWrite. + auto const t0 = std::chrono::steady_clock::now(); mDiskQueueCv.wait(lock, [this, &job] { @@ -530,6 +536,14 @@ void KVCacheTransferManager::enqueueDiskWriteUnstaged( && mDiskInflight.find(job.filename) == mDiskInflight.end()) || mDiskWriterStop; }); + auto const blkNs + = std::chrono::duration_cast(std::chrono::steady_clock::now() - t0).count(); + if (blkNs > 100000) + { + ++mProfEnqBlkCount; + mProfEnqBlkTotNs += static_cast(blkNs); + mProfEnqBlkMaxNs = std::max(mProfEnqBlkMaxNs, static_cast(blkNs)); + } if (mDiskWriterStop) { return; @@ -548,15 +562,57 @@ std::vector KVCacheTransferManager::drainCompletedSpills() return out; } +void KVCacheTransferManager::reportDiskProfLocked() +{ + auto const now = std::chrono::steady_clock::now(); + if (mProfLast.time_since_epoch().count() == 0) + { + mProfLast = now; + return; + } + auto const sec = std::chrono::duration_cast(now - mProfLast).count(); + if (sec < 5) + { + return; + } + auto const ms = [](std::uint64_t ns) { return static_cast(ns) / 1.0e6; }; + TLLM_LOG_INFO( + "[disk-prof] %llds spills=%zu qpeak=%zu | enqBlk n=%zu tot=%.0fms max=%.0fms | " + "load calls=%zu waited=%zu tot=%.0fms max=%.0fms", + static_cast(sec), mProfSpills, mProfQueuePeak, mProfEnqBlkCount, ms(mProfEnqBlkTotNs), + ms(mProfEnqBlkMaxNs), mProfLoadCalls, mProfLoadWaitCount, ms(mProfLoadWaitTotNs), ms(mProfLoadWaitMaxNs)); + mProfSpills = 0; + mProfQueuePeak = 0; + mProfEnqBlkCount = 0; + mProfEnqBlkTotNs = 0; + mProfEnqBlkMaxNs = 0; + mProfLoadCalls = 0; + mProfLoadWaitCount = 0; + mProfLoadWaitTotNs = 0; + mProfLoadWaitMaxNs = 0; + mProfLast = now; +} + void KVCacheTransferManager::waitForDiskSlotWrites(std::string const& filename) { std::unique_lock lock(mDiskMutex); + auto const t0 = std::chrono::steady_clock::now(); mDiskInflightCv.wait(lock, [this, &filename] { auto it = mDiskInflight.find(filename); return it == mDiskInflight.end() || it->second <= 0; }); + auto const wns + = std::chrono::duration_cast(std::chrono::steady_clock::now() - t0).count(); + ++mProfLoadCalls; + if (wns > 100000) + { + ++mProfLoadWaitCount; + mProfLoadWaitTotNs += static_cast(wns); + mProfLoadWaitMaxNs = std::max(mProfLoadWaitMaxNs, static_cast(wns)); + } + reportDiskProfLocked(); } void KVCacheTransferManager::spillToFile(BlockPtr const& srcHostBlock, SizeType32 diskSlot, From 629277ec220472d73c2a1f6868ee233637c41ff6 Mon Sep 17 00:00:00 2001 From: Md Nafis Ul Haque Shifat Date: Mon, 6 Jul 2026 17:14:46 +0000 Subject: [PATCH 11/47] [#16251][perf] disk tier: N background writer threads on the shared queue (TLLM_KV_DISK_WRITERS) Single writer saturated the queue under eviction bursts (qpeak pinned at 1024, multi-second enqueue backpressure -> TTFT spikes). Spawn N workers draining the shared queue; idle workers pull the next job (auto-balanced), per-file in-flight gate keeps same-slot writes serialized. Default 1. Signed-off-by: Md Nafis Ul Haque Shifat --- .../batch_manager/kvCacheTransferManager.h | 9 ++++++++- .../batch_manager/kvCacheTransferManager.cpp | 18 ++++++++++++++---- 2 files changed, 22 insertions(+), 5 deletions(-) diff --git a/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h b/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h index 99fb3c8a3e7a..b0b342d6760d 100644 --- a/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h +++ b/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h @@ -174,7 +174,7 @@ class KVCacheTransferManager std::uint64_t spillId{0}; // >0 => track per-spill completion for the reserved-pool reap }; - std::thread mDiskWriter; + std::vector mDiskWriters; std::mutex mDiskMutex; std::condition_variable mDiskQueueCv; std::condition_variable mDiskInflightCv; @@ -202,6 +202,13 @@ class KVCacheTransferManager auto* e = std::getenv("TLLM_KV_DISK_WRITE_QUEUE"); return e ? std::stoul(e) : 1024UL; }()}; + // Number of background writer threads draining the shared queue (TLLM_KV_DISK_WRITERS, default 1). + std::size_t const mNumDiskWriters{[] + { + auto* e = std::getenv("TLLM_KV_DISK_WRITERS"); + auto n = e ? std::stoul(e) : 1UL; + return n < 1UL ? 1UL : n; + }()}; void diskWriterLoop(); void enqueueDiskWrite(std::string filename, void const* src, std::size_t bytes); diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp index e0f10c077cea..b2385ffb9a68 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp @@ -89,21 +89,31 @@ KVCacheTransferManager::KVCacheTransferManager( TLLM_CHECK(mDeviceId != -1); if (mAsyncDiskStore) { - mDiskWriter = std::thread(&KVCacheTransferManager::diskWriterLoop, this); - TLLM_LOG_INFO("[disk-tier] async store ENABLED (write-queue max=%zu)", mDiskWriteQueueMax); + for (std::size_t i = 0; i < mNumDiskWriters; ++i) + { + mDiskWriters.emplace_back(&KVCacheTransferManager::diskWriterLoop, this); + } + TLLM_LOG_INFO( + "[disk-tier] async store ENABLED (writers=%zu, write-queue max=%zu)", mNumDiskWriters, mDiskWriteQueueMax); } } KVCacheTransferManager::~KVCacheTransferManager() { - if (mDiskWriter.joinable()) + if (!mDiskWriters.empty()) { { std::lock_guard lock(mDiskMutex); mDiskWriterStop = true; } mDiskQueueCv.notify_all(); - mDiskWriter.join(); + for (auto& t : mDiskWriters) + { + if (t.joinable()) + { + t.join(); + } + } } } From 331d6784a12902c79753c04b1bd12478799333c8 Mon Sep 17 00:00:00 2001 From: Md Nafis Ul Haque Shifat Date: Mon, 6 Jul 2026 17:15:55 +0000 Subject: [PATCH 12/47] [#16251][chore] disk tier: drop profiling instrumentation Signed-off-by: Md Nafis Ul Haque Shifat --- .../batch_manager/kvCacheTransferManager.h | 13 ----- .../batch_manager/kvCacheTransferManager.cpp | 56 ------------------- 2 files changed, 69 deletions(-) diff --git a/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h b/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h index b0b342d6760d..5c73ee9b10fa 100644 --- a/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h +++ b/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h @@ -19,7 +19,6 @@ #include "tensorrt_llm/runtime/bufferManager.h" #include "tensorrt_llm/runtime/cudaEvent.h" -#include #include #include #include @@ -184,18 +183,6 @@ class KVCacheTransferManager // completes, publish its id so the block manager can reap the pinned source block. std::unordered_map mSpillRemaining; std::vector mCompletedSpills; - // ---- disk async-store profiling: rate-limited [disk-prof] line, load-side vs write-side stalls ---- - std::size_t mProfSpills{0}; - std::size_t mProfQueuePeak{0}; - std::size_t mProfEnqBlkCount{0}; - std::uint64_t mProfEnqBlkTotNs{0}; - std::uint64_t mProfEnqBlkMaxNs{0}; - std::size_t mProfLoadCalls{0}; - std::size_t mProfLoadWaitCount{0}; - std::uint64_t mProfLoadWaitTotNs{0}; - std::uint64_t mProfLoadWaitMaxNs{0}; - std::chrono::steady_clock::time_point mProfLast{}; - void reportDiskProfLocked(); // caller holds mDiskMutex bool mDiskWriterStop{false}; std::size_t const mDiskWriteQueueMax{[] { diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp index b2385ffb9a68..f016cbf341c2 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp @@ -15,8 +15,6 @@ * limitations under the License. */ -#include -#include #include #include @@ -450,7 +448,6 @@ void KVCacheTransferManager::diskWriterLoop() } job = std::move(mDiskWriteQueue.front()); mDiskWriteQueue.pop(); - mProfQueuePeak = std::max(mProfQueuePeak, mDiskWriteQueue.size() + 1); } mDiskQueueCv.notify_all(); // a producer may be blocked waiting for queue space @@ -489,8 +486,6 @@ void KVCacheTransferManager::diskWriterLoop() mCompletedSpills.push_back(job.spillId); } } - ++mProfSpills; - reportDiskProfLocked(); } mDiskQueueCv.notify_all(); // wake a producer waiting to re-write this slot mDiskInflightCv.notify_all(); // wake a loader waiting on this slot @@ -538,7 +533,6 @@ void KVCacheTransferManager::enqueueDiskWriteUnstaged( { std::unique_lock lock(mDiskMutex); // Same backpressure + per-slot serialization as enqueueDiskWrite. - auto const t0 = std::chrono::steady_clock::now(); mDiskQueueCv.wait(lock, [this, &job] { @@ -546,14 +540,6 @@ void KVCacheTransferManager::enqueueDiskWriteUnstaged( && mDiskInflight.find(job.filename) == mDiskInflight.end()) || mDiskWriterStop; }); - auto const blkNs - = std::chrono::duration_cast(std::chrono::steady_clock::now() - t0).count(); - if (blkNs > 100000) - { - ++mProfEnqBlkCount; - mProfEnqBlkTotNs += static_cast(blkNs); - mProfEnqBlkMaxNs = std::max(mProfEnqBlkMaxNs, static_cast(blkNs)); - } if (mDiskWriterStop) { return; @@ -572,57 +558,15 @@ std::vector KVCacheTransferManager::drainCompletedSpills() return out; } -void KVCacheTransferManager::reportDiskProfLocked() -{ - auto const now = std::chrono::steady_clock::now(); - if (mProfLast.time_since_epoch().count() == 0) - { - mProfLast = now; - return; - } - auto const sec = std::chrono::duration_cast(now - mProfLast).count(); - if (sec < 5) - { - return; - } - auto const ms = [](std::uint64_t ns) { return static_cast(ns) / 1.0e6; }; - TLLM_LOG_INFO( - "[disk-prof] %llds spills=%zu qpeak=%zu | enqBlk n=%zu tot=%.0fms max=%.0fms | " - "load calls=%zu waited=%zu tot=%.0fms max=%.0fms", - static_cast(sec), mProfSpills, mProfQueuePeak, mProfEnqBlkCount, ms(mProfEnqBlkTotNs), - ms(mProfEnqBlkMaxNs), mProfLoadCalls, mProfLoadWaitCount, ms(mProfLoadWaitTotNs), ms(mProfLoadWaitMaxNs)); - mProfSpills = 0; - mProfQueuePeak = 0; - mProfEnqBlkCount = 0; - mProfEnqBlkTotNs = 0; - mProfEnqBlkMaxNs = 0; - mProfLoadCalls = 0; - mProfLoadWaitCount = 0; - mProfLoadWaitTotNs = 0; - mProfLoadWaitMaxNs = 0; - mProfLast = now; -} - void KVCacheTransferManager::waitForDiskSlotWrites(std::string const& filename) { std::unique_lock lock(mDiskMutex); - auto const t0 = std::chrono::steady_clock::now(); mDiskInflightCv.wait(lock, [this, &filename] { auto it = mDiskInflight.find(filename); return it == mDiskInflight.end() || it->second <= 0; }); - auto const wns - = std::chrono::duration_cast(std::chrono::steady_clock::now() - t0).count(); - ++mProfLoadCalls; - if (wns > 100000) - { - ++mProfLoadWaitCount; - mProfLoadWaitTotNs += static_cast(wns); - mProfLoadWaitMaxNs = std::max(mProfLoadWaitMaxNs, static_cast(wns)); - } - reportDiskProfLocked(); } void KVCacheTransferManager::spillToFile(BlockPtr const& srcHostBlock, SizeType32 diskSlot, From c0df24dda33186bfc96a1aa287a3c15970414ed0 Mon Sep 17 00:00:00 2001 From: Md Nafis Ul Haque Shifat Date: Tue, 7 Jul 2026 00:36:21 +0000 Subject: [PATCH 13/47] [#16251][perf] batch_manager: detach disk-tier onboard from the scheduler thread Disk onboards read SSD->GPU synchronously on the scheduler thread (or, with a reader pool, in parallel but still barriered before the forward), stalling every request in the step behind one slow read. Engine (C++): a disk onboard hands each of a block's pool reads to the reader pool and returns; completion is tracked per block, keyed by the matched-identity block id so requests reusing a prefix gate on the same key. The reader makes its copy device-complete before publishing, and GDS/POSIX both flow through the pool. areBlocksReady(requestId) reports when a request's blocks have landed; the read barrier (waitForAllReads/mReadInflight) is removed. Scheduler (Python): a context request whose onboard reads are still in flight is held out of the forward pass and re-checked each step -- the capacity scheduler counts its already-allocated blocks and _context_seq_len skips re-adding it -- then admitted once areBlocksReady(). Tests: detached reads byte-exact incl. concurrent draining; block-manager reuse matches the synchronous path and its readiness gate resolves. Signed-off-by: Md Nafis Ul Haque Shifat --- .../batch_manager/kvCacheManager.h | 12 ++ .../batch_manager/kvCacheTransferManager.h | 53 +++++- .../batch_manager/kvCacheManager.cpp | 41 ++++- .../batch_manager/kvCacheTransferManager.cpp | 165 +++++++++++++++++- .../nanobind/batch_manager/kvCacheManager.cpp | 1 + .../batch_manager/kvCacheManagerTest.cpp | 142 +++++++++++++++ tensorrt_llm/_torch/pyexecutor/py_executor.py | 33 ++++ .../_torch/pyexecutor/resource_manager.py | 11 ++ 8 files changed, 453 insertions(+), 5 deletions(-) diff --git a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h index f31fb7d5f86b..d483a7830fe4 100644 --- a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h +++ b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h @@ -1269,6 +1269,9 @@ class WindowBlockManager void onboardBlock(GenerationRequest& sequence, BlockPtr const& offloadBlock, executor::KvCacheTransferMode mode = executor::KvCacheTransferMode::DRAM, std::string const& directory = ""); + //! \brief True when every block held by requestId has its disk read landed (safe to forward). + [[nodiscard]] bool areBlocksReady(LlmRequest::RequestIdType requestId); + //! \brief Bring block from primary to secondary memory. //! \details Does nothing if block is already in secondary memory. void offloadBlock(BlockPtr const& block, executor::KvCacheTransferMode mode = executor::KvCacheTransferMode::DRAM, @@ -1713,6 +1716,9 @@ class BlockManager void onboardBlock(GenerationRequest& sequence, BlockPtr const& offloadBlock, SizeType32 windowSize, executor::KvCacheTransferMode mode = executor::KvCacheTransferMode::DRAM, std::string const& directory = ""); + //! \brief True when every block held by requestId (across windows) has its disk read landed. + [[nodiscard]] bool areBlocksReady(LlmRequest::RequestIdType requestId); + //! \brief Bring block from primary to secondary memory for window size. //! \details Does nothing if block is already in secondary memory. void offloadBlock(BlockPtr const& block, SizeType32 windowSize, @@ -2538,6 +2544,12 @@ class KVCacheManager : public BaseKVCacheManager return mBlockManager.getNumFreeBlocks(); } + //! \brief True when every KV block held by requestId has its disk read landed (detached onboard). + [[nodiscard]] bool areBlocksReady(LlmRequest::RequestIdType requestId) + { + return mBlockManager.areBlocksReady(requestId); + } + [[nodiscard]] SizeType32 getNumPools() const override { return mBlockManager.getNumPools(); diff --git a/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h b/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h index 5c73ee9b10fa..aff0533a623b 100644 --- a/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h +++ b/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h @@ -24,6 +24,9 @@ #include #include #include +#include +#include +#include namespace tr = tensorrt_llm::runtime; namespace kvc = tensorrt_llm::executor::kv_cache; @@ -85,9 +88,24 @@ class KVCacheTransferManager void spillToFile(BlockPtr const& srcHostBlock, SizeType32 diskSlot, std::vector const& pools, std::string const& directory); - //! \brief Disk tier: read slot files into a GPU-resident block (synchronous POSIX). + //! \brief Disk tier: read a block's slot files into GPU memory. With a reader pool the block's reads are + //! handed off and tracked under \p trackBlockId (POSIX or GDS, decided by the reader); the owning request + //! must be gated on isBlockReadPending() / areBlocksReady() before it is forwarded. With no pool + //! (\p trackBlockId < 0) the read is synchronous on the calling thread. void loadFromFile(BlockPtr const& dstPrimaryBlock, SizeType32 diskSlot, std::vector const& pools, - std::string const& directory); + std::string const& directory, std::int32_t trackBlockId = -1); + + //! \brief Return + clear the block ids whose reads have landed since the last call. + [[nodiscard]] std::vector drainCompletedBlockReads(); + + //! \brief True while block blockId still has an in-flight disk read (not yet safe to read on the GPU). + [[nodiscard]] bool isBlockReadPending(std::int32_t blockId); + + //! \brief True when a reader pool is present, so onboards can be detached from the scheduler thread. + [[nodiscard]] bool asyncDiskReadEnabled() const + { + return mNumDiskReaders >= 1; + } //! \brief Disk tier (unstaged async): spill a host block's bytes by handing the writer the source //! pointers directly (no staging memcpy). The caller MUST keep the source host slot pinned until @@ -202,6 +220,37 @@ class KVCacheTransferManager void enqueueDiskWriteUnstaged(std::string filename, void const* src, std::size_t bytes, std::uint64_t spillId); void waitForDiskSlotWrites(std::string const& filename); + // ---- Disk-tier async ONBOARD: read slot files disk->GPU OFF the scheduler thread ---- + // Gated by TLLM_KV_DISK_READERS (0/unset => synchronous read = current behavior; N>=1 => N reader + // threads drain a shared read queue). A job carries all of one block's pools; a request is forward-safe + // once areBlocksReady() reports every block it holds has landed. See loadFromFile / isBlockReadPending. + struct DiskReadJob + { + std::int32_t blockId{-1}; // block whose readiness this job satisfies (tracking key) + bool useGds{false}; // GDS DMA (true) vs POSIX read + H2D copy (false) + std::vector dsts; // GPU destination pointer, per pool + std::vector files; // disk slot file, per pool + std::vector bytes; // byte count, per pool + }; + + std::size_t const mNumDiskReaders{[] + { + auto* e = std::getenv("TLLM_KV_DISK_READERS"); + return e ? std::stoul(e) : 0UL; + }()}; + std::vector mDiskReaders; + std::mutex mReadMutex; + std::condition_variable mReadQueueCv; + std::queue mReadQueue; + // Blocks with a disk read in flight, keyed by the matched-identity block id so every request reusing that + // prefix gates on the same key. A block is forward-safe once it leaves this set. + std::unordered_set mPendingBlockReads; + std::vector mCompletedBlockReads; + bool mDiskReaderStop{false}; + + void diskReaderLoop(); + void enqueueDiskRead(DiskReadJob job); + // Cumulative transfer statistics, reset on each call to getAndResetTransferStats(). // Protected by mStatsMutex for thread-safe access. mutable std::mutex mStatsMutex; diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp index 36b3d157f794..54c9455310d8 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp @@ -1596,6 +1596,19 @@ void BlockManager::onboardBlock(GenerationRequest& sequence, BlockPtr const& off mWindowBlockManagers.at(windowSize).onboardBlock(sequence, offloadBlock, mode, directory); } +bool BlockManager::areBlocksReady(LlmRequest::RequestIdType requestId) +{ + // Forward-safe only when every window's blocks for this request have their disk reads landed. + for (auto& [windowSize, manager] : mWindowBlockManagers) + { + if (!manager.areBlocksReady(requestId)) + { + return false; + } + } + return true; +} + void WindowBlockManager::onboardBlock(GenerationRequest& sequence, BlockPtr const& offloadBlock, executor::KvCacheTransferMode mode, std::string const& directory) { @@ -1604,8 +1617,13 @@ void WindowBlockManager::onboardBlock(GenerationRequest& sequence, BlockPtr cons // A fresh GPU block; this may itself trigger the GPU->host cascade (unchanged, correct). auto block = getFreeBlock( sequence, executor::KvCacheRetentionConfig::kDefaultRetentionPriority, std::nullopt, mode, directory); - mTransferManager->loadFromFile(block, offloadBlock->getDiskSlot(), mPools, mDiskCachePath); - offloadBlock->swapDiskResidency(block); // matched identity now GPU-resident + // Detach the onboard when a reader pool is present: hand the read off (POSIX or GDS, decided by the + // reader) tagged with the matched identity's block id (stable across swapDiskResidency), so every + // request reusing this prefix gates on the same key. No pool -> synchronous read inside loadFromFile. + auto const trackId = mTransferManager->asyncDiskReadEnabled() ? offloadBlock->getBlockId() + : static_cast(-1); + mTransferManager->loadFromFile(block, offloadBlock->getDiskSlot(), mPools, mDiskCachePath, trackId); + offloadBlock->swapDiskResidency(block); // matched identity now GPU-resident (bytes may still be in flight) offloadBlock->setPriority(executor::KvCacheRetentionConfig::kDefaultRetentionPriority); block->clearRetention(); block->setPriority(executor::KvCacheRetentionConfig::kMinRetentionPriority); // freed card = empty slot again @@ -1636,6 +1654,25 @@ void WindowBlockManager::onboardBlock(GenerationRequest& sequence, BlockPtr cons } } +bool WindowBlockManager::areBlocksReady(LlmRequest::RequestIdType requestId) +{ + // A request is forward-safe only when every block it holds has its disk read landed. Keyed by block id, + // so a request sharing an in-flight onboarded block waits on the same key as the request that started it. + auto it = mAllocatedBlocksPerSeq.find(requestId); + if (it == mAllocatedBlocksPerSeq.end()) + { + return true; + } + for (auto const& block : it->second) + { + if (mTransferManager->isBlockReadPending(block->getBlockId())) + { + return false; + } + } + return true; +} + void BlockManager::offloadBlock( BlockPtr const& block, SizeType32 windowSize, executor::KvCacheTransferMode mode, std::string const& directory) { diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp index f016cbf341c2..228ddf51cfca 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp @@ -94,6 +94,14 @@ KVCacheTransferManager::KVCacheTransferManager( TLLM_LOG_INFO( "[disk-tier] async store ENABLED (writers=%zu, write-queue max=%zu)", mNumDiskWriters, mDiskWriteQueueMax); } + if (asyncDiskReadEnabled()) + { + for (std::size_t i = 0; i < mNumDiskReaders; ++i) + { + mDiskReaders.emplace_back(&KVCacheTransferManager::diskReaderLoop, this); + } + TLLM_LOG_INFO("[disk-tier] async onboard ENABLED (readers=%zu)", mNumDiskReaders); + } } KVCacheTransferManager::~KVCacheTransferManager() @@ -113,6 +121,21 @@ KVCacheTransferManager::~KVCacheTransferManager() } } } + if (!mDiskReaders.empty()) + { + { + std::lock_guard lock(mReadMutex); + mDiskReaderStop = true; + } + mReadQueueCv.notify_all(); + for (auto& t : mDiskReaders) + { + if (t.joinable()) + { + t.join(); + } + } + } } tr::ITensor::SharedPtr KVCacheTransferManager::computeBlockPointer( @@ -569,6 +592,123 @@ void KVCacheTransferManager::waitForDiskSlotWrites(std::string const& filename) }); } +void KVCacheTransferManager::diskReaderLoop() +{ + // Each reader owns its own CUDA stream (H2D copies stay off the null stream and the onboard/offload + // managers) and its own GDS agent (readers never share one). cudaStreamSynchronize (POSIX) or the agent's + // internal wait (GDS) makes each transfer device-complete before the block is published, so per-block + // readiness alone is a sufficient pre-forward gate. + cudaStream_t stream; + TLLM_CUDA_CHECK(cudaStreamCreate(&stream)); + std::vector hostBuffer; + std::shared_ptr gdsAgent; // per-reader GDS agent, created lazily on the first GDS job + while (true) + { + DiskReadJob job; + { + std::unique_lock lock(mReadMutex); + mReadQueueCv.wait(lock, [this] { return !mReadQueue.empty() || mDiskReaderStop; }); + if (mReadQueue.empty()) + { + if (mDiskReaderStop) + { + break; + } + continue; + } + job = std::move(mReadQueue.front()); + mReadQueue.pop(); + } + + // Perform the block's transfer. GDS issues one batched SSD->GPU DMA for all pools; POSIX reads each + // pool into a host buffer and copies it up. Both are device-complete before we publish the block. + if (job.useGds) + { + try + { + if (gdsAgent == nullptr) + { + kvc::BaseAgentConfig config{std::string("GDSReaderAgent"), true, true}; + gdsAgent = kvc::makeLoopbackAgent("nixl", &config); + } + std::vector fileBlobs; + std::vector memoryBlobs; + for (size_t i = 0; i < job.dsts.size(); ++i) + { + fileBlobs.emplace_back(job.files[i], O_RDONLY, 0664, job.bytes[i]); + memoryBlobs.emplace_back(job.dsts[i], job.bytes[i], mDeviceId); + } + kvc::FileDescs fileDescs(std::move(fileBlobs)); + kvc::MemoryDescs memoryDescs(kvc::MemoryType::kVRAM, memoryBlobs); + gdsAgent->executeLoopbackRequest(memoryDescs, fileDescs, /*isOffload=*/false); + } + catch (std::exception const& e) + { + TLLM_LOG_WARNING("disk tier: GDS onboard failed (%s); falling back to POSIX", e.what()); + job.useGds = false; // fall through to POSIX for this job + } + } + if (!job.useGds) + { + for (size_t i = 0; i < job.dsts.size(); ++i) + { + int fd = ::open(job.files[i].c_str(), O_RDONLY); + if (fd < 0) + { + TLLM_LOG_ERROR("[disk-tier] cannot open %s for async read", job.files[i].c_str()); + continue; + } + hostBuffer.resize(job.bytes[i]); + auto const got = ::read(fd, hostBuffer.data(), job.bytes[i]); + ::close(fd); + if (got == static_cast(job.bytes[i])) + { + TLLM_CUDA_CHECK( + cudaMemcpyAsync(job.dsts[i], hostBuffer.data(), job.bytes[i], cudaMemcpyHostToDevice, stream)); + } + else + { + TLLM_LOG_ERROR("[disk-tier] short async read from %s (%zd/%zu)", job.files[i].c_str(), + static_cast(got), job.bytes[i]); + } + } + TLLM_CUDA_CHECK(cudaStreamSynchronize(stream)); + } + + { + std::lock_guard lock(mReadMutex); + // The transfer is device-complete, so publish the block: any request holding it can be forwarded. + mPendingBlockReads.erase(job.blockId); + mCompletedBlockReads.push_back(job.blockId); + } + } + cudaStreamDestroy(stream); +} + +void KVCacheTransferManager::enqueueDiskRead(DiskReadJob job) +{ + { + std::lock_guard lock(mReadMutex); + mPendingBlockReads.insert(job.blockId); // mark pending before the job becomes poppable + mReadQueue.push(std::move(job)); + } + mReadQueueCv.notify_one(); +} + +std::vector KVCacheTransferManager::drainCompletedBlockReads() +{ + std::lock_guard lock(mReadMutex); + std::vector out; + out.swap(mCompletedBlockReads); + return out; +} + +bool KVCacheTransferManager::isBlockReadPending(std::int32_t blockId) +{ + std::lock_guard lock(mReadMutex); + return mPendingBlockReads.find(blockId) != mPendingBlockReads.end(); +} + void KVCacheTransferManager::spillToFile(BlockPtr const& srcHostBlock, SizeType32 diskSlot, std::vector const& pools, std::string const& directory) { @@ -633,7 +773,7 @@ void KVCacheTransferManager::spillToFileUnstaged(BlockPtr const& srcHostBlock, S } void KVCacheTransferManager::loadFromFile(BlockPtr const& dstPrimaryBlock, SizeType32 diskSlot, - std::vector const& pools, std::string const& directory) + std::vector const& pools, std::string const& directory, std::int32_t trackBlockId) { TLLM_CHECK_WITH_INFO(!directory.empty(), "disk tier requires a directory"); auto const idx = getPendingTransferIndex(dstPrimaryBlock); @@ -653,6 +793,29 @@ void KVCacheTransferManager::loadFromFile(BlockPtr const& dstPrimaryBlock, SizeT waitForDiskSlotWrites(diskSlotFilename(directory, diskSlot, poolIdx)); } } + // Detached onboard: when a reader pool is present, package all of the block's pools into one tracked job + // and hand it off. The reader performs the transfer (GDS DMA if enabled, else POSIX read + copy) and makes + // it device-complete before publishing the block; the request is forward-safe once areBlocksReady() sees + // its blocks land. trackBlockId < 0 (no pool) takes the synchronous fallback below. + if (asyncDiskReadEnabled() && trackBlockId >= 0) + { + DiskReadJob job; + job.blockId = trackBlockId; + job.useGds = mDiskUseGds; + for (size_t poolIdx = 0; poolIdx < pools.size(); ++poolIdx) + { + TLLM_CHECK_WITH_INFO( + !pools[poolIdx].layerFirstLayout, "disk tier does not support layer-first layout pools"); + auto ptr = computeBlockPointer(dstPrimaryBlock, pools, poolIdx); + job.dsts.push_back(ptr->data()); + job.files.push_back(diskSlotFilename(directory, diskSlot, poolIdx)); + job.bytes.push_back(static_cast(ptr->getSizeInBytes())); + } + enqueueDiskRead(std::move(job)); + return; + } + + // Synchronous fallback (no reader pool): read inline on the calling thread. if (mDiskUseGds) { // GDS onboard: direct SSD -> GPU VRAM via the NIXL loopback agent (cuFile). diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp index 02a62679ea80..02a611e4eeaf 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp @@ -684,6 +684,7 @@ void tb::kv_cache_manager::KVCacheManagerBindings::initBindings(nb::module_& m) [](tbk::KVCacheManager& self, SizeType32 numRequired, SizeType32 windowSize) { return self.getBlockManager().schedulingHasFreeBlocks(numRequired, windowSize); }, nb::arg("num_required"), nb::arg("window_size"), nb::call_guard()) + .def("are_blocks_ready", &tbk::KVCacheManager::areBlocksReady, nb::arg("request_id")) .def_prop_ro( "is_variable_window", [](tbk::KVCacheManager& self) { return self.getBlockManager().isVariableWindow(); }) // Per-pool introspection: lets Python discover (windowSize, sizePerHead, dtype) per diff --git a/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp b/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp index 888060c60c6c..05cf1724e4c2 100644 --- a/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp @@ -10740,6 +10740,40 @@ TEST_F(KVCacheManagerTest, DiskTierSpillAllRoundTripTest) releaseDiskTierSequence(blockManager, seqA2); } +// Detached onboard (reader pool on): reusing a spilled prefix must produce the SAME observable reuse as the +// synchronous path above, and the request must report forward-safe via areBlocksReady() once its onboard +// lands. Same spill/evict/reuse scenario as DiskTierSpillAllRoundTripTest, only with TLLM_KV_DISK_READERS set. +TEST_F(KVCacheManagerTest, DiskTierDetachedOnboardReadyTest) +{ + setenv("TLLM_KV_DISK_READERS", "4", /*overwrite=*/1); + auto const stream = std::make_shared(); + DiskTierDir dir; + auto blockManagerPtr + = makeDiskTierBlockManager(stream, dir.str(), /*retainedOnly=*/false, /*prim=*/4, /*sec=*/2, /*disk=*/8); + auto& blockManager = *blockManagerPtr; + + auto seqA = addDiskTierSequence(blockManager, 0, iotaTokens(0, 16), std::nullopt); + releaseDiskTierSequence(blockManager, seqA); + churnOnce(blockManager, 1, 100); // evict A's blocks through to disk + + auto seqA2 = addDiskTierSequence(blockManager, 2, iotaTokens(0, 16), std::nullopt); + + // The onboard was handed to the reader pool; the request is forward-safe only once its blocks land. + auto const deadline = std::chrono::steady_clock::now() + std::chrono::seconds(30); + while (!blockManager.areBlocksReady(/*requestId=*/2)) + { + ASSERT_LT(std::chrono::steady_clock::now(), deadline) << "detached onboard never became ready"; + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + + // Same reuse outcome as the synchronous path: >=3 of 4 blocks reused, part served from disk. + EXPECT_GE(seqA2.req->getContextCurrentPosition(), 12); + EXPECT_GE(blockManager.getNumDiskOnboards(), 1u); + + releaseDiskTierSequence(blockManager, seqA2); + unsetenv("TLLM_KV_DISK_READERS"); +} + // Retained-only mode: unmarked victims are discarded at the gate; marked ones spill. TEST_F(KVCacheManagerTest, DiskTierRetainedOnlyGateTest) { @@ -10901,3 +10935,111 @@ TEST_F(KVCacheManagerTest, DiskTierProtectEvictsExpiredTest) EXPECT_GT(blockManager.getNumDiskSpills(), spillsFilled); EXPECT_EQ(blockManager.getNumDiskAdmissionRefused(), 0u); } + +namespace +{ +// Per-slot-distinct byte pattern for the disk round-trip. Distinct across slots so a cross-slot +// mixup (slot i's bytes landing in slot j) is caught, not just intra-block corruption. +float diskPatternValue(int slot, int i) +{ + return static_cast((i * 3 + slot * 101) % 997); +} + +// Spill nSlots distinct-pattern blocks to disk, onboard each into a GPU slot via loadFromFile (detached +// through the reader pool when readers>0, inline when readers==0), wait for per-block readiness, and verify +// every GPU slot holds its EXACT pattern. Unlike DiskTierSpillAllRoundTripTest (which only checks reuse +// accounting), this asserts the onboarded KV *bytes* -- the property the detached reader must preserve. +// TLLM_KV_DISK_ASYNC_STORE is left unset so spills are synchronous and this isolates the read path. +void runDiskOnboardByteRoundTrip(char const* readers, int nSlots) +{ + setenv("TLLM_KV_DISK_READERS", readers, /*overwrite=*/1); + DiskTierDir dir; + int constexpr blockSize = 2048; // floats per slot + + auto bufferManager = tr::BufferManager(std::make_shared()); + auto transferManager = KVCacheTransferManager(bufferManager); // reads TLLM_KV_DISK_READERS at construction + + auto pool = KVCacheBlockPool(0, /*kvFactor=*/2, 0, 0, 0); + pool.primaryPtr = bufferManager.gpu(tr::ITensor::makeShape({nSlots, blockSize}), nvinfer1::DataType::kFLOAT); + bufferManager.setZero(*pool.primaryPtr); + pool.secondaryPtr + = tr::BufferManager::pinned(tr::ITensor::makeShape({nSlots, blockSize}), nvinfer1::DataType::kFLOAT); + + // Fill each host slot with its pattern and spill it to disk slot k (synchronous write). + float* secBase = tr::bufferCast(*pool.secondaryPtr); + for (int k = 0; k < nSlots; ++k) + { + for (int i = 0; i < blockSize; ++i) + { + secBase[static_cast(k) * blockSize + i] = diskPatternValue(k, i); + } + auto src = std::make_shared(k, tk::KVCacheIndex(k, /*isSecondary=*/true)); + transferManager.spillToFile(src, /*diskSlot=*/k, {pool}, dir.str()); + } + + // Ensure the GPU pool's zero-fill has completed before detached readers overwrite the slots. + bufferManager.getStream().synchronize(); + + // Onboard each disk slot into GPU slot k, tracked by k so readers>0 exercises the detached reader pool + // (readers==0 reads inline on this thread). + std::vector dst; + for (int k = 0; k < nSlots; ++k) + { + auto d = std::make_shared(nSlots + k, tk::KVCacheIndex(k, /*isSecondary=*/false)); + transferManager.loadFromFile(d, /*diskSlot=*/k, {pool}, dir.str(), /*trackBlockId=*/k); + dst.push_back(d); + } + // Detached reads land asynchronously; wait until each slot's read has completed. readers==0 never enters + // the pending set, so this returns immediately for the synchronous path. + auto const deadline = std::chrono::steady_clock::now() + std::chrono::seconds(30); + for (int k = 0; k < nSlots; ++k) + { + while (transferManager.isBlockReadPending(k)) + { + ASSERT_LT(std::chrono::steady_clock::now(), deadline) << "onboard did not complete for slot " << k; + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + } + bufferManager.getStream().synchronize(); + + // Read each GPU slot back to host and verify byte-exact against its pattern. + std::vector hostCopy(blockSize); + float const* primBase = tr::bufferCast(*pool.primaryPtr); + for (int k = 0; k < nSlots; ++k) + { + ASSERT_EQ(cudaMemcpy(hostCopy.data(), primBase + static_cast(k) * blockSize, + blockSize * sizeof(float), cudaMemcpyDeviceToHost), + cudaSuccess); + int bad = 0; + for (int i = 0; i < blockSize; ++i) + { + if (hostCopy[i] != diskPatternValue(k, i)) + { + ++bad; + } + } + EXPECT_EQ(bad, 0) << "readers=" << readers << " slot=" << k << ": disk onboard corrupted " << bad << "/" + << blockSize << " elements"; + } + unsetenv("TLLM_KV_DISK_READERS"); +} +} // namespace + +// The detached reader pool must onboard byte-exact KV from disk (4 readers, 4 slots). +TEST_F(KVCacheManagerTest, DiskTierAsyncOnboardByteExactTest) +{ + runDiskOnboardByteRoundTrip(/*readers=*/"4", /*nSlots=*/4); +} + +// More slots than readers => concurrent draining; distinct per-slot patterns catch any cross-slot +// mixup (slot i's bytes landing in slot j). +TEST_F(KVCacheManagerTest, DiskTierAsyncOnboardConcurrentNoMixupTest) +{ + runDiskOnboardByteRoundTrip(/*readers=*/"3", /*nSlots=*/8); +} + +// Control: the synchronous read path (readers=0) is byte-exact too => detached matches inline. +TEST_F(KVCacheManagerTest, DiskTierSyncOnboardByteExactTest) +{ + runDiskOnboardByteRoundTrip(/*readers=*/"0", /*nSlots=*/4); +} diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index cef6ef31e91f..e88bb2b532a7 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -2539,6 +2539,7 @@ def _executor_loop_pp(self): self._handle_dynamic_draft_len(scheduled_batch) self.resource_manager.prepare_resources(scheduled_batch) + self._park_requests_awaiting_onboard(scheduled_batch) # The generation requests that do not have batch_idx # need to be in front of the batch due to the assumptions @@ -3094,6 +3095,36 @@ def wait_on_pp_send_handles(self, send_handles, microbatch_id): send_handles[microbatch_id].wait() send_handles[microbatch_id] = None + def _park_requests_awaiting_onboard(self, scheduled_batch): + """Hold context requests whose detached disk-onboard reads have not landed out of this forward. + Their KV blocks are allocated but not yet filled, so forwarding now would read stale KV. A parked + request stays active and is re-scheduled next step -- the capacity scheduler counts its already + allocated blocks (getRemainingBlocksToCompletion) and _context_seq_len skips re-adding it -- then it + is admitted once are_blocks_ready() reports its blocks landed.""" + kv = getattr(self, "kv_cache_manager", None) + if kv is None or not hasattr(kv, "are_blocks_ready"): + return + ctx = scheduled_batch.context_requests + if not ctx: + return + kept = [] + parked = 0 + for req in ctx: + if kv.are_blocks_ready(req): + req.py_onboard_pending = False + kept.append(req) + else: + req.py_onboard_pending = True + parked += 1 + if parked: + scheduled_batch.reset_context_requests(kept) + self._onboard_parked_total = getattr(self, "_onboard_parked_total", + 0) + parked + if self._onboard_parked_total <= 5 or self._onboard_parked_total % 500 == 0: + logger.info( + "disk-onboard: parked %d context request(s) awaiting KV read (cumulative %d)", + parked, self._onboard_parked_total) + def _handle_dynamic_draft_len(self, scheduled_batch: ScheduledRequests) -> None: """Handle dynamic draft length for the current batch. @@ -3776,6 +3807,7 @@ def _executor_loop(self): self._handle_dynamic_draft_len(scheduled_batch) self.resource_manager.prepare_resources(scheduled_batch) + self._park_requests_awaiting_onboard(scheduled_batch) if self.kv_connector_manager: self.kv_connector_manager.handle_metadata() @@ -4241,6 +4273,7 @@ def _executor_loop_overlap(self): self._handle_dynamic_draft_len(scheduled_batch) self.resource_manager.prepare_resources(scheduled_batch) + self._park_requests_awaiting_onboard(scheduled_batch) if self.kv_connector_manager: self.kv_connector_manager.handle_metadata() diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index a572e4121859..0f312e72ff16 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -728,6 +728,12 @@ def get_needed_resource_to_completion(self, request: LlmRequest) -> int: def _context_seq_len(self, req: LlmRequest, is_cross: bool, is_star_cp: bool) -> Optional[int]: """Return the sequence length to pass to add_sequence_batch, or None to skip this request.""" + if getattr(req, "py_onboard_pending", False): + # Blocks were claimed on a prior step and the detached disk-onboard read is still in flight. + # Do not re-add; the request is held out of the forward until areBlocksReady() + # (see PyExecutor._park_requests_awaiting_onboard). Same "add once, forward later" path + # that chunked context already uses. + return None if is_cross: if (getattr(req, "py_skip_cross_kv_projection", False) or not req.is_first_context_chunk @@ -1400,6 +1406,11 @@ def _pack_beam_cache_indices(beams: List[List[int]]) -> List[int]: packed.append(beam[-1]) return packed + def are_blocks_ready(self, request) -> bool: + """True when every KV block this request holds has its disk read landed (detached onboard). + Always True for requests with no in-flight disk read, so it is cheap for the common case.""" + return self.impl.are_blocks_ready(request.py_request_id) + def get_num_free_blocks(self) -> int: if self.is_linear_attention: value = self.impl.get_kv_cache_stats( From 938f8e6bb684d34a17a1dd0c1caf463057c87176 Mon Sep 17 00:00:00 2001 From: Md Nafis Ul Haque Shifat Date: Tue, 7 Jul 2026 21:17:52 +0000 Subject: [PATCH 14/47] [#16251][feat] batch_manager: reuse-count disk spill gate (TLLM_KV_DISK_MIN_REUSE) Add a per-block reuse counter (mReuseCount, bumped at the prefix-match reuse point in onboardAndAllocateBlocks) and gate the host->disk spill on it: only spill blocks reused >= k times (k = TLLM_KV_DISK_MIN_REUSE env, 0 = off). Sits beside the retained_only gate as an independent 3rd disk mode (spill-all / retention-only / reuse-gate). The counter is block metadata, so it survives GPU/host/disk tier moves like mRetentionExpiry (swapDiskResidency swaps only residency pointers). Signed-off-by: Md Nafis Ul Haque Shifat --- .../batch_manager/kvCacheManager.h | 15 +++++++++++++++ .../batch_manager/kvCacheManager.cpp | 18 ++++++++++++++++++ 2 files changed, 33 insertions(+) diff --git a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h index d483a7830fe4..c0f453e57ab2 100644 --- a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h +++ b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h @@ -487,6 +487,17 @@ class KVCacheBlock : public std::enable_shared_from_this mRetentionExpiry = std::nullopt; } + //! \brief Times this block's content has been reused (prefix-matched); drives the disk reuse-gate. + [[nodiscard]] SizeType32 getReuseCount() const + { + return mReuseCount; + } + + void incReuseCount() + { + ++mReuseCount; + } + void incRefCount(); void decRefCount(); @@ -647,6 +658,10 @@ class KVCacheBlock : public std::enable_shared_from_this // Disk-tier retention deadline (steady clock); nullopt = never marked. std::optional mRetentionExpiry; + + // How many times this block's content has been reused (prefix-matched). Metadata on the block object, + // so it survives tier moves like mRetentionExpiry (swapDiskResidency swaps only residency pointers). + SizeType32 mReuseCount{0}; }; class GenerationRequest diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp index 54c9455310d8..e65cb1e1f7ff 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp @@ -1395,6 +1395,23 @@ BlockPtr WindowBlockManager::reclaimSecondaryBlock() } return victim; // unmarked or expired: discard, exactly as stock } + // Reuse-gate: only spill blocks whose content has been reused >= k times (TLLM_KV_DISK_MIN_REUSE; 0=off). + // Skips one-shot blocks so the SSD tier holds only genuinely-hot prefixes. Env read once, cached. + static SizeType32 const kDiskMinReuse = []() -> SizeType32 + { + auto const* e = std::getenv("TLLM_KV_DISK_MIN_REUSE"); + return e ? static_cast(std::atoi(e)) : 0; + }(); + if (kDiskMinReuse > 0 && victim->getReuseCount() < kDiskMinReuse) + { + ++mDiskGateDropped; + if (mDiskGateDropped == 1 || mDiskGateDropped % 10000 == 0) + { + TLLM_LOG_INFO("[disk-tier] reuse-gate dropped=%zu (k=%d, windowSize=%d)", mDiskGateDropped, + static_cast(kDiskMinReuse), mWindowSize); + } + return victim; // reused fewer than k times: discard, don't spill + } auto diskTarget = claimDiskTarget(); if (diskTarget == nullptr) { @@ -2134,6 +2151,7 @@ SizeType32 WindowBlockManager::onboardAndAllocateBlocks( addBlockToAllBeams(claimed.block, sequence); if (!claimed.isPlaceholder) { + claimed.block->incReuseCount(); // per-block reuse count for the disk reuse-gate ++mReusedBlocks; if (claimed.isPartialMatch) { From ffcb1f569d5873c8907fc52cd78634e7b3965944 Mon Sep 17 00:00:00 2001 From: Md Nafis Ul Haque Shifat Date: Tue, 7 Jul 2026 21:46:35 -0700 Subject: [PATCH 15/47] [#16251][feat] disk tier: tensor-parallel support (per-rank disk dir + cross-rank onboard readiness) Two fixes to run the disk KV-cache tier under TP>1 (validated on gpt-oss-120b TP=2). Per-rank disk cache directory: all ranks share one container filesystem and one disk_cache_path, but the disk slot files (block__pool_

.bin) are named identically across ranks while each rank holds a different KV shard -- so without namespacing the ranks clobber each other's KV on load. Append /rank_ (the global rank, unique across TP and PP) and mkdir it. Note: disk_cache_size is per-rank, so total disk = world_size * disk_cache_size. Cross-rank onboard-readiness sync: the detached disk-onboard reads land at slightly different times per rank, so a per-rank park decision in _park_requests_awaiting_onboard would give ranks different batches and the TP all-reduce in the forward would mismatch (hang/crash). AND the per-request readiness across the TP group -- a cheap tp_allreduce "any pending?" fast path, then tp_allgather when a read is in flight -- so a request is admitted only once its KV has landed on every rank. pp_size==1 => the TP group is the whole world; TP=1 is a no-op (identical to the prior path). Signed-off-by: Md Nafis Ul Haque Shifat --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 37 +++++++++++++++++-- .../_torch/pyexecutor/resource_manager.py | 18 ++++++++- 2 files changed, 49 insertions(+), 6 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index e88bb2b532a7..17dbc0224316 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -3107,10 +3107,38 @@ def _park_requests_awaiting_onboard(self, scheduled_batch): ctx = scheduled_batch.context_requests if not ctx: return + # Local readiness per context request, in batch order. The batch is + # broadcast from the leader and scheduling is deterministic, so ctx is + # identically ordered on every rank. + local_ready = [1 if kv.are_blocks_ready(req) else 0 for req in ctx] + + # TP safety: a context request must be parked on EVERY rank or NONE. Each + # rank's detached disk-onboard reads land at slightly different times, so a + # per-rank park decision would give ranks different batches and the TP + # all-reduce in the forward would mismatch (hang/crash). AND readiness + # across the TP group: admit a request only once its KV has landed on all + # ranks. (pp_size==1 -> the TP group is the whole world; TP=1 is a no-op.) + local_has_pending = 0 if all(local_ready) else 1 + if self.dist.tp_size > 1: + any_pending = self.dist.tp_allreduce(int(local_has_pending), + op=ReduceOp.MAX) + else: + any_pending = local_has_pending + if not any_pending: + # Every rank has all context blocks ready -> keep the full batch. + for req in ctx: + req.py_onboard_pending = False + return + if self.dist.tp_size > 1: + gathered = self.dist.tp_allgather(local_ready) + global_ready = [all(bits) for bits in zip(*gathered)] + else: + global_ready = local_ready + kept = [] parked = 0 - for req in ctx: - if kv.are_blocks_ready(req): + for req, ready in zip(ctx, global_ready): + if ready: req.py_onboard_pending = False kept.append(req) else: @@ -3122,8 +3150,9 @@ def _park_requests_awaiting_onboard(self, scheduled_batch): 0) + parked if self._onboard_parked_total <= 5 or self._onboard_parked_total % 500 == 0: logger.info( - "disk-onboard: parked %d context request(s) awaiting KV read (cumulative %d)", - parked, self._onboard_parked_total) + f"disk-onboard: parked {parked} context request(s) " + f"awaiting KV read (cumulative {self._onboard_parked_total})" + ) def _handle_dynamic_draft_len(self, scheduled_batch: ScheduledRequests) -> None: diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index 0f312e72ff16..8120df9131f9 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -621,8 +621,22 @@ def append_to_kv_heads_per_layer(num_kv_heads_per_layer: List[int], logger.info(f"[disk-tier] {blocks_in_disk_pool} blocks " f"({disk_cache_size / (1 << 30):.1f} GiB)") kwargs['blocks_in_disk_pool'] = blocks_in_disk_pool - kwargs['disk_cache_path'] = getattr(kv_cache_config, "disk_cache_path", - None) or "" + disk_cache_path = getattr(kv_cache_config, "disk_cache_path", + None) or "" + if disk_cache_path and blocks_in_disk_pool: + # Per-rank subdir: all ranks share one container filesystem and one + # disk_cache_path mount, but each holds a different KV shard while the + # disk slot files (block__pool_

.bin) are named identically across + # ranks -> without namespacing the ranks clobber each other's KV on load. + # mapping.rank is the global rank (unique across TP and PP). Note: + # disk_cache_size is per-rank, so total disk = world_size * disk_cache_size. + disk_cache_path = os.path.join(disk_cache_path, + f"rank_{self.mapping.rank}") + os.makedirs(disk_cache_path, exist_ok=True) + logger.info( + f"[disk-tier] per-rank disk cache dir: {disk_cache_path} " + f"(rank {self.mapping.rank}/{self.mapping.world_size})") + kwargs['disk_cache_path'] = disk_cache_path kwargs['disk_cache_retained_only'] = bool( getattr(kv_cache_config, "disk_cache_retained_only", False)) kwargs['disk_cache_protect_unexpired'] = bool( From 09c80a901cbda732e9a4a88e8d1b0ebc9bcc2711 Mon Sep 17 00:00:00 2001 From: Md Nafis Ul Haque Shifat Date: Tue, 7 Jul 2026 21:46:35 -0700 Subject: [PATCH 16/47] [#16251][perf] disk tier: skip areBlocksReady scan when no disk read is in flight areBlocksReady scans every block a request holds via isBlockReadPending, which takes mReadMutex per block, and the scheduler-side park calls it for every context request every step. With no disk read in flight (the common case) that is an O(context-length) mutex-locked scan over blocks that are all trivially ready. Add a lock-free atomic mirror of the in-flight read count (mReadInflightCount, updated under mReadMutex on every pending-set insert/erase) and short-circuit areBlocksReady to return true immediately when it is zero. Behaviour is identical when reads are pending. Signed-off-by: Md Nafis Ul Haque Shifat --- .../batch_manager/kvCacheTransferManager.h | 10 ++++++++++ cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp | 7 +++++++ .../batch_manager/kvCacheTransferManager.cpp | 2 ++ 3 files changed, 19 insertions(+) diff --git a/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h b/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h index aff0533a623b..f3b674aad520 100644 --- a/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h +++ b/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h @@ -19,6 +19,7 @@ #include "tensorrt_llm/runtime/bufferManager.h" #include "tensorrt_llm/runtime/cudaEvent.h" +#include #include #include #include @@ -101,6 +102,12 @@ class KVCacheTransferManager //! \brief True while block blockId still has an in-flight disk read (not yet safe to read on the GPU). [[nodiscard]] bool isBlockReadPending(std::int32_t blockId); + //! \brief Lock-free: true if ANY disk read is in flight, letting areBlocksReady() skip its scan. + [[nodiscard]] bool anyReadPending() const noexcept + { + return mReadInflightCount.load(std::memory_order_acquire) > 0; + } + //! \brief True when a reader pool is present, so onboards can be detached from the scheduler thread. [[nodiscard]] bool asyncDiskReadEnabled() const { @@ -245,6 +252,9 @@ class KVCacheTransferManager // Blocks with a disk read in flight, keyed by the matched-identity block id so every request reusing that // prefix gates on the same key. A block is forward-safe once it leaves this set. std::unordered_set mPendingBlockReads; + // Lock-free mirror of mPendingBlockReads.size() (updated under mReadMutex on every insert/erase) so + // areBlocksReady() can skip its per-block scan + the mutex when nothing is in flight (common case). + std::atomic mReadInflightCount{0}; std::vector mCompletedBlockReads; bool mDiskReaderStop{false}; diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp index e65cb1e1f7ff..e9910ecec256 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp @@ -1673,6 +1673,13 @@ void WindowBlockManager::onboardBlock(GenerationRequest& sequence, BlockPtr cons bool WindowBlockManager::areBlocksReady(LlmRequest::RequestIdType requestId) { + // Fast path: with no disk read in flight (the common case) every block is trivially ready, so skip the + // per-block scan and its mutex entirely. Without this the park scans every block of every context + // request each step -- an O(context-length) mutex-locked host cost even when nothing is onboarding. + if (!mTransferManager->anyReadPending()) + { + return true; + } // A request is forward-safe only when every block it holds has its disk read landed. Keyed by block id, // so a request sharing an in-flight onboarded block waits on the same key as the request that started it. auto it = mAllocatedBlocksPerSeq.find(requestId); diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp index 228ddf51cfca..82b41d868750 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp @@ -679,6 +679,7 @@ void KVCacheTransferManager::diskReaderLoop() std::lock_guard lock(mReadMutex); // The transfer is device-complete, so publish the block: any request holding it can be forwarded. mPendingBlockReads.erase(job.blockId); + mReadInflightCount.store(mPendingBlockReads.size(), std::memory_order_release); mCompletedBlockReads.push_back(job.blockId); } } @@ -690,6 +691,7 @@ void KVCacheTransferManager::enqueueDiskRead(DiskReadJob job) { std::lock_guard lock(mReadMutex); mPendingBlockReads.insert(job.blockId); // mark pending before the job becomes poppable + mReadInflightCount.store(mPendingBlockReads.size(), std::memory_order_release); mReadQueue.push(std::move(job)); } mReadQueueCv.notify_one(); From c128718e6a62ccfdf2ed21d1c2c3bb154afeb31e Mon Sep 17 00:00:00 2001 From: Md Nafis Ul Haque Shifat Date: Wed, 8 Jul 2026 12:13:14 -0700 Subject: [PATCH 17/47] [#16251][fix] batch_manager: bind disk-tier reader thread to its CUDA device Reader threads created their CUDA stream without setting the device, so on TP ranks whose device is not 0 the POSIX H2D onboard targeted the wrong GPU and crashed. Bind the thread to the manager's device at loop entry. Signed-off-by: Md Nafis Ul Haque Shifat --- cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp | 3 +++ 1 file changed, 3 insertions(+) diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp index 82b41d868750..79fffe96e862 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp @@ -598,6 +598,9 @@ void KVCacheTransferManager::diskReaderLoop() // managers) and its own GDS agent (readers never share one). cudaStreamSynchronize (POSIX) or the agent's // internal wait (GDS) makes each transfer device-complete before the block is published, so per-block // readiness alone is a sufficient pre-forward gate. + // Bind this reader thread to the manager's device: a fresh std::thread starts on device 0, so + // without this the stream and POSIX H2D copies target the wrong GPU on TP ranks != 0 (crash). + TLLM_CUDA_CHECK(cudaSetDevice(mDeviceId)); cudaStream_t stream; TLLM_CUDA_CHECK(cudaStreamCreate(&stream)); std::vector hostBuffer; From 1b507acbd7c07e2cceb3945cd7b5197ceec2d4e7 Mon Sep 17 00:00:00 2001 From: Md Nafis Ul Haque Shifat Date: Wed, 8 Jul 2026 12:13:14 -0700 Subject: [PATCH 18/47] [#16251][fix] batch_manager: correct disk-tier queue-integrity validators verifyQueueIntegrity had three validators/labels for four queue levels: the disk level reused the placeholder validator and the loop indexed one past the array. Give each level its own validator and label (disk = isOnDisk). Signed-off-by: Md Nafis Ul Haque Shifat --- cpp/tensorrt_llm/batch_manager/evictionPolicy.cpp | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/cpp/tensorrt_llm/batch_manager/evictionPolicy.cpp b/cpp/tensorrt_llm/batch_manager/evictionPolicy.cpp index d982b1b8d935..2e10a9a7dbd2 100644 --- a/cpp/tensorrt_llm/batch_manager/evictionPolicy.cpp +++ b/cpp/tensorrt_llm/batch_manager/evictionPolicy.cpp @@ -125,10 +125,11 @@ void LRUEvictionPolicy::initializePlaceholders(std::vector& allPlaceho bool LRUEvictionPolicy::verifyQueueIntegrity() const { - static char const* const levelToStr[] = {"primary", "secondary", "placeholder"}; + static char const* const levelToStr[] = {"primary", "secondary", "disk", "placeholder"}; static const std::function levelValidators[] = {[](BlockPtr const& block) { return block->isPrimary(); }, - [](BlockPtr const& block) { return !block->isPrimary(); }, + [](BlockPtr const& block) { return !block->isPrimary() && !block->isOnDisk() && !block->isPlaceholder(); }, + [](BlockPtr const& block) { return block->isOnDisk(); }, [](BlockPtr const& block) { return block->isPlaceholder(); }}; bool queueCompromised = false; for (SizeType32 queueLevel = 0; queueLevel < kNumCacheLevels + 1; queueLevel++) From b54f55a84a8010e98a532f44f578a80e88369e76 Mon Sep 17 00:00:00 2001 From: Md Nafis Ul Haque Shifat Date: Wed, 8 Jul 2026 12:13:14 -0700 Subject: [PATCH 19/47] [#16251][fix] batch_manager: don't partial-reuse disk-resident blocks A disk-resident block cannot serve a partial reuse -- its bytes live in a reserved disk slot, not in pool memory, so the partial-copy path reads garbage. Exclude on-disk blocks from both partial-match routes; full matches stay reusable via the async onboard. Signed-off-by: Md Nafis Ul Haque Shifat --- cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp index e9910ecec256..32151cea5ba9 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp @@ -1818,6 +1818,13 @@ WindowBlockManager::ReuseMatchResult WindowBlockManager::findReusableBlockMatche if (existing.has_value() && *existing) { auto block = *existing; + // A non-full disk block would be flagged a partial match, routing to the broken + // partial-copy path. Stop here: the request keeps the fully-matched prefix (already + // committed via updateSafePrefix) and prefills the rest. A full disk block is unaffected. + if (block->isOnDisk() && !block->isFull()) + { + break; + } candidateMatches.push_back(ReuseMatch{block, numMatchedTokens, !block->isFull(), false}); } else if (mIsSWA) @@ -1847,6 +1854,13 @@ WindowBlockManager::ReuseMatchResult WindowBlockManager::findReusableBlockMatche } auto block = *existing; + // Disk-resident blocks can't serve a partial reuse: the partial-copy path reads pool + // memory, but disk blocks are poolless (reserved disk slots), so the copy reads garbage. + // Skip this candidate; another partial candidate or a fresh prefill covers these tokens. + if (block->isOnDisk()) + { + continue; + } if (copyOnPartialReuse || (!block->hasRefs() && block->isLeaf())) { auto const numMatchedTokens = static_cast(match.key.uniqueTokens.size()); From e2864a866a41409b5d3c84f7bed1a279d439dad3 Mon Sep 17 00:00:00 2001 From: Md Nafis Ul Haque Shifat Date: Wed, 8 Jul 2026 12:13:58 -0700 Subject: [PATCH 20/47] [#16251][fix] batch_manager: hold a disk slot until its onboard read completes onboardBlock enqueued the async read of a disk slot then immediately returned its card to the disk free queue, so the next spill could claim and overwrite the slot file before the read ran -- onboarding the wrong block. Hold the card until the read lands (reclaimed in claimDiskTarget once the read completes) and drop the now-unused completed-reads list. Signed-off-by: Md Nafis Ul Haque Shifat --- .../batch_manager/kvCacheManager.h | 5 +++ .../batch_manager/kvCacheTransferManager.h | 4 --- .../batch_manager/kvCacheManager.cpp | 32 ++++++++++++++++++- .../batch_manager/kvCacheTransferManager.cpp | 9 ------ 4 files changed, 36 insertions(+), 14 deletions(-) diff --git a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h index c0f453e57ab2..dda63a123ed2 100644 --- a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h +++ b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h @@ -1508,6 +1508,11 @@ class WindowBlockManager SizeType32 mReservedHostBlockTarget{0}; std::vector mReservedHostBlocks; std::unordered_map mPendingSpillBlocks; + // Onboarded disk cards whose async read is still in flight, keyed by the read's tracked block id. + // Held out of the disk free queue until the read lands so a spill cannot overwrite the slot file + // mid-read; reapReadPendingReleases() returns each card once its read completes. + std::unordered_map mReadPendingReleases; + void reapReadPendingReleases(); std::uint64_t mUnstagedSpillSeq{0}; //! \brief Pick and claim the disk block to overwrite: empty slots first, then diff --git a/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h b/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h index f3b674aad520..0c282d1b948f 100644 --- a/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h +++ b/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h @@ -96,9 +96,6 @@ class KVCacheTransferManager void loadFromFile(BlockPtr const& dstPrimaryBlock, SizeType32 diskSlot, std::vector const& pools, std::string const& directory, std::int32_t trackBlockId = -1); - //! \brief Return + clear the block ids whose reads have landed since the last call. - [[nodiscard]] std::vector drainCompletedBlockReads(); - //! \brief True while block blockId still has an in-flight disk read (not yet safe to read on the GPU). [[nodiscard]] bool isBlockReadPending(std::int32_t blockId); @@ -255,7 +252,6 @@ class KVCacheTransferManager // Lock-free mirror of mPendingBlockReads.size() (updated under mReadMutex on every insert/erase) so // areBlocksReady() can skip its per-block scan + the mutex when nothing is in flight (common case). std::atomic mReadInflightCount{0}; - std::vector mCompletedBlockReads; bool mDiskReaderStop{false}; void diskReaderLoop(); diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp index 32151cea5ba9..03a0beaab43b 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp @@ -1322,8 +1322,27 @@ void WindowBlockManager::releaseSubtree(BlockPtr const& block) } } +void WindowBlockManager::reapReadPendingReleases() +{ + // Cards whose async onboard read has landed can rejoin the disk free queue: their slot file is no + // longer being read, so a later spill may safely overwrite it. + for (auto it = mReadPendingReleases.begin(); it != mReadPendingReleases.end();) + { + if (mTransferManager->isBlockReadPending(it->first)) + { + ++it; + continue; + } + mEvictionPolicy->releaseBlock(it->second); + it = mReadPendingReleases.erase(it); + } +} + BlockPtr WindowBlockManager::claimDiskTarget() { + // Reclaim onboarded slots whose read has landed before handing out a disk slot, so a slot with an + // in-flight read is never claimable -- its file must not be overwritten while that read is pending. + reapReadPendingReleases(); // Priority order: empty slots (kMin) < unmarked content (kDefault) < retained (kMax), // so getFreeBlock returns empties first and a retained block only when nothing cheaper // (empty slot or unmarked cached block) remains. @@ -1644,7 +1663,18 @@ void WindowBlockManager::onboardBlock(GenerationRequest& sequence, BlockPtr cons offloadBlock->setPriority(executor::KvCacheRetentionConfig::kDefaultRetentionPriority); block->clearRetention(); block->setPriority(executor::KvCacheRetentionConfig::kMinRetentionPriority); // freed card = empty slot again - mEvictionPolicy->releaseBlock(block); // freed card returns to the disk free queue + if (trackId >= 0) + { + // The async read of this slot is still queued/in-flight. Returning the freed card to the disk + // free queue now would let the next spill claim + overwrite the slot file before the read runs + // (read-vs-write race -> the reader onboards the wrong block). Hold the card until the read + // lands; reapReadPendingReleases() (from claimDiskTarget) returns it once !isBlockReadPending. + mReadPendingReleases.emplace(trackId, block); + } + else + { + mEvictionPolicy->releaseBlock(block); // synchronous read already complete: slot safe to reuse + } ++mDiskOnboards; if (mDiskOnboards == 1 || mDiskOnboards % 1000 == 0) { diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp index 79fffe96e862..dcef83d22aa5 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp @@ -683,7 +683,6 @@ void KVCacheTransferManager::diskReaderLoop() // The transfer is device-complete, so publish the block: any request holding it can be forwarded. mPendingBlockReads.erase(job.blockId); mReadInflightCount.store(mPendingBlockReads.size(), std::memory_order_release); - mCompletedBlockReads.push_back(job.blockId); } } cudaStreamDestroy(stream); @@ -700,14 +699,6 @@ void KVCacheTransferManager::enqueueDiskRead(DiskReadJob job) mReadQueueCv.notify_one(); } -std::vector KVCacheTransferManager::drainCompletedBlockReads() -{ - std::lock_guard lock(mReadMutex); - std::vector out; - out.swap(mCompletedBlockReads); - return out; -} - bool KVCacheTransferManager::isBlockReadPending(std::int32_t blockId) { std::lock_guard lock(mReadMutex); From 0f0034663f9f17bf076a87ee220f3ef9ce6942aa Mon Sep 17 00:00:00 2001 From: Md Nafis Ul Haque Shifat Date: Wed, 8 Jul 2026 12:13:58 -0700 Subject: [PATCH 21/47] [#16251][fix] batch_manager: fail loud on disk-tier I/O errors Short reads/writes and open failures on disk slot files were logged but the block was still published, serving or caching corrupt KV. Fail loud on any disk I/O error, matching the synchronous spill path. Signed-off-by: Md Nafis Ul Haque Shifat --- .../batch_manager/kvCacheTransferManager.cpp | 49 +++++++------------ 1 file changed, 18 insertions(+), 31 deletions(-) diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp index dcef83d22aa5..706f8a46ed99 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp @@ -48,7 +48,7 @@ static bool gpuToFilePosix(tr::ITensor::SharedPtr const& srcPtr, std::string con TLLM_CHECK_WITH_INFO(cpyErr == cudaSuccess, "cudaMemcpy to host failed, error=%d", cpyErr); ssize_t written = ::write(fd, hostBuffer.data(), numBytes); - TLLM_CHECK_WITH_INFO(written >= 0, "POSIX write error=%zd", written); + TLLM_CHECK_WITH_INFO(written == numBytes, "POSIX write: short/failed write %zd/%zd", written, numBytes); TLLM_LOG_DEBUG("Wrote %zd bytes to %s (POSIX fallback)", written, filename.c_str()); @@ -65,7 +65,7 @@ static bool fileToGpuPosix(tr::ITensor::SharedPtr const& dstPtr, std::string con std::vector hostBuffer(numBytes); ssize_t bytesRead = ::read(fd, hostBuffer.data(), numBytes); - TLLM_CHECK_WITH_INFO(bytesRead >= 0, "POSIX read error=%zd", bytesRead); + TLLM_CHECK_WITH_INFO(bytesRead == numBytes, "POSIX read: short/failed read %zd/%zd", bytesRead, numBytes); TLLM_LOG_DEBUG("Read %zd bytes from %s (POSIX fallback)", bytesRead, filename.c_str()); @@ -478,20 +478,14 @@ void KVCacheTransferManager::diskWriterLoop() // Unstaged jobs (src != nullptr) read the pinned host slot directly; staged jobs read the copy. void const* data = job.src ? job.src : static_cast(job.staged.data()); int fd = ::open(job.filename.c_str(), O_CREAT | O_WRONLY, 0644); - if (fd >= 0) - { - auto const written = ::pwrite(fd, data, job.bytes, 0); - ::close(fd); - if (written != static_cast(job.bytes)) - { - TLLM_LOG_ERROR("[disk-tier] short async write to %s (%zd/%zu)", job.filename.c_str(), - static_cast(written), job.bytes); - } - } - else - { - TLLM_LOG_ERROR("[disk-tier] cannot open %s for async write", job.filename.c_str()); - } + TLLM_CHECK_WITH_INFO(fd >= 0, + "[disk-tier] cannot open %s for async write; failing loudly rather than leaving a corrupt slot", + job.filename.c_str()); + auto const written = ::pwrite(fd, data, job.bytes, 0); + ::close(fd); + TLLM_CHECK_WITH_INFO(written == static_cast(job.bytes), + "[disk-tier] short async write to %s (%zd/%zu); failing loudly rather than leaving a corrupt slot", + job.filename.c_str(), static_cast(written), job.bytes); { std::lock_guard lock(mDiskMutex); @@ -656,24 +650,17 @@ void KVCacheTransferManager::diskReaderLoop() for (size_t i = 0; i < job.dsts.size(); ++i) { int fd = ::open(job.files[i].c_str(), O_RDONLY); - if (fd < 0) - { - TLLM_LOG_ERROR("[disk-tier] cannot open %s for async read", job.files[i].c_str()); - continue; - } + TLLM_CHECK_WITH_INFO(fd >= 0, + "[disk-tier] cannot open %s for async read; failing loudly rather than serving corrupt KV", + job.files[i].c_str()); hostBuffer.resize(job.bytes[i]); auto const got = ::read(fd, hostBuffer.data(), job.bytes[i]); ::close(fd); - if (got == static_cast(job.bytes[i])) - { - TLLM_CUDA_CHECK( - cudaMemcpyAsync(job.dsts[i], hostBuffer.data(), job.bytes[i], cudaMemcpyHostToDevice, stream)); - } - else - { - TLLM_LOG_ERROR("[disk-tier] short async read from %s (%zd/%zu)", job.files[i].c_str(), - static_cast(got), job.bytes[i]); - } + TLLM_CHECK_WITH_INFO(got == static_cast(job.bytes[i]), + "[disk-tier] short async read from %s (%zd/%zu); failing loudly rather than serving corrupt KV", + job.files[i].c_str(), static_cast(got), job.bytes[i]); + TLLM_CUDA_CHECK( + cudaMemcpyAsync(job.dsts[i], hostBuffer.data(), job.bytes[i], cudaMemcpyHostToDevice, stream)); } TLLM_CUDA_CHECK(cudaStreamSynchronize(stream)); } From 85f61b44289644172668769a76d65d0febbb74c8 Mon Sep 17 00:00:00 2001 From: Md Nafis Ul Haque Shifat Date: Wed, 8 Jul 2026 12:13:58 -0700 Subject: [PATCH 22/47] [#16251][fix] pyexecutor: gate disk onboard-readiness sync to plain TP The cross-rank onboard-readiness allreduce/allgather is only correct under plain tensor parallelism; under attention-DP or pipeline parallelism it deadlocks or drops requests. Gate the collective to plain TP, and refuse the disk tier under attention-DP/PP at KV-manager init. Signed-off-by: Md Nafis Ul Haque Shifat --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 36 +++++++++++++------ .../_torch/pyexecutor/resource_manager.py | 7 ++++ 2 files changed, 32 insertions(+), 11 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 17dbc0224316..cc991a8afd63 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -3118,21 +3118,35 @@ def _park_requests_awaiting_onboard(self, scheduled_batch): # all-reduce in the forward would mismatch (hang/crash). AND readiness # across the TP group: admit a request only once its KV has landed on all # ranks. (pp_size==1 -> the TP group is the whole world; TP=1 is a no-op.) - local_has_pending = 0 if all(local_ready) else 1 - if self.dist.tp_size > 1: - any_pending = self.dist.tp_allreduce(int(local_has_pending), + # Cross-rank AND is only correct for plain TP, where every rank holds the identical batch + # (leader-broadcast + deterministic schedule). Under attention-DP the ranks schedule + # independent batches, and under PP the TP-group readiness misses PP peers' shards -- both + # break the collective (deadlock, or silent drops via zip-truncation). Do the cross-rank sync + # only for plain TP; otherwise gate locally (correct for TP=1, and a no-op under ADP/PP since + # disk onboards are refused there at KV-manager init). + tp_sync = (self.dist.tp_size > 1 and not self.enable_attention_dp + and self.dist.pp_size == 1) + if tp_sync: + any_pending = self.dist.tp_allreduce(0 if all(local_ready) else 1, op=ReduceOp.MAX) - else: - any_pending = local_has_pending - if not any_pending: - # Every rank has all context blocks ready -> keep the full batch. - for req in ctx: - req.py_onboard_pending = False - return - if self.dist.tp_size > 1: + if not any_pending: + for req in ctx: + req.py_onboard_pending = False + return gathered = self.dist.tp_allgather(local_ready) global_ready = [all(bits) for bits in zip(*gathered)] else: + if self.dist.tp_size > 1 and not getattr( + self, "_warned_park_local_only", False): + logger.warning( + "disk-onboard readiness: cross-rank park sync is unsupported under " + "attention-DP/PP; gating locally (disk onboards should be disabled here)" + ) + self._warned_park_local_only = True + if all(local_ready): + for req in ctx: + req.py_onboard_pending = False + return global_ready = local_ready kept = [] diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index 8120df9131f9..5176662590c6 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -613,6 +613,13 @@ def append_to_kv_heads_per_layer(num_kv_heads_per_layer: List[int], max_kv_event_entries=self.event_buffer_max_size) disk_cache_size = getattr(kv_cache_config, "disk_cache_size", None) or 0 + if disk_cache_size and (mapping.enable_attention_dp + or mapping.pp_size > 1): + raise ValueError( + "disk KV-cache tier is not supported with attention-DP or pipeline parallelism " + "(cross-rank onboard-readiness sync assumes identical per-rank batches); got " + f"enable_attention_dp={mapping.enable_attention_dp}, pp_size={mapping.pp_size}. " + "Disable disk_cache_size or use plain tensor parallelism.") blocks_in_disk_pool = 0 if disk_cache_size: max_tokens_disk = disk_cache_size // self.get_cache_bytes_per_token( From 5d29113f579927aed0501706b5de0ea6538e2f83 Mon Sep 17 00:00:00 2001 From: Md Nafis Ul Haque Shifat Date: Wed, 8 Jul 2026 12:13:58 -0700 Subject: [PATCH 23/47] [#16251][fix] pyexecutor: re-check can_queue after parking onboard-pending requests can_queue was computed before parking requests whose disk onboard had not landed; if parking emptied the batch we forwarded on nothing. Re-check can_queue after parking, broadening the existing post-batch-change re-check. Signed-off-by: Md Nafis Ul Haque Shifat --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index cc991a8afd63..234a6424fb70 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -3858,8 +3858,10 @@ def _executor_loop(self): if can_queue: self._kv_connector_start_batch(scheduled_batch) - # if using a kv connector, we need to call can_queue again since scheduled_batch might have changed - if self.kv_connector_manager: + # scheduled_batch may have shrunk since the gate above -- parking context requests whose + # disk onboard has not landed empties it, and a kv connector can change it too. Re-check + # before forwarding so we never forward on an empty batch. + if can_queue: can_queue, _ = self._can_queue(scheduled_batch) if not can_queue: @@ -4324,8 +4326,10 @@ def _executor_loop_overlap(self): if can_queue: self._kv_connector_start_batch(scheduled_batch) - # if using a kv connector, we need to call can_queue again since scheduled_batch might have changed - if self.kv_connector_manager: + # scheduled_batch may have shrunk since the gate above -- parking context requests whose + # disk onboard has not landed empties it, and a kv connector can change it too. Re-check + # before forwarding so we never forward on an empty batch. + if can_queue: can_queue, can_queue_this_rank = self._can_queue( scheduled_batch) From b95254ed0fba2e0dc89c241f428acac0b11c25c5 Mon Sep 17 00:00:00 2001 From: Md Nafis Ul Haque Shifat Date: Wed, 8 Jul 2026 14:12:36 -0700 Subject: [PATCH 24/47] [#16251][perf] disk tier: move onboard slot-write wait off the scheduler into the reader loadFromFile ran waitForDiskSlotWrites on the scheduler even on the detached path, so onboarding a just-spilled slot could park the whole engine behind the write queue. Do that wait on the reader instead -- the owning request stays parked via isBlockReadPending until the read lands, so the scheduler is never blocked. The synchronous fallback still waits inline. Signed-off-by: Md Nafis Ul Haque Shifat --- .../batch_manager/kvCacheTransferManager.cpp | 20 ++++++++++++++++--- 1 file changed, 17 insertions(+), 3 deletions(-) diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp index 706f8a46ed99..3edeada1f2eb 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp @@ -617,6 +617,17 @@ void KVCacheTransferManager::diskReaderLoop() mReadQueue.pop(); } + // A just-spilled slot may still have its write queued; wait it out before reading -- but HERE on the + // reader, not on the scheduler in loadFromFile. The owning request stays parked via isBlockReadPending + // until this read lands, so the engine loop is never blocked behind the writer queue. + if (mAsyncDiskStore) + { + for (auto const& f : job.files) + { + waitForDiskSlotWrites(f); + } + } + // Perform the block's transfer. GDS issues one batched SSD->GPU DMA for all pools; POSIX reads each // pool into a host buffer and copies it up. Both are device-complete before we publish the block. if (job.useGds) @@ -768,8 +779,11 @@ void KVCacheTransferManager::loadFromFile(BlockPtr const& dstPrimaryBlock, SizeT { it->second.synchronize(); } - // Async store may still be persisting this slot; a read must not race the write. - if (mAsyncDiskStore) + // Async store may still be persisting this slot; a read must not race the write. On the detached path + // the reader performs this wait itself (see diskReaderLoop), so it never blocks the scheduler; only the + // synchronous fallback below waits here. + bool const detach = asyncDiskReadEnabled() && trackBlockId >= 0; + if (mAsyncDiskStore && !detach) { for (size_t poolIdx = 0; poolIdx < pools.size(); ++poolIdx) { @@ -780,7 +794,7 @@ void KVCacheTransferManager::loadFromFile(BlockPtr const& dstPrimaryBlock, SizeT // and hand it off. The reader performs the transfer (GDS DMA if enabled, else POSIX read + copy) and makes // it device-complete before publishing the block; the request is forward-safe once areBlocksReady() sees // its blocks land. trackBlockId < 0 (no pool) takes the synchronous fallback below. - if (asyncDiskReadEnabled() && trackBlockId >= 0) + if (detach) { DiskReadJob job; job.blockId = trackBlockId; From 5eec1bfea29a23d96615262ca83d76e7dd3a02b6 Mon Sep 17 00:00:00 2001 From: Md Nafis Ul Haque Shifat Date: Wed, 8 Jul 2026 14:12:36 -0700 Subject: [PATCH 25/47] [#16251][perf] disk tier: retention-aware write-pressure policy (opt-in TLLM_KV_DISK_DROP_ON_PRESSURE) Under write-queue saturation the scheduler blocked in enqueueDiskWrite* waiting for room. Add an opt-in policy (TLLM_KV_DISK_DROP_ON_PRESSURE, default off = prior wait behavior): best-effort (non-retained) spills are shed at the eviction gate instead of stalling; retained spills never drop and bypass the queue cap (their TTL guarantee requires they land). Drop is decided per-spill at the gate, not per-pool in enqueue*, which would otherwise leave partial/corrupt blocks and leak the reserved-pool reap. Signed-off-by: Md Nafis Ul Haque Shifat --- .../batch_manager/kvCacheManager.h | 1 + .../batch_manager/kvCacheTransferManager.h | 10 ++++- .../batch_manager/kvCacheManager.cpp | 19 +++++++++ .../batch_manager/kvCacheTransferManager.cpp | 39 ++++++++++++------- 4 files changed, 54 insertions(+), 15 deletions(-) diff --git a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h index dda63a123ed2..7b34a84ce169 100644 --- a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h +++ b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h @@ -1452,6 +1452,7 @@ class WindowBlockManager bool mDiskProtectUnexpired{false}; std::size_t mDiskSpills{0}; std::size_t mDiskGateDropped{0}; + std::size_t mDiskWritePressureDropped{0}; std::size_t mDiskAdmissionRefused{0}; std::size_t mDiskOnboards{0}; diff --git a/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h b/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h index 0c282d1b948f..d9a2e3ae4337 100644 --- a/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h +++ b/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h @@ -127,6 +127,10 @@ class KVCacheTransferManager return mAsyncDiskStore; } + //! \brief True if the async write queue is at capacity. Used to shed best-effort (non-retained) spills under + //! writer saturation; retained spills bypass the cap and are never shed. + [[nodiscard]] bool diskWriteQueueFull(); + void syncTransfers(); ~KVCacheTransferManager(); @@ -193,6 +197,7 @@ class KVCacheTransferManager std::vector staged; // staged-memcpy path (used when src == nullptr) void const* src{nullptr}; // unstaged path: writer reads this pinned host pointer directly std::uint64_t spillId{0}; // >0 => track per-spill completion for the reserved-pool reap + bool retained{false}; // retained spill: bypasses the queue cap, never dropped }; std::vector mDiskWriters; @@ -220,8 +225,9 @@ class KVCacheTransferManager }()}; void diskWriterLoop(); - void enqueueDiskWrite(std::string filename, void const* src, std::size_t bytes); - void enqueueDiskWriteUnstaged(std::string filename, void const* src, std::size_t bytes, std::uint64_t spillId); + void enqueueDiskWrite(std::string filename, void const* src, std::size_t bytes, bool retained); + void enqueueDiskWriteUnstaged( + std::string filename, void const* src, std::size_t bytes, std::uint64_t spillId, bool retained); void waitForDiskSlotWrites(std::string const& filename); // ---- Disk-tier async ONBOARD: read slot files disk->GPU OFF the scheduler thread ---- diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp index 03a0beaab43b..1d0bb5b73dc8 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp @@ -1431,6 +1431,25 @@ BlockPtr WindowBlockManager::reclaimSecondaryBlock() } return victim; // reused fewer than k times: discard, don't spill } + // Write-pressure shedding (opt-in via TLLM_KV_DISK_DROP_ON_PRESSURE): when the async writer is saturated, + // drop a BEST-EFFORT (non-retained) spill rather than stalling the scheduler behind the write queue. + // Retained blocks are never shed -- their TTL guarantee requires they reach disk (they bypass the cap at + // enqueue). Default (unset/0) preserves prior behavior: best-effort spills wait for queue room. + static bool const kDropOnPressure = [] + { + auto const* e = std::getenv("TLLM_KV_DISK_DROP_ON_PRESSURE"); + return e && std::atoi(e) != 0; + }(); + if (kDropOnPressure && !victim->isRetainedNow() && mTransferManager->diskWriteQueueFull()) + { + ++mDiskWritePressureDropped; + if (mDiskWritePressureDropped == 1 || mDiskWritePressureDropped % 10000 == 0) + { + TLLM_LOG_INFO( + "[disk-tier] write-pressure dropped=%zu (windowSize=%d)", mDiskWritePressureDropped, mWindowSize); + } + return victim; // best-effort spill shed under writer saturation: evict as stock (no scheduler stall) + } auto diskTarget = claimDiskTarget(); if (diskTarget == nullptr) { diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp index 3edeada1f2eb..10fa2f78bc5a 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp @@ -509,24 +509,27 @@ void KVCacheTransferManager::diskWriterLoop() } } -void KVCacheTransferManager::enqueueDiskWrite(std::string filename, void const* src, std::size_t bytes) +void KVCacheTransferManager::enqueueDiskWrite(std::string filename, void const* src, std::size_t bytes, bool retained) { DiskWriteJob job; job.filename = std::move(filename); job.bytes = bytes; + job.retained = retained; job.staged.resize(bytes); std::memcpy(job.staged.data(), src, bytes); // fast host->host copy; frees the slot { std::unique_lock lock(mDiskMutex); - // Backpressure + per-slot serialization: wait for queue room AND no earlier write - // to this same slot still outstanding (so a reused slot cannot race). + // Per-slot serialization (always) + backpressure. A retained spill must land, so it bypasses the + // queue cap (bounded by retained volume); a best-effort spill also waits for room. Best-effort shedding + // under saturation happens upstream at the eviction gate, so a best-effort spill reaching here + then + // finding the queue full is the rare TOCTOU case. mDiskQueueCv.wait(lock, [this, &job] { - return (mDiskWriteQueue.size() < mDiskWriteQueueMax - && mDiskInflight.find(job.filename) == mDiskInflight.end()) - || mDiskWriterStop; + bool const slotClear = mDiskInflight.find(job.filename) == mDiskInflight.end(); + bool const roomOrRetained = job.retained || mDiskWriteQueue.size() < mDiskWriteQueueMax; + return (slotClear && roomOrRetained) || mDiskWriterStop; }); if (mDiskWriterStop) { @@ -539,23 +542,25 @@ void KVCacheTransferManager::enqueueDiskWrite(std::string filename, void const* } void KVCacheTransferManager::enqueueDiskWriteUnstaged( - std::string filename, void const* src, std::size_t bytes, std::uint64_t spillId) + std::string filename, void const* src, std::size_t bytes, std::uint64_t spillId, bool retained) { DiskWriteJob job; job.filename = std::move(filename); job.bytes = bytes; job.src = src; // no staging: the writer reads the pinned host slot directly job.spillId = spillId; + job.retained = retained; { std::unique_lock lock(mDiskMutex); - // Same backpressure + per-slot serialization as enqueueDiskWrite. + // Same per-slot serialization + backpressure as enqueueDiskWrite: retained spills bypass the queue cap + // (never stall the scheduler), best-effort spills also wait for room. mDiskQueueCv.wait(lock, [this, &job] { - return (mDiskWriteQueue.size() < mDiskWriteQueueMax - && mDiskInflight.find(job.filename) == mDiskInflight.end()) - || mDiskWriterStop; + bool const slotClear = mDiskInflight.find(job.filename) == mDiskInflight.end(); + bool const roomOrRetained = job.retained || mDiskWriteQueue.size() < mDiskWriteQueueMax; + return (slotClear && roomOrRetained) || mDiskWriterStop; }); if (mDiskWriterStop) { @@ -586,6 +591,12 @@ void KVCacheTransferManager::waitForDiskSlotWrites(std::string const& filename) }); } +bool KVCacheTransferManager::diskWriteQueueFull() +{ + std::lock_guard lock(mDiskMutex); + return mDiskWriteQueue.size() >= mDiskWriteQueueMax; +} + void KVCacheTransferManager::diskReaderLoop() { // Each reader owns its own CUDA stream (H2D copies stay off the null stream and the onboard/offload @@ -707,6 +718,7 @@ void KVCacheTransferManager::spillToFile(BlockPtr const& srcHostBlock, SizeType3 std::vector const& pools, std::string const& directory) { TLLM_CHECK_WITH_INFO(!directory.empty(), "disk tier requires a directory"); + bool const retained = srcHostBlock->isRetainedNow(); // retained spill must land -> bypasses the queue cap // The victim's bytes may still be the target of an in-flight async GPU->host copy; // wait it out before reading host memory (same event discipline as copyBlock). auto const idx = getPendingTransferIndex(srcHostBlock); @@ -732,7 +744,7 @@ void KVCacheTransferManager::spillToFile(BlockPtr const& srcHostBlock, SizeType3 continue; } // Async path: copy bytes out (frees the slot immediately) and hand to the writer. - enqueueDiskWrite(filename, ptr->data(), bytes); + enqueueDiskWrite(filename, ptr->data(), bytes, retained); } } @@ -741,6 +753,7 @@ void KVCacheTransferManager::spillToFileUnstaged(BlockPtr const& srcHostBlock, S { TLLM_CHECK_WITH_INFO(!directory.empty(), "disk tier requires a directory"); TLLM_CHECK_WITH_INFO(mAsyncDiskStore, "spillToFileUnstaged requires the async writer"); + bool const retained = srcHostBlock->isRetainedNow(); // retained spill must land -> bypasses the queue cap // Same event discipline as spillToFile: the victim's bytes may still be the target of an in-flight // GPU->host copy; block until it lands before the writer reads host memory. auto const idx = getPendingTransferIndex(srcHostBlock); @@ -762,7 +775,7 @@ void KVCacheTransferManager::spillToFileUnstaged(BlockPtr const& srcHostBlock, S auto const bytes = static_cast(ptr->getSizeInBytes()); // The pool buffer outlives the manager, so this raw pointer stays valid after `ptr` (a view) // is destroyed -- the writer reads it later off-thread. - enqueueDiskWriteUnstaged(filename, ptr->data(), bytes, spillId); + enqueueDiskWriteUnstaged(filename, ptr->data(), bytes, spillId, retained); } } From 326420f9e52a85274a9a71c1d9450bd88d32c8ba Mon Sep 17 00:00:00 2001 From: Md Nafis Ul Haque Shifat Date: Wed, 8 Jul 2026 15:21:17 -0700 Subject: [PATCH 26/47] [#16251][perf] disk tier: stage onboard reads through a pinned host buffer The POSIX reader staged each onboard through a pageable std::vector, so cudaMemcpyAsync fell back to the driver's hidden pinned-staging copy -- H2D ran at ~5-8 GB/s and the call was secretly synchronous. Stage through a per-reader pinned buffer instead (grown on demand to one block's pool bytes; pageable fallback if pinning is refused). Each pool reads into its own offset with a single stream sync after the block, so the async DMA never races the next read and pool reads overlap prior pools' DMA. Onboard latency is parked-request TTFT for disk-cache hits. Signed-off-by: Md Nafis Ul Haque Shifat --- .../batch_manager/kvCacheTransferManager.cpp | 52 +++++++++++++++++-- 1 file changed, 47 insertions(+), 5 deletions(-) diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp index 10fa2f78bc5a..dfbb3ba5fa79 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp @@ -608,7 +608,12 @@ void KVCacheTransferManager::diskReaderLoop() TLLM_CUDA_CHECK(cudaSetDevice(mDeviceId)); cudaStream_t stream; TLLM_CUDA_CHECK(cudaStreamCreate(&stream)); - std::vector hostBuffer; + // Per-reader PINNED staging buffer, grown on demand to one job's total pool bytes (fixed after the first + // job). Pinned so cudaMemcpyAsync is a real async DMA (~20+ GB/s) instead of the driver's hidden pageable + // staging copy (~5-8 GB/s, secretly synchronous). pageableFallback is used only if a pinned alloc is refused. + std::uint8_t* pinnedBuf = nullptr; + std::size_t pinnedCap = 0; + std::vector pageableFallback; std::shared_ptr gdsAgent; // per-reader GDS agent, created lazily on the first GDS job while (true) { @@ -669,20 +674,53 @@ void KVCacheTransferManager::diskReaderLoop() } if (!job.useGds) { + // Stage the whole block through the pinned buffer, each pool at its OWN offset, so no region is + // overwritten while its DMA is still reading it (a pinned cudaMemcpyAsync returns before the DMA + // finishes). One stream sync after the loop makes the buffer safe to reuse for the next job. + std::size_t total = 0; + for (auto const b : job.bytes) + { + total += b; + } + if (total > pinnedCap) + { + if (pinnedBuf != nullptr) + { + cudaFreeHost(pinnedBuf); + pinnedBuf = nullptr; + } + if (cudaHostAlloc(reinterpret_cast(&pinnedBuf), total, cudaHostAllocDefault) == cudaSuccess) + { + pinnedCap = total; + } + else + { + pinnedBuf = nullptr; + pinnedCap = 0; + TLLM_LOG_WARNING( + "[disk-tier] pinned host alloc of %zu B refused; falling back to pageable staging", total); + } + } + std::uint8_t* buf = pinnedBuf; + if (buf == nullptr) + { + pageableFallback.resize(total); + buf = pageableFallback.data(); + } + std::size_t off = 0; for (size_t i = 0; i < job.dsts.size(); ++i) { int fd = ::open(job.files[i].c_str(), O_RDONLY); TLLM_CHECK_WITH_INFO(fd >= 0, "[disk-tier] cannot open %s for async read; failing loudly rather than serving corrupt KV", job.files[i].c_str()); - hostBuffer.resize(job.bytes[i]); - auto const got = ::read(fd, hostBuffer.data(), job.bytes[i]); + auto const got = ::read(fd, buf + off, job.bytes[i]); ::close(fd); TLLM_CHECK_WITH_INFO(got == static_cast(job.bytes[i]), "[disk-tier] short async read from %s (%zd/%zu); failing loudly rather than serving corrupt KV", job.files[i].c_str(), static_cast(got), job.bytes[i]); - TLLM_CUDA_CHECK( - cudaMemcpyAsync(job.dsts[i], hostBuffer.data(), job.bytes[i], cudaMemcpyHostToDevice, stream)); + TLLM_CUDA_CHECK(cudaMemcpyAsync(job.dsts[i], buf + off, job.bytes[i], cudaMemcpyHostToDevice, stream)); + off += job.bytes[i]; } TLLM_CUDA_CHECK(cudaStreamSynchronize(stream)); } @@ -694,6 +732,10 @@ void KVCacheTransferManager::diskReaderLoop() mReadInflightCount.store(mPendingBlockReads.size(), std::memory_order_release); } } + if (pinnedBuf != nullptr) + { + cudaFreeHost(pinnedBuf); + } cudaStreamDestroy(stream); } From 04e2ab3faf88d19f38ea2b28d4d769dfb6d572f0 Mon Sep 17 00:00:00 2001 From: Md Nafis Ul Haque Shifat Date: Wed, 8 Jul 2026 15:43:35 -0700 Subject: [PATCH 27/47] [#16251][perf] pyexecutor: skip onboard-readiness park + its TP collective when the disk tier is off _park_requests_awaiting_onboard ran on every context step and, on plain TP>1, issued a per-step object-pickling tp_allreduce -- even with no disk tier, because it only bailed on hasattr(are_blocks_ready), which is always true. That taxed every TP>1 model sharing the image. Gate the park on a flag captured at init (disk_cache_size>0 AND TLLM_KV_DISK_READERS>0): detached onboards are impossible otherwise, so are_blocks_ready() is always true and the park + its collective are pure overhead. No-op for disk deployments (flag true -> unchanged); the skip only affects non-disk / sync-onboard configs. Signed-off-by: Md Nafis Ul Haque Shifat --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 234a6424fb70..dc8261a2e646 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -580,6 +580,14 @@ def __init__( self.max_draft_len = max_draft_len self.max_total_draft_tokens = max_total_draft_tokens self.llm_args = self.model_engine.llm_args + # Detached disk-onboard readiness (and its per-step cross-rank collective in + # _park_requests_awaiting_onboard) is only reachable when BOTH the disk pool and the async reader + # pool are enabled; otherwise are_blocks_ready() is always true, so the park is pure overhead. Capture + # once so non-disk models -- including every TP>1 model sharing this image -- pay nothing for it. + _kvc = getattr(self.llm_args, "kv_cache_config", None) + self._disk_onboard_active = bool( + _kvc is not None and getattr(_kvc, "disk_cache_size", None) + and int(os.environ.get("TLLM_KV_DISK_READERS", "0") or "0") > 0) self.max_stats_len = self.llm_args.max_stats_len self.max_num_tokens = self.llm_args.max_num_tokens self.print_log = self.llm_args.print_iter_log @@ -3101,6 +3109,11 @@ def _park_requests_awaiting_onboard(self, scheduled_batch): request stays active and is re-scheduled next step -- the capacity scheduler counts its already allocated blocks (getRemainingBlocksToCompletion) and _context_seq_len skips re-adding it -- then it is admitted once are_blocks_ready() reports its blocks landed.""" + # No detached disk onboards possible (no disk pool or no async readers) -> are_blocks_ready() is + # always true and there is nothing to synchronize. Skip entirely so non-disk (esp. TP>1) models never + # pay the per-step cross-rank collective below. + if not self._disk_onboard_active: + return kv = getattr(self, "kv_cache_manager", None) if kv is None or not hasattr(kv, "are_blocks_ready"): return From b27c6c94343821498a8678ab027a4dae35da02c9 Mon Sep 17 00:00:00 2001 From: Md Nafis Ul Haque Shifat Date: Thu, 9 Jul 2026 11:56:54 -0700 Subject: [PATCH 28/47] [#16251][fix] disk tier: close two silent-KV-corruption paths on onboard/release Two paths could hand back valid-looking KV that is actually the wrong tokens: 1. releaseBlocks freed a block whose detached disk-onboard read was still in flight. areBlocksReady() parks the forward path on read-readiness, but nothing gated release: a request completing/cancelling mid-onboard returned its DMA-destination block to the free queue, and getFreeBlock has no readiness check -- so a stranger's prefill and the late read collided in the same GPU memory. Mirror the park onto the release path: hold read-pending blocks in mReleaseReadPending until the DMA lands (reapReadPendingReleases frees them). Gated by mNumDiskBlocks>0 && anyReadPending(), so non-disk models are unaffected. 2. reclaimSecondaryBlock's sync (copy) spill returned diskTarget without detaching its displaced disk-cached identity, relying on the caller's cleanup. getFreeBlock detaches; offloadBlock (the other caller) does not -- leaving a radix-tree node pointing at the recycled host slot, so a later prefix match reused freed content. Detach the identity unconditionally in reclaimSecondaryBlock for both callers; getFreeBlock's later detach then becomes a harmless no-op. Validated: kvCacheManagerTest disk-tier suite 11/11 green (spill/onboard/release + byte-exact onboard); the 12 pre-existing non-disk failures are unchanged. Signed-off-by: Md Nafis Ul Haque Shifat --- .../batch_manager/kvCacheManager.h | 4 ++ .../batch_manager/kvCacheManager.cpp | 72 +++++++++++++------ 2 files changed, 55 insertions(+), 21 deletions(-) diff --git a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h index 7b34a84ce169..9a63715682cb 100644 --- a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h +++ b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h @@ -1513,6 +1513,10 @@ class WindowBlockManager // Held out of the disk free queue until the read lands so a spill cannot overwrite the slot file // mid-read; reapReadPendingReleases() returns each card once its read completes. std::unordered_map mReadPendingReleases; + // Blocks whose owning request was released while their detached disk-onboard read was still in flight: + // held out of the free queue until the read lands (reapReadPendingReleases frees them) so the DMA + // destination is never repurposed mid-read. See releaseBlocks / the areBlocksReady park. + std::vector mReleaseReadPending; void reapReadPendingReleases(); std::uint64_t mUnstagedSpillSeq{0}; diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp index 1d0bb5b73dc8..5f8a498e0fed 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp @@ -1336,6 +1336,19 @@ void WindowBlockManager::reapReadPendingReleases() mEvictionPolicy->releaseBlock(it->second); it = mReadPendingReleases.erase(it); } + // Blocks whose owning request was released (completed/cancelled) while their detached onboard read was + // still in flight (see releaseBlocks): free them only once the DMA has landed, so the destination GPU + // memory is never handed to another request mid-read. + for (auto it = mReleaseReadPending.begin(); it != mReleaseReadPending.end();) + { + if (mTransferManager->isBlockReadPending((*it)->getBlockId())) + { + ++it; + continue; + } + mEvictionPolicy->releaseBlock(*it); + it = mReleaseReadPending.erase(it); + } } BlockPtr WindowBlockManager::claimDiskTarget() @@ -1499,27 +1512,30 @@ BlockPtr WindowBlockManager::reclaimSecondaryBlock() { TLLM_LOG_INFO("[disk-tier] spills=%zu (windowSize=%d)", mDiskSpills, mWindowSize); } - if (useUnstaged) + // diskTarget carries a displaced disk-cached identity (it was a cached disk slot before we claimed it + // as a spill target). Detach that identity from the reuse tree UNCONDITIONALLY here: getFreeBlock (one + // caller) would detach it via its own cleanup below, but offloadBlock (the other caller) does NOT -- + // leaving a radix-tree node pointing at the recycled host slot => silent wrong-token reuse. Detaching + // for both callers closes that hole; getFreeBlock's later detach then becomes a harmless no-op. { - // diskTarget now owns the victim's old host slot, which the writer is still reading. It will - // NOT flow through getFreeBlock's cleanup (we return a reserved block instead), so evict any - // displaced disk-block identity it still carries HERE, then pin it until the write drains. + std::lock_guard treeLock(mLookupTree->getMutex()); + if (mEventManager && blockInRadixTree(diskTarget)) { - std::lock_guard treeLock(mLookupTree->getMutex()); - if (mEventManager && blockInRadixTree(diskTarget)) - { - mEventManager->enqueueRemovedEvent(diskTarget, mWindowSize); - } - diskTarget->detachFromLookupNode(); + mEventManager->enqueueRemovedEvent(diskTarget, mWindowSize); } - diskTarget->clearRetention(); + diskTarget->detachFromLookupNode(); + } + diskTarget->clearRetention(); + if (useUnstaged) + { + // Unstaged: diskTarget owns the victim's old host slot, which the writer is still reading. Pin it + // until the write drains so nobody reuses the slot mid-write; hand the caller a clean reserved block. mPendingSpillBlocks[unstagedSpillId] = diskTarget; // pinned: keeps the host slot out of every free list auto reserved = mReservedHostBlocks.back(); // clean, claimed, secondary host block mReservedHostBlocks.pop_back(); return reserved; } - diskTarget->clearRetention(); - return diskTarget; // owns the victim's old host slot; already claimed + return diskTarget; // owns the victim's old host slot; already claimed, identity detached above } BlockPtr WindowBlockManager::getFreeBlock(GenerationRequest& sequence, executor::RetentionPriority priority, @@ -3518,6 +3534,9 @@ std::optional WindowBlockManager::releaseBlocks( TLLM_LOG_DEBUG("%s::releaseBlocks Request %lu, %d blocks stored for reuse", mLogPrefix.c_str(), sequence.getRequestId(), numBlocksStoredForReuse); } + // Fast-path gate: only consult per-block read-pendingness when a detached disk onboard could be in + // flight. Non-disk models (and the steady state with no reads queued) take the original path for free. + bool const guardReadPending = mNumDiskBlocks > 0 && mTransferManager->anyReadPending(); // Iterate all allocated blocks (including placeholder sentinels at OOW positions); // EvictionPolicy::releaseBlock silently skips placeholders. for (auto it = allocatedBlocks.rbegin(); it != allocatedBlocks.rend(); ++it) @@ -3533,16 +3552,27 @@ std::optional WindowBlockManager::releaseBlocks( // mRefCount==0 and are silently ignored by EvictionPolicy::releaseBlock(). if (!block->hasRefs()) { - auto const isDetached = block->isDetached(); - if (isDetached) + // Silent-corruption guard (disk tier), the RELEASE-path mirror of areBlocksReady()'s park: if a + // detached disk-onboard read is still filling this block, returning it to the free queue now lets + // getFreeBlock hand it to another request while the DMA is in flight -- their prefill and the late + // read collide in the same GPU memory. Hold it out until the read lands; reapReadPendingReleases() + // releases it once !isBlockReadPending. + if (guardReadPending && mTransferManager->isBlockReadPending(block->getBlockId())) { - // Detached blocks have no reusable hash chain. Drop the stale owning link before recycling the block; - // otherwise front-queue reuse can join completed request chains into an unbounded ownership chain. - block->setPrevBlockInSeq(nullptr); + mReleaseReadPending.push_back(block); + } + else + { + auto const isDetached = block->isDetached(); + if (isDetached) + { + // Detached blocks have no reusable hash chain. Drop the stale owning link before recycling; + // otherwise front-queue reuse can join completed request chains into an unbounded chain. + block->setPrevBlockInSeq(nullptr); + } + // Send blocks with no reusable state to the front so they evict before cached-for-reuse blocks. + mEvictionPolicy->releaseBlock(block, /*toFront=*/isDetached); } - // Send block to front of free queue if it has no reusable state, - // so detached blocks are evicted before blocks cached for reuse. - mEvictionPolicy->releaseBlock(block, /*toFront=*/isDetached); } } // Remove stored block ids in sequence From 97eee6e6bf6307bbdacc2f379d5c862f9ff21833 Mon Sep 17 00:00:00 2001 From: Md Nafis Ul Haque Shifat Date: Thu, 9 Jul 2026 12:17:31 -0700 Subject: [PATCH 29/47] [#16251][fix] disk tier: reap read-pending releases before the disk-free check reapReadPendingReleases() -- which returns landed onboard cards (and, since the release-vs-onboard fix, release-deferred blocks) to their free queues -- ran only inside claimDiskTarget(). But reclaimSecondaryBlock() reaches claimDiskTarget only after passing a getNumFreeBlocks(kDiskLevel)==0 early-return, and parked cards sit OUTSIDE the free queue. So an onboard-heavy burst that parks every free disk card drives the free count to 0, fires the early-return, and never reaches claimDiskTarget (the only reap site) -- the landed-but-parked cards are never reclaimed and spills wedge permanently, exactly under the load the retained tier is for. Reap right after claiming the victim, before the disk-free check, guarded by mNumDiskBlocks>0. Placed in reclaimSecondaryBlock so it covers both callers (getFreeBlock and offloadBlock); claimDiskTarget's reap stays as a harmless second pass. Validated: kvCacheManagerTest disk-tier suite 11/11 green; the 12 pre-existing non-disk failures are unchanged. Signed-off-by: Md Nafis Ul Haque Shifat --- cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp index 5f8a498e0fed..0268be7cd8db 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp @@ -1414,6 +1414,16 @@ BlockPtr WindowBlockManager::reclaimSecondaryBlock() } auto victim = std::get<0>(mEvictionPolicy->getFreeBlock(kSecondaryLevel)); mEvictionPolicy->claimBlock(victim); + // Reap landed read-pending releases BEFORE testing disk free space. Parked onboarded cards (and + // release-deferred blocks) sit OUTSIDE the free queues until reaped, so an onboard-heavy burst that + // parks every free disk card would drive getNumFreeBlocks(kDiskLevel) to 0, fire the early-return + // below, and never reach claimDiskTarget() -- the only other reap site -- wedging spills permanently. + // Reaping here (covering both reclaimSecondaryBlock callers) lets landed cards rejoin the disk free + // queue so the check sees the true count. + if (mNumDiskBlocks > 0) + { + reapReadPendingReleases(); + } if (mNumDiskBlocks == 0 || victim->getUniqueTokens().empty() || mEvictionPolicy->getNumFreeBlocks(kDiskLevel) == 0) { return victim; // disk tier off, nothing reusable to keep, or disk full: same as before From a9884062ffbe84f208697295f0592b00e765f5a2 Mon Sep 17 00:00:00 2001 From: Md Nafis Ul Haque Shifat Date: Thu, 9 Jul 2026 13:05:07 -0700 Subject: [PATCH 30/47] [#16251][fix] serve: don't pass kv_cache_retention_config to the mm-encoder openai_mm_encoder forwarded kv_cache_retention_config=_disk_retention_config(request) to MultimodalEncoder.generate_async, whose signature is (inputs, sampling_params) -- no such kwarg and no **kwargs -- so every multimodal-encoder request 500s, even with the disk tier off (the value is passed unconditionally). Encoders emit embeddings and hold no KV cache to retain, so drop the kwarg on that path. The four LLM/completion call sites are unchanged (their generate_async accepts it). Signed-off-by: Md Nafis Ul Haque Shifat --- tensorrt_llm/serve/openai_server.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/serve/openai_server.py b/tensorrt_llm/serve/openai_server.py index 5061a53dad7c..c774ebe10172 100644 --- a/tensorrt_llm/serve/openai_server.py +++ b/tensorrt_llm/serve/openai_server.py @@ -1542,9 +1542,10 @@ async def create_mm_embedding_response(promise: RequestOutput): if mm_data is not None: prompt["multi_modal_data"] = mm_data - promise = self.generator.generate_async( - inputs=prompt, - kv_cache_retention_config=_disk_retention_config(request)) + # NB: MultimodalEncoder.generate_async(inputs, sampling_params) takes no + # kv_cache_retention_config (and no **kwargs); encoders emit embeddings and hold no + # KV cache to retain. Passing it here 500s every mm-encoder request, so omit it. + promise = self.generator.generate_async(inputs=prompt) asyncio.create_task(self.await_disconnected(raw_request, promise)) response = await create_mm_embedding_response(promise) From 4e12a209ede218ab2348ece1024070f2517ecc14 Mon Sep 17 00:00:00 2001 From: Md Nafis Ul Haque Shifat Date: Tue, 14 Jul 2026 15:36:45 -0700 Subject: [PATCH 31/47] [#16251][fix] kv cache: shut down disk workers before releasing pools releasePools() freed the GPU/host KV pools while the disk worker threads were still running -- they were stopped only in the ~KVCacheTransferManager destructor, which runs later. Queued onboard reads hold raw GPU dst pointers and unstaged spills hold raw host src pointers into those pools, so a worker could touch freed pool memory during shutdown (use-after-free). Extract the stop+join into an idempotent shutdownDiskWorkers() and invoke it at the top of releasePools(), before the pools are freed; the destructor now delegates to the same operation. Add a shutdown test with queued reads and unstaged writes. Signed-off-by: Md Nafis Ul Haque Shifat --- .../batch_manager/kvCacheTransferManager.h | 7 ++ .../batch_manager/kvCacheManager.cpp | 3 + .../batch_manager/kvCacheTransferManager.cpp | 14 +++- .../batch_manager/kvCacheManagerTest.cpp | 69 +++++++++++++++++++ 4 files changed, 92 insertions(+), 1 deletion(-) diff --git a/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h b/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h index d9a2e3ae4337..d5d90fc60493 100644 --- a/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h +++ b/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h @@ -111,6 +111,10 @@ class KVCacheTransferManager return mNumDiskReaders >= 1; } + //! \brief Idempotently stop+join all disk workers. MUST run before the pools are released: queued + //! reads hold raw GPU dst and unstaged writes raw host src pointers into them (else use-after-free). + void shutdownDiskWorkers(); + //! \brief Disk tier (unstaged async): spill a host block's bytes by handing the writer the source //! pointers directly (no staging memcpy). The caller MUST keep the source host slot pinned until //! \p spillId appears in drainCompletedSpills(). @@ -259,6 +263,9 @@ class KVCacheTransferManager // areBlocksReady() can skip its per-block scan + the mutex when nothing is in flight (common case). std::atomic mReadInflightCount{0}; bool mDiskReaderStop{false}; + // Set once shutdownDiskWorkers() has stopped+joined the disk workers, so releasePools() and the + // destructor do not repeat the work. Only touched on the owning thread. + bool mDiskWorkersShutdown{false}; void diskReaderLoop(); void enqueueDiskRead(DiskReadJob job); diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp index 0268be7cd8db..2875bd49a918 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp @@ -1240,6 +1240,9 @@ void WindowBlockManager::releasePools() { if (mTransferManager) { + // Stop the disk workers before freeing the pools: queued reads / unstaged writes hold raw + // GPU dst / host src pointers into them, so a worker running after release is a use-after-free. + mTransferManager->shutdownDiskWorkers(); mTransferManager->syncTransfers(); } mBufferManager.getStream().synchronize(); diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp index dfbb3ba5fa79..de49f0b2ee5b 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp @@ -104,8 +104,15 @@ KVCacheTransferManager::KVCacheTransferManager( } } -KVCacheTransferManager::~KVCacheTransferManager() +void KVCacheTransferManager::shutdownDiskWorkers() { + // Idempotent stop+join of all disk workers. Queued reads hold raw GPU dst pointers and unstaged + // writes hold raw host src pointers into the pools, so this must run before the pools are freed. + if (mDiskWorkersShutdown) + { + return; + } + mDiskWorkersShutdown = true; if (!mDiskWriters.empty()) { { @@ -138,6 +145,11 @@ KVCacheTransferManager::~KVCacheTransferManager() } } +KVCacheTransferManager::~KVCacheTransferManager() +{ + shutdownDiskWorkers(); +} + tr::ITensor::SharedPtr KVCacheTransferManager::computeBlockPointer( BlockPtr const& block, std::vector const& pools, size_t poolIdx) { diff --git a/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp b/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp index 05cf1724e4c2..140780d3b8b0 100644 --- a/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp @@ -11043,3 +11043,72 @@ TEST_F(KVCacheManagerTest, DiskTierSyncOnboardByteExactTest) { runDiskOnboardByteRoundTrip(/*readers=*/"0", /*nSlots=*/4); } + +// shutdownDiskWorkers() must stop+join every disk worker while reads are queued (holding raw GPU dst +// pointers) and unstaged writes are queued (holding raw host src pointers), so releasePools() can free +// the pools with no use-after-free -- and it must be idempotent. Passing == no hang/crash reaching the +// end (and, under ASAN, no worker access to the released pool). Regression test for the shutdown-order +// bug where workers stopped only in the destructor, after releasePools() had already freed the pools. +TEST_F(KVCacheManagerTest, DiskTierShutdownWithQueuedWorkTest) +{ + setenv("TLLM_KV_DISK_READERS", "4", /*overwrite=*/1); + setenv("TLLM_KV_DISK_WRITERS", "4", /*overwrite=*/1); + setenv("TLLM_KV_DISK_ASYNC_STORE", "1", /*overwrite=*/1); + DiskTierDir dir; + int constexpr blockSize = 2048; // floats per slot + int constexpr nSlots = 8; + + auto bufferManager = tr::BufferManager(std::make_shared()); + auto transferManager = KVCacheTransferManager(bufferManager); // spawns readers + writers from the env above + + auto pool = KVCacheBlockPool(0, /*kvFactor=*/2, 0, 0, 0); + pool.primaryPtr = bufferManager.gpu(tr::ITensor::makeShape({2 * nSlots, blockSize}), nvinfer1::DataType::kFLOAT); + bufferManager.setZero(*pool.primaryPtr); + pool.secondaryPtr + = tr::BufferManager::pinned(tr::ITensor::makeShape({2 * nSlots, blockSize}), nvinfer1::DataType::kFLOAT); + + // Seed host slots and write them to disk (synchronous) so the queued reads below have files to read. + float* secBase = tr::bufferCast(*pool.secondaryPtr); + for (int k = 0; k < nSlots; ++k) + { + for (int i = 0; i < blockSize; ++i) + { + secBase[static_cast(k) * blockSize + i] = diskPatternValue(k, i); + } + auto seed = std::make_shared(k, tk::KVCacheIndex(k, /*isSecondary=*/true)); + transferManager.spillToFile(seed, /*diskSlot=*/k, {pool}, dir.str()); + } + bufferManager.getStream().synchronize(); + + // Queue detached onboard reads: each holds raw GPU dst pointers into the primary pool. + std::vector reads; + for (int k = 0; k < nSlots; ++k) + { + auto d = std::make_shared(nSlots + k, tk::KVCacheIndex(nSlots + k, /*isSecondary=*/false)); + transferManager.loadFromFile(d, /*diskSlot=*/k, {pool}, dir.str(), /*trackBlockId=*/nSlots + k); + reads.push_back(d); + } + // Queue unstaged writes (distinct disk slots): each holds raw host src pointers into the secondary pool. + std::vector writes; + for (int k = 0; k < nSlots; ++k) + { + auto s = std::make_shared(2 * nSlots + k, tk::KVCacheIndex(k, /*isSecondary=*/true)); + transferManager.spillToFileUnstaged( + s, /*diskSlot=*/nSlots + k, {pool}, dir.str(), /*spillId=*/static_cast(k + 1)); + writes.push_back(s); + } + + // Stop+join all workers while that work is outstanding. join() waits out any in-flight job against the + // still-valid pools; queued jobs are abandoned. After this, no worker can touch pool memory. + transferManager.shutdownDiskWorkers(); + transferManager.shutdownDiskWorkers(); // idempotent: second call is a no-op + + // Safe to free the pools now (the releasePools() order). No worker is mid-transfer into them. + pool.primaryPtr->release(); + pool.secondaryPtr->release(); + + // The transfer-manager destructor (at scope end) calls shutdownDiskWorkers() once more -- also a no-op. + unsetenv("TLLM_KV_DISK_READERS"); + unsetenv("TLLM_KV_DISK_WRITERS"); + unsetenv("TLLM_KV_DISK_ASYNC_STORE"); +} From ad59c7979a22ac32943c1f1a9bb898a8815f4f27 Mon Sep 17 00:00:00 2001 From: Md Nafis Ul Haque Shifat Date: Tue, 14 Jul 2026 16:33:49 -0700 Subject: [PATCH 32/47] [#16251][fix] disk tier: bound the write queue for retained spills The async spill enqueue let retained spills bypass mDiskWriteQueueMax (`job.retained || queue.size() < max`). Retention is per request and each staged spill holds a full-block byte vector, so TTL-heavy traffic could grow the pending-write backlog to disk-tier scale in RAM when the SSD fell behind and OOM the process. Gate every spill on queue room in both the staged and unstaged paths, so a full queue backpressures the producer instead of growing RAM; retained spills are still never dropped (that gate is upstream in reclaimSecondaryBlock). Removes the now-unused retained plumbing. Adds a test that bursts well past the cap and asserts the queue stays bounded. Signed-off-by: Md Nafis Ul Haque Shifat --- .../batch_manager/kvCacheTransferManager.h | 9 +- .../batch_manager/kvCacheTransferManager.cpp | 30 +++--- .../batch_manager/kvCacheManagerTest.cpp | 101 ++++++++++++++++++ 3 files changed, 117 insertions(+), 23 deletions(-) diff --git a/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h b/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h index d5d90fc60493..fa1f9e7437b5 100644 --- a/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h +++ b/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h @@ -131,8 +131,8 @@ class KVCacheTransferManager return mAsyncDiskStore; } - //! \brief True if the async write queue is at capacity. Used to shed best-effort (non-retained) spills under - //! writer saturation; retained spills bypass the cap and are never shed. + //! \brief True if the async write queue is at capacity. The eviction gate (reclaimSecondaryBlock) checks it + //! to shed best-effort spills under writer saturation; retained spills are never shed -- they backpressure. [[nodiscard]] bool diskWriteQueueFull(); void syncTransfers(); @@ -201,7 +201,6 @@ class KVCacheTransferManager std::vector staged; // staged-memcpy path (used when src == nullptr) void const* src{nullptr}; // unstaged path: writer reads this pinned host pointer directly std::uint64_t spillId{0}; // >0 => track per-spill completion for the reserved-pool reap - bool retained{false}; // retained spill: bypasses the queue cap, never dropped }; std::vector mDiskWriters; @@ -229,9 +228,9 @@ class KVCacheTransferManager }()}; void diskWriterLoop(); - void enqueueDiskWrite(std::string filename, void const* src, std::size_t bytes, bool retained); + void enqueueDiskWrite(std::string filename, void const* src, std::size_t bytes); void enqueueDiskWriteUnstaged( - std::string filename, void const* src, std::size_t bytes, std::uint64_t spillId, bool retained); + std::string filename, void const* src, std::size_t bytes, std::uint64_t spillId); void waitForDiskSlotWrites(std::string const& filename); // ---- Disk-tier async ONBOARD: read slot files disk->GPU OFF the scheduler thread ---- diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp index de49f0b2ee5b..598bd065adc8 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp @@ -521,27 +521,24 @@ void KVCacheTransferManager::diskWriterLoop() } } -void KVCacheTransferManager::enqueueDiskWrite(std::string filename, void const* src, std::size_t bytes, bool retained) +void KVCacheTransferManager::enqueueDiskWrite(std::string filename, void const* src, std::size_t bytes) { DiskWriteJob job; job.filename = std::move(filename); job.bytes = bytes; - job.retained = retained; job.staged.resize(bytes); std::memcpy(job.staged.data(), src, bytes); // fast host->host copy; frees the slot { std::unique_lock lock(mDiskMutex); - // Per-slot serialization (always) + backpressure. A retained spill must land, so it bypasses the - // queue cap (bounded by retained volume); a best-effort spill also waits for room. Best-effort shedding - // under saturation happens upstream at the eviction gate, so a best-effort spill reaching here + then - // finding the queue full is the rare TOCTOU case. + // Per-slot serialization + bounded-queue backpressure: a retained spill is never dropped (that gate is + // upstream) but still waits here for room, not bypassing the cap, so retained load can't grow RAM unbounded. mDiskQueueCv.wait(lock, [this, &job] { bool const slotClear = mDiskInflight.find(job.filename) == mDiskInflight.end(); - bool const roomOrRetained = job.retained || mDiskWriteQueue.size() < mDiskWriteQueueMax; - return (slotClear && roomOrRetained) || mDiskWriterStop; + bool const hasRoom = mDiskWriteQueue.size() < mDiskWriteQueueMax; + return (slotClear && hasRoom) || mDiskWriterStop; }); if (mDiskWriterStop) { @@ -554,25 +551,24 @@ void KVCacheTransferManager::enqueueDiskWrite(std::string filename, void const* } void KVCacheTransferManager::enqueueDiskWriteUnstaged( - std::string filename, void const* src, std::size_t bytes, std::uint64_t spillId, bool retained) + std::string filename, void const* src, std::size_t bytes, std::uint64_t spillId) { DiskWriteJob job; job.filename = std::move(filename); job.bytes = bytes; job.src = src; // no staging: the writer reads the pinned host slot directly job.spillId = spillId; - job.retained = retained; { std::unique_lock lock(mDiskMutex); - // Same per-slot serialization + backpressure as enqueueDiskWrite: retained spills bypass the queue cap - // (never stall the scheduler), best-effort spills also wait for room. + // Same per-slot serialization + bounded-queue backpressure as enqueueDiskWrite: a retained spill waits + // for room too (never dropped, never bypasses the cap), so retained load stays RAM-bounded. mDiskQueueCv.wait(lock, [this, &job] { bool const slotClear = mDiskInflight.find(job.filename) == mDiskInflight.end(); - bool const roomOrRetained = job.retained || mDiskWriteQueue.size() < mDiskWriteQueueMax; - return (slotClear && roomOrRetained) || mDiskWriterStop; + bool const hasRoom = mDiskWriteQueue.size() < mDiskWriteQueueMax; + return (slotClear && hasRoom) || mDiskWriterStop; }); if (mDiskWriterStop) { @@ -772,7 +768,6 @@ void KVCacheTransferManager::spillToFile(BlockPtr const& srcHostBlock, SizeType3 std::vector const& pools, std::string const& directory) { TLLM_CHECK_WITH_INFO(!directory.empty(), "disk tier requires a directory"); - bool const retained = srcHostBlock->isRetainedNow(); // retained spill must land -> bypasses the queue cap // The victim's bytes may still be the target of an in-flight async GPU->host copy; // wait it out before reading host memory (same event discipline as copyBlock). auto const idx = getPendingTransferIndex(srcHostBlock); @@ -798,7 +793,7 @@ void KVCacheTransferManager::spillToFile(BlockPtr const& srcHostBlock, SizeType3 continue; } // Async path: copy bytes out (frees the slot immediately) and hand to the writer. - enqueueDiskWrite(filename, ptr->data(), bytes, retained); + enqueueDiskWrite(filename, ptr->data(), bytes); } } @@ -807,7 +802,6 @@ void KVCacheTransferManager::spillToFileUnstaged(BlockPtr const& srcHostBlock, S { TLLM_CHECK_WITH_INFO(!directory.empty(), "disk tier requires a directory"); TLLM_CHECK_WITH_INFO(mAsyncDiskStore, "spillToFileUnstaged requires the async writer"); - bool const retained = srcHostBlock->isRetainedNow(); // retained spill must land -> bypasses the queue cap // Same event discipline as spillToFile: the victim's bytes may still be the target of an in-flight // GPU->host copy; block until it lands before the writer reads host memory. auto const idx = getPendingTransferIndex(srcHostBlock); @@ -829,7 +823,7 @@ void KVCacheTransferManager::spillToFileUnstaged(BlockPtr const& srcHostBlock, S auto const bytes = static_cast(ptr->getSizeInBytes()); // The pool buffer outlives the manager, so this raw pointer stays valid after `ptr` (a view) // is destroyed -- the writer reads it later off-thread. - enqueueDiskWriteUnstaged(filename, ptr->data(), bytes, spillId, retained); + enqueueDiskWriteUnstaged(filename, ptr->data(), bytes, spillId); } } diff --git a/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp b/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp index 140780d3b8b0..224ebb3a3815 100644 --- a/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp @@ -92,6 +92,34 @@ class KVCacheTransferManagerTestAccess return transferManager.mPendingWrites.find(tbk::KVCacheTransferManager::getPendingTransferIndex(block)) != transferManager.mPendingWrites.end(); } + + [[nodiscard]] static std::size_t writeQueueSize(tbk::KVCacheTransferManager& transferManager) + { + std::lock_guard lock(transferManager.mDiskMutex); + return transferManager.mDiskWriteQueue.size(); + } + + static void enqueueWrite( + tbk::KVCacheTransferManager& transferManager, std::string filename, void const* src, std::size_t bytes) + { + transferManager.enqueueDiskWrite(std::move(filename), src, bytes); + } + + static void enqueueWriteUnstaged(tbk::KVCacheTransferManager& transferManager, std::string filename, + void const* src, std::size_t bytes, std::uint64_t spillId) + { + transferManager.enqueueDiskWriteUnstaged(std::move(filename), src, bytes, spillId); + } + + // Unblock a producer parked in the enqueue backpressure wait, so the test can join it before teardown. + static void requestWriterStop(tbk::KVCacheTransferManager& transferManager) + { + { + std::lock_guard lock(transferManager.mDiskMutex); + transferManager.mDiskWriterStop = true; + } + transferManager.mDiskQueueCv.notify_all(); + } }; } // namespace tensorrt_llm::testing @@ -11112,3 +11140,76 @@ TEST_F(KVCacheManagerTest, DiskTierShutdownWithQueuedWorkTest) unsetenv("TLLM_KV_DISK_WRITERS"); unsetenv("TLLM_KV_DISK_ASYNC_STORE"); } + +// The disk write queue must stay bounded by mDiskWriteQueueMax under sustained spill load: with the retained +// bypass removed, a producer facing a full queue backpressures instead of growing RAM. Covers staged + unstaged. +namespace +{ +void runWriteQueueBoundedTest(bool unstaged) +{ + using tensorrt_llm::testing::KVCacheTransferManagerTestAccess; + setenv("TLLM_KV_DISK_ASYNC_STORE", "1", /*overwrite=*/1); + setenv("TLLM_KV_DISK_WRITERS", "1", /*overwrite=*/1); // one writer drains; the cap must still bound the backlog + setenv("TLLM_KV_DISK_WRITE_QUEUE", "4", /*overwrite=*/1); + std::size_t constexpr cap = 4; + std::size_t constexpr burst = 200; // far more than the cap: the producer must block once the queue is full + + DiskTierDir dir; // real directory for the writer to drain into (cleaned on exit) + std::vector src(256 * 1024, 0xAB); // large enough that the single writer lags the producer + auto bufferManager = tr::BufferManager(std::make_shared()); + auto transferManager = KVCacheTransferManager(bufferManager); + + // Continuously sample the queue depth while the burst runs and record the peak. + std::atomic observedMax{0}; + std::atomic producing{true}; + std::thread monitor( + [&] + { + while (producing.load()) + { + auto const sz = KVCacheTransferManagerTestAccess::writeQueueSize(transferManager); + auto prev = observedMax.load(); + while (sz > prev && !observedMax.compare_exchange_weak(prev, sz)) + { + } + std::this_thread::sleep_for(std::chrono::microseconds(200)); + } + }); + + // Burst far more spills than the cap. Each enqueue blocks while the queue is full, so the backlog can never + // exceed mDiskWriteQueueMax no matter how far the single writer falls behind -- the bound the reviewer wants. + for (std::size_t k = 0; k < burst; ++k) + { + auto const fn = dir.str() + "/slot" + std::to_string(k); + if (unstaged) + { + KVCacheTransferManagerTestAccess::enqueueWriteUnstaged(transferManager, fn, src.data(), src.size(), k + 1); + } + else + { + KVCacheTransferManagerTestAccess::enqueueWrite(transferManager, fn, src.data(), src.size()); + } + } + producing.store(false); + monitor.join(); + + EXPECT_LE(observedMax.load(), cap) << "write queue exceeded its cap -- spills not backpressured"; + EXPECT_GT(observedMax.load(), 0u) << "queue never filled -- the burst did not actually stress the cap"; + + unsetenv("TLLM_KV_DISK_WRITE_QUEUE"); + unsetenv("TLLM_KV_DISK_WRITERS"); + unsetenv("TLLM_KV_DISK_ASYNC_STORE"); +} +} // namespace + +// Staged path: a full write queue backpressures the producer instead of growing RAM without bound. +TEST_F(KVCacheManagerTest, DiskTierWriteQueueBoundedStagedTest) +{ + runWriteQueueBoundedTest(/*unstaged=*/false); +} + +// Unstaged path: same bound. +TEST_F(KVCacheManagerTest, DiskTierWriteQueueBoundedUnstagedTest) +{ + runWriteQueueBoundedTest(/*unstaged=*/true); +} From 3727a82861362fb361aa69c6db877fe8104656d7 Mon Sep 17 00:00:00 2001 From: Md Nafis Ul Haque Shifat Date: Tue, 14 Jul 2026 17:09:49 -0700 Subject: [PATCH 33/47] [#16251][fix] disk tier: anchor retention deadline at commit, not allocation The disk-retention deadline was stamped when a block was allocated (getFreeBlock), before its KV was computed or committed to the reuse tree. The contract starts the duration when a block is first made available for reuse, so a long prefill/generation could consume part or all of the requested TTL before the block was even reusable. Hold the requested duration on the block while it is active (setPendingRetention) and convert it to an absolute deadline only when the block is committed to the reuse tree (commitRetention, from storeBlocks). Adds a test asserting no deadline at allocation and ~now+TTL after commit. Signed-off-by: Md Nafis Ul Haque Shifat --- .../batch_manager/kvCacheManager.h | 21 ++++++++++ .../batch_manager/kvCacheManager.cpp | 9 +++- .../batch_manager/kvCacheManagerTest.cpp | 42 +++++++++++++++++++ 3 files changed, 71 insertions(+), 1 deletion(-) diff --git a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h index 9a63715682cb..9e920b7331a9 100644 --- a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h +++ b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h @@ -487,6 +487,24 @@ class KVCacheBlock : public std::enable_shared_from_this mRetentionExpiry = std::nullopt; } + //! \brief Record the requested disk-retention duration without starting the clock; commitRetention() + //! turns it into an absolute deadline when the block is first committed to the reuse tree. + void setPendingRetention(std::chrono::milliseconds duration) + { + mPendingRetentionDuration + = mPendingRetentionDuration ? std::max(*mPendingRetentionDuration, duration) : duration; + } + + //! \brief Anchor a pending retention duration to now (called when the block is committed / made reusable). + void commitRetention() + { + if (mPendingRetentionDuration) + { + markRetained(std::chrono::steady_clock::now().time_since_epoch() + *mPendingRetentionDuration); + mPendingRetentionDuration = std::nullopt; + } + } + //! \brief Times this block's content has been reused (prefix-matched); drives the disk reuse-gate. [[nodiscard]] SizeType32 getReuseCount() const { @@ -659,6 +677,9 @@ class KVCacheBlock : public std::enable_shared_from_this // Disk-tier retention deadline (steady clock); nullopt = never marked. std::optional mRetentionExpiry; + // Requested retention duration, held until commit; commitRetention() converts it to mRetentionExpiry. + std::optional mPendingRetentionDuration; + // How many times this block's content has been reused (prefix-matched). Metadata on the block object, // so it survives tier moves like mRetentionExpiry (swapDiskResidency swaps only residency pointers). SizeType32 mReuseCount{0}; diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp index 2875bd49a918..6d0a023f9878 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp @@ -1622,7 +1622,8 @@ BlockPtr WindowBlockManager::getFreeBlock(GenerationRequest& sequence, executor: mEvictionPolicy->claimBlock(block, priority, durationMs); if (auto const diskRetentionMs = sequence.getDiskRetentionMs()) { - block->markRetained(std::chrono::steady_clock::now().time_since_epoch() + *diskRetentionMs); + // Anchor the deadline at commit (storeBlocks), not at allocation. + block->setPendingRetention(*diskRetentionMs); } TLLM_LOG_DEBUG("%s::getFreeBlock - Block %d is now acquired by sequence %d", mLogPrefix.c_str(), block->getBlockId(), sequence.getRequestId()); @@ -3085,6 +3086,12 @@ std::pair> WindowBlockManager::sto } } + // Blocks are now committed to the reuse tree: anchor any pending retention duration to now. + for (auto const& b : storedBlocks) + { + b->commitRetention(); + } + if (mEventManager) { // Linear-attention placeholders can legitimately land in storedBlocks via the diff --git a/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp b/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp index 224ebb3a3815..14b7a6c0eeda 100644 --- a/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp @@ -10852,6 +10852,48 @@ TEST_F(KVCacheManagerTest, DiskTierExpiryAtSpillTest) EXPECT_GE(blockManager.getNumDiskGateDropped(), 2u); } +// The disk-retention deadline must be anchored when the block is committed to the reuse tree +// (storeContextBlocks), not at allocation -- otherwise a long prefill eats into the requested TTL. +TEST_F(KVCacheManagerTest, DiskTierRetentionAnchoredAtCommitTest) +{ + auto const stream = std::make_shared(); + DiskTierDir dir; + auto blockManagerPtr + = makeDiskTierBlockManager(stream, dir.str(), /*retainedOnly=*/false, /*prim=*/4, /*sec=*/2, /*disk=*/8); + auto& blockManager = *blockManagerPtr; + auto const windowSize = blockManager.getPoolWindowSize(0); + auto constexpr ttl = std::chrono::milliseconds(60000); + + auto s = addDiskTierSequence(blockManager, 0, iotaTokens(0, 16), ttl); + auto const blockIds = s.seq->getCacheBlockIds(windowSize).at(0); // beam 0 + ASSERT_FALSE(blockIds.empty()); + + // Allocated but not yet committed: no block may carry an absolute deadline. + for (auto const bid : blockIds) + { + EXPECT_FALSE(blockManager.getBlockById(bid, windowSize)->getRetentionExpiry().has_value()) + << "retention deadline stamped at allocation instead of deferred to commit"; + } + + // Commit to the reuse tree -> the deadline is anchored at ~now + ttl for the committed blocks. + auto const before = std::chrono::steady_clock::now().time_since_epoch(); + blockManager.storeContextBlocks(*s.seq, *s.req); + auto const after = std::chrono::steady_clock::now().time_since_epoch(); + + std::size_t stamped = 0; + for (auto const bid : blockIds) + { + auto const expiry = blockManager.getBlockById(bid, windowSize)->getRetentionExpiry(); + if (expiry.has_value()) + { + ++stamped; + EXPECT_GE(*expiry, before + ttl); // clock started at commit, not earlier at allocation + EXPECT_LE(*expiry, after + ttl); + } + } + EXPECT_GT(stamped, 0u) << "commit did not anchor any retention deadline"; +} + // Bug-1 regression: a re-send that REUSES blocks re-stamps them (max-merge extends). TEST_F(KVCacheManagerTest, DiskTierReuseRestampTest) { From 7177595e4f8fffa819d56b393c9c691390cc8336 Mon Sep 17 00:00:00 2001 From: Md Nafis Ul Haque Shifat Date: Tue, 14 Jul 2026 18:16:32 -0700 Subject: [PATCH 34/47] [#16251][fix] disk tier: contain async worker I/O failures diskWriterLoop and the POSIX branch of diskReaderLoop let TLLM_CHECK throws escape the std::thread entry points, so a disk I/O fault (unwritable or missing/short slot file) took down the whole process via std::terminate. Catch failures per job on the worker threads. A failed spill logs and still runs its completion bookkeeping, so slot waiters and the reserved-pool reap are notified and the writer keeps draining. A failed onboard leaves the block pending -- nothing forwards onto never-filled KV -- and sets a flag that areBlocksReady() re-raises on the scheduler thread, the same path a synchronous disk-read failure already takes into the executor's error handling. Adds tests for both: a failed spill still publishes its spill id and the writer keeps serving jobs; a failed onboard surfaces as a throw from areBlocksReady instead of crashing the reader. Signed-off-by: Md Nafis Ul Haque Shifat --- .../batch_manager/kvCacheTransferManager.h | 10 ++ .../batch_manager/kvCacheManager.cpp | 8 ++ .../batch_manager/kvCacheTransferManager.cpp | 63 ++++++++----- .../batch_manager/kvCacheManagerTest.cpp | 94 +++++++++++++++++++ 4 files changed, 153 insertions(+), 22 deletions(-) diff --git a/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h b/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h index fa1f9e7437b5..ad03b040c8e6 100644 --- a/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h +++ b/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h @@ -105,6 +105,14 @@ class KVCacheTransferManager return mReadInflightCount.load(std::memory_order_acquire) > 0; } + //! \brief Lock-free: true once any async onboard read has failed. The failed block stays in the pending + //! set (nothing forwards onto it); areBlocksReady() re-raises the failure on the scheduler thread, the + //! same path a synchronous disk-read failure takes. + [[nodiscard]] bool anyReadFailed() const noexcept + { + return mAnyReadFailed.load(std::memory_order_acquire); + } + //! \brief True when a reader pool is present, so onboards can be detached from the scheduler thread. [[nodiscard]] bool asyncDiskReadEnabled() const { @@ -261,6 +269,8 @@ class KVCacheTransferManager // Lock-free mirror of mPendingBlockReads.size() (updated under mReadMutex on every insert/erase) so // areBlocksReady() can skip its per-block scan + the mutex when nothing is in flight (common case). std::atomic mReadInflightCount{0}; + // Set by a reader thread when an async onboard read fails; polled by areBlocksReady(). + std::atomic mAnyReadFailed{false}; bool mDiskReaderStop{false}; // Set once shutdownDiskWorkers() has stopped+joined the disk workers, so releasePools() and the // destructor do not repeat the work. Only touched on the owning thread. diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp index 6d0a023f9878..79a50918e792 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp @@ -1752,6 +1752,14 @@ void WindowBlockManager::onboardBlock(GenerationRequest& sequence, BlockPtr cons bool WindowBlockManager::areBlocksReady(LlmRequest::RequestIdType requestId) { + // Surface a failed async onboard here, on the scheduler thread -- the same path a synchronous disk-read + // failure takes. The failed block stays in the pending set, so nothing has forwarded onto its contents; + // the reader thread already logged the details. + if (mTransferManager->anyReadFailed()) + { + TLLM_THROW("disk tier: async onboard read failed (see [disk-tier] error log); failing rather than " + "serving never-filled KV"); + } // Fast path: with no disk read in flight (the common case) every block is trivially ready, so skip the // per-block scan and its mutex entirely. Without this the park scans every block of every context // request each step -- an O(context-length) mutex-locked host cost even when nothing is onboarding. diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp index 598bd065adc8..ee3bacdff905 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp @@ -489,16 +489,25 @@ void KVCacheTransferManager::diskWriterLoop() // The slow, writeback-throttle-prone part runs HERE, off the scheduler thread. // Unstaged jobs (src != nullptr) read the pinned host slot directly; staged jobs read the copy. void const* data = job.src ? job.src : static_cast(job.staged.data()); - int fd = ::open(job.filename.c_str(), O_CREAT | O_WRONLY, 0644); - TLLM_CHECK_WITH_INFO(fd >= 0, - "[disk-tier] cannot open %s for async write; failing loudly rather than leaving a corrupt slot", - job.filename.c_str()); - auto const written = ::pwrite(fd, data, job.bytes, 0); - ::close(fd); - TLLM_CHECK_WITH_INFO(written == static_cast(job.bytes), - "[disk-tier] short async write to %s (%zd/%zu); failing loudly rather than leaving a corrupt slot", - job.filename.c_str(), static_cast(written), job.bytes); + try + { + int fd = ::open(job.filename.c_str(), O_CREAT | O_WRONLY, 0644); + TLLM_CHECK_WITH_INFO(fd >= 0, "[disk-tier] cannot open %s for async write", job.filename.c_str()); + auto const written = ::pwrite(fd, data, job.bytes, 0); + ::close(fd); + TLLM_CHECK_WITH_INFO(written == static_cast(job.bytes), + "[disk-tier] short async write to %s (%zd/%zu)", job.filename.c_str(), + static_cast(written), job.bytes); + } + catch (std::exception const& e) + { + // Contain the I/O error here (else it escapes the thread -> std::terminate). The slot is left + // unwritten; a later onboard of it fails and aborts that request (see diskReaderLoop). + TLLM_LOG_ERROR("[disk-tier] async write to %s failed: %s", job.filename.c_str(), e.what()); + } + // Runs whether the write succeeded or failed, so producers waiting on this slot and the reserved-pool + // reap are always notified. { std::lock_guard lock(mDiskMutex); if (auto it = mDiskInflight.find(job.filename); it != mDiskInflight.end() && --it->second <= 0) @@ -716,21 +725,31 @@ void KVCacheTransferManager::diskReaderLoop() buf = pageableFallback.data(); } std::size_t off = 0; - for (size_t i = 0; i < job.dsts.size(); ++i) + try { - int fd = ::open(job.files[i].c_str(), O_RDONLY); - TLLM_CHECK_WITH_INFO(fd >= 0, - "[disk-tier] cannot open %s for async read; failing loudly rather than serving corrupt KV", - job.files[i].c_str()); - auto const got = ::read(fd, buf + off, job.bytes[i]); - ::close(fd); - TLLM_CHECK_WITH_INFO(got == static_cast(job.bytes[i]), - "[disk-tier] short async read from %s (%zd/%zu); failing loudly rather than serving corrupt KV", - job.files[i].c_str(), static_cast(got), job.bytes[i]); - TLLM_CUDA_CHECK(cudaMemcpyAsync(job.dsts[i], buf + off, job.bytes[i], cudaMemcpyHostToDevice, stream)); - off += job.bytes[i]; + for (size_t i = 0; i < job.dsts.size(); ++i) + { + int fd = ::open(job.files[i].c_str(), O_RDONLY); + TLLM_CHECK_WITH_INFO(fd >= 0, "[disk-tier] cannot open %s for async read", job.files[i].c_str()); + auto const got = ::read(fd, buf + off, job.bytes[i]); + ::close(fd); + TLLM_CHECK_WITH_INFO(got == static_cast(job.bytes[i]), + "[disk-tier] short async read from %s (%zd/%zu)", job.files[i].c_str(), + static_cast(got), job.bytes[i]); + TLLM_CUDA_CHECK( + cudaMemcpyAsync(job.dsts[i], buf + off, job.bytes[i], cudaMemcpyHostToDevice, stream)); + off += job.bytes[i]; + } + TLLM_CUDA_CHECK(cudaStreamSynchronize(stream)); + } + catch (std::exception const& e) + { + // Contain it (a throw escaping this thread -> std::terminate). The block stays pending so no + // request forwards onto never-filled KV; areBlocksReady() re-raises this on the scheduler thread. + TLLM_LOG_ERROR("[disk-tier] async onboard of block %d failed: %s", job.blockId, e.what()); + mAnyReadFailed.store(true, std::memory_order_release); + continue; } - TLLM_CUDA_CHECK(cudaStreamSynchronize(stream)); } { diff --git a/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp b/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp index 14b7a6c0eeda..980bc887abca 100644 --- a/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp @@ -11114,6 +11114,100 @@ TEST_F(KVCacheManagerTest, DiskTierSyncOnboardByteExactTest) runDiskOnboardByteRoundTrip(/*readers=*/"0", /*nSlots=*/4); } +// A failed async spill (unwritable slot file) must be contained on the writer thread: per-slot bookkeeping +// still runs (the spill id is published for the reserved-pool reap) and the worker keeps serving later jobs. +// Before the fix the TLLM_CHECK throw escaped the thread entry point and std::terminate'd the process. +TEST_F(KVCacheManagerTest, DiskTierAsyncWriteFailureContainedTest) +{ + setenv("TLLM_KV_DISK_ASYNC_STORE", "1", /*overwrite=*/1); + setenv("TLLM_KV_DISK_WRITERS", "1", /*overwrite=*/1); + DiskTierDir dir; + int constexpr blockSize = 256; + auto bufferManager = tr::BufferManager(std::make_shared()); + auto transferManager = KVCacheTransferManager(bufferManager); + auto pool = KVCacheBlockPool(0, /*kvFactor=*/2, 0, 0, 0); + pool.primaryPtr = bufferManager.gpu(tr::ITensor::makeShape({2, blockSize}), nvinfer1::DataType::kFLOAT); + pool.secondaryPtr = tr::BufferManager::pinned(tr::ITensor::makeShape({2, blockSize}), nvinfer1::DataType::kFLOAT); + + auto waitForSpill = [&](std::uint64_t spillId) + { + auto const deadline = std::chrono::steady_clock::now() + std::chrono::seconds(30); + while (std::chrono::steady_clock::now() < deadline) + { + auto const done = transferManager.drainCompletedSpills(); + if (std::find(done.begin(), done.end(), spillId) != done.end()) + { + return true; + } + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + return false; + }; + + // Unwritable destination (missing parent dir): the open fails on the writer thread. + auto src = std::make_shared(0, tk::KVCacheIndex(0, /*isSecondary=*/true)); + transferManager.spillToFileUnstaged(src, /*diskSlot=*/0, {pool}, dir.str() + "/missing_subdir", /*spillId=*/1); + EXPECT_TRUE(waitForSpill(1)) << "failed spill did not publish its id (bookkeeping skipped)"; + + // Worker survived: a valid spill afterwards completes normally. + auto src2 = std::make_shared(1, tk::KVCacheIndex(1, /*isSecondary=*/true)); + transferManager.spillToFileUnstaged(src2, /*diskSlot=*/1, {pool}, dir.str(), /*spillId=*/2); + EXPECT_TRUE(waitForSpill(2)) << "writer stopped processing after a failed job"; + + unsetenv("TLLM_KV_DISK_WRITERS"); + unsetenv("TLLM_KV_DISK_ASYNC_STORE"); +} + +// A failed detached onboard read (missing slot file) must not crash the reader thread; the failure surfaces +// as a throw from areBlocksReady() on the calling (scheduler) thread, and no request is unparked onto the +// never-filled block. Before the fix the reader's TLLM_CHECK throw std::terminate'd the process. +TEST_F(KVCacheManagerTest, DiskTierOnboardReadFailureSurfacesTest) +{ + setenv("TLLM_KV_DISK_READERS", "2", /*overwrite=*/1); + auto const stream = std::make_shared(); + DiskTierDir dir; + auto blockManagerPtr + = makeDiskTierBlockManager(stream, dir.str(), /*retainedOnly=*/false, /*prim=*/4, /*sec=*/2, /*disk=*/8); + auto& blockManager = *blockManagerPtr; + + auto seqA = addDiskTierSequence(blockManager, 0, iotaTokens(0, 16), std::nullopt); + releaseDiskTierSequence(blockManager, seqA); + churnOnce(blockManager, 1, 100); // A's deepest blocks spill to disk + ASSERT_GE(blockManager.getNumDiskSpills(), 2u); + + // Break the tier: delete the slot files so the detached onboard reads below fail on the reader threads. + for (auto const& entry : fs::directory_iterator(dir.str())) + { + fs::remove_all(entry.path()); + } + + // Rematch A's prefix: the disk onboards are handed to the reader pool and fail off-thread. + auto seqA2 = addDiskTierSequence(blockManager, 2, iotaTokens(0, 16), std::nullopt); + + bool threw = false; + auto const deadline = std::chrono::steady_clock::now() + std::chrono::seconds(30); + while (!threw && std::chrono::steady_clock::now() < deadline) + { + try + { + (void) blockManager.areBlocksReady(/*requestId=*/2); // false while pending; throws once failed + } + catch (std::exception const&) + { + threw = true; + } + if (!threw) + { + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + } + EXPECT_TRUE(threw) << "failed onboard was not surfaced on the scheduler thread"; + + // Teardown stays clean: the sequence releases without touching the never-filled blocks. + releaseDiskTierSequence(blockManager, seqA2); + unsetenv("TLLM_KV_DISK_READERS"); +} + // shutdownDiskWorkers() must stop+join every disk worker while reads are queued (holding raw GPU dst // pointers) and unstaged writes are queued (holding raw host src pointers), so releasePools() can free // the pools with no use-after-free -- and it must be idempotent. Passing == no hang/crash reaching the From ad2777552d1996290da3d0ab12fa5107fa0d54ec Mon Sep 17 00:00:00 2001 From: Md Nafis Ul Haque Shifat Date: Tue, 14 Jul 2026 18:32:36 -0700 Subject: [PATCH 35/47] [#16251][infra] apply clang-format to the review-fix files Signed-off-by: Md Nafis Ul Haque Shifat --- .../tensorrt_llm/batch_manager/kvCacheTransferManager.h | 3 +-- cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp | 5 +++-- cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp | 4 ++-- cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp | 2 +- 4 files changed, 7 insertions(+), 7 deletions(-) diff --git a/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h b/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h index ad03b040c8e6..7ef174a86013 100644 --- a/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h +++ b/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h @@ -237,8 +237,7 @@ class KVCacheTransferManager void diskWriterLoop(); void enqueueDiskWrite(std::string filename, void const* src, std::size_t bytes); - void enqueueDiskWriteUnstaged( - std::string filename, void const* src, std::size_t bytes, std::uint64_t spillId); + void enqueueDiskWriteUnstaged(std::string filename, void const* src, std::size_t bytes, std::uint64_t spillId); void waitForDiskSlotWrites(std::string const& filename); // ---- Disk-tier async ONBOARD: read slot files disk->GPU OFF the scheduler thread ---- diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp index 79a50918e792..bd4e46f97119 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp @@ -1757,8 +1757,9 @@ bool WindowBlockManager::areBlocksReady(LlmRequest::RequestIdType requestId) // the reader thread already logged the details. if (mTransferManager->anyReadFailed()) { - TLLM_THROW("disk tier: async onboard read failed (see [disk-tier] error log); failing rather than " - "serving never-filled KV"); + TLLM_THROW( + "disk tier: async onboard read failed (see [disk-tier] error log); failing rather than " + "serving never-filled KV"); } // Fast path: with no disk read in flight (the common case) every block is trivially ready, so skip the // per-block scan and its mutex entirely. Without this the park scans every block of every context diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp index ee3bacdff905..2fdbe1307553 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp @@ -496,8 +496,8 @@ void KVCacheTransferManager::diskWriterLoop() auto const written = ::pwrite(fd, data, job.bytes, 0); ::close(fd); TLLM_CHECK_WITH_INFO(written == static_cast(job.bytes), - "[disk-tier] short async write to %s (%zd/%zu)", job.filename.c_str(), - static_cast(written), job.bytes); + "[disk-tier] short async write to %s (%zd/%zu)", job.filename.c_str(), static_cast(written), + job.bytes); } catch (std::exception const& e) { diff --git a/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp b/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp index 980bc887abca..98b96a358f3d 100644 --- a/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp @@ -11290,7 +11290,7 @@ void runWriteQueueBoundedTest(bool unstaged) std::size_t constexpr cap = 4; std::size_t constexpr burst = 200; // far more than the cap: the producer must block once the queue is full - DiskTierDir dir; // real directory for the writer to drain into (cleaned on exit) + DiskTierDir dir; // real directory for the writer to drain into (cleaned on exit) std::vector src(256 * 1024, 0xAB); // large enough that the single writer lags the producer auto bufferManager = tr::BufferManager(std::make_shared()); auto transferManager = KVCacheTransferManager(bufferManager); From ffcb625c14387b956b80b7c1e32d6db49632375f Mon Sep 17 00:00:00 2001 From: Md Nafis Ul Haque Shifat Date: Tue, 14 Jul 2026 18:47:32 -0700 Subject: [PATCH 36/47] [#16251][doc] fix stale disk-cache docstring and comment placement - llm_args: disk_cache_size is used by the v1 disk tier too, not only KV cache manager v2 - kvCacheTransferManager.h: reattach the syncTransfers() doc comment that the inserted disk-tier declarations had orphaned - openai_server: note explicitly that /v1/responses does not expose kv_cache_ttl_seconds Signed-off-by: Md Nafis Ul Haque Shifat --- .../tensorrt_llm/batch_manager/kvCacheTransferManager.h | 6 +++--- tensorrt_llm/llmapi/llm_args.py | 2 +- tensorrt_llm/serve/openai_server.py | 2 ++ 3 files changed, 6 insertions(+), 4 deletions(-) diff --git a/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h b/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h index 7ef174a86013..c94ecf2a23dd 100644 --- a/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h +++ b/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h @@ -82,9 +82,6 @@ class KVCacheTransferManager //! KVCacheManager::addSequenceBatch in every step. void syncWithBufferManager(); - //! \brief Synchronize bufferManager stream with internal streams. This method ensures that prefill and decode - //! kernels for next step will wait for offloading and onboarding work that has already been scheduled. This method - //! must be called after the last call to KVCacheManager::addSequenceBatch in every step. //! \brief Disk tier: write a host-resident block's bytes to its slot files (synchronous POSIX). void spillToFile(BlockPtr const& srcHostBlock, SizeType32 diskSlot, std::vector const& pools, std::string const& directory); @@ -143,6 +140,9 @@ class KVCacheTransferManager //! to shed best-effort spills under writer saturation; retained spills are never shed -- they backpressure. [[nodiscard]] bool diskWriteQueueFull(); + //! \brief Synchronize bufferManager stream with internal streams. This method ensures that prefill and decode + //! kernels for next step will wait for offloading and onboarding work that has already been scheduled. This method + //! must be called after the last call to KVCacheManager::addSequenceBatch in every step. void syncTransfers(); ~KVCacheTransferManager(); diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index c53c779aa2f0..9579c7747ceb 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -3343,7 +3343,7 @@ class KvCacheConfig(StrictBaseModel, PybindMirror): disk_cache_size: Optional[NonNegativeInt] = Field( default=None, description= - "Size of the disk cache in bytes. Only used by KV cache manager v2 in the PyTorch backend." + "Size of the disk cache in bytes. Enables the disk KV cache tier in the PyTorch backend, for both the default KV cache manager and KV cache manager v2." ) disk_cache_path: Optional[str] = Field( default=None, diff --git a/tensorrt_llm/serve/openai_server.py b/tensorrt_llm/serve/openai_server.py index 39f089033d1d..89f4de59a276 100644 --- a/tensorrt_llm/serve/openai_server.py +++ b/tensorrt_llm/serve/openai_server.py @@ -2263,6 +2263,8 @@ async def create_streaming_generator(promise: RequestOutput, ) promise = self.generator.generate_async( inputs=input_tokens, + # ResponsesRequest does not expose kv_cache_ttl_seconds yet, so + # this resolves to None (no disk retention on /v1/responses). kv_cache_retention_config=_disk_retention_config(request), sampling_params=sampling_params, streaming=request.stream, From 5f0afb7127ae38603dfa04cf281a70c69d91f198 Mon Sep 17 00:00:00 2001 From: Md Nafis Ul Haque Shifat Date: Tue, 14 Jul 2026 19:49:41 -0700 Subject: [PATCH 37/47] [#16251][fix] validate kv_cache_ttl_seconds and reject CP with the disk tier - openai_protocol: add ge=0 to kv_cache_ttl_seconds on both request classes; a negative TTL now fails validation instead of behaving like already-expired retention - resource_manager: reject context parallelism (cp_size > 1) in the disk-tier guard; the onboard-readiness sync covers only plain TP - test_disk_tier_args: cover disk_cache_protect_unexpired (default/set/pickle/pydantic->pybind) and negative-TTL rejection Signed-off-by: Md Nafis Ul Haque Shifat --- .../_torch/pyexecutor/resource_manager.py | 10 +++--- tensorrt_llm/serve/openai_protocol.py | 10 +++--- tests/unittest/llmapi/test_disk_tier_args.py | 35 ++++++++++++++++++- 3 files changed, 46 insertions(+), 9 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index 4f7fe3bdf8da..2d1945d7c9a1 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -616,11 +616,13 @@ def append_to_kv_heads_per_layer(num_kv_heads_per_layer: List[int], disk_cache_size = getattr(kv_cache_config, "disk_cache_size", None) or 0 if disk_cache_size and (mapping.enable_attention_dp - or mapping.pp_size > 1): + or mapping.pp_size > 1 or mapping.cp_size > 1): raise ValueError( - "disk KV-cache tier is not supported with attention-DP or pipeline parallelism " - "(cross-rank onboard-readiness sync assumes identical per-rank batches); got " - f"enable_attention_dp={mapping.enable_attention_dp}, pp_size={mapping.pp_size}. " + "disk KV-cache tier is not supported with attention-DP, pipeline parallelism, " + "or context parallelism (cross-rank onboard-readiness sync assumes identical " + "per-rank batches and covers only the plain-TP group); got " + f"enable_attention_dp={mapping.enable_attention_dp}, pp_size={mapping.pp_size}, " + f"cp_size={mapping.cp_size}. " "Disable disk_cache_size or use plain tensor parallelism.") blocks_in_disk_pool = 0 if disk_cache_size: diff --git a/tensorrt_llm/serve/openai_protocol.py b/tensorrt_llm/serve/openai_protocol.py index 662598dc119e..604455c0ed9a 100644 --- a/tensorrt_llm/serve/openai_protocol.py +++ b/tensorrt_llm/serve/openai_protocol.py @@ -521,9 +521,10 @@ class CompletionRequest(OpenAIBaseModel): top_k: int = 0 kv_cache_ttl_seconds: Optional[int] = Field( default=None, + ge=0, description= - "Keep this request's KV cache reusable for this many seconds (disk-tier retention)." - ) + "Keep this request's KV cache reusable for this many seconds (disk-tier " + "retention). 0 or unset means no disk retention.") top_p_min: float = 0.0 min_p: float = 0.0 repetition_penalty: float = 1.0 @@ -872,9 +873,10 @@ class ChatCompletionRequest(OpenAIBaseModel): top_k: int = 0 kv_cache_ttl_seconds: Optional[int] = Field( default=None, + ge=0, description= - "Keep this request's KV cache reusable for this many seconds (disk-tier retention)." - ) + "Keep this request's KV cache reusable for this many seconds (disk-tier " + "retention). 0 or unset means no disk retention.") top_p_min: float = 0.0 min_p: float = 0.0 repetition_penalty: float = 1.0 diff --git a/tests/unittest/llmapi/test_disk_tier_args.py b/tests/unittest/llmapi/test_disk_tier_args.py index 46a5a5375a52..67e012a4f785 100644 --- a/tests/unittest/llmapi/test_disk_tier_args.py +++ b/tests/unittest/llmapi/test_disk_tier_args.py @@ -7,18 +7,23 @@ import datetime import pickle +import pytest + from tensorrt_llm.bindings import executor as tllm def test_kv_cache_config_disk_fields(): cfg = tllm.KvCacheConfig() assert cfg.disk_cache_retained_only is False # default preserves spill-all + assert cfg.disk_cache_protect_unexpired is False # default: LRU may displace unexpired blocks cfg.disk_cache_size = 1 << 30 cfg.disk_cache_path = "/tmp/kvdisk-test" cfg.disk_cache_retained_only = True + cfg.disk_cache_protect_unexpired = True assert cfg.disk_cache_size == 1 << 30 assert cfg.disk_cache_path == "/tmp/kvdisk-test" assert cfg.disk_cache_retained_only is True + assert cfg.disk_cache_protect_unexpired is True def test_kv_cache_config_disk_fields_pickle(): @@ -26,10 +31,12 @@ def test_kv_cache_config_disk_fields_pickle(): cfg.disk_cache_size = 42 cfg.disk_cache_path = "/x" cfg.disk_cache_retained_only = True + cfg.disk_cache_protect_unexpired = True back = pickle.loads(pickle.dumps(cfg)) assert back.disk_cache_size == 42 assert back.disk_cache_path == "/x" assert back.disk_cache_retained_only is True + assert back.disk_cache_protect_unexpired is True def test_retention_config_disk_retention_ms_roundtrip(): @@ -48,13 +55,18 @@ def test_llm_args_kv_cache_config_disk_fields(): from tensorrt_llm.llmapi.llm_args import KvCacheConfig as PydanticKvCacheConfig py_cfg = PydanticKvCacheConfig( - disk_cache_size=1 << 30, disk_cache_path="/tmp/kvdisk-test", disk_cache_retained_only=True + disk_cache_size=1 << 30, + disk_cache_path="/tmp/kvdisk-test", + disk_cache_retained_only=True, + disk_cache_protect_unexpired=True, ) assert py_cfg.disk_cache_retained_only is True + assert py_cfg.disk_cache_protect_unexpired is True cpp_cfg = py_cfg._to_pybind() assert cpp_cfg.disk_cache_size == 1 << 30 assert cpp_cfg.disk_cache_path == "/tmp/kvdisk-test" assert cpp_cfg.disk_cache_retained_only is True + assert cpp_cfg.disk_cache_protect_unexpired is True def test_serve_disk_retention_helper(): @@ -77,3 +89,24 @@ class Bare: # object without the field at all (non-serve request types) pass assert _disk_retention_config(Bare()) is None + + +def test_kv_cache_ttl_seconds_rejects_negative(): + from pydantic import ValidationError + + from tensorrt_llm.serve.openai_protocol import ChatCompletionRequest, CompletionRequest + + with pytest.raises(ValidationError): + CompletionRequest(model="m", prompt="p", kv_cache_ttl_seconds=-1) + with pytest.raises(ValidationError): + ChatCompletionRequest( + model="m", messages=[{"role": "user", "content": "hi"}], kv_cache_ttl_seconds=-1 + ) + # 0 and positive values stay accepted; 0 means "no disk retention" (see _disk_retention_config). + assert ( + CompletionRequest(model="m", prompt="p", kv_cache_ttl_seconds=0).kv_cache_ttl_seconds == 0 + ) + assert ( + CompletionRequest(model="m", prompt="p", kv_cache_ttl_seconds=3600).kv_cache_ttl_seconds + == 3600 + ) From c054a3cf30b3a697e917a3f28e519167f92c783b Mon Sep 17 00:00:00 2001 From: Md Nafis Ul Haque Shifat Date: Tue, 14 Jul 2026 19:59:03 -0700 Subject: [PATCH 38/47] [#16251][infra] dedup spill enqueue paths and scope test env vars - kvCacheTransferManager: extract enqueueDiskWriteCommon() for the lock/wait/push tail shared by the staged and unstaged spill enqueues, so the backpressure policy cannot drift between the two - kvCacheManagerTest: replace raw setenv/unsetenv pairs with a ScopedEnvVar RAII guard, so a gtest ASSERT early-exit cannot leak disk-tier env vars into later tests Signed-off-by: Md Nafis Ul Haque Shifat --- .../batch_manager/kvCacheTransferManager.h | 2 + .../batch_manager/kvCacheTransferManager.cpp | 29 +++------- .../batch_manager/kvCacheManagerTest.cpp | 57 +++++++++++-------- 3 files changed, 42 insertions(+), 46 deletions(-) diff --git a/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h b/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h index c94ecf2a23dd..c3ceec3d77cf 100644 --- a/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h +++ b/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h @@ -238,6 +238,8 @@ class KVCacheTransferManager void diskWriterLoop(); void enqueueDiskWrite(std::string filename, void const* src, std::size_t bytes); void enqueueDiskWriteUnstaged(std::string filename, void const* src, std::size_t bytes, std::uint64_t spillId); + // Shared lock/wait/push tail of the two enqueue paths (per-slot serialization + bounded-queue backpressure). + void enqueueDiskWriteCommon(DiskWriteJob&& job); void waitForDiskSlotWrites(std::string const& filename); // ---- Disk-tier async ONBOARD: read slot files disk->GPU OFF the scheduler thread ---- diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp index 2fdbe1307553..fa339725108b 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp @@ -537,26 +537,7 @@ void KVCacheTransferManager::enqueueDiskWrite(std::string filename, void const* job.bytes = bytes; job.staged.resize(bytes); std::memcpy(job.staged.data(), src, bytes); // fast host->host copy; frees the slot - - { - std::unique_lock lock(mDiskMutex); - // Per-slot serialization + bounded-queue backpressure: a retained spill is never dropped (that gate is - // upstream) but still waits here for room, not bypassing the cap, so retained load can't grow RAM unbounded. - mDiskQueueCv.wait(lock, - [this, &job] - { - bool const slotClear = mDiskInflight.find(job.filename) == mDiskInflight.end(); - bool const hasRoom = mDiskWriteQueue.size() < mDiskWriteQueueMax; - return (slotClear && hasRoom) || mDiskWriterStop; - }); - if (mDiskWriterStop) - { - return; - } - ++mDiskInflight[job.filename]; // mark BEFORE the block becomes loadable - mDiskWriteQueue.push(std::move(job)); - } - mDiskQueueCv.notify_one(); + enqueueDiskWriteCommon(std::move(job)); } void KVCacheTransferManager::enqueueDiskWriteUnstaged( @@ -567,11 +548,15 @@ void KVCacheTransferManager::enqueueDiskWriteUnstaged( job.bytes = bytes; job.src = src; // no staging: the writer reads the pinned host slot directly job.spillId = spillId; + enqueueDiskWriteCommon(std::move(job)); +} +void KVCacheTransferManager::enqueueDiskWriteCommon(DiskWriteJob&& job) +{ { std::unique_lock lock(mDiskMutex); - // Same per-slot serialization + bounded-queue backpressure as enqueueDiskWrite: a retained spill waits - // for room too (never dropped, never bypasses the cap), so retained load stays RAM-bounded. + // Per-slot serialization + bounded-queue backpressure: a retained spill is never dropped (that gate is + // upstream at eviction) but still waits here for room, so retained load can't grow RAM unbounded. mDiskQueueCv.wait(lock, [this, &job] { diff --git a/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp b/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp index 98b96a358f3d..8b41a297b010 100644 --- a/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp @@ -10652,6 +10652,28 @@ class DiskTierDir std::filesystem::path mPath; }; +// Scoped setenv/unsetenv, so a gtest ASSERT early-exit cannot leak the variable into later tests. +class ScopedEnvVar +{ +public: + ScopedEnvVar(char const* name, char const* value) + : mName{name} + { + setenv(name, value, /*overwrite=*/1); + } + + ~ScopedEnvVar() + { + unsetenv(mName); + } + + ScopedEnvVar(ScopedEnvVar const&) = delete; + ScopedEnvVar& operator=(ScopedEnvVar const&) = delete; + +private: + char const* mName; +}; + // Small-pool BlockManager with a disk tier. tokensPerBlock=4; window fits maxBlocksPerSeq. std::unique_ptr makeDiskTierBlockManager(std::shared_ptr const& stream, std::string const& diskPath, bool retainedOnly, SizeType32 blocksInPrimaryPool, SizeType32 blocksInSecondaryPool, @@ -10773,7 +10795,7 @@ TEST_F(KVCacheManagerTest, DiskTierSpillAllRoundTripTest) // lands. Same spill/evict/reuse scenario as DiskTierSpillAllRoundTripTest, only with TLLM_KV_DISK_READERS set. TEST_F(KVCacheManagerTest, DiskTierDetachedOnboardReadyTest) { - setenv("TLLM_KV_DISK_READERS", "4", /*overwrite=*/1); + ScopedEnvVar readersEnv{"TLLM_KV_DISK_READERS", "4"}; auto const stream = std::make_shared(); DiskTierDir dir; auto blockManagerPtr @@ -10799,7 +10821,6 @@ TEST_F(KVCacheManagerTest, DiskTierDetachedOnboardReadyTest) EXPECT_GE(blockManager.getNumDiskOnboards(), 1u); releaseDiskTierSequence(blockManager, seqA2); - unsetenv("TLLM_KV_DISK_READERS"); } // Retained-only mode: unmarked victims are discarded at the gate; marked ones spill. @@ -11022,7 +11043,7 @@ float diskPatternValue(int slot, int i) // TLLM_KV_DISK_ASYNC_STORE is left unset so spills are synchronous and this isolates the read path. void runDiskOnboardByteRoundTrip(char const* readers, int nSlots) { - setenv("TLLM_KV_DISK_READERS", readers, /*overwrite=*/1); + ScopedEnvVar readersEnv{"TLLM_KV_DISK_READERS", readers}; DiskTierDir dir; int constexpr blockSize = 2048; // floats per slot @@ -11091,7 +11112,6 @@ void runDiskOnboardByteRoundTrip(char const* readers, int nSlots) EXPECT_EQ(bad, 0) << "readers=" << readers << " slot=" << k << ": disk onboard corrupted " << bad << "/" << blockSize << " elements"; } - unsetenv("TLLM_KV_DISK_READERS"); } } // namespace @@ -11119,8 +11139,8 @@ TEST_F(KVCacheManagerTest, DiskTierSyncOnboardByteExactTest) // Before the fix the TLLM_CHECK throw escaped the thread entry point and std::terminate'd the process. TEST_F(KVCacheManagerTest, DiskTierAsyncWriteFailureContainedTest) { - setenv("TLLM_KV_DISK_ASYNC_STORE", "1", /*overwrite=*/1); - setenv("TLLM_KV_DISK_WRITERS", "1", /*overwrite=*/1); + ScopedEnvVar asyncStoreEnv{"TLLM_KV_DISK_ASYNC_STORE", "1"}; + ScopedEnvVar writersEnv{"TLLM_KV_DISK_WRITERS", "1"}; DiskTierDir dir; int constexpr blockSize = 256; auto bufferManager = tr::BufferManager(std::make_shared()); @@ -11153,9 +11173,6 @@ TEST_F(KVCacheManagerTest, DiskTierAsyncWriteFailureContainedTest) auto src2 = std::make_shared(1, tk::KVCacheIndex(1, /*isSecondary=*/true)); transferManager.spillToFileUnstaged(src2, /*diskSlot=*/1, {pool}, dir.str(), /*spillId=*/2); EXPECT_TRUE(waitForSpill(2)) << "writer stopped processing after a failed job"; - - unsetenv("TLLM_KV_DISK_WRITERS"); - unsetenv("TLLM_KV_DISK_ASYNC_STORE"); } // A failed detached onboard read (missing slot file) must not crash the reader thread; the failure surfaces @@ -11163,7 +11180,7 @@ TEST_F(KVCacheManagerTest, DiskTierAsyncWriteFailureContainedTest) // never-filled block. Before the fix the reader's TLLM_CHECK throw std::terminate'd the process. TEST_F(KVCacheManagerTest, DiskTierOnboardReadFailureSurfacesTest) { - setenv("TLLM_KV_DISK_READERS", "2", /*overwrite=*/1); + ScopedEnvVar readersEnv{"TLLM_KV_DISK_READERS", "2"}; auto const stream = std::make_shared(); DiskTierDir dir; auto blockManagerPtr @@ -11205,7 +11222,6 @@ TEST_F(KVCacheManagerTest, DiskTierOnboardReadFailureSurfacesTest) // Teardown stays clean: the sequence releases without touching the never-filled blocks. releaseDiskTierSequence(blockManager, seqA2); - unsetenv("TLLM_KV_DISK_READERS"); } // shutdownDiskWorkers() must stop+join every disk worker while reads are queued (holding raw GPU dst @@ -11215,9 +11231,9 @@ TEST_F(KVCacheManagerTest, DiskTierOnboardReadFailureSurfacesTest) // bug where workers stopped only in the destructor, after releasePools() had already freed the pools. TEST_F(KVCacheManagerTest, DiskTierShutdownWithQueuedWorkTest) { - setenv("TLLM_KV_DISK_READERS", "4", /*overwrite=*/1); - setenv("TLLM_KV_DISK_WRITERS", "4", /*overwrite=*/1); - setenv("TLLM_KV_DISK_ASYNC_STORE", "1", /*overwrite=*/1); + ScopedEnvVar readersEnv{"TLLM_KV_DISK_READERS", "4"}; + ScopedEnvVar writersEnv{"TLLM_KV_DISK_WRITERS", "4"}; + ScopedEnvVar asyncStoreEnv{"TLLM_KV_DISK_ASYNC_STORE", "1"}; DiskTierDir dir; int constexpr blockSize = 2048; // floats per slot int constexpr nSlots = 8; @@ -11272,9 +11288,6 @@ TEST_F(KVCacheManagerTest, DiskTierShutdownWithQueuedWorkTest) pool.secondaryPtr->release(); // The transfer-manager destructor (at scope end) calls shutdownDiskWorkers() once more -- also a no-op. - unsetenv("TLLM_KV_DISK_READERS"); - unsetenv("TLLM_KV_DISK_WRITERS"); - unsetenv("TLLM_KV_DISK_ASYNC_STORE"); } // The disk write queue must stay bounded by mDiskWriteQueueMax under sustained spill load: with the retained @@ -11284,9 +11297,9 @@ namespace void runWriteQueueBoundedTest(bool unstaged) { using tensorrt_llm::testing::KVCacheTransferManagerTestAccess; - setenv("TLLM_KV_DISK_ASYNC_STORE", "1", /*overwrite=*/1); - setenv("TLLM_KV_DISK_WRITERS", "1", /*overwrite=*/1); // one writer drains; the cap must still bound the backlog - setenv("TLLM_KV_DISK_WRITE_QUEUE", "4", /*overwrite=*/1); + ScopedEnvVar asyncStoreEnv{"TLLM_KV_DISK_ASYNC_STORE", "1"}; + ScopedEnvVar writersEnv{"TLLM_KV_DISK_WRITERS", "1"}; // one writer drains; the cap must still bound the backlog + ScopedEnvVar writeQueueEnv{"TLLM_KV_DISK_WRITE_QUEUE", "4"}; std::size_t constexpr cap = 4; std::size_t constexpr burst = 200; // far more than the cap: the producer must block once the queue is full @@ -11331,10 +11344,6 @@ void runWriteQueueBoundedTest(bool unstaged) EXPECT_LE(observedMax.load(), cap) << "write queue exceeded its cap -- spills not backpressured"; EXPECT_GT(observedMax.load(), 0u) << "queue never filled -- the burst did not actually stress the cap"; - - unsetenv("TLLM_KV_DISK_WRITE_QUEUE"); - unsetenv("TLLM_KV_DISK_WRITERS"); - unsetenv("TLLM_KV_DISK_ASYNC_STORE"); } } // namespace From 1a551b9f1b93485876b85197b9d571e488cf88fa Mon Sep 17 00:00:00 2001 From: Md Nafis Ul Haque Shifat Date: Sun, 19 Jul 2026 18:36:33 -0700 Subject: [PATCH 39/47] [#16251][chore] minor comment cleanup Signed-off-by: Md Nafis Ul Haque Shifat --- .../tensorrt_llm/batch_manager/kvCacheManager.h | 2 +- .../batch_manager/kvCacheTransferManager.h | 2 +- .../batch_manager/kvCacheManager.cpp | 16 ++++++---------- .../batch_manager/kvCacheTransferManager.cpp | 8 ++++---- .../batch_manager/kvCacheManagerTest.cpp | 16 +++++++--------- 5 files changed, 19 insertions(+), 25 deletions(-) diff --git a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h index 9e920b7331a9..4ced21b325ac 100644 --- a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h +++ b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h @@ -1500,7 +1500,7 @@ class WindowBlockManager private: //! Disk-tier displacement order among retained blocks: exact earliest deadline; - //! among equal deadlines, first-spilled first. Upstream releases and evicts chains + //! among equal deadlines, first-spilled first. Block release and eviction walk chains //! leaf-first, so arrival order at the disk is deepest-first and FIFO here equals //! suffix-first within a chain (and LRU-like across requests). Entries are lazy: //! validated against live block state on pop, stale ones discarded. diff --git a/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h b/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h index c3ceec3d77cf..73279defc6e6 100644 --- a/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h +++ b/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h @@ -243,7 +243,7 @@ class KVCacheTransferManager void waitForDiskSlotWrites(std::string const& filename); // ---- Disk-tier async ONBOARD: read slot files disk->GPU OFF the scheduler thread ---- - // Gated by TLLM_KV_DISK_READERS (0/unset => synchronous read = current behavior; N>=1 => N reader + // Gated by TLLM_KV_DISK_READERS (0/unset => synchronous read, the default; N>=1 => N reader // threads drain a shared read queue). A job carries all of one block's pools; a request is forward-safe // once areBlocksReady() reports every block it holds has landed. See loadFromFile / isBlockReadPending. struct DiskReadJob diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp index bd4e46f97119..1f4ae827327b 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp @@ -1417,12 +1417,8 @@ BlockPtr WindowBlockManager::reclaimSecondaryBlock() } auto victim = std::get<0>(mEvictionPolicy->getFreeBlock(kSecondaryLevel)); mEvictionPolicy->claimBlock(victim); - // Reap landed read-pending releases BEFORE testing disk free space. Parked onboarded cards (and - // release-deferred blocks) sit OUTSIDE the free queues until reaped, so an onboard-heavy burst that - // parks every free disk card would drive getNumFreeBlocks(kDiskLevel) to 0, fire the early-return - // below, and never reach claimDiskTarget() -- the only other reap site -- wedging spills permanently. - // Reaping here (covering both reclaimSecondaryBlock callers) lets landed cards rejoin the disk free - // queue so the check sees the true count. + // Blocks with an in-flight onboard read sit outside the free queues; reap the landed ones + // before the free-space check below. if (mNumDiskBlocks > 0) { reapReadPendingReleases(); @@ -1438,7 +1434,7 @@ BlockPtr WindowBlockManager::reclaimSecondaryBlock() { TLLM_LOG_INFO("[disk-tier] gate dropped=%zu (windowSize=%d)", mDiskGateDropped, mWindowSize); } - return victim; // unmarked or expired: discard, exactly as stock + return victim; // unmarked or expired: discard, as when no disk tier is configured } // Reuse-gate: only spill blocks whose content has been reused >= k times (TLLM_KV_DISK_MIN_REUSE; 0=off). // Skips one-shot blocks so the SSD tier holds only genuinely-hot prefixes. Env read once, cached. @@ -1474,13 +1470,13 @@ BlockPtr WindowBlockManager::reclaimSecondaryBlock() TLLM_LOG_INFO( "[disk-tier] write-pressure dropped=%zu (windowSize=%d)", mDiskWritePressureDropped, mWindowSize); } - return victim; // best-effort spill shed under writer saturation: evict as stock (no scheduler stall) + return victim; // best-effort spill shed under writer saturation: plain eviction, no scheduler stall } auto diskTarget = claimDiskTarget(); if (diskTarget == nullptr) { // Protect mode + disk full of unexpired blocks: refuse so live TTLs are served. - // Victim drops from host exactly as an unmarked block (stock eviction). + // Victim drops from host exactly as an unmarked block would. ++mDiskAdmissionRefused; if (mDiskAdmissionRefused == 1 || mDiskAdmissionRefused % 10000 == 0) { @@ -1504,7 +1500,7 @@ BlockPtr WindowBlockManager::reclaimSecondaryBlock() mTransferManager->spillToFile(victim, diskSlot, mPools, mDiskCachePath); } victim->swapDiskResidency(diskTarget); // victim's identity now disk-resident, tree intact - victim->setDurationMs(std::nullopt); // keep it out of upstream's expiring-block machinery + victim->setDurationMs(std::nullopt); // keep it out of the eviction policy's expiring-block machinery if (victim->isRetainedNow()) { // Parked at max priority: displacement order among retained blocks is decided by diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp index fa339725108b..37f56b80ddc7 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp @@ -555,8 +555,8 @@ void KVCacheTransferManager::enqueueDiskWriteCommon(DiskWriteJob&& job) { { std::unique_lock lock(mDiskMutex); - // Per-slot serialization + bounded-queue backpressure: a retained spill is never dropped (that gate is - // upstream at eviction) but still waits here for room, so retained load can't grow RAM unbounded. + // Per-slot serialization + bounded-queue backpressure: a retained spill is never dropped (the drop + // gate runs at eviction) but still waits here for room, so retained load can't grow RAM unbounded. mDiskQueueCv.wait(lock, [this, &job] { @@ -611,8 +611,8 @@ void KVCacheTransferManager::diskReaderLoop() cudaStream_t stream; TLLM_CUDA_CHECK(cudaStreamCreate(&stream)); // Per-reader PINNED staging buffer, grown on demand to one job's total pool bytes (fixed after the first - // job). Pinned so cudaMemcpyAsync is a real async DMA (~20+ GB/s) instead of the driver's hidden pageable - // staging copy (~5-8 GB/s, secretly synchronous). pageableFallback is used only if a pinned alloc is refused. + // job). Pinned so cudaMemcpyAsync is a real async DMA instead of the driver's hidden pageable + // staging copy, which is effectively synchronous. pageableFallback is used only if a pinned alloc is refused. std::uint8_t* pinnedBuf = nullptr; std::size_t pinnedCap = 0; std::vector pageableFallback; diff --git a/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp b/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp index 92cd49095bc2..894f29895ed1 100644 --- a/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp @@ -10698,11 +10698,9 @@ TEST_F(KVCacheManagerTest, VswaDisaggDtypeMismatchTriggersGuard) } // ============================== Disk tier tests ============================== -// Appended to cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp — reuses -// the file's fixtures (KVCacheManagerTest, KvCacheManagerTestUtil). Covers the -// disk cache tier: spill/onboard round trip, the retained-only gate, retention -// stamping (fresh + reuse), expiry, deadline-ordered displacement, and -// suffix-first FIFO among equal deadlines. +// Covers the disk cache tier: spill/onboard round trip, the retained-only gate, +// retention stamping (fresh + reuse), expiry, deadline-ordered displacement, and +// suffix-first FIFO among equal deadlines. Reuses the file's fixtures. #include #include @@ -11219,8 +11217,8 @@ TEST_F(KVCacheManagerTest, DiskTierSyncOnboardByteExactTest) } // A failed async spill (unwritable slot file) must be contained on the writer thread: per-slot bookkeeping -// still runs (the spill id is published for the reserved-pool reap) and the worker keeps serving later jobs. -// Before the fix the TLLM_CHECK throw escaped the thread entry point and std::terminate'd the process. +// still runs (the spill id is published for the reserved-pool reap) and the worker keeps serving later +// jobs; an escaped throw from the thread entry point would std::terminate the process. TEST_F(KVCacheManagerTest, DiskTierAsyncWriteFailureContainedTest) { ScopedEnvVar asyncStoreEnv{"TLLM_KV_DISK_ASYNC_STORE", "1"}; @@ -11260,8 +11258,8 @@ TEST_F(KVCacheManagerTest, DiskTierAsyncWriteFailureContainedTest) } // A failed detached onboard read (missing slot file) must not crash the reader thread; the failure surfaces -// as a throw from areBlocksReady() on the calling (scheduler) thread, and no request is unparked onto the -// never-filled block. Before the fix the reader's TLLM_CHECK throw std::terminate'd the process. +// as a throw from areBlocksReady() on the calling (scheduler) thread, and no request is unparked onto +// the never-filled block. TEST_F(KVCacheManagerTest, DiskTierOnboardReadFailureSurfacesTest) { ScopedEnvVar readersEnv{"TLLM_KV_DISK_READERS", "2"}; From a2c86287a7b5f7ab0a33b6333eff0db95eeaeaa3 Mon Sep 17 00:00:00 2001 From: Md Nafis Ul Haque Shifat Date: Sun, 19 Jul 2026 18:46:15 -0700 Subject: [PATCH 40/47] [#16251][fix] erase disk deadline entries when their blocks leave the disk A retained block spilled to disk pushes an entry onto mDiskDeadlines, but nothing removed that entry when the block was onboarded back to GPU. The only cleanup was inside the displacement walk, which runs only when the disk is full of retained blocks -- so under normal reuse the stale entries accumulated forever (memory leak, and ever-longer walks once displacement does run). Fix: keep the deadlines in a std::set keyed by (expiry, seq). Each block remembers its key, so its entry is erased with one lookup whenever the block leaves the disk (onboard or displacement). The walk also erases any stale entry it finds. Signed-off-by: Md Nafis Ul Haque Shifat --- .../batch_manager/kvCacheManager.h | 54 ++++++++++++++-- .../batch_manager/kvCacheManager.cpp | 62 ++++++++++++------- .../batch_manager/kvCacheManagerTest.cpp | 20 ++++++ 3 files changed, 109 insertions(+), 27 deletions(-) diff --git a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h index 4ced21b325ac..1ddf7a17b697 100644 --- a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h +++ b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h @@ -456,6 +456,26 @@ class KVCacheBlock : public std::enable_shared_from_this mDiskSlot = diskSlot; } + [[nodiscard]] std::uint64_t getDiskDeadlineSeq() const + { + return mDiskDeadlineSeq; + } + + void setDiskDeadlineSeq(std::uint64_t seq) + { + mDiskDeadlineSeq = seq; + } + + [[nodiscard]] std::chrono::steady_clock::time_point::duration getDiskDeadlineExpiry() const + { + return mDiskDeadlineExpiry; + } + + void setDiskDeadlineExpiry(std::chrono::steady_clock::time_point::duration expiry) + { + mDiskDeadlineExpiry = expiry; + } + //! \brief Exchange tier residency with \p other: pool index and disk slot swap together. //! Used by the host->disk spill and the disk->GPU onboard; the tree-resident identity //! keeps its node/hash while the physical backing moves. @@ -674,6 +694,10 @@ class KVCacheBlock : public std::enable_shared_from_this // Disk cache tier: slot index in the disk file pool; kNoDiskSlot when not disk-resident. SizeType32 mDiskSlot{kNoDiskSlot}; + // (expiry, seq) key of this block's entry in the disk-tier deadline set; seq 0 = no entry. + std::chrono::steady_clock::time_point::duration mDiskDeadlineExpiry{}; + std::uint64_t mDiskDeadlineSeq{0}; + // Disk-tier retention deadline (steady clock); nullopt = never marked. std::optional mRetentionExpiry; @@ -1498,29 +1522,34 @@ class WindowBlockManager return mDiskOnboards; } + [[nodiscard]] std::size_t getDiskDeadlineCount() const + { + return mDiskDeadlines.size(); + } + private: //! Disk-tier displacement order among retained blocks: exact earliest deadline; //! among equal deadlines, first-spilled first. Block release and eviction walk chains //! leaf-first, so arrival order at the disk is deepest-first and FIFO here equals - //! suffix-first within a chain (and LRU-like across requests). Entries are lazy: - //! validated against live block state on pop, stale ones discarded. + //! suffix-first within a chain (and LRU-like across requests). An entry is erased + //! when its block leaves the disk (onboard or displacement). struct DiskDeadline { std::chrono::steady_clock::time_point::duration expiry; std::uint64_t seq; BlockPtr block; - bool operator>(DiskDeadline const& other) const + bool operator<(DiskDeadline const& other) const { if (expiry != other.expiry) { - return expiry > other.expiry; + return expiry < other.expiry; } - return seq > other.seq; // first-arrived pops first among equal deadlines + return seq < other.seq; // first-spilled first among equal deadlines } }; - std::priority_queue, std::greater> mDiskDeadlines; + std::set mDiskDeadlines; std::uint64_t mDiskSpillSeq{0}; // ---- Unstaged async store: reserved host-block pool (env TLLM_KV_DISK_RESERVED_BLOCKS, 0 = off) ---- @@ -1539,6 +1568,9 @@ class WindowBlockManager // destination is never repurposed mid-read. See releaseBlocks / the areBlocksReady park. std::vector mReleaseReadPending; void reapReadPendingReleases(); + + //! \brief Remove \p block's mDiskDeadlines entry, if any. + void eraseDiskDeadline(BlockPtr const& block); std::uint64_t mUnstagedSpillSeq{0}; //! \brief Pick and claim the disk block to overwrite: empty slots first, then @@ -2177,6 +2209,16 @@ class BlockManager return total; } + [[nodiscard]] std::size_t getDiskDeadlineCount() const + { + std::size_t total = 0; + for (auto const& [windowSize, manager] : mWindowBlockManagers) + { + total += manager.getDiskDeadlineCount(); + } + return total; + } + private: std::map mWindowBlockManagers; std::map mWindowSizeToMetadata; diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp index 1f4ae827327b..b0c50e09f1da 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp @@ -1354,6 +1354,17 @@ void WindowBlockManager::reapReadPendingReleases() } } +void WindowBlockManager::eraseDiskDeadline(BlockPtr const& block) +{ + auto const seq = block->getDiskDeadlineSeq(); + if (seq == 0) + { + return; + } + mDiskDeadlines.erase(DiskDeadline{block->getDiskDeadlineExpiry(), seq, nullptr}); + block->setDiskDeadlineSeq(0); +} + BlockPtr WindowBlockManager::claimDiskTarget() { // Reclaim onboarded slots whose read has landed before handing out a disk slot, so a slot with an @@ -1368,28 +1379,33 @@ BlockPtr WindowBlockManager::claimDiskTarget() mEvictionPolicy->claimBlock(candidate); return candidate; } - // Every free disk block is retained: displace the earliest real deadline; among + // Every free disk block is retained: displace the earliest live deadline; among // equal deadlines the first-spilled block goes first (leaf-first arrival makes this - // suffix-first within a chain). Entries are validated against live block state on - // pop; failures mean the entry describes the past (block left disk, was re-stamped, - // or is mid-transfer and about to leave) and are discarded. - while (!mDiskDeadlines.empty()) - { - auto const top = mDiskDeadlines.top(); - mDiskDeadlines.pop(); - if (top.block->isOnDisk() && !top.block->hasRefs() && mEvictionPolicy->isEnqueued(top.block) - && top.block->getRetentionExpiry() == std::optional{top.expiry}) - { - // Protect mode: earliest deadline still future => all live => refuse. - // (Expired blocks sort first, so any would have been evicted above.) - if (mDiskProtectUnexpired && top.block->isRetainedNow()) - { - mDiskDeadlines.push(top); - return nullptr; - } - mEvictionPolicy->claimBlock(top.block); - return top.block; + // suffix-first within a chain). Mid-onboard blocks stay for later; stale entries are + // erased on sight. + for (auto it = mDiskDeadlines.begin(); it != mDiskDeadlines.end();) + { + if (!it->block->isOnDisk()) + { + it = mDiskDeadlines.erase(it); + continue; + } + if (it->block->hasRefs() || !mEvictionPolicy->isEnqueued(it->block)) + { + ++it; + continue; + } + // Protect mode: earliest deadline still future => all live => refuse. + // (Expired blocks sort first, so any would have been evicted above.) + if (mDiskProtectUnexpired && it->block->isRetainedNow()) + { + return nullptr; } + auto block = it->block; + block->setDiskDeadlineSeq(0); + mDiskDeadlines.erase(it); + mEvictionPolicy->claimBlock(block); + return block; } // Heap drained by stale entries; refuse if the remaining candidate is still live. if (mDiskProtectUnexpired && candidate->isRetainedNow()) @@ -1506,7 +1522,10 @@ BlockPtr WindowBlockManager::reclaimSecondaryBlock() // Parked at max priority: displacement order among retained blocks is decided by // mDiskDeadlines (exact deadline, suffix-first within a chain), not queue order. victim->setPriority(executor::KvCacheRetentionConfig::kMaxRetentionPriority); - mDiskDeadlines.push({*victim->getRetentionExpiry(), ++mDiskSpillSeq, victim}); + auto const seq = ++mDiskSpillSeq; + victim->setDiskDeadlineSeq(seq); + victim->setDiskDeadlineExpiry(*victim->getRetentionExpiry()); + mDiskDeadlines.insert(DiskDeadline{victim->getDiskDeadlineExpiry(), seq, victim}); } else { @@ -1704,6 +1723,7 @@ void WindowBlockManager::onboardBlock(GenerationRequest& sequence, BlockPtr cons auto const trackId = mTransferManager->asyncDiskReadEnabled() ? offloadBlock->getBlockId() : static_cast(-1); mTransferManager->loadFromFile(block, offloadBlock->getDiskSlot(), mPools, mDiskCachePath, trackId); + eraseDiskDeadline(offloadBlock); offloadBlock->swapDiskResidency(block); // matched identity now GPU-resident (bytes may still be in flight) offloadBlock->setPriority(executor::KvCacheRetentionConfig::kDefaultRetentionPriority); block->clearRetention(); diff --git a/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp b/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp index 894f29895ed1..14bd5d759080 100644 --- a/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp @@ -11440,3 +11440,23 @@ TEST_F(KVCacheManagerTest, DiskTierWriteQueueBoundedUnstagedTest) { runWriteQueueBoundedTest(/*unstaged=*/true); } + +// A retained block's deadline entry must be erased when the block leaves the disk: spill->reuse +// cycles otherwise accumulate entries forever (disk=64 stays roomy, so displacement never runs). +TEST_F(KVCacheManagerTest, DiskTierDeadlineEntryLeakTest) +{ + auto const stream = std::make_shared(); + DiskTierDir dir; + auto bm = makeDiskTierBlockManager(stream, dir.str(), /*retainedOnly=*/true, /*prim=*/4, /*sec=*/2, /*disk=*/64); + auto& blockManager = *bm; + for (SizeType32 i = 0; i < 100; ++i) + { + auto s = addDiskTierSequence(blockManager, 1000 + i, iotaTokens(0, 16), std::chrono::milliseconds(600000)); + releaseDiskTierSequence(blockManager, s); + churnOnce(blockManager, 2000 + i, static_cast(100 + i)); + auto r = addDiskTierSequence(blockManager, 3000 + i, iotaTokens(0, 16), std::chrono::milliseconds(600000)); + releaseDiskTierSequence(blockManager, r); + } + EXPECT_LE(blockManager.getDiskDeadlineCount(), + static_cast(64)); // never more entries than disk slots; leaked entries grow past it +} From 23c7130df9d7ceb8260a2e0ab8f858cbc7ff16dd Mon Sep 17 00:00:00 2001 From: Md Nafis Ul Haque Shifat Date: Mon, 20 Jul 2026 17:34:52 -0700 Subject: [PATCH 41/47] [#16251][fix] refuse TLLM_KV_DISK_DROP_ON_PRESSURE with tensor parallelism The drop decision depends on each rank's write-queue depth, so ranks could diverge on which blocks remain reusable. Refuse the combination at init, like attention-DP and pipeline parallelism. Signed-off-by: Md Nafis Ul Haque Shifat --- tensorrt_llm/_torch/pyexecutor/resource_manager.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index dfd7949d15e9..dda3abbb12e2 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -645,6 +645,13 @@ def append_to_kv_heads_per_layer(num_kv_heads_per_layer: List[int], f"enable_attention_dp={mapping.enable_attention_dp}, pp_size={mapping.pp_size}, " f"cp_size={mapping.cp_size}. " "Disable disk_cache_size or use plain tensor parallelism.") + if disk_cache_size and mapping.tp_size > 1 and os.environ.get( + "TLLM_KV_DISK_DROP_ON_PRESSURE", "0") not in ("", "0"): + raise ValueError( + "TLLM_KV_DISK_DROP_ON_PRESSURE is not supported with tensor parallelism: " + "the drop decision depends on each rank's write-queue depth at spill time, " + "so ranks diverge on which blocks remain reusable. Unset it or use TP=1." + ) blocks_in_disk_pool = 0 if disk_cache_size: max_tokens_disk = disk_cache_size // self.get_cache_bytes_per_token( From 492a0cfbf62509d7c4fb917fe0af8e4ddba6ee60 Mon Sep 17 00:00:00 2001 From: Md Nafis Ul Haque Shifat Date: Mon, 20 Jul 2026 17:34:52 -0700 Subject: [PATCH 42/47] [#16251][fix] anchor retention decisions to a rank-consistent clock Retention deadlines were anchored and checked against each rank's own wall clock, which could diverge the radix trees across ranks. Now one rank broadcasts its time each iteration and every rank uses it for retention decisions. Signed-off-by: Md Nafis Ul Haque Shifat --- .../batch_manager/kvCacheManager.h | 32 ++++++++++-- .../batch_manager/kvCacheManager.cpp | 16 +++--- .../nanobind/batch_manager/kvCacheManager.cpp | 2 + .../batch_manager/kvCacheManagerTest.cpp | 50 +++++++++++++++---- tensorrt_llm/_torch/pyexecutor/py_executor.py | 18 +++++++ .../_torch/pyexecutor/resource_manager.py | 3 ++ 6 files changed, 99 insertions(+), 22 deletions(-) diff --git a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h index 34984957e7f1..c6d875a5003d 100644 --- a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h +++ b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h @@ -492,9 +492,9 @@ class KVCacheBlock : public std::enable_shared_from_this mRetentionExpiry = mRetentionExpiry ? std::max(*mRetentionExpiry, expiry) : expiry; } - [[nodiscard]] bool isRetainedNow() const + [[nodiscard]] bool isRetainedNow(std::chrono::steady_clock::time_point::duration now) const { - return mRetentionExpiry.has_value() && *mRetentionExpiry > std::chrono::steady_clock::now().time_since_epoch(); + return mRetentionExpiry.has_value() && *mRetentionExpiry > now; } [[nodiscard]] std::optional getRetentionExpiry() const @@ -515,12 +515,12 @@ class KVCacheBlock : public std::enable_shared_from_this = mPendingRetentionDuration ? std::max(*mPendingRetentionDuration, duration) : duration; } - //! \brief Anchor a pending retention duration to now (called when the block is committed / made reusable). - void commitRetention() + //! \brief Anchor a pending retention duration to \p now (called when the block is committed / made reusable). + void commitRetention(std::chrono::steady_clock::time_point::duration now) { if (mPendingRetentionDuration) { - markRetained(std::chrono::steady_clock::now().time_since_epoch() + *mPendingRetentionDuration); + markRetained(now + *mPendingRetentionDuration); mPendingRetentionDuration = std::nullopt; } } @@ -1527,6 +1527,11 @@ class WindowBlockManager return mDiskDeadlines.size(); } + void setRetentionClock(std::chrono::steady_clock::time_point::duration now) + { + mRetentionNow = now; + } + private: //! Disk-tier displacement order among retained blocks: exact earliest deadline; //! among equal deadlines, first-spilled first. Block release and eviction walk chains @@ -1552,6 +1557,9 @@ class WindowBlockManager std::set mDiskDeadlines; std::uint64_t mDiskSpillSeq{0}; + // Retention decisions use this rank-consistent clock (leader-broadcast each iteration). + std::chrono::steady_clock::time_point::duration mRetentionNow{0}; + // ---- Unstaged async store: reserved host-block pool (env TLLM_KV_DISK_RESERVED_BLOCKS, 0 = off) ---- // Target set from the environment in allocatePools. A spill hands out a reserved free host slot // immediately and pins the victim's old slot until its async write drains, avoiding the staging @@ -2219,6 +2227,14 @@ class BlockManager return total; } + void setRetentionClock(std::chrono::steady_clock::time_point::duration now) + { + for (auto& [windowSize, manager] : mWindowBlockManagers) + { + manager.setRetentionClock(now); + } + } + private: std::map mWindowBlockManagers; std::map mWindowSizeToMetadata; @@ -2923,6 +2939,12 @@ class KVCacheManager : public BaseKVCacheManager mBlockManager.refreshBlocks(); } + void setRetentionClock(std::int64_t nowNs) + { + mBlockManager.setRetentionClock(std::chrono::duration_cast( + std::chrono::nanoseconds(nowNs))); + } + void flushIterationEvents() override { mBlockManager.flushIterationEvents(); diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp index ee03bd270fa6..da64dc90de4f 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp @@ -1375,7 +1375,7 @@ BlockPtr WindowBlockManager::claimDiskTarget() // so getFreeBlock returns empties first and a retained block only when nothing cheaper // (empty slot or unmarked cached block) remains. auto candidate = std::get<0>(mEvictionPolicy->getFreeBlock(kDiskLevel)); - if (!candidate->isRetainedNow()) + if (!candidate->isRetainedNow(mRetentionNow)) { mEvictionPolicy->claimBlock(candidate); return candidate; @@ -1398,7 +1398,7 @@ BlockPtr WindowBlockManager::claimDiskTarget() } // Protect mode: earliest deadline still future => all live => refuse. // (Expired blocks sort first, so any would have been evicted above.) - if (mDiskProtectUnexpired && it->block->isRetainedNow()) + if (mDiskProtectUnexpired && it->block->isRetainedNow(mRetentionNow)) { return nullptr; } @@ -1409,7 +1409,7 @@ BlockPtr WindowBlockManager::claimDiskTarget() return block; } // Heap drained by stale entries; refuse if the remaining candidate is still live. - if (mDiskProtectUnexpired && candidate->isRetainedNow()) + if (mDiskProtectUnexpired && candidate->isRetainedNow(mRetentionNow)) { return nullptr; } @@ -1444,7 +1444,7 @@ BlockPtr WindowBlockManager::reclaimSecondaryBlock() { return victim; // disk tier off, nothing reusable to keep, or disk full: same as before } - if (mDiskRetainedOnly && !victim->isRetainedNow()) + if (mDiskRetainedOnly && !victim->isRetainedNow(mRetentionNow)) { ++mDiskGateDropped; if (mDiskGateDropped == 1 || mDiskGateDropped % 10000 == 0) @@ -1479,7 +1479,7 @@ BlockPtr WindowBlockManager::reclaimSecondaryBlock() auto const* e = std::getenv("TLLM_KV_DISK_DROP_ON_PRESSURE"); return e && std::atoi(e) != 0; }(); - if (kDropOnPressure && !victim->isRetainedNow() && mTransferManager->diskWriteQueueFull()) + if (kDropOnPressure && !victim->isRetainedNow(mRetentionNow) && mTransferManager->diskWriteQueueFull()) { ++mDiskWritePressureDropped; if (mDiskWritePressureDropped == 1 || mDiskWritePressureDropped % 10000 == 0) @@ -1518,7 +1518,7 @@ BlockPtr WindowBlockManager::reclaimSecondaryBlock() } victim->swapDiskResidency(diskTarget); // victim's identity now disk-resident, tree intact victim->setDurationMs(std::nullopt); // keep it out of the eviction policy's expiring-block machinery - if (victim->isRetainedNow()) + if (victim->isRetainedNow(mRetentionNow)) { // Parked at max priority: displacement order among retained blocks is decided by // mDiskDeadlines (exact deadline, suffix-first within a chain), not queue order. @@ -2271,7 +2271,7 @@ SizeType32 WindowBlockManager::onboardAndAllocateBlocks( { // Reused blocks bypass getFreeBlock, so the disk-retention stamp must also // happen here; markRetained max-merges, keeping the later deadline. - claimed.block->markRetained(std::chrono::steady_clock::now().time_since_epoch() + *diskRetentionMs); + claimed.block->markRetained(mRetentionNow + *diskRetentionMs); } onboardBlock(sequence, claimed.block, claimResult.mode, claimResult.directory); addBlockToAllBeams(claimed.block, sequence); @@ -3115,7 +3115,7 @@ std::pair> WindowBlockManager::sto // Blocks are now committed to the reuse tree: anchor any pending retention duration to now. for (auto const& b : storedBlocks) { - b->commitRetention(); + b->commitRetention(mRetentionNow); } if (mEventManager) diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp index a77082ee02aa..9fc6a83543e6 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp @@ -660,6 +660,8 @@ void tb::kv_cache_manager::KVCacheManagerBindings::initBindings(nb::module_& m) .value("SELFKONLY", tbk::CacheType::kSELFKONLY); nb::class_(m, "KVCacheManager") + .def("set_retention_clock", &tbk::KVCacheManager::setRetentionClock, nb::arg("now_ns"), + nb::call_guard()) .def(nb::init const&, SizeType32, SizeType32, std::map> const&, SizeType32, SizeType32, std::vector const&, tensorrt_llm::DataType, SizeType32, int64_t, SizeType32, SizeType32, diff --git a/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp b/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp index 8ffead639537..3201c4b0a8c6 100644 --- a/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp @@ -10817,6 +10817,11 @@ VecTokens iotaTokens(TokenIdType first, SizeType32 count) // Adds a sequence; diskTtl set => the request carries a disk retention TTL. // Returns the pair; prepopulated length is readable via req->getContextCurrentPosition(). +std::chrono::steady_clock::time_point::duration retentionTick(std::int64_t ms) +{ + return std::chrono::duration_cast(std::chrono::milliseconds(ms)); +} + DiskTierSeq addDiskTierSequence(BlockManager& blockManager, LlmRequest::RequestIdType id, VecTokens tokens, std::optional diskTtl) { @@ -10962,9 +10967,10 @@ TEST_F(KVCacheManagerTest, DiskTierExpiryAtSpillTest) = makeDiskTierBlockManager(stream, dir.str(), /*retainedOnly=*/true, /*prim=*/4, /*sec=*/2, /*disk=*/8); auto& blockManager = *blockManagerPtr; + blockManager.setRetentionClock(retentionTick(1000)); auto seqP = addDiskTierSequence(blockManager, 0, iotaTokens(0, 16), std::chrono::milliseconds(50)); releaseDiskTierSequence(blockManager, seqP); - std::this_thread::sleep_for(std::chrono::milliseconds(120)); // TTL dead before eviction + blockManager.setRetentionClock(retentionTick(1200)); // TTL dead before eviction churnOnce(blockManager, 1, 100); EXPECT_EQ(blockManager.getNumDiskSpills(), 0u); EXPECT_GE(blockManager.getNumDiskGateDropped(), 2u); @@ -10993,10 +10999,10 @@ TEST_F(KVCacheManagerTest, DiskTierRetentionAnchoredAtCommitTest) << "retention deadline stamped at allocation instead of deferred to commit"; } - // Commit to the reuse tree -> the deadline is anchored at ~now + ttl for the committed blocks. - auto const before = std::chrono::steady_clock::now().time_since_epoch(); + // Commit to the reuse tree -> the deadline is anchored at the retention clock + ttl. + auto const t0 = retentionTick(5000); + blockManager.setRetentionClock(t0); blockManager.storeContextBlocks(*s.seq, *s.req); - auto const after = std::chrono::steady_clock::now().time_since_epoch(); std::size_t stamped = 0; for (auto const bid : blockIds) @@ -11005,8 +11011,7 @@ TEST_F(KVCacheManagerTest, DiskTierRetentionAnchoredAtCommitTest) if (expiry.has_value()) { ++stamped; - EXPECT_GE(*expiry, before + ttl); // clock started at commit, not earlier at allocation - EXPECT_LE(*expiry, after + ttl); + EXPECT_EQ(*expiry, t0 + ttl); // clock started at commit, not earlier at allocation } } EXPECT_GT(stamped, 0u) << "commit did not anchor any retention deadline"; @@ -11022,13 +11027,14 @@ TEST_F(KVCacheManagerTest, DiskTierReuseRestampTest) auto& blockManager = *blockManagerPtr; // Short-TTL cold, then immediate re-send with a long TTL: pure reuse path. + blockManager.setRetentionClock(retentionTick(1000)); auto seqA = addDiskTierSequence(blockManager, 0, iotaTokens(0, 16), std::chrono::milliseconds(50)); releaseDiskTierSequence(blockManager, seqA); auto seqA2 = addDiskTierSequence(blockManager, 1, iotaTokens(0, 16), std::chrono::milliseconds(60000)); EXPECT_GE(seqA2.req->getContextCurrentPosition(), 12); // reused, not recomputed releaseDiskTierSequence(blockManager, seqA2); - std::this_thread::sleep_for(std::chrono::milliseconds(120)); // past the ORIGINAL 50 ms TTL + blockManager.setRetentionClock(retentionTick(1200)); // past the ORIGINAL 50 ms TTL churnOnce(blockManager, 2, 100); // Without the reuse re-stamp the blocks are expired and gated; with it they spill. EXPECT_GE(blockManager.getNumDiskSpills(), 2u); @@ -11036,7 +11042,7 @@ TEST_F(KVCacheManagerTest, DiskTierReuseRestampTest) // Control: same shape, no re-send. Expires and is gated. auto seqB = addDiskTierSequence(blockManager, 3, iotaTokens(200, 16), std::chrono::milliseconds(50)); releaseDiskTierSequence(blockManager, seqB); - std::this_thread::sleep_for(std::chrono::milliseconds(120)); + blockManager.setRetentionClock(retentionTick(1400)); // past seqB's TTL auto const spillsBefore = blockManager.getNumDiskSpills(); churnOnce(blockManager, 4, 300); EXPECT_EQ(blockManager.getNumDiskSpills(), spillsBefore); @@ -11109,12 +11115,13 @@ TEST_F(KVCacheManagerTest, DiskTierProtectEvictsExpiredTest) stream, dir.str(), /*retainedOnly=*/true, /*prim=*/4, /*sec=*/2, /*disk=*/2, /*protectUnexpired=*/true); auto& blockManager = *blockManagerPtr; + blockManager.setRetentionClock(retentionTick(1000)); auto seqA = addDiskTierSequence(blockManager, 0, iotaTokens(0, 16), std::chrono::milliseconds(50)); releaseDiskTierSequence(blockManager, seqA); churnOnce(blockManager, 1, 100); auto const spillsFilled = blockManager.getNumDiskSpills(); EXPECT_GE(spillsFilled, 2u); - std::this_thread::sleep_for(std::chrono::milliseconds(120)); // A's disk blocks now expired + blockManager.setRetentionClock(retentionTick(1200)); // A's disk blocks now expired // New victim: the expired blocks ARE evictable even in protect mode => it spills. auto seqB = addDiskTierSequence(blockManager, 2, iotaTokens(200, 16), std::chrono::milliseconds(600000)); @@ -11475,3 +11482,28 @@ TEST_F(KVCacheManagerTest, DiskTierDeadlineEntryLeakTest) EXPECT_LE(blockManager.getDiskDeadlineCount(), static_cast(64)); // never more entries than disk slots; leaked entries grow past it } + +// Retention decisions must be identical across TP ranks: anchors and expiry checks use the +// injected rank-consistent clock, not each rank's wall clock. Two managers driven with the +// identical op sequence but different real-time pacing must agree on the gate outcome. +TEST_F(KVCacheManagerTest, DiskTierRetentionClockConsistencyTest) +{ + auto const stream = std::make_shared(); + DiskTierDir dirA, dirB; + auto bmA = makeDiskTierBlockManager(stream, dirA.str(), /*retainedOnly=*/true, /*prim=*/4, /*sec=*/2, /*disk=*/8); + auto bmB = makeDiskTierBlockManager(stream, dirB.str(), /*retainedOnly=*/true, /*prim=*/4, /*sec=*/2, /*disk=*/8); + for (auto* bm : {bmA.get(), bmB.get()}) + { + bm->setRetentionClock(retentionTick(1000)); + auto s = addDiskTierSequence(*bm, 1, iotaTokens(0, 16), std::chrono::milliseconds(150)); + releaseDiskTierSequence(*bm, s); + if (bm == bmB.get()) + { + std::this_thread::sleep_for(std::chrono::milliseconds(300)); // real-time skew between ranks + } + bm->setRetentionClock(retentionTick(1100)); // same logical instant on both, inside the TTL + churnOnce(*bm, 2, 100); // retained-only gate decides spill-vs-drop here + } + EXPECT_EQ(bmA->getDiskDeadlineCount(), bmB->getDiskDeadlineCount()); + EXPECT_GT(bmA->getDiskDeadlineCount(), 0u); +} diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 6e8a2d59517e..c56f5160f262 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -594,6 +594,8 @@ def __init__( self._disk_onboard_active = bool( _kvc is not None and getattr(_kvc, "disk_cache_size", None) and int(os.environ.get("TLLM_KV_DISK_READERS", "0") or "0") > 0) + self._disk_tier_active = bool( + _kvc is not None and getattr(_kvc, "disk_cache_size", None)) self.max_stats_len = self.llm_args.max_stats_len self.max_num_tokens = self.llm_args.max_num_tokens self.print_log = self.llm_args.print_iter_log @@ -2612,6 +2614,7 @@ def _executor_loop_pp(self): self._handle_dynamic_draft_len(scheduled_batch) + self._sync_retention_clock() self.resource_manager.prepare_resources(scheduled_batch) self._park_requests_awaiting_onboard(scheduled_batch) @@ -3174,6 +3177,19 @@ def wait_on_pp_send_handles(self, send_handles, microbatch_id): send_handles[microbatch_id].wait() send_handles[microbatch_id] = None + def _sync_retention_clock(self): + """All ranks anchor and evaluate retention against the same timestamp.""" + if not self._disk_tier_active: + return + kv = getattr(self, "kv_cache_manager", None) + if kv is None or not hasattr(kv, "set_retention_clock"): + return + now_ns = time.time_ns() + if self.dist.tp_size > 1 and not self.enable_attention_dp: + now_ns = self.dist.tp_allreduce( + now_ns if self.dist.tp_rank == 0 else 0, op=ReduceOp.MAX) + kv.set_retention_clock(now_ns) + def _park_requests_awaiting_onboard(self, scheduled_batch): """Hold context requests whose detached disk-onboard reads have not landed out of this forward. Their KV blocks are allocated but not yet filled, so forwarding now would read stale KV. A parked @@ -4097,6 +4113,7 @@ def _executor_loop(self): self._handle_dynamic_draft_len(scheduled_batch) + self._sync_retention_clock() self.resource_manager.prepare_resources(scheduled_batch) self._park_requests_awaiting_onboard(scheduled_batch) @@ -4574,6 +4591,7 @@ def _executor_loop_overlap(self): self._handle_dynamic_draft_len(scheduled_batch) + self._sync_retention_clock() self.resource_manager.prepare_resources(scheduled_batch) self._park_requests_awaiting_onboard(scheduled_batch) diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index dda3abbb12e2..23c9c46dfa38 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -1502,6 +1502,9 @@ def are_blocks_ready(self, request) -> bool: Always True for requests with no in-flight disk read, so it is cheap for the common case.""" return self.impl.are_blocks_ready(request.py_request_id) + def set_retention_clock(self, now_ns: int) -> None: + self.impl.set_retention_clock(now_ns) + def get_num_free_blocks(self) -> int: if self.is_linear_attention: value = self.impl.get_kv_cache_stats( From 46e404093c70a648f3e5cae11f2beeaf5d9e5166 Mon Sep 17 00:00:00 2001 From: Md Nafis Ul Haque Shifat Date: Mon, 20 Jul 2026 21:56:19 -0700 Subject: [PATCH 43/47] [#16251][chore] rename a parameter flagged by codespell Signed-off-by: Md Nafis Ul Haque Shifat --- cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h index c6d875a5003d..0772e68b38a0 100644 --- a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h +++ b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h @@ -2939,10 +2939,10 @@ class KVCacheManager : public BaseKVCacheManager mBlockManager.refreshBlocks(); } - void setRetentionClock(std::int64_t nowNs) + void setRetentionClock(std::int64_t timestampNs) { mBlockManager.setRetentionClock(std::chrono::duration_cast( - std::chrono::nanoseconds(nowNs))); + std::chrono::nanoseconds(timestampNs))); } void flushIterationEvents() override From 343512c7bfda83e4d1c53df3b6789d0b960b0d18 Mon Sep 17 00:00:00 2001 From: Md Nafis Ul Haque Shifat Date: Wed, 29 Jul 2026 10:55:11 -0700 Subject: [PATCH 44/47] [#16251][chore] add disk write-pressure getter and fix onboard doc comment Expose mDiskWritePressureDropped via getters at both WindowBlockManager and BlockManager, matching the other disk counters. Split the stacked doc comment so reclaimSecondaryBlock and onboardBlock each carry their own. Signed-off-by: Md Nafis Ul Haque Shifat --- .../batch_manager/kvCacheManager.h | 19 +++++++++++++++++-- 1 file changed, 17 insertions(+), 2 deletions(-) diff --git a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h index 13230f28ba0e..382f3a82439e 100644 --- a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h +++ b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h @@ -1319,13 +1319,13 @@ class WindowBlockManager void setOffsets(kernels::KVCacheIndex* offsetsPtr, tensorrt_llm::Dims const& offsetsShape, SizeType32 beamIdx, SizeType32 blockIdx, KVCacheBlock::IdType blockId) const; - //! \brief Bring offloaded block from secondary to primary memory. - //! \details Does nothing if block is already in primary memory. //! \brief Reclaim a free secondary (host) block for reuse. If the victim still holds //! reusable content and the disk tier has room, spill it to disk first (residency swap; //! the tree-resident identity moves to the disk level). Returns a CLAIMED block. [[nodiscard]] BlockPtr reclaimSecondaryBlock(); + //! \brief Bring offloaded block from secondary to primary memory. + //! \details Does nothing if block is already in primary memory. void onboardBlock(GenerationRequest& sequence, BlockPtr const& offloadBlock, executor::KvCacheTransferMode mode = executor::KvCacheTransferMode::DRAM, std::string const& directory = ""); @@ -1536,6 +1536,11 @@ class WindowBlockManager return mDiskAdmissionRefused; } + [[nodiscard]] std::size_t getNumDiskWritePressureDropped() const + { + return mDiskWritePressureDropped; + } + [[nodiscard]] std::size_t getNumDiskOnboards() const { return mDiskOnboards; @@ -2232,6 +2237,16 @@ class BlockManager return total; } + [[nodiscard]] std::size_t getNumDiskWritePressureDropped() const + { + std::size_t total = 0; + for (auto const& [windowSize, manager] : mWindowBlockManagers) + { + total += manager.getNumDiskWritePressureDropped(); + } + return total; + } + [[nodiscard]] std::size_t getNumDiskOnboards() const { std::size_t total = 0; From 4bcca7baa3d17b99eaed5904b015d64437cdcab3 Mon Sep 17 00:00:00 2001 From: Md Nafis Ul Haque Shifat Date: Wed, 5 Aug 2026 10:53:27 -0700 Subject: [PATCH 45/47] [#16251][fix] disk tier: fix per-manager slot collisions under spec decode A parked disk onboard re-enters as a first context chunk, so eagle3/mtp reuse an existing hidden-state slot instead of re-allocating. And the disk slot directory is namespaced by draft/target, not just rank, so a draft KV manager cannot overwrite the target's slot files. Signed-off-by: Md Nafis Ul Haque Shifat --- tensorrt_llm/_torch/pyexecutor/resource_manager.py | 12 +++++------- tensorrt_llm/_torch/speculative/eagle3.py | 6 +++++- tensorrt_llm/_torch/speculative/mtp.py | 6 +++++- 3 files changed, 15 insertions(+), 9 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index 1bf3fc93504a..ba2a898cb4e2 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -669,14 +669,12 @@ def append_to_kv_heads_per_layer(num_kv_heads_per_layer: List[int], disk_cache_path = getattr(kv_cache_config, "disk_cache_path", None) or "" if disk_cache_path and blocks_in_disk_pool: - # Per-rank subdir: all ranks share one container filesystem and one - # disk_cache_path mount, but each holds a different KV shard while the - # disk slot files (block__pool_

.bin) are named identically across - # ranks -> without namespacing the ranks clobber each other's KV on load. - # mapping.rank is the global rank (unique across TP and PP). Note: + # Namespace by rank and draft/target: slot files restart at 0 per manager, + # and one-model spec decode gives the draft its own manager on the same rank. # disk_cache_size is per-rank, so total disk = world_size * disk_cache_size. - disk_cache_path = os.path.join(disk_cache_path, - f"rank_{self.mapping.rank}") + disk_cache_path = os.path.join( + disk_cache_path, f"rank_{self.mapping.rank}" + + ("_draft" if self.is_draft else "")) os.makedirs(disk_cache_path, exist_ok=True) logger.info( f"[disk-tier] per-rank disk cache dir: {disk_cache_path} " diff --git a/tensorrt_llm/_torch/speculative/eagle3.py b/tensorrt_llm/_torch/speculative/eagle3.py index a1f4f82869bd..f6105fb8367b 100644 --- a/tensorrt_llm/_torch/speculative/eagle3.py +++ b/tensorrt_llm/_torch/speculative/eagle3.py @@ -108,7 +108,11 @@ def prepare_resources(self, scheduled_batch: ScheduledRequests): self.slot_ids = [] for req in context_batch: if req.is_first_context_chunk: - slot_id = self.slot_manager.add_slot(req.request_id) + # A parked disk onboard re-enters as a first context chunk, so the + # request may already own a slot; reuse it rather than re-allocating. + slot_id = self.slot_manager.get_slot(req.request_id) + if slot_id is None: + slot_id = self.slot_manager.add_slot(req.request_id) self.slot_ids.append(slot_id) if self.use_relaxed_acceptance_for_thinking: self.relaxed_delta_pool[slot_id].fill_(0) diff --git a/tensorrt_llm/_torch/speculative/mtp.py b/tensorrt_llm/_torch/speculative/mtp.py index 3dc12a669aa6..b71eccc2a953 100644 --- a/tensorrt_llm/_torch/speculative/mtp.py +++ b/tensorrt_llm/_torch/speculative/mtp.py @@ -85,7 +85,11 @@ def prepare_resources(self, scheduled_batch: ScheduledRequests): # allocate hidden state tensors for req in context_batch: if req.is_first_context_chunk: - slot_id = self.slot_manager.add_slot(req.request_id) + # A parked disk onboard re-enters as a first context chunk, so the + # request may already own a slot; reuse it rather than re-allocating. + slot_id = self.slot_manager.get_slot(req.request_id) + if slot_id is None: + slot_id = self.slot_manager.add_slot(req.request_id) if self.use_relaxed_acceptance_for_thinking: self.mtp_relaxed_delta_pool[slot_id].copy_( 0, non_blocking=True) From cebfeb5e6dd7a4fa0d4715a0338fb1c878ea20d2 Mon Sep 17 00:00:00 2001 From: Md Nafis Ul Haque Shifat Date: Wed, 5 Aug 2026 12:06:13 -0700 Subject: [PATCH 46/47] [#16251][fix] sync retention clock in disagg-gen path; annotate are_blocks_ready The disagg-gen-init path allocates and commits KV before the main-loop retention-clock sync, so it would anchor deadlines against a stale clock; sync there too. Also annotate are_blocks_ready's request parameter and document set_retention_clock. Signed-off-by: Md Nafis Ul Haque Shifat --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 3 +++ tensorrt_llm/_torch/pyexecutor/resource_manager.py | 8 +++++++- 2 files changed, 10 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 8e64fdddc11e..20d580ebd599 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -6057,6 +6057,9 @@ def _pad_attention_dp_dummy_request(self): @nvtx_range("_prepare_disagg_gen_init") def _prepare_disagg_gen_init(self, fitting_disagg_gen_init_requests): + # This path allocates/commits KV before the main-loop sync, so refresh the + # retention clock here too or disagg-gen retention would use a stale value. + self._sync_retention_clock() if fitting_disagg_gen_init_requests: disagg_gen_init_to_prepare = ScheduledRequests() disagg_gen_init_to_prepare.context_requests_last_chunk = fitting_disagg_gen_init_requests diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index 28f6beb18c72..3182c16b331b 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -1521,12 +1521,18 @@ def _pack_beam_cache_indices(beams: List[List[int]]) -> List[int]: packed.append(beam[-1]) return packed - def are_blocks_ready(self, request) -> bool: + def are_blocks_ready(self, request: LlmRequest) -> bool: """True when every KV block this request holds has its disk read landed (detached onboard). Always True for requests with no in-flight disk read, so it is cheap for the common case.""" return self.impl.are_blocks_ready(request.py_request_id) def set_retention_clock(self, now_ns: int) -> None: + """Set the rank-consistent clock used for disk-tier retention decisions. + + Args: + now_ns: Timestamp in nanoseconds, broadcast from the leader so every + rank anchors and evaluates retention against the same value. + """ self.impl.set_retention_clock(now_ns) def get_num_free_blocks(self) -> int: From d9802333e0f46bd3c3b3284c7e74faf90115b5d6 Mon Sep 17 00:00:00 2001 From: Md Nafis Ul Haque Shifat Date: Wed, 12 Aug 2026 14:15:50 -0700 Subject: [PATCH 47/47] [#16251][fix] disk-tier: use getUInt64Env for env knobs; test cleanup --- .../batch_manager/kvCacheTransferManager.h | 23 ++++++------------- .../batch_manager/kvCacheManagerTest.cpp | 11 --------- tests/unittest/llmapi/test_disk_tier_args.py | 3 ++- 3 files changed, 9 insertions(+), 28 deletions(-) diff --git a/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h b/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h index 73279defc6e6..87a4e32853ae 100644 --- a/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h +++ b/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h @@ -15,10 +15,12 @@ */ #include "tensorrt_llm/batch_manager/kvCacheManager.h" +#include "tensorrt_llm/common/envUtils.h" #include "tensorrt_llm/executor/types.h" #include "tensorrt_llm/runtime/bufferManager.h" #include "tensorrt_llm/runtime/cudaEvent.h" +#include #include #include #include @@ -31,6 +33,7 @@ namespace tr = tensorrt_llm::runtime; namespace kvc = tensorrt_llm::executor::kv_cache; +namespace tc = tensorrt_llm::common; #pragma once @@ -222,18 +225,10 @@ class KVCacheTransferManager std::unordered_map mSpillRemaining; std::vector mCompletedSpills; bool mDiskWriterStop{false}; - std::size_t const mDiskWriteQueueMax{[] - { - auto* e = std::getenv("TLLM_KV_DISK_WRITE_QUEUE"); - return e ? std::stoul(e) : 1024UL; - }()}; + std::size_t const mDiskWriteQueueMax{ + std::max(1, tc::getUInt64Env("TLLM_KV_DISK_WRITE_QUEUE").value_or(1024))}; // Number of background writer threads draining the shared queue (TLLM_KV_DISK_WRITERS, default 1). - std::size_t const mNumDiskWriters{[] - { - auto* e = std::getenv("TLLM_KV_DISK_WRITERS"); - auto n = e ? std::stoul(e) : 1UL; - return n < 1UL ? 1UL : n; - }()}; + std::size_t const mNumDiskWriters{std::max(1, tc::getUInt64Env("TLLM_KV_DISK_WRITERS").value_or(1))}; void diskWriterLoop(); void enqueueDiskWrite(std::string filename, void const* src, std::size_t bytes); @@ -255,11 +250,7 @@ class KVCacheTransferManager std::vector bytes; // byte count, per pool }; - std::size_t const mNumDiskReaders{[] - { - auto* e = std::getenv("TLLM_KV_DISK_READERS"); - return e ? std::stoul(e) : 0UL; - }()}; + std::size_t const mNumDiskReaders{tc::getUInt64Env("TLLM_KV_DISK_READERS").value_or(0)}; std::vector mDiskReaders; std::mutex mReadMutex; std::condition_variable mReadQueueCv; diff --git a/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp b/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp index b31ee982d620..69e713d8814e 100644 --- a/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp @@ -111,16 +111,6 @@ class KVCacheTransferManagerTestAccess { transferManager.enqueueDiskWriteUnstaged(std::move(filename), src, bytes, spillId); } - - // Unblock a producer parked in the enqueue backpressure wait, so the test can join it before teardown. - static void requestWriterStop(tbk::KVCacheTransferManager& transferManager) - { - { - std::lock_guard lock(transferManager.mDiskMutex); - transferManager.mDiskWriterStop = true; - } - transferManager.mDiskQueueCv.notify_all(); - } }; } // namespace tensorrt_llm::testing @@ -11756,7 +11746,6 @@ TEST_F(KVCacheManagerTest, DiskTierShutdownWithQueuedWorkTest) pool.secondaryPtr = tr::BufferManager::pinned(tr::ITensor::makeShape({2 * nSlots, blockSize}), tensorrt_llm::DataType::kFLOAT); - // Seed host slots and write them to disk (synchronous) so the queued reads below have files to read. float* secBase = tr::bufferCast(*pool.secondaryPtr); for (int k = 0; k < nSlots; ++k) { diff --git a/tests/unittest/llmapi/test_disk_tier_args.py b/tests/unittest/llmapi/test_disk_tier_args.py index 67e012a4f785..4920e417ac3e 100644 --- a/tests/unittest/llmapi/test_disk_tier_args.py +++ b/tests/unittest/llmapi/test_disk_tier_args.py @@ -1,4 +1,5 @@ -# tests/unittest/llmapi/test_disk_tier_args.py +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 """Disk-tier plumbing unit tests. pydantic->pybind passthrough, retention pickle, and the serve-layer TTL helper.