[#16251][feat] KV cache manager v1: add disk (L3) tier with per-request retention TTL - #16252
Open
nafis271 wants to merge 55 commits into
Open
[#16251][feat] KV cache manager v1: add disk (L3) tier with per-request retention TTL#16252nafis271 wants to merge 55 commits into
nafis271 wants to merge 55 commits into
Conversation
…sk blocks, host->disk spill and disk->GPU onboard Based on v1.3.0rc19. Adds a third KV-cache tier backed by POSIX slot files: - KvCacheConfig gains diskCacheSize/diskCachePath (serialization + nanobind + llm_args passthrough) - Poolless disk blocks at eviction level 2 (placeholders renumber to 3); block residency swaps between pool index and disk slot so the reuse-tree identity never moves - Host-tier victims with reusable content spill to disk on reclaim (spill-all; filtering mode to follow); disk-resident matches onboard file->GPU in onboardBlock - Synchronous POSIX I/O v1; pytorch path only (TRT-engine path passes 0) Validated on gpt-oss-120b: 49824/49836 cached tokens served through a disk round trip with temperature-0 byte-identical output; 300k+ spills and 11k+ onboards under sustained production traffic with zero I/O errors. Signed-off-by: Md Nafis Ul Haque Shifat <nafis@deepinfra.com>
…retention TTL, serve wiring Signed-off-by: Md Nafis Ul Haque Shifat <nafis@deepinfra.com>
… displacement (empties first, exact TTL order, suffix-first FIFO), retention wiring at all serve entry points Signed-off-by: Md Nafis Ul Haque Shifat <nafis@deepinfra.com>
…aking), unit tests (C++ disk-tier suite + isEnqueued + python plumbing), disk-count getters Signed-off-by: Md Nafis Ul Haque Shifat <nafis@deepinfra.com>
…ate/expiry/reuse/deadline + isEnqueued + placeholder level fix) Signed-off-by: Md Nafis Ul Haque Shifat <nafis@deepinfra.com>
…use new blocks when disk full of live TTLs; only expired evictable) + unit tests Signed-off-by: Md Nafis Ul Haque Shifat <nafis@deepinfra.com>
…empty slots before displacing cached blocks) Empty disk slots initialized at default priority (bucket 35) while unmarked spilled content was filed at kMin (bucket 0). getFreeBlock returns the lowest bucket first, so claimDiskTarget displaced just-spilled cached blocks before using free empty slots -> unmarked content never survived to re-request (spill-all onboard ~0% under churn). Restore order empty(kMin) < unmarked(kDefault) < retained(kMax): - evictionPolicy initialize(): disk-level blocks start at kMin (+ setPriority in lockstep so claim/release erase from the correct queue). - reclaimSecondaryBlock: unmarked spilled content -> kDefault (above empties, below retained); evicted LRU only once disk is genuinely full. retained path unchanged. Verified: spill-all onboard 0%->99%, disk-tier unit tests pass, gate+expiry+retained-survival unaffected. Signed-off-by: Md Nafis Ul Haque Shifat <nafis@deepinfra.com>
…served-pool unstaged) Working per-slot async fire-and-forget store: background writer thread + per-slot files + staged memcpy + load-wait. GDS path still present. Snapshot taken before the unstaged reserved-block-pool rework so we can roll back cleanly. Signed-off-by: Md Nafis Ul Haque Shifat <nafis@deepinfra.com>
…in source, hand out reserved slot) Async spill hands the writer the victim host slot directly (no staging memcpy), pins that slot until the write drains, and returns a pre-claimed reserved host block instead. Reaped drained slots rejoin the pool: one out, one in. Gated by TLLM_KV_DISK_RESERVED_BLOCKS (0=off=staged). Signed-off-by: Md Nafis Ul Haque Shifat <nafis@deepinfra.com>
… vs enqueue-backpressure vs queue depth) Signed-off-by: Md Nafis Ul Haque Shifat <nafis@deepinfra.com>
…ared queue (TLLM_KV_DISK_WRITERS) Single writer saturated the queue under eviction bursts (qpeak pinned at 1024, multi-second enqueue backpressure -> TTFT spikes). Spawn N workers draining the shared queue; idle workers pull the next job (auto-balanced), per-file in-flight gate keeps same-slot writes serialized. Default 1. Signed-off-by: Md Nafis Ul Haque Shifat <nafis@deepinfra.com>
Signed-off-by: Md Nafis Ul Haque Shifat <nafis@deepinfra.com>
… scheduler thread Disk onboards read SSD->GPU synchronously on the scheduler thread (or, with a reader pool, in parallel but still barriered before the forward), stalling every request in the step behind one slow read. Engine (C++): a disk onboard hands each of a block's pool reads to the reader pool and returns; completion is tracked per block, keyed by the matched-identity block id so requests reusing a prefix gate on the same key. The reader makes its copy device-complete before publishing, and GDS/POSIX both flow through the pool. areBlocksReady(requestId) reports when a request's blocks have landed; the read barrier (waitForAllReads/mReadInflight) is removed. Scheduler (Python): a context request whose onboard reads are still in flight is held out of the forward pass and re-checked each step -- the capacity scheduler counts its already-allocated blocks and _context_seq_len skips re-adding it -- then admitted once areBlocksReady(). Tests: detached reads byte-exact incl. concurrent draining; block-manager reuse matches the synchronous path and its readiness gate resolves. Signed-off-by: Md Nafis Ul Haque Shifat <nafis@deepinfra.com>
…_KV_DISK_MIN_REUSE) Add a per-block reuse counter (mReuseCount, bumped at the prefix-match reuse point in onboardAndAllocateBlocks) and gate the host->disk spill on it: only spill blocks reused >= k times (k = TLLM_KV_DISK_MIN_REUSE env, 0 = off). Sits beside the retained_only gate as an independent 3rd disk mode (spill-all / retention-only / reuse-gate). The counter is block metadata, so it survives GPU/host/disk tier moves like mRetentionExpiry (swapDiskResidency swaps only residency pointers). Signed-off-by: Md Nafis Ul Haque Shifat <nafis@deepinfra.com>
…k dir + cross-rank onboard readiness) Two fixes to run the disk KV-cache tier under TP>1 (validated on gpt-oss-120b TP=2). Per-rank disk cache directory: all ranks share one container filesystem and one disk_cache_path, but the disk slot files (block_<id>_pool_<p>.bin) are named identically across ranks while each rank holds a different KV shard -- so without namespacing the ranks clobber each other's KV on load. Append /rank_<mapping.rank> (the global rank, unique across TP and PP) and mkdir it. Note: disk_cache_size is per-rank, so total disk = world_size * disk_cache_size. Cross-rank onboard-readiness sync: the detached disk-onboard reads land at slightly different times per rank, so a per-rank park decision in _park_requests_awaiting_onboard would give ranks different batches and the TP all-reduce in the forward would mismatch (hang/crash). AND the per-request readiness across the TP group -- a cheap tp_allreduce "any pending?" fast path, then tp_allgather when a read is in flight -- so a request is admitted only once its KV has landed on every rank. pp_size==1 => the TP group is the whole world; TP=1 is a no-op (identical to the prior path). Signed-off-by: Md Nafis Ul Haque Shifat <nafis@deepinfra.com>
… read is in flight areBlocksReady scans every block a request holds via isBlockReadPending, which takes mReadMutex per block, and the scheduler-side park calls it for every context request every step. With no disk read in flight (the common case) that is an O(context-length) mutex-locked scan over blocks that are all trivially ready. Add a lock-free atomic mirror of the in-flight read count (mReadInflightCount, updated under mReadMutex on every pending-set insert/erase) and short-circuit areBlocksReady to return true immediately when it is zero. Behaviour is identical when reads are pending. Signed-off-by: Md Nafis Ul Haque Shifat <nafis@deepinfra.com>
…s CUDA device Reader threads created their CUDA stream without setting the device, so on TP ranks whose device is not 0 the POSIX H2D onboard targeted the wrong GPU and crashed. Bind the thread to the manager's device at loop entry. Signed-off-by: Md Nafis Ul Haque Shifat <nafis@deepinfra.com>
…validators verifyQueueIntegrity had three validators/labels for four queue levels: the disk level reused the placeholder validator and the loop indexed one past the array. Give each level its own validator and label (disk = isOnDisk). Signed-off-by: Md Nafis Ul Haque Shifat <nafis@deepinfra.com>
…blocks A disk-resident block cannot serve a partial reuse -- its bytes live in a reserved disk slot, not in pool memory, so the partial-copy path reads garbage. Exclude on-disk blocks from both partial-match routes; full matches stay reusable via the async onboard. Signed-off-by: Md Nafis Ul Haque Shifat <nafis@deepinfra.com>
… read completes onboardBlock enqueued the async read of a disk slot then immediately returned its card to the disk free queue, so the next spill could claim and overwrite the slot file before the read ran -- onboarding the wrong block. Hold the card until the read lands (reclaimed in claimDiskTarget once the read completes) and drop the now-unused completed-reads list. Signed-off-by: Md Nafis Ul Haque Shifat <nafis@deepinfra.com>
Short reads/writes and open failures on disk slot files were logged but the block was still published, serving or caching corrupt KV. Fail loud on any disk I/O error, matching the synchronous spill path. Signed-off-by: Md Nafis Ul Haque Shifat <nafis@deepinfra.com>
…lain TP The cross-rank onboard-readiness allreduce/allgather is only correct under plain tensor parallelism; under attention-DP or pipeline parallelism it deadlocks or drops requests. Gate the collective to plain TP, and refuse the disk tier under attention-DP/PP at KV-manager init. Signed-off-by: Md Nafis Ul Haque Shifat <nafis@deepinfra.com>
…ard-pending requests can_queue was computed before parking requests whose disk onboard had not landed; if parking emptied the batch we forwarded on nothing. Re-check can_queue after parking, broadening the existing post-batch-change re-check. Signed-off-by: Md Nafis Ul Haque Shifat <nafis@deepinfra.com>
…scheduler into the reader loadFromFile ran waitForDiskSlotWrites on the scheduler even on the detached path, so onboarding a just-spilled slot could park the whole engine behind the write queue. Do that wait on the reader instead -- the owning request stays parked via isBlockReadPending until the read lands, so the scheduler is never blocked. The synchronous fallback still waits inline. Signed-off-by: Md Nafis Ul Haque Shifat <nafis@deepinfra.com>
… (opt-in TLLM_KV_DISK_DROP_ON_PRESSURE) Under write-queue saturation the scheduler blocked in enqueueDiskWrite* waiting for room. Add an opt-in policy (TLLM_KV_DISK_DROP_ON_PRESSURE, default off = prior wait behavior): best-effort (non-retained) spills are shed at the eviction gate instead of stalling; retained spills never drop and bypass the queue cap (their TTL guarantee requires they land). Drop is decided per-spill at the gate, not per-pool in enqueue*, which would otherwise leave partial/corrupt blocks and leak the reserved-pool reap. Signed-off-by: Md Nafis Ul Haque Shifat <nafis@deepinfra.com>
…host buffer The POSIX reader staged each onboard through a pageable std::vector, so cudaMemcpyAsync fell back to the driver's hidden pinned-staging copy -- H2D ran at ~5-8 GB/s and the call was secretly synchronous. Stage through a per-reader pinned buffer instead (grown on demand to one block's pool bytes; pageable fallback if pinning is refused). Each pool reads into its own offset with a single stream sync after the block, so the async DMA never races the next read and pool reads overlap prior pools' DMA. Onboard latency is parked-request TTFT for disk-cache hits. Signed-off-by: Md Nafis Ul Haque Shifat <nafis@deepinfra.com>
… collective when the disk tier is off _park_requests_awaiting_onboard ran on every context step and, on plain TP>1, issued a per-step object-pickling tp_allreduce -- even with no disk tier, because it only bailed on hasattr(are_blocks_ready), which is always true. That taxed every TP>1 model sharing the image. Gate the park on a flag captured at init (disk_cache_size>0 AND TLLM_KV_DISK_READERS>0): detached onboards are impossible otherwise, so are_blocks_ready() is always true and the park + its collective are pure overhead. No-op for disk deployments (flag true -> unchanged); the skip only affects non-disk / sync-onboard configs. Signed-off-by: Md Nafis Ul Haque Shifat <nafis@deepinfra.com>
…n onboard/release Two paths could hand back valid-looking KV that is actually the wrong tokens: 1. releaseBlocks freed a block whose detached disk-onboard read was still in flight. areBlocksReady() parks the forward path on read-readiness, but nothing gated release: a request completing/cancelling mid-onboard returned its DMA-destination block to the free queue, and getFreeBlock has no readiness check -- so a stranger's prefill and the late read collided in the same GPU memory. Mirror the park onto the release path: hold read-pending blocks in mReleaseReadPending until the DMA lands (reapReadPendingReleases frees them). Gated by mNumDiskBlocks>0 && anyReadPending(), so non-disk models are unaffected. 2. reclaimSecondaryBlock's sync (copy) spill returned diskTarget without detaching its displaced disk-cached identity, relying on the caller's cleanup. getFreeBlock detaches; offloadBlock (the other caller) does not -- leaving a radix-tree node pointing at the recycled host slot, so a later prefix match reused freed content. Detach the identity unconditionally in reclaimSecondaryBlock for both callers; getFreeBlock's later detach then becomes a harmless no-op. Validated: kvCacheManagerTest disk-tier suite 11/11 green (spill/onboard/release + byte-exact onboard); the 12 pre-existing non-disk failures are unchanged. Signed-off-by: Md Nafis Ul Haque Shifat <nafis@deepinfra.com>
…disk-free check reapReadPendingReleases() -- which returns landed onboard cards (and, since the release-vs-onboard fix, release-deferred blocks) to their free queues -- ran only inside claimDiskTarget(). But reclaimSecondaryBlock() reaches claimDiskTarget only after passing a getNumFreeBlocks(kDiskLevel)==0 early-return, and parked cards sit OUTSIDE the free queue. So an onboard-heavy burst that parks every free disk card drives the free count to 0, fires the early-return, and never reaches claimDiskTarget (the only reap site) -- the landed-but-parked cards are never reclaimed and spills wedge permanently, exactly under the load the retained tier is for. Reap right after claiming the victim, before the disk-free check, guarded by mNumDiskBlocks>0. Placed in reclaimSecondaryBlock so it covers both callers (getFreeBlock and offloadBlock); claimDiskTarget's reap stays as a harmless second pass. Validated: kvCacheManagerTest disk-tier suite 11/11 green; the 12 pre-existing non-disk failures are unchanged. Signed-off-by: Md Nafis Ul Haque Shifat <nafis@deepinfra.com>
…ve the disk A retained block spilled to disk pushes an entry onto mDiskDeadlines, but nothing removed that entry when the block was onboarded back to GPU. The only cleanup was inside the displacement walk, which runs only when the disk is full of retained blocks -- so under normal reuse the stale entries accumulated forever (memory leak, and ever-longer walks once displacement does run). Fix: keep the deadlines in a std::set keyed by (expiry, seq). Each block remembers its key, so its entry is erased with one lookup whenever the block leaves the disk (onboard or displacement). The walk also erases any stale entry it finds. Signed-off-by: Md Nafis Ul Haque Shifat <nafis@deepinfra.com>
…parallelism The drop decision depends on each rank's write-queue depth, so ranks could diverge on which blocks remain reusable. Refuse the combination at init, like attention-DP and pipeline parallelism. Signed-off-by: Md Nafis Ul Haque Shifat <nafis@deepinfra.com>
…lock Retention deadlines were anchored and checked against each rank's own wall clock, which could diverge the radix trees across ranks. Now one rank broadcasts its time each iteration and every rank uses it for retention decisions. Signed-off-by: Md Nafis Ul Haque Shifat <nafis@deepinfra.com>
Signed-off-by: Md Nafis Ul Haque Shifat <nafis@deepinfra.com>
Contributor
Author
|
Hi @chienchunhung, thanks again for taking the time to review; those should be fixed now. Would appreciate it if you could take another look when you have time. |
Collaborator
|
[by Codex] @thorjohnsen Could you please review this PR? Thank you! |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h (2)
1517-1542: 🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
mDiskWritePressureDroppedcounter has no public getter.
getNumDiskSpills,getNumDiskGateDropped,getNumDiskAdmissionRefused, andgetNumDiskOnboardsare all exposed at bothWindowBlockManager(1524-1542) and aggregated atBlockManager(2205-2243) level, butmDiskWritePressureDropped(declared at line 1519) has no corresponding getter at either level. Silent write-pressure drops are exactly the kind of disk-tier backpressure signal operators need to diagnose reliability issues; without a getter it's currently unobservable.Also applies to: 2205-2253
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h` around lines 1517 - 1542, Add public getters for mDiskWritePressureDropped in both WindowBlockManager and the aggregated BlockManager interfaces, following the existing getNumDiskSpills and related getter patterns, and ensure BlockManager aggregates and exposes the corresponding counter.
1322-1330: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winDoxygen comment appears merged/misattributed between
reclaimSecondaryBlockandonboardBlock.Lines 1322-1323 ("Bring offloaded block from secondary to primary memory... Does nothing if block is already in primary memory") describe
onboardBlock's behavior, but the new lines 1324-1326 were inserted between that comment andreclaimSecondaryBlock()'s declaration, so the whole stacked block (two\brieftags) now documentsreclaimSecondaryBlock, leavingonboardBlock(line 1329) with no doc comment at all.As per coding guidelines, "use Doxygen comments for new interfaces."📝 Proposed fix to restore correct doc-comment attribution
- //! \brief Bring offloaded block from secondary to primary memory. - //! \details Does nothing if block is already in primary memory. - //! \brief Reclaim a free secondary (host) block for reuse. If the victim still holds + //! \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 = "");🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h` around lines 1322 - 1330, Separate the Doxygen comments so the reclaimSecondaryBlock documentation describes only reclaiming a free secondary block and spilling reusable content to disk, while the “Bring offloaded block…” documentation immediately precedes onboardBlock. Ensure both interfaces retain accurate, distinct documentation and avoid stacked \brief tags.Source: Coding guidelines
🧹 Nitpick comments (1)
cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h (1)
459-467: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueNew sequence-number fields use unsigned types without an approved exception.
mDiskDeadlineSeq/getDiskDeadlineSeq/setDiskDeadlineSeq(459-467),DiskDeadline::seq+mDiskSpillSeq(1560-1577), andmUnstagedSpillSeq/mPendingSpillBlockskey (1586-1601) are all newstd::uint64_tfields that are monotonically-increasing counters, not bitmaps, external-library requirements, or size comparisons.As per coding guidelines, "Prefer signed integers; use unsigned types only for bitmaps, external-library requirements, or unavoidable size comparisons."
Also applies to: 1560-1577, 1586-1601
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h` around lines 459 - 467, Replace the unsigned types used for the sequence counters with approved signed integer types throughout the affected declarations and APIs. Update getDiskDeadlineSeq/setDiskDeadlineSeq, DiskDeadline::seq, mDiskSpillSeq, mUnstagedSpillSeq, and the mPendingSpillBlocks key consistently, including any related initializers or comparisons.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h`:
- Around line 1517-1542: Add public getters for mDiskWritePressureDropped in
both WindowBlockManager and the aggregated BlockManager interfaces, following
the existing getNumDiskSpills and related getter patterns, and ensure
BlockManager aggregates and exposes the corresponding counter.
- Around line 1322-1330: Separate the Doxygen comments so the
reclaimSecondaryBlock documentation describes only reclaiming a free secondary
block and spilling reusable content to disk, while the “Bring offloaded block…”
documentation immediately precedes onboardBlock. Ensure both interfaces retain
accurate, distinct documentation and avoid stacked \brief tags.
---
Nitpick comments:
In `@cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h`:
- Around line 459-467: Replace the unsigned types used for the sequence counters
with approved signed integer types throughout the affected declarations and
APIs. Update getDiskDeadlineSeq/setDiskDeadlineSeq, DiskDeadline::seq,
mDiskSpillSeq, mUnstagedSpillSeq, and the mPendingSpillBlocks key consistently,
including any related initializers or comparisons.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 9bcdd530-40cb-4e20-8459-871e5d912725
📒 Files selected for processing (8)
cpp/include/tensorrt_llm/batch_manager/kvCacheManager.hcpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.hcpp/include/tensorrt_llm/executor/executor.hcpp/tensorrt_llm/batch_manager/kvCacheManager.cppcpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cppcpp/tensorrt_llm/executor/serialization.cppcpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cppcpp/tensorrt_llm/nanobind/executor/executorConfig.cpp
🚧 Files skipped from review as they are similar to previous changes (7)
- cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h
- cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp
- cpp/tensorrt_llm/executor/serialization.cpp
- cpp/tensorrt_llm/nanobind/executor/executorConfig.cpp
- cpp/include/tensorrt_llm/executor/executor.h
- cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp
- cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp
…doc comment Expose mDiskWritePressureDropped via getters at both WindowBlockManager and BlockManager, matching the other disk counters. Split the stacked doc comment so reclaimSecondaryBlock and onboardBlock each carry their own. Signed-off-by: Md Nafis Ul Haque Shifat <nafis@deepinfra.com>
Collaborator
|
[by Codex] @thorjohnsen Friendly review reminder: this PR is awaiting your review. Thanks! |
Contributor
Author
|
@thorjohnsen @nvpohanh @chienchunhung gentle ping - review when you get a chance would be much appreciated 🙏 |
…spec decode A parked disk onboard re-enters as a first context chunk, so eagle3/mtp reuse an existing hidden-state slot instead of re-allocating. And the disk slot directory is namespaced by draft/target, not just rank, so a draft KV manager cannot overwrite the target's slot files. Signed-off-by: Md Nafis Ul Haque Shifat <nafis@deepinfra.com>
Contributor
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
tensorrt_llm/_torch/pyexecutor/resource_manager.py (2)
1504-1510: 📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick winAdd the required public API type and documentation.
Annotate
requestasLlmRequest. Add Google-style documentation forset_retention_clock.Proposed fix
- def are_blocks_ready(self, request) -> bool: + def are_blocks_ready(self, request: LlmRequest) -> bool:🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/pyexecutor/resource_manager.py` around lines 1504 - 1510, Update ResourceManager.are_blocks_ready to annotate its request parameter as LlmRequest, importing the type if needed. Add a Google-style docstring to set_retention_clock describing the now_ns argument and the method’s retention-clock behavior, while preserving the existing delegation.Source: Coding guidelines
1509-1510: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winSynchronize the retention clock before disaggregated generation initialization.
_prepare_disagg_gen_init()calls the KV resource manager directly before_sync_retention_clock(). This path can perform retention admission or eviction with stalemRetentionNow. Synchronize the clock before this call.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/pyexecutor/resource_manager.py` around lines 1509 - 1510, Update the disaggregated generation initialization flow in _prepare_disagg_gen_init() to invoke _sync_retention_clock() before calling the KV resource manager, ensuring retention admission and eviction use the current clock. Keep set_retention_clock() unchanged and preserve the existing initialization sequence after synchronization.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@tensorrt_llm/_torch/pyexecutor/resource_manager.py`:
- Around line 1504-1510: Update ResourceManager.are_blocks_ready to annotate its
request parameter as LlmRequest, importing the type if needed. Add a
Google-style docstring to set_retention_clock describing the now_ns argument and
the method’s retention-clock behavior, while preserving the existing delegation.
- Around line 1509-1510: Update the disaggregated generation initialization flow
in _prepare_disagg_gen_init() to invoke _sync_retention_clock() before calling
the KV resource manager, ensuring retention admission and eviction use the
current clock. Keep set_retention_clock() unchanged and preserve the existing
initialization sequence after synchronization.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: ff523bfb-7a2e-4a62-b40c-7adbdff7fd59
📒 Files selected for processing (3)
tensorrt_llm/_torch/pyexecutor/resource_manager.pytensorrt_llm/_torch/speculative/eagle3.pytensorrt_llm/_torch/speculative/mtp.py
# Conflicts: # cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h # cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp # cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp
… are_blocks_ready The disagg-gen-init path allocates and commits KV before the main-loop retention-clock sync, so it would anchor deadlines against a stale clock; sync there too. Also annotate are_blocks_ready's request parameter and document set_retention_clock. Signed-off-by: Md Nafis Ul Haque Shifat <nafis@deepinfra.com>
Contributor
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
Contributor
There was a problem hiding this comment.
Actionable comments posted: 5
🧹 Nitpick comments (7)
cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h (1)
199-199: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDeclare
mDiskUseGdsasconst.
mDiskUseGdsis read once at construction and never reassigned, likemAsyncDiskStoreon line 203. Mark itbool constfor consistency and to prevent later mutation.As per coding guidelines: "declare unmodified variables as
const".♻️ Proposed change
- bool mDiskUseGds{std::getenv("TLLM_KV_DISK_GDS") != nullptr}; + bool const mDiskUseGds{std::getenv("TLLM_KV_DISK_GDS") != nullptr};🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h` at line 199, Declare the KVCacheTransferManager member mDiskUseGds as const, matching the existing const treatment of other immutable configuration members such as mAsyncDiskStore while preserving its current environment-based initialization.Source: Coding guidelines
tests/unittest/llmapi/test_disk_tier_args.py (1)
51-69: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse an isolated temporary directory for this test.
/tmp/kvdisk-testis shared across test runs and remains after the test completes. A pre-existing file at that path can makeos.makedirs(..., exist_ok=True)fail. Use thetmp_pathfixture and passstr(tmp_path)asdisk_cache_path.Proposed change
-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 +def test_llm_args_kv_cache_config_disk_fields(tmp_path): 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_path=str(tmp_path),🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/llmapi/test_disk_tier_args.py` around lines 51 - 69, Update test_llm_args_kv_cache_config_disk_fields to accept the tmp_path fixture and use it as the disk cache directory instead of the shared /tmp/kvdisk-test path. Remove the manual os.makedirs call, pass str(tmp_path) to disk_cache_path, and update the corresponding assertion to compare against that value.cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp (4)
11122-11123: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueMove the new includes to the top of the file.
#include <filesystem>and#include <thread>appear at Line 11122, after roughly 11000 lines of code. Include directives belong in the file's include block so the translation unit's dependencies are visible in one place.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp` around lines 11122 - 11123, Move the <filesystem> and <thread> includes from their current location near the end of the file into the existing top-of-file include block, keeping the dependencies declared together and removing the duplicate late includes.
11655-11658: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueInitialize the secondary pool before spilling from it.
pool.secondaryPtris allocated as pinned host memory and is never written beforespillToFileUnstagedreads it. The writer therefore copies indeterminate bytes to disk. The test passes either way, but sanitizers report the read. Add astd::memsetor a fill loop, matching whatrunDiskOnboardByteRoundTripdoes at Line 11568.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp` around lines 11655 - 11658, Initialize the pinned host buffer assigned to pool.secondaryPtr before calling spillToFileUnstaged, using std::memset or the same fill approach as runDiskOnboardByteRoundTrip. Ensure the entire allocated secondary pool contains defined bytes before it is read and written to disk.
11445-11452: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winThe control assertion depends on eviction victim selection.
EXPECT_EQ(blockManager.getNumDiskSpills(), spillsBefore)assumes the secondchurnOnceevicts onlyseqB's expired blocks.seqA2's blocks carry a 60 s TTL and were reused earlier in the same test, so if the churn selects any of them as a host victim, they pass the retained-only gate and the spill count increases. That makes this control assertion order-dependent.Consider scoping the control to a fresh
BlockManager, or asserting on a per-sequence signal instead of the global spill counter.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp` around lines 11445 - 11452, Make the control assertion independent of eviction victim selection by isolating it in a fresh BlockManager, or replace the global getNumDiskSpills comparison with a per-sequence signal that verifies seqB is not re-spilled. Update the control setup around addDiskTierSequence, churnOnce, and the EXPECT_EQ assertion while preserving the intended expired-sequence gating behavior.
11221-11226: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueMove the doc comment to the function it describes.
The comment at Line 11221 describes
addDiskTierSequence("Adds a sequence; diskTtl set => ..."), but it sits directly aboveretentionTick. Move it to Line 11228, immediately aboveaddDiskTierSequence.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp` around lines 11221 - 11226, Move the descriptive comment from above retentionTick to immediately above addDiskTierSequence, leaving retentionTick without that unrelated documentation and preserving the comment text.tensorrt_llm/_torch/pyexecutor/py_executor.py (1)
2643-2645: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winRemove the no-op park call from the PP loop.
resource_managerrejects disk KV-cache configurations whenpp_size > 1, and_park_requests_awaiting_onboard()returns when disk onboarding is inactive. The PP loop therefore cannot shrink the batch through this call and does not need the non-PP re-check.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/pyexecutor/py_executor.py` around lines 2643 - 2645, Remove the `_park_requests_awaiting_onboard(scheduled_batch)` call from the PP loop after `prepare_resources`; leave `_sync_retention_clock()` and `resource_manager.prepare_resources(scheduled_batch)` unchanged.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h`:
- Around line 225-236: Add a shared environment-size parsing helper for
TLLM_KV_DISK_WRITE_QUEUE, TLLM_KV_DISK_WRITERS, and mNumDiskReaders that rejects
malformed, negative, and overflowing values, reports the variable name
diagnostically, and clamps each result to its valid range; use a minimum of 1
for the queue capacity and writer count, while preserving 0 as the valid
synchronous-read value for mNumDiskReaders. Update all three member initializers
to use the helper and add the required <cstring>, <stdexcept>, and <algorithm>
includes.
In `@cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp`:
- Around line 11837-11856: The bounded-queue test runs its burst synchronously
without a timeout, so a stalled writer can hang indefinitely. In
runWriteQueueBoundedTest at
cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp lines 11837-11856,
move burst production to a separate thread, join it with a deadline, and call
KVCacheTransferManagerTestAccess::requestWriterStop on timeout so the producer
unblocks before assertions; retain requestWriterStop at lines 115-123 because
this change uses it.
- Around line 11744-11770: The seed spills in the setup around
KVCacheTransferManager and spillToFile are asynchronous because
TLLM_KV_DISK_ASYNC_STORE is enabled before manager construction. Make the setup
description accurate by disabling or scoping out asynchronous disk storage
during seeding so spillToFile completes inline, then restore the asynchronous
setting before exercising queued reads; alternatively update the comment to
describe the actual queued behavior if synchronous seeding is not required.
In `@tensorrt_llm/_torch/pyexecutor/py_executor.py`:
- Around line 592-601: Update the environment parsing used for disk-reader
activation and construction to strictly validate TLLM_KV_DISK_READERS as a
non-negative decimal value, rejecting malformed strings such as “4x” and
negative values such as “-1”. Reuse the validated value consistently for
_disk_onboard_active and both disk-reader components, while preserving the
existing unset-default behavior of zero.
In `@tests/unittest/llmapi/test_disk_tier_args.py`:
- Around line 1-5: Add the repository-standard NVIDIA copyright and SPDX header
at the beginning of test_disk_tier_args.py, before its module docstring, using
2026 as the latest meaningful modification year.
---
Nitpick comments:
In `@cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h`:
- Line 199: Declare the KVCacheTransferManager member mDiskUseGds as const,
matching the existing const treatment of other immutable configuration members
such as mAsyncDiskStore while preserving its current environment-based
initialization.
In `@cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp`:
- Around line 11122-11123: Move the <filesystem> and <thread> includes from
their current location near the end of the file into the existing top-of-file
include block, keeping the dependencies declared together and removing the
duplicate late includes.
- Around line 11655-11658: Initialize the pinned host buffer assigned to
pool.secondaryPtr before calling spillToFileUnstaged, using std::memset or the
same fill approach as runDiskOnboardByteRoundTrip. Ensure the entire allocated
secondary pool contains defined bytes before it is read and written to disk.
- Around line 11445-11452: Make the control assertion independent of eviction
victim selection by isolating it in a fresh BlockManager, or replace the global
getNumDiskSpills comparison with a per-sequence signal that verifies seqB is not
re-spilled. Update the control setup around addDiskTierSequence, churnOnce, and
the EXPECT_EQ assertion while preserving the intended expired-sequence gating
behavior.
- Around line 11221-11226: Move the descriptive comment from above retentionTick
to immediately above addDiskTierSequence, leaving retentionTick without that
unrelated documentation and preserving the comment text.
In `@tensorrt_llm/_torch/pyexecutor/py_executor.py`:
- Around line 2643-2645: Remove the
`_park_requests_awaiting_onboard(scheduled_batch)` call from the PP loop after
`prepare_resources`; leave `_sync_retention_clock()` and
`resource_manager.prepare_resources(scheduled_batch)` unchanged.
In `@tests/unittest/llmapi/test_disk_tier_args.py`:
- Around line 51-69: Update test_llm_args_kv_cache_config_disk_fields to accept
the tmp_path fixture and use it as the disk cache directory instead of the
shared /tmp/kvdisk-test path. Remove the manual os.makedirs call, pass
str(tmp_path) to disk_cache_path, and update the corresponding assertion to
compare against that value.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 82d6dff9-97ea-41a5-b346-456005f5c37a
📒 Files selected for processing (23)
cpp/include/tensorrt_llm/batch_manager/evictionPolicy.hcpp/include/tensorrt_llm/batch_manager/kvCacheManager.hcpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.hcpp/include/tensorrt_llm/executor/executor.hcpp/tensorrt_llm/batch_manager/evictionPolicy.cppcpp/tensorrt_llm/batch_manager/kvCacheManager.cppcpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cppcpp/tensorrt_llm/executor/kvCacheConfig.cppcpp/tensorrt_llm/executor/kvCacheRetentionConfig.cppcpp/tensorrt_llm/executor/serialization.cppcpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cppcpp/tensorrt_llm/nanobind/executor/executorConfig.cppcpp/tensorrt_llm/nanobind/executor/request.cppcpp/tests/unit_tests/batch_manager/evictionPolicyTest.cppcpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpptensorrt_llm/_torch/pyexecutor/py_executor.pytensorrt_llm/_torch/pyexecutor/resource_manager.pytensorrt_llm/_torch/speculative/eagle3.pytensorrt_llm/_torch/speculative/mtp.pytensorrt_llm/llmapi/llm_args.pytensorrt_llm/serve/openai_protocol.pytensorrt_llm/serve/openai_server.pytests/unittest/llmapi/test_disk_tier_args.py
🚧 Files skipped from review as they are similar to previous changes (19)
- cpp/include/tensorrt_llm/batch_manager/evictionPolicy.h
- tensorrt_llm/_torch/speculative/eagle3.py
- cpp/tensorrt_llm/executor/kvCacheRetentionConfig.cpp
- tensorrt_llm/_torch/speculative/mtp.py
- tensorrt_llm/serve/openai_protocol.py
- cpp/tests/unit_tests/batch_manager/evictionPolicyTest.cpp
- tensorrt_llm/llmapi/llm_args.py
- cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp
- cpp/tensorrt_llm/nanobind/executor/request.cpp
- cpp/tensorrt_llm/executor/kvCacheConfig.cpp
- tensorrt_llm/_torch/pyexecutor/resource_manager.py
- cpp/tensorrt_llm/nanobind/executor/executorConfig.cpp
- cpp/include/tensorrt_llm/executor/executor.h
- tensorrt_llm/serve/openai_server.py
- cpp/tensorrt_llm/batch_manager/evictionPolicy.cpp
- cpp/tensorrt_llm/executor/serialization.cpp
- cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp
- cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h
- cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp
Collaborator
|
[by Codex] @thorjohnsen Could you please review PR #16252 for the KV-cache manager changes? Thanks! |
thorjohnsen
reviewed
Aug 11, 2026
thorjohnsen
left a comment
thorjohnsen
left a comment
Collaborator
There was a problem hiding this comment.
I assume onboarding from disk takes a long time, during which the GPUs sit idle. Wouldn't it make more sense to onboard via host from disk, disk -> host -> gpu? Scheduler sees the next sequence to be launched has reusable blocks on disk and start prefetching those to host memory, sequence only launches after prefetching has finished.
Contributor
Author
|
Thanks for taking a look @thorjohnsen ! The onboard is asynchronous, so the GPU does not actually sit idle during a disk read as the other requests keep being processed, so the GPU stays busy. For the transfer path, the async POSIX path stages disk - > pinned host - > GPU and the GDS path does a direct write to GPU skipping the host (though i have only verified GDS functionally, not benchmarked its throughput). I did think about the prefetching idea and I think it could be a followup to this pr. |
Collaborator
|
Thanks, I see that the scheduler does not wait for onboarding from disk, the only real cost is that the GPU target blocks are held until the request can start. This will reduce amount of available KV cache slightly, but I agree it can be postponed to a follow-up PR. |
nvpohanh
reviewed
Aug 14, 2026
| /// @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). |
Collaborator
There was a problem hiding this comment.
@nafis271 Could you check if you can switch to kvCacheManagerV2 which already has this support? If you want more information about kvCacheManagerV2, please let us know. Thanks!
We are planning to migrate from V1 to V2 very soon.
Contributor
Author
There was a problem hiding this comment.
Thanks for letting me know! The reason we implemented it on V1 is that when we wanted this explicit user-provided TTL cache at DeepInfra, most of our models ran on V1 and we could not migrate due to performance and compatibility issues with V2. But I understand it makes sense to move to V2 if that is the long term plan.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary by CodeRabbit
kv_cache_ttl_seconds.disk_cache_sizedisk_cache_pathdisk_cache_retained_onlydisk_cache_protect_unexpiredDev Engineer Review
KVCacheManager,BlockManager,WindowBlockManager, nanobind, and Python resource management.KVCacheBlocktracks disk slots, deadlines, retention state, and reuse counts.areBlocksReady()prevents forwarding before detached disk reads complete and reports asynchronous read failures.PyExecutorparks pending context requests and synchronizes tensor-parallel readiness where supported.QA Engineer Review
Test changes
cpp/tests/unit_tests/batch_manager/evictionPolicyTest.cppLRUPolicyTest.IsEnqueuedTest.cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpptests/unittest/llmapi/test_disk_tier_args.pykv_cache_ttl_secondsvalues.Coverage
tests/integration/test_lists/,test-db/, andqa/was not confirmed.test-db/orqa/.Verdict: needs follow-up.
Description
Closes #16251.
The v1 KV cache block manager currently supports two tiers: GPU memory (primary) and
host memory (secondary). This PR adds a third, disk-backed tier. When a block is
evicted from the secondary (host) tier it is spilled to disk instead of being dropped,
and is onboarded back to GPU on a later prefix hit — extending KV reuse well beyond
DRAM (a disk read is far cheaper than recomputing a long prefix).
It also adds a user-provided retention TTL (
kv_cache_ttl_seconds, per request): theengine keeps that request's KV blocks reusable on disk for the requested duration and
tries to honor it unless the disk tier is out of space.
KV cache manager v2 already has a three-tier (disk) cache, but today it must run with
the overlap scheduler disabled, its performance with the disk tier is not
production-level (see Benchmarks below), and it has no user-provided retention TTL.
This PR brings a production-oriented disk tier — with a retention TTL — to the default v1
manager.
The tier is off by default; nothing is added to the scheduling hot path when
disk_cache_sizeis unset.How it works
existing radix reuse tree. Disk blocks are poolless — a block id + a fixed per-rank
disk slot, no KV memory pool — so disk capacity is independent of GPU/host memory. A new
kDiskLevelis added to the eviction policy (evictionPolicy.cpp).bytes are written to that block's disk slot instead of being dropped; the block stays in
the reuse tree, now disk-resident (
reclaimSecondaryBlock→spillToFileinkvCacheManager.cpp/kvCacheTransferManager.cpp). A small pool of pre-reserved hostblocks (
RESERVED_BLOCKS) lets a spill hand back a free slot immediately — theevicted block drains to disk in the background — so the scheduler never waits on a copy;
if that pool empties under a burst, the block is copied out instead.
back into a GPU block. Reads are asynchronous with per-block readiness: the scheduler
parks the request until its blocks are ready and then resumes it, never blocking on disk
I/O (
onboardBlock/loadFromFile; park logic inpy_executor.py).bounded shared queue (
diskWriterLoop/diskReaderLoop); the scheduler enqueues andcontinues, with backpressure when the queue is full.
retention TTL is set.
kv_cache_ttl_secondsstamps a request's committed blocks with anexpiry (
mRetentionExpiry); on disk they are protected from eviction (deadline heapmDiskDeadlines) until they expire. The stamp is orthogonal — it never changesGPU/host eviction, so retained traffic can't crowd normal traffic out of the fast tiers.
ranks.
Configuration (
kv_cache_config)disk_cache_sizeNonedisk_cache_pathNonedisk_cache_size > 0.disk_cache_retained_onlyfalsedisk_cache_protect_unexpiredfalseEnvironment variables (tuning)
TLLM_KV_DISK_ASYNC_STORETLLM_KV_DISK_WRITERS1TLLM_KV_DISK_READERS00= no dedicated readers.TLLM_KV_DISK_WRITE_QUEUE1024TLLM_KV_DISK_RESERVED_BLOCKS0(off)0uses the default spill path. Size it against how far the writers can fall behind.TLLM_KV_DISK_MIN_REUSE00= spill all).TLLM_KV_DISK_DROP_ON_PRESSURE0(off)TLLM_KV_DISK_GDSServe API (retention TTL)
kv_cache_ttl_seconds(optional int) on chat-completion and completion requests: "Keepthis request's KV cache reusable for this many seconds (disk-tier retention)." Omitted =
today's best-effort LRU reuse.
Benchmarks
Setup. gpt-oss-120b, single B300, TP=1. Both arms share: 100 GiB host tier + 300 GiB
disk tier,
free_gpu_memory_fraction=0.8, block reuse on / partial-reuse off, MoE-TRTLLMbackend,
MAX_UTILIZATIONscheduler, chunked prefill. v1 (this PR): overlap scheduleron; asynchronous spill/onboard with 4 reader + 4 writer threads, a reserved-block pool,
spill-all — every evicted block written to disk (
disk_cache_retained_only=false, reusegate off), no drop-on-pressure.
v2:
use_kv_cache_manager_v2=true; overlap scheduler off (required — overlap-on crashesunder load); disk prefetch tested at both
disk_prefetch_num_reqs=2and=32, withidentical results (~250 s TTFT either way) — V2's collapse is not a prefetch-tuning artifact.
Load generator: LMCache LMBenchmark, the
synthetic-multi-round-qa/long_input_short_output_run.shworkload — 120 concurrent users,12 rounds, ~50k-token chat history, 1000-token system prompt, 100-token answers, at qps
1.5 for 500 s — a working set that overflows GPU + host into the disk tier. Identical
workload for both arms; figures are LMBenchmark's per-interval report windows.
warmup (TTFT transiently ~16 s while the 300 GB tier writes out), it converges to a
steady ~0.8–1.2 s TTFT while sustaining the full ~85k in-tok/s offered load.
in that stable mode it cannot sustain the load: once the working set overflows into
disk, throughput collapses to ~1/5 of offered and per-window TTFT climbs monotonically
past 250 s — a backlog that never stabilizes.
Test Coverage
cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp— disk-tier suite: byte-exactspill/onboard round-trip,
disk_cache_retained_onlyadmission gate, TTL / deadlineeviction order,
disk_cache_protect_unexpired, theTLLM_KV_DISK_MIN_REUSEreuse gate,and disk-tier queue-integrity validation.
cpp/tests/unit_tests/batch_manager/evictionPolicyTest.cpp— three-level eviction policy.tests/unittest/llmapi/test_disk_tier_args.py— disk config validation and theserve-layer
kv_cache_ttl_seconds→ retention-config translation.PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title. (This PR isapi-compatible: new optional config fields + an optional request field; no behavior change when unset.)Any new dependencies have been scanned for license and vulnerabilities (none added)
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.