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 4ce5fa94b131..6857b8d055b8 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 @@ -67,6 +68,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 +438,104 @@ 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; + } + + [[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. + void swapDiskResidency(std::shared_ptr const& other) + { + std::swap(mMemoryPoolBlockIndex, other->mMemoryPoolBlockIndex); + 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(std::chrono::steady_clock::time_point::duration now) const + { + return mRetentionExpiry.has_value() && *mRetentionExpiry > now; + } + + [[nodiscard]] std::optional getRetentionExpiry() const + { + return mRetentionExpiry; + } + + void clearRetention() + { + 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 \p now (called when the block is committed / made reusable). + void commitRetention(std::chrono::steady_clock::time_point::duration now) + { + if (mPendingRetentionDuration) + { + markRetained(now + *mPendingRetentionDuration); + mPendingRetentionDuration = 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(); @@ -590,6 +690,23 @@ 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}; + + // (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; + + // 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}; }; class GenerationRequest @@ -738,6 +855,11 @@ class GenerationRequest return mKvCacheRetentionConfig.getTransferMode(); } + [[nodiscard]] std::optional getDiskRetentionMs() const + { + return mKvCacheRetentionConfig.getDiskRetentionMs(); + } + [[nodiscard]] std::string const& getDirectory() const { return mKvCacheRetentionConfig.getDirectory(); @@ -875,8 +997,10 @@ class WindowBlockManager explicit WindowBlockManager(tensorrt_llm::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, + 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, @@ -1208,11 +1332,19 @@ class WindowBlockManager void setOffsets(kernels::KVCacheIndex* offsetsPtr, tensorrt_llm::Dims const& offsetsShape, SizeType32 beamIdx, SizeType32 blockIdx, KVCacheBlock::IdType blockId) const; + //! \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 = ""); + //! \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, @@ -1391,6 +1523,104 @@ class WindowBlockManager // Number of blocks in pools SizeType32 mNumPrimaryBlocks; SizeType32 mNumSecondaryBlocks; + SizeType32 mNumDiskBlocks; + std::string mDiskCachePath; + bool mDiskRetainedOnly{false}; + 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}; + +public: + [[nodiscard]] std::size_t getNumDiskSpills() const + { + return mDiskSpills; + } + + [[nodiscard]] std::size_t getNumDiskGateDropped() const + { + return mDiskGateDropped; + } + + [[nodiscard]] std::size_t getNumDiskAdmissionRefused() const + { + return mDiskAdmissionRefused; + } + + [[nodiscard]] std::size_t getNumDiskWritePressureDropped() const + { + return mDiskWritePressureDropped; + } + + [[nodiscard]] std::size_t getNumDiskOnboards() const + { + return mDiskOnboards; + } + + [[nodiscard]] std::size_t getDiskDeadlineCount() const + { + 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 + //! 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). 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 + { + if (expiry != other.expiry) + { + return expiry < other.expiry; + } + return seq < other.seq; // first-spilled first among equal deadlines + } + }; + + 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 + // 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; + // 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; + // 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(); + + //! \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 + //! unmarked content, then the retained block with the earliest deadline. + [[nodiscard]] BlockPtr claimDiskTarget(); // List of allocated blocks for each sequences std::unordered_map> mAllocatedBlocksPerSeq; @@ -1531,7 +1761,8 @@ class BlockManager SizeType32 indexerKCacheQuantBlockSize = 128, SizeType32 indexerKCacheIndexHeadDim = 0, bool indexerKCacheUseFp4 = false, std::optional> const& indexerKCacheLayerMask = std::nullopt, std::optional linearAttentionMetadata = std::nullopt, - std::vector const& poolConfigurations = {}); + std::vector const& poolConfigurations = {}, SizeType32 blocksInDiskPool = 0, + std::string const& diskCachePath = {}, bool diskRetainedOnly = false, bool diskProtectUnexpired = false); [[nodiscard]] bool isEnableIndexerKCache() const { @@ -1617,6 +1848,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, @@ -2000,6 +2234,77 @@ 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 getNumDiskAdmissionRefused() const + { + std::size_t total = 0; + for (auto const& [windowSize, manager] : mWindowBlockManagers) + { + total += manager.getNumDiskAdmissionRefused(); + } + 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; + for (auto const& [windowSize, manager] : mWindowBlockManagers) + { + total += manager.getNumDiskOnboards(); + } + return total; + } + + [[nodiscard]] std::size_t getDiskDeadlineCount() const + { + std::size_t total = 0; + for (auto const& [windowSize, manager] : mWindowBlockManagers) + { + total += manager.getDiskDeadlineCount(); + } + 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; std::vector mLayerToWindowSize; @@ -2346,7 +2651,8 @@ class KVCacheManager : public BaseKVCacheManager SizeType32 indexerKCacheIndexHeadDim = 0, bool indexerKCacheUseFp4 = false, std::optional> const& indexerKCacheLayerMask = std::nullopt, std::optional linearAttentionMetadata = std::nullopt, - std::vector const& poolConfigurations = {}); + std::vector const& poolConfigurations = {}, SizeType32 blocksInDiskPool = 0, + 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, @@ -2361,7 +2667,8 @@ class KVCacheManager : public BaseKVCacheManager SizeType32 indexerKCacheIndexHeadDim = 0, bool indexerKCacheUseFp4 = false, std::optional> const& indexerKCacheLayerMask = std::nullopt, std::optional linearAttentionMetadata = std::nullopt, - std::vector const& poolConfigurations = {}); + std::vector const& poolConfigurations = {}, SizeType32 blocksInDiskPool = 0, + 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, @@ -2376,7 +2683,8 @@ class KVCacheManager : public BaseKVCacheManager SizeType32 indexerKCacheIndexHeadDim = 0, bool indexerKCacheUseFp4 = false, std::optional> const& indexerKCacheLayerMask = std::nullopt, std::optional linearAttentionMetadata = std::nullopt, - std::vector const& poolConfigurations = {}); + std::vector const& poolConfigurations = {}, SizeType32 blocksInDiskPool = 0, + 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, @@ -2387,7 +2695,8 @@ class KVCacheManager : public BaseKVCacheManager SizeType32 indexerKCacheIndexHeadDim = 0, bool indexerKCacheUseFp4 = false, std::optional> const& indexerKCacheLayerMask = std::nullopt, std::optional linearAttentionMetadata = std::nullopt, - std::vector const& poolConfigurations = {}); + std::vector const& poolConfigurations = {}, SizeType32 blocksInDiskPool = 0, + std::string const& diskCachePath = {}, bool diskRetainedOnly = false, bool diskProtectUnexpired = false); ~KVCacheManager() override = default; @@ -2417,6 +2726,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(); @@ -2702,6 +3017,12 @@ class KVCacheManager : public BaseKVCacheManager mBlockManager.refreshBlocks(); } + void setRetentionClock(std::int64_t timestampNs) + { + mBlockManager.setRetentionClock(std::chrono::duration_cast( + std::chrono::nanoseconds(timestampNs))); + } + void flushIterationEvents() override { mBlockManager.flushIterationEvents(); diff --git a/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h b/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h index 9e9a07b19d8e..87a4e32853ae 100644 --- a/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h +++ b/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h @@ -15,12 +15,25 @@ */ #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 +#include +#include +#include +#include +#include +#include + namespace tr = tensorrt_llm::runtime; namespace kvc = tensorrt_llm::executor::kv_cache; +namespace tc = tensorrt_llm::common; #pragma once @@ -72,11 +85,71 @@ class KVCacheTransferManager //! KVCacheManager::addSequenceBatch in every step. void syncWithBufferManager(); + //! \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 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::int32_t trackBlockId = -1); + + //! \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 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 + { + 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(). + 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; + } + + //! \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(); + //! \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(); + //! \brief Get transfer stats accumulated since last call, and reset the counters. [[nodiscard]] KvCacheTransferStats getAndResetTransferStats(); @@ -124,6 +197,79 @@ 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; // 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::vector mDiskWriters; + std::mutex mDiskMutex; + std::condition_variable mDiskQueueCv; + 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{ + 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{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); + 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 ---- + // 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 + { + 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{tc::getUInt64Env("TLLM_KV_DISK_READERS").value_or(0)}; + 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; + // 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. + bool mDiskWorkersShutdown{false}; + + void diskReaderLoop(); + void enqueueDiskRead(DiskReadJob job); // Cumulative transfer statistics, reset on each call to getAndResetTransferStats(). // Protected by mStatsMutex for thread-safe access. diff --git a/cpp/include/tensorrt_llm/executor/executor.h b/cpp/include/tensorrt_llm/executor/executor.h index acc0efe18966..ca488c550251 100644 --- a/cpp/include/tensorrt_llm/executor/executor.h +++ b/cpp/include/tensorrt_llm/executor/executor.h @@ -636,6 +636,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. @@ -647,7 +649,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: @@ -663,6 +665,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 @@ -1100,6 +1106,14 @@ 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); + [[nodiscard]] bool getDiskCacheRetainedOnly() const; + void setDiskCacheRetainedOnly(bool diskCacheRetainedOnly); + [[nodiscard]] bool getDiskCacheProtectUnexpired() const; + void setDiskCacheProtectUnexpired(bool diskCacheProtectUnexpired); void fillEmptyFieldsFromRuntimeDefaults(tensorrt_llm::runtime::RuntimeDefaults const& runtimeDefaults); @@ -1159,6 +1173,19 @@ 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 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/evictionPolicy.cpp b/cpp/tensorrt_llm/batch_manager/evictionPolicy.cpp index 42b5bd9e6879..2e10a9a7dbd2 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,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 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). @@ -67,17 +78,26 @@ 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 < sizes[cacheLevel]; blockId++) + 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]); } - mNumFreeBlocksPerLevel[cacheLevel] = sizes[cacheLevel]; - startIdx += sizes[cacheLevel]; + mNumFreeBlocksPerLevel[cacheLevel] = sizesPadded[cacheLevel]; + startIdx += sizesPadded[cacheLevel]; } mSecondaryOffloadMinPriority = secondaryOffloadMinPriority.value_or(kDefaultSecondaryOffloadMinPriority); @@ -105,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++) @@ -160,6 +181,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 18a8d83e828c..738cd8d4b24e 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp @@ -595,7 +595,8 @@ BlockManager::BlockManager(std::vector const& numKvHeadsPerLayer, Si SizeType32 indexerKCacheIndexHeadDim, bool indexerKCacheUseFp4, std::optional> const& indexerKCacheLayerMask, std::optional linearAttentionMetadata, - std::vector const& poolConfigurations) + std::vector const& poolConfigurations, SizeType32 blocksInDiskPool, + std::string const& diskCachePath, bool diskRetainedOnly, bool diskProtectUnexpired) : mNumLayers{static_cast(numKvHeadsPerLayer.size())} , mTokensPerBlock{tokensPerBlock} , mEventManager{std::move(eventManager)} @@ -698,10 +699,12 @@ 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, - enablePartialReuse, copyOnPartialReuse, kvCacheConnectorManager, mLookupTree, mLoopbackAgent, - enableIndexerKCache, indexerKCacheQuantBlockSize, indexerKCacheIndexHeadDim, indexerKCacheUseFp4, - indexerKCacheLayerMask, + allottedSecondaryBlocks, + (SizeType32(windowSize) == blocksPerWindow.rbegin()->first) ? blocksInDiskPool : SizeType32{0}, + diskCachePath, diskRetainedOnly, diskProtectUnexpired, maxNumSequences, stream, cacheType, + secondaryOffloadMinPriority, mEventManager, enablePartialReuse, copyOnPartialReuse, kvCacheConnectorManager, + mLookupTree, mLoopbackAgent, enableIndexerKCache, indexerKCacheQuantBlockSize, indexerKCacheIndexHeadDim, + indexerKCacheUseFp4, indexerKCacheLayerMask, LinearAttentionMetadata::hasLinearCache(windowSize) ? linearAttentionMetadata : std::nullopt, numPlaceholderBlocks); } @@ -751,8 +754,10 @@ BlockManager::BlockManager(std::vector const& numKvHeadsPerLayer, Si WindowBlockManager::WindowBlockManager(tensorrt_llm::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, + 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, @@ -763,6 +768,10 @@ WindowBlockManager::WindowBlockManager(tensorrt_llm::DataType dtype, SizeType32 , mWindowSize{windowSize} , mNumPrimaryBlocks{blocksInPrimaryPool} , mNumSecondaryBlocks{blocksInSecondaryPool} + , mNumDiskBlocks{blocksInDiskPool} + , mDiskCachePath{diskCachePath} + , mDiskRetainedOnly{diskRetainedOnly} + , mDiskProtectUnexpired{diskProtectUnexpired} , mBufferManager{std::move(stream)} , mSchedulingNumFreeBlocks{0} , mTokensPerBlock{tokensPerBlock} @@ -861,7 +870,7 @@ WindowBlockManager::WindowBlockManager(tensorrt_llm::DataType dtype, SizeType32 } // 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})); @@ -871,6 +880,22 @@ WindowBlockManager::WindowBlockManager(tensorrt_llm::DataType dtype, SizeType32 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); + // 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) + { + 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 @@ -897,15 +922,16 @@ WindowBlockManager::WindowBlockManager(tensorrt_llm::DataType dtype, SizeType32 } 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) { @@ -1216,6 +1242,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() @@ -1230,6 +1281,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(); @@ -1314,6 +1368,247 @@ 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); + } + // 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); + } +} + +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 + // 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. + auto candidate = std::get<0>(mEvictionPolicy->getFreeBlock(kDiskLevel)); + if (!candidate->isRetainedNow(mRetentionNow)) + { + mEvictionPolicy->claimBlock(candidate); + return candidate; + } + // 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). 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(mRetentionNow)) + { + 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(mRetentionNow)) + { + return nullptr; + } + mEvictionPolicy->claimBlock(candidate); + return candidate; +} + +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); + // 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(); + } + 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 + } + if (mDiskRetainedOnly && !victim->isRetainedNow(mRetentionNow)) + { + ++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, 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. + 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 + } + // 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(mRetentionNow) && 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: 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 would. + ++mDiskAdmissionRefused; + if (mDiskAdmissionRefused == 1 || mDiskAdmissionRefused % 10000 == 0) + { + TLLM_LOG_INFO("[disk-tier] admission refused=%zu (windowSize=%d)", mDiskAdmissionRefused, mWindowSize); + } + return victim; + } + 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 the eviction policy's expiring-block machinery + 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. + victim->setPriority(executor::KvCacheRetentionConfig::kMaxRetentionPriority); + auto const seq = ++mDiskSpillSeq; + victim->setDiskDeadlineSeq(seq); + victim->setDiskDeadlineExpiry(*victim->getRetentionExpiry()); + mDiskDeadlines.insert(DiskDeadline{victim->getDiskDeadlineExpiry(), seq, victim}); + } + else + { + // 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; + if (mDiskSpills == 1 || mDiskSpills % 1000 == 0) + { + TLLM_LOG_INFO("[disk-tier] spills=%zu (windowSize=%d)", mDiskSpills, mWindowSize); + } + // 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. + { + std::lock_guard treeLock(mLookupTree->getMutex()); + if (mEventManager && blockInRadixTree(diskTarget)) + { + mEventManager->enqueueRemovedEvent(diskTarget, mWindowSize); + } + 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; + } + return diskTarget; // owns the victim's old host slot; already claimed, identity detached above +} + BlockPtr WindowBlockManager::getFreeBlock(GenerationRequest& sequence, executor::RetentionPriority priority, std::optional durationMs, executor::KvCacheTransferMode mode, std::string const& directory, bool wantPlaceholder) @@ -1334,7 +1629,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, @@ -1380,9 +1675,15 @@ 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()) + { + // 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()); @@ -1439,10 +1740,58 @@ 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) { std::lock_guard lock(mLookupTree->getMutex()); + 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); + // 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); + eraseDiskDeadline(offloadBlock); + 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 + 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) + { + TLLM_LOG_INFO("[disk-tier] onboards=%zu (windowSize=%d)", mDiskOnboards, mWindowSize); + } + return; + } if (!offloadBlock->isPlaceholder() && !offloadBlock->isPrimary()) { auto block = getFreeBlock( @@ -1462,6 +1811,41 @@ 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. + 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); + 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) { @@ -1481,7 +1865,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); @@ -1584,6 +1968,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) @@ -1613,6 +2004,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()); @@ -1914,10 +2312,17 @@ 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(mRetentionNow + *diskRetentionMs); + } onboardBlock(sequence, claimed.block, claimResult.mode, claimResult.directory); addBlockToAllBeams(claimed.block, sequence); if (!claimed.isPlaceholder) { + claimed.block->incReuseCount(); // per-block reuse count for the disk reuse-gate ++mReusedBlocks; if (claimed.isPartialMatch) { @@ -2794,6 +3199,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(mRetentionNow); + } + if (mEventManager) { // Linear-attention placeholders can legitimately land in storedBlocks via the @@ -3281,6 +3692,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) @@ -3296,16 +3710,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 @@ -3349,13 +3774,15 @@ KVCacheManager::KVCacheManager(SizeType32 numLayers, SizeType32 numKvHeads, Size bool enableIndexerKCache, SizeType32 indexerKCacheQuantBlockSize, SizeType32 indexerKCacheIndexHeadDim, bool indexerKCacheUseFp4, std::optional> const& indexerKCacheLayerMask, std::optional linearAttentionMetadata, - std::vector const& poolConfigurations) + std::vector const& poolConfigurations, SizeType32 blocksInDiskPool, + 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, - indexerKCacheLayerMask, linearAttentionMetadata, poolConfigurations) + indexerKCacheLayerMask, linearAttentionMetadata, poolConfigurations, blocksInDiskPool, diskCachePath, + diskRetainedOnly, diskProtectUnexpired) { } @@ -3369,13 +3796,15 @@ KVCacheManager::KVCacheManager(std::vector const& numKvHeadsPerLayer SizeType32 indexerKCacheQuantBlockSize, SizeType32 indexerKCacheIndexHeadDim, bool indexerKCacheUseFp4, std::optional> const& indexerKCacheLayerMask, std::optional linearAttentionMetadata, - std::vector const& poolConfigurations) + std::vector const& poolConfigurations, SizeType32 blocksInDiskPool, + 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, indexerKCacheLayerMask, linearAttentionMetadata, poolConfigurations) + indexerKCacheUseFp4, indexerKCacheLayerMask, linearAttentionMetadata, poolConfigurations, blocksInDiskPool, + diskCachePath, diskRetainedOnly, diskProtectUnexpired) { } @@ -3389,7 +3818,8 @@ KVCacheManager::KVCacheManager(std::vector const& numKvHeadsPerLayer SizeType32 indexerKCacheQuantBlockSize, SizeType32 indexerKCacheIndexHeadDim, bool indexerKCacheUseFp4, std::optional> const& indexerKCacheLayerMask, std::optional linearAttentionMetadata, - std::vector const& poolConfigurations) + std::vector const& poolConfigurations, SizeType32 blocksInDiskPool, + std::string const& diskCachePath, bool diskRetainedOnly, bool diskProtectUnexpired) : mMaxBeamWidth(maxBeamWidth) , mDataType(dtype) , mMaxAttentionWindow(*std::max_element(maxAttentionWindowVec.begin(), maxAttentionWindowVec.end())) @@ -3402,7 +3832,8 @@ KVCacheManager::KVCacheManager(std::vector const& numKvHeadsPerLayer mChunkSize, cacheType, secondaryOffloadMinPriority, std::move(eventManager), enablePartialReuse, copyOnPartialReuse, std::move(kvCacheConnectorManager), std::nullopt, enableIndexerKCache, indexerKCacheQuantBlockSize, indexerKCacheIndexHeadDim, indexerKCacheUseFp4, indexerKCacheLayerMask, - linearAttentionMetadata, poolConfigurations) + linearAttentionMetadata, 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} { @@ -3432,13 +3863,15 @@ KVCacheManager::KVCacheManager(SizeType32 numLayers, SizeType32 numKvHeads, Size SizeType32 indexerKCacheQuantBlockSize, SizeType32 indexerKCacheIndexHeadDim, bool indexerKCacheUseFp4, std::optional> const& indexerKCacheLayerMask, std::optional linearAttentionMetadata, - std::vector const& poolConfigurations) + std::vector const& poolConfigurations, SizeType32 blocksInDiskPool, + 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, indexerKCacheLayerMask, - linearAttentionMetadata, poolConfigurations) + linearAttentionMetadata, poolConfigurations, blocksInDiskPool, diskCachePath, diskRetainedOnly, + diskProtectUnexpired) { } diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp index 33ad5b4a1675..d08606f2342c 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" @@ -48,7 +49,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 +66,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()); @@ -85,6 +86,69 @@ KVCacheTransferManager::KVCacheTransferManager( { TLLM_CUDA_CHECK(cudaGetDevice(&mDeviceId)); TLLM_CHECK(mDeviceId != -1); + if (mAsyncDiskStore) + { + 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); + } + 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); + } +} + +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()) + { + { + std::lock_guard lock(mDiskMutex); + mDiskWriterStop = true; + } + mDiskQueueCv.notify_all(); + for (auto& t : mDiskWriters) + { + if (t.joinable()) + { + t.join(); + } + } + } + if (!mDiskReaders.empty()) + { + { + std::lock_guard lock(mReadMutex); + mDiskReaderStop = true; + } + mReadQueueCv.notify_all(); + for (auto& t : mDiskReaders) + { + if (t.joinable()) + { + t.join(); + } + } + } +} + +KVCacheTransferManager::~KVCacheTransferManager() +{ + shutdownDiskWorkers(); } tr::ITensor::SharedPtr KVCacheTransferManager::computeBlockPointer( @@ -392,6 +456,469 @@ 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::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. + // 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()); + 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) + { + 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 + } +} + +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 + enqueueDiskWriteCommon(std::move(job)); +} + +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; + enqueueDiskWriteCommon(std::move(job)); +} + +void KVCacheTransferManager::enqueueDiskWriteCommon(DiskWriteJob&& job) +{ + { + std::unique_lock lock(mDiskMutex); + // 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] + { + 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(); +} + +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); + mDiskInflightCv.wait(lock, + [this, &filename] + { + auto it = mDiskInflight.find(filename); + return it == mDiskInflight.end() || it->second <= 0; + }); +} + +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 + // 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)); + // 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 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; + 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(); + } + + // 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) + { + 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) + { + // 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; + try + { + 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; + } + } + + { + 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); + } + } + if (pinnedBuf != nullptr) + { + cudaFreeHost(pinnedBuf); + } + cudaStreamDestroy(stream); +} + +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(); +} + +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) +{ + 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); + 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); + } +} + +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, std::int32_t trackBlockId) +{ + 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(); + } + // 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) + { + 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 (detach) + { + 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). + // 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"); + auto ptr = computeBlockPointer(dstPrimaryBlock, pools, poolIdx); + fileToGpuPosix(ptr, diskSlotFilename(directory, diskSlot, poolIdx)); + } +} + void KVCacheTransferManager::syncWithBufferManager() { tr::CudaEvent readyForOffloadEvent; diff --git a/cpp/tensorrt_llm/executor/kvCacheConfig.cpp b/cpp/tensorrt_llm/executor/kvCacheConfig.cpp index 54ddc199e6d3..24461e79af28 100644 --- a/cpp/tensorrt_llm/executor/kvCacheConfig.cpp +++ b/cpp/tensorrt_llm/executor/kvCacheConfig.cpp @@ -202,6 +202,46 @@ 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; +} + +bool KvCacheConfig::getDiskCacheRetainedOnly() const +{ + return mDiskCacheRetainedOnly; +} + +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/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 eff5c62b92ec..0ccbf11cdabc 100644 --- a/cpp/tensorrt_llm/executor/serialization.cpp +++ b/cpp/tensorrt_llm/executor/serialization.cpp @@ -1361,10 +1361,18 @@ 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); + // retained-only flag is read after construction below - 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); + kvCacheConfig.setDiskCacheRetainedOnly(su::deserialize(is)); + kvCacheConfig.setDiskCacheProtectUnexpired(su::deserialize(is)); + return kvCacheConfig; } void Serialization::serialize(KvCacheConfig const& kvCacheConfig, std::ostream& os) @@ -1382,6 +1390,10 @@ 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); + su::serialize(kvCacheConfig.getDiskCacheRetainedOnly(), os); + su::serialize(kvCacheConfig.getDiskCacheProtectUnexpired(), os); } size_t Serialization::serializedSize(KvCacheConfig const& kvCacheConfig) @@ -1401,6 +1413,10 @@ 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()); + totalSize += su::serializedSize(kvCacheConfig.getDiskCacheRetainedOnly()); + totalSize += su::serializedSize(kvCacheConfig.getDiskCacheProtectUnexpired()); return totalSize; } @@ -1799,8 +1815,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) @@ -1810,6 +1829,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) @@ -1820,6 +1840,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 222f2e528cc3..ed3767da0553 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp @@ -678,13 +678,16 @@ 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, bool, tbk::CacheType, std::optional, std::shared_ptr, bool, bool, std::shared_ptr, bool, SizeType32, SizeType32, bool, std::optional> const&, - std::optional, std::vector const&>(), + std::optional, 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"), @@ -696,13 +699,15 @@ void tb::kv_cache_manager::KVCacheManagerBindings::initBindings(nb::module_& m) nb::arg("indexer_k_cache_index_head_dim") = 0, nb::arg("indexer_k_cache_use_fp4") = false, nb::arg("indexer_k_cache_layer_mask").none() = std::nullopt, nb::arg("linear_attention_metadata").none() = std::nullopt, - nb::arg("pool_configurations") = std::vector{}, - nb::call_guard()) + 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::arg("disk_cache_protect_unexpired") = false, nb::call_guard()) .def( "scheduling_has_free_blocks", [](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/tensorrt_llm/nanobind/executor/executorConfig.cpp b/cpp/tensorrt_llm/nanobind/executor/executorConfig.cpp index acd33c0df769..4b82f68ade8d 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(), self.getDiskCacheRetainedOnly(), self.getDiskCacheProtectUnexpired()); }; auto kvCacheConfigSetstate = [](tle::KvCacheConfig& self, nb::tuple const& state) { - if (state.size() != 14) + if (state.size() != 18) { throw std::runtime_error("Invalid state!"); } @@ -128,6 +129,10 @@ 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])); + self.setDiskCacheRetainedOnly(nb::cast(state[16])); + self.setDiskCacheProtectUnexpired(nb::cast(state[17])); }; nb::class_(m, "KvCacheConfig") .def(nb::init const&, std::optional> const&, @@ -152,6 +157,12 @@ 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("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/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/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 764d3236c609..69e713d8814e 100644 --- a/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp @@ -93,6 +93,24 @@ 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); + } }; } // namespace tensorrt_llm::testing @@ -11085,3 +11103,801 @@ TEST_F(KVCacheManagerTest, VswaDisaggDtypeMismatchTriggersGuard) EXPECT_THROW({ CacheTransBufferManager const ctbm(kvCacheManager.get(), /*maxNumTokens=*/std::nullopt); }, tensorrt_llm::common::TllmException); } + +// ============================== Disk tier tests ============================== +// 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 + +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; +}; + +// 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, + SizeType32 blocksInDiskPool, bool protectUnexpired = false) +{ + 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}, tensorrt_llm::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, + /*indexerKCacheLayerMask=*/std::nullopt, /*linearAttentionMetadata=*/std::nullopt, + /*poolConfigurations=*/std::vector{}, blocksInDiskPool, diskPath, retainedOnly, + protectUnexpired); + 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(). +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) +{ + 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); +} + +// 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) +{ + ScopedEnvVar readersEnv{"TLLM_KV_DISK_READERS", "4"}; + 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); +} + +// 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; + + blockManager.setRetentionClock(retentionTick(1000)); + auto seqP = addDiskTierSequence(blockManager, 0, iotaTokens(0, 16), std::chrono::milliseconds(50)); + releaseDiskTierSequence(blockManager, seqP); + blockManager.setRetentionClock(retentionTick(1200)); // TTL dead before eviction + churnOnce(blockManager, 1, 100); + EXPECT_EQ(blockManager.getNumDiskSpills(), 0u); + 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 the retention clock + ttl. + auto const t0 = retentionTick(5000); + blockManager.setRetentionClock(t0); + blockManager.storeContextBlocks(*s.seq, *s.req); + + std::size_t stamped = 0; + for (auto const bid : blockIds) + { + auto const expiry = blockManager.getBlockById(bid, windowSize)->getRetentionExpiry(); + if (expiry.has_value()) + { + ++stamped; + EXPECT_EQ(*expiry, t0 + ttl); // clock started at commit, not earlier at allocation + } + } + 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) +{ + 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. + 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); + + 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); + + // 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); + blockManager.setRetentionClock(retentionTick(1400)); // past seqB's TTL + 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); +} + +// 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; + auto blockManagerPtr = makeDiskTierBlockManager( + stream, dir.str(), /*retainedOnly=*/true, /*prim=*/4, /*sec=*/2, /*disk=*/4, /*protectUnexpired=*/true); + auto& blockManager = *blockManagerPtr; + + // 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; + + 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); + 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)); + releaseDiskTierSequence(blockManager, seqB); + churnOnce(blockManager, 3, 300); + 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) +{ + ScopedEnvVar readersEnv{"TLLM_KV_DISK_READERS", readers}; + 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}), tensorrt_llm::DataType::kFLOAT); + bufferManager.setZero(*pool.primaryPtr); + pool.secondaryPtr + = tr::BufferManager::pinned(tr::ITensor::makeShape({nSlots, blockSize}), tensorrt_llm::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"; + } +} +} // 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); +} + +// 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; 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"}; + ScopedEnvVar writersEnv{"TLLM_KV_DISK_WRITERS", "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}), tensorrt_llm::DataType::kFLOAT); + pool.secondaryPtr + = tr::BufferManager::pinned(tr::ITensor::makeShape({2, blockSize}), tensorrt_llm::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"; +} + +// 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. +TEST_F(KVCacheManagerTest, DiskTierOnboardReadFailureSurfacesTest) +{ + ScopedEnvVar readersEnv{"TLLM_KV_DISK_READERS", "2"}; + 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); +} + +// 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) +{ + 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; + + 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}), tensorrt_llm::DataType::kFLOAT); + bufferManager.setZero(*pool.primaryPtr); + pool.secondaryPtr + = tr::BufferManager::pinned(tr::ITensor::makeShape({2 * nSlots, blockSize}), tensorrt_llm::DataType::kFLOAT); + + 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. +} + +// 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; + 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 + + 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"; +} +} // 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); +} + +// 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 +} + +// 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 bcd1363138d7..e22bd9947163 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -616,6 +616,16 @@ 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._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 @@ -2773,7 +2783,9 @@ 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) # The generation requests that do not have batch_idx # need to be in front of the batch due to the assumptions @@ -3349,6 +3361,97 @@ 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 + 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 + 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.) + # 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) + 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 = [] + parked = 0 + for req, ready in zip(ctx, global_ready): + if ready: + 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( + 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: """Handle dynamic draft length for the current batch. @@ -4229,7 +4332,9 @@ 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) if self.kv_connector_manager: self.kv_connector_manager.handle_metadata() @@ -4237,8 +4342,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: @@ -4772,7 +4879,9 @@ 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) if self.kv_connector_manager: self.kv_connector_manager.handle_metadata() @@ -4780,8 +4889,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) @@ -6633,6 +6744,9 @@ def _pad_empty_attention_dp_batch( @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 827bd5e534b7..9c7fdbf1c3f1 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -685,6 +685,49 @@ 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 + if disk_cache_size and (mapping.enable_attention_dp + or mapping.pp_size > 1 or mapping.cp_size > 1): + raise ValueError( + "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.") + 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( + ) + 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 + disk_cache_path = getattr(kv_cache_config, "disk_cache_path", + None) or "" + if disk_cache_path and blocks_in_disk_pool: + # 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}" + + ("_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} " + 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( + 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 @@ -980,6 +1023,12 @@ def maybe_fit_token_budget(self, 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 req.py_request_id in self._preprepared_dummy_request_ids: return None if is_cross: @@ -1732,6 +1781,20 @@ def _pack_beam_cache_indices(beams: List[List[int]]) -> List[int]: packed.append(beam[-1]) return packed + 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: if self.is_linear_attention: value = self.impl.get_kv_cache_stats( diff --git a/tensorrt_llm/_torch/speculative/eagle3.py b/tensorrt_llm/_torch/speculative/eagle3.py index d497862e0c8e..dd6cb6362339 100644 --- a/tensorrt_llm/_torch/speculative/eagle3.py +++ b/tensorrt_llm/_torch/speculative/eagle3.py @@ -111,7 +111,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 5496ef78ce1d..8c1fef3d25d9 100644 --- a/tensorrt_llm/_torch/speculative/mtp.py +++ b/tensorrt_llm/_torch/speculative/mtp.py @@ -75,7 +75,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) diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index bf1afd2b97f7..aa52ad22e478 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -3804,13 +3804,23 @@ 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, 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." + ) + 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= @@ -4022,6 +4032,13 @@ 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 "" + 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') diff --git a/tensorrt_llm/serve/openai_protocol.py b/tensorrt_llm/serve/openai_protocol.py index ef5fd54cc875..88afd7026e0e 100644 --- a/tensorrt_llm/serve/openai_protocol.py +++ b/tensorrt_llm/serve/openai_protocol.py @@ -550,6 +550,12 @@ 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, + ge=0, + description= + "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 @@ -910,6 +916,12 @@ 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, + ge=0, + description= + "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/tensorrt_llm/serve/openai_server.py b/tensorrt_llm/serve/openai_server.py index 1c38c3c90634..9f9b58927978 100644 --- a/tensorrt_llm/serve/openai_server.py +++ b/tensorrt_llm/serve/openai_server.py @@ -2,6 +2,7 @@ import array import asyncio import base64 +import datetime as _dt import functools import json import os @@ -49,6 +50,7 @@ from tensorrt_llm.llmapi.disagg_utils import (DisaggClusterConfig, MetadataServerConfig, ServerRole) from tensorrt_llm.llmapi.llm import LLM, RequestOutput +from tensorrt_llm.llmapi.llm_utils import KvCacheRetentionConfig as _KvRetention from tensorrt_llm.llmapi.reasoning_parser import ReasoningParserFactory from tensorrt_llm.llmapi.thinking_budget import \ add_thinking_budget_logits_processor @@ -128,6 +130,20 @@ def _is_visual_gen_instance(obj) -> bool: visual_gen = sys.modules.get("tensorrt_llm.visual_gen.visual_gen") return visual_gen is not None and isinstance(obj, visual_gen.VisualGen) + +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 # msgspec msgpack is an opt-in transport for the disagg orchestrator->worker @@ -1707,6 +1723,7 @@ async def chat_stream_generator( promise = self.generator.generate_async( inputs=generate_inputs, + kv_cache_retention_config=_disk_retention_config(request), sampling_params=sampling_params, _postproc_params=postproc_params if self.postproc_worker_enabled else None, @@ -1852,7 +1869,10 @@ async def create_mm_embedding_response(promise: RequestOutput): if mm_item_order: prompt["mm_item_order"] = mm_item_order - promise = self.generator.generate_async(inputs=prompt, ) + # 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) @@ -2044,6 +2064,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, @@ -2186,6 +2207,7 @@ async def create_streaming_generator(promise: RequestOutput, # Generate promise = self.generator.generate_async( inputs=harmony_tokens, + kv_cache_retention_config=_disk_retention_config(request), sampling_params=sampling_params, _postproc_params=postproc_params if self.postproc_worker_enabled else None, @@ -2343,6 +2365,9 @@ 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, _postproc_params=postproc_params 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..4920e417ac3e --- /dev/null +++ b/tests/unittest/llmapi/test_disk_tier_args.py @@ -0,0 +1,113 @@ +# 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. +""" + +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(): + cfg = tllm.KvCacheConfig() + 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(): + 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, + 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(): + 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 + + +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 + )