From 4b767a827daac055dc551e9a6736eb02a67c38bc Mon Sep 17 00:00:00 2001 From: Yao Yao Date: Wed, 28 Oct 2026 12:38:57 +0000 Subject: [PATCH 01/12] [None][feat] KVCM2 core: sparse attention arena, host-backed pools and Batch The KVCM2-only subset of the sparse-attention offload work, split out so the cache-manager core can be reviewed and CI'd without its Python consumers. Sparse attention needs a small, per-request GPU working set of KV units chosen afresh each step, while the history it is chosen from lives in host memory. This adds: - A per-request cell arena owned by `Batch`. Each request reserves one lease, and every buffer of a layer mirrors the same entry index. - Host-backed sparse pools. Sparsity is a property of a pool, so a pool group's key carries it alongside the coalescing group and size; the GPU tier can then host-back exactly the sparse pools without the life cycle having to be wholly one or the other. - Resolve, fetch and write-back kernels. The resolve maintains residency and emits a selection-level miss list; the copy applies it per buffer; the flush only ever reads residency. - Selection groups: layers driven by one indexer share a resident set, a recency order and one resolve, while keeping their own mirrored cells. `fetchSparsePages` is the fused form and rejects a layer whose selection is shared; `resolveSparseSelection` + `fetchResolvedPages` serve that case. Supporting changes in the same tree: a `StrongIndex`/`TypedVec`/`EnumMap` family for index type-safety that also works in device code, a named-buffer table for the `Batch` staging block, and a parallel initial host-tier fill built on this branch's target/step backing model. Ordering contract worth a reviewer's attention: the miss list and per-slot cell mapping are one per life cycle, so the next resolve of any selection must be ordered after every `fetchResolvedPages()` consuming the current one -- with an event where the consumer is on another stream. Signed-off-by: Yao Yao --- .../kv_cache_manager_v2/AGENTS.md | 129 +- .../kv_cache_manager_v2/CMakeLists.txt | 5 +- .../INDEX_CONVERSION_DESIGN.md | 275 + .../SPARSE_ATTENTION_DESIGN.md | 1417 ++++ .../kv_cache_manager_v2/batch.cpp | 2130 +++++- .../batch_manager/kv_cache_manager_v2/batch.h | 1020 ++- .../kv_cache_manager_v2/batchedPageCopy.cu | 14 +- .../kv_cache_manager_v2/batchedPageCopy.h | 29 +- .../kv_cache_manager_v2/coldPageCodec.h | 5 + .../kv_cache_manager_v2/common.h | 7 +- .../kv_cache_manager_v2/config.cpp | 161 +- .../kv_cache_manager_v2/config.h | 145 +- .../kv_cache_manager_v2/exceptions.h | 5 +- .../indexConversionKernels.cu | 260 + .../indexConversionKernels.h | 104 + .../kv_cache_manager_v2/introspection.cpp | 495 +- .../kv_cache_manager_v2/introspection.h | 97 + .../kv_cache_manager_v2/kvCache.cpp | 785 +-- .../kv_cache_manager_v2/kvCache.h | 293 +- .../kv_cache_manager_v2/kvCacheManager.cpp | 378 +- .../kv_cache_manager_v2/kvCacheManager.h | 288 +- .../kv_cache_manager_v2/lifeCycleRegistry.cpp | 16 +- .../kv_cache_manager_v2/lifeCycleRegistry.h | 32 +- .../kv_cache_manager_v2/page.cpp | 188 +- .../batch_manager/kv_cache_manager_v2/page.h | 52 +- .../kv_cache_manager_v2/pendingStats.h | 4 +- .../kv_cache_manager_v2/sparseCache.cpp | 243 + .../kv_cache_manager_v2/sparseCache.h | 224 + .../kv_cache_manager_v2/sparseCacheKernels.cu | 1075 +++ .../kv_cache_manager_v2/sparseCacheKernels.h | 455 ++ .../kv_cache_manager_v2/stagingBuffer.cpp | 16 +- .../kv_cache_manager_v2/stagingBuffer.h | 20 +- .../kv_cache_manager_v2/storage/config.cpp | 38 +- .../kv_cache_manager_v2/storage/config.h | 34 + .../kv_cache_manager_v2/storage/core.cpp | 1039 ++- .../kv_cache_manager_v2/storage/core.h | 442 +- .../kv_cache_manager_v2/storageManager.cpp | 458 +- .../kv_cache_manager_v2/storageManager.h | 110 +- .../kv_cache_manager_v2/utils/generator.h | 348 + .../kv_cache_manager_v2/utils/hostMem.cpp | 181 +- .../kv_cache_manager_v2/utils/hostMem.h | 66 +- .../utils/hostMemBacking.cpp | 306 +- .../utils/hostMemBacking.h | 227 +- .../kv_cache_manager_v2/utils/math.h | 13 +- .../kv_cache_manager_v2/utils/sharedPtr.h | 2 +- .../utils/staticHostMem.cpp | 5 +- .../kv_cache_manager_v2/utils/staticHostMem.h | 10 + .../kv_cache_manager_v2/utils/typedIndex.h | 149 +- .../batch_manager/kvCacheManagerV2.cpp | 1052 ++- .../unit_tests/batch_manager/CMakeLists.txt | 12 + .../kvCacheManagerV2ColdPageTest.cpp | 2805 +------- .../kvCacheManagerV2ConcurrencyTest.cpp | 37 - .../kvCacheManagerV2ConfigTest.cpp | 85 +- .../kvCacheManagerV2GeneratorTest.cpp | 208 + .../kvCacheManagerV2HostMemBackingTest.cpp | 266 +- .../kvCacheManagerV2HostMemTest.cpp | 109 +- .../kvCacheManagerV2StagingBufferTest.cu | 2 +- .../kvCacheManagerV2StatsTest.cpp | 10 +- .../kvCacheManagerV2HostMemNumaTest.cpp | 6 +- .../runtime/kv_cache_manager_v2/AGENTS.md | 4 +- .../runtime/kv_cache_manager_v2/__init__.py | 19 +- .../runtime/kv_cache_manager_v2/__init__.pyi | 1149 +++- .../kv_cache_manager_v2/_introspection.py | 330 +- .../kv_cache_manager_v2_tests/fake_engine.py | 24 +- .../test_kv_cache_manager_v2.py | 5868 ++++++++++++++++- 65 files changed, 20038 insertions(+), 5743 deletions(-) create mode 100644 cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/INDEX_CONVERSION_DESIGN.md create mode 100644 cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/SPARSE_ATTENTION_DESIGN.md create mode 100644 cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/indexConversionKernels.cu create mode 100644 cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/indexConversionKernels.h create mode 100644 cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/sparseCache.cpp create mode 100644 cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/sparseCache.h create mode 100644 cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/sparseCacheKernels.cu create mode 100644 cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/sparseCacheKernels.h create mode 100644 cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/generator.h create mode 100644 cpp/tests/unit_tests/batch_manager/kvCacheManagerV2GeneratorTest.cpp diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/AGENTS.md b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/AGENTS.md index d9c143c1fc1b..32d6371a3666 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/AGENTS.md +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/AGENTS.md @@ -18,6 +18,11 @@ The current C++ headers and tests are the source of truth. Historical Python behavior remains useful for compatibility testing, but migration-era proposed layouts or ownership models must not override the implementation. +`SPARSE_ATTENTION_DESIGN.md` covers how a sparse buffer's pages are held in host +memory and read through a per-request GPU cache, and `INDEX_CONVERSION_DESIGN.md` +how page indices become the compact arena indices attention kernels take. Read +them before changing `batch.*`, `sparseCache.*` or the index conversion. + ## Layout - `common.h`, `config.h`, `tokenIdExt.h`, and `exceptions.h`: shared types, @@ -25,7 +30,12 @@ layouts or ownership models must not override the implementation. - `blockRadixTree.*`: the shared prefix-reuse tree and SHA-256 block keys. - `page.*`, `kvCache.*`, and `kvCacheManager.*`: page lifecycle, per-request cache state, and the top-level manager. -- `batch.*`: stable request rows, dirty tracking, and raw GPU metadata publication. +- `batch.*`: the caches stepped together and their device-side mirror — row + ownership, the per-layer-group page tables a sparse fetch reads, and the + per-row watermarks a step derives its counts from. +- `sparseCache.*`, `sparseCacheKernels.*`, and `indexConversionKernels.*`: the + per-request GPU cache over host-backed sparse pages, and the conversion from + page indices to the compact arena indices attention kernels take. - `storage/`, `storageManager.*`, `evictionController.*`, and `copyEngine.*`: pools, eviction ownership, migration, and data movement. - `lifeCycleRegistry.*`: layer-group/lifecycle mapping, including attention and @@ -46,6 +56,23 @@ layouts or ownership models must not override the implementation. - Keep hot-path checks under `TLLM_CHECK_DEBUG*` or `gDebug` when they are too expensive for release builds. Do not remove debug invariants simply because production does not execute them. +- Check with `TLLM_CHECK*` rather than a hand-rolled `if (bad) { throw ...; }`. + It is the convention across the project, and it keeps a condition and the + message explaining it in one place that reads as a check. + The exception is a failure whose *type* is the contract. The bindings make a + Python class of each one in `exceptions.h`, so callers -- and the poison latch + -- act on which class they caught; throw those directly: + `OutOfMemoryError` and its `Host`/`Disk`/`Cu` forms, `LogicError`, + `AssertionError`, `CuError`, `CorruptedError`, `ResourceBusyError`, + `OutOfPagesError`. + Validation at the binding boundary that throws `std::invalid_argument` or + `std::out_of_range` is legacy and should migrate to `TLLM_CHECK*` as it is + touched, so that KVCM2 ends up on the same footing as the rest of the repo. + Migrate deliberately rather than in passing: nanobind maps those to + `ValueError` and `IndexError` where `TLLM_CHECK*` raises `RuntimeError`, and + the suites assert the first two in well over a hundred places, so each + conversion carries its own test updates and is a Python API change under the + nanobind rule on preserving exception types. - Keep the core independent of Python. C++ storage, copying, CUDA, hashing, and lifecycle operations call C++ APIs directly; Python adaptation belongs in nanobind. @@ -59,8 +86,9 @@ must eventually be `close()`d. Its normal flow is: 1. Match the input `TokenSpan` against `BlockRadixTree` within a `ReuseScope`. 2. Allocate request-local pages for unmatched blocks through `StorageManager`. -3. Lock or restore required pages before model execution: GPU for writable and - dense pages, level-1 host memory for cold sparse history. +3. Lock or restore required pages before model execution. Every page locks at the + hot level; a sparse pool's hot allocation is host memory, which is a property + of that pool rather than of the level or of the group holding it. 4. Commit completed blocks to the tree, making their immutable pages available for later requests; `stopCommitting()` finalizes this process. 5. Suspend or close the request, returning pages to holding/eviction ownership. @@ -70,23 +98,42 @@ may have sliding-window and sink-token rules; SSM lifecycles represent a recurrent-state checkpoint. A pool-group index is a storage-layout index and is not interchangeable with a layer ID or lifecycle ID. -Attention lifecycle identity also includes `isSparse`. All buffers in one -attention layer must agree on this flag. Sparse buffers require `HOST_MEM` at -level 1 and are invalid for SSM layers. Sparse and dense GPU pools stay separate -even when their slot sizes match; compatible sparse buffers still coalesce. +Sparsity is a property of a pool, not of a lifecycle: a layer's buffers may be +partly sparse, and the pool each lands in carries the flag. Sparse buffers are +invalid for SSM layers. Sparse and dense pools stay separate even when their slot +sizes match, which is what lets a sparse pool be backed by host memory while a +dense pool of the same group and level is device memory; the group reports +GPU_MEM either way. Compatible sparse buffers still coalesce. Sparse buffers place no requirement on the cold tiers. `StorageManager` coordinates GPU, host, and disk cache levels. It allocates slots, schedules pages for eviction, migrates pages between levels, and resizes pools. `CopyEngine` performs the actual batched transfers; C++ code calls it directly and must not round-trip through Python bindings. -`Batch` groups requests from one manager across all layer groups. Request changes -invalidate their stable rows; `publish()` uploads final raw indices and eligible -history counts, including post-rollback state. Device addresses remain fixed. -Publish and wait for readiness outside graph capture; after submitting readers -or replaying a graph, call `recordRead()` before mutating requests or membership. -Publication waits for offload and prior readers, and retains each staging buffer -until its upload completes. DLPack views keep the allocation alive, not KV pages. +`Batch` groups the `KvCache`s stepped together and owns their row order. Rows are +slots rather than positions in a compacted list, so removing a member leaves a +hole instead of shifting later rows and the shapes and addresses a captured graph +baked in stay valid. Membership is exclusive and carries the sparse arena lease, +so a request joins only once it is ready to be stepped; `move()` relocates a +member between rows without ending membership, which `removeKvCache()` followed +by `addKvCache()` would. A member may not be resumed, suspended or closed while +seated, because those write page indices through the Batch's staged row. + +`publish()` uploads a step's page tables and watermarks from pinned staging and +returns with the copy still in flight, so every later host write to a member's +page indices — a `KvCache::resize()`, a commit, a beam fork, an SWA recycle, +anything that locks or unlocks a page — must be separated from it by +`waitPublished()`, or it races the copy and publishes a torn table. The Batch's +own mutators are exempt: `resize()`, `move()`, `addKvCache()` and +`removeKvCache()` each wait for themselves. `publish()` throws rather than waits +when the previous publish has not been waited on, because waiting there would +cover only its own write and hide a member already mutated into the hazard; the +lock and unlock paths carry a debug check for the member side of the same rule. + +A sparse member writes back the units its history passes through +`flushSparseUnits()`, and must have done so for the history it is giving up +before it leaves, since removal returns the arena lease and discards whatever the +GPU cache held. The dependency direction is broadly: @@ -96,6 +143,7 @@ types/config/exceptions -> storage pools + eviction + copy engine + radix tree -> pages + storage manager -> per-request KvCache + -> Batch + sparse cache -> KvCacheManager -> nanobind API ``` @@ -121,8 +169,9 @@ pages and performs final commit-state bookkeeping. ### Page status - `LOCKED`: pinned at the current storage level; neither eviction nor dropping - is permitted. GPU locks support every lifecycle. Level-1 `HOST_MEM` locks - support sparse attention only; disk pages cannot be locked. + is permitted. Every lifecycle locks at the hot level, sparse or not. A sparse + lifecycle's hot pool group is host memory, but that is a property of the + allocation rather than of the level, so locking does not distinguish them. - `HELD`: eviction is allowed, but dropping is not. - `DROPPABLE`: both eviction and dropping are allowed. @@ -132,21 +181,13 @@ they establish write completion, migration ordering, and safe reuse across streams. A stream change for an active cache intentionally synchronizes the new stream with the old one. -`batchedLockPages()` takes an explicit destination per page and restores cold -pages before locking, deduplicating pages shared by multiple owners. Resume, -prefetch, and prefix rebasing preserve host-resident sparse history. A partial -prefix copies from the source's actual tier into a private GPU page without -moving a shared host source. Rollback records the original lock level, and -`ScratchSlotLock` remains GPU-only. This does not itself demote GPU history; -offload requires a separate transfer and GPU-slot ownership handoff. - -Prefill admission, prefetch, and rebasing promote a host-locked sparse page into -one shared GPU slot. Promotion waits for all owners' prior work and finished -readers, updates every owner's page indices and metadata version, and releases -the source host slot with a copy-completion fence. Allocation or copy failure -preserves host ownership. Decoding owners retain deferred-offload work so that -`Batch.publish()` retries demotion after the prefill owner enters decode, -suspends, or closes, even without history growth. +`batchedLockPages()` restores cold pages to the hot level before locking them, +deduplicating pages shared by multiple owners. The destination is always the hot +level: an earlier revision carried one per page, for sparse lifecycles that +resided at level 1, and that is gone along with the residency split. A partial +prefix copies from the source's actual tier into a private hot page without +moving a shared source. `ScratchSlotLock` remains GPU-only. This does not itself +demote history; offload requires a separate transfer and slot ownership handoff. ## Ownership and lifetime @@ -163,11 +204,6 @@ KvCache |- KvCacheManager (shared; cache keeps manager alive) `- per-beam/per-block page holders and locks -Batch -|- KvCacheManager (shared) -|- request rows (non-owning; exclusive membership, driven by one owning thread) -`- fixed device metadata and event-protected staging buffers - BlockRadixTree `- roots -> child Blocks (strong ownership through next maps) `- lifecycle page entries (raw observer links) @@ -188,10 +224,6 @@ Eviction controller strong ones merely to simplify access. - `KvCache` keeps its `KvCacheManager` alive. The manager's registry of living caches must not create the reverse strong-reference cycle. -- Closing a request removes its `Batch` row; closing a batch detaches its live - requests without closing them. Removed rows stay dirty until publication clears - them. Batch destruction waits for uploads and recorded readers before freeing - memory; exported arrays can extend the allocation lifetime beyond `close()`. - A committed page is referenced by the radix tree without making the tree its permanent owner. Eviction queues may be the only strong owner of a droppable page, so never store a raw pointer past the operation that obtained it. @@ -295,6 +327,23 @@ physically duplicate the whole prompt; ownership splits at the prompt tail: Keep allocation granularity, stream ordering, and source/destination lifetime valid through asynchronous copies. +## Testing a change here + +The Python suite (`tests/unittest/kv_cache_manager_v2_tests/`) does **not** build this +package's C++ tests, so a green Python run says nothing about whether they still compile: +it loads a prebuilt `.so` and never sees the headers. Removing a symbol the C++ tests use +therefore breaks them silently. Build and run them too: + +```bash +cd cpp/build +for t in kvCacheManagerV2ConfigTest kvCacheManagerV2ColdPageTest \ + kvCacheManagerV2ColdPageCopyTest kvCacheManagerV2TypedIndexTest \ + kvCacheManagerV2HostMemTest kvCacheManagerV2StatsTest \ + kvCacheManagerV2ConcurrencyTest kvCacheManagerV2DigestPoolTest; do + cmake --build . --target "$t" -j"$(nproc)" || echo "$t: BUILD FAIL" +done +``` + ## Interfaces, bindings, and build - `TokenSpan` is non-owning. The caller retains its backing storage for the full diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/CMakeLists.txt b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/CMakeLists.txt index abf8ffb9ed82..8655e1cb7ff5 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/CMakeLists.txt +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/CMakeLists.txt @@ -45,7 +45,6 @@ set(KV_CACHE_HOST_MEMORY_SRCS kv_cache_manager_v2/utils/poison.cpp) set(KV_CACHE_MANAGER_V2_SRCS - kv_cache_manager_v2/batch.cpp kv_cache_manager_v2/common.cpp kv_cache_manager_v2/config.cpp kv_cache_manager_v2/cudaVirtMem.cpp @@ -69,6 +68,10 @@ set(KV_CACHE_MANAGER_V2_SRCS kv_cache_manager_v2/page.cpp kv_cache_manager_v2/storageManager.cpp kv_cache_manager_v2/introspection.cpp + kv_cache_manager_v2/batch.cpp + kv_cache_manager_v2/sparseCache.cpp + kv_cache_manager_v2/sparseCacheKernels.cu + kv_cache_manager_v2/indexConversionKernels.cu kv_cache_manager_v2/kvCache.cpp kv_cache_manager_v2/kvCacheManager.cpp ${SHA256_SRCS}) diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/INDEX_CONVERSION_DESIGN.md b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/INDEX_CONVERSION_DESIGN.md new file mode 100644 index 000000000000..5b8f800d8833 --- /dev/null +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/INDEX_CONVERSION_DESIGN.md @@ -0,0 +1,275 @@ +# Index conversion in Batch + +Status: implemented. The conversion, the index modes, the table specs and the scratch +substitution all live in `Batch`. + +A `Batch` publishes page indices as the cache stores them -- slot ids, level-relative, +with `kBadPageIndex` preserved. Attention kernels need something else: a plain table of +indices into their own address space. Converting between the two is what this document +describes, and the point of putting it in `Batch` is that the conversion is the second half +of a transfer `Batch` already performs. + +## Why it belongs here + +`publish()` already has two stages: a host-to-device fill into `mDeviceStaging` on a stream +of the batch's own, then a copy into the tables kernels read, on the caller's stream. The +second stage is a plain copy only because nothing yet transforms the indices. The runtime's +dense path has the identical shape -- raw device buffer, then a conversion into the table a +captured graph reads -- so the conversion is not a new stage but the one already there, +doing work. + +Folding it in removes a full pass: today the staging is copied wholesale and the runtime +separately reads its own raw table to convert it. One kernel reading the staging once +replaces both. + +## The conversion + +``` +entry = slotId * scale + concatOffset +``` + +`scale` is `getPageIndexScale(layer, role)`, which resolves to the number of buffers +coalesced into that pool's page -- a slot maps to `scale` consecutive page indices, one per +coalesced buffer. `concatOffset` selects among them. + +`kBadPageIndex` passes through unchanged. It has to: a sparse table numbers pure-history +and input-token entries in different levels' pools, so only the sentinel separates an absent +block from slot zero. + +`concatOffset` is **derived, never configured**. A caller names the buffers that must sit in +one table; `Batch` takes the differences of their `getMemPoolBaseAddress()` values and +refuses at construction if the separation is not constant and exactly divisible. The caller +states intent; the offset is a consequence. + +## Three independent axes + +| axis | decided by | +|---|---| +| buffers **within** one spec | always encoded in the entry -- this is the buffer dimension | +| **across layers** | the index mode | +| **across specs**, same layer | equal `scale` | + +The third is what lets a quantised layer's `{K,V}` table serve `{KScale,VScale}` as well: +both pools coalesce two buffers, the pool group mirrors their slots, so the entries coincide +and one table serves both through different base pointers. Mirrored pool management is what +makes this detectable rather than accidental. + +## Index mode + +``` +SHARED one table per life cycle; base pointer is per layer +PER_LAYER one table per layer; base pointer is the shared pool-group base +``` + +Entries are the same formula in both. What differs is how many tables exist and where the +per-layer distinction lives -- in the base under SHARED, in the table's identity under +PER_LAYER. + +Without scratch substitution one kernel covers both modes: PER_LAYER is SHARED writing to +more destinations with a different offset per layer, and a descriptor per table expresses +that uniformly. Scratch is the only thing that needs its own kernel, and only under +PER_LAYER -- it substitutes a different page for one layer, so entries stop being the affine +function of the slot id that the conversion applies. + +Which means: with entries identical across the layers of a life cycle, the only thing that +can make layer *L* differ from *L+1* is **scratch-slot substitution**. That is why +`KvCache::supportsIndexMode(SHARED)` is `!hasScratchSlots()`, and why a life cycle that +cannot hold scratch -- anything that is not an `AttnLifeCycle` with a window, since +`computeScratchRange` returns empty for those -- can use SHARED even where scratch reuse is +enabled globally. + +### Choosing it + +Per life cycle, from the manager, at construction: + +```cpp +//! Const: configuration only, never member state. +//! PER_LAYER -> true +//! SHARED -> true if this life cycle can never hold scratch +//! nullopt if it might +std::optional KvCacheManager::supportsIndexMode(PageIndexMode, LifeCycleId) const; + +Batch(..., std::optional mode = {}); +``` + +An explicit `mode` applies to every life cycle. Unset, each life cycle takes SHARED where +the manager answers `true` and PER_LAYER otherwise -- preferring SHARED, which is the +smaller table. + +Naming SHARED is allowed where the manager answers `nullopt`, and naming it is the only way +to get it there. The manager answers for every request a life cycle could carry, so a +sliding window leaves it undecided; a batch that will hold only decode members knows more +than that, because scratch covers the blocks a prefill chunk adds and a one-token request +adds none. A named mode is therefore a claim about this batch's members rather than about +the configuration, and `publish()` holds the caller to it on every step. A derived mode is +never allowed to outrun the manager, since nothing would be checking it. + +The decision is fixed at construction because a converted table's shape is baked into a +captured graph. It cannot follow a member's scratch state, which varies per step. + +`KvCache::supportsIndexMode(PageIndexMode, LifeCycleId)` exists only to assert: at publish, +for each member and each SHARED life cycle, check the member holds no scratch slots there. +Where the mode was derived, a violation means a request acquired scratch inside a SHARED +batch, which the prefill/decode split says cannot happen. Where it was named, this is the +check that makes naming it safe. The check is an `empty()` test per (member, SHARED life +cycle), so it stays in release builds. + +## Requests and tables + +A spec names buffers that must be addressable from one table under one base -- a kernel +requirement, not an optimisation. Requests are per layer, since a `BufferId` is +`(layer, role)`. + +```cpp +using IndexTableId = StrongIndex; + +//! Buffers concatenated into one table's buffer dimension, under one base. +//! Order defines the dimension. +struct IndexTableSpec { std::vector buffers; }; + +IndexTableId numIndexTables() const; +//! Every spec this table satisfies. Each spec's buffer order is its own mapping +//! into the buffer dimension, so K and KScale both sit at entry 0. +std::vector const& indexTable(IndexTableId) const; +CUdeviceptr indexTableAddress(IndexTableId) const; +//! What a kernel adds to a converted index for this buffer in this table. +MemAddress indexTableBase(IndexTableId, LayerId, DataRole) const; +//! Which table serves this buffer, and its slot in that table's buffer dimension. +std::pair indexTableSlot(LayerId, DataRole) const; +``` + +So `{{K},{V}}` on one table means K and V read identical entries through different bases. +`{{K,V},{KScale,VScale}}` means each pair is concatenated under one base, and the two pairs +share storage. + +`indexTableSlot` is the normal query: a kernel asks where its indices are and at what entry, +binds the address and base, and never iterates the spec list. `numIndexTables` and +`indexTable` answer the introspection question -- how many tables exist and what each covers. + +Every spec merged onto one table has the same buffer count; unequal sizes could not have +coinciding entries. A spec spanning life cycles or scales is a hard error at +construction: no single table can serve it, and splitting silently would break the caller's +assumption that one table addresses the set. + +A merged set's per-layer tables are allocated contiguously, so a caller that wants to bind +them as one strided tensor with a plane per layer can. The descriptor does not mention this, +so a caller that does not care is not obliged to. + +## Buffers + +| | holds | read by | +|---|---|---| +| `mStaging` | raw `[lc][row][block]`, the row state, the per-row leases, the scratch slot tables; pinned host | the upload | +| `mDeviceStaging` | the same block, device side | the conversion | +| `mIndexTableMemory` | the converted tables | the captured kernels | +| `mMirror` | the row state and the leases | the captured kernels | + +The converted tables are their own allocation, carved by a geometry that comes from the mode +and the merge classes, so `indexTableAddress()` hands out offsets into it and a merged set's +per-layer tables stay contiguous. + +Buffers share an allocation only where one transfer can carry the whole of it, since what +that buys is a driver call off the step's critical path. The staged block is one +host-to-device copy and `mMirror` one device-to-device copy, so each is worth packing; the +converted tables are produced by the conversion kernel rather than copied, so packing them +with either would save no call and cost the offset arithmetic. + +`mMirror` carries no raw indices. With sparse life cycles converted too, nothing captured +reads the raw form, so the mirror is the staged block's tail -- the row state and the leases, +the two things a captured kernel reads that are not converted indices -- and `mDeviceStaging` +is the only raw copy on the device. `gpuCoalescedPageTable()` reads it there, which is both +where the truth is and where the conversion reads it. + +That makes the mirror a tail rather than a copy of the block, so a staged element offset +addresses it once rebased by `_historyOffset()`, and the per-step device-to-device copy +carries the row state and the leases alone. The raw indices dominate the block -- a row's +page table is `maxBlocksPerSeq` entries where its row state is a handful -- so what used to +be the larger of the two per-step transfers is now the smaller by orders of magnitude. + +Batch's own device footprint rises: the converted tables carry the buffer dimension, and +under PER_LAYER the layer dimension, where the raw mirror carried neither. The system total +falls, because the converted tables move from the runtime to Batch rather than being added, +and two raw buffers stop existing -- Batch's raw mirror and the runtime's +`_device_kv_cache_block_offsets_input`. + +## What this does to the sparse path + +The copy reaches a cold page as `hostBase + converted * pageStride`, where `pageStride` is +one buffer and `hostBase` is the base of the index table the buffer was declared in. The +per-buffer offset is folded into the converted value, so neither a row's published state, +the arena's shared metadata nor the cache shape carries one. The pool's slot size plays no part: what a converted index counts is whole buffers. +The GPU-side cache stops needing to know that a page coalesces several buffers, which was +never its concern -- it caches a buffer, not a page. + +A run of buffers is copied in one launch from one base, so each buffer reads its own entry +of the table: `entry = buffer - bufferBase`. That is the same `numEntries` dimension a +`{K, V}` spec already produced. + +A table is laid out `[numRows][numEntries][blocksPerRow]`, so a row's entries are adjacent +and a buffer strides between rows by `numEntries * blocksPerRow`. Entry-major would have +been the obvious shape -- it makes one buffer's rows contiguous, which is what the copy +walks -- but it puts `numRows * blocksPerRow` between a request's K and V. Attention +kernels take the pair together: XQA's `kvCachePageList` is +`[batchSize][beamWidth][2][maxNbPagesPerSeq]`, K's page list immediately followed by V's. +Row-major is that shape, so a table can be handed to such a kernel as it stands rather than +repacked per step, and the copy pays an ordinary strided read instead. + +Both producers of page indices read that one table, at different extents. The fetch passes +it from row zero across the batch. The write-back passes it positioned at its own row, +entries still striding by the whole table's rows -- it used to stage a private slice of raw +slot ids instead, which was a second addressing mode for the same kernel. + +Two constraints follow and are checked when the Batch is built, because a run that cannot be +addressed produces wrong page addresses at step time rather than an error: + +- every sparse buffer must be named by some spec -- a fetch with no converted indices + would read the wrong pages rather than fail; +- a `(layer, pool)` run must sit at consecutive entries of **one** spec. Split across + two, a single launch would need two base pointers. + +That shape is not a choice -- it follows from the pool layout -- so +`KvCacheManager::sparseIndexTableSpecs()` hands it out. A Batch still has to be given it: +the declaration stays explicit, and only the re-derivation is saved. + +## Considered: converting per layer into a reused buffer + +PER_LAYER costs one table per layer, where SHARED costs one per life cycle. A layer's +converted indices are dead as soon as its attention has read them, and next step's page +table differs anyway, so the memory looks reusable: convert on demand, one layer at a time, +into a single buffer the caller recycles -- and overlap the conversion with the preceding +layer's MLP or MoE. The batch would take its specs in groups declared at construction, and +a new call would convert one group into a caller-provided buffer. + +It is not worth it, because the memory it reclaims is not the memory that matters. + +What forces PER_LAYER is scratch substitution, and scratch is a prefill phenomenon: +`computeScratchRange` intersects the blocks stale at capacity with the blocks the current +chunk is adding, and in decode the only block being added is the newest, which is still in +window. A decode request therefore holds no scratch, and a decode-only batch answers +`supportsIndexMode(SHARED)` for every life cycle. So the batch that dominates memory -- +many concurrent rows, decode -- already takes the smaller table, and the batch that needs +PER_LAYER is the prefill one, whose row count is small and whose tables are correspondingly +cheap. + +Reaching that does need something, but something smaller: prefill and decode want different +index modes, and the mode is fixed per batch at construction, so they have to be two batches +running in parallel, each publishing its own page table. That is expressible today -- the +attention op already takes `AttentionInputType::ContextOnly` and `GenerationOnly`, so the +two batches are two op calls rather than one, and each call reads its own +`maxBlocksPerSequence`, so the two page tables need not agree on width. What it does cost is +the host-side readers that index one aggregated block-offset buffer by absolute batch row -- +speculative decoding's draft swap and token relocation, the sparse path, KV-cache sharing -- +each of which would have to choose a table by phase. + +So the cheaper route to the same saving is the mode split, not on-demand conversion. If this +is revisited, note that the two interact: with prefill and decode separated, the prefill +batch is the only PER_LAYER one left, and per-layer conversion would be reclaiming a +fraction of an already small allocation. + +## Not in scope + +Exposing the conversion as a plugin interface. The shape it would take is a two-method +interface -- one declaring what the batch must materialise, one consuming the staging -- so +that a model can substitute a lazy per-layer converter or one fused into its attention +kernel. Neither exists yet, and the default must be the safe one regardless, so the +interface waits until something needs to vary. diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/SPARSE_ATTENTION_DESIGN.md b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/SPARSE_ATTENTION_DESIGN.md new file mode 100644 index 000000000000..b5fa4bf1061b --- /dev/null +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/SPARSE_ATTENTION_DESIGN.md @@ -0,0 +1,1417 @@ +# Sparse attention in KVCacheManagerV2 + +Target: the **decode instance of a disaggregated deployment**. + +**Implemented**: the resolve, copy, index and flush-list kernels; `Batch`; the per-request +reservation and device state; the batch-wide write-back as `Batch::flushSparseUnits()`, +which `resize()` checks has run rather than performing itself; `fetch_sparse_pages` and +`prepare_sparse_write_indices`; sparse buffers across more than one pool; sparsity as a pool +attribute, so one layer may hold both sparse and dense buffers. + +**Pending**: nothing. Compression is a separate feature, configured through +`KvCacheCompressionConfig` rather than through anything here, and needs nothing +sparse-specific -- see *Compression*. + +Out of scope, by decision rather than by omission: prefill, which a separate instance serves, +and speculative decoding — see *Supported scope*; and per-KV-head selection, which no +per-`(request, layer)` resident set can express — see *Device-side cache*. Kernel fusion +remains future work; so do multi-beam requests. `ISparseUnitGather` is retired — see +*Extraction interface*. + +Read `AGENTS.md` in this directory first for the storage, locking and concurrency model +this builds on. + +## Problem + +Sparse-attention decoding reads only a small, per-step, per-layer subset of a sequence's +history. Keeping all of that history in GPU memory wastes the resource the sparse +algorithm exists to conserve: a request's memory bill still grows with its full context, +so decoding hits a capacity wall long before it runs out of compute. + +The goal is to hold history on the host tier and let GPU residency be bounded by the +selection width rather than by context length. + +Two existing properties make this cheap: + +- The host tier is registered `CU_MEMHOSTREGISTER_PORTABLE | CU_MEMHOSTREGISTER_DEVICEMAP` + (`utils/hostMem.cpp:305`), so a kernel can read a host page directly. + `isGpuAccessibleMemory` (`storageManager.cpp:138`) already relies on this for the + cold-page codec, which encodes and decodes against host memory with no staging buffer. +- Selection output is a device array. If everything the fetch path needs is also on the + device, the host never enters the loop, and the operation is stream-ordered and + capturable in a CUDA graph. + +## Vocabulary + +| term | meaning | +| --- | --- | +| page | the KVCM2 allocation, locking and migration unit: `tokens_per_block` tokens | +| buffer | one layer's slice of a page; several coalesce into one page | +| selection unit | what the backend selects; a factor of `tokens_per_block`, never larger | +| in-flight | blocks above the watermark: written but not yet flushed | +| units per page | selection units a page holds, `tokens_per_block / sparse_selection_unit_tokens` | +| pure history | a block whose ordinal is below the history watermark set by `resize()` | + +## Scope + +Sparse-attention families differ in what they select, and the difference is load-bearing: + +| family | selection unit | models | +| --- | --- | --- | +| MSA | block, 16 × 128 tokens (`minimax_m3/common.py:70`) | MiniMax-M3 | +| NSA, MoBA | block | DeepSeek-V4-Flash, Kimi | +| Quest | page | training-free, any model | +| DSA | **token**, k = 2048 | DeepSeek-V3.2, V4.1, GLM-5.x | + +Block-level families work with the selection unit equal to one page and need nothing +special. DSA needs a sub-page unit, which this design supports through the selection-unit +mechanism below, at the cost of a finer cache entry. First implementation should target +the block-level families; DSA follows without an API change. + +DeepSeek-V4's compressed pools are worth noting as a middle case: per-layer +`compress_ratios` with `compressed_block_sizes = tokens_per_block // ratio` +(`deepseek_v4/cache_manager.py:370`), so at ratio 128 a compressed page holds exactly one +entry and the selection unit is already a page. + +## Shape of the design + +1. **No new tier or medium.** What varies is the *memory kind of a pool group*: a sparse + life cycle's level-0 pool group is backed by host memory, a dense one by GPU memory. + Level 1 stays an ordinary codec target for both. +2. **`BufferConfig.is_sparse` is a per-buffer flag** and selects that backing. A sparse + buffer requires a lifecycle of its own. Sparse buffers of the same size still coalesce + with each other. +3. **Every page of a sparse buffer locks at level 0**, like every other page. Sparse pages + are ordinary hot pages — coalesced, addressed by page index, reachable through + `get_mem_pool_base_address`, and visible to the transceiver — that happen to live in + host memory. +4. **All GPU residency is the per-request cache**, which is a dedicated arena rather than a + storage tier. A sparse buffer has no separately managed GPU pages: both the block being + written and the history a step selects occupy cache cells, drawn from the same lease and + replaced by the same LRU. +5. **The host page is authoritative once the watermark passes the block.** A block reaches + its host page through `Batch::flushSparseUnits()`, which the caller runs before the + `resize()` that raises `history_length` past it; `resize()` checks that flush has + happened rather than performing it. After it the cache entry is evictable. Before it the + entry holds the only copy and is pinned. +6. **Extraction is pluggable**, because KVCM2 buffers are opaque bytes and a sub-page unit + cannot be located without layout knowledge. + +## Locking and flush + +`acquirePageIndex` already writes the page's *current* slot id with no level assumption +(`page.cpp:401`): + +```cpp +int old = kvc->updateBasePageIndex( + mUser.beamIndex, mUser.ordinal, mUser.lifeCycle, slotIdToPageIndexValue(pg.slotId())); +``` + +so `getBasePageIndices()` reports hot slot ids with no change, and because every page of a +sparse buffer locks there, the table is **uniform**: every entry is a level-0 slot id, as +for a dense buffer. Nothing resolves against the cache arena, and no caller has to know +which rule applies to which ordinal. + +A block's data reaches its host page by a **flush**: a copy on the Batch's stream that does +**not** evict the GPU entry, which stays cached and merely stops being pinned. + +The governing rule is that **a block must be flushed before it can be selected**. Selectable +means below the watermark, so the watermark move is the deadline: everything it passes must +already be in its host page. + +### Who flushes + +`Batch::flushSparseUnits(newHistoryLengths)`, given the per-row history the caller is about +to set. The Batch records it; `resize()` then *checks* against that record rather than +flushing anything itself. + +The flush used to live inside `resize()`, on the grounds that it was the single chokepoint +through which `setHistoryLength` and `setCapacity` both pass. That reasoning held, but it +bought the chokepoint at the price of hiding an ordering requirement: the write-back reads +converted page indices, which only a publish produces, so a flush buried in `resize()` was +silently depending on when the last publish happened. Naming the flush makes that a stated +precondition of a call the caller issues. + +It is also what lets the write-back be **batch-wide**. Per request, it was one scan and one +copy launch per pool *per row*, plus two synchronous host-to-device copies to tell the +kernel which units moved. Driven from the Batch it is one scan over `(layers, rows)` and one +copy launch per pool for the whole step, with a single small upload of per-row ranges. + +The check `resize()` performs is `newHistory <= recorded`, not equality: a step may resize a +row more than once, and a row whose growth failed rolls back to a history below the one that +was flushed for it. Flushing early is always safe — a flush is a copy, and a block flushed +before a failed resize stays correctly flushed. Only flushing late is not. + +`commit()` needs no flush site of its own: the manager maintains +`num_committed_tokens <= history_length <= capacity`, so a committed block is already at or +below the watermark and was flushed before the watermark passed it. + +`history_length` therefore keeps two jobs: it defines selection eligibility, and it is the +flush deadline. + +### How much a step may cross + +The write-back's work list is sized by `maxTotalInputTokens`, the tokens one step may ingest +across every row, which the runtime already bounds for its own batching. A row advancing `t` +tokens crosses at most `t / unitTokens + 1` units -- the `+1` for the partial unit its +watermark starts inside -- so one layer's list needs + +``` +maxTotalInputTokens / unitTokens + maxBatchSize +``` + +entries, and `flushSparseUnits` rejects a move whose units exceed that, summed over rows. + +The bound is an **aggregate** one, and it is usable only because the list is compacted. Were +it indexed by row, each row's slice would have to allow for that row consuming the entire +token budget, and the product over rows would be hundreds of times larger. The same property +is what makes the guard right for in-flight batching: one prefilling row may cross the whole +step's worth of units while every other row crosses none, which a per-row limit would refuse +no matter how idle the rest of the batch was. + +Left unset, the bound falls back to `maxBatchSize * maxInFlightUnits`. That is today's decode +behaviour -- a step advances capacity one token per row, so a row crosses at most one unit +boundary -- and it is what the sparse path currently supports; see "Supported scope". + +### Ordering + +`flushSparseUnits` stages its per-row ranges on a helper stream and launches the write-back +on the Batch's stream: + +``` +helper: waitEvent(flushRead) // the previous step finished reading the ranges + cuMemcpyBatchAsync(ranges -> device, CU_MEMCPY_SRC_ACCESS_ORDER_DURING_API_CALL) + staged = record(helper) +main: waitEvent(staged) + scan the flush lists, then one write-back copy per pool + flushRead = record(main) +``` + +The helper stream buys both things, dominated by the first. `DURING_API_CALL` consumes the +host array *before the call returns*, so issuing the copy on the Batch's stream would hold +the host until the previous step's graph replay reached it; on a helper stream the source is +captured at once and the call returns promptly. The device write then overlaps that replay +too, since the main stream reaches its wait on `staged` only after draining what it already +had queued — though with two ints per row that part is worth a microsecond, where the +host-side decoupling is worth whatever the replay has left. + +The kernel stays on the Batch's stream because the cells it reads are written there. The +wait is a stream wait, not a host sync: the hazard is the previous step's kernel still +reading the same destination, which is ordinary stream ordering. (Do not set +`CU_MEMCPY_FLAG_PREFER_OVERLAP_WITH_COMPUTE` on that copy — it breaks `DURING_API_CALL` +semantics; driver bug 6718200, see `batchedPageCopy.cu`.) + +The write-back reads the index table as of the **last publish**. That is current enough +because the move is clamped to the capacity the step started from, so the blocks it touches +were allocated by an earlier resize and published then. + +### Unflushed entries carry no state + +An entry holds either a fetched copy of a history page or the live, still-growing block. +The first is read-only by construction — fetch only ever targets ordinals below the +watermark — and the second is written by definition. So the two populations are +distinguished by the ordinal alone: + +``` +unflushed <=> resident_ordinal >= history_length / unit_tokens +``` + +Nothing records this. History is monotonic (`resize()` rejects a decrease), so an ordinal +that was history stays history and a resident ordinal at or above the watermark cannot be a +stale fetched copy. The flush recorded for a `resize()` covers exactly +`[old_history, new_history)`, so +the flushed prefix always equals the watermark and there is no second high-water mark to +keep. A per-entry dirty bitmap would be derived state that can only ever disagree with the +ordinal it was derived from. + +Prefill and decode share one manager under in-flight batching, so the unflushed set is not +one entry: it is every in-flight block, which during prefill is the whole chunk. Two facts +keep this from being a general write-back cache: + +- **Unflushed entries are private.** Only finalized pages enter the radix tree, so an + unflushed block has exactly one owner: no aliasing, no cross-request coherence, no + concurrent writers. +- **The unflushed set is at most one unit**, not discovered — see `max_input_length` below. + Reservation covers it, so a fetch never meets a cache with nothing to evict. + +Eviction eligibility is therefore one comparison against the watermark. There is no +evict-an-unflushed-entry path to implement: such an entry is simply not a candidate, and +the reservation guarantees enough candidates remain. + +## Per-request reservation + +Two separate reservations, easily confused. + +**Hot pages.** Every block of the sequence has a locked hot page, so the count is +`divUp(capacity, tokens_per_block)` — the existing per-sequence locking logic, entirely +unchanged. Only the backing medium differs. This is the payoff of uniform locking: nothing +new to size, and no notion of "history pages" distinct from the rest. + +**Cache cells.** `max_cached_units` is the whole cache, counting both populations: + +``` +cells = max_cached_units one buffer's cells +lease = num_buffers * divUp(cells, expansion) contiguous cells a request takes +``` + +The two populations are disjoint — the selection is below the watermark, everything in +flight is above it — so the cache must hold a full selection *plus* every in-flight unit at +once. That is the only constraint, and it is what makes the second population derivable +rather than declared. + +An in-flight unit's data lives only in its cache cell until the watermark passes it, so the +cache must hold every unit above the watermark. Within the supported scope that is **at most +one**: a step advances capacity one token past history, and one token straddles at most one +unit boundary. + +*Sizing*: two to four times the selection width is the usual starting point. A cache only as +wide as the selection retains nothing between steps, so what it holds above the selection is +what survives to be hit next step — it is all reuse, not headroom. + +### `max_input_length` + +A bound on `capacity - history_length`, i.e. on how far capacity may outrun the watermark. +**One token**, and read-only. + +It is a property of the supported shape rather than of how wide a cache is. A sparse request +contributes exactly one query token per step, because the fetch resolves one selection per +request while an attention kernel indexes its selection per query token; the two agree only +at one token each. This is the guard that keeps prefill and speculative decoding out — see +*Supported scope* below. + +It is one token rather than anything derived from the reservation's width. What the cache +*could* pin says nothing about what a step produces: the fetch cannot serve a multi-token +step at any width. + +The lease therefore depends on configuration alone. It is taken when the cache is seated in +a `Batch` and returned when it leaves, so a request holds cells exactly while it is being +stepped: one still receiving its disaggregated history, or sitting a step out, holds none. +Seating cannot fail for want of a cell, because the arena is sized from the same +`constraints` the Batch takes its rows from -- a free row therefore has a free lease, and +the fetch path never has to handle exhaustion. Because every request of a life cycle needs +exactly the same number of cells, the arena is a fixed-size-block free list rather than a +general allocator. + +Leaving returns the lease with every entry free, so it discards whatever the GPU cache held. +`flushSparseUnits()` is what preserves anything worth keeping, and a Batch that churns its +membership every step cold-starts every request's cache. + +`resize()` still checks the gap against this bound. It is the host-side counterpart of the +resolve kernel's in-flight overflow report: one catches a request outgrowing its cache +before the step, the other if it somehow reaches the device. + +### Why the cache is a flat array of cells + +A lease is one contiguous run of buffer-sized cells, and an entry resolves to one by +arithmetic alone: + +``` +cell = buffer * divUp(max_cached_units, expansion) + entry / expansion +address = lease_base + cell * buffer_bytes +unit = entry % expansion which unit within the cell +``` + +No page table, no slot indirection, no coalescing. Coalescing earns nothing inside the +cache because different layers select different units, so a slot's cells would never be +filled together; the cache took whole coalesced slots only while it drew them from the +storage allocator, which deals in slots. A dedicated arena has no such constraint, and +dropping it removes a page lookup from both the copy and the index kernel. + +The consequence worth stating is that **a cache cell is not a page**. It holds one buffer's +unit, identified by `(buffer, entry index)` rather than by a block ordinal. This breaks an +assumption that runs through KVCM2 — that a page is both the unit of allocation *and* the +unit of identity — which is precisely why the cache is not a storage tier: it would have to +answer "which block does this page hold", and it has no answer. + +## Device-side cache + +Per `(request, layer)`, managed entirely by kernels: a resident set, an LRU queue and a +free list. + +**Scope is per `(request, layer)`, never per KV head.** Every head of a layer shares one +resident set, because a selection unit is the tokens it spans across all of them. An +algorithm that lets different heads select different units is not expressible here, and this +is a real exclusion rather than a hypothetical: Rocket's selection indices are +`[num_kv_heads, total_tokens]` and MiniMax-M3's are `[tokens, num_kv_heads, topk]` — both +choose per head. Serving them would need a resident set per `(request, layer, head)`, which +multiplies the metadata and the reservation by the head count. + +**Scope is per request.** Two requests sharing a prefix do not share a cached copy. This +matches HiSparse and closes a soundness hole by construction: entries are reachable only +through their own request's page table, so a host slot can never be recycled underneath a +surviving entry. A global, provenance-keyed cache would give cross-request sharing for +free via the radix tree, but needs generation counters or an invalidation pass; deferred. + +**Replacement is true LRU**, not "evict whatever was not selected this step". Entries that +hit are promoted above newly fetched misses, so a unit selected repeatedly outranks a +first-time selection and one-off selections leave first. Retention across steps is the +point: HiSparse reports a large miss-rate reduction from holding hot entries beyond the +current top-k, which is only possible when the cache is larger than the selection width. + +**Residency is tracked by a reverse map** — a dense array indexed by slot `0..N-1` holding +the resident unit ordinal, N entries per `(request, layer)` — not a forward map from +ordinal to slot. The two scale with different things: + +| | size | at 1 M context, 256 requests, 64 layers | +| --- | --- | --- | +| forward, ordinal → slot | batch × layers × **max_ordinals** | hundreds of MB, and unbounded in context | +| reverse, slot → ordinal | batch × layers × **N** | proportional to cache already bought | + +The forward map is proportional to context length, the quantity this feature exists to +stop paying for. The dense reverse form is also what the LRU order and free list want. + +### Probe index + +Answering "is this ordinal resident, and where" against a dense reverse map would be O(N) +per selection. The fetch kernel therefore builds a hash in shared memory at the top of each +launch and discards it at the end. Two properties follow from it being ephemeral: there are +no deletions, so no tombstones and none of the maintenance that made a persistent device +hash map unattractive; and nothing has to stay consistent across steps or graph replays. + +**The hash is over the required ordinals** -- the in-flight units followed by the selection +-- and the resident set is then streamed past it. So the table is sized by `k` rather than +by the cache's capacity `N`: + +```cpp +// Index the required ordinals, then match the resident set against that index. +for (int32_t i = threadIdx.x; i < required; i += kThreads) { + int32_t ord = i < inFlight ? historyUnits + i : selected[i - inFlight]; + if (ord < 0) continue; + uint32_t pos = slotOf(ord, mask); + int32_t key = atomicCAS(&keys[pos], kEmptyKey, ord); + while (key != kEmptyKey && key != ord) { pos = (pos + 1U) & mask; key = atomicCAS(...); } +} +__syncthreads(); + +for (int32_t s = threadIdx.x; s < entries; s += kThreads) { // the resident set + int32_t ord = residentOrdinal[s]; + cls[s] = kClassUntouched | (ord >= historyUnits ? kPinned : 0); + previousLru[s] = lruOrder[s]; + if (ord == kFreeOrdinal) continue; + ... find ord in the table ... + if (key == ord) vals[pos] = s; // the resident entry claims its query +} +``` + +Open addressing with linear probing, capacity `next_pow2(2 * required)` so the load factor +stays at or below one half, which also guarantees an unsuccessful probe terminates on an +empty entry. Keys and values are separate arrays rather than packed 64-bit entries: a +`__syncthreads()` separates build from match, so no thread reads an entry mid-insert and the +value needs no atomicity, keeping the CAS 32-bit. Multiplicative hashing rather than +identity, because ordinals are small and dense, and dense runs make an unsuccessful probe +walk the whole run. + +The resident sweep is also where per-entry recency and pin state are cached, so the pass +the match needs is the pass the eviction ranking needs; see "Capacity and eviction". + +**This reverses an earlier decision, deliberately.** The index was first built over the +*resident* set, keyed by ordinal and valued by slot, sized `next_pow2(2N)`. That was argued +for on the grounds that hashing the selection would leave the miss list unknown until the +resident walk finished. The argument does not hold: the resident-keyed build had to walk +every slot and cross a `__syncthreads()` before its first probe, so the misses were never +known any earlier. What the choice really bought was one fewer pass -- the probe returned a +hit's slot directly, where the selection-keyed form needs a third pass of size `k` to read +the slots back out of the table. + +What the reversal costs, and why it is still worth it: + +- **Shared memory falls from O(N) to O(k).** At the DSA geometry the table drops from 8192 + buckets to 2048. This is the whole point: `N` now costs linear per-entry scratch rather + than a `next_pow2`-rounded table, so the quantity bounding how large a cache one launch + can serve is the selection width, not the capacity. +- **A third pass of size `k`** resolves hits from the table. Shared-memory only, and small + against the resident sweep. +- **The index build no longer precedes the dependency wait.** The resident-keyed build read + only the layer's own residency table, so it could run while the producing kernel was still + draining; hashing the selection cannot, because the selection is what the producer + produces. Only shared-memory initialisation now sits above the wait. + +Measured on DSA at the geometry above, the resolve kernel fell from 1.468 to 0.833 ms per +decode step summed across 21 sparse layers. + +A reservation-only call -- `prepare_sparse_write_indices`, which has no selection -- skips +the table entirely and runs a separate kernel that scans the recency order for its single +in-flight ordinal. See "Writing the in-flight unit". + +## Extraction interface + +KVCM2 buffers are raw bytes with no layout information, so locating a selection unit inside +one requires the owner's knowledge. Two mechanisms, chosen per buffer, serving two +different situations. + +### Declaring where a unit sits + +A unit is the tokens it spans across *every* KV head, so where its bytes sit follows from +two numbers on `BufferConfig`, both consumed once at configuration: + +```cpp +std::optional sparseSelectionUnitTokens; // unset: a unit is a whole page +int sparseUnitChunks = 1; // pieces a unit is gathered from +``` + +`sparseUnitChunks` is the page's dimension order expressed as a count: + +- **1** for a sequence-major page (`[token][head][dim]`), where a unit's tokens carry their + heads with them and the whole unit is one contiguous range. Also 1 for any single-head + page — a compressed or pooled latent — where the two orders are the same bytes. +- **numHeads** for a head-major page (`[head][token][dim]`), where each head contributes + its own run and a unit is that many strided pieces. + +Everything else follows, because a page is densely packed: + +``` +unitsPerPage = tokens_per_block / sparseSelectionUnitTokens +chunkStride = size / sparseUnitChunks +chunkBytes = size / (sparseUnitChunks * unitsPerPage) +unitStride = chunkBytes +``` + +The four-field `SelectionUnitLayout` the copy kernel reads is derived from these at +configuration, so the kernel is unchanged and the config surface is two integers. The +derivation assumes dense packing, so a page with an alignment gap — between heads, or +between units — is not expressible; no model in this repository has one. + +**Why this is not `tokens_per_block_override`.** That knob reinterprets a buffer's page for +*storage*, and stretches its public page index space by the same factor. Sparse selection +needs only the cache's own granularity, so it gets its own field and leaves the buffer's +pages, page indices and index space on the manager's `tokens_per_block`. Setting both on +one buffer is rejected. `units per page` and `PageIndexConverter`'s `expansion` are +distinct quantities: the sparse code says `unitsPerPage`, and `expansion` belongs to the +index converter alone. + +**What `sparseUnitChunks > 1` is for.** Every sparse model in the repository today passes 1: +DeepSeek V4 asserts `num_kv_heads == 1`, and GLM's kpool, QSA's index-K and MiniMax's +index-K are all single-head and sequence-major. The field earns its place on QSA, whose +selection is shared across heads (its index-K has one head) while the buffer attention then +gathers from is head-major with ordinary GQA heads — the one combination that needs strided +units. Rocket and MiniMax-M3 under MSA are also head-major with several heads, but both +select *per head*, which this cache excludes for a different reason (see *Device-side +cache*); strided units would not rescue them. + +### The gather interface, and why there isn't one + +Earlier revisions planned an `ISparseUnitGather` — a per-buffer, stream-taking device call +that would extract a unit from a representation arithmetic could not address. It existed for +exactly one reason: the fetch read cold pages, so a codec that compressed or permuted them +put the unit somewhere a layout value could not describe. + +Reading hot pages removes the reason. The hot representation is KVCM2's own, always in +ordinary slot layout, and a codec by construction only transforms data on its way *out* of +that tier. A compressed level 1 is therefore invisible to the fetch, and no plugin boundary +is needed on this path at all. + +This also retires the LTO-IR/nvJitLink fusion work that existed only to inline such a +gatherer back into the resolve. The extension point is gone rather than deferred, which is +worth more than the code it would have taken to build. + +### Kernel structure + +The descriptor path crosses **no plugin boundary**. KVCM2 resolves selections, updates LRU, +allocates cells and performs the copies itself, because a layout is data it can interpret +rather than code it must call. + +It is two launches — a resolve and a copy — but that split is internal and chosen, not +forced: the resolve writes a miss list the copy consumes, which is what lets the PCIe +copies be issued before the recency bookkeeping. Fusing them would be a local optimisation, +not an interface change. + +There is no second path to fuse with, so the structure is final rather than a first stage. + +## Compression + +Compressing history is attractive — it is the bulk of the footprint, written once and read +rarely — and with sparse history on the hot tier it needs **nothing sparse-specific**. A +codec already transforms pages crossing the hot boundary; a sparse life cycle's pages now +cross it like any other, so level 1 can hold a compressed form and the existing +`IKvCacheColdPageCodec` does the work. + +The fetch is unaffected: it reads hot pages, which are never the compressed form. + +`SparseHistoryPassesThroughTheCodec` pins this: a sparse life cycle's page is encoded on +demotion and decoded on promotion, with the test first asserting that the life cycle's hot +pool group really is `HOST_MEM` so that it cannot pass by accidentally describing a dense +one. Nothing in the fetch path is involved, which is the point -- it reads hot pages, which +are never the compressed form. + +**The two are decoupled deliberately.** Compression is configured by +`KvCacheCompressionConfig`, not by `SparseAttentionConfig`, on the stated grounds that +compression changes *which KV is stored* while sparsity changes *which KV is read*. A +storage-bound method provides a native cold-page codec that KVCM2 retains and invokes +(`_util.py`, `cold_page_codec_provider`); an iteration-driven one runs on the +resource-manager cycle instead. Neither knows anything about this package. + +Two algorithms ship: `quantization_for_cold_page`, which stores host and disk pages as +NVFP4 while leaving the active GPU cache alone, and `triattention`, which evicts tokens +periodically during decode. + +So "can sparse history be compressed" reduces to "does a sparse page reach the codec", and +it does. Worth stating precisely which pages, because two true sentences sound like one: +sparse history is *hot* -- host-backed, but hot -- and a cold-page codec compresses pages +on their way to a **cold** tier. Sparse history is therefore compressed once demoted, at +which point the codec applies with nothing sparse-specific, and not before. + +What remains is the question the codec has always posed and this design does not answer: +whether a lossy codec is acceptable for KV history, and under what accuracy budget. + +## Fetch + +One call per **layer** per step, on the caller's stream — not per layer group. A layer's +selection is normally computed from the previous layer's output, so a step's selections are +not all known at its start; the call has to be the unit the caller can actually fill. + +Within a layer, every sparse buffer shares the call. They select the same units, so they +share one metadata row, and one resolve yields the entry list that addresses each buffer's +own cell. Their agreement is structural rather than a consequence of running identical +logic on identical input. + +### Selection groups + +A model that runs one indexer across a run of layers hands them all the *same* selection +every step. Resolving it per layer computes the same answer as many times as there are +layers, each against its own resident set and recency order — and they cannot diverge, +because `initialiseLease` clears every row alike and identical input then evolves +identically. So the copies are of one answer, paid for in kernel launches and in +`entriesPerSelection` of metadata per layer. + +Those layers instead name one `selection_id`. Members of a selection share **one** resident +set, one recency order, one in-flight reservation and one miss list. The cells stay per +layer and are *mirrored*, exactly as a layer's K and k-scale mirror today: the entry index +is shared, the bytes behind it are not. `entriesPerSelection` is the metadata's first +dimension, and `layerSelection[layer]` is where a layer reads it — the identity where every +layer selects for itself. + +The id is opaque, not a `BufferId`. KVCM2 never uses *which* buffer produced a selection, so +naming one would state a claim it does not rely on, and that claim goes stale for a model +with several indexer buffers per layer or one that computes a selection without caching +anything. The validations worth having are expressible over an opaque id: members must +agree on `num_selected_blocks`, `max_cached_units` and `unit_tokens`, must share a life +cycle (the metadata they share is the life cycle's), and an explicit id may not name a layer +that set none of its own — silence is never shared. + +The call then splits in two. `resolveSparseSelection()` runs **once** per selection, and +`fetchResolvedPages()` runs **per layer**, each copying the same misses into its own cells +and writing its own index list. `fetchSparsePages()` remains as the fused form and serves a +layer whose selection is its own; it rejects a shared one rather than mis-serving it, since +a second resolve would count the first's admissions as hits, issue no copies, and leave the +second layer's cells unfilled. + +The copy kernel needed no change for any of this. It consumes a work list plus a buffer +range, and nothing in it depended on which layer produced the list. Neither did the +write-back: `sparseResolveKernel` is the only writer of `residentOrdinal`/`lruOrder`, and +`sparseFlushListKernel` only reads them — so collapsing the resolve is sufficient, and a +per-layer flush against shared metadata stays correct. + +Layers that are *not* tied keep a call each, which is what starts the first layer's copies +as early as possible instead of waiting for the whole run. + +The tied layers stay separate layers. The alternative — declaring their buffers as one +`AttentionLayerConfig` with distinct roles — also yields one metadata row and one call, but +it merges the copies too, so nothing can be fetched until every member's bytes are ready. +It also misdescribes the model: GLM-5.2's four indexer-sharing layers are four layers, and +only the full-indexer one declares an `INDEX_KEY` buffer, the other three being wholly +sparse. It is the **life cycle** that needs a dense buffer, not each layer. A selection +group keeps the layers apart for the copy and shares only what is actually shared. + +So a call covers one layer's buffers and the copy needs no way to name a subset of them. +A caller wanting a layer's copies split further would be asking for the inverse of the +configuration choice above, having already declared those buffers as one layer. + +Each call then: + +1. read the selection: unit ordinals, uniform width, `-1` padded; +2. build the probe index over the required ordinals — in-flight units, then the selection; +3. stream the resident set past it: a match is a hit with its slot, the rest are misses; +4. allocate a slot per miss from the free list, evicting LRU when empty; +5. extract and copy each miss from the device-mapped host page; +6. promote hits above new misses in recency order; +7. write the output index list: one index per selection slot, in the selection's own + order, `kBadPageIndex` where a slot is padded. + +The output is compact rather than a table indexed by unit ordinal, because that is the +shape every sparse attention consumer in this repository takes. A token-sparse kernel is +handed the index list in place of its page table — `attentionOp.cpp` assigns +`kvPageIdxPtr = sparse_attn_indices`, and the trtllm-gen K/V TMA descriptors are rebuilt +as a flat token array so one "page" is one token — while a block-sparse kernel takes the +list beside a page table it still uses. None indexes a table by ordinal, so producing one +would only oblige the caller to compact it again. + +One index per buffer, though. A layer's buffers share the resolve but not their cells, so +the list carries a buffer dimension: `[row][buffer][selection]`, which is also the single +tensor covering K and V that attention kernels tend to want. + +Steps 5 and 6 are independent. Because the probe in step 3 resolves misses directly, the +copies — the only step that touches PCIe, and so the long pole — can be issued before the +recency bookkeeping runs rather than after it. + +No host involvement, so the call enqueues and returns. Every shape is static across steps. + +The selection is not the only thing that needs a cell. The block being written needs one +too, and needs it before the selection takes any victim. Both are resolved by the same +probe, over one virtual list: the in-flight ordinals `[history_units, capacity_units)` +first, then the selection. An in-flight unit differs in exactly one way — a miss allocates +its cell but never joins the copy list, because its host page holds nothing to fetch until +the watermark passes it. + +### Writing the in-flight unit + +The unit between the two watermarks is the one the model is about to produce. `fetch_sparse_pages` +makes it resident and pins it, but never reported *where*, so nothing could write it. +`prepare_sparse_write_indices` closes that: it runs the resolve with no selection and converts +the in-flight slot into the same index space the fetch produces. + +It is a separate call because the orders differ. A fetch needs the selection, which the +indexer computes from the previous layer's output; a write needs only the watermarks, so it +can run first — which is the order the model requires. Idempotent while the watermarks hold: +the resolve probes residency before taking a victim, so the fetch's own resolve finds the +entry and reuses it. + +Nothing is fetched for it. An in-flight unit's host page is not written until the watermark +passes it, so the resolve excludes in-flight ordinals from the copy list (`ordinal < +historyUnits`) — copying would read stale bytes over what the model just wrote. + +`out` is `[batch][buffers]`, one entry per request: `resize` bounds `capacity - history_length` +to a token, so at most one unit is in flight. `BAD_PAGE_INDEX` is an ordinary outcome rather +than an error — with 4-token units a unit completes one step in four, which is also the only +step on which a ratio-4 compressor has produced anything to write. + +### Supported scope: one query token per request + +The fetch is a decode-phase call. It resolves **one selection per request**, while a sparse +attention kernel indexes its selection **per query token** — trtllm-gen's +`ptrSparseMlaTopKLens` is `[numTokensQ]`, "where each tokenQ has a corresponding topK +length". The two agree only when a request contributes one query token per step, so that is +the supported shape, enforced by `max_input_length` and by `fetch_sparse_pages` requiring one +selection row per batch row. + +Excluded, deliberately: + +- **Prefill.** Served by a separate prefill instance under disaggregated serving, which needs + no sparsity here: it computes the prompt and hands over the KV cache. +- **Speculative decoding and MTP.** Generation, but several query tokens per request per + step. + +Both are excluded for the same reason, and it is not the row count. The resolve makes units +*resident* in a request's cache, which holds `max_cached_units`; residency is a property of +the request, not of a query token. N query tokens each choosing their own units have a union +of up to N times the selection width, so the reservation could no longer cover a call, and a +unit evicted to make room for a later query token would invalidate an index already returned +for an earlier one. Taking one row per token would make that reachable, not fix it. + +**What else is sized by this, and would need revisiting.** The restriction reaches further +than the fetch. `maxInputLength()` returns one token, which makes +`_sparseCacheMaxInFlightUnits` one, which in turn sizes the write-back's per-row flush list +and bounds how far one watermark move may travel. Those adapt on their own if the constant +rises -- `flushSparseUnits()` refuses a span wider than the headroom, and the scan asserts +the list cannot overrun. What does *not* adapt is throughput: the write-back is tuned on the +assumption that a row has at most one unit to move, so it is given a single warp per +request. Correct for any count, but once a flush list is long that warp is the long pole and +wants the same treatment as the fetch. + +This is a difference in kind from the sparse implementations in this repository, not a gap in +degree. Their sparse prefill is real — the context path reuses the generation kernel with +per-token indices — but their indices point straight at a GPU-resident KV pool, and the +in-chunk tokens are reached by writing K/V to the cache before attention and forcing a local +window in as ordinary indices. There is nothing to stage, so no per-request cache is needed. +Here history is on host, which is exactly what makes prefill hard. + +### Arriving with a history: the bulk watermark jump + +A resumed request is the ordinary case under disaggregated serving, and it moves history from +zero to the whole prompt in one `resize()`. Nothing is written back, because nothing was ever +in flight: those units reached their host pages directly rather than through a cell. + +The write-back is therefore bounded by the capacity the resize started from, not by how far +history travelled: + +``` +flush range = [old_history, min(new_history, old_capacity)) in units +``` + +Only a unit below the old capacity can have been in the cache at all. Without that clamp the +range is the whole prompt, which no reservation could cover, and the resize fails on a +write-back that had nothing to do. + +That jump has to happen **before the request joins a Batch**. `Batch::resize` refuses a +history beyond the capacity its row already had, because capacity granted by a call makes +room for tokens written after it — so history can only cover tokens whose space was already +there. A resumed request sets capacity and history together from nothing, which is the one +shape that rule excludes. Seeding it on the way in costs nothing: nothing was in flight, so +there is nothing to write back and no converted table is needed yet. + +### Capacity and eviction + +Capacity is guaranteed by reservation, not back-pressure: a fetch can always satisfy its +own layer, because the reservation covers a full selection width. Within a step a fetch may +evict entries of layers already processed but never one a later kernel still needs, which +LRU recency delivers — layer N's entries are most-recently-used when layer N+1 fetches. + +An entry is evictable only when its unit ordinal is **below the watermark**. That is what +protects the block being written, and it is a correctness predicate rather than a policy +hint — but it is derived, not stored: see *Unflushed entries carry no state*. + +### Preconditions the caller owns + +Selections must name only pure-history units; `num_blocks` carries each request's eligible +count. + +A step's order is the caller's too, and it is now explicit rather than implied: + +1. `flushSparseUnits(newHistory)` — writes back what the watermark is about to pass. It + reads the index table as of the last publish, which covers the blocks in range. +2. `resize(...)` per row, or `Batch::resize` — refused if it would take history past what + was flushed. +3. `publish()` — converts this step's indices for the fetch to read. +4. the fetch, inside the captured graph. + +Steps 1 and 3 both read or write Batch-owned device state, so neither may be skipped on a +step whose watermark moves. A history advance with no flush before it is refused; a flush +with no publish behind it would write through indices for blocks that did not exist yet, +which the clamp to the starting capacity is what prevents. + +Naming an ineligible one is no longer a memory-safety problem. The table is uniform, so the +ordinal resolves to a real host page of the same request — but an unflushed block's host +page has no valid contents yet, so the result is stale data rather than unrelated memory. +A debug check against the flushed state can catch it; silence would be a wrong answer, not +corruption. + +### Measuring the hit rate + +A hit is a selected unit the cache already held; a miss is one the fetch copied from its +host page. The ratio is what a deployment tunes the cache size against, so it is collected +rather than inferred. + +**Both figures come from the resolve, because only the resolve can produce either.** The +obvious denominator — rows times selection width, which the host knows for free — counts +things that are neither. A selection entry may be padded, and padding is negative and lives +on the device, chosen by whatever selected the units. An in-flight unit is resolved in the +same pass but is not a fetch: its host page is not written until the watermark passes it, +so there is nothing to copy and nothing to miss. Counting either as a hit inflates the rate, +and does so worst where histories are short — early decode, short prompts — which is +exactly the regime someone reads the rate to understand. + +So the resolve counts the selection entries that named real history, in the pass that +already classifies them. The miss half is `claimed`, the run the row takes in the compacted +list, which it computes anyway. The two are produced together and agree by construction on +what they count; hits are their difference and are never stored. + +**Each row adds its own contribution, rather than anything reading a total afterwards.** +This is the part that is easy to get wrong. The step's counters — `work.total` and anything +beside it — are cleared once per `publish()`, not once per fetch, which is what lets one +clear cover every layer without a non-kernel node between each resolve and the work before +it. A consumer downstream that read those counters and folded them into a cumulative pair +would therefore be correct only if every layer were fetched exactly once per publish. That +is what a step does, but the interface permits more, and a second fetch of a layer would +silently re-add the first one's totals. Accumulating per CTA from the values the row holds +in registers has no such precondition, costs two atomics per row, and needs no ordering +against anything. + +It also keeps the whole mechanism inside one kernel: nothing in the copy, no programmatic +dependency to reason about, and no need to exclude the write-back — which shares the copy +kernel and whose moved units are due to be written rather than units the cache failed to +hold. The in-flight resolve contributes nothing on its own, since it has no selection and +its ordinals sit above the watermark. + +64-bit counters: a step at B = 256 with a 2048-wide selection resolves millions of units per +layer, so a 32-bit pair would wrap within a few hundred steps — and a wrapped miss count +makes the hit rate *improve*, which is the failure a reader would trust. + +**There is no enable switch.** The cost is two atomics per row per fetched layer. A switch +would cost nothing to run either, but it would bake into a captured graph, so turning +statistics on mid-run would keep reporting zeros until the graph was captured again — a +disabled feature indistinguishable from a broken one. + +`peekHitMissCounter()` reads the pair on a temporary stream and synchronises only that, so +it never stalls the step. The pair may be a few launches stale, which does not move a ratio; +a caller wanting the two to agree exactly synchronises the device first. Because the read +races a running fetch and the counters are separate atomics, a copy can land between them — +the result is clamped rather than checked, since a diagnostic must not be able to fail the +cache. Per-row detail is already available separately in `missCount`. + +## One arena, one index space + +The cache arena holds nothing but cells — selected history and the unit being written alike. +That is what lets the fetch emit a **single** index list covering both, and therefore why +existing attention kernels need no change: a kernel still indexes one list against one base. + +Indices are unit-granular, and over the buffer's **arena region** — one per hot pool, see +*Sparse buffers across more than one pool*: + +``` +index = cell * expansion + unit + = (lease * cells_per_lease + slot_in_pool * cells_per_buffer + entry / expansion) + * expansion + entry % expansion +``` + +where `slot_in_pool` is the buffer's index among those of the life cycle sharing its pool, +and `cells_per_lease` is that pool's buffers times `cells_per_buffer`. + +A sparse buffer therefore has **two** base addresses, answering different questions: + +| accessor | names | consumer | +| --- | --- | --- | +| `get_mem_pool_base_address` | the buffer's pages, hot tier | transceiver, KV transfer, anything addressing blocks | +| `get_sparse_cache_base_address` | the cache arena region of the buffer's pool | the attention kernel, against the fetch's index list | + +`get_mem_pool_base_address` keeps meaning "this buffer's pages" for every buffer, sparse or +dense, which is what lets the transceiver work unchanged on hot, coalesced pages. The arena +is a second allocation holding the only GPU-resident copy a kernel can read, so it gets a +second accessor rather than overloading the first. The consumer's side is otherwise +unchanged: one base, one index list, the expanded index space `PageIndexConverter` already +produces. + +`tokens_per_block_override` plays no part in this. It is the index converter's own +mechanism: it scales `BufferConfig.size` and stretches the buffer's public page index space +by `expansion`, while allocation, locking and migration stay on the global +`tokens_per_block`. It is rejected on a sparse buffer, which says what it means with +`sparse_selection_unit_tokens` and leaves its index space alone. `unitsPerPage` and +`expansion` are distinct quantities and the sparse code never uses the latter. + +## Batch + +Every device array here is indexed by batch row, and every call must agree on which row a +request occupies. `Batch` owns that: membership, row order, the page table, `num_blocks`, +and the staging mirror. It is global, not per layer group, and carries one page table per +layer group. + +It also owns three things the sparse path depends on: + +- **the stream.** Taken at construction; joining binds the member to it, and moving a member + off it is refused. A member's own work and the Batch's transfers have to be ordered + against each other, and a stream per call left two sources of truth for which one that is. + It is also what makes the write-back's ordering free rather than checked, since a member + records its events on the stream the flush is launched on. +- **the converted index tables** (`INDEX_CONVERSION_DESIGN.md`), which both the fetch and + the write-back read. +- **the write-back**, as `flushSparseUnits()` — see "Locking and flush". + +**The upload is whole-table by decision.** `publish()` stages every row and sends the table +in one `cuMemcpyHtoDAsync` on the upload stream, which overlaps compute. Sending only the +rows whose locks moved would replace that single driver call with one per row, and the cost +of those launches outweighs the bytes they save. The table is large -- four layer groups, +about 32 MiB a step at a 1M `max_seq_len` and B = 256 -- but size is not the figure that +decides it. + +**Dirty rows — planned, not built.** What dirty tracking would still save is the host-side +restaging, not the transfer: a `_KVCache` knows its `Batch` and row (it does, for detach), +any lock-changing operation could mark that row dirty, and `Batch` could restage only those, +following the existing `mark_stats_dirty` / `clear_stats_dirty` / +`get_dirty_stats_kv_cache_ids` pattern. Until then a step restages the whole table, which is +a cost rather than a wrong answer. + +Rejecting direct mutation instead was considered and rejected: `commit` changes locks and +runs every decode step, so strictness would require a batched twin for every mutator before +anything could ship. Dirty tracking degrades into a slower step rather than a wrong answer. +A debug-mode check can still flag direct mutation while batched. + +**Batched resize** resizes every member and publishes once at the end. It may initially +forward to `_KVCache.resize()`; the committed contract is row order and +publish-on-completion, which survives the later coalescing work. The result is per-request, +not a single flag, since `_KVCache.resize()` already rolls back per cache; a table published +after a partial failure must reflect the post-rollback state. + +**Rows are slots.** Removing a request leaves a hole to be reused rather than shifting later +rows, which would invalidate the whole table and destabilise the shapes graph capture needs. +The `_torch` layer already works this way with `index_mapper_capacity` and a `copy_idx` +gather. Membership is exclusive: a row already occupied is refused rather than rebound, +because nothing downstream would notice the substitution. + +**A row owns its member.** `mRows` holds a `shared_ptr`, so a seated cache cannot be +destroyed under the Batch and a caller may keep only the Batch; `removeKvCache()` hands the +reference back, which makes "the Batch kept it alive" and "here it is back" one operation. +The binding pinned every cache ever added for the Batch's whole life and released none of +them, which was the wrong shape as well as a leak: what a row needs is a reference for +exactly as long as it is seated. `KvCache` points back with a raw `mBatch`, so there is no +cycle, and `_detach()` -- the path for a member destroyed while seated -- is now unreachable +by construction and says so with a fatal check rather than carrying teardown that could +never run. + +## Sizing + +Per sparse buffer: **`num_selected_blocks`** (the selection width `k`) and **N** (the +reservation upper bound, in the same units). + +N should be `2k`–`4k`, not `k`. `k` is the correctness floor — it guarantees a fetch can +satisfy its layer — but it leaves no room for retention, and retention across steps is +where the hit rate comes from. HiSparse finds `B ∈ [2k, 4k]` the useful range. + +There is no separate budget for in-flight blocks. The sparse pool's GPU tier holds nothing +but cache cells, and an in-flight block occupies one of the same cells the selection draws +on, so `max_cached_units` covers both: + +``` +cache reservation = max_batch_size * N (N in [2k, 4k]) +in-flight headroom = 1 unit one query token per step +``` + +with `max_batch_size` from `constraints`. The surplus above `k` is retention, not headroom, +so `N` is sized by one demand rather than two: retention wants +`2k`–`4k`, and a prefill chunk wants enough cells to pin the whole chunk at once. + +## Python surface + +Declared or planned in `tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi`: + +- `BufferConfig.is_sparse: bool`, `BufferConfig.num_selected_blocks: int`, the reservation + bound, `BufferConfig.sparse_selection_unit_tokens: int | None` and + `BufferConfig.sparse_unit_chunks: int` — construction-time, per buffer. +- `KVCacheManager.is_sparse(layer_id, data_role) -> bool`. +- `Batch` — membership, row order, device arrays, `resize()`, dirty-row publication. +- `KVCacheManager.fetch_sparse_pages(batch, buffers, selected, out, stream)`. +- `Batch.sparse_out_slots(buffers) -> list[int]` — the slice of `out` each buffer's + indices land in, for both the fetch and the write. + +Publishing the page table is `Batch`'s job. It must not reuse the dense path's block +offsets: those are scaled and fold `BAD_PAGE_INDEX` to slot zero, whereas a sparse table +numbers history and input-token entries in different levels' pools and needs the sentinel +to distinguish an absent block from slot zero. + +Device arrays are typed by a `DeviceArray` protocol over DLPack, so callers pass torch +tensors without this package depending on torch. + +## Implementation notes + +Slot release is already level-aware: `Page::~Page()` releases with the page's own +`cacheLevel`, and the `kHotLevel` in `ScratchSlotLock::unlock()` is correct, since SWA +prefill scratch is GPU-only by definition. + +- **`_lockHeldBlocks` carries no level hazard.** The OOM rollback for `resize` re-locks at + the hot level, which is where every buffer locks, + so a page table's entries and the level its pages lock at cannot disagree. Nothing carries + a per-target lock level: there is no `queryLockLevel`, no `BatchedLockTarget::cacheLevel`, + no `StaleBackup::cacheLevel` and no `residencyLevel`. +- **`_unlockStaleBlocks` is not the flush hook.** It exists for SWA blocks outside the + window, whose content is *dead*; sparse history is live and will be read. Same trigger and + range arithmetic, different disposition — routing sparse layers through it drops the data + instead of flushing it. It also gates on `windowSize.has_value()`, so for a full-attention + sparse layer it fires on nothing. +- **`_shortcutSetHistoryLength` skips the flush for exactly our configuration.** It reports + "unchanged" when no SWA stale range moved, and for a full-attention layer + `windowSize.has_value()` is false, so nothing is ever reported as moved: it assigns + `mHistoryLength` and returns true, and `resize()` returns early before reaching its body. + A flush placed in that body never runs. The shortcut must either carry the flush or be + declined while a sparse flush is pending. +- **Hook `resize()`, not `setHistoryLength()`.** The latter is `resize(nullopt, hist)` and + `setCapacity()` is `resize(cap, nullopt)`, so `resize()` is the only place the watermark + moves — and a caller may move capacity and history in the same call. +- **`commit()` is not a flush site at all.** The manager maintains + `num_committed_tokens <= history_length <= capacity`, so a committed block is already at + or below the watermark and was flushed by the `resize()` that passed it. Where `commit()` + advances the watermark it does so by calling `resize()`. + +Also: the host pool is far larger than the GPU pool, so a page index must be widened to 64 +bits before multiplying by the page stride. MSA already does this +(`triton_sparse_decode.py`); QSA and XQA should be audited. + +## Planned: a selection unit size per layer + +A layer declares its fetch path with `sparse_cache`, and layers sharing a life cycle must +currently declare the same one. The arena is per life cycle and its metadata is a single +stride — a selection's rows live at `residentOrdinal + slot * entriesPerSelection` — so one +entry count serves every layer, and `cellsPerBuffer` follows from `max_cached_units` and the +unit size together. + +That is an implementation limit, not a property of the design. Lifting it: + +- **Size a cell from the largest unit of the life cycle**, or simply from a whole page. +- **A layer with a smaller unit subdivides that cell.** This is the mechanism a cell already + uses: it holds `unitsPerPage` units, addressed as `cell * unitsPerPage + unit`. A layer + whose unit is half the cell's just has twice the divisor. +- **Require unit sizes to be powers of two**, so every such division is exact and a buffer's + bytes split without remainder at any level. +- The resident set and the recency order are **already per selection**, so they need no + change; only the entry count and the divisor become per layer rather than per life cycle. + +Until then a mismatch is rejected at construction. It was previously taken as a `max` across +the life cycle, which quietly served one layer's shape to another. + +## The output list merges across mirrored pools + +The arena mirrors the hot pool group, and a lease takes the same cell index in every pool +of it. So **mirrored pools** -- equal `buffersInPool`, hence equal `cellsPerLease`, with +`cellsPerBuffer` one value for the whole shape -- resolve a selection to the same index. +`out` carries one slice per distinct index table entry rather than one per buffer, with +the caller applying each pool's own base. + +For a life cycle whose pool group holds a dense indexer and two sparse pools -- `{K,V}` +and `{kscale,vscale}` -- `sparseIndexTableSpecs()` emits one spec per `(pool, layerRow)` +and `Batch::_buildIndexTables` folds the two into **one** table, because their `scale` and +`concatOffsets` agree. The layer's four buffers therefore take the slices `[0, 1, 0, 1]`: +`K` with its scale, `V` with its scale. A scale shared by `K` and `V` together is a pool +of one buffer, whose `cellsPerLease` differs, so it keeps a slice of its own: `[0, 1, 2]`. + +`Batch::sparseOutSlots` is the query, and `fetchSparsePages`, +`fetchResolvedPages` and `prepareSparseWriteIndices` all follow it. They have to agree, or a caller would need one +rule for reading a fetch and another for a write. A run whose slices an earlier one +already wrote is not launched at all -- the saved work and the saved writes are the point +-- and the entries of `out` no buffer maps to are left untouched rather than packed out, +so `out`'s shape stays `len(buffers)` however the buffers happen to merge. + +The grouping follows the caller's `IndexTableSpec`, so two buffers merge only where one +table was asked for over both. The merge is not read off the dense tables alone, though: +those merge on the hot pool's `scale` and `concatOffsets`, which is not the same quantity +as the arena's geometry -- a pool carrying only some of the group's layers has the same +scale and a different `cellsPerLease`. `sparseOutSlots` checks the arena geometry agrees +before sharing a slice, and raises rather than aliasing two different cells onto one. + +Rests on a property only a second lease can observe: +`TestSparseFetch::test_mirrored_pools_resolve_a_selection_to_one_index` uses one, because +at lease zero the `leaseIndex * cellsPerLease` term vanishes and pools of unequal width +are indistinguishable. + +The rest of the arrangement is what makes one base per slice enough: + +- **A base per pool, in both spaces.** `getSparseCacheBaseAddress` returns + `arena.base(poolIndex)` -- "per pool, not per buffer" -- so `{K,V}` share one arena base + and `{kscale,vscale}` another, exactly as the two pools share one hot-pool base each. +- **The within-pool offset rides in the index**, as `slotInPool * cellsPerBuffer`. That is + the concat-offset arrangement, so one base does serve a pool's buffers. +- **The output order is the plan's**, because `_sparseLayerSpan` rejects any other: + "buffers must be exactly this layer's sparse buffers, in the cache's order". + +## Sparse buffers across more than one pool + +A life cycle's sparse buffers need not agree on size. The **quantised shape** — block scales +as a separate, smaller `DataRole` — works, which matters because it is a real configuration +rather than an exotic one. + +Pools partition a life cycle's buffers by size, and the copy addresses one pool with one set +of strides: `bufferBytes`, the hot slot stride and the unit layout are all per pool. The +design follows from making the arena **mirror the hot pool group** rather than inventing a +layout of its own: + +- **One region per pool**, sized `maxLeases * buffersInPool * cellsPerBuffer * cellBytes`. + A lease is an index rather than a cell offset, because a pool's cells-per-lease is its own. + A pool index no sparse buffer of the life cycle uses keeps its place in the list and sizes + to nothing: the arena is indexed by the storage pool index rather than by a dense numbering + of its own, and an empty region costs only the entry that says where the next one starts. +- **Buffers ordered `(layer, pool)`.** A layer's buffers stay consecutive, which the fetch + needs to name them as a range; and within a layer a pool's buffers are consecutive too, + which the copy needs for the same reason. `bufferLayer` stays non-decreasing. +- **One launch per maximal run of consecutive buffers sharing a pool.** When the buffers are + alike that is a single run spanning the whole life cycle, so the uniform case pays nothing + for the general one; the quantised shape yields one run per layer and pool. The fetch + clips each run to the layer it resolved; the write-back takes them whole, because its work + list is per layer and a run already spans as many layers as share its pool. +- **`get_sparse_cache_base_address` is role-dependent**, naming the region of the buffer's + pool. It already took a role. + +Every per-pool quantity is a launch scalar — `deviceBase`, `bufferBytes`, `cellsPerLease`, +`layout` — so the kernels keep their shape. What they gained is `bufferSlotInPool`, a +per-buffer map alongside `bufferLayer`, replacing the assumption that a buffer's position +among *all* the life cycle's buffers is also its position in its pool. A per-buffer offset +within the cold page sat here too, until the converted index absorbed it; see +`INDEX_CONVERSION_DESIGN.md`. + +Buffers sharing a pool must still share a selection unit layout. A pool fixes a cell's size +but not how its units are strided within it, and the copy walks a run with one set of +strides; the check is per run rather than per life cycle. + +**The runs share one launch.** Nothing orders one pool's copy against another's — their +cells are disjoint and the converted page table they share is read-only — while each is +small and latency-bound, so issuing them back to back on one stream would expose each one's +latency in turn. So one grid runs them all, and what differs between them lives in a +**device-resident descriptor**. + +`gridDim.y` walks the buffers of every pool as **one flat range**, with a constant map from +buffer to pool. The alternative — pools on one grid axis, buffers-per-pool on another — +has to size the buffer axis for the widest pool, so every buffer a narrower pool lacks +costs an idle block. Pools of one life cycle are often alike, and then the two are the same +shape; they stop being alike as soon as a pool holds more buffers per layer than another, +which a key-and-value pool beside a scale pool already does. + +The descriptors are built once, with the Batch. Every field comes from the buffer plan and +the index tables, both fixed at construction, so nothing is written per step and a captured +graph replays a copy with no host involvement at all. The fetch needs one row per layer — a +layer has its own index-table entry and its own slice of each pool's buffers — while the +write-back spans every layer at once and needs only the one. + +The vector width is part of the descriptor, chosen from the layout's **sizes alone**. The +two base pointers are checked against it where the descriptor is built, because a pool whose +base is misaligned is a configuration error worth naming rather than a case to carry into +the kernel. A block switches on the width once and calls a device function templated on it; +`blockIdx.z` is block-uniform, so the switch costs no divergence. + +One grid rather than several is also what lets **every** pool carry the launch dependency +below. A per-pool fan-out across streams could not: a programmatic dependency follows a +stream, not an event. + +### Launching before the producer finishes + +All three kernels are enrolled in programmatic dependent launch when the environment enables +it, and in each the order of the prologue is the point: a block does everything that is not +its predecessor's output *above* `cudaGridDependencySynchronize()`, so that by the time the +predecessor drains the block is resident, addressed, and as far along as it can be. + +How far along differs by kernel. The resolve gets least out of it: its probe index is keyed +by the required ordinals, and the selection is exactly what the producer produces, so only +shared-memory initialisation sits above the wait. That was not always so -- an index built +over the resident set read only the layer's own residency table, which nothing before a +resolve touches, so the whole build could run while the predecessor drained. Keying the +table by `k` rather than by the cache's capacity was worth giving that up; see +"Probe index". The copy resolves its pool, its +buffer and its work list from descriptors and block coordinates; the flush-list scan +addresses its row from the staged view. Each then waits for the one thing that is genuinely +produced upstream. + +**The work-list counters are cleared once a step, in `publish()`**, not once per producer. +That is not a micro-optimisation: a clear is a memset, and a memset is not a kernel, so one +standing between every resolve and whatever ran before it would sever the launch dependency +at each of a model's layers -- for DeepSeek-V4, sixty-one times a step. + +What makes one clear enough is that the counters are per layer and per direction, so no two +producers in a step share one. The fetch keeps a *single* list, reused from layer to layer, +because a layer's copy has consumed it before the next layer's resolve overwrites it; only +the counter has to read zero when a resolve starts. Per-layer lists would have cost the +layer count times the list, which is megabytes, against a counter's four bytes. + +A step that somehow skipped `publish()` would have its second resolve add on top of the +first's count, and the claim's bound check would fail rather than quietly overrun the list. + +Global **writes** stay below the wait in all three, even where the datum looks private. A +block that has not yet waited may still be running beside a grid that could read what it +would overwrite, and the cost of keeping the rule is a few stores. + +The producers release as early as they correctly can. The resolve's work list is final at +the end of its eviction loop, and everything after that — the recency rebuild, three block +scans over the resident set — is its own bookkeeping that nothing downstream reads; it calls +`cudaTriggerProgrammaticLaunchCompletion()` there and hands the copy that whole tail as +overlap. The flush-list scan has no such tail, since the list is the last thing it writes, +so its release only brings the copy's launch forward. + +The copy releases immediately after its own wait. That looks wrong and is not: the kernel +that follows it on the stream builds the attention index table from the resolve's slots and +never reads a cell, so it must not be ordered behind a copy it does not depend on. + +## Where sparse data lives + +A sparse buffer's pages are ordinary hot pages that happen to be host memory, device-mapped +so the fetch reads them without a host round trip. Cells come from a per-life-cycle arena, +leased per request; the fetch reads hot pages at `BufferAttr::offset`. + +A sparse deployment places no requirement on its cold tiers. History is hot, so level 1 may +be host, disk, or absent. + +Memory kind is a property of the **pool**, not of the level and not of the life cycle. The +GPU tier allocates a sparse pool from host memory and a dense one from device memory, in the +same group; `AttnLifeCycle` carries no `isSparse`. + +A layer may therefore hold both a sparse buffer and a dense one, which is what a model +wants: an indexer and the data it selects over are one attention computation and belong in +one layer. Sparsity constrains pool formation and nothing above it, so two layers differing +only in sparsity share a life cycle — one layer group, one page table in `Batch`. + +**The codec is the hot↔cold transform**, not the level 0 → level 1 one. +`StorageManager::migrate` computes `srcIsHot = srcLevel == kHotLevel` and performs a plain +byte copy when source and destination are on the same side, invoking the codec only across +that boundary. Sparse history is hot, so it crosses that boundary on its way to a cold tier +like any other page: compression needs nothing sparse-specific. + +The arena is sized from `constraints`, which already declares the batch shape `Batch` derives +its rows from, so the cells are a static reservation rather than an on-demand draw on the +shared pool. + +### Sizing a host-backed pool + +Sparsity is a property of a **pool**, not of a life cycle or of a pool group. A pool is +uniformly sparse or uniformly dense, because buffer grouping keys on `(size, isSparse)`; +the GPU tier allocates the sparse ones from host memory, device-mapped, and reports +`GPU_MEM` for the group either way. + +That it reports GPU is sound because nothing reads the hot side's tier. The codec path keeps +only `coldTier = srcIsHot ? dstTier : srcTier`, and the copy-engine path needs both ends +hot, which never happens; host memory is device-accessible, so a mixed group needs no second +tier. + +**A sparse pool's bytes are excluded from the device quota.** A slot is one row spanning +every pool of its group, so once a group holds at least one device-backed pool, the quota +sizes that pool and the sparse pools inherit the same row count. The quota therefore charges +only the dense pools. + +This is why a pool group must hold **at least one dense pool**: a wholly sparse group would +have no device bytes per slot and so no slot count to derive. That is a validation, not a +heuristic. + +**But the row count is not free.** Inheriting it gave the sparse pools no bound of their own: +their host memory grew with whatever the device quota bought, and the only way to hold it +down was to shrink the device cache that everything else shares. +`KVCacheManagerConfig::hostSparsePoolQuota` is that bound, and the allocator takes it as a +second budget over a disjoint set of the same pools rather than as a clamp applied afterwards. + +The two budgets **partition** a group's pools — dense charged to the device quota, sparse to +the host budget, each in its own allocation grain — so one slot count settles both at once. +`ratioToSlotCountList` therefore takes every pool's size once, plus an index list naming the +pools each budget answers for. Three properties follow, and the first is the reason this is +one pass rather than two: + +- **A group bound by one budget releases what it does not use of the other.** Clamping the + slot count afterwards could only lower it; the device bytes a host-bound group then stopped + using went to nobody, and a group with no sparse pool at all -- which can never be + host-bound -- was left short for no reason. +- **Both budgets and `min_slots` are kept exactly; the pool ratio is what gives way.** The + floors come off the top before anything is shared, so no division of the rest can break one. +- **Their only conflict is a configuration error, and it is caught once.** The floors can + exceed a budget, but that names neither the ratio nor the device quota -- the two things a + rebalance changes. So it is rejected the first time the allocator runs and cannot begin to + fail later. + +The split belongs to the hot level alone and is not a property a level can report about +itself: a cold group is a single pool holding dense and sparse buffers together, and a GPU +tier may itself be a cold level. The pools a quota pays for are therefore derived as the +complement of the ones the host budget names -- never passed alongside them -- so the two +cannot overlap, and no pool can be left for neither to pay. + +#### The quota excludes sparse pools; the ratio counts them + +These are deliberately different currencies, and the asymmetry looks like a bug until the +third party is in view. + +- **The quota** (`roundedNumBytes(cacheTier())`) counts only the pools on the level's own + medium. Host-backed bytes consume none of the device memory it measures. +- **The ratio** (`ratioList()`, and the profiler's) weighs *every* pool, host-backed + included. It has to: it is projected onto cold levels, where a single pool holds dense and + sparse buffers together, and it is compared against the achieved layout, which is measured + the same way. + +The consequence is easy to miss. `KvCacheManager::_needAdjustment` compares the layout's +current ratio against the target `constrainPoolGroupRatio` produces, and rebalances when they +differ. `constrainPoolGroupRatio` derives that target by *running the allocator* and +normalising the bytes it would produce -- so it must run it with the **same budgets**, the +sparse host quota included. Omit the budget there and the target describes a layout the +allocator will not produce: every comparison mismatches, and a rebalance re-triggers on every +check even though the profiler's answer never moved. The ratio, the measurement and the +constraint all have to speak all-bytes for that loop to settle. + +What the asymmetry *does* require is that the division be driven by all-bytes demand, not +that the ratio be re-expressed per budget. Dividing the device quota alone by an all-bytes +ratio does not realise that ratio: a sparse-heavy group takes a device share its device cost +per slot does not warrant, and converts it into more slots than asked for. + +The release step does not rescue this. It returns the device grains a group cannot use +*because the host budget held it* -- so where the host budget is generous, nothing is held, +nothing is released, and the over-share stands. The fix belongs in how the budgets are +divided: demand is computed from each group's all-bytes cost per slot, so the achieved +all-bytes shares track the ratio, and the per-medium budgets act as caps on that demand. +Re-expressing the ratio itself in a single medium's currency would be wrong for the reasons +above -- it has to stay comparable with what `ratioList()` measures and with what a cold +level inherits. + +#### How the division runs + +The allocator is a search over one scalar, and all-bytes is what that scalar buys. + +Write `A(pg, n)` for what `n` slots of group `pg` cost in all-bytes — every pool rounded up +to the grain of the medium backing it, which is exactly `CacheLevelStorage::allMediaBytes()`. For +a *level* of all-bytes spending `L`, each group's demand is + + demand(pg, L) = max( minSlots[pg], maxSlotsForAllBytes(L · r[pg]) ) + +i.e. the most slots whose all-bytes cost fits that group's share of `L`, never below its +floor. `demand` is monotonic in `L`, so each budget admits a largest feasible level, found by +bisection: `levelQuota` for the device quota, `levelSecond` for the host budget. Both must +hold, so the search is run against both and the binding one wins. + +Three details make the result exact rather than approximately right. + +- **Floors come off the top.** Each group starts at `minSlots`, snapped up by `freeRun` to + the largest count costing no more grains in any medium — a count short of what it has + already paid for is not the count a re-derivation from the same bytes returns, and that + disagreement is the grains↔slots round-trip bug in its original form. Host-backed pools + count here even when no host quota is set: a pool nothing is bounding is still real memory, + and letting the snap ignore it grows the floor to whatever the device grain absorbs. +- **A group frozen by one budget releases the other.** When the host budget pins a group, its + slot count stops rising, and the device grains its demand would have claimed are returned to + the search rather than reserved. This is what makes one pass strictly better than a quota + followed by a clamp. +- **A bounded greedy finish spends the remainder.** Bisection lands on a level, not on an + integral allocation; leftover grains in either budget are then handed out one group at a + time, each step choosing the group whose resulting `A(pg, n)/r[pg]` is lowest. That is a + leximin step on achieved all-bytes shares: equal shares mean the achieved ratio *is* `r`, + so raising the currently worst-served group is the move that gets closest to it. + +The fixed-point property falls out of measuring demand and achievement in the same currency: +feed the allocator the normalised all-bytes ratio of its own output and it returns the same +slot counts, which is what stops `_needAdjustment` from re-firing on an unchanged profile. +`test_feeding_back_the_layouts_own_ratio_changes_nothing` pins it, including the case where a +host-backed pool is most of a group's bytes and none of its quota. + +A level's quota measures one medium, which is why host-backed bytes cannot be drawn from it: +they consume none of the device bytes it counts, so charging them would cap sparse history at +a GPU-sized number — the opposite of the point — while leaving those device bytes +unallocated. Attributing them to a *host tier's* quota instead was considered and rejected: +`cacheTiers` is a positional list, so "the host tier" may be several or none, and a hot +pool's size would then change when someone added a cold tier that has nothing to do with it. + +**Coverage of the data path.** `TestSparseDataPath` asserts the bytes themselves on both +sides of the copy: a pattern written to a hot page must appear in the cell the fetch names, +and a pattern written to a cell must appear in its hot page once the watermark passes it. + +This was worth building rather than arguing about. Offsetting the page base by one slot and +re-running showed **57 other sparse tests still passing** -- residency, index shape and +counts are all insensitive to which allocation the copy reads, which is exactly the mistake +the move to the hot tier could have introduced. + +**A cell can never outlive its page's lock.** Moving sparse pages to the hot level made them +eligible for demotion for the first time, which raised the question of whether a page could be +encoded while the cache still held a cell copied from it. It cannot: + +- An `ACTIVE` request locks the pages it requires (`AGENTS.md`, *Per-request cache*), and + `LOCKED` forbids eviction outright (*Page status*). +- Eviction is the only route to migration, and migration across the hot boundary is the only + thing that invokes the codec. +- The lease is held inside the locks' lifetime, not alongside it: seating requires a resumed + cache, which has already locked the pages it requires, and leaving returns the lease. A + cell therefore exists only while its page is locked, with room to spare at both ends. + +`commit()` does downgrade a lock to a holder, but only for a page it is discarding in favour +of an existing tree page -- not one the request goes on to read. + +The one configuration where a live request holds unlocked pages is a **sparse layer with a +sliding window**, where `_unlockStaleBlocks` releases out-of-window blocks. It is safe, and +not by caller discipline: the indexer that produces the selection applies the same window, so +a block the window has made stale is never named. Nothing reads an unlocked page, so nothing +can read one that has been encoded or recycled. + +That is worth stating because it is the *selector*, not this package, that makes it true. +`fetch_sparse_pages` cannot detect an ineligible ordinal — the page table is uniform, so a +stale ordinal still resolves to a real page of the same request. The combination is unusual +in any case: sparsity already bounds how much history a step reads. + +## Settled by construction + +Recorded because each looks like an open question and is not: + +- **`num_blocks` lives on the device, in `Batch`.** Deriving it per call from the watermark + would produce a host value the kernel cannot use, costing an H2D per layer. +- **Eviction never copies.** An entry is evictable only once flushed, so evicting is a + free-list push. Flushes are triggered by the watermark alone, never by eviction pressure. +- **Pinning is not a general scheme.** The pinned set is the in-flight unit, private + because unflushed blocks are not in the radix tree, and bounded to one by + `max_input_length`. Ordering between the rest still comes from LRU recency alone. +- **No stale-key invalidation.** Per-request cache scope makes it impossible. +- **`tokens_per_block_override` is not the sparse granularity lever**, and is rejected on a + sparse buffer. `sparse_selection_unit_tokens` is — see *Extraction interface*. + +## Resolved + +- **This *is* decode-only, deliberately.** Sharing one manager between prefill and decode + is wrong in kind, not in size: a sparse kernel indexes its selection per query token, and a + multi-token step's selections cannot all be made resident in one per-request cache. The + decode instance of a disaggregated deployment is the supported target; see *Supported + scope*. +- **A selection unit never exceeds `tokens_per_block`.** Sub-page units are supported + through `expansion`; a unit spanning several pages is rejected, which retires the + questions of slot contiguity and k-at-a-time eviction. +- **`N` is a uniform upper bound**, not per-request. Ragged device arrays would cost the + static shapes graph capture needs. +- **History never moves backwards.** `resize()` rejects any decrease outright + (`kvCache.cpp`: "History length cannot be decreased"), and the only other assignment + initialises a fresh request from a prefix match. A flushed block can therefore never be + rewound into, with no rule of our own required. + +## References + +- HiSparse: [arXiv:2608.07009](https://arxiv.org/abs/2608.07009), + [SGLang docs](https://docs.sglang.io/docs/advanced_features/hisparse_guide), + [LMSYS blog](https://www.lmsys.org/blog/2026-04-10-sglang-hisparse/) +- MiniMax Sparse Attention: [arXiv:2606.13392](https://arxiv.org/abs/2606.13392) +- DeepSeek-V4: [arXiv:2606.19348](https://arxiv.org/pdf/2606.19348) diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/batch.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/batch.cpp index ce368f81c1e4..7c297b422860 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/batch.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/batch.cpp @@ -16,418 +16,2056 @@ */ #include "kv_cache_manager_v2/batch.h" + +#include "kv_cache_manager_v2/batchedPageCopy.h" + #include "kv_cache_manager_v2/exceptions.h" +#include "kv_cache_manager_v2/indexConversionKernels.h" #include "kv_cache_manager_v2/kvCache.h" #include "kv_cache_manager_v2/kvCacheManager.h" -#include "kv_cache_manager_v2/utils/funcGuard.h" -#include "kv_cache_manager_v2/utils/optionalGilRelease.h" +#include "kv_cache_manager_v2/sparseCacheKernels.h" +#include "kv_cache_manager_v2/storageManager.h" +#include "kv_cache_manager_v2/utils/math.h" + +#include "tensorrt_llm/common/cudaUtils.h" #include -#include -#include -#include +#include +#include namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 { + namespace { -void checkOutsideCapture(CudaStream stream) +//! Largest batch and longest sequence the declared constraints describe, or zero when none +//! were declared. +std::pair deriveBounds(KVCacheManagerConfig const& config) { - CUstreamCaptureStatus status; - cuCheck(cuStreamIsCapturing(reinterpret_cast(stream), &status)); - if (status != CU_STREAM_CAPTURE_STATUS_NONE) + int rows = 0; + int capacity = 0; + for (auto const& desc : config.constraints) { - throw LogicError("Batch publication and reader fences must run outside CUDA graph capture"); + rows = std::max(rows, static_cast(desc.kvCaches.size())); + for (auto const& kvCache : desc.kvCaches) + { + capacity = std::max(capacity, kvCache.capacity); + } } + return {rows, capacity == 0 ? 0 : divUp(capacity, config.tokensPerBlock)}; } } // namespace -Batch::Batch(std::shared_ptr manager, int maxRows, int maxBlocks, int maxBeamWidth) - : mManager(std::move(manager)) - , mMaxRows(maxRows) - , mMaxBlocks(maxBlocks) - , mMaxBeamWidth(maxBeamWidth) +Batch::Batch(KvCacheManager& manager, CUstream stream, std::vector indexTables, + std::optional maxBatchSize, std::optional maxBlocksPerSeq, std::optional mode, + std::optional maxTotalInputTokens) + : mManager(&manager) + , mCudaStream(stream) + , mMaxTotalInputTokens(maxTotalInputTokens) { - KVCM2_API_GUARD(); - if (!mManager || maxRows <= 0 || maxBlocks <= 0 || maxBeamWidth != 1) + TLLM_CHECK_WITH_INFO(!maxTotalInputTokens.has_value() || *maxTotalInputTokens > 0, + "Batch: maxTotalInputTokens must be positive when given"); + auto const [derivedRows, derivedBlocks] = deriveBounds(manager.config()); + mMaxBatchSize = maxBatchSize.value_or(derivedRows); + mMaxBlocksPerSeq = maxBlocksPerSeq.value_or(derivedBlocks); + TLLM_CHECK_WITH_INFO(mMaxBatchSize > 0 && mMaxBlocksPerSeq > 0, + "Batch needs a row and block bound; pass them explicitly when the manager declares no constraints"); + + mRows = TypedVec>(BatchRow{mMaxBatchSize}, nullptr); + + _buildSparseArenas(); + + // Prefer SHARED: one table per life cycle rather than one per layer. The manager + // answers from configuration alone, so this cannot drift with what a member is doing. + mIndexModes = TypedVec(mManager->numLifeCycles(), PageIndexMode::PER_LAYER); + mIndexModeNamed = mode.has_value(); + for (LifeCycleId lcId{0}; lcId < mManager->numLifeCycles(); ++lcId) + { + mIndexModes[lcId] + = mode.value_or(manager.supportsIndexMode(PageIndexMode::SHARED, lcId) == std::optional(true) + ? PageIndexMode::SHARED + : PageIndexMode::PER_LAYER); + } + + // After the modes, which decide where a life cycle puts its base, and before the staging + // block, whose tail is sized by how many scratch slots the tables say a row can hold. + _buildIndexTables(indexTables); + _checkSparseRunsAreAddressable(); + auto tables = _layoutIndexTables(); + auto const scratchElements = tables.value().value_or(0); + _layoutBufs(scratchElements); + mUploadBytes = static_cast(mBuf[BufName::kScratch].end); + + // Page tables and watermarks are published together, so a kernel cannot read one step's + // table beside another step's lengths -- an out-of-range read rather than a stale + // answer. + mStaging = StaticHostMem(mUploadBytes); + auto* staging = reinterpret_cast(mStaging.address()); + // A row that no member occupies must still read as "no such block" rather than as slot + // zero, since the fetch has no other way to tell them apart. + std::ranges::fill(_buf(BufName::kPageTables), kBadPageIndex.value()); + // Watermarks of a free row are zero: a request with nothing eligible. + for (auto const name : {BufName::kHistory, BufName::kCapacity}) { - throw std::invalid_argument("Batch requires a manager, positive dimensions and beam width 1"); + std::ranges::fill(_buf(name), 0); } - auto const apiLock = mManager->lockExclusive(); - mNumLayerGroups = mManager->lifeCycles().size().value(); - size_t const rows = static_cast(mNumLayerGroups) * mMaxRows * mMaxBeamWidth; - size_t const columns = static_cast(mMaxBlocks) + 1; - if (rows == 0 || rows > std::numeric_limits::max() / sizeof(int32_t) / columns) + // No row holds a lease yet, which is what an unpublished batch should read as. Zero + // would name arena slot zero. + std::ranges::fill(_buf(BufName::kLeases), kNoSparseLease); + + // The scratch slot metadata is the block's tail, so one transfer per step carries it + // along with the page tables rather than needing a second. + mScratchStaging = _buf(BufName::kScratch).data(); + if (scratchElements > 0) { - throw std::invalid_argument("Batch metadata dimensions overflow"); + // A zero range leaves every block on the plain path, which is what a member holding + // no scratch looks like. The slots behind the ranges hold the sentinel instead, + // since zero there would name slot zero. + auto const ranges + = 2 * static_cast(mManager->numLifeCycles().value()) * static_cast(mMaxBatchSize); + std::fill(mScratchStaging, mScratchStaging + ranges, 0); + std::fill(mScratchStaging + ranges, mScratchStaging + scratchElements, kBadPageIndex.value()); } - mTableElements = rows * mMaxBlocks; - mTotalBytes = rows * columns * sizeof(int32_t); - mRows.resize(mMaxRows, nullptr); - mDirty.resize(mMaxRows, 1); - cuCheck(cuCtxGetDevice(&mDeviceId)); - CUdeviceptr ptr = 0; - cuCheck(cuMemAlloc(&ptr, mTotalBytes)); - mDeviceMemory.reset(reinterpret_cast(ptr)); - // Two generations cover the usual publication/consumer overlap without - // allocating or registering pinned memory in the publication path. - mUploads.push_back({std::make_unique(mTotalBytes)}); - mUploads.push_back({std::make_unique(mTotalBytes)}); + + // The device staging takes the whole block, so the host-to-device crossing is one + // linear copy that can overlap the previous step. The mirror takes only the tail a + // captured kernel reads: the page tables and the scratch are read from the staging, by + // the conversion, which runs fresh each publish and so needs no snapshot of its own. + mMirror = allocateDeviceMemory(_mirrorBytes()); + mDeviceStaging = allocateDeviceMemory(mUploadBytes); + auto* const deviceStaging = reinterpret_cast(mDeviceStaging.get()); + cuCheck(cuMemcpyHtoDAsync(reinterpret_cast(mMirror.get()), _buf(BufName::kHistory).data(), + _mirrorBytes(), cudaStream())); + cuCheck(cuMemcpyHtoDAsync(reinterpret_cast(deviceStaging), staging, mUploadBytes, cudaStream())); + + // Hands the layout its slice of the block above and lets it build the descriptors, + // which name both that block's head and its scratch tail. + tables.send(reinterpret_cast(_deviceBufAt(BufName::kScratch))); + TLLM_CHECK_WITH_INFO(tables.done(), "The index table layout asked for memory twice"); + + { + // Each generator yields what it needs and then fills the slice it is handed, so + // one allocation serves them all even though the last count is unknown until the + // last generator has run. + std::array layouts{_layoutStepScratch(), _layoutFlushRanges()}; + size_t total = 0; + for (auto& layout : layouts) + { + total += layout.value().value_or(0); + } + mDeviceArena = allocateDeviceMemory(total); + auto* cursor = mDeviceArena.get(); + for (auto& layout : layouts) + { + auto const elements = layout.value().value_or(0); + layout.send(cursor); + cursor += elements; + // One round each: a body that asked again would be destroyed suspended, with + // the rest of its layout silently unrun. + TLLM_CHECK_WITH_INFO(layout.done(), "A layout generator asked for memory twice"); + } + } + + // Last: a descriptor names where in the index tables a buffer's converted indices sit. + _allocateCopyDesc(); + + // Everything above is staged on the Batch's stream. A consumer reading this Batch from + // another stream has nothing ordering it against that, so the state is made visible here. + cuCheck(cuStreamSynchronize(cudaStream())); } -Batch::~Batch() +void Batch::_allocateCopyDesc() { - KVCM2_POISON_ON_EXCEPT( - [this]() + // One descriptor per buffer, serving both directions, for every life cycle. Small + // enough to share an allocation and a lifetime. + std::vector host; + std::vector offset(static_cast(mManager->numLifeCycles().value()), 0); + + for (LifeCycleId lcId{0}; lcId < mManager->numLifeCycles(); ++lcId) + { + auto& scratch = mScratch[lcId]; + if (scratch.selectionWidth == 0) { - OptionalGilRelease const gilRelease; - close(); - auto const apiLock = mManager->lockExclusive(); - mReady.synchronize(); - for (auto const& reader : mReaders) + continue; + } + auto const runs = _sparseCopyTemplates(lcId); + auto const plan = mManager->_sparseBufferPlan(lcId); + scratch.numBuffers = static_cast(plan.size()); + offset[static_cast(lcId.value())] = host.size(); + + for (auto const& run : runs) + { + // A run is consecutive buffers of one pool, which can span layers: a pool holds + // every layer's buffer of one size. An index table spec covers one layer's + // run of a pool (sparseIndexTableSpecs), so each buffer names its own table, + // already offset to its entry, and the base that table's indices are relative to. + for (int32_t i = 0; i < run.numBuffers; ++i) { - reader.synchronize(); + auto const buffer = run.bufferBase + i; + auto const& id = plan[static_cast(buffer)].id; + auto const table = sparseIndexRun(id.layerId, id.role); + SparseCopyDesc desc; + desc.hostBase = static_cast(table.base); + desc.hostPageStride = run.hostPageStride; + // Pre-offset to this buffer's entry, so the kernel indexes the row alone. + desc.pageTable = reinterpret_cast(table.table); + desc.pageTableStride = static_cast(table.rowStride); + // Not the row stride: a row now spans every entry of the table, while the + // blocks this run may address are still one row's worth of one entry. + desc.numBlocks = static_cast(mMaxBlocksPerSeq); + // The buffer's own slot in the pool, folded in so the kernel need not read + // it through a view it loads per work item. + desc.deviceBase = run.deviceBase + + static_cast(_sparseArena(lcId).shape().bufferSlotInPool[static_cast(buffer)]) + * run.cellsPerBuffer * run.bufferBytes; + desc.leaseStride = static_cast(run.cellsPerLease) * run.bufferBytes; + desc.cellBytes = run.bufferBytes; + desc.layout = run.layout; + desc.width = run.width; + + auto const width = static_cast(sparseCopyWidthBytes(desc.width)); + TLLM_CHECK_WITH_INFO(desc.hostBase % width == 0 && desc.deviceBase % width == 0, + "A sparse buffer's bases must be aligned to the %llu bytes its layout copies at a time; host " + "base %llu, device base %llu", + static_cast(width), static_cast(desc.hostBase), + static_cast(desc.deviceBase)); + host.push_back(desc); } - for (auto const& upload : mUploads) + } + } + + if (host.empty()) + { + return; + } + mCopyDesc = allocateDeviceMemory(host.size()); + auto const* device = mCopyDesc.get(); + cuCheck(cuMemcpyHtoDAsync( + reinterpret_cast(device), host.data(), sizeof(SparseCopyDesc) * host.size(), cudaStream())); + for (LifeCycleId lcId{0}; lcId < mManager->numLifeCycles(); ++lcId) + { + auto& scratch = mScratch[lcId]; + if (scratch.selectionWidth != 0) + { + scratch.copyDesc = device + offset[static_cast(lcId.value())]; + } + } + // `host` does not outlive this scope, so the copy out of it has to have landed. + cuCheck(cuStreamSynchronize(cudaStream())); +} + +int32_t Batch::_flushCapacity(LifeCycleId lifeCycle) const +{ + // A row advancing `t` tokens crosses at most `t / unitTokens + 1` units -- the `+1` for + // the partial unit its watermark starts inside -- so the batch crosses at most + // `maxTotalInputTokens / unitTokens + maxBatchSize`. The bound is an aggregate one, and + // usable only because the list is compacted: indexed by row, every row's slice would + // have to allow for that row taking the whole budget. + int const perPage = mManager->sparseCacheUnitsPerPage(lifeCycle); + int const unitTokens = perPage > 0 ? mManager->config().tokensPerBlock / perPage : 0; + if (mMaxTotalInputTokens.has_value() && unitTokens > 0) + { + return *mMaxTotalInputTokens / unitTokens + mMaxBatchSize; + } + // Without a budget, a step is assumed to advance capacity by the in-flight headroom in + // every row at once, which is what a decode step does. + return mMaxBatchSize * mManager->_sparseCacheMaxInFlightUnits(lifeCycle); +} + +Batch::LayoutGenerator Batch::_layoutStepScratchForLifecycle(LifeCycleId lcId) +{ + int const width = mManager->_sparseCacheSelectionWidth(lcId); + if (!mManager->hasSparseFetchPath(lcId)) + { + co_return; + } + auto const rows = static_cast(mMaxBatchSize); + auto& scratch = mScratch[lcId]; + scratch.selectionWidth = width; + scratch.maxInFlightUnits = mManager->_sparseCacheMaxInFlightUnits(lcId); + scratch.numLayers = mManager->_sparseCacheLayerCount(lcId); + auto const layers = static_cast(scratch.numLayers); + // One list for the fetch, holding the worst case of every row missing its whole + // selection; one per layer for the write-back, each holding the units the step's token + // budget can carry across a watermark. The fetch's is the life cycle's, not the + // selection's: every layer of the selection just resolved reads it, and the next + // resolve of any selection here overwrites it. + scratch.fetchWork.capacity = static_cast(rows) * width; + scratch.fetchWork.numLists = 1; + scratch.flushWork.capacity = _flushCapacity(lcId); + scratch.flushWork.numLists = scratch.numLayers; + + auto const listEntries = [](SparseWorkList const& work) + { return static_cast(work.capacity) * static_cast(work.numLists); }; + auto const fetchEntries = listEntries(scratch.fetchWork); + auto const flushEntries = listEntries(scratch.flushWork); + auto const outSlot = rows * static_cast(width); + auto const inFlightSlot = rows * static_cast(scratch.maxInFlightUnits); + auto const flushCount = rows * layers; + + int32_t* cursor = co_yield outSlot + rows + inFlightSlot + 3 * fetchEntries + 3 * flushEntries + flushCount; + + auto const carve = [&cursor](size_t count) + { + auto* ptr = cursor; + cursor += count; + return ptr; + }; + scratch.outSlot = carve(outSlot); + scratch.missCount = carve(rows); + scratch.inFlightSlot = carve(inFlightSlot); + auto const carveList = [&carve](SparseWorkList& work, size_t entries) + { + work.row = carve(entries); + work.ordinal = carve(entries); + work.slot = carve(entries); + }; + carveList(scratch.fetchWork, fetchEntries); + carveList(scratch.flushWork, flushEntries); + scratch.flushCount = carve(flushCount); + // -1 so that "no write-back has run" stays distinguishable from "one ran and found + // nothing to move". + cuCheck(cuMemsetD32Async(reinterpret_cast(scratch.flushCount), 0xFFFFFFFFU, flushCount, cudaStream())); +} + +Batch::LayoutGenerator Batch::_layoutStepScratch() +{ + mScratch = TypedVec(mManager->numLifeCycles(), SparseStepScratch{}); + + // The counters stay here rather than in the leaves: they lead the block so that one + // clear covers every life cycle's, which a per-life-cycle slice would scatter. + std::vector> leaves; + size_t mains = 0; + for (LifeCycleId lcId{0}; lcId < mManager->numLifeCycles(); ++lcId) + { + auto leaf = _layoutStepScratchForLifecycle(lcId); + auto const needed = leaf.value(); + if (!needed.has_value()) + { + continue; + } + mains += *needed; + // A counter per layer for each of the two lists, even though the fetch's list is + // one: the list is reused from selection to selection, but the counter has to read + // zero when a resolve starts. Sized by layers because selections cannot outnumber + // them. + mCounterElements += 2 * static_cast(mScratch[lcId].numLayers); + leaves.emplace_back(lcId, std::move(leaf)); + } + if (leaves.empty()) + { + co_return; + } + + int32_t* base = co_yield mains + mCounterElements; + + // Its own allocation: the step's clear covers mCounters, and these have to survive it. + mHitMiss = allocateDeviceMemory(1); + cuCheck(cuMemsetD32Async( + reinterpret_cast(mHitMiss.get()), 0U, sizeof(SparseHitMiss) / sizeof(uint32_t), cudaStream())); + + mCounters = base; + auto* counter = base; + auto* cursor = base + mCounterElements; + for (auto& [lcId, leaf] : leaves) + { + auto& scratch = mScratch[lcId]; + auto const layers = static_cast(scratch.numLayers); + // The fetch takes one counter per selection against its single list, so a resolve + // names its own by offsetting this; the write-back's counters match its lists one + // to one. + scratch.fetchWork.total = counter; + counter += layers; + scratch.flushWork.total = counter; + counter += layers; + scratch.hitMiss = mHitMiss.get(); + + auto const elements = *leaf.value(); + leaf.send(cursor); + cursor += elements; + } +} + +SparseStepScratch const& Batch::sparseStepScratch(LifeCycleId lcId) const +{ + TLLM_CHECK_WITH_INFO(lcId < mScratch.size(), "Batch::sparseStepScratch: no such layer group"); + return mScratch[lcId]; +} + +void Batch::setCudaStream(CUstream stream) +{ + if (stream == cudaStream()) + { + return; + } + // The Batch's own queued work -- the staging copy, the index conversion, the counter + // resets -- is ordered on the old stream alone, and a Batch with no members still has + // it. Each member's move below syncs that member, which is a different question. + CachedCudaEvent handover(reinterpret_cast(cudaStream())); + handover.waitInStream(reinterpret_cast(stream)); + mCudaStream = stream; + for (auto const& kvCache : mRows) + { + if (kvCache != nullptr) + { + // The private form: the public one refuses to move a member at all. + kvCache->_setCudaStream(stream); + } + } +} + +Batch::~Batch() +{ + // Both buffers are about to go, and transfers on the upload stream still name them. + _drain(); + KVCM2_POISON_ON_EXCEPT( + [this] + { + for (auto const& kvCache : mRows) { - upload.completion.synchronize(); + if (kvCache != nullptr) + { + removeKvCache(*kvCache); + } } }); } -void Batch::checkOpen() const +size_t Batch::_rowElements() const noexcept { - if (mClosed) + return static_cast(mMaxBatchSize) * static_cast(mMaxBlocksPerSeq); +} + +PageIndexMode Batch::indexMode(LifeCycleId lifeCycleId) const +{ + TLLM_CHECK_WITH_INFO( + lifeCycleId.value() >= 0 && lifeCycleId < mManager->numLifeCycles(), "Batch::indexMode: unknown layer group"); + return mIndexModes[lifeCycleId]; +} + +void Batch::_buildIndexTables(std::vector const& specs) +{ + // The offset of a buffer relative to a chosen base, in page-index units. Mirrors the + // runtime's kv_offset: a base difference divided by the page stride, which must divide + // exactly or the two buffers do not sit a whole number of pages apart. + auto pageOffset = [this](BufferId const& buffer, MemAddress base) { - throw LogicError("Batch is closed"); + auto const bufferBase = mManager->getMemPoolBaseAddress(buffer.layerId, buffer.role, PageIndexMode::SHARED); + auto const stride = static_cast(mManager->getPageStride(buffer.layerId, buffer.role)); + TLLM_CHECK_WITH_INFO(stride > 0, "Page stride must be positive to offset a buffer within its slot"); + auto const delta = static_cast(bufferBase) - static_cast(base); + TLLM_CHECK_WITH_INFO(delta % stride == 0, + "Buffers sharing an index table must sit a whole number of pages apart; they cannot share one"); + return static_cast(delta / stride); + }; + + for (auto const& spec : specs) + { + TLLM_CHECK_WITH_INFO(!spec.buffers.empty(), "An index table spec names no buffers"); + + auto const& first = spec.buffers.front(); + auto const lifeCycle = mManager->_bufferAttr(first.layerId, first.role).lifeCycleId; + // Substitution replaces one layer's page, so every layer of the life cycle can no + // longer read one table. SHARED is exactly the claim that they can. + // Derived SHARED has to be provable from configuration, because nothing else would + // check it. A named one is the caller's claim about this batch's members, which + // publish() tests against each of them, so it is allowed to outrun what the manager + // can prove about requests in general. + bool const mayHoldScratch = !mManager->supportsIndexMode(PageIndexMode::SHARED, lifeCycle).has_value(); + TLLM_CHECK_WITH_INFO(!mayHoldScratch || mIndexModes[lifeCycle] == PageIndexMode::PER_LAYER || mIndexModeNamed, + "A life cycle that can hold scratch slots needs PER_LAYER index tables; its layers stop sharing one " + "set of converted indices as soon as a scratch slot is substituted. Name SHARED on the Batch to " + "assert no member of it will hold scratch there -- as a decode-only batch does -- and publish() " + "will hold you to it"); + auto const scale = mManager->getPageIndexScale(first.layerId, first.role); + auto const mode = mIndexModes[lifeCycle]; + + // Under SHARED the base carries the per-layer offset, so the spec's own offsets + // are relative to its first buffer. Under PER_LAYER the base is the pool group's + // and carries none, so every offset is absolute within the slot. Derived the same + // way every reader derives it, so the offsets below and the base a caller is + // handed later cannot disagree. + auto const base = _specBase(spec); + + std::vector offsets; + offsets.reserve(spec.buffers.size()); + for (auto const& buffer : spec.buffers) + { + auto const& attr = mManager->_bufferAttr(buffer.layerId, buffer.role); + // The reference converter emits `expansion` entries per block. Nothing in the + // tree sets tokensPerBlockOverride, and every attention backend refuses an + // expansion of anything but one. + TLLM_CHECK_WITH_INFO(attr.expansion == 1, + "Index conversion does not implement page index expansion, so a buffer with a tokens-per-block " + "override cannot have converted index tables"); + TLLM_CHECK_WITH_INFO(attr.lifeCycleId == lifeCycle, + "An index table spec spans life cycles; their slot ids differ, so no one table can serve it"); + TLLM_CHECK_WITH_INFO(mManager->getPageIndexScale(buffer.layerId, buffer.role) == scale, + "An index table spec spans page index scales; their converted indices differ, so no one table " + "can serve it"); + offsets.push_back(pageOffset(buffer, base)); + } + + // Specs coincide when their entries do. Under PER_LAYER the table is per layer, + // so specs for different layers never merge however alike their entries look. + auto const sameTable = [&](IndexTable const& table) + { + if (table.lifeCycle != lifeCycle || table.scale != scale || table.concatOffsets != offsets) + { + return false; + } + return mode == PageIndexMode::SHARED || table.served.front().buffers.front().layerId == first.layerId; + }; + + auto found = std::ranges::find_if(mIndexTables, sameTable); + if (found == mIndexTables.end()) + { + mIndexTables.push_back(IndexTable{{}, lifeCycle, scale, offsets}); + found = std::prev(mIndexTables.end()); + } + IndexTableId const tableId{static_cast(std::distance(mIndexTables.begin(), found))}; + found->served.push_back(spec); + + for (size_t slot = 0; slot < spec.buffers.size(); ++slot) + { + auto const& buffer = spec.buffers[slot]; + BufferSlot const resolved{tableId, static_cast(slot)}; + auto const [it, inserted] = mBufferSlots.emplace(buffer, resolved); + // The base is not compared because it cannot differ here: two specs reach one + // table only when their concatOffsets agree, and an offset is measured from + // the spec's own base, so equal offsets at equal entries mean equal bases. + TLLM_CHECK_WITH_INFO(inserted || (it->second.table == resolved.table && it->second.entry == resolved.entry), + "A buffer appears in two index table specs that resolved differently. Specs may " + "not overlap: name a buffer in one spec only, and read a subset of a table by " + "slicing it at the entry indexTableSlot() gives that buffer"); + } } } -void Batch::checkPublished() const +Batch::LayoutGenerator Batch::_layoutFlushRanges() { - checkOpen(); - if (!mPublished) + mFlushedHistory = TypedVec(BatchRow{mMaxBatchSize}, 0); + + bool anySparse = false; + for (LifeCycleId lcId{0}; lcId < mManager->numLifeCycles(); ++lcId) { - throw LogicError("Batch metadata changed; publish before consumption"); + anySparse = anySparse || mManager->sparseCacheMaxUnits(lcId) > 0; } + if (!anySparse) + { + co_return; + } + + auto const elements + = 2 * static_cast(mMaxBatchSize) * static_cast(mManager->numLifeCycles().value()); + mFlushRanges = co_yield elements; + // An empty range asks for nothing, which is what a row that never flushes should read. + cuCheck(cuMemsetD32Async(reinterpret_cast(mFlushRanges), 0, elements, cudaStream())); + mFlushRangeStaging.assign(elements, 0); +} + +int32_t const* Batch::flushFromUnits(LifeCycleId lifeCycle) const +{ + auto* base = mFlushRanges; + return base + 2 * static_cast(mMaxBatchSize) * static_cast(lifeCycle.value()); +} + +int32_t const* Batch::flushToUnits(LifeCycleId lifeCycle) const +{ + return flushFromUnits(lifeCycle) + mMaxBatchSize; } -void Batch::markDirty(int row) noexcept +int Batch::flushedHistory(BatchRow row) const { - mDirty[row] = 1; - mPublished = false; + TLLM_CHECK_WITH_INFO(row >= 0 && row < mMaxBatchSize, "Batch::flushedHistory: row is beyond max_batch_size"); + return mFlushedHistory[row]; } -int Batch::add(KvCache& cache, std::optional row) +void Batch::flushSparseUnits(std::vector const& newHistoryLengths) { - KVCM2_API_GUARD(); - auto const apiLock = mManager->lockExclusive(); - checkOpen(); - if (&cache.manager() != mManager.get() || cache.isClosed()) + TLLM_CHECK_WITH_INFO( + newHistoryLengths.size() == mRows.stdSize(), "Batch::flushSparseUnits: one history length per row"); + if (!mFlushRanges) + { + // No sparse life cycle, so nothing is ever in flight and the record is vacuous. + for (BatchRow row{0}; row < mRows.size(); ++row) + { + mFlushedHistory[row] = newHistoryLengths[toSizeT(row)]; + } + return; + } + + // Every row is checked before any work is issued, so a move no reservation could cover + // leaves the batch as it was. + bool anyWork = false; + for (LifeCycleId lcId{0}; lcId < mManager->numLifeCycles(); ++lcId) + { + int const unitsPerPage = mManager->sparseCacheUnitsPerPage(lcId); + if (!mManager->hasSparseFetchPath(lcId)) + { + continue; + } + int const unitTokens = mManager->config().tokensPerBlock / unitsPerPage; + auto const base = static_cast(2 * mMaxBatchSize * lcId.value()); + + // Bounded in aggregate rather than per row, which is what the compacted list allows + // and what in-flight batching needs: one prefilling row may cross the whole step's + // worth of units while the rest cross none. + int64_t crossed = 0; + + for (BatchRow row{0}; row < mRows.size(); ++row) + { + auto const& kvCache = mRows[row]; + int fromUnit = 0; + int toUnit = 0; + if (kvCache != nullptr) + { + TLLM_CHECK_WITH_INFO(newHistoryLengths[toSizeT(row)] >= kvCache->historyLength(), + "Batch::flushSparseUnits: history length cannot be decreased"); + // Clamped to the capacity this row starts from: a unit at or above it was + // never in flight, so it has no cell and nothing to write back. + fromUnit = kvCache->historyLength() / unitTokens; + toUnit = std::min(newHistoryLengths[toSizeT(row)], kvCache->capacity()) / unitTokens; + + crossed += toUnit - fromUnit; + anyWork = anyWork || toUnit > fromUnit; + } + mFlushRangeStaging[base + toSizeT(row)] = fromUnit; + mFlushRangeStaging[base + static_cast(mMaxBatchSize) + toSizeT(row)] = toUnit; + } + + TLLM_CHECK_WITH_INFO(crossed <= mScratch[lcId].flushWork.capacity, + "Batch::flushSparseUnits: the batch crossed more units than the step's token budget allows, so " + "the write-back list could not name them all; construct the Batch with maxTotalInputTokens"); + } + + for (BatchRow row{0}; row < mRows.size(); ++row) + { + mFlushedHistory[row] = newHistoryLengths[toSizeT(row)]; + } + if (!anyWork) + { + // Nothing crossed anywhere, so the ranges on the device are stale but unread. + return; + } + + auto* const upload = mUploadStream.handle(); + if (!mFlushRead.isClosed()) { - throw LogicError("Batch members must be live requests from the same manager"); + // The previous step's write-back read these same ranges; its reads must retire + // before the copy below overwrites them. + mFlushRead.waitInStream(reinterpret_cast(upload)); + mFlushRead.close(); } - if (cache.mPageStorageBatch != nullptr) + copyToDeviceSyncSource( + reinterpret_cast(mFlushRanges), mFlushRangeStaging.data(), mFlushRangeStaging.size(), upload); + auto staged = mUploadStream.recordEvent(); + staged.waitInStream(reinterpret_cast(cudaStream())); + + for (LifeCycleId lcId{0}; lcId < mManager->numLifeCycles(); ++lcId) { - if (cache.mPageStorageBatch == this && (!row || row == cache.mPageStorageRow)) + if (!mManager->hasSparseFetchPath(lcId)) { - return *cache.mPageStorageRow; + continue; + } + writeBackSparseUnits(lcId, cudaStream()); + } + mFlushRead = CachedCudaEvent{reinterpret_cast(cudaStream())}; +} + +void Batch::_checkSparseRunsAreAddressable() const +{ + for (LifeCycleId lcId{0}; lcId < mManager->numLifeCycles(); ++lcId) + { + auto const plan = mManager->_sparseBufferPlan(lcId); + if (plan.empty()) + { + continue; + } + for (size_t begin = 0; begin < plan.size();) + { + size_t end = begin; + while (end < plan.size() && plan[end].poolIndex == plan[begin].poolIndex + && plan[end].layerRow == plan[begin].layerRow) + { + ++end; + } + + auto const slotOf = [this](SparseBufferPlan const& entry) + { + auto const found = mBufferSlots.find(entry.id); + TLLM_CHECK_WITH_INFO(found != mBufferSlots.end(), + "A sparse buffer has no index table spec; the fetch reads converted indices, so every " + "sparse buffer needs one"); + return found->second; + }; + + auto const baseOf = [this](SparseBufferPlan const& entry) + { return _specBase(_specOf(entry.id.layerId, entry.id.role, "Batch::_checkSparseRunsAreAddressable")); }; + + auto const first = slotOf(plan[begin]); + auto const firstBase = baseOf(plan[begin]); + for (size_t i = begin + 1; i < end; ++i) + { + auto const slot = slotOf(plan[i]); + // Both, because neither implies the other: a merged table serves several + // specs, each with its own base, so one table is not one base pointer. + TLLM_CHECK_WITH_INFO(slot.table == first.table && baseOf(plan[i]) == firstBase, + "Sparse buffers sharing a pool are copied in one launch, from one base pointer, so they must " + "share one index table spec"); + TLLM_CHECK_WITH_INFO(slot.entry == first.entry + static_cast(i - begin), + "Sparse buffers sharing a pool must sit at consecutive entries of their request, in the order " + "the pool holds them"); + } + begin = end; } - throw LogicError("A request can belong to only one batch and row"); } - int const index = row.value_or(static_cast(std::find(mRows.begin(), mRows.end(), nullptr) - mRows.begin())); - if (index < 0 || index >= mMaxRows || mRows[index] != nullptr) +} + +Batch::LayoutGenerator Batch::_layoutIndexTables() +{ + if (mIndexTables.empty()) + { + co_return; + } + + // Scratch-capable life cycles sort last so that one launch covers the plain tables and + // one the substituting ones, rather than a kernel branching per table. + auto const mayHoldScratch = [this](IndexTable const& table) + { return !mManager->supportsIndexMode(PageIndexMode::SHARED, table.lifeCycle).has_value(); }; + + std::vector order(mIndexTables.size()); + std::iota(order.begin(), order.end(), size_t{0}); + std::ranges::stable_partition(order, [&](size_t index) { return !mayHoldScratch(mIndexTables[index]); }); + + std::vector sorted; + std::vector remap(mIndexTables.size()); + sorted.reserve(mIndexTables.size()); + for (size_t position = 0; position < order.size(); ++position) + { + remap[order[position]] = static_cast(position); + sorted.push_back(std::move(mIndexTables[order[position]])); + } + mIndexTables = std::move(sorted); + for (auto& [key, slot] : mBufferSlots) + { + slot.table = IndexTableId{remap[static_cast(slot.table.value())]}; + } + mNumPlainTables = static_cast(std::count_if( + mIndexTables.begin(), mIndexTables.end(), [&](IndexTable const& table) { return !mayHoldScratch(table); })); + + // The most slots a request could need for one life cycle, over this batch's own tables: + // the stride every (row, life cycle) of the scratch staging gets. Zero where no member's + // step can outrun the window, and so adds no already-stale block. + for (size_t index = static_cast(mNumPlainTables); index < mIndexTables.size(); ++index) + { + auto const slots + = mManager->_scratchSlotsFor(mIndexTables[index].lifeCycle, mMaxBlocksPerSeq, mMaxTotalInputTokens); + mMaxScratchSlots = std::max(mMaxScratchSlots, slots); + } + + // With no slot any member can reach there is nothing to substitute, so every table takes + // the plain path and the scratch kernel never runs over a tail that is not there. + if (mMaxScratchSlots == 0) + { + mNumPlainTables = static_cast(mIndexTables.size()); + } + + // Everything above only sized the scratch tail; everything below names addresses within + // the block the caller allocates from it. + mScratchDevice = co_yield _scratchElements(); + + auto const rowBlocks = static_cast(mMaxBatchSize) * static_cast(mMaxBlocksPerSeq); + size_t elements = 0; + for (auto& table : mIndexTables) + { + table.elementOffset = elements; + elements += table.concatOffsets.size() * rowBlocks; + mMaxTableEntries = std::max(mMaxTableEntries, static_cast(table.concatOffsets.size())); + } + + mIndexTableMemory = allocateDeviceMemory(elements); + auto const memory = reinterpret_cast(mIndexTableMemory.get()); + // A table a publish has not reached yet must read as absent rather than as slot zero. + cuCheck(cuMemsetD32Async(memory, static_cast(kBadPageIndex.value()), elements, cudaStream())); + + // The descriptors and the offsets they point at share one allocation: the offsets are a + // handful of ints per table and would otherwise be an allocation each. + size_t offsetElements = 0; + for (auto const& table : mIndexTables) { - throw std::invalid_argument("Batch row is unavailable"); + offsetElements += table.concatOffsets.size(); } - if (cache.numBlocks().value() > mMaxBlocks || cache.beamWidth().value() > mMaxBeamWidth) + size_t const descBytes = sizeof(IndexTableDesc) * mIndexTables.size(); + size_t const offsetsBytes = sizeof(int32_t) * offsetElements; + + mIndexTableDescs = allocateDeviceMemory(descBytes + offsetsBytes); + auto const descs = reinterpret_cast(mIndexTableDescs.get()); + + std::vector offsetHost; + offsetHost.reserve(offsetElements); + std::vector descHost; + descHost.reserve(mIndexTables.size()); + + auto const offsetsBase = descs + descBytes; + auto const rawBase = _deviceBufAt(BufName::kPageTables); + for (auto const& table : mIndexTables) { - throw std::invalid_argument("Request exceeds batch dimensions"); + IndexTableDesc desc; + // The raw slot ids of this table's life cycle, in the staging the upload filled. + desc.raw = reinterpret_cast( + rawBase + sizeof(int32_t) * _rowElements() * static_cast(table.lifeCycle.value())); + desc.converted = reinterpret_cast(memory + sizeof(int32_t) * table.elementOffset); + desc.offsets = reinterpret_cast(offsetsBase + sizeof(int32_t) * offsetHost.size()); + desc.numBuffers = static_cast(table.concatOffsets.size()); + desc.scale = table.scale; + if (static_cast(descHost.size()) >= mNumPlainTables) + { + auto const row = static_cast(mMaxBatchSize); + auto const lc = static_cast(table.lifeCycle.value()); + auto const begs = reinterpret_cast(mScratchDevice); + auto const groups = static_cast(mManager->numLifeCycles().value()); + desc.scratchBegs = reinterpret_cast(begs + sizeof(int32_t) * lc * row); + desc.scratchEnds = reinterpret_cast(begs + sizeof(int32_t) * (groups * row + lc * row)); + desc.scratchSlots = reinterpret_cast( + begs + sizeof(int32_t) * (2 * groups * row + lc * row * static_cast(mMaxScratchSlots))); + desc.scratchSlotStride = mMaxScratchSlots; + auto const& firstBuffer = table.served.front().buffers.front(); + desc.scratchPagesPerBlock + = mManager->getPageIndexConverter(firstBuffer.layerId, firstBuffer.role).scratchPagesPerBlock; + } + descHost.push_back(desc); + offsetHost.insert(offsetHost.end(), table.concatOffsets.begin(), table.concatOffsets.end()); } - mRows[index] = &cache; - cache.mPageStorageBatch = this; - cache.mPageStorageRow = index; - cache.onPageStorageChanged(); - return index; + + cuCheck(cuMemcpyHtoDAsync(descs, descHost.data(), descBytes, cudaStream())); + cuCheck(cuMemcpyHtoDAsync(offsetsBase, offsetHost.data(), offsetsBytes, cudaStream())); + // Neither host vector outlives this scope, so the copies out of them have to have landed. + cuCheck(cuStreamSynchronize(cudaStream())); } -void Batch::remove(KvCache& cache) +void Batch::_clearIndexRow(BatchRow row) { - auto const apiLock = mManager->lockExclusive(); - if (cache.mPageStorageBatch == nullptr) + if (mIndexTables.empty()) { return; } - if (cache.mPageStorageBatch != this) + auto const* const descs = reinterpret_cast(mIndexTableDescs.get()); + launchClearIndexRow(descs, static_cast(mIndexTables.size()), mMaxTableEntries, row.value(), + mMaxBlocksPerSeq, mManager->_smCount(), cudaStream()); +} + +void Batch::_convertIndices(CUstream stream) +{ + if (mIndexTables.empty()) { - throw LogicError("Request belongs to another batch"); + return; } - int const row = *cache.mPageStorageRow; - mRows[row] = nullptr; - markDirty(row); - cache.mPageStorageBatch = nullptr; - cache.mPageStorageRow.reset(); - cache.onPageStorageChanged(); + auto const* const descs = reinterpret_cast(mIndexTableDescs.get()); + auto const* const capacity = reinterpret_cast(_mirrorAt(BufName::kCapacity)); + int32_t const tokensPerBlock = mManager->config().tokensPerBlock; + launchIndexConversion(descs, mNumPlainTables, mMaxTableEntries, mMaxBatchSize, mMaxBlocksPerSeq, + static_cast(mMaxBlocksPerSeq), capacity, tokensPerBlock, mManager->_smCount(), stream); + launchIndexConversionWithScratch(descs + mNumPlainTables, + static_cast(mIndexTables.size()) - mNumPlainTables, mMaxTableEntries, mMaxBatchSize, mMaxBlocksPerSeq, + static_cast(mMaxBlocksPerSeq), capacity, tokensPerBlock, mManager->_smCount(), stream); +} + +CUdeviceptr Batch::indexTableAddress(IndexTableId tableId) const +{ + TLLM_CHECK_WITH_INFO(tableId.value() >= 0 && tableId < numIndexTables(), "Batch::indexTableAddress: unknown table"); + return reinterpret_cast(mIndexTableMemory.get()) + + sizeof(int32_t) * mIndexTables[static_cast(tableId.value())].elementOffset; +} + +ConvertedIndexTable Batch::convertedIndexTable(IndexTableId tableId) const +{ + TLLM_CHECK_WITH_INFO( + tableId.value() >= 0 && tableId < numIndexTables(), "Batch::convertedIndexTable: unknown table"); + ConvertedIndexTable view; + view.data = indexTableAddress(tableId); + view.numRows = mMaxBatchSize; + // A table's entries are its requests' buffers: the offsets it converts with, one each. + view.numBuffers = static_cast(mIndexTables[static_cast(tableId.value())].concatOffsets.size()); + view.blocksPerRow = mMaxBlocksPerSeq; + return view; +} + +Batch::SparseIndexRun Batch::sparseIndexRun(LayerId layerId, DataRole const& role) const +{ + auto const& slot = _bufferSlot(layerId, role, "Batch::sparseIndexRun"); + auto const tableId = slot.table; + SparseIndexRun run; + // Row-major in the entries, so a row costs the whole table's entries and an entry costs + // one row of blocks. + auto const numBuffers = static_cast(mIndexTables[static_cast(tableId.value())].concatOffsets.size()); + run.bufferStride = mMaxBlocksPerSeq; + run.rowStride = numBuffers * mMaxBlocksPerSeq; + run.table = indexTableAddress(tableId) + sizeof(int32_t) * static_cast(slot.entry) * run.bufferStride; + run.base = _specBase(_specOf(layerId, role, "Batch::sparseIndexRun")); + return run; } -void Batch::close() +size_t Batch::_scratchElements() const noexcept { - auto const apiLock = mManager->lockExclusive(); - if (mClosed || Poison::poisoned()) + // No table takes the substituting path, so nothing would read a range or a slot. + if (mMaxScratchSlots == 0) + { + return 0; + } + auto const rows = static_cast(mMaxBatchSize); + auto const groups = static_cast(mManager->numLifeCycles().value()); + // [begs][ends][slots], each indexed by life cycle then row. + return 2 * groups * rows + groups * rows * static_cast(mMaxScratchSlots); +} + +void Batch::_stageScratch() +{ + if (mMaxScratchSlots == 0) { return; } - for (auto* cache : mRows) + auto const rows = static_cast(mMaxBatchSize); + auto const groups = static_cast(mManager->numLifeCycles().value()); + auto* const begs = mScratchStaging; + auto* const ends = begs + groups * rows; + + // Only the ranges: a member's slot ids reach its row as they change, through the + // buffer bound when it joined. + for (BatchRow row{0}; row < mRows.size(); ++row) { - if (cache != nullptr) + auto const& kvCache = mRows[row]; + for (LifeCycleId lcId{0}; lcId < mManager->numLifeCycles(); ++lcId) { - remove(*cache); + auto const flat = static_cast(lcId.value()) * rows + toSizeT(row); + + // A free row, a suspended member, or one holding no scratch here all leave an + // empty range, which the kernel reads as "every block takes the plain path". + auto const range = kvCache != nullptr && kvCache->isActive() ? kvCache->getSwaScratchRange(lcId) + : HalfOpenRange{}; + begs[flat] = range.beg.value(); + ends[flat] = range.end.value(); + if (range.empty()) + { + continue; + } + TLLM_CHECK_DEBUG_WITH_INFO(range.end.value() - range.beg.value() + <= mManager->_maxScratchBlocks(lcId, mMaxBlocksPerSeq, mMaxTotalInputTokens), + "A request's scratch range outruns the blocks the batch was sized for"); } } - mClosed = true; - mPublished = false; } -size_t Batch::rowOffset(int group, int row) const noexcept +IndexTableId Batch::numIndexTables() const noexcept +{ + return IndexTableId{static_cast(mIndexTables.size())}; +} + +std::vector const& Batch::indexTable(IndexTableId tableId) const { - return (static_cast(group) * mMaxRows + row) * mMaxBeamWidth * mMaxBlocks; + TLLM_CHECK_WITH_INFO(tableId.value() >= 0 && tableId < numIndexTables(), "Batch::indexTable: unknown table"); + return mIndexTables[static_cast(tableId.value())].served; } -size_t Batch::countOffset(int group, int row) const noexcept +Batch::BufferSlot const& Batch::_bufferSlot(LayerId layerId, DataRole const& role, char const* api) const { - return mTableElements + (static_cast(group) * mMaxRows + row) * mMaxBeamWidth; + auto const found = mBufferSlots.find(BufferId{layerId, role}); + TLLM_CHECK_WITH_INFO(found != mBufferSlots.end(), std::string(api) + ": no index table spec named this buffer"); + return found->second; } -std::vector Batch::dirtyRows() const +IndexTableId Batch::indexTableOf(IndexTableSpec const& spec) const { - auto const apiLock = mManager->lockShared(); - checkOpen(); - std::vector rows; - for (int row = 0; row < mMaxRows; ++row) + for (size_t i = 0; i < mIndexTables.size(); ++i) { - if (mDirty[row]) + for (auto const& served : mIndexTables[i].served) { - rows.push_back(row); + if (served.buffers == spec.buffers) + { + return IndexTableId{static_cast(i)}; + } } } - return rows; + throw std::out_of_range("Batch::indexTableOf: no index table was built for this spec"); } -std::vector Batch::publish(CudaStream stream) +IndexTableSpec const& Batch::_specOf(LayerId layerId, DataRole const& role, char const* api) const { - KVCM2_API_GUARD(); - auto const apiLock = mManager->lockExclusive(); - checkOpen(); - auto const cudaStream = reinterpret_cast(stream); - checkOutsideCapture(stream); - // Owner release can unblock history without changing a request's watermark or metadata version. - // Retry before collecting dirty rows: one shared-page move can invalidate several rows. - for (auto* cache : mRows) + auto const& slot = _bufferSlot(layerId, role, api); + BufferId const wanted{layerId, role}; + for (auto const& served : mIndexTables[static_cast(slot.table.value())].served) { - if (cache != nullptr && cache->isActive() && cache->mIsDecoding && cache->mHasDeferredSparseOffload) + if (std::ranges::find(served.buffers, wanted) != served.buffers.end()) { - cache->_offloadSparseHistory({0, 0}, cache->mHistoryLength); + return served; } } - auto rows = dirtyRows(); - if (rows.empty()) + // Unreachable while construction records a slot only for a buffer a spec names. + throw LogicError(std::string(api) + ": a buffer's table does not serve the spec that named it"); +} + +MemAddress Batch::_specBase(IndexTableSpec const& spec) const +{ + auto const& first = spec.buffers.front(); + auto const lifeCycle = mManager->_bufferAttr(first.layerId, first.role).lifeCycleId; + return mManager->getMemPoolBaseAddress(first.layerId, first.role, mIndexModes[lifeCycle]); +} + +PoolSpec Batch::poolSpecFor(IndexTableSpec const& spec) const +{ + // Membership first: a spec the Batch never built has no base. + static_cast(indexTableOf(spec)); + + auto const& first = spec.buffers.front(); + PoolSpec out; + out.base = _specBase(spec); + out.pageStride = mManager->getPageStride(first.layerId, first.role); + for (auto const& buffer : spec.buffers) { - mReady.waitInStream(stream); - return rows; + // One stride for the spec, because a reader has only one: the entries carry each + // buffer's offset in pages, which is a count of this stride and of no other. + TLLM_CHECK_WITH_INFO( + static_cast(mManager->getPageStride(buffer.layerId, buffer.role)) == out.pageStride, + "An index table spec spans page strides, so one base cannot address its buffers"); } + return out; +} + +std::pair Batch::indexTableSlot(LayerId layerId, DataRole const& role) const +{ + auto const& slot = _bufferSlot(layerId, role, "Batch::indexTableSlot"); + return {slot.table, slot.entry}; +} - struct RequestSnapshot +MemAddress Batch::indexTableBase(LayerId layerId, DataRole const& role) const +{ + return _specBase(_specOf(layerId, role, "Batch::indexTableBase")); +} + +void Batch::_drain() noexcept +{ + KVCM2_POISON_ON_EXCEPT( + [this] + { + waitPublished(); + if (!mDeviceStagingConsumed.isClosed()) + { + mDeviceStagingConsumed.synchronize(); + mDeviceStagingConsumed.close(); + } + }); +} + +void Batch::_layoutBufs(size_t scratchElements) noexcept +{ + auto const rows = static_cast(mMaxBatchSize); + auto const groups = static_cast(mManager->numLifeCycles().value()); + + struct Spec { - KvCache* cache; - uint64_t version; - std::vector groups; + BufName name; + size_t bytes; + size_t align; }; - std::vector snapshots; - snapshots.reserve(rows.size()); - for (int row : rows) + std::array const specs{{ + {BufName::kPageTables, sizeof(int32_t) * _rowElements() * groups, alignof(int32_t)}, + {BufName::kHistory, sizeof(int32_t) * rows, alignof(int32_t)}, + {BufName::kCapacity, sizeof(int32_t) * rows, alignof(int32_t)}, + {BufName::kLeases, sizeof(SparseLeaseIndex) * rows, alignof(SparseLeaseIndex)}, + {BufName::kScratch, sizeof(int32_t) * scratchElements, alignof(int32_t)}, + }}; + + ptrdiff_t offset = 0; + for (auto const& spec : specs) { - auto* cache = mRows[row]; - RequestSnapshot request{cache, cache ? cache->pageStorageVersion() : 0, {}}; - if (cache != nullptr) + auto const align = static_cast(std::max(spec.align, kBufMinAlign)); + offset = (offset + align - 1) / align * align; + mBuf[spec.name] = HalfOpenRange{offset, offset + static_cast(spec.bytes)}; + offset = mBuf[spec.name].end; + } +} + +SparseLeaseIndex const* Batch::_leaseMirror() const noexcept +{ + return reinterpret_cast(sparseLeases()); +} + +size_t Batch::_mirrorBytes() const noexcept +{ + // From the histories through the leases: the page tables ahead of them are read from the + // staging, and the scratch behind them is the device staging's own tail. + return static_cast(mBuf[BufName::kLeases].end - mBuf[BufName::kHistory].beg); +} + +CUdeviceptr Batch::_deviceBufAt(BufName name) const noexcept +{ + return reinterpret_cast(mDeviceStaging.get()) + static_cast(mBuf[name].beg); +} + +CUdeviceptr Batch::_mirrorAt(BufName name) const noexcept +{ + // The mirror carries the block from the histories on, so a buffer addresses it rebased. + return reinterpret_cast(mMirror.get()) + + static_cast(mBuf[name].beg - mBuf[BufName::kHistory].beg); +} + +void Batch::_stageRowState() +{ + auto const history = _buf(BufName::kHistory); + auto const capacity = _buf(BufName::kCapacity); + auto const leases = _buf(BufName::kLeases); + + // A row's page table is mMaxBlocksPerSeq entries wide, and the index conversion walks it + // as far as the capacity published here reaches. + auto const capacityLimit = static_cast(mMaxBlocksPerSeq) * mManager->config().tokensPerBlock; + + for (BatchRow row{0}; row < mRows.size(); ++row) + { + auto const& kvCache = mRows[row]; + history[toSizeT(row)] = kvCache == nullptr ? 0 : kvCache->historyLength(); + capacity[toSizeT(row)] = kvCache == nullptr ? 0 : kvCache->capacity(); + TLLM_CHECK_WITH_INFO(capacity[toSizeT(row)] <= capacityLimit, + "Batch: row %d has capacity %d tokens, past the %lld its page table holds", row.value(), + capacity[toSizeT(row)], static_cast(capacityLimit)); + + // A free row, or a member holding no reservation, publishes no lease, which the + // kernels read as "nothing to resolve here". One slot serves every life cycle. + leases[toSizeT(row)] = kvCache == nullptr ? kNoSparseLease : leaseOf(row); + } +} + +void Batch::_seedRowWatermarks(BatchRow row, int historyLength, int capacity) +{ + auto const index = toSizeT(row); + _buf(BufName::kHistory)[index] = historyLength; + _buf(BufName::kCapacity)[index] = capacity; +} + +int32_t* Batch::_stagingRow(LifeCycleId lifeCycleId, BatchRow row) const +{ + auto* base = _buf(BufName::kPageTables).data(); + return base + static_cast(_rowElements()) * lifeCycleId.value() + + static_cast(mMaxBlocksPerSeq) * row.value(); +} + +int32_t* Batch::_scratchSlotRow(LifeCycleId lifeCycleId, BatchRow row) const +{ + auto const rows = static_cast(mMaxBatchSize); + auto const groups = static_cast(mManager->numLifeCycles().value()); + auto const flat = static_cast(lifeCycleId.value()) * rows + toSizeT(row); + return mScratchStaging + 2 * groups * rows + flat * static_cast(mMaxScratchSlots); +} + +BatchRow Batch::addKvCache(KvCache& kvCache, std::optional requestedRow) +{ + if (kvCache.mBatch != nullptr) + { + throw LogicError("A KvCache belongs to at most one Batch"); + } + if (kvCache.mManager.get() != mManager) + { + // Rows, arena leases and index tables are all this manager's. A member from another + // would be addressed through them and read memory that was never its own. + throw LogicError("A KvCache may only join a Batch of the manager that created it"); + } + if (kvCache.beamWidth() != BeamIndex{1}) + { + // Only beam 0 is mirrored below, so a wider request would publish one beam's pages + // and have every beam read them. + throw LogicError("Batch does not support beam search yet; a member must have beam width 1"); + } + if (!kvCache.isActive()) + { + // Membership means the member is being stepped, which is what lets its arena lease, + // its row and its active state be one thing rather than three that can disagree. + throw LogicError("A KvCache must be resumed before it joins a Batch"); + } + // Binding a row writes it, so the last publish must be done reading it. + waitPublished(); + BatchRow row{0}; + if (!requestedRow.has_value()) + { + auto const free = std::find(mRows.begin(), mRows.end(), nullptr); + if (free == mRows.end()) { - if (cache->numBlocks().value() > mMaxBlocks || cache->beamWidth().value() != mMaxBeamWidth) - { - throw std::invalid_argument("Request exceeds batch dimensions"); - } - request.groups.reserve(mNumLayerGroups); - for (int group = 0; group < mNumLayerGroups; ++group) - { - request.groups.push_back(cache->getPageStorageSnapshot(LayerGroupId{group})); - } + throw LogicError("Batch is full; it was sized for max_batch_size rows"); } - snapshots.push_back(std::move(request)); + row = BatchRow{static_cast(std::distance(mRows.begin(), free))}; } + else + { + row = *requestedRow; + TLLM_CHECK_WITH_INFO( + row != kNoBatchRow && row < mMaxBatchSize, "Batch::addKvCache: row is beyond max_batch_size"); + if (mRows[row] != nullptr) + { + throw LogicError("Batch::addKvCache: row is already occupied"); + } + } + + // The bind copies the current indices in and keeps writing there until the member + // leaves, so the row needs no refresh of its own and a step that only advances its + // members has nothing to rebind. + for (LifeCycleId lc{0}; lc < mManager->numLifeCycles(); ++lc) + { + // Beam 0 only; addKvCache rejects anything wider. + kvCache.setBasePageIndexBuf(kDefaultBeamIndex, lc, _stagingRow(lc, row), mMaxBlocksPerSeq); + if (mMaxScratchSlots > 0) + { + kvCache.setScratchSlotBuf(lc, _scratchSlotRow(lc, row), mMaxScratchSlots); + } + } + + _seedRowWatermarks(row, kvCache.historyLength(), kvCache.capacity()); + // A seating member brings no in-flight units: its lease is fresh, so everything its + // history names already lives in the pages behind it. Seeding from the history is what + // lets the row be reused, since the departing member leaves this equal to its own. + mFlushedHistory[row] = kvCache.historyLength(); + + // Joining is a stream change, so it carries the same cross-stream sync a direct move + // would: work the member already enqueued has to retire before the Batch's stream + // reaches anything that reads it. + kvCache._setCudaStream(cudaStream()); + + // Owning from here: the member cannot be destroyed under the Batch, so a caller may + // keep only the Batch and get the reference back when it leaves. + mRows[row] = kvCache.shared_from_this(); + kvCache.mBatch = this; + kvCache.mBatchRow = row; + return row; +} + +std::shared_ptr const& Batch::operator[](BatchRow row) const +{ + return mRows.at(row); +} - // A busy staging buffer cannot be overwritten by the CPU just by queueing a - // stream wait. Reuse only completed buffers; retain each in-flight generation. - std::unique_ptr staging; - for (auto it = mUploads.begin(); it != mUploads.end();) +std::shared_ptr Batch::removeKvCache(BatchRow row) +{ + TLLM_CHECK_WITH_INFO(row >= 0 && row < mMaxBatchSize, "Batch::removeKvCache: row is beyond max_batch_size"); + auto seated = mRows[row]; + return seated ? removeKvCache(*seated) : nullptr; +} + +std::shared_ptr Batch::removeKvCache(KvCache& kvCache) +{ + if (kvCache.mBatch != this) + { + throw LogicError("KvCache is not a member of this Batch"); + } + waitPublished(); + BatchRow const row = *kvCache.mBatchRow; + // Checked before anything is touched, so a rejected removal leaves the batch as it was. + // + // Flushing a history the member never takes up is pointless, so the two agreeing is what + // a member that has written back everything it means to keep looks like. The leases + // released below discard the GPU cache, and the next occupant seeds its own watermark + // from its own history, so anything still pending here has no later chance to be noticed. + if (mManager->hasSparseCache()) + { + TLLM_CHECK_WITH_INFO(kvCache.historyLength() == mFlushedHistory[row], + "A sparse KvCache leaving row %d has history %d against %d written back; call " + "Batch::flushSparseUnits() with the current history before removing it", + row.value(), kvCache.historyLength(), mFlushedHistory[row]); + } + // Held until the end: clearing the row below drops the Batch's reference, and the rest + // of this runs against the member. + auto leaving = std::move(mRows[row]); + // Membership ends first, so the unbinding below goes through the same setter a caller + // would use. Nothing binds a seated member's buffer, including the Batch. + kvCache.mBatch = nullptr; + kvCache.mBatchRow = std::nullopt; + for (LifeCycleId lc{0}; lc < mManager->numLifeCycles(); ++lc) { - if (it->completion.queryComplete()) + // Hands the indices back to the cache's own storage before the row is reused. + kvCache.setBasePageIndexBuf(kDefaultBeamIndex, lc, nullptr, 0); + if (mMaxScratchSlots > 0) { - staging = std::move(it->staging); - mUploads.erase(it); - break; + kvCache.setScratchSlotBuf(lc, nullptr, 0); + auto* slots = _scratchSlotRow(lc, row); + std::fill(slots, slots + mMaxScratchSlots, kBadPageIndex.value()); } - else + auto* staged = _stagingRow(lc, row); + std::fill(staged, staged + mMaxBlocksPerSeq, kBadPageIndex.value()); + } + _clearIndexRow(row); + // The row keeps its lease and hands it back clean, as it does the converted indices + // above: flushSparseUnits() has written back anything worth keeping, so what is left + // names units no later occupant has any claim on. + if (mManager->hasSparseCache()) + { + for (LifeCycleId lc{0}; lc < mManager->numLifeCycles(); ++lc) { - ++it; + _sparseArena(lc).initialiseLease(leaseOf(row), cudaStream()); } } - if (!staging) + _seedRowWatermarks(row, 0, 0); + mRows[row] = nullptr; + return leaving; +} + +void Batch::move(BatchRow dstRow, BatchRow srcRow) +{ + TLLM_CHECK_WITH_INFO(dstRow >= 0 && dstRow < mMaxBatchSize && srcRow >= 0 && srcRow < mMaxBatchSize, + "Batch::move: row is beyond max_batch_size"); + // Checked before anything is touched, so a rejected move leaves the batch as it was. + // dstRow == srcRow fails here too: a row cannot be occupied and empty at once. + if (mRows[srcRow] == nullptr) + { + throw LogicError("Batch::move: source row holds no member"); + } + if (mRows[dstRow] != nullptr) { - staging = std::make_unique(mTotalBytes); + throw LogicError("Batch::move: destination row is already occupied"); } - auto* host = reinterpret_cast(staging->address()); - for (size_t i = 0; i < rows.size(); ++i) + + waitPublished(); + + for (LifeCycleId lc{0}; lc < mManager->numLifeCycles(); ++lc) { - for (int group = 0; group < mNumLayerGroups; ++group) + auto* const from = _stagingRow(lc, srcRow); + auto* const to = _stagingRow(lc, dstRow); + // Swapped rather than copied so the vacated row ends up holding what the empty + // destination held: every entry kBadPageIndex, which is what a free row must publish. + std::swap_ranges(from, from + mMaxBlocksPerSeq, to); + if (mMaxScratchSlots > 0) { - int32_t* indices = host + rowOffset(group, rows[i]); - std::fill_n(indices, mMaxBlocks, kBadPageIndex.value()); - int32_t& count = host[countOffset(group, rows[i])]; - count = 0; - if (snapshots[i].cache != nullptr) - { - auto const& snapshot = snapshots[i].groups[group]; - std::copy(snapshot.basePageIndices().begin(), snapshot.basePageIndices().end(), indices); - count = snapshot.eligibleHistoryBlocks(); - } + // The scratch slots are a second per-row table the member writes through, and + // _stageScratch() publishes them from the row it occupies. They travel with it + // for the same reason its page indices do. + auto* const scratchFrom = _scratchSlotRow(lc, srcRow); + auto* const scratchTo = _scratchSlotRow(lc, dstRow); + std::swap_ranges(scratchFrom, scratchFrom + mMaxScratchSlots, scratchTo); } } - mUploads.push_back({std::move(staging)}); - auto& upload = mUploads.back(); - auto fence = FuncGuard( - [&]() + std::swap(mRows[dstRow], mRows[srcRow]); + std::swap(mFlushedHistory[dstRow], mFlushedHistory[srcRow]); + // The member keeps the cells it already filled, so its lease travels with it rather + // than its contents being copied into the destination row's. + std::swap(mRowLease[dstRow], mRowLease[srcRow]); + + auto& kvCache = *mRows[dstRow]; + kvCache.mBatchRow = dstRow; + for (LifeCycleId lc{0}; lc < mManager->numLifeCycles(); ++lc) + { + // The row's contents have already moved. Rebind its sole writer without copying from + // the old row or ending membership -- mBatch and the arena lease are untouched. + kvCache.mBasePageIndices[kDefaultBeamIndex][lc] = Span{_stagingRow(lc, dstRow), mMaxBlocksPerSeq}; + if (mMaxScratchSlots > 0) { - mReady = CachedCudaEvent(stream); - upload.completion = mReady; - }); - mReady.waitInStream(stream); - for (auto const& reader : mReaders) + kvCache.setScratchSlotBuf(lc, _scratchSlotRow(lc, dstRow), mMaxScratchSlots); + } + } + _seedRowWatermarks(dstRow, kvCache.historyLength(), kvCache.capacity()); + _seedRowWatermarks(srcRow, 0, 0); + _clearIndexRow(srcRow); +} + +void Batch::_detach(KvCache& kvCache) noexcept +{ + if (kvCache.mBatch != this || !kvCache.mBatchRow.has_value()) { - reader.waitInStream(stream); + return; } - mReaders.clear(); - for (auto const& request : snapshots) + // A seated member holds a reference from mRows, so its destructor cannot run while it + // is one. Reaching here means a row was cleared without unbinding the member, leaving a + // row that publishes from an object about to stop existing. + KVCM2_CHECK_FATAL_WITH_INFO( + false, "A KvCache was destroyed while seated in a Batch, so its row outlived the cache it names"); +} + +std::vector Batch::resize( + std::vector> const& capacities, std::vector> const& historyLengths) +{ + TLLM_CHECK_WITH_INFO(capacities.size() == mRows.stdSize() && historyLengths.size() == mRows.stdSize(), + "Batch::resize: one capacity and history length per row"); + + // Every row is checked before any is touched, so a bad request leaves the batch as it + // was rather than part-way through. + for (BatchRow row{0}; row < mRows.size(); ++row) { - for (auto const& snapshot : request.groups) + auto const& kvCache = mRows[row]; + if (kvCache == nullptr || !historyLengths[toSizeT(row)].has_value()) { - snapshot.waitReady(stream); + continue; } + int const newHist = *historyLengths[toSizeT(row)]; + TLLM_CHECK_WITH_INFO(newHist >= kvCache->historyLength(), "Batch::resize: history length cannot be decreased"); + // Against the capacity that already existed, not the one being asked for. Capacity + // this call grants makes room for tokens written after it, so history can only + // cover tokens whose space was there while they were being written. + TLLM_CHECK_WITH_INFO(newHist <= kvCache->capacity(), + "Batch::resize: history length cannot exceed the capacity the row already had"); } - auto const device = reinterpret_cast(mDeviceMemory.get()); - for (int row : rows) + + // Before the first member is touched, since each resize writes its staging row. + waitPublished(); + + std::vector results(mRows.stdSize(), 0); + for (BatchRow row{0}; row < mRows.size(); ++row) { - for (int group = 0; group < mNumLayerGroups; ++group) + auto const& kvCache = mRows[row]; + if (kvCache == nullptr) + { + continue; + } + if (kvCache->resize(capacities[toSizeT(row)], historyLengths[toSizeT(row)])) { - size_t const offset = rowOffset(group, row); - cuCheck(cuMemcpyHtoDAsync( - device + offset * sizeof(int32_t), host + offset, mMaxBlocks * sizeof(int32_t), cudaStream)); - size_t const count = countOffset(group, row); - cuCheck(cuMemcpyHtoDAsync(device + count * sizeof(int32_t), host + count, sizeof(int32_t), cudaStream)); + results[toSizeT(row)] = 1; + continue; + } + // The growth failed and rolled itself back, but the history advance is still both + // legal -- it was checked against the capacity this row kept -- and free: the stale + // range and the scratch range are each monotone in history, so raising it can only + // release slots. Taking it anyway is what lets the units below it be written back + // and their cells released, which is the pressure that made the growth fail. + if (capacities[toSizeT(row)].has_value() && historyLengths[toSizeT(row)].has_value()) + { + kvCache->resize(std::nullopt, historyLengths[toSizeT(row)]); } } - fence.run(); - for (size_t i = 0; i < rows.size(); ++i) + // Once, after every member has settled: a partial failure leaves the rolled-back state, + // which is what the table should describe. + publish(); + return results; +} + +void Batch::publish() +{ + // A step's work-list counters start here, all of them at once: per layer and per + // direction they are independent, and one clear at the head of a step keeps a + // non-kernel node out from between every resolve and whatever ran before it. + // + // On the Batch's stream, like the staging below it. A fetch issued on another stream + // therefore has to be ordered after this call -- but it already did, because it reads + // the page tables and watermarks this same call stages. + if (mCounterElements > 0) { - auto const& request = snapshots[i]; - if (request.cache != nullptr) + TLLM_CUDA_CHECK(cudaMemsetAsync( + mCounters, 0, sizeof(int32_t) * mCounterElements, reinterpret_cast(cudaStream()))); + } + + // Refusing rather than waiting, because by this point a wait would be too late to help. + // Everything that writes the staged rows has to be separated from the copy reading them: + // the members' own mutations, which happened before this call, and the watermarks and + // views staged just below. Waiting here would cover only the second, and would do it + // while quietly hiding that the first had already raced. + TLLM_CHECK_WITH_INFO(mPublished.isClosed(), + "Batch::publish: the previous publish is still reading the staged rows. Call waitPublished() before " + "whatever wrote them -- a direct resize(), a commit(), a beam fork, an SWA recycle, anything that locks " + "or unlocks a page -- and not merely before this call: moving the wait here silences this check while " + "leaving that earlier write racing the copy. Batch's own mutators, resize() and move() among them, " + "already wait."); + + // A SHARED life cycle gives every layer the same converted indices, which only holds + // while no member substitutes scratch slots there. Prefill acquires those and decode + // does not, so a derived mode would have been PER_LAYER and reaching here means the + // phases were mixed after it was fixed. A named SHARED makes this the check that holds + // the caller to the claim: it is what lets a decode-only batch take the smaller table + // on a life cycle the manager cannot clear for requests in general. + // Life cycle outermost: the rows are walked only for one the manager left undecided, + // since every other answer passes for every member. + for (LifeCycleId lcId{0}; lcId < mManager->numLifeCycles(); ++lcId) + { + if (mManager->supportsIndexMode(mIndexModes[lcId], lcId) == std::optional(true)) + { + continue; + } + for (auto const& kvCache : mRows) { - TLLM_CHECK(request.cache->acknowledgePageStorage(request.version)); + TLLM_CHECK_WITH_INFO(kvCache == nullptr || kvCache->supportsIndexMode(mIndexModes[lcId], lcId), + "Batch::publish: a member holds scratch slots in a SHARED life cycle; its layers no longer share " + "one converted index table"); } - mDirty[rows[i]] = 0; } - mPublished = true; - return rows; + + // The members write their own page indices, but not their watermarks; collect those + // now so the single transfer below carries a matched pair. + _stageRowState(); + _stageScratch(); + + auto const staging = reinterpret_cast(mDeviceStaging.get()); + auto* const upload = mUploadStream.handle(); + + // The fill below overwrites what the previous step's copy read into the final tables. + // Those run on different streams, so nothing but this event orders them. + if (!mDeviceStagingConsumed.isClosed()) + { + mDeviceStagingConsumed.waitInStream(reinterpret_cast(upload)); + mDeviceStagingConsumed.close(); + } + + // One transfer for the whole block rather than one per changed row, or one per buffer: + // a single driver call on the upload stream overlaps compute, where the alternatives + // cost a launch apiece on the step's critical path. + cuCheck(cuMemcpyHtoDAsync(staging, reinterpret_cast(mStaging.address()), mUploadBytes, upload)); + mPublished = mUploadStream.recordEvent(); + + // The staging now holds this step's tables, so carrying the mirror across is all that + // is left on the Batch's stream -- one copy, because the mirror is its tail verbatim. + mPublished.waitInStream(reinterpret_cast(cudaStream())); + cuCheck(cuMemcpyDtoDAsync( + reinterpret_cast(mMirror.get()), _deviceBufAt(BufName::kHistory), _mirrorBytes(), cudaStream())); + _convertIndices(cudaStream()); + mDeviceStagingConsumed = CachedCudaEvent{reinterpret_cast(cudaStream())}; +} + +void Batch::waitPublished() +{ + if (mPublished.isClosed()) + { + return; + } + mPublished.synchronize(); + // Closing returns the event to the pool and makes every later call free. + mPublished.close(); +} + +CoalescedPageTable Batch::coalescedPageTable(LifeCycleId lifeCycleId) const +{ + TLLM_CHECK_WITH_INFO(lifeCycleId.value() >= 0 && lifeCycleId < mManager->numLifeCycles(), + "Batch::coalescedPageTable: unknown layer group"); + CoalescedPageTable view; + view.data = _stagingRow(lifeCycleId, BatchRow{0}); + view.numRows = mMaxBatchSize; + view.blocksPerRow = mMaxBlocksPerSeq; + return view; +} + +GpuCoalescedPageTable Batch::gpuCoalescedPageTable(LifeCycleId lifeCycleId) const +{ + TLLM_CHECK_WITH_INFO(lifeCycleId.value() >= 0 && lifeCycleId < mManager->numLifeCycles(), + "Batch::gpuCoalescedPageTable: unknown layer group"); + GpuCoalescedPageTable view; + // From the staging: the mirror carries only what a captured kernel reads, and the raw + // indices are read there by the conversion rather than copied on. + view.data = _deviceBufAt(BufName::kPageTables) + + sizeof(int32_t) * _rowElements() * static_cast(lifeCycleId.value()); + view.numRows = mMaxBatchSize; + view.blocksPerRow = mMaxBlocksPerSeq; + return view; +} + +CUdeviceptr Batch::sparseLeases() const +{ + return _mirrorAt(BufName::kLeases); +} + +WatermarkView Batch::watermarks() const +{ + WatermarkView view; + view.historyLength = _mirrorAt(BufName::kHistory); + view.capacity = _mirrorAt(BufName::kCapacity); + view.numRows = mMaxBatchSize; + return view; +} + +SparseHitMiss Batch::peekHitMissCounter() const +{ + if (!mHitMiss) + { + return {}; + } + // A temporary stream, synced alone: the Batch's own stream carries the step, and + // waiting on that would turn a statistics read into a pipeline stall. + SparseHitMiss host; + TemporaryCudaStream temp{{}}; + cuCheck(cuMemcpyDtoHAsync(&host, reinterpret_cast(mHitMiss.get()), sizeof(host), temp.get())); + cuCheck(cuStreamSynchronize(temp.get())); + return host; +} + +void Batch::resetHitMissCounter() +{ + if (mHitMiss) + { + TLLM_CUDA_CHECK( + cudaMemsetAsync(mHitMiss.get(), 0, sizeof(SparseHitMiss), reinterpret_cast(cudaStream()))); + } +} + +size_t Batch::sparseArenaBytes() const +{ + size_t total = 0; + for (LifeCycleId lc{0}; lc < mManager->numLifeCycles(); ++lc) + { + total += _sparseArena(lc).numBytes(); + } + return total; +} + +int Batch::numKvCaches() const noexcept +{ + return static_cast( + std::count_if(mRows.begin(), mRows.end(), [](std::shared_ptr const& row) { return row != nullptr; })); +} + +bool Batch::isContiguous() const noexcept +{ + // Nothing occupied past the first free row is the same thing as a prefix. + auto const firstFree = std::find(mRows.begin(), mRows.end(), nullptr); + return std::none_of(firstFree, mRows.end(), [](std::shared_ptr const& row) { return row != nullptr; }); } -void Batch::waitReady(CudaStream stream) const +std::optional Batch::rowOf(KvCache const& kvCache) const { - KVCM2_API_GUARD(); - auto const apiLock = mManager->lockShared(); - checkPublished(); - checkOutsideCapture(stream); - mReady.waitInStream(stream); + return kvCache.mBatch == this ? kvCache.mBatchRow : std::nullopt; } -void Batch::recordRead(CudaStream stream) +void Batch::prepareSparseWriteIndices( + std::vector const& buffers, int32_t* out, std::optional streamOpt) { - KVCM2_API_GUARD(); - auto const apiLock = mManager->lockExclusive(); - checkOpen(); - checkOutsideCapture(stream); - std::erase_if(mReaders, [](auto const& event) { return event.queryComplete(); }); - CachedCudaEvent completion(stream); - mReaders.push_back(completion); - for (auto* cache : mRows) + CUstream const stream = streamOpt.value_or(cudaStream()); + TLLM_CHECK_WITH_INFO(out != nullptr, "prepareSparseWriteIndices: out must be a device pointer"); + + auto const span = mManager->_sparseLayerSpan(buffers, "prepareSparseWriteIndices"); + LifeCycleId const lcId = span.lifeCycle; + + auto const& scratch = sparseStepScratch(lcId); + TLLM_CHECK_WITH_INFO( + scratch.inFlightSlot != nullptr, "prepareSparseWriteIndices: this layer group carries no sparse cache"); + + int const unitsPerPage = mManager->sparseCacheUnitsPerPage(lcId); + auto const leases = _leaseMirror(); + auto const arena = _sparseArena(lcId).view(); + auto const marks = watermarks(); + auto const numRows = maxBatchSize(); + + // A resolve with no selection. It makes the in-flight unit resident -- reusing its entry + // when a previous step already pinned it -- and does nothing else: an in-flight unit's + // host page is not written until the watermark passes it, so there is nothing to copy. + // + // The selection and miss arrays are therefore left null, which the resolve tolerates only + // because it records a miss against the copy list under `ordinal < historyUnits`, and an + // in-flight ordinal never satisfies it. A change that made that write unconditional would + // fault here rather than quietly copying stale bytes over what the model wrote. + SparseResolveArgs resolve; + resolve.leases = leases; + resolve.arena = arena; + resolve.numRows = numRows; + resolve.selectionSlot = span.selectionSlot; + resolve.selectionWidth = 0; + resolve.historyTokens = reinterpret_cast(marks.historyLength); + resolve.capacityTokens = reinterpret_cast(marks.capacity); + resolve.unitTokens = mManager->config().tokensPerBlock / unitsPerPage; + resolve.inFlightSlot = scratch.inFlightSlot; + resolve.maxInFlightUnits = scratch.maxInFlightUnits; + resolve.missCount = scratch.missCount; + // No selection on this path, so there is nothing to resolve into a slot. + resolve.selected = nullptr; + resolve.outSlot = nullptr; + launchSparseResolve(resolve, stream); + + // The in-flight list stands in for a selection: the resolve writes kNoSparseSlot past the + // span, so a row with nothing in flight needs no separate signal. + // + // Shared the same way the fetch shares, and for the same reason: buffers whose arena + // indices coincide are resolved once. + auto const slots = sparseOutSlots(buffers); + std::vector written(static_cast(span.numBuffers), false); + for (auto const& copy : _sparseCopyTemplates(lcId)) { - if (cache != nullptr && cache->isActive()) + auto const runBase = std::max(copy.bufferBase, span.bufferBase); + auto const runEnd = std::min(copy.bufferBase + copy.numBuffers, span.bufferBase + span.numBuffers); + if (runBase >= runEnd) + { + continue; + } + int const firstSlot = slots[static_cast(runBase - span.bufferBase)]; + if (written[static_cast(firstSlot)]) { - completion.waitInStream(reinterpret_cast(cache->cudaStream())); + continue; } + for (int32_t buffer = runBase; buffer < runEnd; ++buffer) + { + written[static_cast(slots[static_cast(buffer - span.bufferBase)])] = true; + } + SparseOutIndexArgs indices; + indices.leases = leases; + indices.arena = arena; + indices.numRows = numRows; + indices.selected = nullptr; + indices.outSlot = scratch.inFlightSlot; + // One entry per request: resize() bounds capacity - history to a single token, so at + // most one unit is ever in flight. The scratch row is the reservation's headroom, + // which is wider, hence the separate stride. + indices.selectionWidth = 1; + indices.outSlotStride = scratch.maxInFlightUnits; + indices.out = out; + // The row stride is the caller's: `out` keeps one entry per buffer, and an entry + // no buffer maps to is never written. + indices.outBuffers = span.numBuffers; + indices.outBase = firstSlot; + indices.bufferBase = runBase; + indices.numBuffers = runEnd - runBase; + indices.cellsPerBuffer = copy.cellsPerBuffer; + indices.cellsPerLease = copy.cellsPerLease; + indices.unitsPerPage = unitsPerPage; + launchSparseOutIndices(indices, stream); } } -std::vector> Batch::resize(std::vector> const& capacities, - std::vector> const& historyLengths, CudaStream stream) +void Batch::writeBackSparseUnits(LifeCycleId lifeCycle, CUstream stream) +{ + auto const leases = _leaseMirror(); + auto const arena = _sparseArena(lifeCycle).view(); + auto const numRows = maxBatchSize(); + + // Which units each row crossed. One scan covers every (layer, row); a row whose + // watermark did not move reports none. + auto const& scratch = sparseStepScratch(lifeCycle); + SparseFlushListArgs listArgs; + listArgs.leases = leases; + listArgs.arena = arena; + listArgs.numRows = numRows; + listArgs.fromUnits = flushFromUnits(lifeCycle); + listArgs.toUnits = flushToUnits(lifeCycle); + listArgs.work = scratch.flushWork; + listArgs.outCount = scratch.flushCount; + launchSparseFlushList(listArgs, stream); + + // One launch for every pool and every layer: the descriptors the Batch built already + // say where each pool's buffers and index tables are, and a flush list is per layer. + SparseUnitCopyArgs copy; + copy.leases = leases; + copy.arena = arena; + copy.numRows = numRows; + copy.copyDesc = scratch.copyDesc; + copy.bufferBase = 0; + copy.numBuffers = scratch.numBuffers; + copy.direction = SparseCopyDirection::kWriteBack; + copy.unitsPerPage = mManager->sparseCacheUnitsPerPage(lifeCycle); + // The work list is every layer of this row, which is what a write-back crosses at once, + // unlike a fetch, which resolves one layer across the whole batch. + copy.work = scratch.flushWork; + copy.gridCtas = sparseCopyGridCtas(scratch.flushWork.capacity, mManager->_smCount()); + launchSparseUnitCopy(copy, stream); +} + +std::vector Batch::sparseOutSlots(std::vector const& buffers) const { - KVCM2_API_GUARD(); - auto const apiLock = mManager->lockExclusive(); - checkOpen(); - checkOutsideCapture(stream); - if (capacities.size() != mRows.size() || historyLengths.size() != mRows.size()) + auto const span = mManager->_sparseLayerSpan(buffers, "sparseOutSlots"); + auto const runs = _sparseCopyTemplates(span.lifeCycle); + auto const& shape = _sparseArena(span.lifeCycle).shape(); + + // What decides whether two buffers sharing a dense table also share an arena index. + // cellsPerBuffer is one value for the whole shape, so a lease spans the same cells in + // two pools exactly when they hold the same number of buffers. + auto geometry = [&](int32_t buffer) { - throw std::invalid_argument("Batch resize arguments must be indexed by stable row"); - } - for (int row = 0; row < mMaxRows; ++row) + for (auto const& run : runs) + { + if (buffer >= run.bufferBase && buffer < run.bufferBase + run.numBuffers) + { + return std::tuple{static_cast(run.cellsPerLease), + static_cast(run.cellsPerBuffer), shape.bufferSlotInPool[static_cast(buffer)]}; + } + } + throw LogicError("sparseOutSlots: a sparse buffer belongs to no pool run"); + }; + + std::vector> tableEntries; + std::vector> geometries; + std::vector slots; + slots.reserve(buffers.size()); + for (size_t i = 0; i < buffers.size(); ++i) { - if ((capacities[row] - && (*capacities[row] < 0 - || static_cast(*capacities[row]) - > static_cast(mMaxBlocks) * mManager->tokensPerBlock())) - || (historyLengths[row] && *historyLengths[row] < 0) - || (mRows[row] == nullptr && (capacities[row] || historyLengths[row]))) + auto const [tableId, entry] = indexTableSlot(buffers[i].layerId, buffers[i].role); + auto const key = std::make_pair(tableId.value(), entry); + auto const geo = geometry(span.bufferBase + static_cast(i)); + auto const found = std::find(tableEntries.begin(), tableEntries.end(), key); + if (found == tableEntries.end()) { - throw std::invalid_argument("Invalid batch resize capacity, history, or empty row"); + tableEntries.push_back(key); + geometries.push_back(geo); + slots.push_back(static_cast(tableEntries.size()) - 1); + continue; } + auto const at = static_cast(std::distance(tableEntries.begin(), found)); + // The dense tables merge on the hot pool's scale and concat offsets; the arena has + // its own geometry and has to agree separately. It normally does, but the two are + // not the same quantity -- a pool carrying only some of the group's layers has the + // same scale and a different cellsPerLease -- and taking the dense answer there + // would alias two different cells onto one slice. + TLLM_CHECK_WITH_INFO(geometries[at] == geo, + "sparseOutSlots: buffers share an index table entry but not their arena geometry, so one " + "slice cannot serve both"); + slots.push_back(static_cast(at)); } - std::vector> results(mMaxRows); - for (int row = 0; row < mMaxRows; ++row) + return slots; +} + +void Batch::fetchSparsePages(std::vector const& buffers, int32_t const* selected, int selectionWidth, + int32_t* out, std::optional streamOpt) +{ + auto const span = mManager->_sparseLayerSpan(buffers, "fetchSparsePages"); + TLLM_CHECK_WITH_INFO(mManager->_selectionLayerCount(span.lifeCycle, span.selectionSlot) == 1, + "fetchSparsePages: this layer shares its selection_id, so its selection is resolved once for every " + "layer naming it; call resolveSparseSelection() then fetchResolvedPages()"); + resolveSparseSelection(buffers, selected, selectionWidth, streamOpt); + fetchResolvedPages(buffers, selected, selectionWidth, out, streamOpt); +} + +void Batch::resolveSparseSelection(std::vector const& buffers, int32_t const* selected, int selectionWidth, + std::optional streamOpt) +{ + CUstream const stream = streamOpt.value_or(cudaStream()); + TLLM_CHECK_WITH_INFO(selected != nullptr, "resolveSparseSelection: selected must be a device pointer"); + + auto const span = mManager->_sparseLayerSpan(buffers, "resolveSparseSelection"); + LifeCycleId const lcId = span.lifeCycle; + + auto const& scratch = sparseStepScratch(lcId); + TLLM_CHECK_WITH_INFO( + scratch.outSlot != nullptr, "resolveSparseSelection: this layer group carries no sparse cache"); + TLLM_CHECK_WITH_INFO(selectionWidth == scratch.selectionWidth, + "resolveSparseSelection: selectionWidth must match the configured numSelectedBlocks"); + + int const unitsPerPage = mManager->sparseCacheUnitsPerPage(lcId); + auto const marks = watermarks(); + + SparseResolveArgs resolve; + resolve.leases = _leaseMirror(); + resolve.arena = _sparseArena(lcId).view(); + resolve.numRows = maxBatchSize(); + resolve.selectionSlot = span.selectionSlot; + resolve.selected = selected; + resolve.selectionWidth = selectionWidth; + resolve.historyTokens = reinterpret_cast(marks.historyLength); + resolve.capacityTokens = reinterpret_cast(marks.capacity); + resolve.unitTokens = mManager->config().tokensPerBlock / unitsPerPage; + resolve.outSlot = scratch.outSlot; + resolve.inFlightSlot = scratch.inFlightSlot; + resolve.maxInFlightUnits = scratch.maxInFlightUnits; + resolve.work = _sparseFetchWork(scratch, span.selectionSlot); + resolve.missCount = scratch.missCount; + // Every row adds its own hits and misses as it resolves them. + resolve.hitMiss = scratch.hitMiss; + launchSparseResolve(resolve, stream); +} + +void Batch::fetchResolvedPages(std::vector const& buffers, int32_t const* selected, int selectionWidth, + int32_t* out, std::optional streamOpt) +{ + CUstream const stream = streamOpt.value_or(cudaStream()); + TLLM_CHECK_WITH_INFO( + selected != nullptr && out != nullptr, "fetchResolvedPages: selected and out must be device pointers"); + + auto const span = mManager->_sparseLayerSpan(buffers, "fetchResolvedPages"); + LifeCycleId const lcId = span.lifeCycle; + auto const bufferBase = span.bufferBase; + auto const numBuffers = span.numBuffers; + + auto const& scratch = sparseStepScratch(lcId); + TLLM_CHECK_WITH_INFO(scratch.outSlot != nullptr, "fetchResolvedPages: this layer group carries no sparse cache"); + TLLM_CHECK_WITH_INFO(selectionWidth == scratch.selectionWidth, + "fetchResolvedPages: selectionWidth must match the configured numSelectedBlocks"); + + int const unitsPerPage = mManager->sparseCacheUnitsPerPage(lcId); + auto const leases = _leaseMirror(); + auto const arena = _sparseArena(lcId).view(); + auto const numRows = maxBatchSize(); + auto const layerWork = _sparseFetchWork(scratch, span.selectionSlot); + + // The misses the selection produced, applied to each of this layer's buffers. Layers + // sharing the selection each run this and each fill their own mirror of the entry. + // + // One launch for every pool of this layer. The descriptors the Batch built already + // clip each pool to the layer's buffers and name its index table, so nothing here + // depends on which layer is being fetched beyond choosing the buffer range. + SparseUnitCopyArgs copy; + copy.leases = leases; + copy.arena = arena; + copy.numRows = numRows; + // The same descriptors as the write-back: a run never spans layers, so it already + // names this layer's index table. Only the buffer range below differs. + copy.copyDesc = scratch.copyDesc; + copy.bufferBase = bufferBase; + copy.numBuffers = numBuffers; + copy.direction = SparseCopyDirection::kFetch; + copy.unitsPerPage = unitsPerPage; + copy.work = layerWork; + copy.gridCtas = sparseCopyGridCtas(scratch.fetchWork.capacity, mManager->_smCount()); + launchSparseUnitCopy(copy, stream); + + // The index table the attention kernel reads is derived from the resolve's slots alone: + // it neither reads a cell nor waits for one to be filled. The copy releases its launch + // dependency as soon as it is addressed, so this follows it on the stream without + // waiting for it to finish. + auto const plan = mManager->_sparseBufferPlan(lcId); + // Buffers whose indices coincide share a slice, so a run that only repeats slices an + // earlier one already wrote is not launched at all. A layer's data and scale pools + // mirror, which is exactly that case. + auto const slots = sparseOutSlots(buffers); + std::vector written(static_cast(numBuffers), false); + for (auto const& pool : _sparseCopyTemplates(lcId)) { - if (auto* cache = mRows[row]) + auto const runBase = std::max(pool.bufferBase, bufferBase); + auto const runEnd = std::min(pool.bufferBase + pool.numBuffers, bufferBase + numBuffers); + if (runBase >= runEnd) { - results[row] = cache->resize(capacities[row], historyLengths[row]); + continue; } + // A run's buffers are consecutive entries of one table, so their slices are + // consecutive too and the first one says where the run writes. + int const firstSlot = slots[static_cast(runBase - bufferBase)]; + if (written[static_cast(firstSlot)]) + { + continue; + } + for (int32_t buffer = runBase; buffer < runEnd; ++buffer) + { + written[static_cast(slots[static_cast(buffer - bufferBase)])] = true; + } + SparseOutIndexArgs indices; + indices.leases = leases; + indices.arena = arena; + indices.numRows = numRows; + indices.selected = selected; + indices.outSlot = scratch.outSlot; + indices.selectionWidth = selectionWidth; + indices.out = out; + // The row stride is the caller's, not the slice count: `out` keeps one entry per + // buffer, and a slice no buffer maps to is simply never written. + indices.outBuffers = numBuffers; + indices.outBase = firstSlot; + indices.bufferBase = runBase; + indices.numBuffers = runEnd - runBase; + indices.cellsPerBuffer = pool.cellsPerBuffer; + indices.cellsPerLease = pool.cellsPerLease; + indices.unitsPerPage = unitsPerPage; + launchSparseOutIndices(indices, stream); } - publish(stream); - return results; } -MemAddress Batch::pageTableAddress(LayerGroupId group) const +void Batch::_buildSparseArenas() { - if (group.value() < 0 || group.value() >= mNumLayerGroups) + LifeCycleId const numLc = mManager->numLifeCycles(); + mSparseArenas = TypedVec(numLc); + + // One lease per row this Batch can seat, which is the concurrency the arenas serve: a + // row's member is stepped, and stepping is what reads the cells. + int const maxLeases = mMaxBatchSize; + + // Every row owns a lease for the Batch's lifetime. move() permutes the mapping, so a + // member carries its cells between rows without them being copied. + mRowLease.resize(BatchRow{maxLeases}); + for (int row = 0; row < maxLeases; ++row) { - throw std::out_of_range("Invalid batch layer group"); + mRowLease[BatchRow{row}] = SparseLeaseIndex{row}; + } + + for (LifeCycleId lcId{0}; lcId < numLc; ++lcId) + { + if (!mManager->hasSparseFetchPath(lcId)) + { + continue; + } + auto const plan = mManager->_sparseBufferPlan(lcId); + TLLM_CHECK(!plan.empty()); + + auto const cellsPerBuffer = mManager->_sparseCacheSubSlotsPerBuffer(lcId); + SparseCacheInit shape{ + .numLayers = mManager->_sparseCacheLayerCount(lcId), + .numSelections = mManager->_sparseCacheGroupCount(lcId), + .numBuffers = mManager->sparseCacheBufferCount(lcId), + .entriesPerSelection = cellsPerBuffer * mManager->sparseCacheUnitsPerPage(lcId), + .cellsPerBuffer = cellsPerBuffer, + .layerSelection = mManager->_sparseLayerSelection(lcId), + }; + (void) checkedSparseResolveSharedBytes(shape.entriesPerSelection, mManager->_sparseCacheSelectionWidth(lcId), + mManager->_sparseCacheMaxInFlightUnits(lcId)); + + // The arena mirrors the hot pool group: one region per pool the life cycle's sparse + // buffers draw on, sized by that pool's cell and the buffers sharing it. + size_t numPools = 0; + for (auto const& entry : plan) + { + numPools = std::max(numPools, static_cast(entry.poolIndex.value()) + 1); + } + shape.pools.assign(numPools, {0, 0}); + for (auto const& entry : plan) + { + shape.bufferLayer.push_back(entry.layerRow); + shape.bufferSlotInPool.push_back(entry.slotInPool); + auto& [cellBytes, buffers] = shape.pools[static_cast(entry.poolIndex.value())]; + cellBytes = entry.size; + ++buffers; + } + // A pool index the life cycle's sparse buffers never use keeps its place in the list, + // because the arena is indexed by the storage pool index rather than by a dense + // numbering of its own. It stays {0, 0}: holding no cells, it takes no space, and its + // base is simply where the next pool's region starts. + mSparseArenas[lcId] = SparseCellArena(shape, maxLeases); } - return reinterpret_cast(mDeviceMemory.get()) + rowOffset(group.value(), 0) * sizeof(int32_t); } -MemAddress Batch::numBlocksAddress(LayerGroupId group) const +std::vector Batch::_sparseCopyTemplates(LifeCycleId lifeCycle) const { - if (group.value() < 0 || group.value() >= mNumLayerGroups) + auto const plan = mManager->_sparseBufferPlan(lifeCycle); + TLLM_CHECK(!plan.empty()); + auto const& arena = mSparseArenas[lifeCycle]; + + std::vector templates; + for (size_t begin = 0; begin < plan.size();) { - throw std::out_of_range("Invalid batch layer group"); + auto const pool = plan[begin].poolIndex; + size_t end = begin; + while (end < plan.size() && plan[end].poolIndex == pool) + { + ++end; + } + + // A pool holds buffers of one size, but not necessarily of one unit layout, and the + // copy walks a run with one set of strides. + auto const& layout = plan[begin].layout; + for (size_t i = begin; i < end; ++i) + { + auto const& other = plan[i].layout; + if (other.numChunks != layout.numChunks || other.chunkBytes != layout.chunkBytes + || other.chunkStride != layout.chunkStride || other.unitStride != layout.unitStride) + { + throw LogicError("Sparse buffers sharing a pool must share a selection unit layout"); + } + } + + SparseCopyRun args; + // The page side is the life cycle's own hot pool: host memory, registered device-mapped, + // so the kernel addresses it directly. + // + // A converted index counts whole buffers, not coalesced slots, so the stride is one + // buffer and the pool's slot size never enters. The base belongs to the index table + // the run was declared in, which only the caller's Batch knows, so it is left unset + // here. + args.hostPageStride = static_cast(plan[begin].size); + // The hot side is the cache arena, not a storage pool. + args.deviceBase = arena.base(pool); + args.bufferBytes = static_cast(arena.bufferBytes(pool)); + args.cellsPerBuffer = mManager->_sparseCacheSubSlotsPerBuffer(lifeCycle); + args.cellsPerLease = arena.cellsPerLease(pool); + args.layout = layout; + args.bufferBase = static_cast(begin); + args.numBuffers = static_cast(end - begin); + args.width = sparseCopyWidth(args.layout, args.hostPageStride, args.bufferBytes); + templates.push_back(args); + begin = end; } - return reinterpret_cast(mDeviceMemory.get()) + countOffset(group.value(), 0) * sizeof(int32_t); + return templates; +} + +MemAddress Batch::getSparseCacheBaseAddress(LayerId layerId, DataRole const& role) const +{ + LifeCycleId const lcId = mManager->getLayerGroupId(layerId); + auto const& arena = mSparseArenas[lcId]; + auto const plan = mManager->_sparseBufferPlan(lcId); + auto const it = std::ranges::find_if( + plan, [&](SparseBufferPlan const& entry) { return entry.id.layerId == layerId && entry.id.role == role; }); + TLLM_CHECK_WITH_INFO( + !arena.empty() && it != plan.end(), "getSparseCacheBaseAddress: this buffer has no sparse fetch path"); + return static_cast(arena.base(it->poolIndex)); } } // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/batch.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/batch.h index af91ee482345..e02bbe82bf74 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/batch.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/batch.h @@ -19,14 +19,19 @@ #include "kv_cache_manager_v2/common.h" #include "kv_cache_manager_v2/lifeCycleRegistry.h" +#include "kv_cache_manager_v2/sparseCacheKernels.h" #include "kv_cache_manager_v2/stagingBuffer.h" #include "kv_cache_manager_v2/utils/cudaEvent.h" +#include "kv_cache_manager_v2/utils/generator.h" #include "kv_cache_manager_v2/utils/staticHostMem.h" +#include + #include -#include +#include #include #include +#include #include namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 @@ -35,103 +40,978 @@ namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 class KvCache; class KvCacheManager; -//! Stable request rows and raw device metadata, shared across all layer groups. +//! One layer group's page table as the device sees it. Always packed: the Batch allocates +//! it, so a row is exactly `blocksPerRow` entries. +struct GpuCoalescedPageTable +{ + CUdeviceptr data = 0; + int32_t numRows = 0; + int32_t blocksPerRow = 0; +}; + +//! One index table's converted page indices as the device sees them, laid out +//! [numRows][numBuffers][blocksPerRow]. A buffer reads the entry `index_table_slot` gives +//! it. +//! +//! Row-major in the entries so a request's buffers sit together: an attention kernel takes +//! a key and value page list as one pair -- XQA's is +//! [batchSize][beamWidth][2][maxNbPagesPerSeq] -- and this is that shape, so a table can be +//! handed over as it stands instead of being repacked every step. +struct ConvertedIndexTable +{ + CUdeviceptr data = 0; + int32_t numRows = 0; + //! One per buffer of the requests this table serves; they differ only by the offset + //! baked into them. + int32_t numBuffers = 0; + int32_t blocksPerRow = 0; +}; + +//! What a converted index means: the address it is relative to, and what it counts. +//! +//! A property of the spec rather than of the table it landed on. Specs whose converted +//! entries coincide are merged onto one table, and reading that one set of entries through +//! different bases is exactly what keeps a quantised layer's data and its block scales +//! apart. So a table has as many of these as it serves specs, and only the spec picks one. +struct PoolSpec +{ + //! What a converted index is added to. + MemAddress base = 0; + //! Bytes a converted index counts, which is the pool's page stride. + size_t pageStride = 0; +}; + +//! One lifecycle's page table as the host writes it, in the Batch's pinned staging. +struct CoalescedPageTable +{ + //! Points into the pinned staging block, which publish() uploads asynchronously. A host + //! write here races that copy and is in any case overwritten by the next resize(). + int32_t const* data = nullptr; + int32_t numRows = 0; + int32_t blocksPerRow = 0; +}; + +//! Identifies one converted index table of a Batch. +using IndexTableId = StrongIndex; + +//! Device scratch one lifecycle's step works through: what a fetch resolves into before +//! the copy and the index scatter read it back, and what a write-back lists before its own +//! copy reads it. +//! +//! Batch-owned: both consumers are one launch spanning every row, and every extent is a +//! configuration constant, so one allocation serves every step and the step allocates +//! nothing -- which is what lets a CUDA graph capture it. +struct SparseStepScratch +{ + int32_t* outSlot = nullptr; //!< [numRows][selectionWidth] + int32_t* missCount = nullptr; //!< [numRows] + int32_t* inFlightSlot = nullptr; //!< [numRows][maxInFlightUnits] + int32_t selectionWidth = 0; + int32_t maxInFlightUnits = 0; + + //! The misses of one selection, compacted across rows. A fetch resolves one selection + //! at a time, so a single list serves the whole batch -- and every layer sharing that + //! selection, each copying the same units into its own cells. + SparseWorkList fetchWork; + + //! The write-back's counterpart. A write-back crosses every layer of a row at once and + //! the copy that consumes it walks buffers, so it takes one list per layer rather than + //! the one the fetch folds away. + SparseWorkList flushWork; + int32_t* flushCount = nullptr; //!< [numRows][numLayers] + int32_t numLayers = 0; + + //! The Batch's cumulative fetch statistics, the same pointer for every life cycle. + //! Carried here so a fetch reaches it the way it reaches the rest of its arguments. + SparseHitMiss* hitMiss = nullptr; + + //! What differs between the pools of a copy, on the device for the kernel to index by + //! `blockIdx.z`. Built and uploaded once with the Batch: every field comes from the + //! buffer plan and the index tables, both fixed at construction, so a captured graph + //! replays a copy with no host write at all. + //! + //! [numBuffers] one copy descriptor per BufferId, serving both directions. BufferIds + //! are ordered (layer, pool), so one already names its own layer's index-table entry + //! and needs no layer dimension; the fetch and the write-back then want the same + //! descriptors and differ only in which of them they launch over. + SparseCopyDesc const* copyDesc = nullptr; + int32_t numBuffers = 0; +}; + +//! Per-row watermarks, in tokens. Everything a step needs about how far a request has got +//! derives from these two, at whatever granularity the consumer works in: //! -//! Driven by one owning thread, like its member requests. Membership is non-owning: -//! closing/destroying a request detaches it; closing/destroying a batch detaches its -//! live members without closing them. Mutable state uses the manager's API lock. -//! Publish outside graph capture, then waitReady on a consumer stream. Submit all -//! reads and call recordRead before mutating requests, rows, or tables. -class Batch : public std::enable_shared_from_this +//! num_blocks = historyLength / tokens_per_block (valid page table entries) +//! history_units = historyLength / unit_tokens (selectable units) +//! capacity_units = capacity / unit_tokens (the in-flight span's end) +//! +//! Publishing the source rather than the three derived counts keeps one truth; the divisors +//! are per-layer-group constants a kernel already has. +//! +//! A free row reads zero, so it describes a request with nothing eligible. +struct WatermarkView +{ + CUdeviceptr historyLength = 0; + CUdeviceptr capacity = 0; + int32_t numRows = 0; +}; + +// --------------------------------------------------------------------------- +// Batch — the KvCaches stepped together, and their device-side mirror. +// +// Owns the row order and the per-layer-group page tables the sparse fetch reads. Rows are +// slots, not positions in a compacted list: removing a member leaves a hole rather than +// shifting later rows, so array shapes and addresses survive a CUDA graph capture. +// +// Membership is exclusive. A KvCache detaches itself on destruction, so a dropped request +// cannot leave a dangling row. +// --------------------------------------------------------------------------- +class Batch { public: - Batch(std::shared_ptr manager, int maxRows, int maxBlocks, int maxBeamWidth = 1); + //! @param indexTables what consumers need addressable together. Requests whose + //! converted indices coincide are merged onto one table, which is how a quantised + //! layer's data and block scales come to share one. + //! + //! @param maxBatchSize rows, i.e. the most requests that can be stepped together. + //! @param maxBlocksPerSeq columns, i.e. the longest sequence a row can describe. + //! + //! Both bound device arrays whose shape must not change between steps, so neither can + //! be grown later. Left unset they derive from the manager's `constraints`, which + //! throws when none were declared. + //! + //! @param mode page index mode, applied to every life cycle. Left unset, each life + //! cycle takes the smaller table -- SHARED where the manager says it is always + //! available, PER_LAYER otherwise. + //! + //! Set, it overrides that answer rather than being checked against it, and SHARED may + //! be named where KvCacheManager::supportsIndexMode() returns nullopt. That answer + //! covers every request the life cycle could carry, so a sliding window leaves it + //! undecided; naming SHARED is the caller's narrower claim that no member of *this* + //! batch will hold scratch slots there. A decode-only batch is the case it exists for: + //! scratch covers the blocks a prefill chunk adds, and a one-token request adds none. + //! + //! The claim is validated, not trusted: every publish() checks each member against + //! KvCache::supportsIndexMode(SHARED, lc) and throws if one holds scratch after all. + //! Seat only members that keep it -- the throw arrives mid-step, with the batch already + //! built around the smaller table. + //! + //! The mode is fixed here because a converted table's shape is baked into a captured + //! graph. It cannot follow a member's scratch state, which varies per step. + //! + //! @param stream the stream every member runs on: a member's own work and the Batch's + //! transfers have to be ordered against each other. setCudaStream() moves the Batch and + //! its members together, which is the only way it changes. The caller keeps it -- a + //! Batch neither creates nor destroys it, and it must outlive the Batch. + //! + //! @param maxTotalInputTokens tokens one step may ingest across every row, as the + //! runtime's own batching bound gives it. It sizes the write-back's work list: a row + //! advancing `t` tokens crosses at most `t / unitTokens + 1` units, so the batch + //! crosses at most `maxTotalInputTokens / unitTokens + maxBatchSize` of them. + //! + //! Left unset it falls back to `maxBatchSize * maxInFlightUnits`, which is + //! `maxBatchSize` -- a decode step advances capacity one token per row, so one row + //! crosses at most one unit boundary. Pass it once a step can ingest more than that, + //! which is what chunked prefill under in-flight batching does. + explicit Batch(KvCacheManager& manager, CUstream stream, std::vector indexTables, + std::optional maxBatchSize = std::nullopt, std::optional maxBlocksPerSeq = std::nullopt, + std::optional mode = std::nullopt, std::optional maxTotalInputTokens = std::nullopt); + + //! The stream the Batch and every member run on. + [[nodiscard]] CUstream cudaStream() const noexcept + { + return mCudaStream; + } + + //! Moves the Batch and every member onto `stream`. + //! + //! The new stream is made to wait on what the old one already has queued, so the + //! caller need not synchronize: the device arrays and the staging buffer are the same + //! objects either side of the move, and work that names them stays ordered. + //! + //! A KvCache's own stream setter refuses while it is a member. + //! + //! Like the constructor, this takes no ownership: the caller's stream must outlive the + //! Batch, or be replaced here before it is destroyed. + void setCudaStream(CUstream stream); + + //! Page index mode of one life cycle, as chosen at construction. + [[nodiscard]] PageIndexMode indexMode(LayerGroupId layerGroupId) const; + + // ---- Converted index tables -------------------------------------------- + + [[nodiscard]] IndexTableId numIndexTables() const noexcept; + + //! Every request one table satisfies. A request's buffer order is its own mapping into + //! the buffer dimension, so the first buffer of each sits at entry zero. + [[nodiscard]] std::vector const& indexTable(IndexTableId tableId) const; + + //! Which table serves this spec. The spec must be one the Batch was built with. + //! + //! Equivalent specs resolve to the same id -- a layer's data and its block scales name + //! one table between them. The spec's buffers occupy entries [0, buffers.size()) in its + //! own order. + [[nodiscard]] IndexTableId indexTableOf(IndexTableSpec const& spec) const; + + //! How to turn this spec's converted indices into addresses. + //! + //! Keyed by the spec because the base is: a merged table serves several specs and each + //! reads its entries through its own pool. The spec must be one the Batch was built + //! with, and its buffers must share a page stride -- they sit in one pool, so one base + //! serves them all. + //! + //! The result is stable for the Batch's lifetime, so a captured graph may bake it. + [[nodiscard]] PoolSpec poolSpecFor(IndexTableSpec const& spec) const; + + //! One table's converted indices: int32[max_batch][entries][blocks]. + //! + //! A row carries every entry before the next row starts, which is what puts a spec's + //! key and value planes side by side. So one entry's plane is strided, not contiguous: + //! it begins at `entry * blocks` elements in and repeats every `entries * blocks`. + //! + //! The address is stable for the Batch's lifetime, so a captured graph may bake it. + [[nodiscard]] ConvertedIndexTable convertedIndexTable(IndexTableId tableId) const; + ~Batch(); Batch(Batch const&) = delete; Batch& operator=(Batch const&) = delete; - //! Attach a request at a chosen row or the first free row. Membership is exclusive. - int add(KvCache& cache, std::optional row = std::nullopt); - //! Detach a member, leaving its row dirty so the next publication clears it. - void remove(KvCache& cache); - //! Detach all members. Exported arrays retain their allocation until the last owner dies. - void close(); - - //! Retry deferred sparse offloads, then upload final dirty rows and counts; return uploaded rows. - //! Offload failures propagate and retain pending work for retry. - //! Staging buffers are retained until their asynchronous copies complete. - std::vector publish(CudaStream stream); - //! Wait for publication and KV readiness. Reject unpublished changes. - void waitReady(CudaStream stream) const; - //! Fence submitted reads before table overwrite and request storage release. - void recordRead(CudaStream stream); - //! Resize by stable row slot, then publish once. Empty rows return nullopt. - std::vector> resize(std::vector> const& capacities, - std::vector> const& historyLengths, CudaStream stream); - - std::vector dirtyRows() const; - - int maxRows() const noexcept + //! Seat a cache in `row`, or in the lowest free row when none is named. + //! + //! Seat a request only once it is ready to be stepped. Seating takes the sparse arena + //! lease, so a request that holds pages without being stepped -- a disaggregated one + //! still receiving its KV, or one sitting out a step -- costs the arena nothing while it + //! stays out. Nothing here checks it: an early seat costs arena pressure, not a throw. + //! + //! Throws when the row is taken or out of range, the batch is full, the cache already + //! belongs to a batch, or its beam width exceeds one -- only beam 0 is mirrored, so a + //! wider request would have every beam read beam 0's pages. + BatchRow addKvCache(KvCache& kvCache, std::optional row = std::nullopt); + + //! Free the cache's row, leaving a hole, and hand back the reference the Batch held. + //! Throws when it is not a member. + //! + //! Returns the arena lease, which discards whatever the GPU cache held: a lease comes back + //! with every entry free. Call flushSparseUnits() first for anything worth keeping. + std::shared_ptr removeKvCache(KvCache& kvCache); + + //! The same, by the row addKvCache returned. Null when nothing is seated there. + std::shared_ptr removeKvCache(BatchRow row); + + //! Move the member at `srcRow` to the empty `dstRow`, leaving `srcRow` free. + //! + //! Membership does not lapse: the arena lease, the GPU cache contents behind it and the + //! stream binding all survive, which is what separates this from removeKvCache() followed + //! by addKvCache(). That pair would hand the lease back -- a lease returns with every + //! entry free -- and so would discard residency the member is still entitled to. + //! + //! Throws when either row is out of range, when `srcRow` holds nothing, or when `dstRow` + //! is occupied -- including `dstRow == srcRow`, which is both at once. + //! + //! Waits for the previous staging upload before touching either row. Device tables retain + //! their addresses; the move reaches them on the next publish. + void move(BatchRow dstRow, BatchRow srcRow); + + //! Resize every member, then publish once. + //! + //! Waits on the previous publish before touching any member. Per-request operations with + //! no Batch equivalent -- commit(), a beam fork, an SWA recycle -- still need an explicit + //! waitPublished(). + //! + //! One result per row, in row order; false for a free row. A member that fails rolls + //! itself back, so a batch can partially fail and the table published afterwards + //! describes the post-rollback state. + //! + //! `capacities` and `historyLengths` are indexed by row, so both must be maxBatchSize + //! long. A nullopt leaves that dimension alone, matching KvCache::resize. + //! Rejects, before touching any member, a history length that falls or that outruns the + //! capacity its row already had: capacity granted by this call makes room for tokens + //! written after it, so history can only cover tokens whose space was already there. + //! + //! A row whose growth fails still takes its history advance, which is legal by that same + //! check and needs no slots -- both the stale range and the scratch range are monotone + //! in history, so raising it only releases. Its result is still 0: the row is not at the + //! state that was asked for. + std::vector resize( + std::vector> const& capacities, std::vector> const& historyLengths); + + //! Copy the staged rows to the device, enqueuing on the Batch's stream. + //! + //! The staging rows are written by the members themselves, so they are always current; + //! what this adds is a stream-ordered snapshot. Without it the host would be mutating + //! the table a previously enqueued fetch is still reading. + //! + //! Two stages. The host-to-device fill runs on an upload stream of this Batch's own, so + //! it can overlap the previous step; the Batch's stream then carries one linear + //! device-to-device copy into the tables the fetch reads. Everything expensive is in the + //! first stage, and only the second has to wait for anything. + //! + //! **The second stage must be ordered after the kernels that read the previous + //! publish.** Those live in a captured graph, so the tables cannot be double-buffered + //! out of the hazard -- the capture baked their addresses, which is why the staging + //! sits upstream of them instead. The Batch's stream being the one the graph replays on + //! is what provides that ordering. A torn publish is not a stale answer: the cache + //! views carry pointers, so it can pair one request's residency with another's lease. + //! + //! The copy is asynchronous and reads pinned memory, so it is still in flight when this + //! returns. **Every later host write to a member's page indices must be separated from + //! it by waitPublished()** -- a KvCache::resize(), a commit(), a beam fork, an SWA + //! recycle, anything that locks or unlocks a page. Writing sooner races the copy and + //! publishes a torn table. This Batch's own mutators are not in that list: resize(), + //! move(), addKvCache() and removeKvCache() each wait for themselves. + //! + //! Throws when the previous publish has not been waited on, which is how that rule is + //! enforced: this call stages the watermarks and views itself, so it would be writing + //! into the same hazard. It refuses rather than waiting because a wait here covers only + //! its own write -- a member mutated before the call has already raced by then, and + //! waiting would hide it. + void publish(); + + //! Write back the units each row's watermark is about to pass, given the history the + //! caller is about to set. One per row; a row may repeat its current history, which + //! asks for nothing. + //! + //! **Call this before advancing history, not after.** A block must be in its host page + //! before it can be selected, and selectable means below the watermark; resize() refuses + //! a history this has not covered. + //! + //! Bounded by what was in flight, not by how far history travels: only a unit that was + //! in flight has a cell holding its one copy, so a resumed request moving history across + //! a whole prompt writes back nothing. What is refused is a batch whose rows together + //! cross more units than the write-back list can name -- the list is one per layer and + //! compacted across rows, sized from `maxTotalInputTokens`, so one prefilling row may + //! use the budget several decoding rows leave unspent. + //! + //! Reads the index table as of the last publish, which covers the blocks in range for + //! the same reason -- they were allocated by an earlier step. + void flushSparseUnits(std::vector const& newHistoryLengths); + + //! Where each of `buffers` reads its indices in fetchSparsePages()' output. + //! + //! Buffers whose converted indices coincide share a slice. A layer's data and scale + //! pools mirror -- equal buffer counts, so a lease spans the same cells in both -- and + //! a selection then resolves to one index serving each, the caller applying the base + //! of the pool it is reading. Buffers that do not coincide keep slices of their own: a + //! single scale shared by key and value is its own index table, so it gets one. + //! + //! One entry per buffer, in the order given. A slice no buffer maps to is never + //! written and holds whatever was there: read a buffer's indices at its own slot and + //! nowhere else. + [[nodiscard]] std::vector sparseOutSlots(std::vector const& buffers) const; + + //! Resolve a batch's selections against each buffer's per-request history cache. + //! + //! Device-driven and stream-ordered throughout: the selection is read from device + //! memory, so the host neither inspects it nor decides anything, and the call enqueues + //! onto `stream` and returns. Every shape is static across steps, so a CUDA graph can + //! capture it. + //! + //! `buffers` must be exactly one layer's sparse buffers. One layer, not one layer + //! group, because a layer's selection is its own unless it names a selection_id: an + //! indexer of its own reads the previous layer's output, so a step's selections are not + //! all known at its start. Layers that do share one resolve it through + //! resolveSparseSelection() instead, which this refuses. The whole layer, not + //! a subset, because the resolve marks entries resident on behalf of every buffer at + //! once: filling only some would leave the rest stale while the metadata claimed + //! otherwise. + //! + //! Only for a layer whose selection is its own. Where several layers name one + //! selection_id their selection is known for all of them at once, and resolving it once + //! per layer would have the later resolves count the earlier ones' admissions as hits + //! and leave their cells unfilled; those layers call resolveSparseSelection() once and + //! fetchResolvedPages() per layer instead. This rejects them rather than quietly + //! mis-serving them. + //! + //! `selected` is [maxBatchSize][selectionWidth] unit ordinals, negative where padded, + //! and must name only units below the history watermark -- a unit at or above it has + //! not been written back, so fetching it would copy stale bytes. That is a wrong + //! answer rather than a memory error, and the call cannot detect it, so it is the + //! selector's responsibility. + //! + //! `out` is [maxBatchSize][buffers.size()][selectionWidth], written in full, in the + //! expanded index space getPageIndexConverter() produces: entry k is the index holding + //! `selected[row][k]`, or kBadPageIndex where that slot carries no selection. + //! + //! Relative to getSparseCacheBaseAddress() for the buffer, not getMemPoolBaseAddress() + //! and not a cache level: the pool names the buffer's pages, which are host memory for a + //! sparse life cycle, while the cells a kernel reads live in the arena. + //! + //! Compact, in the selection's own order, rather than a table indexed by unit ordinal, + //! because that is the shape sparse attention kernels take. See + //! SPARSE_ATTENTION_DESIGN.md's "Fetch". + //! + //! `stream` defaults to the Batch's own, which is the only correct choice while the + //! layers are fetched in order: a layer's selection comes from the one before it. Name + //! a stream to run ahead of that -- a model that fetches several layers early can do so + //! on its own -- or to be captured, since a captured fetch must name the capture stream. + //! + //! Naming one moves the ordering to the caller, once per step rather than once per call: + //! `stream` must wait on the Batch's stream after the publish() that prepared the step, + //! and the Batch's stream must wait on `stream` before any call that begins the next one. + //! Fetches within a step need no join between them, which is what naming a stream buys. + //! Capturing them and replaying the graph on the Batch's stream satisfies both. + //! + //! Both halves are load-bearing: the resolve reads Batch state the publish wrote, and + //! writes a row's residency back into state the next step rewrites on the Batch's stream. + void fetchSparsePages(std::vector const& buffers, int32_t const* selected, int selectionWidth, + int32_t* out, std::optional stream = std::nullopt); + + //! Resolve one selection, admitting what it names and listing what has to be copied. + //! + //! The first half of fetchSparsePages, split out for the layers that share a selection: + //! they resolve once between them and then fetch a layer at a time. `buffers` names any + //! one member layer -- the residency and the miss list belong to the selection, not to + //! the layer that happened to name it -- on the same terms as fetchSparsePages. + //! + //! It leaves two things in the life cycle's step scratch: the miss list the copies walk + //! and the per-slot cell mapping the index writes read. Both belong to the life cycle, + //! not to the selection, so the next resolve of any selection in it overwrites them and + //! must be ordered after every fetchResolvedPages() that consumes the current one. One + //! stream gives that ordering; across streams it takes an event. A model gets it by + //! running a group's layers together, as the SWA scratch already assumes of the + //! currently executing layer. + void resolveSparseSelection(std::vector const& buffers, int32_t const* selected, int selectionWidth, + std::optional stream = std::nullopt); + + //! Copy one layer's share of a resolved selection into its cells, and write its indices. + //! + //! The second half of fetchSparsePages. `selected` must be the array the matching + //! resolveSparseSelection() read, since the indices written are those units'; `out` is + //! as fetchSparsePages describes it, for this layer's buffers alone. + void fetchResolvedPages(std::vector const& buffers, int32_t const* selected, int selectionWidth, + int32_t* out, std::optional stream = std::nullopt); + + //! Give the in-flight unit of each request a cell, and report where to write it. + //! + //! The unit between the two watermarks is the one the model is about to produce. It has + //! no host copy yet -- a unit reaches its page only when resize() carries the watermark + //! past it -- so nothing is fetched here: the call assigns a cell and leaves its bytes + //! to the model. + //! + //! Needs only the watermarks, not the selection, so it runs before fetchSparsePages. + //! + //! Idempotent while the watermarks and the reservation hold: the resolve probes + //! residency before it takes a victim, so a later call -- including the fetch's own + //! resolve -- finds the entry and reuses it rather than assigning a second. The entry + //! is pinned for as long as its unit is in flight, so no fetch can evict it. + //! + //! `buffers` is one layer's sparse buffers, exactly as for fetchSparsePages. + //! + //! `out` is [maxBatchSize][buffers.size()], written in full, in the same expanded index + //! space fetchSparsePages produces and relative to the same per-pool arena region: + //! kBadPageIndex where the row holds no reservation or has nothing in flight. + //! `stream` defaults to the Batch's own, on the same terms as fetchSparsePages. + void prepareSparseWriteIndices( + std::vector const& buffers, int32_t* out, std::optional stream = std::nullopt); + + //! The history each row has been written back to. A row that has never needed a flush + //! reads its own history, so a caller may always compare against it. + [[nodiscard]] int flushedHistory(BatchRow row) const; + + //! Block until the last publish() has finished reading the staging rows. + //! + //! A no-op when nothing is in flight, so a caller may use it freely at the top of a step. + void waitPublished(); + + //! One layer group's page table on the host, as a publish will upload it. Read-only: + //! the entries are the members' to write, through their page locks. + [[nodiscard]] CoalescedPageTable coalescedPageTable(LayerGroupId layerGroupId) const; + + //! Per-row history length and capacity, published by the same transfer as the page + //! tables. They have to travel together: the lengths bound which entries of a row are + //! valid, so a kernel pairing one step's table with another step's lengths would read + //! past the end of a row that has since been reused. + [[nodiscard]] WatermarkView watermarks() const; + + //! Row of a member; nullopt when this cache is not one. + [[nodiscard]] std::optional rowOf(KvCache const& kvCache) const; + + //! Members seated, which is the rows occupied. + [[nodiscard]] int numKvCaches() const noexcept; + + //! Cumulative fetch statistics, in selection units, since the Batch was built or the + //! counters were last reset. + //! + //! A hit is a selected unit the GPU cache already held; a miss is one the fetch had to + //! copy from its host page. Padded selection entries and in-flight units are neither, + //! so they reach neither figure -- which is why this comes from the resolve rather than + //! from a count the caller could keep. + //! + //! Reads back on a temporary stream and syncs only that, so the counters may be a few + //! launches stale. That does not move a ratio; a caller wanting the two to agree + //! exactly synchronises the device first and calls this after. + [[nodiscard]] SparseHitMiss peekHitMissCounter() const; + + //! Zero the cumulative pair. Enqueued on the Batch's stream, so it orders against the + //! fetches around it rather than tearing one. + void resetHitMissCounter(); + + //! Device bytes every life cycle's sparse arena holds, zero without sparse offload. + //! + //! The arenas sit outside the cache tiers' quota, so a caller budgeting device memory + //! must account for them separately. + [[nodiscard]] size_t sparseArenaBytes() const; + + //! Whether the occupied rows are exactly [0, numKvCaches()). + //! + //! A Batch does not require it. Removing a member leaves its row free and the next join + //! takes the lowest free one, so a hole is legal and transient. + [[nodiscard]] bool isContiguous() const noexcept; + + [[nodiscard]] int maxBatchSize() const noexcept + { + return mMaxBatchSize; + } + + [[nodiscard]] int maxBlocksPerSeq() const noexcept + { + return mMaxBlocksPerSeq; + } + + //! Base of the sparse cache arena region backing this buffer, for a consumer reading + //! units the fetch resolved. Throws when the buffer has no sparse fetch path. + //! + //! Role-dependent: the arena mirrors the hot pool group, so buffers of different sizes + //! sit in regions of their own and an index the fetch emits is relative to the region of + //! the buffer it was emitted for. + //! + //! Distinct from KvCacheManager::getMemPoolBaseAddress, which names the buffer's *pages* + //! for every buffer alike -- that is what the transceiver and the KV transfer path + //! address, and it keeps one meaning whether or not a buffer is sparse. The arena is a + //! second allocation holding the only GPU-resident copy a kernel can read, and it is the + //! Batch's because it is sized by the rows the Batch seats. + [[nodiscard]] MemAddress getSparseCacheBaseAddress(LayerId layerId, DataRole const& role) const; + + //! The member seated in `row`, or null when that row is free. + //! + //! The inverse of rowOf(). Throws when `row` is out of range, like removeKvCache(int): + //! a free row is an ordinary answer, so null means empty and never means "no such row". + //! + //! The reference stays valid for the Batch's lifetime -- the row vector is sized once at + //! construction -- but what it holds changes as members are seated and removed, so it is + //! a view of the slot rather than of a member. + [[nodiscard]] std::shared_ptr const& operator[](BatchRow row) const; + +private: + //! Reads the private step state below in order to report it. + friend class KvCacheIntrospection; + + //! Where one table's converted indices live. + [[nodiscard]] CUdeviceptr indexTableAddress(IndexTableId tableId) const; + + //! How a sparse copy addresses the run of buffers beginning at (layerId, role). + //! + //! The run's buffers occupy consecutive entries of one request, which construction + //! checked, so naming the first is enough: the rest follow by BufferId. + struct SparseIndexRun + { + //! The run's first buffer, for row zero. + CUdeviceptr table = 0; + //! What a converted index is added to. Shared by the whole run. + MemAddress base = 0; + //! Elements between consecutive buffers, and between consecutive rows. + int bufferStride = 0; + int rowStride = 0; + }; + + [[nodiscard]] SparseIndexRun sparseIndexRun(LayerId layerId, DataRole const& role) const; + + //! Which table serves this buffer, and its entry in that table's buffer dimension. + [[nodiscard]] std::pair indexTableSlot(LayerId layerId, DataRole const& role) const; + + //! What a converted index for this buffer is added to. + [[nodiscard]] MemAddress indexTableBase(LayerId layerId, DataRole const& role) const; + + //! Write back every row's crossed units for one life cycle, from the ranges the Batch + //! staged. One scan over (layer, row), then one copy per pool covering the whole batch. + //! + //! `stream` carries the same obligations as fetchSparsePages' when it is not the + //! Batch's own; unlike there it has no default, so a caller naming the Batch's stream + //! is opting into the simple case rather than out of it. + void writeBackSparseUnits(LifeCycleId lifeCycle, CUstream stream); + + //! One layer group's page table on the device, which the index conversion reads. + //! + //! In the device staging rather than the mirror, because that is where the conversion + //! reads it and nothing captured reads the raw form at all. So it holds the staged step, + //! not the published one, and it is live: the next publish() refills it. + //! + //! Read it between waitPublished() and the next publish(). Before the first, the upload + //! that fills it may not have landed; after the second is enqueued, the upload stream is + //! free to overwrite it. The event that guards the staging against that overwrite covers + //! the Batch's own readers, which are enqueued before it is recorded, and cannot cover a + //! caller's. + [[nodiscard]] GpuCoalescedPageTable gpuCoalescedPageTable(LayerGroupId layerGroupId) const; + + //! Per-row sparse cache state, as the fetch kernel sees it. + //! + //! A device int32 array indexed by row, holding the arena slot each row's member took. + //! A free row, or a member whose cache is not reserved, reads as kNoSparseLease, which + //! the kernel skips. + //! + //! One array rather than one per layer group: a member is seated in every sparse life + //! cycle or none, and the arenas are sized alike, so a slot addresses the request in + //! all of them. An index rather than the pointers it implies, because a lease's cells + //! and its metadata are both arithmetic from the bases a launch carries once in + //! SparseArenaView. The slot changes whenever a member resumes or suspends, which is + //! why it is restaged on every publish alongside the watermarks. + [[nodiscard]] CUdeviceptr sparseLeases() const; + + //! Scratch for one layer group's fetch. Null pointers when the group carries no sparse + //! cache, which is what makes "nothing to fetch here" cheap to detect. + [[nodiscard]] SparseStepScratch const& sparseStepScratch(LayerGroupId layerGroupId) const; + + //! The unit ordinals bounding what each row's history watermark crossed in the last + //! flushSparseUnits(), for one life cycle: row `r` crossed [from[r], to[r]). + //! + //! Two parallel arrays of maxBatchSize rather than an array of pairs: the scan reads + //! one row per thread and takes the bounds as separate pointers, so each stays + //! coalesced. A struct per row would stride both reads by two. + //! + //! Ordinals are in this life cycle's units -- tokensPerBlock / unitsPerPage -- so one + //! token count quantises differently per life cycle. That is why these are per life + //! cycle where flushedHistory(), which is in tokens, is not. + //! + //! Device memory, written by flushSparseUnits() and read by the scan it launches. + //! Stale unless that flush found work: a step where no row crossed a boundary skips + //! the upload, which is safe only because nothing reads these then. + [[nodiscard]] int32_t const* flushFromUnits(LifeCycleId lifeCycle) const; + + //! @see flushFromUnits + [[nodiscard]] int32_t const* flushToUnits(LifeCycleId lifeCycle) const; + + friend class KvCache; + + //! Called by a member being destroyed. Unreachable while seated -- the Batch holds a + //! reference, so the destructor cannot run -- and kept as the assertion of that. + void _detach(KvCache& kvCache) noexcept; + + //! Whether the last publish() is still reading the staging rows. + //! + //! For a seated member to check before writing its page indices, which live in those rows. + [[nodiscard]] bool _publishInFlight() const noexcept { - return mMaxRows; + return !mPublished.isClosed(); } - int maxBlocks() const noexcept + //! Host-side page-index row for one (layer group, row), inside the pinned staging block. + [[nodiscard]] int32_t* _stagingRow(LifeCycleId lifeCycleId, BatchRow row) const; + //! A (life cycle, row)'s slice of the scratch slot staging. Only valid once the + //! staging exists, i.e. when mMaxScratchSlots is positive. + [[nodiscard]] int32_t* _scratchSlotRow(LifeCycleId lifeCycleId, BatchRow row) const; + + //! Block until every transfer this Batch issued has retired, so its buffers can go. + void _drain() noexcept; + + [[nodiscard]] size_t _rowElements() const noexcept; + + //! The buffers the staged block carries, in layout order. Each is placed for its own element, so + //! the one holding leases follows that index rather than the int32_t beside it. + enum class BufName : int + { + kPageTables = 0, + kHistory, + kCapacity, + kLeases, + kScratch, + kCount + }; + + static constexpr std::size_t kBufCount = static_cast(BufName::kCount); + + //! The least a buffer is aligned to; one whose element asks for more gets more. The + //! floor is here so a kernel may read any of them with the widest load the hardware + //! offers, without the layout being revisited to allow it. + static constexpr std::size_t kBufMinAlign = 16; + + //! Place and measure every buffer once, each from its own element and alignment. + void _layoutBufs(size_t scratchElements) noexcept; + + //! A buffer, as what it holds and how much of it. The extent travels with the pointer, + //! so nothing reads or fills past the buffer's own bytes; the padding that aligns the + //! next one is outside it. + template + [[nodiscard]] std::span _buf(BufName name) const noexcept { - return mMaxBlocks; + auto* const base = reinterpret_cast(reinterpret_cast(mStaging.address()) + mBuf[name].beg); + return std::span{base, static_cast(mBuf[name].length()) / sizeof(T)}; } - int maxBeamWidth() const noexcept + //! Where a buffer sits in the device staging, which carries the whole block as uploaded. + //! The conversion reads the page tables and the scratch from there, not from the mirror. + [[nodiscard]] CUdeviceptr _deviceBufAt(BufName name) const noexcept; + + //! Seed a row's staged watermarks, so a row reads as its occupant rather than as + //! whatever the last one left. + void _seedRowWatermarks(BatchRow row, int historyLength, int capacity); + + //! Refresh the staged watermarks and cache views from the members. KvCache reaches the + //! staging only through setBasePageIndexBuf, for its page indices. + void _stageRowState(); + + //! A buffer's layout: it yields how many int32 elements it needs and resumes with + //! the memory, so one allocation can serve several of them. + using LayoutGenerator = Generator; + + //! Yields the elements one life cycle's sparse scratch needs, then carves the slice + //! it is handed. Knows nothing of the counters, which are the driver's. + LayoutGenerator _layoutStepScratchForLifecycle(LifeCycleId lcId); + + //! Sums what the life cycles need, yields that plus the counter block, and hands each + //! life cycle its slice. Drives the leaves above while being driven itself. + LayoutGenerator _layoutStepScratch(); + + //! Entries one layer's write-back list must hold, from the step's token budget. + [[nodiscard]] int32_t _flushCapacity(LifeCycleId lifeCycle) const; + + //! Fill the per-buffer copy descriptors and upload them. Runs once, after the index + //! tables exist, because a descriptor names where in them a buffer's indices sit. + void _allocateCopyDesc(); + + //! Device room for the per-row unit ranges a write-back reads: + //! [life cycle][2][maxBatchSize], as [from][to]. Allocated with the fetch scratch, + //! which has the same lifetime. + //! Yields the elements the write-back ranges need, then clears the slice. + LayoutGenerator _layoutFlushRanges(); + + //! Points into mDeviceArena; null when no life cycle holds a sparse cache. + int32_t* mFlushRanges = nullptr; + //! Staged host-side, then handed to a copy that consumes it before returning. + std::vector mFlushRangeStaging; + //! Per row, the history its units have been written back to. + TypedVec mFlushedHistory; + //! Recorded after the write-back, so the next call's staging can wait for its reads. + CachedCudaEvent mFlushRead = CachedCudaEvent::makeNull(); + + //! The same array where a kernel reads it, which is what the launches are handed. + [[nodiscard]] SparseLeaseIndex const* _leaseMirror() const noexcept; + + //! Bytes of the mirror: the staged block from the histories through the leases, which is + //! everything a captured kernel reads. The page tables ahead of them are read from the + //! staging, by the conversion, so carrying them across would serve nothing. + [[nodiscard]] size_t _mirrorBytes() const noexcept; + + //! A staged element offset, as an address in the mirror. + [[nodiscard]] CUdeviceptr _mirrorAt(BufName name) const noexcept; + + //! Size the arenas from mMaxBatchSize, so every row the Batch can seat has a lease. + void _buildSparseArenas(); + + //! Everything a sparse copy needs that follows from configuration alone: the two pools, + //! the unit layout and the cell geometry. A caller fills in whose cache it is copying + //! into and which work list to read. + //! + //! One entry per maximal run of consecutive buffers sharing a hot pool. Buffers are + //! ordered (layer, pool), so a life cycle whose sparse buffers are alike yields a single + //! run, and the quantised shape -- data and block scales in pools of their own -- yields + //! one per layer and pool. A caller that wants only some buffers clips each run to them. + [[nodiscard]] std::vector _sparseCopyTemplates(LifeCycleId lifeCycle) const; + + //! The fetch's single work list, with the counter belonging to `selectionSlot`. + //! + //! One list serves every selection of a step because each is consumed, by the copies of + //! the layers resolving within it, before the next resolve overwrites it; the counters + //! are per selection because they are cleared once per publish rather than between + //! resolves, so a second resolve on a used counter would claim past the list. + [[nodiscard]] static SparseWorkList _sparseFetchWork(SparseStepScratch const& scratch, int32_t selectionSlot) { - return mMaxBeamWidth; + SparseWorkList work = scratch.fetchWork; + work.total += selectionSlot; + return work; } - int numLayerGroups() const noexcept + //! This life cycle's arena of GPU cells, from which a member takes its lease when it + //! joins. Empty when the life cycle has no sparse fetch path configured. + //! + //! Owned here rather than by StorageManager: the cells must stay GPU-resident while the + //! life cycle's hot pool group becomes host memory, and a cell is not a page, so it has + //! no answer to the question every storage slot must answer -- which block it holds. + //! The arena slot this row's member holds. Serves every sparse life cycle at once: the + //! arenas are sized alike, so a slot that exists in one exists in all. + [[nodiscard]] SparseLeaseIndex leaseOf(BatchRow row) const noexcept { - return mNumLayerGroups; + return mRowLease[row]; } - int deviceId() const noexcept + [[nodiscard]] SparseCellArena& _sparseArena(LifeCycleId lifeCycle) { - return mDeviceId; + return mSparseArenas[lifeCycle]; } - //! Internal device views: [row, beam, block] and [row, beam], respectively. - //! Addresses stay fixed for the batch lifetime; callers must obey the read contract. - MemAddress pageTableAddress(LayerGroupId group) const; - MemAddress numBlocksAddress(LayerGroupId group) const; + [[nodiscard]] SparseCellArena const& _sparseArena(LifeCycleId lifeCycle) const + { + return mSparseArenas[lifeCycle]; + } -private: - friend class KvCache; - void markDirty(int row) noexcept; - void checkOpen() const; - void checkPublished() const; - size_t rowOffset(int group, int row) const noexcept; - size_t countOffset(int group, int row) const noexcept; + //! Outlives the Batch. + KvCacheManager* const mManager; + CUstream mCudaStream = nullptr; + int mMaxBatchSize = 0; + int mMaxBlocksPerSeq = 0; + std::optional mMaxTotalInputTokens; + + //! A seated member cannot be destroyed under the Batch, so a caller may keep only the + //! Batch. KvCache points back with a raw `mBatch`, so there is no cycle. + TypedVec> mRows; + + //! The arena lease this row owns, for the Batch's lifetime. A permutation of + //! [0, mMaxBatchSize): one lease per row, so a free row always has a free lease and + //! nothing has to be allocated when a member joins. move() swaps two entries, which is + //! what lets a member change rows without its cells being copied. + TypedVec mRowLease; + + //! Indexed by life cycle; empty for one with no sparse fetch path. + TypedVec mSparseArenas; + + //! Cumulative fetch statistics, across every life cycle and step. Outside mStepScratch + //! because the step's clear must not touch it. + CudaUniqPtr mHitMiss; + + //! Device storage behind every life cycle's per-BufferId copy descriptors. + CudaUniqPtr mCopyDesc; + + //! Every life cycle's work-list counters. A step clears them once, in publish(); a + //! producer then raises its own and the copy that consumes the list reads it. + //! + //! Not owning: they lead `mStepScratch`, which is what makes one clear reach all of + //! them, so that allocation frees them. + int32_t* mCounters = nullptr; + size_t mCounterElements = 0; - struct Upload + //! Per life cycle, fixed at construction. Only scratch-slot substitution can make one + //! layer's indices differ from another's in the same life cycle, so a life cycle that + //! cannot hold scratch takes SHARED and the single smaller table that implies. + TypedVec mIndexModes; + + //! Whether the caller named the mode rather than letting it be derived. publish() + //! validates a named SHARED against each member, every step. + bool mIndexModeNamed = false; + + //! One converted index table: the specs it serves, and what a consumer needs to read + //! it. `concatOffsets` is in page-index units, parallel to a request's buffers. + struct IndexTable { - std::unique_ptr staging; - CachedCudaEvent completion = CachedCudaEvent::makeNull(); + std::vector served; + LifeCycleId lifeCycle{0}; + int scale = 1; + std::vector concatOffsets; + //! Element offset into mIndexTableMemory. + size_t elementOffset = 0; }; - std::shared_ptr mManager; - int mMaxRows; - int mMaxBlocks; - int mMaxBeamWidth; - int mNumLayerGroups = 0; - int mDeviceId = 0; - size_t mTableElements = 0; - size_t mTotalBytes = 0; - CudaUniqPtr mDeviceMemory; - std::vector mRows; - std::vector mDirty; - std::list mUploads; - CachedCudaEvent mReady = CachedCudaEvent::makeNull(); - std::vector mReaders; - bool mPublished = false; - bool mClosed = false; + std::vector mIndexTables; + + //! Where one buffer's indices are. + struct BufferSlot + { + IndexTableId table{0}; + int entry = 0; + }; + + std::map mBufferSlots; + + //! The slot naming this buffer. Throws naming `api` when no index table spec has it. + [[nodiscard]] BufferSlot const& _bufferSlot(LayerId layerId, DataRole const& role, char const* api) const; + + //! The spec that named this buffer. Exactly one does: specs may not overlap. + [[nodiscard]] IndexTableSpec const& _specOf(LayerId layerId, DataRole const& role, char const* api) const; + + //! What a converted index of `spec` is added to: its first buffer's pool base, under + //! the index mode this Batch chose. Every buffer of the spec shares it. + [[nodiscard]] MemAddress _specBase(IndexTableSpec const& spec) const; + + //! Validate the requests, merge those whose converted indices coincide, and record + //! where each buffer landed. + void _buildIndexTables(std::vector const& specs); + + //! Refuse a sparse life cycle whose buffers no single launch can address. + //! + //! A sparse copy walks one (layer, pool) run of buffers per launch, from one base + //! pointer, so that run has to sit at consecutive entries of one request. Splitting it + //! across requests would need two bases in one launch; declaring only part of it would + //! leave the rest with no converted indices at all. Checked here because a run that + //! cannot be addressed produces wrong page addresses at step time rather than an error. + void _checkSparseRunsAreAddressable() const; + + //! Sorts the tables, settles how many scratch slots a row may hold, and yields what + //! that costs; resumes with the uploaded block's scratch tail and builds the tables + //! and the device-side descriptors the conversion reads. + LayoutGenerator _layoutIndexTables(); + + //! Stage each member's scratch block range, for the life cycles that can hold scratch. + void _stageScratch(); + + //! Elements of the scratch staging. + [[nodiscard]] size_t _scratchElements() const noexcept; + + //! Host and device sides of the scratch metadata, laid out [begs][ends][slots]. + //! + //! Both point into the tail of the uploaded block rather than owning anything, which + //! is what lets one transfer carry them along with the page tables. Addressable only + //! when mMaxScratchSlots is positive; the tail is empty otherwise. + //! Byte extent of each buffer, measured once at construction. A lookup where the value + //! was arithmetic over members, which the compiler cannot fold. + EnumMap, kBufCount> mBuf{}; + + int32_t* mScratchStaging = nullptr; + int32_t* mScratchDevice = nullptr; + //! Slot ids a request may hold for one life cycle, which bounds the staging. + int mMaxScratchSlots = 0; + //! Tables whose life cycle can hold scratch sort last, so one launch covers each kind. + int mNumPlainTables = 0; + + //! Convert this step's raw slot ids into every table, on `stream`. + void _convertIndices(CUstream stream); + + //! Reset a vacated row's converted entries, so nothing it held outlives its occupant. + void _clearIndexRow(BatchRow row); + + //! All converted tables, carved by IndexTable::elementOffset. + CudaUniqPtr mIndexTableMemory; + //! One IndexTableDesc per table, and the offsets they point at. + CudaUniqPtr<> mIndexTableDescs; + int mMaxTableEntries = 0; + + //! Pinned so the members' index writes land here with no extra copy, and so the + //! publish is a single asynchronous transfer. Holds the page tables, the row state and + //! the views, then the scratch metadata. + StaticHostMem mStaging; + + //! Bytes of that block, and so of each step's transfer. + size_t mUploadBytes = 0; + + //! The row state and the leases, as a captured kernel reads them. Its address is baked + //! into a captured graph, so it is written only on the caller's stream, after the replay + //! that read the previous publish. + //! + //! The staged block's tail, not all of it: the page tables leading that block are read + //! from mDeviceStaging by the conversion, and what a captured kernel reads of them is + //! the converted table instead. + CudaUniqPtr<> mMirror; + + //! A device-side copy of the whole staged block, filled from the host ahead of time. + //! + //! Nothing captured reads it, so its fill may overlap the previous step. + CudaUniqPtr<> mDeviceStaging; + + //! Carries the host-to-device fill, off the caller's stream so it can run early. + CachedCudaStream mUploadStream; + + //! One allocation holding every layer group's fetch scratch, carved per group. + //! One device block split between the layout generators. The pointers they hand out + //! point into it, so it outlives every one of them. + CudaUniqPtr mDeviceArena; + TypedVec mScratch; + + //! Recorded after the host-to-device fill, so the host can tell when the staged rows + //! are free to write again. Null when no publish is in flight. + CachedCudaEvent mPublished = CachedCudaEvent::makeNull(); + + //! Recorded after the device staging has been read into the final tables. The next + //! fill must wait on it, since that fill runs on another stream and would otherwise + //! overwrite the staging while this step's copy is still reading it. + CachedCudaEvent mDeviceStagingConsumed = CachedCudaEvent::makeNull(); }; } // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/batchedPageCopy.cu b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/batchedPageCopy.cu index bf6732be540b..d253bd67d0b9 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/batchedPageCopy.cu +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/batchedPageCopy.cu @@ -987,20 +987,14 @@ void BatchedPageCopier::launchKernel(PoolCopyArgs const& args, CopyDirection dir TLLM_CUDA_CHECK(cudaGetLastError()); } -namespace detail +void copyBytesToDeviceSyncSource(CUdeviceptr dst, void const* src, size_t numBytes, CUstream stream) { - -void copyPageIndicesToDevice(CUdeviceptr dst, PageIndexPair const* src, size_t numPageIndices, CUstream stream) -{ - if (numPageIndices == 0) + if (numBytes == 0) { return; } - TLLM_CHECK_WITH_INFO( - numPageIndices <= std::numeric_limits::max() / sizeof(PageIndexPair), "Page-index array is too large"); - TLLM_CHECK_WITH_INFO(dst != 0 && src != nullptr, "Page-index copy requires valid source and destination"); + TLLM_CHECK_WITH_INFO(dst != 0 && src != nullptr, "Host-to-device staging requires a source and a destination"); - size_t numBytes = numPageIndices * sizeof(PageIndexPair); CUdeviceptr srcAddress = reinterpret_cast(src); CUmemcpyAttributes attributes{}; attributes.srcAccessOrder = CU_MEMCPY_SRC_ACCESS_ORDER_DURING_API_CALL; @@ -1013,6 +1007,4 @@ void copyPageIndicesToDevice(CUdeviceptr dst, PageIndexPair const* src, size_t n TLLM_CU_CHECK(cuMemcpyBatchAsync(&dst, &srcAddress, &numBytes, 1, &attributes, &firstCopy, 1, stream)); } -} // namespace detail - } // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/batchedPageCopy.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/batchedPageCopy.h index e86d1f7cda5b..a35f6c3ce7d1 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/batchedPageCopy.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/batchedPageCopy.h @@ -21,8 +21,11 @@ #include +#include "tensorrt_llm/common/assert.h" + #include #include +#include namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 { @@ -173,17 +176,21 @@ class BatchedPageCopier KernelConfig mOnboard{}; }; -namespace detail -{ - -//! Copies ephemeral host page indices into device memory while leaving the destination update -//! asynchronous. +//! Stages a host array into device memory, reading the source during the call. //! -//! Companion to BatchedPageCopier: it exists only because pageIndexLocation() may require the -//! index array to live in device memory, so callers on that path must stage it there first. This -//! moves the array that *describes* the copy, not the pages themselves. -void copyPageIndicesToDevice(CUdeviceptr dst, PageIndexPair const* src, size_t numPageIndices, CUstream stream); - -} // namespace detail +//! The source is read before this returns, so it need not be pinned, need not outlive the +//! call, and is left unchanged. The destination update stays asynchronous. That asymmetry is +//! why callers issue this on a helper stream: on a stream busy replaying the previous step, +//! the call could not complete until the stream reached it. +void copyBytesToDeviceSyncSource(CUdeviceptr dst, void const* src, size_t numBytes, CUstream stream); + +//! `count` values of `T`, on the terms above. +template +void copyToDeviceSyncSource(CUdeviceptr dst, T const* src, size_t count, CUstream stream) +{ + TLLM_CHECK_WITH_INFO( + count <= std::numeric_limits::max() / sizeof(T), "Host-to-device staging array is too large"); + copyBytesToDeviceSyncSource(dst, src, count * sizeof(T), stream); +} } // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/coldPageCodec.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/coldPageCodec.h index 7e67ecc1e679..372ef9e68784 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/coldPageCodec.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/coldPageCodec.h @@ -35,6 +35,11 @@ struct PoolDesc PoolIndex poolIndex{0}; MemAddress baseAddress = 0; size_t slotBytes = 0; + + //! What this pool's pages actually are, which is its group's tier except for a sparse + //! pool of the GPU tier: those are host-backed, and the group still reports GPU_MEM + //! because its other pools are not. + CacheTier storageMedium = CacheTier::GPU_MEM; }; //! Describes one hot pool group and the lifecycles that share its physical layout. diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/common.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/common.h index 4d867eb906a4..3471804fcf47 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/common.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/common.h @@ -78,7 +78,6 @@ enum class PageIndexMode : int using CacheLevel = StrongIndex; // The kernel-facing hot level; colder levels may also use GPU memory. inline constexpr CacheLevel kHotLevel{0}; -inline constexpr CacheLevel kSparseHistoryLevel{1}; // Sparse history requires HOST_MEM at this level. // Opaque request identifier shared with the rest of the batch manager. using RequestIdType = tensorrt_llm::batch_manager::RequestIdType; @@ -96,6 +95,12 @@ using LayerId = int; // Raw CUDA stream handle (CUstream cast to uintptr_t). using CudaStream = uintptr_t; +// Row of a Batch: the seat a member occupies while it is stepped. Its arena lease and its +// page table are both found by it, and it is the Batch's alone -- a row of one Batch names +// nothing in another. +using BatchRow = StrongIndex; +inline constexpr BatchRow kNoBatchRow{}; + // Index of a beam in beam-search. using BeamIndex = StrongIndex; inline constexpr BeamIndex kDefaultBeamIndex{0}; diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.cpp index 8a6e20701d73..49c618b07669 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.cpp @@ -18,14 +18,67 @@ #include "kv_cache_manager_v2/config.h" #include "kv_cache_manager_v2/exceptions.h" +#include #include +#include #include #include +#include #include namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 { +void SelectionUnitLayout::validate(size_t unitSize) const +{ + TLLM_CHECK_WITH_INFO( + numChunks != 0 && chunkBytes != 0, "SelectionUnitLayout: numChunks and chunkBytes must be non-zero"); + TLLM_CHECK_WITH_INFO(static_cast(numChunks) * chunkBytes == unitSize, + "SelectionUnitLayout: numChunks * chunkBytes must equal one unit's bytes"); + TLLM_CHECK_WITH_INFO(numChunks <= 1 || chunkStride >= chunkBytes, + "SelectionUnitLayout: chunkStride must not overlap consecutive chunks"); + TLLM_CHECK_WITH_INFO( + unitStride >= chunkBytes, "SelectionUnitLayout: unitStride must not overlap consecutive units"); +} + +void BufferConfig::validate() const +{ + TLLM_CHECK_WITH_INFO( + isSparse || sparseUnitChunks == 1, "BufferConfig: sparseUnitChunks set on a non-sparse buffer"); + // Both would name a unit count, and nothing reconciles them. + TLLM_CHECK_WITH_INFO(!isSparse || !tokensPerBlockOverride.has_value(), + "BufferConfig: tokensPerBlockOverride is not supported on a sparse buffer; use the layer's " + "sparse_cache.unit_tokens"); + TLLM_CHECK_WITH_INFO(sparseUnitChunks > 0, "BufferConfig: sparseUnitChunks must be positive"); +} + +void SparseCacheConfig::validate(int tokensPerBlock) const +{ + TLLM_CHECK_WITH_INFO(numSelectedBlocks > 0 && maxCachedUnits > 0, + "SparseCacheConfig: numSelectedBlocks and maxCachedUnits must be positive; omit sparse_cache " + "entirely for a layer with no fetch path"); + // Equality leaves no cell for the unit being written, so the request could not take a + // single token. The cache holds a full selection plus everything in flight, and those + // populations are disjoint: selection is below the watermark, in flight is above it. + TLLM_CHECK_WITH_INFO(maxCachedUnits > numSelectedBlocks, + "SparseCacheConfig: maxCachedUnits must exceed numSelectedBlocks, to leave room for units in " + "flight"); + TLLM_CHECK_WITH_INFO(unitTokens > 0 && tokensPerBlock % unitTokens == 0, + "SparseCacheConfig: unitTokens must be a divisor of tokensPerBlock"); +} + +void AttentionLayerConfig::validate() const +{ + detail::validateNoDuplicateBufferRoles(buffers); + for (auto const& buf : buffers) + { + buf.validate(); + } + TLLM_CHECK_WITH_INFO(!sparseCache.has_value() + || std::any_of(buffers.begin(), buffers.end(), [](BufferConfig const& b) { return b.isSparse; }), + "AttentionLayerConfig: sparse_cache is set on a layer with no sparse buffer to serve"); +} + void DiskCacheTierConfig::validate() const { if (quota == 0) @@ -38,6 +91,82 @@ void DiskCacheTierConfig::validate() const } } +namespace +{ + +//! What a layer's selection goes by, or nullopt where the layer has no sparse cache. +std::optional> selectionOf(LayerConfig const& layer) +{ + return std::visit( + [](auto const& cfg) -> std::optional> + { + if constexpr (std::is_same_v, AttentionLayerConfig>) + { + if (cfg.sparseCache.has_value()) + { + return std::pair{cfg.sparseCache->effectiveSelectionId(cfg.layerId), &*cfg.sparseCache}; + } + } + return std::nullopt; + }, + layer); +} + +//! An unset id reads as the layer's own, so the two namespaces are one and a layer that +//! named nothing can be named by someone else. Reject that, and require the layers of one +//! selection to agree on what they are selecting. +void _validateSelectionIds(std::vector const& layers) +{ + std::map byId; + std::set silent; + + for (auto const& layer : layers) + { + auto const selection = selectionOf(layer); + if (!selection.has_value()) + { + continue; + } + auto const [id, cache] = *selection; + if (!cache->selectionId.has_value()) + { + silent.insert(id); + } + + auto const [it, fresh] = byId.try_emplace(id, cache); + if (fresh) + { + continue; + } + auto const& first = *it->second; + if (first.numSelectedBlocks != cache->numSelectedBlocks || first.maxCachedUnits != cache->maxCachedUnits + || first.unitTokens != cache->unitTokens) + { + throw AssertionError( + "KVCacheManagerConfig: layers sharing a selection_id must agree on " + "num_selected_blocks, max_cached_units and unit_tokens"); + } + } + + for (auto const& layer : layers) + { + auto const selection = selectionOf(layer); + if (!selection.has_value() || !selection->second->selectionId.has_value()) + { + continue; + } + // Naming a layer that named nothing would share its selection without its consent. + if (silent.count(*selection->second->selectionId) != 0) + { + throw AssertionError( + "KVCacheManagerConfig: selection_id names a layer that set none of its own; " + "a layer that does not name a selection never shares one"); + } + } +} + +} // namespace + void KVCacheManagerConfig::validate() const { if (swaScratchReuse.has_value()) @@ -67,6 +196,13 @@ void KVCacheManagerConfig::validate() const std::visit([](auto const& cfg) { cfg.validate(); }, tier); } + // Zero asks for host-backed pools with no room for a single slot. It is not a way to turn + // them off: whether a pool is host-backed is a property of the buffer, not of this quota. + if (hostSparsePoolQuota.has_value() && *hostSparsePoolQuota == 0) + { + throw AssertionError("KVCacheManagerConfig: host_sparse_pool_quota must be positive when set"); + } + // Check for duplicate layer ids. std::set seenLayerIds; for (auto const& layer : layers) @@ -81,10 +217,6 @@ void KVCacheManagerConfig::validate() const } for (auto const& buf : cfg.buffers) { - if (buf.isSparse && (cacheTiers.size() < 2 || cacheTierOf(cacheTiers[1]) != CacheTier::HOST_MEM)) - { - throw std::invalid_argument("Sparse buffers require cache level 1 to use HOST_MEM"); - } if (buf.tokensPerBlockOverride.has_value() && (*buf.tokensPerBlockOverride <= 0 || tokensPerBlock % *buf.tokensPerBlockOverride != 0)) { @@ -93,10 +225,31 @@ void KVCacheManagerConfig::validate() const "tokensPerBlock"); } } + if constexpr (std::is_same_v, AttentionLayerConfig>) + { + if (cfg.sparseCache.has_value()) + { + cfg.sparseCache->validate(tokensPerBlock); + int const unitsPerPage = cfg.sparseCache->unitsPerPage(tokensPerBlock); + for (auto const& buf : cfg.buffers) + { + // Reject an unusable layout now, not at the first fetch, when + // serving has begun. A zero-size buffer has no layout to be + // unusable, and is dropped before the storage layer sees it, + // which is how a dense one of no size is treated too. + if (buf.isSparse && buf.size != 0) + { + static_cast(buf.sparseLayout(unitsPerPage)); + } + } + } + } }, layer); } + _validateSelectionIds(layers); + // SSM-specific validation. bool hasSSM = false; for (auto const& layer : layers) diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.h index 19600e0fd52a..8415192f24bf 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.h @@ -134,10 +134,44 @@ inline std::optional cacheTierMaxQuota(CacheTierConfig const& cfg) // Buffer configuration (one KV buffer inside an attention layer). // --------------------------------------------------------------------------- +// --------------------------------------------------------------------------- +// SelectionUnitLayout — where one selection unit's bytes sit inside a buffer. +// +// Buffer-relative, carrying no page offset: it describes the internal shape of the buffer, +// which is the part the model integration knows. The manager supplies the page offset from +// BufferAttr. +// +// Describes both sides of the copy: the fetch reads hot pages in ordinary slot layout and +// writes cells with the same structure, so one set of strides serves both, only the unit +// index differing. Codecs do not apply -- a codec transforms pages leaving the hot tier. +// --------------------------------------------------------------------------- +struct SelectionUnitLayout +{ + uint32_t numChunks; //!< One per head for a head-major page; 1 otherwise. + uint32_t chunkBytes; //!< Contiguous bytes of one unit within one chunk. + uint32_t chunkStride; //!< Distance between consecutive chunks. + uint32_t unitStride; //!< Distance between consecutive selection units. + + //! Derive the layout of a densely packed page. + //! + //! @param pageBytes One buffer's bytes for a whole page, i.e. BufferConfig::size. + //! @param unitsPerPage Selection units the page holds. + //! @param numChunks Pieces a unit is gathered from; see BufferConfig::sparseUnitChunks. + static SelectionUnitLayout derive(size_t pageBytes, int unitsPerPage, int numChunks) + { + auto const chunkStride = static_cast(pageBytes / static_cast(numChunks)); + auto const chunkBytes = chunkStride / static_cast(unitsPerPage); + return SelectionUnitLayout{static_cast(numChunks), chunkBytes, chunkStride, chunkBytes}; + } + + //! @param unitSize Bytes of one selection unit, i.e. size / unitsPerPage. + void validate(size_t unitSize) const; +}; + struct BufferConfig { DataRole role; - size_t size = 0; // bytes per page (without expansion) + size_t size = 0; // bytes per page // If set, overrides tokens_per_block for this buffer. // Must be a divisor of KVCacheManagerConfig::tokensPerBlock. @@ -148,6 +182,86 @@ struct BufferConfig //! Buffers may coalesce only when their group and expanded size match. //! This does not change their lifecycle or shared slot allocation. int coalescingGroup = 0; + //! Pieces a selection unit is gathered from within its page. Sparse buffers only. + //! + //! - **1** for a sequence-major page (`[token][head][dim]`), where a unit is one + //! contiguous range; also 1 for a single-head page, where the two orders coincide. + //! - **numHeads** for a head-major page (`[head][token][dim]`), where each head + //! contributes its own contiguous run. + //! + //! The rest of the layout follows from this, the layer's unit size and the page being + //! densely packed, so a padded page is not expressible. + int sparseUnitChunks = 1; + + void validate() const; + + //! Where this buffer's selection units sit inside a page, given how many the layer + //! divides a page into. Throws when the page cannot be divided as declared. + [[nodiscard]] SelectionUnitLayout sparseLayout(int unitsPerPage) const + { + auto const pieces = static_cast(sparseUnitChunks) * static_cast(unitsPerPage); + TLLM_CHECK_WITH_INFO(size != 0 && pieces != 0 && size % pieces == 0, + "BufferConfig: size must divide evenly into sparseUnitChunks * units-per-page pieces"); + return SelectionUnitLayout::derive(size, unitsPerPage, sparseUnitChunks); + } +}; + +// --------------------------------------------------------------------------- +// SparseCacheConfig — one layer's sparse fetch path. +// +// Every sparse buffer of a layer shares one resolve, one resident set and one recency +// order. Each Batch row's lease holds [selection][entry] metadata: layers sharing a +// selectionId share this metadata while keeping separate data cells; otherwise each +// layer has its own selection. +// +// Its presence is what configures the fetch path. A sparse buffer in a layer without it +// still gets sparse storage and locking; it simply has no GPU cache to read through. +// --------------------------------------------------------------------------- +struct SparseCacheConfig +{ + //! Selection width: units a sparse-attention step selects for this layer. + int numSelectedBlocks; + + //! Cells one buffer gets in a request's GPU cache, counting both populations: the + //! history it caches and every unit still in flight. + //! + //! Two to four times numSelectedBlocks is the usual choice: a cache only as wide as the + //! selection retains nothing between steps, so what it holds above the selection is + //! what survives to be hit next step. Only one unit of it is ever in flight, because a + //! sparse request contributes one query token per step, so the rest is all reuse. + int maxCachedUnits; + + //! Original tokens one selection unit spans. A whole page is `tokensPerBlock`. + //! + //! Must divide KVCacheManagerConfig::tokensPerBlock; the quotient is how many units a + //! page holds. It sets the sparse cache's granularity only: pages, page indices and + //! index spaces all keep the manager's tokensPerBlock. + int unitTokens; + + //! Layers whose selection is one selection, named so they can say so. + //! + //! A model that runs one indexer across several layers hands them all the same units + //! every step. Naming that lets those layers share one resident set and one resolve, + //! where otherwise each repeats the other's answer against its own copy of it. + //! + //! Unset means this layer selects for itself, and is read as the layer's own id -- so a + //! layer that says nothing never shares, and no explicit id may name a layer that said + //! nothing. Layers that do share must agree on the rest of this config. + std::optional selectionId; + + //! The id this layer's selection goes by, which is its own when it named none. + [[nodiscard]] int effectiveSelectionId(LayerId layerId) const noexcept + { + return selectionId.value_or(layerId); + } + + //! Selection units one page holds, given the manager's tokensPerBlock. + [[nodiscard]] int unitsPerPage(int tokensPerBlock) const + { + return tokensPerBlock / unitTokens; + } + + void validate(int tokensPerBlock) const; }; // --------------------------------------------------------------------------- @@ -206,23 +320,15 @@ struct AttentionLayerConfig // Private working pages never constrain prefix matching or enter the tree. AttentionReusePolicy reusePolicy = AttentionReusePolicy::REQUIRED; + //! The sparse fetch path for this layer's sparse buffers, if it has one. + std::optional sparseCache; [[nodiscard]] std::optional windowSize() const noexcept { return slidingWindowSize; } - void validate() const - { - detail::validateNoDuplicateBufferRoles(buffers); - for (auto const& buf : buffers) - { - if (buf.isSparse != buffers.front().isSparse) - { - throw std::invalid_argument("Sparse and non-sparse buffers cannot share an attention layer lifecycle"); - } - } - } + void validate() const; }; // --------------------------------------------------------------------------- @@ -249,6 +355,7 @@ struct SsmLayerConfig { throw std::invalid_argument("Sparse buffers are only supported for attention layers"); } + buf.validate(); } } }; @@ -450,6 +557,20 @@ struct KVCacheManagerConfig // (a text-only deployment forbids a request claiming otherwise). Default false. bool textOnly = false; + // Host memory the GPU tier's host-backed pools may take, in bytes. Unset, the storage + // manager sets it at construction to 80% of the host memory the host tiers' quotas leave, + // and keeps that value. + // + // A sparse buffer's pages are host memory, so the device quota does not size them: a slot + // spans every pool of its group, so their row count follows from the device-backed pools. + // This only ever lowers a slot count, and pool groups share it by the same ratio that + // divides the device quota. + // + // Honored exactly, as is the feasibility floor `constraints` implies; the pool ratio is + // what gives way between them. They can only conflict if the floors alone exceed this + // budget, which construction rejects. + std::optional hostSparsePoolQuota; + // Publish a finalized partial block for reuse when committing stops. // Beam search disables this while retaining full-block reuse. bool enablePartialCommit = true; diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/exceptions.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/exceptions.h index 9f76e36576a3..3578ccb6fb42 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/exceptions.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/exceptions.h @@ -61,7 +61,10 @@ void poisonOnExcept(char const* context, F&& func) noexcept } // --------------------------------------------------------------------------- -// Exception hierarchy (mirrors _exceptions.py) +// Exception hierarchy. The bindings make a Python class of each one +// (nb::exception in nanobind/batch_manager/kvCacheManagerV2.cpp) and +// runtime/kv_cache_manager_v2/__init__.py re-exports them, so the class a +// Python caller catches is this one. // --------------------------------------------------------------------------- class OutOfMemoryError : public std::runtime_error diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/indexConversionKernels.cu b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/indexConversionKernels.cu new file mode 100644 index 000000000000..f9e1fc59c2d3 --- /dev/null +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/indexConversionKernels.cu @@ -0,0 +1,260 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "indexConversionKernels.h" + +#include "common.h" +#include "kv_cache_manager_v2/utils/math.h" + +#include "tensorrt_llm/common/cudaUtils.h" + +#include +#include +#include +#include + +namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 +{ + +namespace +{ + +constexpr int kThreads = 256; + +//! Blocks each thread converts, as a shift, chosen so the grid lands near one wave. +//! +//! Only the block extent scales with the sequence; the others are the batch's rows and its +//! buffers, and a prefill step leaves both tiny -- one row of a million tokens is two CTAs +//! if a CTA takes a whole row. Splitting a row across CTAs fixes that, and giving each +//! thread more blocks folds the split back once the rows and buffers already fill the +//! machine, so a decode batch keeps the single-chunk shape it has now. +int32_t lg2BlocksPerThread(int32_t numTables, int32_t maxBuffers, int32_t numRows, int32_t maxBlocksPerRow, + int32_t smCount, int32_t ctasPerSm) noexcept +{ + int64_t const perChunk = int64_t{numTables} * maxBuffers * numRows; + int64_t const wave = int64_t{std::max(smCount, 1)} * std::max(ctasPerSm, 1); + int64_t const chunks = std::max(1, wave / std::max(perChunk, 1)); + int64_t const blocks = divUp(int64_t{maxBlocksPerRow}, chunks); + // Smallest shift with kThreads << lg2 >= blocks, which is ceil(log2()) of the blocks a + // thread would take. + int64_t const perThread = divUp(blocks, int64_t{kThreads}); + return perThread <= 1 ? 0 : std::bit_width(static_cast(perThread - 1)); +} + +//! Resident CTAs per SM for `kernel`, as the compiler built it. +//! +//! Asked rather than assumed: the answer folds in the register pressure a constant cannot +//! know, and it is what says when splitting a row stops buying occupancy. Each launcher +//! caches its own, since the two kernels need not agree. +int32_t queryCtasPerSm(void const* kernel) noexcept +{ + int blocks = 0; + if (cudaOccupancyMaxActiveBlocksPerMultiprocessor(&blocks, kernel, kThreads, 0) != cudaSuccess) + { + return 1; + } + return std::max(blocks, 1); +} + +//! CTAs the block extent needs once each thread takes `1 << lg2` of them. +int32_t chunkCount(int32_t maxBlocksPerRow, int32_t lg2) noexcept +{ + return divUp(maxBlocksPerRow, kThreads << lg2); +} + +//! Shift behind a power-of-two tokens-per-block, rejecting anything the kernel cannot divide by. +int32_t lg2TokensPerBlock(int32_t tokensPerBlock) +{ + TLLM_CHECK_WITH_INFO(tokensPerBlock > 0 && std::popcount(static_cast(tokensPerBlock)) == 1, + "index conversion: tokensPerBlock must be a power of two, got %d", tokensPerBlock); + return std::countr_zero(static_cast(tokensPerBlock)); +} + +//! Blocks a row's tokens occupy, which is as far as a consumer may read it. +__device__ inline int32_t rowBlockCount(int32_t capacity, int32_t lg2Tpb) +{ + return static_cast((static_cast(capacity) + (1U << lg2Tpb) - 1U) >> lg2Tpb); +} + +//! One CTA per (table, row, buffer); threads walk the blocks of that row. +//! +//! Buffer is the x extent because CTAs are rasterised there first: the buffers of one row +//! read the same raw row, and write spans that abut, where rows-first would stream the whole +//! raw table between reuses and stride the writes by a buffer's worth of blocks. +__global__ __launch_bounds__(kThreads) void indexConversionKernel(IndexTableDesc const* descs, int32_t numRows, + int32_t maxBuffers, int32_t maxBlocksPerRow, int32_t rawRowStride, int32_t const* capacity, int32_t lg2Tpb, + int32_t lg2Blocks) +{ + assert(gridDim.z == static_cast(numRows)); + int32_t const table = static_cast(blockIdx.x) / maxBuffers; + int32_t const buffer = static_cast(blockIdx.x) % maxBuffers; + auto const& desc = descs[table]; + if (buffer >= desc.numBuffers) + { + return; + } + int32_t const row = static_cast(blockIdx.z); + + int32_t const* raw = desc.raw + static_cast(row) * rawRowStride; + // Row-major in the entries: a row's entries are adjacent, so a request's K and V + // sit together the way an attention kernel's page list wants them. + int32_t* out = desc.converted + (static_cast(row) * desc.numBuffers + buffer) * maxBlocksPerRow; + int32_t const offset = desc.offsets[buffer]; + + int32_t const rowBlocks = rowBlockCount(capacity[row], lg2Tpb); + assert(rowBlocks <= maxBlocksPerRow); + int32_t const begin = static_cast(blockIdx.y) * (kThreads << lg2Blocks); + int32_t const end = min(begin + (kThreads << lg2Blocks), rowBlocks); + for (int32_t block = begin + static_cast(threadIdx.x); block < end; block += kThreads) + { + int32_t const slot = raw[block]; + // An absent block stays absent: folding it to a safe slot would read page zero. + out[block] = slot == kBadPageIndex.value() ? kBadPageIndex.value() : slot * desc.scale + offset; + } +} + +//! One CTA per (table, row, buffer), laid out as above, but scratch blocks take their slot +//! from the request rather than the page table. +__global__ __launch_bounds__(kThreads) void indexConversionScratchKernel(IndexTableDesc const* descs, int32_t numRows, + int32_t maxBuffers, int32_t maxBlocksPerRow, int32_t rawRowStride, int32_t const* capacity, int32_t lg2Tpb, + int32_t lg2Blocks) +{ + assert(gridDim.z == static_cast(numRows)); + int32_t const table = static_cast(blockIdx.x) / maxBuffers; + int32_t const buffer = static_cast(blockIdx.x) % maxBuffers; + auto const& desc = descs[table]; + if (buffer >= desc.numBuffers) + { + return; + } + int32_t const row = static_cast(blockIdx.z); + + int32_t const* raw = desc.raw + static_cast(row) * rawRowStride; + int32_t* out = desc.converted + (static_cast(row) * desc.numBuffers + buffer) * maxBlocksPerRow; + int32_t const offset = desc.offsets[buffer]; + + int32_t const scratchBeg = desc.scratchBegs[row]; + int32_t const scratchEnd = desc.scratchEnds[row]; + int32_t const* scratchSlots = desc.scratchSlots + static_cast(row) * desc.scratchSlotStride; + + int32_t const rowBlocks = rowBlockCount(capacity[row], lg2Tpb); + assert(rowBlocks <= maxBlocksPerRow); + int32_t const begin = static_cast(blockIdx.y) * (kThreads << lg2Blocks); + int32_t const end = min(begin + (kThreads << lg2Blocks), rowBlocks); + for (int32_t block = begin + static_cast(threadIdx.x); block < end; block += kThreads) + { + if (block >= scratchBeg && block < scratchEnd) + { + // The slot cursor advances by the sub-pages each scratch block occupies. The + // layer's own offset wraps inside the slot: a scratch slot is shared, so the + // sub-page positions past it are already spoken for. + int32_t const total = (block - scratchBeg) * desc.scratchPagesPerBlock; + int32_t const slotId = scratchSlots[total / desc.scale]; + out[block] = slotId == kBadPageIndex.value() + ? kBadPageIndex.value() + : slotId * desc.scale + (total % desc.scale + offset) % desc.scale; + continue; + } + int32_t const slot = raw[block]; + out[block] = slot == kBadPageIndex.value() ? kBadPageIndex.value() : slot * desc.scale + offset; + } +} + +//! Reset one row across every table, laid out like the conversion above. +__global__ __launch_bounds__(kThreads) void clearIndexRowKernel( + IndexTableDesc const* descs, int32_t maxBuffers, int32_t row, int32_t maxBlocksPerRow, int32_t lg2Blocks) +{ + int32_t const table = static_cast(blockIdx.x) / maxBuffers; + int32_t const buffer = static_cast(blockIdx.x) % maxBuffers; + auto const& desc = descs[table]; + if (buffer >= desc.numBuffers) + { + return; + } + + int32_t* out = desc.converted + (static_cast(row) * desc.numBuffers + buffer) * maxBlocksPerRow; + int32_t const begin = static_cast(blockIdx.y) * (kThreads << lg2Blocks); + int32_t const end = min(begin + (kThreads << lg2Blocks), maxBlocksPerRow); + for (int32_t block = begin + static_cast(threadIdx.x); block < end; block += kThreads) + { + out[block] = kBadPageIndex.value(); + } +} + +} // namespace + +void launchClearIndexRow(IndexTableDesc const* descs, int32_t numTables, int32_t maxBuffers, int32_t row, + int32_t maxBlocksPerRow, int32_t smCount, CUstream stream) +{ + if (numTables <= 0 || maxBuffers <= 0 || maxBlocksPerRow <= 0) + { + return; + } + // One launch for every table: under PER_LAYER a memset apiece would cost more in driver + // calls than the row is worth in bytes. + static int32_t const ctasPerSm = queryCtasPerSm(reinterpret_cast(clearIndexRowKernel)); + int32_t const lg2Blocks = lg2BlocksPerThread(numTables, maxBuffers, 1, maxBlocksPerRow, smCount, ctasPerSm); + dim3 const grid(static_cast(numTables * maxBuffers), + static_cast(chunkCount(maxBlocksPerRow, lg2Blocks)), 1U); + clearIndexRowKernel<<>>(descs, maxBuffers, row, maxBlocksPerRow, lg2Blocks); + TLLM_CUDA_CHECK(cudaGetLastError()); +} + +void launchIndexConversion(IndexTableDesc const* descs, int32_t numTables, int32_t maxBuffers, int32_t numRows, + int32_t maxBlocksPerRow, int32_t rawRowStride, int32_t const* capacity, int32_t tokensPerBlock, int32_t smCount, + CUstream stream) +{ + if (numTables <= 0 || maxBuffers <= 0 || numRows <= 0 || maxBlocksPerRow <= 0) + { + return; + } + TLLM_CHECK_WITH_INFO(capacity != nullptr, "launchIndexConversion: capacity watermarks are required"); + int32_t const lg2Tpb = lg2TokensPerBlock(tokensPerBlock); + // Table and buffer share the x extent so that a table's buffers are adjacent CTAs: they + // read the same raw row, and the reuse is then two CTAs apart whatever the sequence + // length. The buffer extent is the widest table's, since descs lives on the device and + // the host cannot read it back to size the grid per table; a narrower table's surplus + // CTAs return at once. Rows are slowest, sharing only the descriptor. + static int32_t const ctasPerSm = queryCtasPerSm(reinterpret_cast(indexConversionKernel)); + int32_t const lg2Blocks = lg2BlocksPerThread(numTables, maxBuffers, numRows, maxBlocksPerRow, smCount, ctasPerSm); + dim3 const grid(static_cast(numTables * maxBuffers), + static_cast(chunkCount(maxBlocksPerRow, lg2Blocks)), static_cast(numRows)); + indexConversionKernel<<>>( + descs, numRows, maxBuffers, maxBlocksPerRow, rawRowStride, capacity, lg2Tpb, lg2Blocks); + TLLM_CUDA_CHECK(cudaGetLastError()); +} + +void launchIndexConversionWithScratch(IndexTableDesc const* descs, int32_t numTables, int32_t maxBuffers, + int32_t numRows, int32_t maxBlocksPerRow, int32_t rawRowStride, int32_t const* capacity, int32_t tokensPerBlock, + int32_t smCount, CUstream stream) +{ + if (numTables <= 0 || maxBuffers <= 0 || numRows <= 0 || maxBlocksPerRow <= 0) + { + return; + } + TLLM_CHECK_WITH_INFO(capacity != nullptr, "launchIndexConversionWithScratch: capacity watermarks are required"); + int32_t const lg2Tpb = lg2TokensPerBlock(tokensPerBlock); + static int32_t const ctasPerSm = queryCtasPerSm(reinterpret_cast(indexConversionScratchKernel)); + int32_t const lg2Blocks = lg2BlocksPerThread(numTables, maxBuffers, numRows, maxBlocksPerRow, smCount, ctasPerSm); + dim3 const grid(static_cast(numTables * maxBuffers), + static_cast(chunkCount(maxBlocksPerRow, lg2Blocks)), static_cast(numRows)); + indexConversionScratchKernel<<>>( + descs, numRows, maxBuffers, maxBlocksPerRow, rawRowStride, capacity, lg2Tpb, lg2Blocks); + TLLM_CUDA_CHECK(cudaGetLastError()); +} + +} // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/indexConversionKernels.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/indexConversionKernels.h new file mode 100644 index 000000000000..3695fab5299e --- /dev/null +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/indexConversionKernels.h @@ -0,0 +1,104 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include + +#include + +namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 +{ + +//! One converted index table, as the conversion sees it. +//! +//! A table holds `numBuffers` values per (row, block): one per buffer of the specs it +//! serves, differing only by the offset baked into them. Merged specs share these values +//! and are told apart by their base pointers, which the conversion never sees. +struct IndexTableDesc +{ + //! Raw slot ids for this table's life cycle: [numRows][rawRowStride]. + int32_t const* raw; + //! Destination: [numRows][numBuffers][maxBlocksPerRow], so a row's buffers are adjacent. + int32_t* converted; + //! Offset added to each buffer's indices, in page-index units. Length numBuffers. + int32_t const* offsets; + int32_t numBuffers; + //! Slot id to page index, i.e. how many buffers the pool coalesces. + int32_t scale; + + //! Scratch metadata for this table's life cycle, null when it cannot hold scratch. + //! + //! A scratch block's slot does not come from the page table at all: it comes from the + //! request's own scratch slots, indexed by how far into the scratch range the block + //! sits. + int32_t const* scratchBegs = nullptr; //!< [numRows] + int32_t const* scratchEnds = nullptr; //!< [numRows] + int32_t const* scratchSlots = nullptr; //!< [numRows][scratchSlotStride] + int32_t scratchSlotStride = 0; + //! Sub-pages a scratch block occupies, which is what advances the slot cursor. + int32_t scratchPagesPerBlock = 1; +}; + +//! Reset one row's entries to kBadPageIndex in every table. +//! +//! The conversion stops at a row's capacity, so what an earlier occupant wrote above it +//! would otherwise outlive it. Clearing on the way out is what keeps the entries past a +//! capacity reading as absent. +void launchClearIndexRow(IndexTableDesc const* descs, int32_t numTables, int32_t maxBuffers, int32_t row, + int32_t maxBlocksPerRow, int32_t smCount, CUstream stream); + +//! Convert raw slot ids into kernel page indices for every table at once. +//! +//! One launch covers every table: under PER_LAYER there is a table per layer, and per-launch +//! driver cost would dominate the few arithmetic operations per entry. +//! +//! `entry = slot * scale + offset`, with kBadPageIndex passed through unchanged -- a sparse +//! table numbers pure-history and input-token entries in different levels' pools, so only +//! the sentinel separates an absent block from slot zero. +//! +//! A row is converted only as far as its capacity reaches; the entries past it read as +//! absent because a vacated row is cleared. +//! +//! @param descs device array of numTables descriptors +//! @param maxBuffers widest table's numBuffers; narrower tables return immediately +//! @param maxBlocksPerRow entries a row reserves, and the row stride of `converted`. No +//! capacity may exceed what it holds; the caller is required to have +//! checked that, since the kernels index on it unguarded. +//! @param rawRowStride elements between consecutive rows of a raw table +//! @param capacity per-row token watermark, numRows entries, on the launch stream +//! @param tokensPerBlock must be a power of two; the launcher rejects anything else so the +//! kernel divides by shifting +//! @param smCount multiprocessors to fill; a row's blocks are split across CTAs only +//! until the grid covers them, since a prefill batch has few rows and +//! few buffers and would otherwise leave the machine idle +void launchIndexConversion(IndexTableDesc const* descs, int32_t numTables, int32_t maxBuffers, int32_t numRows, + int32_t maxBlocksPerRow, int32_t rawRowStride, int32_t const* capacity, int32_t tokensPerBlock, int32_t smCount, + CUstream stream); + +//! The same, for tables whose life cycle substitutes scratch slots. +//! +//! A scratch block takes its slot from the request instead of the page table, and wraps the +//! layer's offset within the slot rather than adding it. +//! +//! Scratch is a PER_LAYER-only concern. Under SHARED every layer of a life cycle reads one +//! table, which substitution would break. +void launchIndexConversionWithScratch(IndexTableDesc const* descs, int32_t numTables, int32_t maxBuffers, + int32_t numRows, int32_t maxBlocksPerRow, int32_t rawRowStride, int32_t const* capacity, int32_t tokensPerBlock, + int32_t smCount, CUstream stream); + +} // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/introspection.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/introspection.cpp index 7067d40b9772..b4eade03c7e4 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/introspection.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/introspection.cpp @@ -17,11 +17,14 @@ #include "kv_cache_manager_v2/introspection.h" +#include "kv_cache_manager_v2/batch.h" #include "kv_cache_manager_v2/blockRadixTree.h" #include "kv_cache_manager_v2/kvCache.h" #include "kv_cache_manager_v2/kvCacheManager.h" #include "kv_cache_manager_v2/page.h" +#include "kv_cache_manager_v2/sparseCacheKernels.h" +#include "kv_cache_manager_v2/stagingBuffer.h" #include "kv_cache_manager_v2/storageManager.h" #include "kv_cache_manager_v2/utils/hostMemBacking.h" #include "kv_cache_manager_v2/utils/math.h" @@ -31,11 +34,13 @@ #include +#include #include #include #include #include #include +#include #include #include #include @@ -115,6 +120,488 @@ TypedVec KvCacheIntrospection::computeSlotsForBatch(K return manager.storage().computePoolGroupSlotsForBatch(batch, tokensPerBlock, swaScratchReuse); } +std::vector KvCacheIntrospection::sparseCacheResidentOrdinals(KvCache const& kvCache, int lcId) +{ + LifeCycleId const id{lcId}; + // The arena belongs to the Batch the lease was taken from, so an unseated cache has + // nothing resident to report. + if (id >= kvCache.mManager->numLifeCycles() || kvCache.mBatch == nullptr) + { + return {}; + } + return kvCache.mBatch->_sparseArena(id).peekResidentOrdinals(kvCache.mBatch->leaseOf(*kvCache.mBatchRow)); +} + +std::optional KvCacheIntrospection::sparseCacheBaseCell(KvCache const& kvCache, int lcId, int pool) +{ + LifeCycleId const id{lcId}; + if (id >= kvCache.mManager->numLifeCycles() || kvCache.mBatch == nullptr + || kvCache.mBatch->_sparseArena(id).empty()) + { + return std::nullopt; + } + // Relative to the pool's own arena region, which is what getSparseCacheBaseAddress + // reports for a buffer drawing on it. + auto const& arena = kvCache.mBatch->_sparseArena(id); + return kvCache.mBatch->leaseOf(*kvCache.mBatchRow).value() * arena.cellsPerLease(PoolIndex{pool}); +} + +std::vector KvCacheIntrospection::batchPageTable(Batch const& batch, int layerGroupId) +{ + // The Batch's own stream, not the null one: it is created non-blocking, so a null-stream + // read is not ordered against the conversion that fills what is read here. + CUstream const stream = batch.cudaStream(); + auto const view = batch.gpuCoalescedPageTable(LifeCycleId{layerGroupId}); + std::vector out(static_cast(view.numRows) * view.blocksPerRow); + cuCheck(cuMemcpyDtoHAsync(out.data(), view.data, sizeof(int32_t) * out.size(), stream)); + cuCheck(cuStreamSynchronize(stream)); + return out; +} + +std::pair KvCacheIntrospection::batchIndexTableSlot(Batch const& batch, int layerId, DataRole const& role) +{ + auto const [table, entry] = batch.indexTableSlot(LayerId{layerId}, role); + return {table.value(), entry}; +} + +MemAddress KvCacheIntrospection::batchIndexTableBase(Batch const& batch, int layerId, DataRole const& role) +{ + return batch.indexTableBase(LayerId{layerId}, role); +} + +std::vector KvCacheIntrospection::batchIndexTable(Batch const& batch, int tableId) +{ + // The Batch's own stream, not the null one: it is created non-blocking, so a null-stream + // read is not ordered against the conversion that fills what is read here. + CUstream const stream = batch.cudaStream(); + IndexTableId const typed{tableId}; + auto const numBuffers = static_cast(batch.indexTable(typed).front().buffers.size()); + auto const perEntry = static_cast(batch.maxBatchSize()) * static_cast(batch.maxBlocksPerSeq()); + std::vector out(numBuffers * perEntry); + cuCheck(cuMemcpyDtoHAsync(out.data(), batch.indexTableAddress(typed), sizeof(int32_t) * out.size(), stream)); + cuCheck(cuStreamSynchronize(stream)); + return out; +} + +GpuCoalescedPageTable KvCacheIntrospection::batchGpuCoalescedPageTable(Batch const& batch, int layerGroupId) +{ + return batch.gpuCoalescedPageTable(LifeCycleId{layerGroupId}); +} + +std::vector KvCacheIntrospection::batchSparseEntriesPerRow(Batch const& batch, int layerGroupId) +{ + // The Batch's own stream, not the null one: it is created non-blocking, so a null-stream + // read is not ordered against the conversion that fills what is read here. + CUstream const stream = batch.cudaStream(); + auto const rows = static_cast(batch.maxBatchSize()); + std::vector leases(rows); + cuCheck(cuMemcpyDtoHAsync(leases.data(), batch.sparseLeases(), sizeof(int32_t) * rows, stream)); + cuCheck(cuStreamSynchronize(stream)); + + // Every lease of a life cycle has the same shape, so the arena answers for all of them. + auto const entriesPerLayer = batch._sparseArena(LifeCycleId{layerGroupId}).view().entriesPerSelection; + std::vector out; + out.reserve(rows); + for (auto const lease : leases) + { + out.push_back(lease == kNoSparseLease.value() ? 0 : entriesPerLayer); + } + return out; +} + +std::pair, std::vector> KvCacheIntrospection::batchWatermarks(Batch const& batch) +{ + // The Batch's own stream, not the null one: it is created non-blocking, so a null-stream + // read is not ordered against the conversion that fills what is read here. + CUstream const stream = batch.cudaStream(); + auto const view = batch.watermarks(); + std::vector history(static_cast(view.numRows)); + std::vector capacity(static_cast(view.numRows)); + cuCheck(cuMemcpyDtoHAsync(history.data(), view.historyLength, sizeof(int32_t) * history.size(), stream)); + cuCheck(cuMemcpyDtoHAsync(capacity.data(), view.capacity, sizeof(int32_t) * capacity.size(), stream)); + cuCheck(cuStreamSynchronize(stream)); + return {std::move(history), std::move(capacity)}; +} + +std::vector KvCacheIntrospection::sparseFlushCounts(KvCache const& kvCache, int lcId) +{ + LifeCycleId const id{lcId}; + if (id >= kvCache.mManager->numLifeCycles() || kvCache.mBatch == nullptr) + { + return {}; + } + // The counters are written on the Batch's stream, which a caller may make non-blocking; + // a null-stream read would not be ordered against it. + CUstream const stream = kvCache.mBatch->cudaStream(); + // The lists are the step's, not the request's, so this reads the Batch's row for it. + auto const& scratch = kvCache.mBatch->sparseStepScratch(LifeCycleId{lcId}); + auto const layers = static_cast(scratch.numLayers); + std::vector out(layers); + cuCheck(cuMemcpyDtoHAsync(out.data(), + reinterpret_cast(scratch.flushCount + toSizeT(*kvCache.mBatchRow) * layers), + sizeof(int32_t) * layers, stream)); + cuCheck(cuStreamSynchronize(stream)); + return out; +} + +namespace +{ + +CudaUniqPtr<> deviceCopy(std::vector const& values) +{ + CUstream const stream = nullptr; + // Null for an empty array, which cuMemAlloc refuses and the kernels accept: a probe may + // declare a width of zero, and every array sized from it is then empty. + auto owned = allocateDeviceMemory(sizeof(int32_t) * values.size()); + if (values.empty()) + { + return owned; + } + cuCheck(cuMemcpyHtoDAsync( + reinterpret_cast(owned.get()), values.data(), sizeof(int32_t) * values.size(), stream)); + cuCheck(cuStreamSynchronize(stream)); + return owned; +} + +std::vector readBack(CudaUniqPtr<> const& buffer, size_t count) +{ + CUstream const stream = nullptr; + std::vector out(count); + if (count == 0) + { + return out; + } + cuCheck( + cuMemcpyDtoHAsync(out.data(), reinterpret_cast(buffer.get()), sizeof(int32_t) * count, stream)); + cuCheck(cuStreamSynchronize(stream)); + return out; +} + +} // namespace + +KvCacheIntrospection::SparseResolveProbe KvCacheIntrospection::sparseResolve( + std::vector const& residentOrdinal, std::vector const& lruOrder, int numLayers, + int entriesPerLayer, std::vector const& selected, int selectionWidth, + std::vector const& historyUnits, std::vector const& capacityUnits, int maxInFlightUnits, + std::vector const& inactiveRows) +{ + CUstream const stream = nullptr; + auto const numRows = static_cast(historyUnits.size()); + auto const perRowEntries = static_cast(numLayers) * entriesPerLayer; + auto const entryCount = perRowEntries * static_cast(numRows); + auto const perRowSelection = static_cast(numLayers) * selectionWidth; + auto const selectionCount = perRowSelection * static_cast(numRows); + TLLM_CHECK_WITH_INFO(maxInFlightUnits >= 0, "sparseResolve: maxInFlightUnits counts units and cannot be negative"); + TLLM_CHECK_WITH_INFO(numRows != 0 && capacityUnits.size() == historyUnits.size() + && residentOrdinal.size() == entryCount && lruOrder.size() == entryCount + && selected.size() == selectionCount, + "sparseResolve: input sizes must match the declared shape"); + + auto selection = deviceCopy(selected); + auto outSlot = deviceCopy(std::vector(selectionCount, 0)); + auto missRow = deviceCopy(std::vector(selectionCount, kNoSparseSlot)); + auto missOrdinal = deviceCopy(std::vector(selectionCount, kNoSparseSlot)); + auto missSlot = deviceCopy(std::vector(selectionCount, kNoSparseSlot)); + auto missTotal = deviceCopy(std::vector(static_cast(numLayers), 0)); + auto const countCount = static_cast(numLayers) * static_cast(numRows); + auto missCount = deviceCopy(std::vector(countCount, 0)); + auto history = deviceCopy(historyUnits); + auto capacity = deviceCopy(capacityUnits); + auto const inFlightCount = countCount * static_cast(maxInFlightUnits); + auto inFlight = deviceCopy(std::vector(inFlightCount, kNoSparseSlot)); + + // One metadata block per lease, its resident ordinals then its recency order, which is + // how the arena lays a lease out. The probe takes and returns each array whole, so they + // are interleaved on the way in and separated on the way out. + auto const leaseInts = 2 * perRowEntries; + std::vector metadataHost(leaseInts * static_cast(numRows)); + for (int row = 0; row < numRows; ++row) + { + auto const from = perRowEntries * static_cast(row); + auto const to = leaseInts * static_cast(row); + std::copy_n(residentOrdinal.begin() + from, perRowEntries, metadataHost.begin() + to); + std::copy_n(lruOrder.begin() + from, perRowEntries, metadataHost.begin() + to + perRowEntries); + } + auto metadata = deviceCopy(metadataHost); + + // Row r takes lease r, except an inactive row, which holds none for the kernel to skip. + std::vector leaseHost(static_cast(numRows)); + for (int row = 0; row < numRows; ++row) + { + leaseHost[static_cast(row)] + = std::find(inactiveRows.begin(), inactiveRows.end(), row) != inactiveRows.end() ? kNoSparseLease.value() + : row; + } + auto leases = deviceCopy(leaseHost); + + SparseArenaView arena; + arena.metadataBase = metadata.get(); + arena.leaseStride = sizeof(int32_t) * leaseInts; + arena.lruOffset = static_cast(perRowEntries); + arena.numLayers = numLayers; + // One metadata row per layer: this probe drives a model where every layer selects for + // itself, which is what a resolve sees whenever no selection_id is shared. + arena.numSelections = numLayers; + arena.entriesPerSelection = entriesPerLayer; + // This probe drives the metadata, not the addressing. + arena.bufferSlotInPool = nullptr; + arena.bufferLayer = nullptr; + arena.layerSelection = nullptr; + + // One launch per layer, as a step does: a layer's selection is not known until the + // previous layer has run. The arrays are therefore layer-major, so each launch sees a + // contiguous [row][...] slice of its own. Same stream, so the layers stay ordered. + for (int layer = 0; layer < numLayers; ++layer) + { + auto const selectionOffset = static_cast(layer) * numRows * selectionWidth; + auto const countOffset = static_cast(layer) * numRows; + + SparseResolveArgs args; + args.leases = reinterpret_cast(leases.get()); + args.arena = arena; + args.numRows = numRows; + args.selectionSlot = layer; + args.selected = reinterpret_cast(selection.get()) + selectionOffset; + args.selectionWidth = selectionWidth; + // The hook works in units, so one token is one unit. + args.historyTokens = reinterpret_cast(history.get()); + args.capacityTokens = reinterpret_cast(capacity.get()); + args.unitTokens = 1; + args.inFlightSlot + = reinterpret_cast(inFlight.get()) + countOffset * static_cast(maxInFlightUnits); + args.maxInFlightUnits = maxInFlightUnits; + args.outSlot = reinterpret_cast(outSlot.get()) + selectionOffset; + // One list per layer, each the width of the whole batch's selection: the worst case + // is every row missing everything it asked for. + args.work.row = reinterpret_cast(missRow.get()) + selectionOffset; + args.work.ordinal = reinterpret_cast(missOrdinal.get()) + selectionOffset; + args.work.slot = reinterpret_cast(missSlot.get()) + selectionOffset; + args.work.total = reinterpret_cast(missTotal.get()) + layer; + args.work.capacity = numRows * selectionWidth; + args.work.numLists = 1; + args.missCount = reinterpret_cast(missCount.get()) + countOffset; + + launchSparseResolve(args, stream); + } + cuCheck(cuStreamSynchronize(stream)); + + SparseResolveProbe result; + auto const metadataOut = readBack(metadata, leaseInts * static_cast(numRows)); + result.residentOrdinal.resize(entryCount); + result.lruOrder.resize(entryCount); + for (int row = 0; row < numRows; ++row) + { + auto const from = leaseInts * static_cast(row); + auto const to = perRowEntries * static_cast(row); + std::copy_n(metadataOut.begin() + from, perRowEntries, result.residentOrdinal.begin() + to); + std::copy_n(metadataOut.begin() + from + perRowEntries, perRowEntries, result.lruOrder.begin() + to); + } + result.outSlot = readBack(outSlot, selectionCount); + result.missRow = readBack(missRow, selectionCount); + result.missOrdinal = readBack(missOrdinal, selectionCount); + result.missSlot = readBack(missSlot, selectionCount); + result.missTotal = readBack(missTotal, static_cast(numLayers)); + result.missCount = readBack(missCount, countCount); + result.inFlightSlot = inFlightCount > 0 ? readBack(inFlight, inFlightCount) : std::vector{}; + return result; +} + +std::vector KvCacheIntrospection::sparseUnitCopy(std::vector const& hostPool, uint64_t hostPageStride, + std::vector const& pageTable, std::vector const& deviceArena, int baseCell, int numBlocks, + int cellsPerBuffer, uint64_t bufferBytes, uint32_t numChunks, uint32_t chunkBytes, uint32_t chunkStride, + uint32_t unitStride, int unitsPerPage, std::vector const& missOrdinal, + std::vector const& missSlot, std::vector const& missCount, int selectionWidth, + std::vector const& bufferLayer, bool writeBack, int gridCtas) +{ + // The row's lease is what says whether it holds a reservation, so a negative one asks + // the kernel to skip the very copy this probe exists to run. + TLLM_CHECK_WITH_INFO( + baseCell >= 0, "KvCacheIntrospection::sparseUnitCopy: baseCell names an arena cell and cannot be negative"); + CUstream const stream = nullptr; + auto source = allocateDeviceMemory(hostPool.size()); + auto const hostPtr = reinterpret_cast(source.get()); + cuCheck(cuMemcpyHtoDAsync(hostPtr, hostPool.data(), hostPool.size(), stream)); + + auto destination = allocateDeviceMemory(deviceArena.size()); + auto const devicePtr = reinterpret_cast(destination.get()); + cuCheck(cuMemcpyHtoDAsync(devicePtr, deviceArena.data(), deviceArena.size(), stream)); + + auto table = deviceCopy(pageTable); + // The caller describes work as [list][selectionWidth] with a count per list, which is + // the shape a test can write by hand. Compact it here, because that is what the kernel + // consumes: a dense run per list, with the row carried by the entry. + auto const numLists = bufferLayer.empty() ? 1 : bufferLayer.back() + 1; + auto const perList = static_cast(std::max(selectionWidth, 1)); + std::vector flatRow; + std::vector flatOrdinal; + std::vector flatSlot; + std::vector totals(static_cast(numLists), 0); + flatRow.reserve(missOrdinal.size()); + flatOrdinal.reserve(missOrdinal.size()); + flatSlot.reserve(missOrdinal.size()); + for (int list = 0; list < numLists; ++list) + { + auto const count = static_cast(missCount.at(static_cast(list))); + for (size_t j = 0; j < perList; ++j) + { + auto const index = static_cast(list) * perList + j; + // The hook drives one request, so every entry names row zero. + flatRow.push_back(0); + flatOrdinal.push_back(j < count ? missOrdinal.at(index) : kNoSparseSlot); + flatSlot.push_back(j < count ? missSlot.at(index) : kNoSparseSlot); + } + totals[static_cast(list)] = static_cast(count); + } + + auto workRowBuf = deviceCopy(flatRow); + auto missOrd = deviceCopy(flatOrdinal); + auto missSlotBuf = deviceCopy(flatSlot); + auto totalBuf = deviceCopy(totals); + auto bufferLayerBuf = deviceCopy(bufferLayer); + + // The kernel reads the lease and the per-buffer maps, so the hook assembles the arena a + // single request's cells would be addressed within. The copy reads no metadata block, so + // the arena names none. + std::vector slotInPool(bufferLayer.size()); + std::iota(slotInPool.begin(), slotInPool.end(), 0); + auto slotInPoolBuf = deviceCopy(slotInPool); + + // The probe drives one pool, so a lease of one cell and a buffer-slot identity reproduce + // the address the production path computes from a lease index and a pool slot. + auto leases = deviceCopy(std::vector(1, baseCell)); + SparseArenaView arena; + arena.bufferSlotInPool = reinterpret_cast(slotInPoolBuf.get()); + arena.bufferLayer = reinterpret_cast(bufferLayerBuf.get()); + // This probe drives the addressing, not the metadata. + arena.metadataBase = nullptr; + arena.leaseStride = 0; + arena.lruOffset = 0; + arena.numLayers = 0; + arena.numSelections = 0; + arena.entriesPerSelection = 0; + arena.layerSelection = nullptr; + + // One descriptor per buffer, as the real path builds them: the hook drives the + // addressing, so every buffer shares a pool and differs only by its slot in it. + std::vector descs(bufferLayer.size()); + for (size_t i = 0; i < descs.size(); ++i) + { + auto& desc = descs[i]; + desc.hostBase = hostPtr; + desc.hostPageStride = static_cast(hostPageStride); + desc.pageTable = reinterpret_cast(table.get()) + i * static_cast(numBlocks); + desc.pageTableStride = static_cast(numBlocks); + desc.numBlocks = static_cast(numBlocks); + desc.deviceBase = devicePtr + static_cast(i) * static_cast(cellsPerBuffer) * bufferBytes; + desc.leaseStride = static_cast(bufferBytes); + desc.cellBytes = static_cast(bufferBytes); + desc.layout.numChunks = numChunks; + desc.layout.chunkBytes = chunkBytes; + desc.layout.chunkStride = chunkStride; + desc.layout.unitStride = unitStride; + desc.width = sparseCopyWidth(desc.layout, desc.hostPageStride, bufferBytes); + } + auto descOwner = allocateDeviceMemory(sizeof(SparseCopyDesc) * descs.size()); + auto const descPtr = reinterpret_cast(descOwner.get()); + cuCheck(cuMemcpyHtoDAsync(descPtr, descs.data(), sizeof(SparseCopyDesc) * descs.size(), stream)); + + SparseUnitCopyArgs args; + args.leases = reinterpret_cast(leases.get()); + args.arena = arena; + args.numRows = 1; + args.copyDesc = reinterpret_cast(descPtr); + args.bufferBase = 0; + args.numBuffers = static_cast(descs.size()); + args.direction = writeBack ? SparseCopyDirection::kWriteBack : SparseCopyDirection::kFetch; + args.unitsPerPage = unitsPerPage; + args.work.row = reinterpret_cast(workRowBuf.get()); + args.work.ordinal = reinterpret_cast(missOrd.get()); + args.work.slot = reinterpret_cast(missSlotBuf.get()); + args.work.total = reinterpret_cast(totalBuf.get()); + args.work.capacity = static_cast(flatOrdinal.size()); + args.work.numLists = numLists; + args.gridCtas = gridCtas; + + launchSparseUnitCopy(args, stream); + cuCheck(cuStreamSynchronize(stream)); + + // Write-back moves into the cold pool, so that is the side worth returning. + CUdeviceptr const result = writeBack ? hostPtr : devicePtr; + size_t const resultBytes = writeBack ? hostPool.size() : deviceArena.size(); + std::vector out(resultBytes); + cuCheck(cuMemcpyDtoHAsync(out.data(), result, resultBytes, stream)); + cuCheck(cuStreamSynchronize(stream)); + return out; +} + +KvCacheIntrospection::SparseFlushList KvCacheIntrospection::sparseFlushList(std::vector const& residentOrdinal, + int numLayers, int entriesPerLayer, int fromUnits, int toUnits, int maxFlushUnits) +{ + CUstream const stream = nullptr; + auto const entryCount = static_cast(numLayers) * entriesPerLayer; + TLLM_CHECK_WITH_INFO( + residentOrdinal.size() == entryCount, "sparseFlushList: input size must match the declared shape"); + auto const listCount = static_cast(numLayers) * maxFlushUnits; + + auto from = deviceCopy(std::vector(1, fromUnits)); + auto to = deviceCopy(std::vector(1, toUnits)); + auto outRow = deviceCopy(std::vector(listCount, kNoSparseSlot)); + auto outOrdinal = deviceCopy(std::vector(listCount, kNoSparseSlot)); + auto outSlot = deviceCopy(std::vector(listCount, kNoSparseSlot)); + auto outTotal = deviceCopy(std::vector(static_cast(numLayers), 0)); + auto outCount = deviceCopy(std::vector(static_cast(numLayers), 0)); + + // The scan spans the batch, so the hook stands one row in for a single request, holding + // one lease laid out as the arena lays one: its resident ordinals, then its recency + // order, which the scan does not read. + std::vector metadataHost(2 * entryCount, 0); + std::copy(residentOrdinal.begin(), residentOrdinal.end(), metadataHost.begin()); + auto metadata = deviceCopy(metadataHost); + auto leases = deviceCopy(std::vector(1, 0)); + + // A layer per metadata row, in order: the probe drives unshared selections. + std::vector layerSelectionHost(static_cast(numLayers)); + std::iota(layerSelectionHost.begin(), layerSelectionHost.end(), 0); + auto layerSelection = deviceCopy(layerSelectionHost); + + SparseArenaView arena; + arena.metadataBase = metadata.get(); + arena.leaseStride = sizeof(int32_t) * 2 * entryCount; + arena.lruOffset = static_cast(entryCount); + arena.numLayers = numLayers; + arena.numSelections = numLayers; + arena.entriesPerSelection = entriesPerLayer; + // This probe drives the metadata, not the addressing. + arena.bufferSlotInPool = nullptr; + arena.bufferLayer = nullptr; + arena.layerSelection = reinterpret_cast(layerSelection.get()); + + SparseFlushListArgs args; + args.leases = reinterpret_cast(leases.get()); + args.arena = arena; + args.numRows = 1; + args.fromUnits = reinterpret_cast(from.get()); + args.toUnits = reinterpret_cast(to.get()); + args.work.row = reinterpret_cast(outRow.get()); + args.work.ordinal = reinterpret_cast(outOrdinal.get()); + args.work.slot = reinterpret_cast(outSlot.get()); + args.work.total = reinterpret_cast(outTotal.get()); + args.work.capacity = maxFlushUnits; + args.work.numLists = numLayers; + args.outCount = reinterpret_cast(outCount.get()); + + launchSparseFlushList(args, stream); + cuCheck(cuStreamSynchronize(stream)); + + SparseFlushList result; + result.row = readBack(outRow, listCount); + result.ordinal = readBack(outOrdinal, listCount); + result.slot = readBack(outSlot, listCount); + result.total = readBack(outTotal, static_cast(numLayers)); + result.count = readBack(outCount, static_cast(numLayers)); + return result; +} + bool KvCacheIntrospection::allTreePagesDroppable(KvCacheManager& manager) { for (auto const& [_, root] : manager.radixTree().roots()) @@ -205,13 +692,13 @@ std::optional KvCacheIntrospection::probeHostMemPlacement(size HostMemBackingOptions options; options.allowRemoteNumaFallback = allowRemoteNumaFallback; - auto backing = std::make_unique(options); + size_t const span = roundUp(size, hostMemSizeGranularity()); + auto backing = std::make_unique(span, options); HostMemPlacement placement{}; placement.gpuNumaNode = runtime::TopologyDetector::getInstance().getCurrentGpuNumaId(); - size_t const span = roundUp(size, backing->commitGranularity()); - MemAddress const base = backing->reserve(span); + MemAddress const base = backing->address(); int mode = -1; if (::get_mempolicy(&mode, nullptr, 0, reinterpret_cast(base), MPOL_F_ADDR) == 0) @@ -236,7 +723,7 @@ std::optional KvCacheIntrospection::probeHostMemPlacement(size // so the device has to be reselected before binding by GPU. TLLM_CUDA_CHECK(cudaSetDevice(device)); runtime::TopologyDetector::getInstance().bindThreadByCurrentGpu(); - backing->commit(0, span); + backing->resizeTo(span); std::memset(reinterpret_cast(base), 0, span); } catch (...) diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/introspection.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/introspection.h index b77898ea7a16..0d5be223805d 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/introspection.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/introspection.h @@ -31,6 +31,9 @@ namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 { +//! Defined in batch.h, which this header does not pull in. +struct GpuCoalescedPageTable; + //! Where a host-tier allocation's pages ended up, and under which policy. struct HostMemPlacement { @@ -61,6 +64,100 @@ class KvCacheIntrospection static std::optional committedPageIsLinked(KvCache const& kvCache, int ordinal, int lcId); static bool allTreePagesDroppable(KvCacheManager& manager); + // The device-resident state of a request's sparse history cache for + // one life cycle. Synchronises and copies back, so it is test-only. Empty when the + // life cycle has no cache reserved. + static std::vector sparseCacheResidentOrdinals(KvCache const& kvCache, int lcId); + + // The first cell of this request's lease on the life cycle's arena; + // nullopt when it holds none. A lease is one contiguous run, so its first cell + // identifies it. + static std::optional sparseCacheBaseCell(KvCache const& kvCache, int lcId, int pool); + + // Units the last write-back moved per buffer, -1 where none has run. + static std::vector sparseFlushCounts(KvCache const& kvCache, int lcId); + + // A Batch's published device page table for one layer group, row-major + // [max_batch][max_blocks_per_seq]. Synchronises. + static std::vector batchPageTable(Batch const& batch, int layerGroupId); + + //! (table, entry) serving one buffer. A consumer names a table by the + //! spec it asked for and slices it at the buffer's position in that spec; this is the + //! same answer reached from the buffer, which is what the tests compare. + static std::pair batchIndexTableSlot(Batch const& batch, int layerId, DataRole const& role); + + //! What a converted index for one buffer is added to. + static MemAddress batchIndexTableBase(Batch const& batch, int layerId, DataRole const& role); + + //! One converted index table, row-major [max_batch][entries][blocks_per_seq]. + //! Reads the device copy, so it shows what the last publish() converted. Synchronises. + static std::vector batchIndexTable(Batch const& batch, int tableId); + + //! One layer group's coalesced page table on the device. + //! + //! The view, not a copy: unlike the hooks around it this neither synchronises nor + //! reads back. Its callers check the DLPack device-array contract and that the device + //! pointer is stable across resizes, both of which a host copy would erase. The host + //! table, which is what a consumer actually reads, is Batch::coalescedPageTable(). + static GpuCoalescedPageTable batchGpuCoalescedPageTable(Batch const& batch, int layerGroupId); + + // entriesPerLayer of the lease each row holds, zero where it holds none. Non-zero + // exactly where a member holds a reservation for that layer group. Synchronises. + static std::vector batchSparseEntriesPerRow(Batch const& batch, int layerGroupId); + + // A Batch's published per-row watermarks as (history_length, capacity). + // Synchronises. + static std::pair, std::vector> batchWatermarks(Batch const& batch); + + // Run the sparse resolve kernel over caller-supplied state, so the + // probe index, victim choice and recency rebuild can be exercised without a request. + // Synchronises. Inputs are [layer][entry] and [layer][selectionWidth] row-major. + struct SparseResolveProbe + { + std::vector residentOrdinal; //!< [layer][entry] after the resolve + std::vector lruOrder; //!< [layer][entry] after the resolve + std::vector outSlot; //!< [layer][width] entry per selection + //! The compacted work list, one per layer: [layer][rows * width], of which the + //! first missTotal[layer] are valid, in whatever order the CTAs raced to. + std::vector missRow; + std::vector missOrdinal; + std::vector missSlot; + std::vector missTotal; //!< [layer] entries the list holds + std::vector missCount; //!< [layer][row] entries that row contributed + std::vector inFlightSlot; //!< [layer][maxInFlightUnits] + }; + + static SparseResolveProbe sparseResolve(std::vector const& residentOrdinal, + std::vector const& lruOrder, int numLayers, int entriesPerLayer, std::vector const& selected, + int selectionWidth, std::vector const& historyUnits, std::vector const& capacityUnits, + int maxInFlightUnits, std::vector const& inactiveRows); + + // Run the sparse unit copy over caller-supplied pools, so the address + // arithmetic can be checked against a known byte pattern. Returns the destination pool + // after the copy. Synchronises. + static std::vector sparseUnitCopy(std::vector const& hostPool, uint64_t hostPageStride, + std::vector const& pageTable, std::vector const& deviceArena, int baseCell, int numBlocks, + int cellsPerBuffer, uint64_t bufferBytes, uint32_t numChunks, uint32_t chunkBytes, uint32_t chunkStride, + uint32_t unitStride, int unitsPerPage, std::vector const& missOrdinal, + std::vector const& missSlot, std::vector const& missCount, int selectionWidth, + std::vector const& bufferLayer, bool writeBack, int gridCtas); + + // The entries holding units a watermark move has just passed. + // Synchronises. + struct SparseFlushList + { + //! The compacted list, one per layer: [layer][capacity], of which the first + //! total[layer] are valid, in whatever order the CTAs raced to. + std::vector row; + std::vector ordinal; + std::vector slot; + std::vector total; //!< [layer] + std::vector count; //!< [layer] entries this row contributed + }; + + static SparseFlushList sparseFlushList(std::vector const& residentOrdinal, int numLayers, + int entriesPerLayer, int fromUnits, int toUnits, int maxFlushUnits); + // White-box hook: minimum per-pool-group slot counts to support a BatchDesc. // Reaches StorageManager::computePoolGroupSlotsForBatch() (private) via friendship. static TypedVec computeSlotsForBatch(KvCacheManager& manager, BatchDesc const& batch, diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp index ee5ed631b8ec..cb791c821abf 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp @@ -18,15 +18,19 @@ #include "kv_cache_manager_v2/kvCache.h" #include "kv_cache_manager_v2/batch.h" #include "kv_cache_manager_v2/blockRadixTree.h" +#include "kv_cache_manager_v2/coldPageCodec.h" #include "kv_cache_manager_v2/common.h" #include "kv_cache_manager_v2/exceptions.h" #include "kv_cache_manager_v2/kvCacheManager.h" +#include "kv_cache_manager_v2/sparseCacheKernels.h" #include "kv_cache_manager_v2/storageManager.h" +#include "kv_cache_manager_v2/utils/funcGuard.h" #include "kv_cache_manager_v2/utils/math.h" #include "kv_cache_manager_v2/utils/optionalGilRelease.h" #include "tensorrt_llm/common/assert.h" #include +#include #include #include #include @@ -144,6 +148,11 @@ KvCache::~KvCache() // waiting on the GIL to run a Python callback. Drop the GIL first if we happen to // hold it. OptionalGilRelease const gilRelease; + // Before close(), so a Batch never holds a row pointing at a half-torn cache. + if (mBatch != nullptr) + { + mBatch->_detach(*this); + } close(); }); } @@ -194,21 +203,20 @@ std::vector KvCache::getPagePriorities(LayerGroupId layerGroupId, Beam { throw std::out_of_range("Beam index is outside the request's beam width"); } - if (layerGroupId.value() < 0 || LifeCycleId{layerGroupId.value()} >= mManager->storage().numLifeCycles()) + if (layerGroupId.value() < 0 || layerGroupId >= mManager->storage().numLifeCycles()) { throw std::out_of_range("Layer group is outside the configured life cycles"); } std::vector result; result.reserve(mBlocks.stdSize()); - bool const isSsm - = std::holds_alternative(mManager->storage().getLifeCycle(LifeCycleId{layerGroupId.value()})); + bool const isSsm = std::holds_alternative(mManager->storage().getLifeCycle(layerGroupId)); for (auto const& block : mBlocks) { // Shared prompt blocks keep only beam 0; writable tails have a row per beam. auto const pageBeamIndex = block.pages.size() == BeamIndex{1} ? kDefaultBeamIndex : beamIndex; auto const* page = blockPageGetPage(block.pages[pageBeamIndex].at(layerGroupId)).get(); if (isSsm && block.treeBlock) - page = block.treeBlock->getPage(LifeCycleId{layerGroupId.value()}); + page = block.treeBlock->getPage(layerGroupId); // Preserve ordinals, including absent pages, just like getAggregatedPageIndices(validOnly=false). result.push_back(page ? page->priority() : kPriorityDefault); } @@ -419,115 +427,14 @@ SharedPtr KvCache::_page(BlockOrdinal ordinal, BeamIndex beamIdx, LifeCycl return blockPageGetPage(blockPage); } -CacheLevel KvCache::_lockLevel(Page const& page, BlockOrdinal ordinal) const -{ - bool const readOnly = page.isCommitted() - || (ordinal != kBadBlockOrdinal && ordinal < BlockOrdinal{mHistoryLength / mTokensPerBlock}); - return mIsDecoding && readOnly ? page.queryLockLevel() : kHotLevel; -} - -void KvCache::_offloadSparseHistory(HalfOpenRange range, int historyLength) -{ - TLLM_CHECK_DEBUG(range.end <= BlockOrdinal{historyLength / mTokensPerBlock}); - if (mHasDeferredSparseOffload) - { - range = {0, historyLength / mTokensPerBlock}; - } - bool deferred = false; - std::vector> pages; - for (auto const& [lcId, lc] : mManager->lifeCycles()) - { - auto const* attn = std::get_if(&lc); - if (!attn || !attn->isSparse) - continue; - if (range.end > mBlocks.size()) - throw std::invalid_argument("Sparse history must already have allocated pages"); - for (BlockOrdinal ord = range.beg; ord < range.end; ++ord) - { - for (BeamIndex bi{0}; bi < mBlocks[ord].pages.size(); ++bi) - { - auto page = _page(ord, bi, lcId); - TLLM_CHECK_DEBUG(page && page->status() == PageStatus::LOCKED); - if (page->cacheLevel == kSparseHistoryLevel) - continue; - auto const lock = page->holder.lock()->uniqLock.lock(); - bool needsGpu = false; - for (auto const& owner : lock->owners()) - { - int const ownerHistory = owner.kvCache == this ? historyLength : owner.kvCache->historyLength(); - if ((owner.kvCache != this && !owner.kvCache->isDecoding()) - || owner.ordinal >= BlockOrdinal{ownerHistory / owner.kvCache->tokensPerBlock()}) - { - needsGpu = true; - break; - } - } - if (needsGpu) - { - deferred = true; - continue; - } - pages.push_back(std::move(page)); - } - } - } - // Preserve retry work if allocation or copy submission fails. - mHasDeferredSparseOffload = true; - if (!pages.empty()) - { - int const oldHistoryLength = mHistoryLength; - auto restoreHistory = FuncGuard([&]() { mHistoryLength = oldHistoryLength; }); - mHistoryLength = historyLength; - offloadSparsePages(pages); - } - mHasDeferredSparseOffload = deferred; -} - -void KvCache::_publishHistoryLength(int historyLength) -{ - if (mIsDecoding && historyLength / mTokensPerBlock != mHistoryLength / mTokensPerBlock) - onPageStorageChanged(); - mHistoryLength = historyLength; -} - -bool KvCache::enterDecode() -{ - KVCM2_API_GUARD(); - auto const apiLock = mManager->lockExclusive(); - if (!isActive()) - throw LogicError("Decode admission requires an active request"); - if (mIsDecoding && !mHasDeferredSparseOffload) - return true; - try - { - _offloadSparseHistory({0, mHistoryLength / mTokensPerBlock}, mHistoryLength); - } - catch (OutOfPagesError const&) - { - return false; - } - if (!mIsDecoding) - { - mIsDecoding = true; - onPageStorageChanged(); - } - return true; -} - -void KvCache::offloadSparsePages(std::vector> const& pages) +std::optional KvCache::maxInputLength() const { - KVCM2_API_GUARD(); - auto const apiLock = mManager->lockExclusive(); - if (!isActive()) - { - throw LogicError("Sparse history offload requires an active request"); - } - MigrationRecorder const migrationRecorder - = [this](std::vector> const& sources, std::vector const& slots, CacheLevel srcLevel, - CacheLevel dstLevel) { _recordMigratedSlots(sources, slots, srcLevel, dstLevel); }; - DropRecorder const dropRecorder = [this](std::vector> const& dropped, CacheLevel level) - { _recordDroppedPages(dropped, level); }; - storageManager()->offloadSparsePages(*this, pages, migrationRecorder, dropRecorder); + // One token, not the reservation's spare width. A sparse request contributes exactly one + // query token per step, because the fetch resolves one selection per request while an + // attention kernel indexes its selection per query token; the two agree only at one + // token each. Prefill and speculative decoding break that and are out of scope, so the + // bound is a property of the supported shape rather than of how wide a cache is. + return mManager->hasSparseCache() ? std::optional{1} : std::nullopt; } void KvCache::activate() @@ -537,7 +444,8 @@ void KvCache::activate() mFinishEvent.reset(); - // Prefill restores GPU storage; decode can retain cold sparse history on host. + // Residency is a property of the life cycle: every page of a sparse buffer locks at the + // host tier, written or not, and reaches the kernels through the per-request cache. auto activePages = _activePages(); std::vector targets; targets.reserve(activePages.size()); @@ -555,14 +463,7 @@ void KvCache::activate() } auto& holder = std::get>(*bp); TLLM_CHECK_DEBUG(holder); - // A reused partial block is only a copy source. resume() replaces it with a private GPU - // page before execution, so another owner's host lock need not be moved. - bool const partialCopySource = mNeverResumed && ap.ordinal != kBadBlockOrdinal - && numCommittedTokens() % mTokensPerBlock != 0 - && ap.ordinal == BlockOrdinal{numCommittedTokens() / mTokensPerBlock}; - CacheLevel const level - = partialCopySource ? holder->page->queryLockLevel() : _lockLevel(*holder->page, ap.ordinal); - targets.push_back({holder->page, ap.beamIdx, ap.ordinal, ap.lcId, level}); + targets.push_back({holder->page, ap.beamIdx, ap.ordinal, ap.lcId}); } { @@ -581,26 +482,24 @@ void KvCache::activate() } } -bool KvCache::resume(std::optional stream, std::optional isDecoding, - std::optional> optionalReuseGroups) +bool KvCache::resume(std::optional stream, std::optional> optionalReuseGroups) { KVCM2_API_GUARD(); TLLM_CHECK(mStatus == Status::SUSPENDED); + if (mBatch != nullptr) + { + throw LogicError("A KvCache in a Batch cannot be resumed; remove it from the Batch first"); + } // Set stream first (mirrors Python: self.cuda_stream = cuda_stream). if (stream.has_value()) { - setCudaStream(*stream); + _setCudaStream(*stream); } TLLM_CHECK_WITH_INFO(mCudaStream.has_value(), "cuda_stream is never set"); TLLM_CHECK_DEBUG(!mFinishEvent.has_value()); auto const apiLock = mManager->lockExclusive(); - bool const oldIsDecoding = mIsDecoding; - if (mIsDecoding && isDecoding == false) - throw std::invalid_argument("Cannot return a decoding cache to prefill"); - auto restorePhase = FuncGuard([&]() { mIsDecoding = oldIsDecoding; }); - mIsDecoding = isDecoding.value_or(mIsDecoding); if (mNeverResumed) { _selectOptionalGroups(optionalReuseGroups.value_or(std::vector{})); @@ -645,25 +544,6 @@ bool KvCache::resume(std::optional stream, std::optional isDecod // Pre-allocate GPU slots for deferred copies (partial blocks + SSM) and scratch slots // before locking, so we never end up in a state where pages are locked but we can't allocate. TypedVec> deferredSlots(numLc); - bool deferredCopiesStarted = false; - auto releaseDeferredSlots = FuncGuard( - [&]() - { - std::optional completion; - for (LifeCycleId lc{0}; lc < numLc; ++lc) - { - if (deferredSlots[lc].has_value() && deferredSlots[lc]->hasValidSlot()) - { - if (deferredCopiesStarted) - { - if (!completion) - completion.emplace(reinterpret_cast(cudaStream())); - deferredSlots[lc]->readyEvent = *completion; - } - storageMgr.releaseSlot(lc, kHotLevel, std::move(*deferredSlots[lc])); - } - } - }); TypedVec> privateOrdinals(numLc); TypedVec> privateSlots(numLc); auto releasePrivateSlots = FuncGuard( @@ -724,11 +604,10 @@ bool KvCache::resume(std::optional stream, std::optional isDecod numSlotsNeeded[lc] += scratchDeltaCounts[lc]; } - // Only allocate if any slots are needed. - bool anyNeeded = std::any_of(numSlotsNeeded.begin(), numSlotsNeeded.end(), [](SlotCount n) { return n > 0; }); - if (anyNeeded) + bool const anyResident + = std::any_of(numSlotsNeeded.begin(), numSlotsNeeded.end(), [](SlotCount n) { return n > 0; }); { - TypedVec> tmpSlots; + TypedVec> tmpSlots(numLc); try { MigrationRecorder const migrationRecorder @@ -736,7 +615,10 @@ bool KvCache::resume(std::optional stream, std::optional isDecod CacheLevel dstLevel) { _recordMigratedSlots(pages, slots, srcLevel, dstLevel); }; DropRecorder const dropRecorder = [this](std::vector> const& pages, CacheLevel cacheLevel) { _recordDroppedPages(pages, cacheLevel); }; - tmpSlots = storageMgr.newGpuSlots(numSlotsNeeded, migrationRecorder, dropRecorder); + if (anyResident) + { + tmpSlots = storageMgr.newGpuSlots(numSlotsNeeded, migrationRecorder, dropRecorder); + } } catch (OutOfPagesError const&) { @@ -764,13 +646,13 @@ bool KvCache::resume(std::optional stream, std::optional isDecod deferredSlots[lc] = std::move(tmpSlots[lc].back()); tmpSlots[lc].pop_back(); } + // Remaining slots are scratch slots. auto& scratchSlots = mScratchSlots[lc]; scratchSlots.reserve(scratchSlots.size() + tmpSlots[lc].size()); for (auto& slot : tmpSlots[lc]) { - scratchSlots.emplace_back(std::move(slot), *this, lc, - /*skipWait=*/true); + scratchSlots.emplaceBack(std::move(slot), *this, lc, /*skipWait=*/true); scratchReadyEvents.push_back(scratchSlots.back().slot().readyEvent.handle()); } } @@ -788,12 +670,15 @@ bool KvCache::resume(std::optional stream, std::optional isDecod } catch (OutOfPagesError const&) { + for (LifeCycleId lc{0}; lc < numLc; ++lc) + { + if (deferredSlots[lc].has_value()) + storageMgr.releaseSlot(lc, kHotLevel, std::move(*deferredSlots[lc])); + } // Scratch slots stay in mScratchSlots — they'll be freed by close() inside // a recordEventScope, matching Python behavior. return false; } - mStatus = Status::ACTIVE; - auto rollbackActivation = FuncGuard([&]() { _deactivate(); }); // Private recovery pages have no reusable source. Their contents are // initialized by the caller before inference, after all allocations succeed. @@ -873,7 +758,6 @@ bool KvCache::resume(std::optional stream, std::optional isDecod copiedSources[lcIdx] = lock->page(); CacheLevel const sourceLevel = lock->page()->cacheLevel; - deferredCopiesStarted = true; storageMgr.copySlotData(lcIdx, kHotLevel, sourceLevel, newSlot.slotId(), lock->page()->slotId(), cudaStr); if ((!ssmLcId.has_value() || lcIdx != *ssmLcId) && (recordManagerStats || recordRequestStats)) { @@ -960,27 +844,15 @@ bool KvCache::resume(std::optional stream, std::optional isDecod mOptionalCandidates.clear(); // Deferred copies survive a failed decode admission and must not be repeated on retry. mPartialCopyRetentionHolds.clear(); + // A new cache is SUSPENDED until this resume() activates it, so only one that had + // already resumed counts as a preemption recovery. + bool const firstActivation = mNeverResumed; mNeverResumed = false; - if (mIsDecoding) - { - try - { - _offloadSparseHistory({0, mHistoryLength / mTokensPerBlock}, mHistoryLength); - } - catch (OutOfPagesError const&) - { - return false; - } - onPageStorageChanged(); - } - rollbackActivation.cancel(); - restorePhase.cancel(); - // Only a previously admitted request counts as a preemption recovery. - if (mHasResumed && _shouldRecordStats()) + mStatus = Status::ACTIVE; + if (!firstActivation && _shouldRecordStats()) { mManager->recordRequestResumed(); } - mHasResumed = true; return true; } @@ -1003,12 +875,11 @@ bool KvCache::prefetch(CacheLevel target) { continue; } - CacheLevel const destination = std::max(target, _lockLevel(*page, activePage.ordinal)); - if (page->cacheLevel < destination || !seen.insert(page.get()).second) + if (page->cacheLevel < target || !seen.insert(page.get()).second) { continue; } - pagesByDestination.at(destination).push_back(std::move(page)); + pagesByDestination.at(target).push_back(std::move(page)); } try @@ -1039,21 +910,17 @@ bool KvCache::prefetch(CacheLevel target) void KvCache::suspend() { KVCM2_API_GUARD(); + // A member is removed from its Batch before it suspends; see the declaration. + if (mBatch != nullptr) + { + throw LogicError("A KvCache in a Batch cannot be suspended; remove it from the Batch first"); + } auto const apiLock = mManager->lockExclusive(); TLLM_CHECK_DEBUG(mStatus == Status::ACTIVE); mOptionalCandidates.clear(); TLLM_CHECK_DEBUG(_checkSanity()); TLLM_CHECK_DEBUG(!mFinishEvent.has_value()); - _deactivate(); - if (_shouldRecordStats()) - { - mManager->recordRequestSuspended(); - } -} - -void KvCache::_deactivate() -{ // Copy data from external buffers back to internal vectors (mirrors Python's suspend). for (BeamIndex bi{0}; bi < mBeamWidth; ++bi) for (LifeCycleId lcId{0}; lcId < mBasePageIndices[bi].size(); ++lcId) @@ -1079,11 +946,13 @@ void KvCache::_deactivate() } // Free scratch slots inside scope — unlock() needs finishEvent(). // Mirrors Python: _free_scratch_slots() inside `with self._record_event()`. - _freeScratchSlots(); + _freeReservedGpuSlots(); } mStatus = Status::SUSPENDED; - mHasDeferredSparseOffload = false; - onPageStorageChanged(); + if (_shouldRecordStats()) + { + mManager->recordRequestSuspended(); + } } void KvCache::close() @@ -1098,6 +967,12 @@ void KvCache::close() TLLM_CHECK_DEBUG(_checkSanity()); if (mStatus == Status::CLOSED) return; + // Clearing the blocks below writes the staged row of a seated member; the destructor + // detaches first, so close() never arrives here seated. + if (mBatch != nullptr) + { + throw LogicError("A KvCache in a Batch cannot be closed; remove it from the Batch first"); + } mOptionalCandidates.clear(); discardPendingStats(); @@ -1129,14 +1004,7 @@ void KvCache::close() auto scope = recordEventScope(); _clearBlocks(); } - if (mPageStorageBatch != nullptr) - { - mPageStorageBatch->remove(*this); - } mStatus = Status::CLOSED; - mHasDeferredSparseOffload = false; - mPageStorageRow.reset(); - onPageStorageChanged(); mManager->unregisterKvCache(this); } @@ -1378,7 +1246,7 @@ void KvCache::_clearBlocks() // Drop last block first (mirrors Python: while self._blocks: self._blocks.pop()). while (!mBlocks.empty()) mBlocks.pop_back(); - _freeScratchSlots(); + _freeReservedGpuSlots(); // Clear SSM blocks. auto ssmLcId = mManager->lifeCycles().ssmLifeCycleId(); if (ssmLcId.has_value()) @@ -1640,11 +1508,46 @@ bool KvCache::resize(std::optional capacity, std::optional historyLeng throw std::invalid_argument("History length cannot be decreased"); if (newCap < newHist) throw std::invalid_argument("History length cannot exceed capacity"); - if (mIsDecoding && newHist > mCapacity) + + // A block must be in its host page before it can be selected, and selectable means + // below the watermark. The write-back is the Batch's, so this only checks that one has + // covered the move. + // + // Only a unit that was in flight has its one copy in a cell, and in flight means below + // the capacity this resize starts from -- so a request arriving with a history, whose + // units reached their host pages directly, crosses nothing and needs no flush. That is + // also what lets a cache be seeded before it joins a Batch. + if (newHist > mHistoryLength) + { + bool crosses = false; + for (LifeCycleId lc{0}; mBatch != nullptr && lc < mManager->numLifeCycles() && !crosses; ++lc) + { + if (mManager->sparseCacheMaxUnits(lc) == 0) + { + continue; + } + int const unitTokens = mTokensPerBlock / mManager->sparseCacheUnitsPerPage(lc); + crosses = std::min(newHist, mCapacity) / unitTokens > mHistoryLength / unitTokens; + } + if (crosses) + { + // Not equality: a step may resize a row more than once, and a row whose growth + // failed rolls back to a history below the one flushed for it. + int const flushed = mBatch != nullptr ? mBatch->flushedHistory(*mBatchRow) : 0; + TLLM_CHECK_WITH_INFO(newHist <= flushed, + "A sparse KvCache cannot advance history past what has been written back; call " + "Batch::flushSparseUnits() with the new history first"); + } + } + + // Blocks above the watermark are in flight: their data lives only in sparse cache cells, + // so the gap cannot exceed what the reservation can pin. The host-side counterpart of + // the resolve kernel's in-flight overflow report. + if (newCap != newHist) { - for (auto const& [lcId, lc] : mManager->lifeCycles()) - if (auto const* attn = std::get_if(&lc); attn && attn->isSparse) - throw std::invalid_argument("Sparse decode history cannot include unallocated input tokens"); + auto const bound = maxInputLength(); + TLLM_CHECK_WITH_INFO(!bound.has_value() || newCap - newHist <= *bound, + "capacity - history_length exceeds what the sparse cache can hold in flight"); } // Scratch reuse: enforce constraint. @@ -1660,21 +1563,12 @@ bool KvCache::resize(std::optional capacity, std::optional historyLeng } bool const recordGenerationAllocStats = mGenerationAllocReady && newCap > mCapacity; - if (!enableScratch && divUp(newCap, mTokensPerBlock) == divUp(mCapacity, mTokensPerBlock)) + if (!enableScratch && _canShortcutCapacity(newCap) && _canShortcutHistoryLength(newHist)) { - try - { - if (_shortcutSetHistoryLength(newHist)) - { - mCapacity = newCap; - _refreshGenerationAllocReady(); - return true; - } - } - catch (OutOfPagesError const&) - { - return false; - } + mCapacity = newCap; + mHistoryLength = newHist; + _refreshGenerationAllocReady(); + return true; } BlockOrdinal oldNumBlocks{divUp(mCapacity, mTokensPerBlock)}; @@ -1697,33 +1591,16 @@ bool KvCache::resize(std::optional capacity, std::optional historyLeng auto ssmLcId = mManager->lifeCycles().ssmLifeCycleId(); auto backupHolders = _unlockStaleBlocks(newHist); - // Compute scratch deltas. - auto [excessScratchSlots, deltaScratchSlots, scratchRanges] = _takeExcessScratchSlots(newCap, newHist); - auto restoreLocks = FuncGuard( - [&]() - { - _recoverExcessScratchSlots(excessScratchSlots); - _lockHeldBlocks(backupHolders); - }); - if (newNumBlocks < oldNumBlocks) { TLLM_CHECK_DEBUG_WITH_INFO(!hasScratchSlots(), "Cannot shrink while scratch slots exist"); - try - { - if (mIsDecoding) - _offloadSparseHistory({mHistoryLength / mTokensPerBlock, newHist / mTokensPerBlock}, newHist); - } - catch (OutOfPagesError const&) - { - return false; - } - restoreLocks.cancel(); _subtractPendingAllocationRange(newNumBlocks, oldNumBlocks); auto scope = recordEventScope(); _decreaseCapacity(newNumBlocks); } + auto [excessScratchSlots, deltaScratchSlots, scratchRanges] = _takeExcessScratchSlots(newCap, newHist); + if (newNumBlocks >= oldNumBlocks) { // Compute new normal slots needed per lifecycle. @@ -1785,6 +1662,8 @@ bool KvCache::resize(std::optional capacity, std::optional historyLeng } catch (OutOfPagesError const&) { + _recoverExcessScratchSlots(excessScratchSlots); + _lockHeldBlocks(backupHolders); return false; } } @@ -1793,27 +1672,6 @@ bool KvCache::resize(std::optional capacity, std::optional historyLeng newSlots.resize(numLc); } - // Reserve GPU growth before demoting history. A failed allocation or copy submission must - // leave the old capacity, watermark, and page ownership usable for a retry. - auto releaseNewSlots = FuncGuard( - [&]() - { - for (LifeCycleId lc{0}; lc < numLc; ++lc) - for (auto& slot : newSlots[lc]) - mManager->storage().releaseSlot(lc, kHotLevel, std::move(slot)); - }); - try - { - if (mIsDecoding) - _offloadSparseHistory({mHistoryLength / mTokensPerBlock, newHist / mTokensPerBlock}, newHist); - } - catch (OutOfPagesError const&) - { - return false; - } - releaseNewSlots.cancel(); - restoreLocks.cancel(); - // Wait on newly allocated slots. { std::vector readyEvents; @@ -1869,7 +1727,7 @@ bool KvCache::resize(std::optional capacity, std::optional historyLeng { auto slot = std::move(slots[lc].back()); slots[lc].pop_back(); - mScratchSlots[lc].emplace_back(std::move(slot), *this, lc, /*skipWait=*/true); + mScratchSlots[lc].emplaceBack(std::move(slot), *this, lc, /*skipWait=*/true); } } @@ -1943,7 +1801,7 @@ bool KvCache::resize(std::optional capacity, std::optional historyLeng backup.holder->page->claimRetention(std::nullopt, std::nullopt); } mCapacity = newCap; - _publishHistoryLength(newHist); + mHistoryLength = newHist; _refreshGenerationAllocReady(); TLLM_CHECK_DEBUG(_checkSanity()); return true; @@ -1963,13 +1821,21 @@ void KvCache::setCapacity(int cap) void KvCache::setHistoryLength(int hist) { - if (!resize(std::nullopt, hist)) - throw OutOfPagesError("Not enough pages to advance history"); + bool success = resize(std::nullopt, hist); + TLLM_CHECK(success); + (void) success; } -bool KvCache::_shortcutSetHistoryLength(int newHist) +bool KvCache::_canShortcutCapacity(int newCap) const { - if (newHist == mHistoryLength && !mHasDeferredSparseOffload) + if (newCap == mCapacity) + return true; + return divUp(newCap, mTokensPerBlock) == divUp(mCapacity, mTokensPerBlock); +} + +bool KvCache::_canShortcutHistoryLength(int newHist) const +{ + if (newHist == mHistoryLength) return true; // Check if stale range changes for any lifecycle. for (auto [lcId, lc] : mManager->lifeCycles()) @@ -1996,9 +1862,6 @@ bool KvCache::_shortcutSetHistoryLength(int newHist) if (changed) return false; } - if (mIsDecoding) - _offloadSparseHistory({mHistoryLength / mTokensPerBlock, newHist / mTokensPerBlock}, newHist); - _publishHistoryLength(newHist); return true; } @@ -2258,7 +2121,7 @@ std::vector KvCache::_unlockStaleBlocks(int newHistoryLeng } TLLM_CHECK_DEBUG(std::holds_alternative(bp)); auto holder = blockPageGetPage(bp)->hold(); - ret.push_back({ord, bi, lcIdx, holder, holder->page->cacheLevel}); + ret.push_back({ord, bi, lcIdx, holder}); bp = holdForCommit ? BlockPage{std::move(holder)} : BlockPage{std::monostate{}}; } } @@ -2271,7 +2134,7 @@ void KvCache::_lockHeldBlocks(std::vector const& backup) std::vector targets; targets.reserve(backup.size()); for (auto const& b : backup) - targets.push_back({b.holder->page, b.beamIdx, b.ordinal, b.lcId, b.cacheLevel}); + targets.push_back({b.holder->page, b.beamIdx, b.ordinal, b.lcId}); auto locks = batchedLockPages(*this, targets); for (size_t i = 0; i < locks.size(); ++i) @@ -2400,16 +2263,6 @@ void KvCache::_commitBlock(int ord, bool isLast, bool commitSsm, bool moveSsm) // Existing block: rebase — reuse existing block's committed pages. // Mirrors Python's `elif tree_block.is_full and allow_seq_rebasing and is_full` path. std::vector reuseTasks; - std::vector missingPages; - std::vector> originalPages; - std::vector originalLocks; - auto restorePages = FuncGuard( - [&]() - { - for (auto& [lc, bp] : originalPages) - sb.pages[kDefaultBeamIndex][lc] = std::move(bp); - _lockHeldBlocks(originalLocks); - }); for (LifeCycleId lc{0}; lc < numLc; ++lc) { if (ssmLcId.has_value() && lc == *ssmLcId) @@ -2432,7 +2285,40 @@ void KvCache::_commitBlock(int ord, bool isLast, bool commitSsm, bool moveSsm) bool isLocked = std::holds_alternative(bp); if (existingPage == nullptr) { - missingPages.push_back(lc); + // Existing page gone — put our uncommitted page into the tree block. + if (auto* lock = std::get_if(&bp)) + { + auto up = dynamicPointerCast(lock->page()); + if (up) + { + bp = std::monostate{}; + auto committed = up->convertToCommitted(newBlock, finishEvent(), numTokens); + if (newBlock->eventSink) + { + newBlock->eventSink->addStoredLifeCycle(*newBlock, lc); + } + bp = isLocked + ? BlockPage{committed->lock(*this, kDefaultBeamIndex, static_cast(ord), lc)} + : BlockPage{committed->hold()}; + } + } + else if (auto* holder = std::get_if>(&bp)) + { + if (*holder) + { + auto up = dynamicPointerCast((*holder)->page); + if (up) + { + bp = std::monostate{}; + auto committed = up->convertToCommitted(newBlock, finishEvent(), numTokens); + if (newBlock->eventSink) + { + newBlock->eventSink->addStoredLifeCycle(*newBlock, lc); + } + bp = committed->hold(); + } + } + } } else { @@ -2440,14 +2326,10 @@ void KvCache::_commitBlock(int ord, bool isLast, bool commitSsm, bool moveSsm) if (isLocked) { auto holder = blockPageGetPage(bp)->hold(); - originalLocks.push_back( - {BlockOrdinal{ord}, kDefaultBeamIndex, lc, holder, holder->page->cacheLevel}); bp = std::move(holder); } - originalPages.emplace_back(lc, std::move(bp)); - bp = std::monostate{}; - reuseTasks.push_back({existingPage->sharedFromThis(), kDefaultBeamIndex, static_cast(ord), - lc, _lockLevel(*existingPage, static_cast(ord))}); + reuseTasks.push_back( + {existingPage->sharedFromThis(), kDefaultBeamIndex, static_cast(ord), lc}); } } if (!reuseTasks.empty()) @@ -2458,24 +2340,6 @@ void KvCache::_commitBlock(int ord, bool isLast, bool commitSsm, bool moveSsm) LifeCycleId lc = reuseTasks[ri].lifeCycle; sb.pages[kDefaultBeamIndex][lc] = std::move(locks[ri]); } - if (mIsDecoding) - _offloadSparseHistory({ord, ord + 1}, mHistoryLength); - } - restorePages.cancel(); - // Publish missing lifecycle pages only after migrations and offload can no longer fail. - for (LifeCycleId lc : missingPages) - { - auto& bp = sb.pages[kDefaultBeamIndex][lc]; - auto up = dynamicPointerCast(blockPageGetPage(bp)); - if (!up) - continue; - bool const isLocked = std::holds_alternative(bp); - bp = std::monostate{}; - auto committed = up->convertToCommitted(newBlock, finishEvent(), numTokens); - if (newBlock->eventSink) - newBlock->eventSink->addStoredLifeCycle(*newBlock, lc); - bp = isLocked ? BlockPage{committed->lock(*this, kDefaultBeamIndex, BlockOrdinal{ord}, lc)} - : BlockPage{committed->hold()}; } // Don't clear SSM storage on rebase — the existing block may have a valid snapshot. sb.treeBlock = newBlock; @@ -2562,8 +2426,7 @@ void KvCache::commit(TokenSpan tokens, bool isEnd) bool const commitMinSnapshot = mManager->commitMinSnapshot(); auto ssmLcId = mManager->lifeCycles().ssmLifeCycleId(); - int const oldNumCommittedTokens = numCommittedTokens(); - int const numCommitted = oldNumCommittedTokens + static_cast(tokens.size()); + int const numCommitted = static_cast(mCommittedTokens.size()) + static_cast(tokens.size()); if (commitMinSnapshot) { if (mHistoryLength != static_cast(mCommittedTokens.size()) && mHistoryLength != numCommitted) @@ -2591,14 +2454,6 @@ void KvCache::commit(TokenSpan tokens, bool isEnd) bool const hasNewFullBlocks = newNumFullBlocks > numCommittedBlocksBefore; if (hasNewFullBlocks || hasPartialSnapshot) { - // Keep successfully published blocks, but do not retain newly appended tokens for a block - // whose rebase failed. Otherwise close() would retry that failed rebase during teardown. - auto restoreTokens = FuncGuard( - [&]() - { - int const publishedTokens = std::min(numCommitted, mNumCommittedBlocks * mTokensPerBlock); - mCommittedTokens.resize(std::max(oldNumCommittedTokens, publishedTokens)); - }); // Block whose end is the last committed token — where the SSM snapshot lives. int const ssmSnapshotOrdinal = (numCommitted - 1) / mTokensPerBlock; // Wrapped in recordEventScope() so SharedPageLock::unlock() shares one finish @@ -2625,7 +2480,6 @@ void KvCache::commit(TokenSpan tokens, bool isEnd) _snapshotPartialBlockToTree(partialOrdinal, /*commitSsm=*/ssmLcId.has_value()); } } - restoreTokens.cancel(); } if (isEnd && mCommitState != CommitState::USER_STOP) @@ -2770,8 +2624,8 @@ std::unique_ptr KvCache::planCommittedBlockDrop() BlockOrdinal windowStart; if (auto const* attn = std::get_if(&lc)) { - // Full-attention and sparse history may still be needed by later turns. - if (attn->isSparse || !attn->windowSize.has_value()) + // Full-attention blocks may still be needed by later turns. + if (!attn->windowSize.has_value()) continue; auto const staleRange = _getStaleRange(numCommittedTokens(), lc); windowStart = std::min(staleRange.end, end); @@ -3287,14 +3141,7 @@ bool KvCache::_checkSanity() const } else { - if (mStatus == Status::ACTIVE) - { - TLLM_CHECK_DEBUG(std::holds_alternative(bp)); - } - else - { - TLLM_CHECK_DEBUG(std::holds_alternative>(bp)); - } + TLLM_CHECK_DEBUG(std::holds_alternative(bp)); auto page = blockPageGetPage(bp); TLLM_CHECK_DEBUG(dynamicPointerCast(page) != nullptr); } @@ -3309,128 +3156,6 @@ bool KvCache::_checkSanity() const // Page index tables // --------------------------------------------------------------------------- -void PageStorageSnapshot::waitReady(CudaStream stream) const -{ - for (auto const& event : mReadyEvents) - event.waitInStream(stream); -} - -void KvCache::onPageStorageChanged() noexcept -{ - ++mPageStorageVersion; - mPageStorageDirty = true; - if (mPageStorageBatch != nullptr) - { - mPageStorageBatch->markDirty(*mPageStorageRow); - } -} - -uint64_t KvCache::pageStorageVersion() const -{ - auto const apiLock = mManager->lockShared(); - return mPageStorageVersion; -} - -bool KvCache::pageStorageDirty() const -{ - auto const apiLock = mManager->lockShared(); - return mPageStorageDirty; -} - -bool KvCache::acknowledgePageStorage(uint64_t version) -{ - KVCM2_API_GUARD(); - auto const apiLock = mManager->lockExclusive(); - if (version != mPageStorageVersion) - return false; - mPageStorageDirty = false; - return true; -} - -void KvCache::bindPageStorageRow(std::optional row) -{ - KVCM2_API_GUARD(); - auto const apiLock = mManager->lockExclusive(); - if (mPageStorageBatch != nullptr) - { - throw LogicError("Use Batch membership APIs to change a batched request's row"); - } - if (row && (*row < 0 || mStatus == Status::CLOSED)) - throw LogicError("Page storage rows must be nonnegative and bound to a live request"); - mPageStorageRow = row; - onPageStorageChanged(); -} - -std::optional KvCache::pageStorageRow() const -{ - auto const apiLock = mManager->lockShared(); - return mPageStorageRow; -} - -PageStorageSnapshot KvCache::getPageStorageSnapshot(LayerGroupId lgId, BeamIndex beamIdx) const -{ - // Page::status() temporarily acquires a holder with non-atomic reference counts. - auto const apiLock = mManager->lockExclusive(); - auto const& buf = mBasePageIndices.at(beamIdx).at(lgId); - PageStorageSnapshot snapshot; - snapshot.mVersion = mPageStorageVersion; - snapshot.mRow = mPageStorageRow; - auto const numBlocks = mBlocks.stdSize(); - if (numBlocks != 0) - { - std::visit([&](auto const& indices) - { snapshot.mBasePageIndices.assign(indices.data(), indices.data() + numBlocks); }, - buf); - } - snapshot.mCacheLevels.resize(numBlocks); - if (!isActive()) - { - // Held pages may be evicted or relocated; only active locks expose usable slot indices. - std::fill(snapshot.mBasePageIndices.begin(), snapshot.mBasePageIndices.end(), kBadPageIndex.value()); - return snapshot; - } - - auto const* attn = std::get_if(&mManager->lifeCycles()[lgId]); - int const completeSparseHistory = mIsDecoding && attn && attn->isSparse ? mHistoryLength / mTokensPerBlock : 0; - snapshot.mReadyEvents.reserve(numBlocks); - for (BlockOrdinal ord{0}; ord < mBlocks.size(); ++ord) - { - auto const& pages = mBlocks[ord].pages; - if (beamIdx >= pages.size()) - { - continue; - } - int const index = snapshot.mBasePageIndices[toSizeT(ord)]; - auto const& page = blockPageGetPage(pages.at(beamIdx).at(lgId)); - // The scalar count exposes only the contiguous host prefix, stopping at any deferred GPU page. - if (ord.value() == snapshot.mEligibleHistoryBlocks && ord.value() < completeSparseHistory && page - && index != kBadPageIndex.value() && page->cacheLevel == kSparseHistoryLevel) - { - TLLM_CHECK_WITH_INFO(page->status() == PageStatus::LOCKED && page->hasValidSlot() - && index == slotIdToPageIndexValue(page->slotId()), - "Eligible sparse history must have a locked host mapping"); - ++snapshot.mEligibleHistoryBlocks; - } - if (index == kBadPageIndex.value()) - continue; - // Dense SWA scratch indices refer to GPU slots without a Page object. - snapshot.mCacheLevels[toSizeT(ord)] = page ? page->cacheLevel : kHotLevel; - if (page) - snapshot.mReadyEvents.push_back(page->readyEvent); - } - return snapshot; -} - -void KvCache::recordPageStorageRead(CudaStream stream) -{ - KVCM2_API_GUARD(); - auto const apiLock = mManager->lockExclusive(); - if (!isActive()) - throw LogicError("Page storage reads must finish submission before the request is deactivated"); - CachedCudaEvent completion(stream); - completion.waitInStream(reinterpret_cast(cudaStream())); -} - void KvCache::_checkPageIndexBufferCapacity(BlockOrdinal newNumBlocks) const { for (auto const& beamIndices : mBasePageIndices) @@ -3450,8 +3175,6 @@ void KvCache::_checkPageIndexBufferCapacity(BlockOrdinal newNumBlocks) const void KvCache::_resizePageIndexBuffers(BlockOrdinal newNumBlocks) { - // External buffers keep their full allocation; the number of published entries still changes. - onPageStorageChanged(); for (BeamIndex bi{0}; bi < mBeamWidth; ++bi) { for (LifeCycleId lcId{0}; lcId < mBasePageIndices[bi].size(); ++lcId) @@ -3489,8 +3212,24 @@ int KvCache::updateBasePageIndex(BeamIndex bi, BlockOrdinal ord, LifeCycleId lc, { if (ord == kBadBlockOrdinal) return kBadPageIndex.value(); // SSM pages use BAD_BLOCK_ORDINAL + // Every page-index write in the package arrives here, from SharedPageLock acquiring or + // releasing. For a seated member the buffer below is the Batch's pinned staging row, so + // writing while a publish still reads it uploads a torn table -- one that can pair one + // request's residency with another's lease. Checked here rather than at each public + // entry point because this is the one place they all funnel through, so a page-touching + // API added later is covered without being enumerated. + // + // Gated on the package's cached gDebug rather than TLLM_CHECK_DEBUG_WITH_INFO, whose + // gate is an out-of-line call into another TU. This runs once per page per life cycle. + if (TLLM_UNLIKELY(gDebug)) + { + TLLM_CHECK_WITH_INFO(mBatch == nullptr || !mBatch->_publishInFlight(), + "A publish is still reading the staged page indices of Batch row %d; call Batch::waitPublished() before " + "an operation that locks or unlocks a seated member's pages", + mBatchRow.value_or(kNoBatchRow).value()); + } auto& buf = mBasePageIndices[bi][lc]; - int const old = std::visit( + return std::visit( [&](auto& b) -> int { using T = std::decay_t; @@ -3510,9 +3249,6 @@ int KvCache::updateBasePageIndex(BeamIndex bi, BlockOrdinal ord, LifeCycleId lc, } }, buf); - if (old != value) - onPageStorageChanged(); - return old; } Span KvCache::getBasePageIndices(LayerGroupId lgId, BeamIndex beamIdx) const @@ -3563,7 +3299,13 @@ std::vector KvCache::getAggregatedPageIndices(LayerGroupId lgId, BeamIndex void KvCache::setBasePageIndexBuf(BeamIndex beamIdx, LayerGroupId lgId, int32_t* buf, int len) { - auto const apiLock = mManager->lockExclusive(); + if (mBatch != nullptr) + { + throw LogicError( + "setBasePageIndexBuf: this KvCache is a Batch member, so its page index buffer is the Batch's row; " + "remove it from the Batch first"); + } + auto& slot = mBasePageIndices[beamIdx][lgId]; BlockOrdinal const numBlocks = mBlocks.size(); @@ -3576,7 +3318,6 @@ void KvCache::setBasePageIndexBuf(BeamIndex beamIdx, LayerGroupId lgId, int32_t* auto const n = std::min(toSizeT(numBlocks), static_cast(ext->len)); std::vector vec(ext->ptr, ext->ptr + n); slot = std::move(vec); - onPageStorageChanged(); } // If already a vector, nothing to do. return; @@ -3605,7 +3346,6 @@ void KvCache::setBasePageIndexBuf(BeamIndex beamIdx, LayerGroupId lgId, int32_t* std::copy(oldData, oldData + copyLen, buf); std::fill(buf + copyLen, buf + len, kBadPageIndex.value()); slot = Span{buf, len}; - onPageStorageChanged(); } int KvCache::getSsmBlockBaseIndex(LayerGroupId lgId, BeamIndex beamIdx) const @@ -3636,12 +3376,12 @@ bool KvCache::isSwaScratchReuseEnabled() const noexcept return mEnableSwaScratchReuse; } -bool KvCache::supportsIndexMode(PageIndexMode mode) const +bool KvCache::supportsIndexMode(PageIndexMode mode, LifeCycleId lcId) const { switch (mode) { case PageIndexMode::PER_LAYER: return true; - case PageIndexMode::SHARED: return !hasScratchSlots(); + case PageIndexMode::SHARED: return mScratchSlots.at(lcId).empty(); } return false; } @@ -3674,20 +3414,116 @@ int KvCache::_swaScratchMaxRewindLen() const return cfg->maxRewindLen; } -std::optional KvCache::getScratchDesc(LayerGroupId lgId) const +void ScratchSlots::_mirrorBack() +{ + auto const index = mSlots.size() - 1; + auto const value = slotIdToPageIndexValue(mSlots.back().slot().slotId()); + std::visit( + [&](auto& buf) + { + using T = std::decay_t; + if constexpr (std::is_same_v>) + { + if (buf.size() <= index) + { + buf.resize(index + 1, kBadPageIndex.value()); + } + } + TLLM_CHECK_WITH_INFO(index < static_cast(buf.size()), + "Scratch slot count %zu exceeds the %zu the batch staged for this life cycle", index + 1, + static_cast(buf.size())); + buf[index] = value; + }, + mIndices); +} + +void ScratchSlots::_mirrorSentinel(size_t from, std::optional to) +{ + auto const mirror = std::visit( + [](auto& buf) { + return Span{buf.data(), static_cast(buf.size())}; + }, + mIndices); + auto const end = to.value_or(static_cast(mirror.len)); + TLLM_CHECK_DEBUG_WITH_INFO(from <= end && end <= static_cast(mirror.len), + "Scratch mirror range [%zu, %zu) outruns its %d entries", from, end, mirror.len); + std::fill(mirror.ptr + from, mirror.ptr + end, kBadPageIndex.value()); +} + +ScratchSlotLock ScratchSlots::popBack() +{ + ScratchSlotLock lock = std::move(mSlots.back()); + mSlots.pop_back(); + _mirrorSentinel(mSlots.size(), mSlots.size() + 1); + return lock; +} + +void ScratchSlots::unlockAndClear() +{ + for (auto& lock : mSlots) + { + lock.unlock(); + } + // The mirror already reads sentinel past the live count. + auto const live = mSlots.size(); + mSlots.clear(); + _mirrorSentinel(0, live); +} + +void ScratchSlots::setBuf(int32_t* buf, int len) +{ + auto const live = indices(); + if (buf == nullptr || len == 0) + { + std::vector owned(mSlots.size(), kBadPageIndex.value()); + std::copy(live.data(), live.data() + live.len, owned.begin()); + mIndices = std::move(owned); + return; + } + TLLM_CHECK_WITH_INFO(static_cast(len) >= mSlots.size(), + "setScratchSlotBuf: buffer of %d is shorter than the %zu live scratch slots", len, mSlots.size()); + std::vector seed(live.data(), live.data() + live.len); + std::fill(buf, buf + len, kBadPageIndex.value()); + std::ranges::copy(seed, buf); + mIndices = Span{buf, len}; +} + +Span ScratchSlots::indices() const +{ + auto const result = std::visit( + [this](auto const& buf) -> Span { + return {buf.data(), static_cast(mSlots.size())}; + }, + mIndices); + if (TLLM_UNLIKELY(gDebug)) + { + for (size_t i = 0; i < mSlots.size(); ++i) + { + TLLM_CHECK(result[static_cast(i)] == slotIdToPageIndexValue(mSlots[i].slot().slotId())); + } + } + return result; +} + +HalfOpenRange KvCache::getSwaScratchRange(LayerGroupId lgId) const +{ + // Lock-free: scratch slots are exclusively owned by this KvCache, and defrag-caused page + // movement happens only when every KvCache is suspended, which means mScratchSlots is empty. + TLLM_CHECK_WITH_INFO(isActive(), "getSwaScratchRange() requires an ACTIVE KvCache"); + return _getScratchRange(mManager->lifeCycles().getLifeCycle(lgId)); +} + +Span KvCache::getScratchSlots(LayerGroupId lgId) const +{ + return mScratchSlots[lgId].indices(); +} + +void KvCache::setScratchSlotBuf(LayerGroupId lgId, int32_t* buf, int len) { - // Lock-free: scratch slots are exclusively owned by this KvCache and defrag-cased page movement - // happens only when all KvCache instances are suspended, which means mScratchSlots is empty. - TLLM_CHECK_WITH_INFO(isActive(), "getScratchDesc() requires an ACTIVE KvCache"); - auto const& lc = mManager->lifeCycles().getLifeCycle(lgId); - auto sr = _getScratchRange(lc); - if (!sr) - return std::nullopt; - std::vector slotIds; - slotIds.reserve(mScratchSlots[lgId].size()); - for (auto const& lock : mScratchSlots[lgId]) - slotIds.push_back(slotIdToPageIndexValue(lock.slot().slotId())); - return ScratchDesc{sr, std::move(slotIds)}; + TLLM_CHECK_WITH_INFO(mBatch == nullptr, + "setScratchSlotBuf: this KvCache is a Batch member, so its scratch slots are the Batch's row; " + "remove it from the Batch first"); + mScratchSlots[lgId].setBuf(buf, len); } void KvCache::setEnableSwaScratchReuse(bool enable) @@ -3764,8 +3600,7 @@ KvCache::DeltaScratchSlots KvCache::_takeExcessScratchSlots(int capacity, int hi { for (int i = 0; i < -delta; ++i) { - result.excess[lcIdx].push_back(std::move(mScratchSlots[lcIdx].back())); - mScratchSlots[lcIdx].pop_back(); + result.excess[lcIdx].push_back(mScratchSlots[lcIdx].popBack()); } } } @@ -3778,19 +3613,17 @@ void KvCache::_recoverExcessScratchSlots(TypedVec const& blockPageGetPage(BlockPage const& bp) noexcept return sNull; } +// --------------------------------------------------------------------------- +// ScratchSlots — one life cycle's scratch slot locks, and their page-index mirror. +// +// The mirror is positional: entry `i` holds the page index of the `i`th lock, which is +// what the conversion kernel reads when a block falls inside the scratch range. Every +// length change goes through this class, so the two stay in step. +// +// Entries past the live count read kBadPageIndex; zero would name slot zero. +// --------------------------------------------------------------------------- +class ScratchSlots +{ +public: + template + void emplaceBack(Args&&... args) + { + mSlots.emplace_back(std::forward(args)...); + _mirrorBack(); + } + + [[nodiscard]] ScratchSlotLock popBack(); + + //! Unlock every slot, drop them, and return the mirror to the sentinel, so it names no + //! page that has gone back to the pool. + void unlockAndClear(); + + void reserve(size_t n) + { + mSlots.reserve(n); + } + + //! Mirror into `buf` instead of the internal vector, seeding it from the live slots. + //! A null `buf` reverts to the internal vector, carrying the current contents back. + void setBuf(int32_t* buf, int len); + + //! Page indices of the live slots, in order. + [[nodiscard]] Span indices() const; + + [[nodiscard]] size_t size() const noexcept + { + return mSlots.size(); + } + + [[nodiscard]] bool empty() const noexcept + { + return mSlots.empty(); + } + + [[nodiscard]] ScratchSlotLock const& back() const + { + return mSlots.back(); + } + + [[nodiscard]] ScratchSlotLock const& operator[](size_t i) const + { + return mSlots[i]; + } + + [[nodiscard]] auto begin() const noexcept + { + return mSlots.begin(); + } + + [[nodiscard]] auto end() const noexcept + { + return mSlots.end(); + } + +private: + void _mirrorBack(); + //! Fill [from, to) of the mirror with the sentinel; unset `to` means its end. + void _mirrorSentinel(size_t from, std::optional to = std::nullopt); + + std::vector mSlots; + //! Owned when the vector, borrowed when the span; the span's length is the Batch's + //! per-row stride, so it holds sentinel past `mSlots.size()`. + std::variant, Span> mIndices; +}; + // --------------------------------------------------------------------------- // SeqBlock — one block-slot in a KvCache sequence. // pages[beamIdx][lcId] tracks who holds/locks each page. @@ -155,56 +235,6 @@ class PlannedDropHandle std::optional>> mPageRefs; }; -// Independent host metadata for one request's layer group and beam. Events retain copy readiness, -// not storage ownership: use the indices only while the request is active and its version matches. -class PageStorageSnapshot -{ -public: - uint64_t version() const noexcept - { - return mVersion; - } - - std::optional row() const noexcept - { - return mRow; - } - - std::vector const& basePageIndices() const noexcept - { - return mBasePageIndices; - } - - // A missing level accompanies BAD_PAGE_INDEX. Valid indices address slots in the indicated level. - std::vector> const& cacheLevels() const noexcept - { - return mCacheLevels; - } - - // Contiguous complete sparse history on host; stops at the first missing or GPU mapping. - int eligibleHistoryBlocks() const noexcept - { - return mEligibleHistoryBlocks; - } - - std::vector const& readyEvents() const noexcept - { - return mReadyEvents; - } - - // Queue copy-completion dependencies without blocking the CPU or publishing device metadata. - void waitReady(CudaStream stream) const; - -private: - friend class KvCache; - uint64_t mVersion = 0; - std::optional mRow; - std::vector mBasePageIndices; - std::vector> mCacheLevels; - int mEligibleHistoryBlocks = 0; - std::vector mReadyEvents; -}; - // Availability at the claimed prefix, indexed by LayerGroupId. OPTIONAL // candidates are held until the first successful resume or close. // Describes the initial prefix claim, not subsequent request progress. @@ -261,23 +291,21 @@ class KvCache : public std::enable_shared_from_this // Resume: check utilization and lock active pages at their required storage levels. // Optionally sets a new CUDA stream; if nullopt, uses the existing one. - // isDecoding defaults to the current phase. A cache starts in prefill and cannot return to it - // after decode admission. Set true when admitting a suspended request directly to decode. // Returns false if utilization too high or out of memory. - bool resume(std::optional stream = std::nullopt, std::optional isDecoding = std::nullopt, + // + // Refused while the cache belongs to a Batch: locking pages writes through the bound + // Span, which is the Batch's staged row, so a member is resumed before it joins. + bool resume(std::optional stream = std::nullopt, std::optional> optionalReuseGroups = std::nullopt); - // Enter decode only after prefill has submitted its final KV accesses. Offloads complete sparse - // history, deferring pages still needed on GPU by another owner. Retries deferred pages even - // with an unchanged watermark. Returns false on host OOM. - bool enterDecode(); - - bool isDecoding() const noexcept - { - return mIsDecoding; - } - // Suspend: detach from CUDA stream, unlock pages → PageHolder. + // + // Refused while the cache belongs to a Batch, for the reason resume() is. + // + // Returns the sparse cache arena lease, which **discards whatever the GPU cache held**. + // A lease comes back with every entry free, so a unit whose only copy was in a cell can + // no longer be addressed: call Batch::flushSparseUnits() first for anything worth + // keeping. void suspend(); // Close: release all blocks back to KvCacheManager. @@ -304,37 +332,23 @@ class KvCache : public std::enable_shared_from_this // Returns true on success, false when the resize ran out of pages and was rolled back. bool resize(std::optional capacity, std::optional historyLength = std::nullopt); + // ---- Sparse history cache ---------------------------------------------- + + //! Bound on capacity - history_length: how far capacity may outrun the watermark. + //! + //! Derived, not declared. The cache holds a whole selection plus every unit still in + //! flight, so whatever the reservation has above the selection width is the in-flight + //! headroom, and that headroom in tokens is this bound. Size it through + //! BufferConfig::maxCachedUnits. + //! + //! nullopt when no sparse buffer has a cache, where a bound on in-flight blocks + //! describes nothing. The smallest bound across sparse life cycles when several differ. + [[nodiscard]] std::optional maxInputLength() const; + // Convenience: set only capacity or history length. void setCapacity(int capacity); void setHistoryLength(int historyLength); - //! Internal explicit demotion of complete, locked sparse history. Takes the manager's exclusive lock. - //! Duplicate pages and pages already in host history are ignored. Does not advance history length. - //! Caller must ensure every owner has finished the execution phase that requires these pages on GPU. - void offloadSparsePages(std::vector> const& pages); - - // CPU-side invalidation for page indices, levels, readiness, eligibility and row/buffer bindings. - // Several changes leave one pending refresh. Acknowledging an older version never clears it. - uint64_t pageStorageVersion() const; - bool pageStorageDirty() const; - // Acknowledge only after refreshing every group/beam from snapshots of this same version. - bool acknowledgePageStorage(uint64_t version); - - // Associate a consumer's stable row with this request; nullopt detaches it. Every bind - // requires a refresh, including reuse of the same row number by a new consumer. Close detaches it. - // Batch members must use Batch::add/remove instead of rebinding directly. - void bindPageStorageRow(std::optional row); - std::optional pageStorageRow() const; - - // Copies raw indices (no expansion or BAD-to-zero conversion) and readiness under the API lock. - // Eligibility is zero for prefill, inactive requests and non-sparse layer groups. - PageStorageSnapshot getPageStorageSnapshot(LayerGroupId lgId, BeamIndex beamIdx = kDefaultBeamIndex) const; - - // Call after submitting reads on another stream, before mutating/suspending/closing this cache. - // Joins those reads into the request stream so its normal unlock/commit fences protect storage. - // Snapshot acquisition, read submission and this call belong to the request's owning thread. - void recordPageStorageRead(CudaStream stream); - // ---- Committing tokens ------------------------------------------------- // Commit tokens: finalises the oldest uncommitted block and makes it @@ -355,7 +369,7 @@ class KvCache : public std::enable_shared_from_this // Get base page indices (slot_id) for beamIdx × layerGroupId. // Returns a non-owning Span into the page-index buffer (owned by this KvCache, or by the // caller when set via setBasePageIndexBuf). The span is valid until the next resize(), - // setBasePageIndexBuf() or close(); its contents are also rewritten by suspend()/resume() and offload. + // setBasePageIndexBuf() or close(); its contents are also rewritten by suspend()/resume(). Span getBasePageIndices(LayerGroupId lgId, BeamIndex beamIdx = kDefaultBeamIndex) const; // Get aggregated (slot-level) page indices for one layer group + beam. @@ -367,6 +381,11 @@ class KvCache : public std::enable_shared_from_this // Zero-copy page index buffer: copy current base indices into [buf, buf+len) // and arrange that future updateBasePageIndex calls write there too. // Pass buf=nullptr / len=0 to revert to the internal vector. + // + // Refused while this cache belongs to a Batch: membership binds every layer group to a + // row of the Batch's table, and a second buffer would take the writes with it, leaving + // that row holding whatever it held when the cache joined. Remove it from the Batch + // first, or let the Batch be the one holding the buffer worth writing to. void setBasePageIndexBuf(BeamIndex beamIdx, LayerGroupId lgId, int32_t* buf, int len); // ---- Introspection ----------------------------------------------------- @@ -376,6 +395,12 @@ class KvCache : public std::enable_shared_from_this return mStatus; } + // Row this cache occupies in its Batch; nullopt when it belongs to none. + std::optional batchRow() const noexcept + { + return mBatchRow; + } + CommitState commitState() const noexcept { return mCommitState; @@ -465,7 +490,7 @@ class KvCache : public std::enable_shared_from_this // Plan dropping SWA blocks needed only by the next conversation turn. // // The plan covers committed pages in each SWA life cycle's current attention - // window. Sparse history, full-attention, and attention-sink blocks are excluded because + // window. Full-attention and attention-sink blocks are excluded because // later turns may still need them. An SSM life cycle contributes its final block. Must be // called after stopCommitting(). Returns nullptr without creating a plan if // any required SWA page is unavailable. Mirrors Python's @@ -507,9 +532,20 @@ class KvCache : public std::enable_shared_from_this CUstream cudaStream() const; + // A member runs on its Batch's stream; moving it would order the Batch's transfers and + // the member's own work against different streams. + void setCudaStream(CUstream stream) + { + if (mBatch != nullptr) + { + throw LogicError("A KvCache in a Batch runs on the Batch's stream; remove it first to move it"); + } + _setCudaStream(stream); + } + // Mirrors Python's cuda_stream setter: if already on a stream AND active, // make the new stream wait for the old one before switching (cross-stream sync). - void setCudaStream(CUstream stream) + void _setCudaStream(CUstream stream) { if (mCudaStream.has_value()) { @@ -569,8 +605,21 @@ class KvCache : public std::enable_shared_from_this // ---- SWA scratch slot management ------------------------------------------ - // Return scratch metadata for a layer group, or nullopt if no scratch blocks. - std::optional getScratchDesc(LayerGroupId lgId) const; + // Block ordinal range this life cycle serves from scratch slots; empty when none of it + // does. Read it alongside getScratchSlots(), which names the slots it covers. + HalfOpenRange getSwaScratchRange(LayerGroupId lgId) const; + + // Page indices of this life cycle's scratch slots, in slot order. + // Valid until the next scratch slot change or setScratchSlotBuf() for this life cycle. + Span getScratchSlots(LayerGroupId lgId) const; + + // Zero-copy scratch slot buffer: copy the current indices into [buf, buf+len) and + // arrange that future scratch slot changes write there too. Pass buf=nullptr / len=0 + // to revert to internal storage. + // + // Refused while this cache belongs to a Batch, which binds the buffer itself; see + // setBasePageIndexBuf. + void setScratchSlotBuf(LayerGroupId lgId, int32_t* buf, int len); // True if any lifecycle has scratch slots allocated. bool hasScratchSlots() const; @@ -591,41 +640,26 @@ class KvCache : public std::enable_shared_from_this // requires the already-committed tokens to be digest-free (verified by a scan). void setTextOnly(bool textOnly); - // Whether the given page index mode is supported (SHARED requires no scratch slots). - bool supportsIndexMode(PageIndexMode mode) const; + // Whether the mode is supported for one life cycle: SHARED requires that it hold no + // scratch slots. + bool supportsIndexMode(PageIndexMode mode, LifeCycleId lcId) const; // ---- Internal callbacks (called by SharedPageLock) ---------------------- - // Caller holds the manager's exclusive API lock, including when updating another owner's table. int updateBasePageIndex(BeamIndex bi, BlockOrdinal ord, LifeCycleId lc, int value); std::optional id; // opaque identifier (mirrors Python's id field) private: friend class KvCacheIntrospection; - friend class UniqPageLock; friend class Batch; friend std::vector batchedLockPages( KvCache& kvCache, std::vector const& targets); - void onPageStorageChanged() noexcept; - // Activate: lock active pages at their required levels. mCudaStream must already be set. // Internal — called by resume(). Not public (mirrors Python where activate() doesn't exist). void activate(); - // Release active locks and scratch slots without recording a scheduler suspension. - void _deactivate(); - - // Prefill and writable pages require GPU storage. Decode keeps cold sparse history on host. - CacheLevel _lockLevel(Page const& page, BlockOrdinal ordinal) const; - - // Offload GPU pages in the supplied complete-history range and retry deferred history. - // Pages stay on GPU until they belong to every live owner's complete decode history. - // The candidate watermark is visible only under the exclusive API lock until offload succeeds. - void _offloadSparseHistory(HalfOpenRange range, int historyLength); - void _publishHistoryLength(int historyLength); - // Internal helpers. // Turn the per-block cache levels observed while holding the matched pages into logical token // counts. Called at the end of _setupForReuse, which collects them in the same walk. @@ -663,7 +697,6 @@ class KvCache : public std::enable_shared_from_this BeamIndex beamIdx; LifeCycleId lcId; SharedPtr holder; - CacheLevel cacheLevel; }; // Unlock stale SWA blocks. Returns backup holders for rollback. @@ -684,7 +717,11 @@ class KvCache : public std::enable_shared_from_this std::vector _activePages() const; SharedPtr _page(BlockOrdinal ordinal, BeamIndex beamIdx, LifeCycleId lcId) const; - bool _shortcutSetHistoryLength(int historyLength); + //! Whether the new value can be taken without touching blocks. Both are asked before + //! either is applied, so a resize that only one of them can shortcut leaves neither + //! changed and the slow path still sees the original state. + [[nodiscard]] bool _canShortcutCapacity(int capacity) const; + [[nodiscard]] bool _canShortcutHistoryLength(int historyLength) const; bool _shouldRecordManagerStats() const; bool _shouldRecordRequestStats() const; bool _shouldRecordStats() const; @@ -761,8 +798,8 @@ class KvCache : public std::enable_shared_from_this // Recover previously taken excess scratch slots back into mScratchSlots. void _recoverExcessScratchSlots(TypedVec>& excess); - // Release all scratch slots back to storage. - void _freeScratchSlots(); + // Release every GPU slot this request reserves: SWA scratch and sparse history cache. + void _freeReservedGpuSlots(); // Whether any lifecycle would require scratch blocks at the current state. bool _wouldUseSwaScratchBlocks() const; @@ -784,9 +821,6 @@ class KvCache : public std::enable_shared_from_this BeamIndex mBeamWidth; int mCapacity; int mHistoryLength; - bool mIsDecoding = false; - // Retry by scanning current blocks; deferred work does not retain pages or other requests. - bool mHasDeferredSparseOffload = false; std::optional mExpectedPromptLength; bool mGenerationAllocReady = false; @@ -796,10 +830,11 @@ class KvCache : public std::enable_shared_from_this using LifeCyclePageIndexBuffers = TypedVec; using BeamPageIndexBuffers = TypedVec; BeamPageIndexBuffers mBasePageIndices; - Batch* mPageStorageBatch = nullptr; // Non-owning; both destructors detach membership. - uint64_t mPageStorageVersion = 0; - bool mPageStorageDirty = true; - std::optional mPageStorageRow; + + // Set while this cache is a Batch member; the Batch clears both on removal and the + // destructor detaches, so a dropped request cannot leave a row pointing at it. + Batch* mBatch = nullptr; + std::optional mBatchRow; TypedVec mBlocks; std::vector>> mPartialCopyRetentionHolds; @@ -822,7 +857,6 @@ class KvCache : public std::enable_shared_from_this // SSM pages: [beamIdx][lcId] — always initialized (empty entries = monostate). BeamBlockPages mSsmBlocks; bool mNeverResumed = true; - bool mHasResumed = false; // Successful admission, independent of completed deferred copies. PendingStats mPendingStats; @@ -836,7 +870,12 @@ class KvCache : public std::enable_shared_from_this // SWA scratch slot support. bool mEnableSwaScratchReuse = false; - TypedVec> mScratchSlots; + TypedVec mScratchSlots; + + // This request's slot in the sparse cache arenas, which addresses its cells and the + // device-resident residency, recency and write-back scratch describing them in every + // life cycle at once. Taken when the cache joins a Batch, returned on removal, + // suspend() and close(). Empty unless some life cycle has its fetch path configured. }; } // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.cpp index 1ec31e1bc0ae..77a7e427b91f 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.cpp @@ -16,8 +16,10 @@ */ #include "kv_cache_manager_v2/kvCacheManager.h" +#include "kv_cache_manager_v2/batch.h" #include "kv_cache_manager_v2/common.h" #include "kv_cache_manager_v2/exceptions.h" +#include "kv_cache_manager_v2/sparseCacheKernels.h" #include "kv_cache_manager_v2/storage/config.h" #include "kv_cache_manager_v2/storage/core.h" #include "kv_cache_manager_v2/utils/math.h" @@ -28,7 +30,7 @@ #include #include #include -#include +#include #include namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 @@ -46,12 +48,13 @@ static double nowSeconds() // PageIndexConverter // --------------------------------------------------------------------------- -std::vector PageIndexConverter::operator()( - std::vector const& baseIndices, std::optional indexMode, ScratchDesc const* scratch) const +std::vector PageIndexConverter::operator()(std::vector const& baseIndices, + std::optional indexMode, HalfOpenRange scratchRange, + Span scratchSlots) const { if (!indexMode.has_value()) { - if (scratch) + if (!scratchRange.empty()) { throw std::invalid_argument("index_mode must be provided when scratch is active"); } @@ -67,13 +70,16 @@ std::vector PageIndexConverter::operator()( for (BlockOrdinal ordinal{0}; ordinal < BlockOrdinal{static_cast(baseIndices.size())}; ++ordinal) { int index; - if (scratch && scratch->range.contains(ordinal)) + if (scratchRange.contains(ordinal)) { - // Scratch block: slot IDs come from ScratchDesc, not base_indices. - int blockPos = ordinal - scratch->range.beg; + // Scratch block: its slot comes from the request's own scratch slots, not from + // base_indices. + int blockPos = ordinal - scratchRange.beg; int totalOffset = blockPos * scratchPages; int slotIdx = totalOffset / scale; - int slotId = scratch->slotIds[static_cast(slotIdx)]; + TLLM_CHECK_DEBUG_WITH_INFO(slotIdx < scratchSlots.size(), "Scratch slot %d past the %d the range describes", + slotIdx, scratchSlots.size()); + int slotId = scratchSlots[slotIdx]; int offset = totalOffset % scale; index = slotId * scale + (offset + appliedLayerOffset) % scale; } @@ -95,7 +101,7 @@ std::vector PageIndexConverter::operator()( std::vector PageIndexConverter::operator()(int baseIndex) const { - return operator()(std::vector{baseIndex}, std::nullopt, nullptr); + return operator()(std::vector{baseIndex}); } // --------------------------------------------------------------------------- @@ -127,6 +133,13 @@ KvCacheManager::KvCacheManager(KVCacheManagerConfig const& config, std::shared_p std::move(coldPageCodec), mConfig.swaScratchReuse, mConfig.typicalStep, mConfig.constraints, mConfig.initialPoolRatio, mEventSink, mConfig.maxUtilForResume, mConfig.coldTierOffloadMinPriority); + _computeSparseCacheShape(); + + // The storage above has a device context up, which this needs. + CUdevice device{}; + cuCheck(cuCtxGetDevice(&device)); + mSmCount = sparseCopyMultiprocessors(device); + mTargetRatioListHot = _currentHotRatio(); mTargetRatioListCold = _currentColdRatios(); _resetIterationPeakNumBlocks(); @@ -139,6 +152,252 @@ KvCacheManager::~KvCacheManager() KVCM2_POISON_ON_EXCEPT([this]() { shutdown(); }); } +void KvCacheManager::_computeSparseCacheShape() +{ + LifeCycleId const numLc = mStorage->numLifeCycles(); + mSparseCacheShape = TypedVec(numLc, SparseCacheShape{}); + mSelectionSlot.clear(); + mSparseLayerSelection = TypedVec>(numLc); + + // Selection ids are the caller's, so they are dense only by luck; these give each one + // a slot in its life cycle's metadata, and catch an id used across two of them. + std::unordered_map selectionLifeCycle; + std::map> slotsOf; + + for (auto const& layer : mConfig.layers) + { + auto const* attn = std::get_if(&layer); + if (attn == nullptr) + { + continue; + } + if (!attn->sparseCache.has_value()) + { + continue; + } + LifeCycleId const lcId = mLifeCycles.getId(makeLifeCycle(layer, mConfig.tokensPerBlock)); + auto const& cache = *attn->sparseCache; + int const unitsPerPage = cache.unitsPerPage(mConfig.tokensPerBlock); + + // Layers of one life cycle must currently declare the same sparse_cache. The arena + // is per life cycle and its metadata is one stride: a selection's rows are found at + // `residentOrdinal + slot * entriesPerSelection`, so a single entry count serves + // them all, and cellsPerBuffer follows from maxCachedUnits and unitsPerPage alike. + // See SPARSE_ATTENTION_DESIGN.md. + auto& shape = mSparseCacheShape[lcId]; + if (shape.layers > 0) + { + TLLM_CHECK_WITH_INFO(shape.maxUnits == cache.maxCachedUnits && shape.selection == cache.numSelectedBlocks + && shape.unitsPerPage == unitsPerPage, + "Layers sharing a life cycle must declare the same sparse_cache; they share one arena"); + } + shape.maxUnits = cache.maxCachedUnits; + shape.selection = cache.numSelectedBlocks; + shape.unitsPerPage = unitsPerPage; + + bool layerCounted = false; + for (auto const& buf : attn->buffers) + { + // A zero-size buffer has no layout and never reaches storage, so the plan below + // has no entry for it. Counting it here would size this metadata for a buffer + // nothing writes. + if (!buf.isSparse || buf.size == 0) + { + continue; + } + ++shape.buffers; + if (!layerCounted) + { + ++shape.layers; + layerCounted = true; + } + } + if (!layerCounted) + { + continue; + } + + // A selection is a life cycle's, because the metadata it indexes is. Layers that + // name one from either side of a life cycle boundary are naming two things. + int const selectionId = cache.effectiveSelectionId(attn->layerId); + auto const [owner, firstHere] = selectionLifeCycle.try_emplace(selectionId, lcId); + TLLM_CHECK_WITH_INFO(firstHere || owner->second == lcId, + "Layers sharing a selection_id must share a life cycle; selection %d spans two", selectionId); + + auto& slots = slotsOf[lcId]; + auto const [slot, fresh] = slots.try_emplace(selectionId, static_cast(slots.size())); + mSelectionSlot[attn->layerId] = slot->second; + // Indexed by the layer's dense row, which is what the arena and the write-back's + // grid address it by; _sparseBufferPlan numbers it from this same walk. + mSparseLayerSelection[lcId].push_back(static_cast(slot->second)); + if (fresh) + { + ++shape.groups; + } + } +} + +int KvCacheManager::_selectionSlotOf(LayerId layerId) const +{ + auto const it = mSelectionSlot.find(layerId); + TLLM_CHECK_WITH_INFO(it != mSelectionSlot.end(), "This layer carries no sparse selection"); + return it->second; +} + +std::vector const& KvCacheManager::_sparseLayerSelection(LifeCycleId lifeCycle) const +{ + return mSparseLayerSelection[lifeCycle]; +} + +int KvCacheManager::_selectionLayerCount(LifeCycleId lifeCycle, int slot) const +{ + auto const& selections = mSparseLayerSelection[lifeCycle]; + return static_cast(std::count(selections.begin(), selections.end(), slot)); +} + +int KvCacheManager::_sparseCacheSubSlotsPerBuffer(LifeCycleId lifeCycle) const noexcept +{ + int const maxUnits = sparseCacheMaxUnits(lifeCycle); + if (maxUnits <= 0) + { + return 0; + } + return divUp(maxUnits, sparseCacheUnitsPerPage(lifeCycle)); +} + +int KvCacheManager::_sparseCacheMaxInFlightUnits(LifeCycleId lifeCycle) const noexcept +{ + int const cells = _sparseCacheSubSlotsPerBuffer(lifeCycle) * sparseCacheUnitsPerPage(lifeCycle); + if (cells <= 0) + { + return 0; + } + // A step advances capacity one token past history, and one token straddles at most one + // unit boundary, so at most one unit is ever in flight. That is the whole of the + // headroom: everything else the reservation holds is available to the selection. + return 1; +} + +std::vector KvCacheManager::_sparseBufferPlan(LifeCycleId lifeCycle) const +{ + // Layers outermost, as _computeSparseCacheShape walks them. + std::vector plan; + std::optional previousLayer; + int32_t layerRow = -1; + for (auto const& layer : mConfig.layers) + { + auto const* attn = std::get_if(&layer); + if (attn == nullptr || mLifeCycles.getId(makeLifeCycle(layer, mConfig.tokensPerBlock)) != lifeCycle) + { + continue; + } + auto const firstOfLayer = plan.size(); + for (auto const& buf : attn->buffers) + { + // A sparse buffer in a layer with no sparse_cache has storage and locking but + // no fetch path, so it contributes no cells and no layout. + // Zero-size as in _computeSparseCacheShape: no layout to derive, and + // sparseLayout() rejects one rather than inventing it. + if (!buf.isSparse || buf.size == 0 || !attn->sparseCache.has_value()) + { + continue; + } + if (!previousLayer.has_value() || attn->layerId != *previousLayer) + { + previousLayer = attn->layerId; + ++layerRow; + } + auto const attr = _bufferAttr(attn->layerId, buf.role); + plan.push_back(SparseBufferPlan{BufferId{attn->layerId, buf.role}, attr.poolIndex, attr.size, + static_cast(attr.offset), + buf.sparseLayout(attn->sparseCache->unitsPerPage(mConfig.tokensPerBlock)), layerRow, 0}); + } + // Group a layer's buffers by pool so a (layer, pool) run is consecutive too; the copy + // addresses one pool per launch. Stable, so declaration order survives within a pool. + std::stable_sort(plan.begin() + static_cast(firstOfLayer), plan.end(), + [](SparseBufferPlan const& a, SparseBufferPlan const& b) { return a.poolIndex < b.poolIndex; }); + } + + std::map nextSlot; + for (auto& entry : plan) + { + entry.slotInPool = nextSlot[entry.poolIndex]++; + } + return plan; +} + +std::vector KvCacheManager::_sparseBufferIds(LifeCycleId lifeCycle) const +{ + std::vector ids; + for (auto const& entry : _sparseBufferPlan(lifeCycle)) + { + ids.push_back(entry.id); + } + return ids; +} + +std::vector KvCacheManager::sparseIndexTableSpecs() const +{ + std::vector specs; + for (LifeCycleId lcId{0}; lcId < mStorage->numLifeCycles(); ++lcId) + { + auto const plan = _sparseBufferPlan(lcId); + for (size_t begin = 0; begin < plan.size();) + { + size_t end = begin; + while (end < plan.size() && plan[end].poolIndex == plan[begin].poolIndex + && plan[end].layerRow == plan[begin].layerRow) + { + ++end; + } + IndexTableSpec spec; + for (size_t i = begin; i < end; ++i) + { + spec.buffers.push_back(plan[i].id); + } + specs.push_back(std::move(spec)); + begin = end; + } + } + return specs; +} + +KvCacheManager::SparseLayerSpan KvCacheManager::_sparseLayerSpan( + std::vector const& buffers, char const* api) const +{ + TLLM_CHECK_WITH_INFO(!buffers.empty(), std::string(api) + ": no buffers"); + + LayerId const layerId = buffers.front().layerId; + for (auto const& id : buffers) + { + TLLM_CHECK_WITH_INFO(id.layerId == layerId, std::string(api) + ": every buffer must belong to one layer"); + } + + SparseLayerSpan span; + span.lifeCycle = getLayerGroupId(layerId); + auto const all = _sparseBufferIds(span.lifeCycle); + // Buffers are grouped by layer, so this layer's are a contiguous run. + auto const first + = std::find_if(all.begin(), all.end(), [layerId](BufferId const& id) { return id.layerId == layerId; }); + auto const last = std::find_if(first, all.end(), [layerId](BufferId const& id) { return id.layerId != layerId; }); + TLLM_CHECK_WITH_INFO(std::equal(first, last, buffers.begin(), buffers.end()), + std::string(api) + ": buffers must be exactly this layer's sparse buffers, in the cache's order"); + + span.bufferBase = static_cast(std::distance(all.begin(), first)); + span.numBuffers = static_cast(std::distance(first, last)); + // The metadata row is the layer's index among those carrying a cache, which is the + // number of distinct layers before it in the same walk. + for (auto it = all.begin(); it != first; ++it) + { + if (it == all.begin() || it->layerId != std::prev(it)->layerId) + { + ++span.layerInGroup; + } + } + span.selectionSlot = _selectionSlotOf(layerId); + return span; +} + void KvCacheManager::_checkNoLivingKvCaches(char const* api) const { TLLM_CHECK_WITH_INFO(mLivingKvCaches.empty(), @@ -298,6 +557,23 @@ MemAddress KvCacheManager::getMemPoolBaseAddress( return addr; } +bool KvCacheManager::isSparse(LayerId layerId, DataRole role) const +{ + auto const& attr = mStorage->getBufferAttr(layerId, role); + auto const pgIdx = mStorage->getPoolGroupIndex(kHotLevel, attr.lifeCycleId); + return mStorage->poolMedium(kHotLevel, pgIdx, attr.poolIndex) == CacheTier::HOST_MEM; +} + +size_t KvCacheManager::getNumBytes(CacheLevel level, CacheTier medium) const +{ + return mStorage->getNumBytes(level, medium); +} + +std::optional KvCacheManager::sparseBackingMedium() const +{ + return mStorage->sparseBackingMedium(); +} + int KvCacheManager::getPageStride(LayerId layerId, DataRole role) const { auto const& attr = mStorage->getBufferAttr(layerId, role); @@ -323,13 +599,6 @@ int KvCacheManager::getPageIndexScale(LayerId layerId, DataRole role) const return mStorage->mSlotToPageIndices.at(attr.lifeCycleId).at(attr.poolIndex); } -bool KvCacheManager::isSparse(LayerId layerId, DataRole role) const -{ - auto const& attr = mStorage->getBufferAttr(layerId, role); - auto const* attn = std::get_if(&mLifeCycles[attr.lifeCycleId]); - return attn && attn->isSparse; -} - PageIndexConverter KvCacheManager::getPageIndexConverter(LayerId layerId, DataRole role) const { auto const& attr = mStorage->getBufferAttr(layerId, role); @@ -340,12 +609,82 @@ PageIndexConverter KvCacheManager::getPageIndexConverter(LayerId layerId, DataRo return PageIndexConverter{scale, attr.expansion, offset, scratchPages}; } -std::optional KvCacheManager::supportsIndexMode(PageIndexMode mode) const +int KvCacheManager::_scratchSlotsFor( + LifeCycleId lcId, int maxBlocksPerSeq, std::optional maxTotalInputTokens) const +{ + auto const blocks = _maxScratchBlocks(lcId, maxBlocksPerSeq, maxTotalInputTokens); + if (blocks <= 0) + { + return 0; + } + // The same fraction KvCache::_takeExcessScratchSlots converts with, so what a request + // takes and what the batch stages for it cannot disagree. It is the widest ratio any + // of the life cycle's (layer, pool) pairs reaches, and a narrower one would leave the + // widest layer's slots past the end of the row. + return mStorage->slotUtilFracMax(lcId).ceilMul(blocks); +} + +int KvCacheManager::maxScratchSlots(int maxBlocksPerSeq) const +{ + int result = 0; + for (LifeCycleId lcId{0}; lcId < numLifeCycles(); ++lcId) + { + result = std::max(result, _scratchSlotsFor(lcId, maxBlocksPerSeq)); + } + return result; +} + +int KvCacheManager::_maxScratchBlocks( + LifeCycleId lcId, int maxBlocksPerSeq, std::optional maxTotalInputTokens) const +{ + if (!mConfig.enableSwaScratchReuse() || maxBlocksPerSeq <= 0) + { + return 0; + } + auto const* attn = std::get_if(&lifeCycles().getLifeCycle(lcId)); + if (attn == nullptr || !attn->windowSize.has_value()) + { + return 0; + } + auto const tpb = tokensPerBlock(); + auto const rewind = mConfig.swaScratchReuse->maxRewindLen; + + // Capacity-limited: a request whose history is empty and whose capacity fills the page + // table. Deferred to getStaleRange rather than restated, since which blocks are stale + // turns on the sink count and on the in-flight token the window always keeps. + auto const capacity = static_cast(maxBlocksPerSeq) * tpb; + auto const stale = attn->getStaleRange(static_cast(std::max(0, capacity - rewind)), tpb); + auto bound = static_cast(stale.end.value() - stale.beg.value()); + + if (maxTotalInputTokens.has_value()) + { + // The subtraction is in tokens, before the division: converting to blocks first and + // subtracting after would lose the alignment this is exact at. The bound is tight -- + // the count is largest when the stale end falls on a block boundary, which is the + // case this quotient gives. + auto const reach = static_cast(*maxTotalInputTokens) - *attn->windowSize - rewind + 1; + bound = std::min(bound, reach / tpb); + } + return static_cast(std::max(0, bound)); +} + +std::optional KvCacheManager::supportsIndexMode(PageIndexMode mode, LifeCycleId lcId) const { switch (mode) { case PageIndexMode::PER_LAYER: return true; - case PageIndexMode::SHARED: return mConfig.enableSwaScratchReuse() ? std::optional(std::nullopt) : true; + case PageIndexMode::SHARED: + { + if (!mConfig.enableSwaScratchReuse()) + { + return true; + } + // A life cycle with no sliding window yields an empty scratch range whatever the + // watermarks, so no request can ever hold scratch slots for it. + auto const* attn = std::get_if(&lifeCycles().getLifeCycle(lcId)); + bool const mayHoldScratch = attn != nullptr && attn->windowSize.has_value(); + return mayHoldScratch ? std::optional(std::nullopt) : true; + } } return std::nullopt; } @@ -504,7 +843,8 @@ bool KvCacheManager::resize(CacheLevel level, size_t quota, bool bestEfforts) size_t KvCacheManager::getQuota(CacheLevel level) const { auto const apiLock = lockShared(); - return mStorage->mLevels.at(level).storage->totalQuota(); + auto const& storage = *mStorage->mLevels.at(level).storage; + return storage.roundedNumBytes(storage.cacheTier()); } // ---- Statistics ---------------------------------------------------------- diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.h index 98cff3e0d1cf..3d2a9ec9f36f 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.h @@ -25,11 +25,13 @@ #include "kv_cache_manager_v2/kvCache.h" #include "kv_cache_manager_v2/lifeCycleRegistry.h" #include "kv_cache_manager_v2/movingAverage.h" +#include "kv_cache_manager_v2/sparseCache.h" #include "kv_cache_manager_v2/stats.h" #include "kv_cache_manager_v2/storageManager.h" #include "kv_cache_manager_v2/utils/poison.h" #include "kv_cache_manager_v2/utils/reentrantSharedMutex.h" +#include #include #include #include @@ -61,22 +63,6 @@ struct AggregatedPageDesc std::vector buffers; // constituent buffers in offset order }; -// --------------------------------------------------------------------------- -// ScratchDesc — scratch metadata for one layer group of one sequence. -// Scratch blocks store ephemeral KV data using shared coalesced slots. -// Mirrors _kv_cache_manager.py::ScratchDesc. -// --------------------------------------------------------------------------- -struct ScratchDesc -{ - HalfOpenRange range; // block ordinal range [beg, end) - std::vector slotIds; // scratch slot IDs, length = ceil(numScratchBlocks / scale) - - explicit operator bool() const noexcept - { - return static_cast(range); - } -}; - // --------------------------------------------------------------------------- // PageIndexConverter — convert base page index → kernel page indices. // --------------------------------------------------------------------------- @@ -89,15 +75,57 @@ struct PageIndexConverter // Convert a sequence of base page indices to per-layer page indices. // indexMode: SHARED (default) or PER_LAYER. When scratch is active, must be PER_LAYER. - // scratch: optional scratch descriptor from KvCache::getScratchDesc(). + // scratchRange: block ordinals served from scratch slots, as KvCache::getSwaScratchRange() + // reports them; empty means none are. scratchSlots are their page indices, + // from KvCache::getScratchSlots(). // Mirrors _kv_cache_manager.py::PageIndexConverter.__call__. std::vector operator()(std::vector const& baseIndices, - std::optional indexMode = std::nullopt, ScratchDesc const* scratch = nullptr) const; + std::optional indexMode = std::nullopt, HalfOpenRange scratchRange = {}, + Span scratchSlots = {}) const; // Backward-compatible single-index overload. std::vector operator()(int baseIndex) const; }; +//! What one life cycle's sparse fetch path is shaped like, as _computeSparseCacheShape +//! derives it from the configuration. +struct SparseCacheShape +{ + //! Cells a request caches per buffer, counting history and everything in flight. + int maxUnits = 0; + //! Sub-slots one page divides into. One when the life cycle has no fetch path, so the + //! page and the unit coincide. + int unitsPerPage = 1; + //! Sparse buffers this life cycle has. + int buffers = 0; + //! Units a step selects per buffer. + int selection = 0; + //! Layers of this life cycle that carry a sparse cache. + int layers = 0; + //! Selections those layers resolve between them. One where each selects for itself, + //! fewer where several share one, and what the metadata's first dimension counts. + int groups = 0; +}; + +//! One sparse buffer of a life cycle, with everything the cache derives from it. +//! +//! Ordered layers-outermost, then by pool within a layer, so that both a layer's buffers and +//! a (layer, pool) run are consecutive. The fetch needs the first, because it resolves one +//! layer; the copy needs the second, because a pool fixes a cell's size, the slot stride it +//! is read from, and the unit layout, so one launch serves one pool. +struct SparseBufferPlan +{ + BufferId id; + PoolIndex poolIndex; //!< Which hot pool holds its bytes; pools partition by size. + size_t size; //!< One cell. + uint32_t pageOffset; //!< Its offset within its hot slot. + SelectionUnitLayout layout; + int32_t layerRow; //!< Metadata row: the dense index of its layer in the group. + int32_t slotInPool; //!< Its index among the life cycle's buffers sharing its pool. +}; + +class Batch; + // --------------------------------------------------------------------------- // KvCacheManager — top-level KV cache manager. // Mirrors Python's KVCacheManager. @@ -177,8 +205,19 @@ class KvCacheManager : public std::enable_shared_from_this int getPageStride(LayerId layerId, DataRole role) const; size_t getPageIndexUpperBound(LayerId layerId, DataRole role) const; - // Whether this buffer belongs to a sparse-attention lifecycle. Unknown buffers throw. - bool isSparse(LayerId layerId, DataRole role) const; + //! Whether this buffer was declared sparse, i.e. whether its pages are host memory + //! read through the sparse GPU cache rather than device memory read directly. + [[nodiscard]] bool isSparse(LayerId layerId, DataRole role) const; + + //! Bytes `level` holds in pools whose pages are `medium`. A device level's own tier and + //! its sparse backing medium are accounted separately, the quota paying only for the + //! former, so a caller counting what it pins asks for each in turn. + [[nodiscard]] size_t getNumBytes(CacheLevel level, CacheTier medium) const; + + //! The medium a sparse buffer's pages are -- host memory the GPU reads through a device + //! mapping -- or nullopt when no life cycle is sparse. Pass it to + //! getNumBytes(kHotLevel, ...) to count what sparse buffers hold. + [[nodiscard]] std::optional sparseBackingMedium() const; // Scale factor: base_page_index * scale → kernel page index. int getPageIndexScale(LayerId layerId, DataRole role) const; @@ -217,9 +256,26 @@ class KvCacheManager : public std::enable_shared_from_this return mConfig.enableSwaScratchReuse(); } - // Whether managed KV caches support the given page index mode. - // Returns true/false for a definitive answer, nullopt for per-instance check. - std::optional supportsIndexMode(PageIndexMode mode) const; + //! The most scratch slots one request can need, over every life cycle and buffer, + //! given the widest page table a row will get. Zero where nothing can hold scratch. + //! + //! Sizes scratch staging; a caller that stages slots itself asks rather than re-deriving + //! which blocks can be stale. + [[nodiscard]] int maxScratchSlots(int maxBlocksPerSeq) const; + + //! Whether managed KV caches support the given page index mode for one life cycle, + //! which is the granularity a Batch decides at. True means every request supports it; + //! nullopt means it depends on the individual KvCache, which must then be asked. Never + //! false, so an engaged optional means support. + //! + //! Only scratch-slot substitution can make one layer's converted indices differ from + //! another's in the same life cycle, so SHARED is definitive for a life cycle that can + //! never hold scratch -- anything computeScratchRange() returns empty for, which is + //! every life cycle without a sliding window. + //! + //! Configuration only: the answer never depends on what a request is currently doing, + //! which is what lets a Batch fix its mode at construction. + std::optional supportsIndexMode(PageIndexMode mode, LifeCycleId lcId) const; bool allowSeqRebasing() const noexcept { @@ -336,11 +392,76 @@ class KvCacheManager : public std::enable_shared_from_this return mConfig; } + //! The shape of a life cycle's fetch path, or a default one for a life cycle with none. + [[nodiscard]] SparseCacheShape const& _sparseCacheShape(LifeCycleId lifeCycle) const noexcept + { + static constexpr SparseCacheShape kNoFetchPath{}; + return lifeCycle < mSparseCacheShape.size() ? mSparseCacheShape[lifeCycle] : kNoFetchPath; + } + + //! Cells a request's sparse cache holds per buffer of this life cycle, counting both + //! populations: the history it caches and everything still in flight. Zero when the life + //! cycle is not sparse or its fetch path is unconfigured. + //! + //! The reservation follows from this alone, so it is a property of the manager. What it + //! leaves above the selection width is the headroom KvCache::maxInputLength() reports. + [[nodiscard]] int sparseCacheMaxUnits(LifeCycleId lifeCycle) const noexcept + { + return _sparseCacheShape(lifeCycle).maxUnits; + } + + //! Whether this life cycle fetches sparsely, which is what decides if a step carries any + //! sparse work for it. A configured path has a positive selection width and at least one + //! unit per page, so the reservation alone answers it. + [[nodiscard]] bool hasSparseFetchPath(LifeCycleId lifeCycle) const noexcept + { + return sparseCacheMaxUnits(lifeCycle) > 0; + } + + //! Whether any life cycle has a configured sparse fetch path. A property of the + //! configuration, so it is the same for every cache this manager holds, and it decides + //! whether max_input_length is required. + [[nodiscard]] bool hasSparseCache() const noexcept + { + return std::ranges::any_of(mSparseCacheShape, [](SparseCacheShape const& shape) { return shape.maxUnits > 0; }); + } + + //! Selection units carried by one buffer of this life cycle, i.e. tokens_per_block + //! divided by the unit's token count. One when the unit is a whole block. + [[nodiscard]] int sparseCacheUnitsPerPage(LifeCycleId lifeCycle) const noexcept + { + return _sparseCacheShape(lifeCycle).unitsPerPage; + } + + //! Sparse buffers this life cycle has, one per (layer, role). Cells and cold-page + //! offsets are per buffer, so this is what the copy walks. + //! + //! Because a life cycle's buffers of one size coalesce into that pool's slot, this is + //! also the cells a reserved slot yields, which is what makes the cell count and the + //! cache count the same number. + [[nodiscard]] int sparseCacheBufferCount(LifeCycleId lifeCycle) const noexcept + { + return _sparseCacheShape(lifeCycle).buffers; + } + + //! The index table specs every sparse life cycle needs, one per (layer, pool) run. + //! + //! A Batch must still be given these -- it refuses a sparse buffer it was told nothing + //! about, since a fetch with no converted indices would read the wrong pages rather than + //! fail -- but their shape follows from the pool layout, so a caller need not re-derive it. + [[nodiscard]] std::vector sparseIndexTableSpecs() const; + LifeCycleRegistry const& lifeCycles() const noexcept { return mLifeCycles; } + //! Life cycles the manager was built with; the registry never grows after construction. + LifeCycleId numLifeCycles() const noexcept + { + return mLifeCycles.size(); + } + BlockRadixTree& radixTree() noexcept { return *mRadixTree; @@ -417,8 +538,81 @@ class KvCacheManager : public std::enable_shared_from_this // White-box introspection (incl. test-only auto-tuner state mutation) reaches // private members directly rather than widening the public API. friend class KvCacheIntrospection; + //! Reaches _sparseLayerSpan to resolve a layer's buffers for its own step. + friend class Batch; private: + //! Units a step selects per buffer of this life cycle. The cache must hold a whole + //! selection on top of everything in flight, so this is what the in-flight headroom is + //! measured against. + [[nodiscard]] int _sparseCacheSelectionWidth(LifeCycleId lifeCycle) const noexcept + { + return _sparseCacheShape(lifeCycle).selection; + } + + //! Layers this life cycle has that carry a sparse cache. + //! + //! The cache metadata is per layer rather than per buffer, because every sparse buffer + //! of a layer selects the same units: one resident set and one recency order serve them + //! all, and the entry index resolved from it addresses each buffer's own cell. + [[nodiscard]] int _sparseCacheLayerCount(LifeCycleId lifeCycle) const noexcept + { + return _sparseCacheShape(lifeCycle).layers; + } + + //! Selections this life cycle's layers resolve between them. The metadata is one set + //! per selection, not per layer: layers handed the same units resolve to the same + //! residency, so keeping a set each would be keeping copies of one answer. + [[nodiscard]] int _sparseCacheGroupCount(LifeCycleId lifeCycle) const noexcept + { + return _sparseCacheShape(lifeCycle).groups; + } + + //! Where this layer's selection sits in its life cycle's metadata. Layers sharing a + //! selection_id share the slot, which is the whole of what sharing means here. + [[nodiscard]] int _selectionSlotOf(LayerId layerId) const; + + //! [layers of this life cycle] each layer's metadata row, by its dense layer index -- + //! the arena's `layerSelection`, and what the write-back's per-layer scan reads. + [[nodiscard]] std::vector const& _sparseLayerSelection(LifeCycleId lifeCycle) const; + + //! Layers resolving within one slot. One unless a selection_id is shared, which is what + //! the fused fetch entry point refuses to serve. + [[nodiscard]] int _selectionLayerCount(LifeCycleId lifeCycle, int slot) const; + + //! Sub-slots a request's sparse cache reserves per buffer of this life cycle, i.e. the + //! cell count divided by the expansion. A property of configuration, not of any one + //! request, which is what lets a batch size its fetch scratch before any member exists. + [[nodiscard]] int _sparseCacheSubSlotsPerBuffer(LifeCycleId lifeCycle) const noexcept; + + //! Bound on how far capacity can lead history, in units: whatever the cache holds + //! beyond a full selection. One definition, because the resolve's in-flight array and + //! the write-back's flush list are bounded by the same headroom. + [[nodiscard]] int _sparseCacheMaxInFlightUnits(LifeCycleId lifeCycle) const noexcept; + + //! This life cycle's sparse buffers, in the order the cache numbers them. + [[nodiscard]] std::vector _sparseBufferPlan(LifeCycleId lifeCycle) const; + + //! Just their ids, in the same order. + [[nodiscard]] std::vector _sparseBufferIds(LifeCycleId lifeCycle) const; + + //! Multiprocessors of the device this manager runs on, for sizing a copy's grid. + [[nodiscard]] int32_t _smCount() const noexcept + { + return mSmCount; + } + + //! Behind a driver call, and the sparse copy wants it per launch, so it is read once at + //! construction. + int32_t mSmCount = 0; + + //! Storage attributes of one buffer: its pool, its offset within that pool's slot, its + //! expanded size and its expansion. + [[nodiscard]] BufferAttr const& _bufferAttr(LayerId layerId, DataRole const& role) const + { + return mStorage->getBufferAttr(layerId, role); + } + // First member, so the registration covers the whole lifetime: it is taken before any state // this manager could leave behind exists, and dropped after ~KvCacheManager has run. PoisonHold mPoisonHold; @@ -426,6 +620,42 @@ class KvCacheManager : public std::enable_shared_from_this //! Guards all mutable state reachable from this manager. See the scope note above. mutable ReentrantSharedMutex mApiMutex; + //! One layer's run within its life cycle's sparse buffer list, as the sparse entry + //! points name it: they take a layer's buffers and address them as a range. + struct SparseLayerSpan + { + LifeCycleId lifeCycle{0}; + int32_t bufferBase = 0; + int32_t numBuffers = 0; + int32_t layerInGroup = 0; //!< The dense index of the layer within its life cycle. + int32_t selectionSlot = 0; //!< Its metadata row, which layers sharing a selection share. + }; + + [[nodiscard]] SparseLayerSpan _sparseLayerSpan(std::vector const& buffers, char const* api) const; + + //! The most blocks of one request that can take scratch slots in this life cycle, + //! given the widest page table a batch will give a row and, where the caller knows it, + //! the most tokens one step may ingest. + //! + //! Blocks still in the window or held as sinks are never stale and so never scratch, + //! which is what makes this smaller than the page table. + //! + //! `maxTotalInputTokens` usually bounds it far harder. Scratch is the blocks a step + //! *adds* that are *already* stale at capacity, so a step ingesting T tokens can make + //! at most `(T - window - rewind + 1) / tpb` of them -- the window and rewind come off + //! in tokens, before the division. Under chunked prefill that is the chunk rather than + //! the context, and it is zero whenever a chunk cannot outrun the window. + //! + //! Pass it only where the caller truly bounds a step; the Batch's decode-shaped + //! default would under-count a batch that then runs prefill. + [[nodiscard]] int _maxScratchBlocks( + LifeCycleId lcId, int maxBlocksPerSeq, std::optional maxTotalInputTokens = std::nullopt) const; + + //! Slots a life cycle's worst request needs: its scratch blocks at the widest + //! sub-pages-per-slot ratio any of its layers reaches. Zero where it can hold none. + [[nodiscard]] int _scratchSlotsFor( + LifeCycleId lcId, int maxBlocksPerSeq, std::optional maxTotalInputTokens = std::nullopt) const; + // Throw unless every KvCache has been closed. `api` names the caller so the message // points at the mistake rather than at whatever breaks later. void _checkNoLivingKvCaches(char const* api) const; @@ -437,6 +667,9 @@ class KvCacheManager : public std::enable_shared_from_this PeakBlockStatsByCacheLevel _currentBlockStatsByCacheLevel() const; void _resetIterationPeakNumBlocks(std::optional cacheLevel = std::nullopt); + + // Fill the sparse cache shape tables from the configured buffers. Called once at construction. + void _computeSparseCacheShape(); void _updateIterationPeakNumBlocks(); // Current per-pool-group utilization ratios for the hot and cold representations. @@ -445,6 +678,15 @@ class KvCacheManager : public std::enable_shared_from_this KVCacheManagerConfig mConfig; LifeCycleRegistry mLifeCycles; + + // All indexed by life cycle; see the sparseCache* accessors. + TypedVec mSparseCacheShape; + + //! Each sparse layer's slot in its life cycle's selection metadata, dense from zero. + std::unordered_map mSelectionSlot; + + //! The same, by life cycle and dense layer row, which is how a kernel indexes it. + TypedVec> mSparseLayerSelection; std::shared_ptr mEventSink; std::shared_ptr mStorage; std::shared_ptr mRadixTree; diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.cpp index a4c57cacea07..f32204aba3e4 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.cpp @@ -40,9 +40,7 @@ LifeCycle makeLifeCycle(LayerConfig const& layer, int tokensPerBlock) } else { - bool const isSparse = !cfg.buffers.empty() && cfg.buffers.front().isSparse; - return AttnLifeCycle::make( - cfg.slidingWindowSize, cfg.numSinkTokens, tokensPerBlock, isSparse, cfg.reusePolicy); + return AttnLifeCycle::make(cfg.slidingWindowSize, cfg.numSinkTokens, tokensPerBlock, cfg.reusePolicy); } }, layer); @@ -90,16 +88,4 @@ LifeCycleId LifeCycleRegistry::getId(LifeCycle const& lc) const return it->second; } -LifeCycleId LifeCycleRegistry::size() const noexcept -{ - check(); - return mLifeCycleList.size(); -} - -inline void LifeCycleRegistry::check() const -{ - TLLM_CHECK_DEBUG_WITH_INFO( - toSizeT(mLifeCycleList.size()) == mLifeCycleIdMap.size(), "corrupted life cycle registry"); -} - } // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.h index bdf550e6fcc1..55cebefbf206 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.h @@ -40,7 +40,6 @@ struct AttnLifeCycle { std::optional windowSize; // nullopt = no sliding window int numSinkBlocks = 0; // divUp(numSinkTokens, tokensPerBlock) - bool isSparse = false; AttentionReusePolicy reusePolicy = AttentionReusePolicy::REQUIRED; bool isPrivate() const noexcept @@ -62,8 +61,7 @@ struct AttnLifeCycle { int numBlocks = divUp(historyLength, tokensPerBlock); BlockOrdinal start{std::min(numBlocks, numSinkBlocks)}; - // Sparse selection may revisit any history block, including outside the sliding window. - if (isSparse || !windowSize.has_value()) + if (!windowSize.has_value()) return {start, start}; // `+ 1` is intentional: attention always runs for >= 1 in-flight input // token at position `historyLength`, so the live window is @@ -74,8 +72,7 @@ struct AttnLifeCycle bool operator==(AttnLifeCycle const& o) const noexcept { - return windowSize == o.windowSize && numSinkBlocks == o.numSinkBlocks && isSparse == o.isSparse - && reusePolicy == o.reusePolicy; + return windowSize == o.windowSize && numSinkBlocks == o.numSinkBlocks && reusePolicy == o.reusePolicy; } bool operator<(AttnLifeCycle const& o) const noexcept @@ -88,22 +85,18 @@ struct AttnLifeCycle { return numSinkBlocks < o.numSinkBlocks; } - if (isSparse != o.isSparse) - { - return isSparse < o.isSparse; - } return reusePolicy < o.reusePolicy; } static AttnLifeCycle make(std::optional ws, std::optional numSinkTokens, int tokensPerBlock, - bool isSparse = false, AttentionReusePolicy reusePolicy = AttentionReusePolicy::REQUIRED) + AttentionReusePolicy reusePolicy = AttentionReusePolicy::REQUIRED) { TLLM_CHECK_DEBUG(tokensPerBlock > 0); TLLM_CHECK_DEBUG(!ws.has_value() || *ws > 0); TLLM_CHECK_DEBUG(!numSinkTokens.has_value() || *numSinkTokens >= 0); TLLM_CHECK_DEBUG((!numSinkTokens.has_value() || *numSinkTokens == 0) || ws.has_value()); int sinkBlocks = divUp(numSinkTokens.value_or(0), tokensPerBlock); - return AttnLifeCycle{ws, sinkBlocks, isSparse, reusePolicy}; + return AttnLifeCycle{ws, sinkBlocks, reusePolicy}; } }; @@ -199,7 +192,11 @@ class LifeCycleRegistry LifeCycleId getId(LifeCycle const& lc) const; // Number of unique life cycles. - LifeCycleId size() const noexcept; + [[nodiscard]] LifeCycleId size() const noexcept + { + check(); + return mLifeCycleList.size(); + } // Iteration over LifeCycles in registration order. TypedVec const& getAll() const noexcept @@ -283,7 +280,16 @@ class LifeCycleRegistry } private: - void check() const; + //! Gated on the package's cached gDebug: TLLM_CHECK_DEBUG_WITH_INFO's own gate is an + //! out-of-line call into another TU, which would undo size() being inline. + void check() const + { + if (TLLM_UNLIKELY(gDebug)) + { + TLLM_CHECK_WITH_INFO( + toSizeT(mLifeCycleList.size()) == mLifeCycleIdMap.size(), "corrupted life cycle registry"); + } + } TypedVec mLifeCycleList; std::map mLifeCycleIdMap; diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.cpp index d13e1af69ee7..5a82032341fd 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.cpp @@ -23,8 +23,7 @@ #include "tensorrt_llm/common/assert.h" -#include -#include +#include namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 { @@ -133,21 +132,6 @@ SharedPageLock Page::lock(KvCache& kvCache, BeamIndex beamIndex, BlockOrdinal or return hold()->lock(kvCache, beamIndex, ordinal, lc, skipWait); } -CacheLevel Page::queryLockLevel() const -{ - if (cacheLevel == kHotLevel) - { - return kHotLevel; - } - auto const* attention = std::get_if(&manager->getLifeCycle(lifeCycle)); - if (attention != nullptr && attention->isSparse && kSparseHistoryLevel < manager->numCacheLevels() - && manager->cacheTier(kSparseHistoryLevel) == CacheTier::HOST_MEM) - { - return kSparseHistoryLevel; - } - return kHotLevel; -} - // --------------------------------------------------------------------------- // CommittedPage // --------------------------------------------------------------------------- @@ -352,12 +336,10 @@ SharedPageLock PageHolder::lock( UniqPageLock::UniqPageLock(SharedPtr h) : holder(std::move(h)) { - if (holder->page->cacheLevel != holder->page->queryLockLevel()) + if (holder->page->cacheLevel != kHotLevel) { - throw LogicError("Pages can only be locked on GPU or, for sparse attention, in level-1 host memory"); + throw LogicError("Pages can only be locked at the hot level"); } - // Preserve readiness if the first shared lock fails before it can register an owner. - finishEvents.push_back(holder->page->readyEvent); } UniqPageLock::~UniqPageLock() @@ -366,8 +348,7 @@ UniqPageLock::~UniqPageLock() [this] { Page& p = *page(); - TLLM_CHECK_DEBUG(p.cacheLevel == p.queryLockLevel() && !p.scheduledForEviction()); - TLLM_CHECK_DEBUG(mOwners.empty()); + TLLM_CHECK_DEBUG(p.cacheLevel == kHotLevel && !p.scheduledForEviction()); // Set readyEvent to the merged finish events of all readers. For committed (read-only) // pages, this means the next reader will wait for prior reads to complete, which is // unnecessary but correct. See the CommittedPage comment in page.h for rationale. @@ -403,97 +384,6 @@ void UniqPageLock::notifyFinish(CachedCudaEvent event) } } -void UniqPageLock::prepareSparseOffload(KvCache const& requestingCache) -{ - Page const& p = *page(); - auto const* attn = std::get_if(&p.manager->getLifeCycle(p.lifeCycle)); - if (!attn || !attn->isSparse || !p.hasValidSlot() - || (p.cacheLevel != kHotLevel && p.cacheLevel != kSparseHistoryLevel) - || p.manager->numCacheLevels() <= kSparseHistoryLevel - || p.manager->cacheTier(kSparseHistoryLevel) != CacheTier::HOST_MEM) - { - throw LogicError("Offload requires a locked sparse attention page on GPU or in host history"); - } - if (p.isCommitted() && static_cast(p).numTokensInBlock != requestingCache.tokensPerBlock()) - { - throw LogicError("Cannot offload a partial committed page"); - } - bool requestingOwner = false; - for (auto const& owner : mOwners) - { - if (!owner.kvCache->isActive() || owner.lifeCycle != p.lifeCycle || owner.ordinal < BlockOrdinal{0} - || owner.ordinal >= BlockOrdinal{owner.kvCache->historyLength() / owner.kvCache->tokensPerBlock()}) - { - throw LogicError("Cannot offload a page outside an owner's complete history"); - } - requestingOwner |= owner.kvCache == &requestingCache; - } - if (!requestingOwner) - { - throw LogicError("The offloading request must own a lock on the page"); - } - finishEvents.reserve(1); -} - -void UniqPageLock::prepareSparsePromotion() -{ - Page const& p = *page(); - if (!p.hasValidSlot() || p.cacheLevel != kSparseHistoryLevel || p.queryLockLevel() != kSparseHistoryLevel - || mOwners.empty()) - { - throw LogicError("Promotion requires a locked sparse host page"); - } - for (auto const& owner : mOwners) - { - if (!owner.kvCache->isActive() || owner.lifeCycle != p.lifeCycle) - { - throw LogicError("Promotion requires active owners of the same sparse page"); - } - } - finishEvents.reserve(1); -} - -void UniqPageLock::recordMigrationEvent(CachedCudaEvent const& event) -{ - // The copy stream already waited for every event being replaced here. - page()->readyEvent = event; - finishEvents.clear(); - finishEvents.push_back(event); - // A rejected copy can change readiness without changing the source slot. - for (auto const& owner : mOwners) - owner.kvCache->onPageStorageChanged(); -} - -Slot UniqPageLock::moveToCacheLevel(CacheLevel destination, Slot&& slot) -{ - Page& p = *page(); - TLLM_CHECK_DEBUG(!p.scheduledForEviction()); - TLLM_CHECK_DEBUG((p.cacheLevel == kHotLevel && destination == kSparseHistoryLevel) - || (p.cacheLevel == kSparseHistoryLevel && destination == kHotLevel)); - Slot source = p.exchangeSlot(std::move(slot)); - p.cacheLevel = destination; - for (auto const& owner : mOwners) - { - int const old = owner.kvCache->updateBasePageIndex( - owner.beamIndex, owner.ordinal, owner.lifeCycle, slotIdToPageIndexValue(p.slotId())); - TLLM_CHECK_DEBUG(old == slotIdToPageIndexValue(source.slotId())); - if (destination == kHotLevel && owner.kvCache->mIsDecoding) - { - // Retry even without history growth once the prefill owner releases its GPU requirement. - owner.kvCache->mHasDeferredSparseOffload = true; - } - owner.kvCache->onPageStorageChanged(); - } - return source; -} - -void UniqPageLock::removeOwner(LockOwner const& owner) -{ - auto const it = std::find(mOwners.begin(), mOwners.end(), owner); - TLLM_CHECK_DEBUG(it != mOwners.end()); - mOwners.erase(it); -} - SharedPtr const& UniqPageLock::page() const { TLLM_CHECK_DEBUG(holder && holder->page); @@ -516,14 +406,9 @@ SharedPageLock::SharedPageLock(SharedPtr ul, KvCache& kvCache, Bea , mUser{&kvCache, beamIndex, ordinal, lc} { if (!skipWait) - { page()->readyEvent.waitInStream(reinterpret_cast(kvCache.cudaStream())); - } - mUniqLock->mOwners.push_back(mUser); - auto rollbackOwner = FuncGuard([this]() { mUniqLock->removeOwner(mUser); }); acquirePageIndex(); - rollbackOwner.cancel(); } SharedPageLock::~SharedPageLock() @@ -564,7 +449,6 @@ SharedPtr SharedPageLock::unlock() mUniqLock->notifyFinish(mUser.kvCache->finishEvent()); releasePageIndex(); - mUniqLock->removeOwner(mUser); auto p = page(); // copy shared_ptr before reset mUniqLock.reset(); return p; @@ -617,29 +501,20 @@ std::vector batchedLockPages(KvCache& kvCache, std::vector destinations; - TypedVec>> pagesByLevel(storeMgr->numCacheLevels()); - std::vector> lockedHostPages; + // + // Every page is locked at the hot level, whatever its pool group is made of. A sparse + // life cycle's hot pages are host memory, but that is a property of the allocation + // rather than of the level, so locking does not see it. + std::unordered_set seen; + std::vector> pages; + pages.reserve(targets.size()); for (auto const& target : targets) { auto const& page = target.page; - CacheLevel const level = target.cacheLevel; - if ((level != kHotLevel && level != page->queryLockLevel()) || page->cacheLevel < level) - { - throw LogicError("Invalid destination for page locking; offload requires a separate handoff"); - } - auto const [it, inserted] = destinations.emplace(page.get(), level); - if (!inserted && it->second != level) - throw LogicError("Conflicting lock levels for a shared page"); - if (inserted) - { - pagesByLevel[level].push_back(page); - if (page->status() == PageStatus::LOCKED && page->cacheLevel != level) - { - TLLM_CHECK_DEBUG(level == kHotLevel && page->cacheLevel == kSparseHistoryLevel); - lockedHostPages.push_back(page); - } - } + if (page->status() == PageStatus::LOCKED && page->cacheLevel != kHotLevel) + throw LogicError("Cannot migrate a page locked by another owner"); + if (seen.insert(page.get()).second) + pages.push_back(page); } // Protect every destination group while any group is allocating. On failure, @@ -647,36 +522,27 @@ std::vector batchedLockPages(KvCache& kvCache, std::vectorscheduledForEviction() && storeMgr->isEvictable(*page)) - storeMgr->scheduleForEviction(*page); + for (auto const& page : pages) + if (!page->scheduledForEviction() && storeMgr->isEvictable(*page)) + storeMgr->scheduleForEviction(*page); }); - for (auto const& pages : pagesByLevel) - for (auto const& page : pages) - if (page->scheduledForEviction()) - storeMgr->excludeFromEviction(*page); + for (auto const& page : pages) + if (page->scheduledForEviction()) + storeMgr->excludeFromEviction(*page); MigrationRecorder const migrationRecorder = [&kvCache](std::vector> const& pages, std::vector const& slots, CacheLevel srcLevel, CacheLevel dstLevel) { kvCache._recordMigratedSlots(pages, slots, srcLevel, dstLevel); }; DropRecorder const dropRecorder = [&kvCache](std::vector> const& pages, CacheLevel cacheLevel) { kvCache._recordDroppedPages(pages, cacheLevel); }; - for (CacheLevel level{0}; level < pagesByLevel.size(); ++level) + if (!pages.empty()) { - auto const& pages = pagesByLevel[level]; - if (pages.empty()) - continue; - TypedVec requirements(storeMgr->numPoolGroups(level), 0); + TypedVec requirements(storeMgr->numPoolGroups(kHotLevel), 0); for (auto const& page : pages) - if (page->cacheLevel != level) - ++requirements[storeMgr->getPoolGroupIndex(level, page->lifeCycle)]; - storeMgr->prepareFreeSlots(level, requirements, migrationRecorder, dropRecorder); - if (level == kHotLevel && !lockedHostPages.empty()) - { - storeMgr->promoteSparsePages(kvCache.cudaStream(), lockedHostPages, migrationRecorder, dropRecorder); - } - storeMgr->batchedMigrate(level, pages, migrationRecorder); + if (page->cacheLevel != kHotLevel) + ++requirements[storeMgr->getPoolGroupIndex(kHotLevel, page->lifeCycle)]; + storeMgr->prepareFreeSlots(kHotLevel, requirements, migrationRecorder, dropRecorder); + storeMgr->batchedMigrate(kHotLevel, pages, migrationRecorder); } // Wait for all ready events on KvCache's stream (deduplicated). diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.h index b10075de042f..4fa34dc549a2 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.h @@ -87,10 +87,6 @@ class Page : public Slot, public EnableSharedFromThis // Prevent the page from being dropped (returns/creates a PageHolder). SharedPtr hold(); - //! Return the lock destination for read-only use, preserving pages already at the hot GPU level. - //! Cold sparse pages use host memory; callers must force writable pages to GPU. - [[nodiscard]] CacheLevel queryLockLevel() const; - // Acquire a shared lock at the current level. The caller handles any migration. // skip_wait: caller guarantees the page is ready on kvCache's stream. SharedPageLock lock( @@ -198,17 +194,6 @@ class PageHolder : public EnableSharedFromThis WeakPtr uniqLock; // non-null → LOCKED }; -//! Identifies one live SharedPageLock independently of the lock object's address. -struct LockOwner -{ - KvCache* kvCache; - BeamIndex beamIndex; - BlockOrdinal ordinal; - LifeCycleId lifeCycle; - - bool operator==(LockOwner const&) const = default; -}; - // --------------------------------------------------------------------------- // UniqPageLock — locks a page to prevent eviction (LOCKED status). // Owns finish events from all SharedPageLocks it issued. @@ -232,31 +217,19 @@ class UniqPageLock : public EnableSharedFromThis // Append a finish event, merging when count exceeds 32 to prevent unbounded growth. void notifyFinish(CachedCudaEvent event); - //! Validate complete sparse history for every owner and prepare non-allocating completion updates. - void prepareSparseOffload(KvCache const& requestingCache); - - //! Validate a locked host page and prepare non-allocating completion updates for shared promotion. - void prepareSparsePromotion(); - - //! Record a copy ordered after page readiness, finished readers, and all live owners' prior work. - void recordMigrationEvent(CachedCudaEvent const& event); - - //! Publish a GPU/host handoff to every owner and return the fenced source slot. Caller holds the API lock. - [[nodiscard]] Slot moveToCacheLevel(CacheLevel destination, Slot&& slot); - - std::vector const& owners() const noexcept - { - return mOwners; - } - SharedPtr holder; std::vector finishEvents; +}; -private: - friend class SharedPageLock; - void removeOwner(LockOwner const& owner); - - std::vector mOwners; +// --------------------------------------------------------------------------- +// LockOwner — identifies who holds a SharedPageLock. +// --------------------------------------------------------------------------- +struct LockOwner +{ + KvCache* kvCache; + BeamIndex beamIndex; + BlockOrdinal ordinal; + LifeCycleId lifeCycle; }; // --------------------------------------------------------------------------- @@ -297,7 +270,7 @@ class SharedPageLock }; // --------------------------------------------------------------------------- -// BatchedLockTarget — page, owner, and intended storage level for a lock. +// BatchedLockTarget — one page to lock, named by beam, ordinal and life cycle. // --------------------------------------------------------------------------- struct BatchedLockTarget { @@ -305,7 +278,6 @@ struct BatchedLockTarget BeamIndex beamIndex; BlockOrdinal ordinal; LifeCycleId lifeCycle; - CacheLevel cacheLevel = kHotLevel; }; // --------------------------------------------------------------------------- @@ -318,7 +290,7 @@ std::vector batchedLockPages(KvCache& kvCache, std::vector= 0); - count = std::min(count, std::max(int64_t{0}, maxTokens)); + count = std::min(count, maxTokens); } [[nodiscard]] CountsByLevel const& cachedTokensByLevel() const noexcept @@ -337,7 +337,7 @@ class PendingStats static PendingStatsDelta allocationDelta( PendingAllocationSegment const& segment, BlockOrdinal blockBegin, BlockOrdinal blockEnd) { - int64_t const numBlocks = static_cast(std::max(0, blockEnd - blockBegin)) * segment.beamWidth; + int64_t const numBlocks = static_cast(blockEnd - blockBegin) * segment.beamWidth; int64_t const missedBlocks = segment.countAsMissed ? numBlocks : 0; int64_t const connectorReusedBlocks = segment.countAsConnectorReused ? numBlocks : 0; PendingStatsDelta delta; diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/sparseCache.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/sparseCache.cpp new file mode 100644 index 000000000000..b92ea80706e8 --- /dev/null +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/sparseCache.cpp @@ -0,0 +1,243 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "kv_cache_manager_v2/sparseCache.h" + +#include "kv_cache_manager_v2/exceptions.h" +#include "kv_cache_manager_v2/utils/math.h" + +#include +#include + +namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 +{ +namespace +{ + +constexpr size_t kAlignment = 256; + +size_t alignedBytes(size_t bytes) noexcept +{ + return roundUp(bytes, kAlignment); +} + +} // namespace + +namespace +{ + +//! Per-lease metadata, in carve order. Arrays shared by every lease of a life cycle live +//! on the arena instead. +struct LeaseLayout +{ + size_t ordinal = 0; + size_t lru = 0; + + [[nodiscard]] size_t total() const noexcept + { + return ordinal + lru; + } +}; + +LeaseLayout leaseLayout(SparseCacheInit const& shape) +{ + auto const entries = static_cast(shape.numSelections) * shape.entriesPerSelection; + LeaseLayout out; + out.ordinal = alignedBytes(sizeof(int32_t) * entries); + out.lru = alignedBytes(sizeof(int32_t) * entries); + return out; +} + +} // namespace + +// --------------------------------------------------------------------------- +// SparseCellArena +// --------------------------------------------------------------------------- + +SparseCellArena::SparseCellArena(SparseCacheInit const& shape, int maxLeases) + : mShape(shape) +{ + CUstream const stream = nullptr; + TLLM_CHECK_WITH_INFO(shape.numLayers > 0 && shape.numSelections > 0 && shape.numBuffers > 0 + && shape.entriesPerSelection > 0 && shape.cellsPerBuffer > 0, + "SparseCellArena requires a positive layer, selection, buffer, entry and cell count"); + TLLM_CHECK_WITH_INFO(shape.numSelections <= shape.numLayers, + "SparseCellArena: a selection is shared by layers, so there cannot be more of them than layers"); + TLLM_CHECK_WITH_INFO( + !shape.pools.empty() && maxLeases > 0, "SparseCellArena requires at least one pool and a positive lease count"); + TLLM_CHECK_WITH_INFO(shape.bufferLayer.size() == static_cast(shape.numBuffers), + "SparseCellArena requires one layer index per buffer"); + TLLM_CHECK_WITH_INFO(shape.bufferSlotInPool.size() == static_cast(shape.numBuffers), + "SparseCellArena requires one pool slot per buffer"); + TLLM_CHECK_WITH_INFO(shape.layerSelection.size() == static_cast(shape.numLayers), + "SparseCellArena requires one selection index per layer"); + for (size_t i = 0; i < shape.bufferLayer.size(); ++i) + { + TLLM_CHECK_WITH_INFO(shape.bufferLayer[i] >= 0 && shape.bufferLayer[i] < shape.numLayers, + "SparseCellArena: a buffer's layer index is out of range"); + TLLM_CHECK_WITH_INFO(i == 0 || shape.bufferLayer[i] >= shape.bufferLayer[i - 1], + "SparseCellArena: buffers must be grouped by layer"); + } + for (auto const selection : shape.layerSelection) + { + TLLM_CHECK_WITH_INFO(selection >= 0 && selection < shape.numSelections, + "SparseCellArena: a layer's selection index is out of range"); + } + for (auto const& [cellBytes, buffers] : shape.pools) + { + // {0, 0} is a pool index no sparse buffer of this life cycle uses: it holds its place + // so that a storage pool index indexes this arena directly, and carries no cells. Any + // other combination describes a pool that cannot be addressed. + TLLM_CHECK_WITH_INFO((cellBytes == 0) == (buffers == 0), + "SparseCellArena: a pool needs both a cell size and buffers drawing on it, or neither"); + TLLM_CHECK_WITH_INFO(buffers >= 0, "SparseCellArena: a pool cannot have a negative buffer count"); + } + + auto const layout = leaseLayout(mShape); + mMaxLeases = maxLeases; + mMetadataBytes = layout.total(); + + // One region per pool, each aligned so that a buffer's cell indices stay relative to its own + // pool's base. + std::vector regionOffset; + size_t cellBytes = 0; + for (auto const& [bytes, buffers] : mShape.pools) + { + regionOffset.push_back(cellBytes); + cellBytes += alignedBytes(static_cast(maxLeases) * buffers * mShape.cellsPerBuffer * bytes); + } + + auto const entries = static_cast(mShape.numSelections) * mShape.entriesPerSelection; + size_t const layerBytes = alignedBytes(sizeof(int32_t) * static_cast(mShape.numBuffers)); + size_t const slotBytes = alignedBytes(sizeof(int32_t) * static_cast(mShape.numBuffers)); + size_t const selectionBytes = alignedBytes(sizeof(int32_t) * static_cast(mShape.numLayers)); + size_t const identityBytes = alignedBytes(sizeof(int32_t) * entries); + size_t const sharedBytes = layerBytes + slotBytes + selectionBytes + identityBytes; + + mNumBytes = cellBytes + sharedBytes + static_cast(maxLeases) * mMetadataBytes; + + CUdeviceptr base = 0; + cuCheck(cuMemAlloc(&base, mNumBytes)); + mMemory = CudaUniqPtr<>{reinterpret_cast(base)}; + + for (auto offset : regionOffset) + { + mPoolBase.push_back(base + offset); + } + auto* shared = reinterpret_cast(base) + cellBytes; + mSharedLayer = reinterpret_cast(shared); + mSharedSlotInPool = reinterpret_cast(shared + layerBytes); + mSharedLayerSelection = reinterpret_cast(shared + layerBytes + slotBytes); + mIdentity = reinterpret_cast(shared + layerBytes + slotBytes + selectionBytes); + mMetadataBase = shared + sharedBytes; + + cuCheck(cuMemcpyHtoDAsync(reinterpret_cast(mSharedLayer), mShape.bufferLayer.data(), + sizeof(int32_t) * mShape.bufferLayer.size(), stream)); + cuCheck(cuMemcpyHtoDAsync(reinterpret_cast(mSharedSlotInPool), mShape.bufferSlotInPool.data(), + sizeof(int32_t) * mShape.bufferSlotInPool.size(), stream)); + cuCheck(cuMemcpyHtoDAsync(reinterpret_cast(mSharedLayerSelection), mShape.layerSelection.data(), + sizeof(int32_t) * mShape.layerSelection.size(), stream)); + + // The recency order every lease starts from: a per-selection identity permutation. + std::vector identity(entries); + for (int32_t selection = 0; selection < mShape.numSelections; ++selection) + { + auto const begin = identity.begin() + static_cast(selection) * mShape.entriesPerSelection; + std::iota(begin, begin + mShape.entriesPerSelection, 0); + } + cuCheck(cuMemcpyHtoDAsync( + reinterpret_cast(mIdentity), identity.data(), sizeof(int32_t) * identity.size(), stream)); + + // An unused cell is a clean one, so the arena is born that way and a row hands its + // lease back clean. Ordered after the identity staged above, which this copies from. + for (int slot = 0; slot < mMaxLeases; ++slot) + { + initialiseLease(SparseLeaseIndex{slot}, stream); + } + + // The arena has no stream of its own, so the staging above is ordered on the null + // stream only. `identity` also dies with this scope. Both want it landed here. + cuCheck(cuStreamSynchronize(stream)); +} + +void SparseCellArena::initialiseLease(SparseLeaseIndex slot, CUstream stream) +{ + if (empty()) + { + return; + } + TLLM_CHECK_WITH_INFO(slot != kNoSparseLease && slot < mMaxLeases, + "SparseCellArena::initialiseLease: slot %d is outside [0, %d)", slot.value(), mMaxLeases); + int const index = slot.value(); + + auto const layout = leaseLayout(mShape); + auto* bytes = mMetadataBase + static_cast(index) * mMetadataBytes; + size_t carved = 0; + auto carve = [&](size_t size) + { + auto* ptr = reinterpret_cast(bytes + carved); + carved += size; + return ptr; + }; + + auto* const residentOrdinal = carve(layout.ordinal); + auto* const lruOrder = carve(layout.lru); + + auto const entries = static_cast(mShape.numSelections) * static_cast(mShape.entriesPerSelection); + cuCheck(cuMemsetD32Async(reinterpret_cast(residentOrdinal), + static_cast(kBadBlockOrdinal.value()), entries, stream)); + cuCheck(cuMemcpyDtoDAsync(reinterpret_cast(lruOrder), reinterpret_cast(mIdentity), + sizeof(int32_t) * entries, stream)); +} + +SparseArenaView SparseCellArena::view() const noexcept +{ + SparseArenaView out; + out.metadataBase = mMetadataBase; + out.leaseStride = mMetadataBytes; + // Both regions are carved from one lease block, the ordinals first; see leaseLayout(). + out.lruOffset = static_cast(leaseLayout(mShape).ordinal / sizeof(int32_t)); + out.bufferSlotInPool = mSharedSlotInPool; + out.bufferLayer = mSharedLayer; + out.layerSelection = mSharedLayerSelection; + out.numLayers = mShape.numLayers; + out.numSelections = mShape.numSelections; + out.entriesPerSelection = mShape.entriesPerSelection; + return out; +} + +std::vector SparseCellArena::peekResidentOrdinals(SparseLeaseIndex slot) const +{ + CUstream const stream = nullptr; + if (empty() || slot == kNoSparseLease) + { + return {}; + } + // A resolve writes this metadata on the Batch's stream, which is non-blocking and so + // unordered against the one below. Only a device-wide sync covers it; nothing here + // knows which stream last touched the slot. + cuCheck(cuCtxSynchronize()); + auto const arena = view(); + std::vector out(arena.numEntries()); + cuCheck(cuMemcpyDtoHAsync(out.data(), + reinterpret_cast(arena.metadataBase + static_cast(slot.value()) * arena.leaseStride), + sizeof(int32_t) * out.size(), stream)); + cuCheck(cuStreamSynchronize(stream)); + return out; +} + +} // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/sparseCache.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/sparseCache.h new file mode 100644 index 000000000000..ddd95d21e45b --- /dev/null +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/sparseCache.h @@ -0,0 +1,224 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "kv_cache_manager_v2/stagingBuffer.h" +#include "kv_cache_manager_v2/storage/config.h" +#include "kv_cache_manager_v2/utils/typedIndex.h" + +#include +#include +#include +#include +#include + +namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 +{ + +// --------------------------------------------------------------------------- +// SparseLeaseIndex and SparseArenaView — the device-resident state of one request's sparse +// history cache for one life cycle, as the fetch kernel sees it: a slot per row, and the +// geometry that slot is read within. +// +// A reserved slot is a row of buffer-sized cells, and any layer's unit may occupy any cell; +// a cell holds `unitsPerPage` units when a selection unit is smaller than a buffer. An entry +// is identified by (layer, entry index) and resolved to a cell arithmetically, so nothing +// ties a layer to a fixed position. +// +// The metadata is laid out [selection][entry]. It is per SELECTION, not per buffer or even +// per layer: everything selecting the same units resolves to the same residency, so one +// resident set and one recency order serve them all, and the entry index the resolve +// produces addresses each buffer's own cell. A layer selecting for itself is the singleton +// case of that; layers sharing a selection_id share the row. +// --------------------------------------------------------------------------- +//! The arena slot a row's member holds, or kNoSparseLease when the row carries none. +//! +//! This is the whole of a row's published cache state: everything else a kernel needs is +//! either arithmetic from this index -- a lease's cells and its metadata both sit at a +//! fixed stride from the arena's bases, see "Why the cache is a flat array of cells" in +//! SPARSE_ATTENTION_DESIGN.md -- or a constant of the life cycle, which SparseArenaView +//! carries once per launch rather than once per row. +using SparseLeaseIndex = StrongIndex; +inline constexpr SparseLeaseIndex kNoSparseLease{}; + +//! One life cycle's arena, as a kernel addresses it. Fixed for the arena's lifetime, so a +//! launch takes it by value and no row carries a copy. +struct SparseArenaView +{ + //! Lease i's metadata begins at `metadataBase + i * leaseStride`: first its resident + //! unit ordinal per entry, kBadBlockOrdinal where the entry is free, then at + //! `lruOffset` its entry indices in recency order, most-recently-used first. + //! + //! The ordinals are dense rather than keyed: a fetch builds a hash over them in shared + //! memory to answer "is this ordinal resident, and where". + std::byte* metadataBase; + uint64_t leaseStride; + int32_t lruOffset; + + //! [numBuffers] each buffer's index among those sharing its pool: where its cells sit + //! inside that pool's arena region. + int32_t const* bufferSlotInPool; + + //! [numBuffers] the layer each buffer belongs to, i.e. which work list carries it. + //! Non-decreasing: the buffers of a layer are consecutive, which is what lets a + //! one-layer call name its buffers as a range rather than a set. + int32_t const* bufferLayer; + + //! [numLayers] the metadata row each layer resolves within. The identity where every + //! layer selects for itself; layers sharing a selection share an entry here, which is + //! the whole of what sharing one means to a kernel. + int32_t const* layerSelection; + + //! Layers in this life cycle. Not the metadata's first dimension -- it is the extent of + //! the write-back's grid and of its per-(row, layer) counts. + int32_t numLayers; + + //! Selections those layers resolve between them: the metadata's first dimension. + int32_t numSelections; + int32_t entriesPerSelection; + + [[nodiscard]] constexpr int32_t numEntries() const noexcept + { + return numSelections * entriesPerSelection; + } +}; + +// --------------------------------------------------------------------------- +// SparseCacheInit — everything a request's cache state is built from, all of it +// fixed for the life of the reservation. +// --------------------------------------------------------------------------- +struct SparseCacheInit +{ + int32_t numLayers; + //! Metadata rows those layers share between them, each the residency of one selection. + int32_t numSelections; + int32_t numBuffers; + int32_t entriesPerSelection; + //! Cells one buffer gets, i.e. entriesPerSelection / unitsPerPage. A lease is this many + //! cells per buffer, laid out buffer-major. + int32_t cellsPerBuffer; + //! [numBuffers] the layer each buffer belongs to, non-decreasing. + std::vector bufferLayer; + //! [numLayers] the metadata row each layer resolves within, every row named by at + //! least one layer. + std::vector layerSelection; + //! [numBuffers] each buffer's index among those sharing its pool. + std::vector bufferSlotInPool; + //! One entry per pool of the life cycle: cell size, and buffers drawing on it. + std::vector> pools; +}; + +class SparseCellArena; + +// --------------------------------------------------------------------------- +// SparseCellArena — one life cycle's GPU cells, and the metadata that describes them. +// +// A dedicated allocation rather than storage slots, so the life cycle's hot pool group can +// be host memory without the cache following it there. Every request of a life cycle needs +// the same number of cells, so this is a fixed-size-block free list. +// +// One block holds the cells of every lease followed by their metadata. Cells come first so +// that a cell index is relative to the arena's base, which is what the fetch reports and +// what get_sparse_cache_base_address returns for a sparse buffer. +// --------------------------------------------------------------------------- +class SparseCellArena +{ +public: + SparseCellArena() = default; + + //! @param shape per-request metadata and cell geometry, fixed by configuration, + //! including one entry per pool of the life cycle. + //! @param maxLeases concurrent requests this life cycle can serve. + //! + //! One region per pool, mirroring the hot pool group: pools partition the life cycle's + //! buffers by size, and a cell is one buffer, so cells cannot be uniform across them. + SparseCellArena(SparseCacheInit const& shape, int maxLeases); + + ~SparseCellArena() = default; + SparseCellArena(SparseCellArena&&) noexcept = default; + SparseCellArena& operator=(SparseCellArena&&) noexcept = default; + SparseCellArena(SparseCellArena const&) = delete; + SparseCellArena& operator=(SparseCellArena const&) = delete; + + [[nodiscard]] bool empty() const noexcept + { + return mMemory == nullptr; + } + + //! Base of one pool's cell region: the base the fetch's indices for a buffer are relative to. + [[nodiscard]] CUdeviceptr base(PoolIndex pool) const + { + return mPoolBase.at(pool); + } + + [[nodiscard]] uint64_t bufferBytes(PoolIndex pool) const + { + return mShape.pools.at(static_cast(pool.value())).first; + } + + //! Cells one lease takes in this pool's region: its buffers times the cells each gets. + [[nodiscard]] int cellsPerLease(PoolIndex pool) const + { + return mShape.pools.at(static_cast(pool.value())).second * mShape.cellsPerBuffer; + } + + //! Device bytes this arena holds: every pool's cell region, the state its leases share, + //! and one metadata block per lease -- the whole allocation. + [[nodiscard]] size_t numBytes() const noexcept + { + return mNumBytes; + } + + //! Clear `slot`'s metadata on `stream`, which is what taking it means here: every + //! ordinal becomes kBadBlockOrdinal and the recency order the identity permutation. + //! The slot is the Batch row's own, and one row's slot serves every life cycle. + void initialiseLease(SparseLeaseIndex slot, CUstream stream); + + //! Copy a slot's resident ordinals back to the host, for debugging and introspection + //! only. Not on any serving path: it synchronises the device to pick up whatever a + //! resolve left on the Batch's stream, which no stream of the caller's orders it + //! against, and so reports a moment rather than a state the caller can act on. + [[nodiscard]] std::vector peekResidentOrdinals(SparseLeaseIndex slot) const; + + //! How a kernel addresses this arena's leases. Constant once constructed, so a launch + //! carries one of these rather than one descriptor per row. + [[nodiscard]] SparseArenaView view() const noexcept; + + [[nodiscard]] SparseCacheInit const& shape() const noexcept + { + return mShape; + } + +private: + CudaUniqPtr<> mMemory; + SparseCacheInit mShape; + //! Slots this arena was sized for, which every lease index must fall inside. + int mMaxLeases = 0; + size_t mMetadataBytes = 0; + size_t mNumBytes = 0; + TypedVec mPoolBase; + std::byte* mMetadataBase = nullptr; + + //! Shared by every lease of the life cycle, so a resume copies none of them. + int32_t* mSharedLayer = nullptr; + int32_t* mSharedSlotInPool = nullptr; + int32_t* mSharedLayerSelection = nullptr; + int32_t* mIdentity = nullptr; +}; + +} // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/sparseCacheKernels.cu b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/sparseCacheKernels.cu new file mode 100644 index 000000000000..45700f319f91 --- /dev/null +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/sparseCacheKernels.cu @@ -0,0 +1,1075 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "sparseCacheKernels.h" + +#include "exceptions.h" + +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/common/envUtils.h" + +#include + +#include +#include + +#include +#include + +namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 +{ + +namespace +{ + +//! Enforce a precondition that correct use cannot violate. A debug build names the condition. +#define KVCM2_SPARSE_REQUIRE(cond) \ + do \ + { \ + if (!(cond)) \ + { \ + assert(!"KVCM2 sparse invariant violated: " #cond); \ + __trap(); \ + } \ + } while (0) + +constexpr int kThreads = 256; +constexpr uint32_t kWarpSize = 32; +constexpr int32_t kWarpsPerCopyCta = kThreads / static_cast(kWarpSize); +constexpr int32_t kEmptyKey = -1; +//! kBadPageIndex as a plain int32, since device code cannot use the strong typedef. +constexpr int32_t kBadPageIndexValue = kBadPageIndex.value(); +constexpr int32_t kFreeOrdinal = kBadBlockOrdinal.value(); + +//! A lease's metadata block: its resident ordinals, with its recency order at `lruOffset`. +//! +//! Every lease of a life cycle is the same size, so the block is arithmetic from the +//! arena's base -- the same addressing cells already use, which is why a row publishes an +//! index rather than the pointers it would imply. +__device__ inline int32_t* leaseMetadata(SparseArenaView const& arena, SparseLeaseIndex lease) +{ + return reinterpret_cast(arena.metadataBase + static_cast(lease.value()) * arena.leaseStride); +} + +//! Recency classes, ordered: a hit outranks a unit fetched this step, which outranks an +//! entry the step did not touch. This is what makes a repeatedly selected unit survive a +//! one-off selection. +constexpr uint8_t kClassHit = 0; +constexpr uint8_t kClassFetched = 1; +constexpr uint8_t kClassUntouched = 2; +// Untouched in-flight entries are pinned; a hit is protected by its recency class. +constexpr uint8_t kPinned = 4; + +using BlockScan = cub::BlockScan; + +__host__ __device__ inline size_t alignUp16(size_t bytes) noexcept +{ + return (bytes + 15U) & ~static_cast(15U); +} + +//! Byte offsets of the shared-memory regions one CTA carves. Computed identically on the +//! host, which sizes the launch, and in the kernel, which carves it. Every region is +//! 16-byte aligned: cls is a byte array, so an unaligned int32 region after it faults. +struct SmemLayout +{ + size_t keys = 0; + size_t vals = 0; + size_t cls = 0; + size_t scratch = 0; + size_t order = 0; + size_t lru = 0; + size_t missIdx = 0; + size_t missOrd = 0; + size_t total = 0; +}; + +//! @param required Length of the virtual list the probe walks: in-flight units first, +//! then the selection. +__host__ __device__ inline SmemLayout smemLayout(int32_t tableSize, int32_t entries, int32_t required) noexcept +{ + SmemLayout layout; + size_t offset = 0; + layout.keys = offset; + offset = alignUp16(offset + sizeof(int32_t) * static_cast(tableSize)); + layout.vals = offset; + offset = alignUp16(offset + sizeof(int32_t) * static_cast(tableSize)); + layout.cls = offset; + offset = alignUp16(offset + sizeof(uint8_t) * static_cast(entries)); + layout.scratch = offset; + offset = alignUp16(offset + sizeof(int32_t) * static_cast(entries)); + layout.order = offset; + offset = alignUp16(offset + sizeof(int32_t) * static_cast(entries)); + layout.lru = offset; + offset = alignUp16(offset + sizeof(int32_t) * static_cast(entries)); + layout.missIdx = offset; + offset = alignUp16(offset + sizeof(int32_t) * static_cast(required)); + layout.missOrd = offset; + offset = alignUp16(offset + sizeof(int32_t) * static_cast(required)); + layout.total = offset; + return layout; +} + +[[nodiscard]] constexpr int32_t nextPow2(int32_t value) noexcept +{ + int32_t result = 1; + while (result < value) + { + result <<= 1; + } + return result; +} + +//! Programmatic dependent launch, where the hardware has it. +//! +//! Both calls are hints, not contracts: a wait with no launch dependency returns at once, +//! and a release with no dependent kernel does nothing. So a kernel carrying them is correct +//! whether or not its launch enrolled in PDL, and compiling them out below Hopper only falls +//! back to ordinary stream ordering. +__device__ inline void gridDependencyWait() +{ +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 + cudaGridDependencySynchronize(); +#endif +} + +//! Let a dependent grid start. Only sound once everything it reads from this one is written. +__device__ inline void gridDependencyRelease() +{ +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 + cudaTriggerProgrammaticLaunchCompletion(); +#endif +} + +//! Fibonacci hashing. Unit ordinals are small and dense, and identity hashing turns a dense +//! run into one long probe chain. +__device__ inline uint32_t slotOf(int32_t key, uint32_t mask) noexcept +{ + return (static_cast(key) * 0x9E3779B1U) & mask; +} + +//! Exclusive sum of zero-or-one flags over [0, n), single CTA, result in `flags`. +//! Each warp scans a contiguous range before the ranges exchange their totals. +//! Returns the total. `flags` is shared memory; the scan is in place. +__device__ int32_t blockExclusiveSum(int32_t* flags, int32_t n) +{ + constexpr int32_t kWarps = kThreads / kWarpSize; + __shared__ int32_t warpTotals[kWarps]; + int32_t const lane = static_cast(threadIdx.x % kWarpSize); + int32_t const warp = static_cast(threadIdx.x / kWarpSize); + int32_t const tiles = (n + kThreads - 1) / kThreads; + int32_t const begin = warp * tiles * kWarpSize; + uint32_t const precedingLanes = (1U << lane) - 1U; + int32_t carry = 0; + for (int32_t tile = 0; tile < tiles; ++tile) + { + int32_t const idx = begin + tile * kWarpSize + lane; + uint32_t const votes = __ballot_sync(0xffffffffU, idx < n && flags[idx] != 0); + if (idx < n) + { + flags[idx] = carry + __popc(votes & precedingLanes); + } + carry += __popc(votes); + } + if (lane == 0) + { + warpTotals[warp] = carry; + } + __syncthreads(); + + int32_t offset = 0; + int32_t total = 0; +#pragma unroll + for (int32_t w = 0; w < kWarps; ++w) + { + int32_t const count = warpTotals[w]; + offset += w < warp ? count : 0; + total += count; + } + for (int32_t tile = 0; tile < tiles; ++tile) + { + int32_t const idx = begin + tile * kWarpSize + lane; + if (idx < n) + { + flags[idx] += offset; + } + } + __syncthreads(); + return total; +} + +//! Reserve one in-flight unit without building a selection probe index. +__global__ __launch_bounds__(kThreads) void sparsePrepareWriteKernel(SparseResolveArgs args) +{ + int32_t const row = static_cast(blockIdx.x); + SparseLeaseIndex const lease = args.leases[row]; + if (lease == kNoSparseLease) + { + return; + } + + gridDependencyWait(); + int32_t const historyUnits = args.historyTokens[row] / args.unitTokens; + int32_t const inFlight = args.capacityTokens[row] / args.unitTokens - historyUnits; + KVCM2_SPARSE_REQUIRE(inFlight >= 0 && inFlight <= 1); + if (threadIdx.x == 0) + { + args.missCount[row] = 0; + args.inFlightSlot[row] = kNoSparseSlot; + } + if (inFlight == 0) + { + return; + } + + int32_t const entries = args.arena.entriesPerSelection; + auto* const metadata = leaseMetadata(args.arena, lease); + auto* resident = metadata + static_cast(args.selectionSlot) * entries; + auto* order = metadata + args.arena.lruOffset + static_cast(args.selectionSlot) * entries; + extern __shared__ int32_t previousOrder[]; + using Reduce = cub::BlockReduce; + __shared__ typename Reduce::TempStorage temp; + __shared__ int32_t matchPosition; + __shared__ int32_t chosenPosition; + __shared__ int32_t chosenSlot; + + int32_t match = -1; + int32_t victim = -1; + for (int32_t p = static_cast(threadIdx.x); p < entries; p += kThreads) + { + int32_t const slot = order[p]; + int32_t const ordinal = resident[slot]; + previousOrder[p] = slot; + if (ordinal == historyUnits) + { + match = p; + } + if (ordinal == kFreeOrdinal || ordinal < historyUnits) + { + victim = p; + } + } + int32_t const matched = Reduce(temp).Reduce(match, cuda::maximum<>{}); + if (threadIdx.x == 0) + { + matchPosition = matched; + } + __syncthreads(); + int32_t const lastVictim = Reduce(temp).Reduce(victim, cuda::maximum<>{}); + if (threadIdx.x == 0) + { + chosenPosition = matchPosition >= 0 ? matchPosition : lastVictim; + KVCM2_SPARSE_REQUIRE(chosenPosition >= 0); + chosenSlot = previousOrder[chosenPosition]; + resident[chosenSlot] = historyUnits; + args.inFlightSlot[row] = chosenSlot; + } + __syncthreads(); + + // Move the touched entry to the front, retaining every other entry's order. + for (int32_t p = static_cast(threadIdx.x); p < chosenPosition; p += kThreads) + { + order[p + 1] = previousOrder[p]; + } + if (threadIdx.x == 0) + { + order[0] = chosenSlot; + } +} + +//! One CTA per request row of the resolved selection. +//! +//! Shared memory is carved by hand because the probe index is sized at launch; see +//! SmemLayout for the regions and their alignment. +__global__ __launch_bounds__(kThreads) void sparseResolveKernel(SparseResolveArgs args, int32_t tableSize) +{ + extern __shared__ unsigned char smem[]; + int32_t const selection = args.selectionSlot; + int32_t const row = static_cast(blockIdx.x); + SparseLeaseIndex const lease = args.leases[row]; + if (lease == kNoSparseLease) + { + // A free row, or a member that has suspended and given its reservation back. Its + // outputs keep whatever the caller initialised them to. + return; + } + int32_t const entries = args.arena.entriesPerSelection; + int32_t const width = args.selectionWidth; + uint32_t const mask = static_cast(tableSize) - 1U; + + int32_t const historyUnits = args.historyTokens[row] / args.unitTokens; + int32_t const inFlight = args.capacityTokens[row] / args.unitTokens - historyUnits; + // resize() rejects a gap wider than the cache can pin, so the check cannot fire. + KVCM2_SPARSE_REQUIRE(inFlight >= 0 && inFlight <= args.maxInFlightUnits); + // The probe walks one virtual list: the in-flight units, which must be resident before + // the selection takes any victim, then the selection itself. + int32_t const required = inFlight + width; + + SmemLayout const layout = smemLayout(tableSize, entries, args.maxInFlightUnits + width); + auto* keys = reinterpret_cast(smem + layout.keys); + auto* vals = reinterpret_cast(smem + layout.vals); + auto* cls = reinterpret_cast(smem + layout.cls); + auto* scratch = reinterpret_cast(smem + layout.scratch); + auto* order = reinterpret_cast(smem + layout.order); + auto* previousLru = reinterpret_cast(smem + layout.lru); + auto* missIdx = reinterpret_cast(smem + layout.missIdx); + auto* missOrd = reinterpret_cast(smem + layout.missOrd); + + __shared__ typename BlockScan::TempStorage temp; + __shared__ int32_t numMiss; + __shared__ int32_t numCopy; + __shared__ int32_t numResolved; + + // The view's own arrays belong to this row's allocation, so only the selection indexes + // them. + auto* const metadata = leaseMetadata(args.arena, lease); + int32_t* residentOrdinal = metadata + static_cast(selection) * entries; + int32_t* lruOrder = metadata + args.arena.lruOffset + static_cast(selection) * entries; + // The batch-wide arrays are indexed by row alone: the call is one selection, and every + // buffer resolving within it shares the single result. + auto const selectionBase = static_cast(row) * width; + auto const inFlightBase = static_cast(row) * args.maxInFlightUnits; + int32_t const* selected = args.selected + selectionBase; + int32_t* outSlot = args.outSlot + selectionBase; + int32_t* inFlightSlot = args.inFlightSlot + inFlightBase; + + if (threadIdx.x == 0) + { + numMiss = 0; + numCopy = 0; + numResolved = 0; + } + for (int32_t i = static_cast(threadIdx.x); i < tableSize; i += kThreads) + { + keys[i] = kEmptyKey; + vals[i] = kNoSparseSlot; + } + __syncthreads(); + + // The selection may be produced by the preceding grid. Only shared-memory + // initialization is independent of that producer. + gridDependencyWait(); + + // Index the required ordinals, then match the resident set against that index. + // The index size follows the selection width rather than the cache capacity. + for (int32_t i = static_cast(threadIdx.x); i < required; i += kThreads) + { + int32_t const ord = i < inFlight ? historyUnits + i : selected[i - inFlight]; + if (ord < 0) + { + continue; + } + uint32_t pos = slotOf(ord, mask); + int32_t key = atomicCAS(&keys[pos], kEmptyKey, ord); + while (key != kEmptyKey && key != ord) + { + pos = (pos + 1U) & mask; + key = atomicCAS(&keys[pos], kEmptyKey, ord); + } + } + __syncthreads(); + + for (int32_t s = static_cast(threadIdx.x); s < entries; s += kThreads) + { + int32_t const ord = residentOrdinal[s]; + cls[s] = kClassUntouched | (ord >= historyUnits ? kPinned : 0); + previousLru[s] = lruOrder[s]; + if (ord == kFreeOrdinal) + { + continue; + } + uint32_t pos = slotOf(ord, mask); + int32_t key = keys[pos]; + while (key != kEmptyKey && key != ord) + { + pos = (pos + 1U) & mask; + key = keys[pos]; + } + if (key == ord) + { + vals[pos] = s; + } + } + + // Slots past the in-flight span say "no unit in flight", so a consumer can read this + // array without also being told how many units there are. + for (int32_t i = static_cast(threadIdx.x) + inFlight; i < args.maxInFlightUnits; i += kThreads) + { + inFlightSlot[i] = kNoSparseSlot; + } + __syncthreads(); + + // Probe once per required unit. A hit resolves immediately; a miss is recorded against + // its position in the virtual list and resolved once a victim is chosen. + int32_t resolvedMine = 0; + for (int32_t i = static_cast(threadIdx.x); i < required; i += kThreads) + { + int32_t const ord = i < inFlight ? historyUnits + i : selected[i - inFlight]; + if (ord < 0) + { + outSlot[i - inFlight] = kNoSparseSlot; + continue; + } + resolvedMine += i < inFlight ? 0 : 1; + uint32_t pos = slotOf(ord, mask); + int32_t key = keys[pos]; + while (key != kEmptyKey && key != ord) + { + pos = (pos + 1U) & mask; + key = keys[pos]; + } + if (key == ord && vals[pos] != kNoSparseSlot) + { + int32_t const slot = vals[pos]; + (i < inFlight ? inFlightSlot[i] : outSlot[i - inFlight]) = slot; + cls[slot] = kClassHit; + } + else + { + int32_t const j = atomicAdd(&numMiss, 1); + missOrd[j] = ord; + missIdx[j] = i; + } + } + if (resolvedMine > 0) + { + atomicAdd(&numResolved, resolvedMine); + } + __syncthreads(); + + // Claim this row's run in the compacted list now that the miss set is final, so the + // eviction loop below can write straight into it. Only a miss below the watermark is + // copied, which is the same test the loop applies, so counting it here costs one pass + // over a list already in shared memory and saves staging the entries anywhere. + __shared__ int32_t copyBase; + { + int32_t mine = 0; + for (int32_t j = static_cast(threadIdx.x); j < numMiss; j += kThreads) + { + mine += missOrd[j] < historyUnits ? 1 : 0; + } + int32_t scanned = 0; + int32_t claimed = 0; + BlockScan(temp).ExclusiveSum(mine, scanned, claimed); + if (threadIdx.x == 0) + { + copyBase = claimed > 0 ? atomicAdd(args.work.total, claimed) : 0; + // The list is sized for every row missing its whole selection, so a claim can + // only run past it if the caller sized it from a different width. + KVCM2_SPARSE_REQUIRE(copyBase + claimed <= args.work.capacity); + if (args.hitMiss != nullptr) + { + atomicAdd(&args.hitMiss->resolved, static_cast(numResolved)); + atomicAdd(&args.hitMiss->misses, static_cast(claimed)); + } + } + } + __syncthreads(); + + // Entries that may be given away: free ones, and resident ones this step did not select + // that are not holding in-flight data. + for (int32_t p = static_cast(threadIdx.x); p < entries; p += kThreads) + { + int32_t const slot = previousLru[p]; + scratch[p] = cls[slot] == kClassUntouched ? 1 : 0; + } + __syncthreads(); + + // Victims come from the tail, least recently used first, so rank from the tail. + int32_t const totalEvictable = blockExclusiveSum(scratch, entries); + // The reservation covers a full selection plus every in-flight unit, and the in-flight + // ones are resolved first, so there is always a cell for each remaining miss. + KVCM2_SPARSE_REQUIRE(totalEvictable >= numMiss); + __syncthreads(); + + for (int32_t p = static_cast(threadIdx.x); p < entries; p += kThreads) + { + int32_t const slot = previousLru[p]; + bool const evictable = cls[slot] == kClassUntouched; + if (!evictable) + { + continue; + } + // scratch[p] is the count of evictable entries before p, so this is its rank from + // the tail. The numMiss least recently used take the misses, in selection order. + int32_t const rankFromTail = totalEvictable - 1 - scratch[p]; + if (rankFromTail < numMiss) + { + int32_t const j = numMiss - 1 - rankFromTail; + int32_t const ordinal = missOrd[j]; + int32_t const position = missIdx[j]; + residentOrdinal[slot] = ordinal; + (position < inFlight ? inFlightSlot[position] : outSlot[position - inFlight]) = slot; + cls[slot] = kClassFetched; + // Only history has anything to fetch. An in-flight unit's host page is not + // written until the watermark passes it, so copying from it would read stale + // bytes over the data the model is about to write. + if (ordinal < historyUnits) + { + auto const c = copyBase + atomicAdd(&numCopy, 1); + args.work.row[c] = row; + args.work.ordinal[c] = ordinal; + args.work.slot[c] = slot; + } + } + } + __syncthreads(); + + // The work list and slots are final. The copy can perform its independent setup + // while this block rebuilds recency, then wait before consuming the work list. + gridDependencyRelease(); + + // Rebuild recency: hits, then units fetched this step, then everything untouched, each + // group keeping its previous relative order. Each warp ranks all three classes + // within a contiguous range; class totals give the offsets between ranges. + constexpr int32_t kWarps = kThreads / kWarpSize; + __shared__ int32_t classTotals[3][kWarps]; + int32_t const lane = static_cast(threadIdx.x % kWarpSize); + int32_t const warp = static_cast(threadIdx.x / kWarpSize); + int32_t const tiles = (entries + kThreads - 1) / kThreads; + int32_t const begin = warp * tiles * kWarpSize; + uint32_t const precedingLanes = (1U << lane) - 1U; + int32_t carry[3] = {}; + for (int32_t tile = 0; tile < tiles; ++tile) + { + int32_t const p = begin + tile * kWarpSize + lane; + uint8_t const c = p < entries ? cls[previousLru[p]] & ~kPinned : 3; +#pragma unroll + for (uint8_t group = 0; group < 3; ++group) + { + uint32_t const votes = __ballot_sync(0xffffffffU, c == group); + if (c == group) + { + order[p] = carry[group] + __popc(votes & precedingLanes); + } + carry[group] += __popc(votes); + } + } + if (lane == 0) + { +#pragma unroll + for (uint8_t c = 0; c < 3; ++c) + { + classTotals[c][warp] = carry[c]; + } + } + __syncthreads(); + + int32_t offsets[3] = {}; + int32_t base = 0; +#pragma unroll + for (uint8_t c = 0; c < 3; ++c) + { + int32_t count = 0; + int32_t before = 0; +#pragma unroll + for (int32_t w = 0; w < kWarps; ++w) + { + int32_t const total = classTotals[c][w]; + before += w < warp ? total : 0; + count += total; + } + offsets[c] = base + before; + base += count; + } + for (int32_t tile = 0; tile < tiles; ++tile) + { + int32_t const p = begin + tile * kWarpSize + lane; + if (p < entries) + { + int32_t const slot = previousLru[p]; + uint8_t const c = cls[slot] & ~kPinned; + int32_t const offset = c == kClassHit ? offsets[0] : (c == kClassFetched ? offsets[1] : offsets[2]); + scratch[offset + order[p]] = slot; + } + } + __syncthreads(); + + for (int32_t p = static_cast(threadIdx.x); p < entries; p += kThreads) + { + lruOrder[p] = scratch[p]; + } + if (threadIdx.x == 0) + { + args.missCount[row] = numCopy; + } +} + +//! Copy `bytes` from `src` to `dst` with one warp, `sizeof(Vec)` at a time. +//! +//! The width is the launcher's choice, made from the layout and the addresses it produces, +//! so there is no alignment test here and no tail: `bytes` divides by `sizeof(Vec)` and both +//! pointers are aligned to it. +template +__device__ inline void copyChunkWarp(char* dst, char const* src, uint32_t bytes, uint32_t lane) +{ + auto* out = reinterpret_cast(dst); + auto const* in = reinterpret_cast(src); + uint32_t const count = bytes / static_cast(sizeof(Vec)); + for (uint32_t i = lane; i < count; i += kWarpSize) + { + out[i] = in[i]; + } +} + +//! One pool's work, with a vector width fixed at compile time. +//! +//! Split from the kernel so the width is switched on once per block: one launch covers +//! every pool, so there is no per-pool launch to template. +template +__device__ void copySparseUnits(SparseUnitCopyArgs const& args, SparseCopyDesc const& desc, int32_t list, int32_t first) +{ + uint32_t const lane = threadIdx.x % kWarpSize; + int32_t const stride = static_cast(gridDim.x) * kWarpsPerCopyCta; + auto const listBase = static_cast(list) * args.work.capacity; + int32_t const count = args.work.total[list]; + + for (int32_t j = first; j < count; j += stride) + { + auto const index = listBase + j; + int32_t const row = args.work.row[index]; + int32_t const ordinal = args.work.ordinal[index]; + int32_t const slot = args.work.slot[index]; + + // Read per item rather than per warp: the list is dense across rows, so consecutive + // entries need not share one. A producer's CTA claims a contiguous run, so in + // practice a warp's items mostly do. + SparseLeaseIndex const lease = args.leases[row]; + + // A unit ordinal names a unit within a block; a cache entry names a unit within a + // sub-slot. Both divide by the same units-per-page. + int32_t const block = ordinal / args.unitsPerPage; + int32_t const srcUnit = ordinal % args.unitsPerPage; + int32_t const dstUnit = slot % args.unitsPerPage; + if (block >= static_cast(desc.numBlocks)) + { + // Outside the extent this call was given; the warp's other items may not be. + continue; + } + + // The descriptor's pointer already stands at this buffer's entry, so the table + // index is the row's run of blocks and nothing else. + auto const hostPage + = static_cast(desc.pageTable[static_cast(row) * desc.pageTableStride + block]); + + // The cold side is reached by ordinal through the page table, the hot side by + // arithmetic from the lease: the arena hands out a contiguous run, so there is no + // page to look up, and the buffer's own slot in the pool is already in deviceBase. + // Direction only chooses which of the two is read. + auto* cold = reinterpret_cast( + desc.hostBase + hostPage * desc.hostPageStride + static_cast(srcUnit) * desc.layout.unitStride); + auto* hot = reinterpret_cast(desc.deviceBase + static_cast(lease.value()) * desc.leaseStride + + static_cast(slot / args.unitsPerPage) * desc.cellBytes + + static_cast(dstUnit) * desc.layout.unitStride); + bool const fetching = args.direction == SparseCopyDirection::kFetch; + char const* src = fetching ? cold : hot; + char* dst = fetching ? hot : cold; + + for (uint32_t c = 0; c < desc.layout.numChunks; ++c) + { + auto const offset = static_cast(c) * desc.layout.chunkStride; + copyChunkWarp(dst + offset, src + offset, desc.layout.chunkBytes, lane); + } + } +} + +//! One warp per work item, walking the compacted list with a grid-stride loop. `blockIdx.y` +//! is the buffer, `blockIdx.x` the work. +__global__ __launch_bounds__(kThreads) void sparseUnitCopyKernel(SparseUnitCopyArgs args) +{ + // Broadcast from lane zero before anything can diverge: threadIdx.x varies within the + // warp, so the compiler cannot otherwise see that this quotient does not, and every lane + // must still be present for the shuffle. + uint32_t warpInCta = threadIdx.x / kWarpSize; + warpInCta = __shfl_sync(0xffffffffU, warpInCta, 0); + int32_t const globalWarp = static_cast(blockIdx.x) * kWarpsPerCopyCta + static_cast(warpInCta); + + // Everything down to the dependency wait is the producer's business only in that it + // names where to put the result: the descriptors are built once with the Batch, and the + // buffer and list follow from the block's own coordinates. Resolving it here is what the + // launch dependency is for -- the block is resident and addressed before the kernel it + // waits on has finished. + int32_t const buffer = args.bufferBase + static_cast(blockIdx.y); + SparseCopyDesc const desc = args.copyDesc[buffer]; + // The buffer is this block's, so its layer -- and therefore its list -- is uniform and + // known before any entry is read. A buffer's layer is dense from zero within its group, + // and a multi-list caller has one list per layer of that group in the same order, so the + // layer is the list. The index is zero whenever the caller passed one list, which is what + // a fetch does: it resolves a single layer. + int32_t const list = args.work.numLists > 1 ? args.arena.bufferLayer[buffer] : 0; + + // The work list and its total are the producer's, so this is as far as the block gets + // before the kernel that fills them has run. + gridDependencyWait(); + + // Nothing launched after this reads a cell, so a successor need not wait for the copy + // itself -- only for the block to have got this far. Releasing here is what keeps the + // index kernel that follows on the stream from being ordered behind a copy it does not + // depend on. + gridDependencyRelease(); + + switch (desc.width) + { + case SparseCopyWidth::kWord4: copySparseUnits(args, desc, list, globalWarp); break; + case SparseCopyWidth::kWord2: copySparseUnits(args, desc, list, globalWarp); break; + case SparseCopyWidth::kWord: copySparseUnits(args, desc, list, globalWarp); break; + case SparseCopyWidth::kByte: copySparseUnits(args, desc, list, globalWarp); break; + } +} + +//! One CTA per (row, buffer). +__global__ __launch_bounds__(kThreads) void sparseOutIndexKernel(SparseOutIndexArgs args) +{ + int32_t const row = static_cast(blockIdx.y); + int32_t const buffer = args.bufferBase + static_cast(blockIdx.x); + auto* out = args.out + + (static_cast(row) * args.outBuffers + args.outBase + blockIdx.x) + * static_cast(args.selectionWidth); + + SparseLeaseIndex const lease = args.leases[row]; + bool const reserved = lease != kNoSparseLease; + auto const base = static_cast(row) * args.selectionWidth; + int32_t const slotStride = args.outSlotStride > 0 ? args.outSlotStride : args.selectionWidth; + auto const slotBase = static_cast(row) * slotStride; + + // Every slot is written, so a stale index from a previous step can never be read back. + // One pass does it, because the output is the selection's own shape: a padded slot and + // an unreserved row both just take the sentinel here rather than needing a prior fill. + for (int32_t i = static_cast(threadIdx.x); i < args.selectionWidth; i += kThreads) + { + int32_t const slot = reserved ? args.outSlot[slotBase + i] : kNoSparseSlot; + // Without a selection array the slot alone says whether the entry is live, which is + // what the write-index path passes: its in-flight list is already self-describing. + if ((args.selected != nullptr && args.selected[base + i] < 0) || slot < 0) + { + out[i] = kBadPageIndexValue; + continue; + } + + int32_t const cell = lease.value() * args.cellsPerLease + + args.arena.bufferSlotInPool[buffer] * args.cellsPerBuffer + slot / args.unitsPerPage; + out[i] = cell * args.unitsPerPage + slot % args.unitsPerPage; + } +} + +//! One CTA per (layer, row). The units wanted form a contiguous ordinal range, so a scan of the +//! reverse map finds them without a probe index. +__global__ __launch_bounds__(kThreads) void sparseFlushListKernel(SparseFlushListArgs args) +{ + int32_t const row = static_cast(blockIdx.y); + SparseLeaseIndex const lease = args.leases[row]; + if (lease == kNoSparseLease) + { + // A free row, or a member that has suspended: no reservation to write back from. + return; + } + int32_t const layer = static_cast(blockIdx.x); + int32_t const entries = args.arena.entriesPerSelection; + + // The row's view is staged with the watermarks and is nobody's output here, so a block + // can address itself before waiting. The range itself is uploaded by the call that asked + // for this write-back, so reading it is where the block has to stop. + gridDependencyWait(); + + int32_t const from = args.fromUnits[row]; + int32_t const to = args.toUnits[row]; + if (from >= to) + { + // This row's watermark did not move, so nothing of its crossed, but the grid is + // sized for the widest row. + if (threadIdx.x == 0) + { + args.outCount[static_cast(row) * args.arena.numLayers + layer] = 0; + } + return; + } + + __shared__ typename BlockScan::TempStorage temp; + // The entries scanned are the layer's selection's, which several layers may share; the + // list claimed below is the layer's own, because the copy that consumes it walks the + // layer's buffers and each holds its own mirror of a shared entry. + int32_t const* residentOrdinal + = leaseMetadata(args.arena, lease) + static_cast(args.arena.layerSelection[layer]) * entries; + auto const crossed = [&](int32_t slot) + { + int32_t const ord = residentOrdinal[slot]; + return ord >= from && ord < to; + }; + + // Two passes over the resident set, because the run cannot be claimed until its length + // is known and the entries are too many to stage in shared memory: a prefilling row + // crosses as many units as its prompt holds. The set is the per-request cache, which is + // small enough that scanning it twice costs less than staging it once. + __shared__ int32_t count; + __shared__ int32_t base; + if (threadIdx.x == 0) + { + count = 0; + } + __syncthreads(); + + { + int32_t mine = 0; + for (int32_t s = static_cast(threadIdx.x); s < entries; s += kThreads) + { + mine += crossed(s) ? 1 : 0; + } + int32_t scanned = 0; + int32_t claimed = 0; + BlockScan(temp).ExclusiveSum(mine, scanned, claimed); + if (threadIdx.x == 0) + { + base = claimed > 0 ? atomicAdd(args.work.total + layer, claimed) : 0; + // Sized from the step's token budget, which bounds how far every row's watermark + // can move together. Running past it means the budget the Batch was given is not + // the one the caller honoured. + KVCM2_SPARSE_REQUIRE(base + claimed <= args.work.capacity); + args.outCount[static_cast(row) * args.arena.numLayers + layer] = claimed; + } + } + __syncthreads(); + + auto const listBase = static_cast(layer) * args.work.capacity + base; + for (int32_t s = static_cast(threadIdx.x); s < entries; s += kThreads) + { + if (!crossed(s)) + { + continue; + } + auto const j = listBase + atomicAdd(&count, 1); + args.work.row[j] = row; + args.work.ordinal[j] = residentOrdinal[s]; + args.work.slot[j] = s; + } + + // The list is the last thing this kernel writes, so the release only brings the copy's + // launch forward. + __syncthreads(); + gridDependencyRelease(); +} + +} // namespace + +//! Throw unless the work list is fully described. +void checkWorkList(SparseWorkList const& work, char const* api) +{ + TLLM_CHECK_WITH_INFO( + work.row != nullptr && work.ordinal != nullptr && work.slot != nullptr && work.total != nullptr, + std::string(api) + ": the work list's device arrays must be provided"); + TLLM_CHECK_WITH_INFO(work.capacity > 0 && work.numLists > 0, + std::string(api) + ": the work list's capacity and list count must be positive"); +} + +void launchSparseOutIndices(SparseOutIndexArgs const& args, CUstream stream) +{ + TLLM_CHECK_WITH_INFO(args.numRows > 0 && args.numBuffers > 0 && args.selectionWidth > 0, + "launchSparseOutIndices: every extent must be positive"); + TLLM_CHECK_WITH_INFO(args.leases != nullptr && args.outSlot != nullptr && args.out != nullptr, + "launchSparseOutIndices: every device array must be provided"); + TLLM_CHECK_WITH_INFO( + args.cellsPerBuffer > 0 && args.cellsPerLease > 0 && args.unitsPerPage > 0 && args.bufferBase >= 0, + "launchSparseOutIndices: the cell geometry must be positive"); + TLLM_CHECK_WITH_INFO(args.outBase >= 0 && args.outBase + args.numBuffers <= args.outBuffers, + "launchSparseOutIndices: this call's buffers must fit the output row"); + TLLM_CHECK_WITH_INFO( + args.outSlotStride >= 0 && (args.outSlotStride == 0 || args.outSlotStride >= args.selectionWidth), + "launchSparseOutIndices: the slot stride must cover what this call reads"); + + dim3 const grid(static_cast(args.numBuffers), static_cast(args.numRows)); + sparseOutIndexKernel<<(stream)>>>(args); + TLLM_CUDA_CHECK(cudaGetLastError()); +} + +void launchSparseFlushList(SparseFlushListArgs const& args, CUstream stream) +{ + TLLM_CHECK_WITH_INFO( + args.numRows > 0 && args.arena.numLayers > 0, "launchSparseFlushList: rows and layers must be positive"); + TLLM_CHECK_WITH_INFO(args.arena.layerSelection != nullptr, + "launchSparseFlushList: the scan reads each layer's selection, so the arena must name it"); + TLLM_CHECK_WITH_INFO( + args.leases != nullptr && args.fromUnits != nullptr && args.toUnits != nullptr && args.outCount != nullptr, + "launchSparseFlushList: every device array must be provided"); + checkWorkList(args.work, "launchSparseFlushList"); + TLLM_CHECK_WITH_INFO(args.work.numLists == args.arena.numLayers, + "launchSparseFlushList: a write-back needs one work list per layer"); + dim3 const grid(static_cast(args.arena.numLayers), static_cast(args.numRows)); + tensorrt_llm::common::launchWithPdlWhenEnabled( + "sparseFlushList", sparseFlushListKernel, grid, kThreads, 0, reinterpret_cast(stream), args); +} + +SparseCopyWidth sparseCopyWidth(SelectionUnitLayout const& layout, uint64_t hostPageStride, uint64_t bufferBytes) +{ + // Sizes only. Every offset a lane adds to a base is a sum of these, so a width dividing + // all of them leaves every access aligned relative to the base -- and the bases + // themselves are checked once, on the host, where a misaligned pool is a configuration + // error rather than a case to carry into the kernel. + auto const fits = [&](uint64_t width) + { + auto const divides = [width](uint64_t v) { return v % width == 0; }; + return divides(hostPageStride) && divides(bufferBytes) && divides(layout.unitStride) + && divides(layout.chunkStride) && divides(layout.chunkBytes); + }; + if (fits(sizeof(uint4))) + { + return SparseCopyWidth::kWord4; + } + if (fits(sizeof(uint2))) + { + return SparseCopyWidth::kWord2; + } + if (fits(sizeof(uint32_t))) + { + return SparseCopyWidth::kWord; + } + return SparseCopyWidth::kByte; +} + +size_t sparseCopyWidthBytes(SparseCopyWidth width) noexcept +{ + switch (width) + { + case SparseCopyWidth::kWord4: return sizeof(uint4); + case SparseCopyWidth::kWord2: return sizeof(uint2); + case SparseCopyWidth::kWord: return sizeof(uint32_t); + case SparseCopyWidth::kByte: break; + } + return 1; +} + +void launchSparseUnitCopy(SparseUnitCopyArgs const& args, CUstream stream) +{ + TLLM_CHECK_WITH_INFO( + args.unitsPerPage > 0 && args.numRows > 0 && args.gridCtas > 0 && args.numBuffers > 0 && args.bufferBase >= 0, + "launchSparseUnitCopy: rows, buffers, grid CTAs and units-per-page must be positive"); + TLLM_CHECK_WITH_INFO(args.copyDesc != nullptr && args.leases != nullptr, + "launchSparseUnitCopy: every device array must be provided"); + checkWorkList(args.work, "launchSparseUnitCopy"); + TLLM_CHECK_WITH_INFO(args.work.numLists <= 1 || args.arena.bufferLayer != nullptr, + "launchSparseUnitCopy: a per-layer work list needs the buffer-to-layer map"); + + dim3 const grid(static_cast(args.gridCtas), static_cast(args.numBuffers)); + tensorrt_llm::common::launchWithPdlWhenEnabled( + "sparseUnitCopy", sparseUnitCopyKernel, grid, kThreads, 0, reinterpret_cast(stream), args); +} + +int32_t sparseCopyMultiprocessors(CUdevice device) +{ + int smCount = 0; + cuCheck(cuDeviceGetAttribute(&smCount, CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT, device)); + return std::max(smCount, 1); +} + +int32_t sparseCopyGridCtas(int32_t workCapacity, int32_t smCount) noexcept +{ + constexpr int32_t kCtasPerSm = 4; + + int32_t const full = std::max(smCount, 1) * kCtasPerSm; + int32_t const needed = (workCapacity + kWarpsPerCopyCta - 1) / kWarpsPerCopyCta; + return std::max(1, std::min(full, needed)); +} + +size_t sparseResolveSharedBytes(int32_t entriesPerSelection, int32_t selectionWidth, int32_t maxInFlightUnits) +{ + int32_t const required = maxInFlightUnits + selectionWidth; + return smemLayout(nextPow2(2 * required), entriesPerSelection, required).total; +} + +namespace +{ + +//! Dynamic shared memory a block of the resolve kernel may take, which it is opted in to. +//! +//! The device is asked once and the kernel opted in at that maximum once: the opt-in is a +//! ceiling, and a block's occupancy follows the bytes its launch asks for rather than the +//! ceiling, so granting the most the device allows costs nothing. +int32_t sparseResolveSharedLimit() +{ + static int32_t const limit = [] + { + CUdevice device = 0; + cuCheck(cuCtxGetDevice(&device)); + int deviceMax = 0; + cuCheck(cuDeviceGetAttribute(&deviceMax, CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK_OPTIN, device)); + // The opt-in names dynamic bytes, and what the kernel declares statically is spent + // from the same per-block budget, so only the rest of it can be asked for. + cudaFuncAttributes attr{}; + TLLM_CUDA_CHECK(cudaFuncGetAttributes(&attr, reinterpret_cast(&sparseResolveKernel))); + int const value = deviceMax - static_cast(attr.sharedSizeBytes); + TLLM_CUDA_CHECK(cudaFuncSetAttribute( + reinterpret_cast(&sparseResolveKernel), cudaFuncAttributeMaxDynamicSharedMemorySize, value)); + return static_cast(value); + }(); + return limit; +} + +} // namespace + +size_t checkedSparseResolveSharedBytes(int32_t entriesPerSelection, int32_t selectionWidth, int32_t maxInFlightUnits) +{ + size_t const shared = sparseResolveSharedBytes(entriesPerSelection, selectionWidth, maxInFlightUnits); + int32_t const limit = sparseResolveSharedLimit(); + TLLM_CHECK_WITH_INFO(shared <= static_cast(limit), + "Sparse resolve: the cache needs " + std::to_string(shared) + + " bytes of shared memory, above this device's limit of " + std::to_string(limit) + + "; reduce max_cached_units or num_selected_blocks, or raise the selection unit size"); + return shared; +} + +void launchSparseResolve(SparseResolveArgs const& args, CUstream stream) +{ + TLLM_CHECK_WITH_INFO(args.arena.numSelections > 0 && args.arena.entriesPerSelection > 0 && args.selectionWidth >= 0 + && args.numRows > 0, + "launchSparseResolve: rows, selections and entries must be positive"); + // A zero-width resolve is the write-index path: it makes the in-flight unit resident and + // nothing else, so it needs neither a selection to read nor a miss list to fill. + TLLM_CHECK_WITH_INFO(args.selectionWidth != 0 || args.maxInFlightUnits != 0, + "launchSparseResolve: a resolve with no selection must have an in-flight span"); + TLLM_CHECK_WITH_INFO(args.selectionSlot >= 0 && args.selectionSlot < args.arena.numSelections, + "launchSparseResolve: selectionSlot must name a selection of the life cycle"); + bool const selectionArrays = args.selectionWidth == 0 || (args.selected != nullptr && args.outSlot != nullptr); + TLLM_CHECK_WITH_INFO(args.leases != nullptr && selectionArrays && args.missCount != nullptr + && args.historyTokens != nullptr && args.capacityTokens != nullptr + && (args.maxInFlightUnits <= 0 || args.inFlightSlot != nullptr), + "launchSparseResolve: every device array must be provided"); + TLLM_CHECK_WITH_INFO(args.maxInFlightUnits >= 0, "launchSparseResolve: maxInFlightUnits must not be negative"); + if (args.selectionWidth > 0) + { + checkWorkList(args.work, "launchSparseResolve"); + TLLM_CHECK_WITH_INFO( + args.work.numLists == 1, "launchSparseResolve: a fetch resolves one selection, so it fills one list"); + } + TLLM_CHECK_WITH_INFO(args.unitTokens > 0, "launchSparseResolve: unitTokens must be positive"); + + if (args.selectionWidth == 0 && args.maxInFlightUnits == 1 + && sizeof(int32_t) * static_cast(args.arena.entriesPerSelection) <= 32U * 1024U) + { + size_t const shared = sizeof(int32_t) * static_cast(args.arena.entriesPerSelection); + tensorrt_llm::common::launchWithPdlWhenEnabled("sparsePrepareWrite", sparsePrepareWriteKernel, + dim3(static_cast(args.numRows)), dim3(kThreads), shared, reinterpret_cast(stream), + args); + return; + } + + // entriesPerSelection belongs to the life cycle, so one size and one grid serve every row. + int32_t const tableSize = nextPow2(2 * (args.maxInFlightUnits + args.selectionWidth)); + size_t const shared + = checkedSparseResolveSharedBytes(args.arena.entriesPerSelection, args.selectionWidth, args.maxInFlightUnits); + tensorrt_llm::common::launchWithPdlWhenEnabled("sparseResolve", sparseResolveKernel, + dim3(static_cast(args.numRows)), dim3(kThreads), shared, reinterpret_cast(stream), args, + tableSize); +} + +} // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/sparseCacheKernels.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/sparseCacheKernels.h new file mode 100644 index 000000000000..cde52f078d65 --- /dev/null +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/sparseCacheKernels.h @@ -0,0 +1,455 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "kv_cache_manager_v2/sparseCache.h" + +#include + +#include +#include + +namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 +{ + +//! Written to an output slot that carries no selection, i.e. selection padding. +inline constexpr int32_t kNoSparseSlot = -1; + +//! These kernels enforce their preconditions by trapping, not by reporting: a violation is +//! a KVCM2 logic error, and the trap destroys the context, so no status could be read back. +//! A debug build names the condition before halting. + +//! A list of units for the copy to move, compacted across rows. +//! +//! Each producer CTA claims a contiguous run with one `atomicAdd` on its list's counter and +//! writes its entries there, so the consumer walks a dense array: its grid follows the +//! machine rather than the worst case, and no warp is tied to a row that had no work. The +//! cost is that position no longer says which row an entry belongs to, so the entry says it. +//! +//! A fetch resolves one selection per launch and uses a single list, which every layer of +//! that selection then copies from. A write-back spans every layer of the life cycle at +//! once, so it takes one list per layer, chosen by the copied buffer's layer. +//! +//! Entry order is whatever the producer's CTAs raced to, which nothing depends on: the +//! entries name disjoint cells, so they may be applied in any order. +struct SparseWorkList +{ + //! [numLists][capacity] the row, the unit and the entry holding it. + int32_t* row = nullptr; + int32_t* ordinal = nullptr; + int32_t* slot = nullptr; + //! [numLists] entries written so far. Zeroed before the producer runs; the consumer + //! bounds its loop by it, which is why the count never crosses to the host. + int32_t* total = nullptr; + + //! Entries one list holds, and the distance between consecutive lists. Zero for a single + //! list. + int32_t capacity = 0; + //! Lists behind `total`, so a producer can bound the layer it indexes. + int32_t numLists = 1; +}; + +//! Fetch statistics, as the device counter and as its host copy. +//! +//! The two figures the resolve produces, not the two a reader wants: `hits()` is their +//! difference and is never stored. `misses` is the row's claimed run in the compacted miss +//! list -- the selected units that were not resident and so had to be copied from their +//! host pages. +//! +//! `unsigned long long` rather than `uint64_t` because the device side accumulates with +//! `atomicAdd`, which has no overload for `unsigned long` on LP64. +//! +//! 64-bit because a step resolves millions of units per layer and a 32-bit pair would wrap. +struct SparseHitMiss +{ + //! Selection entries that named real history. Padding and in-flight units are neither + //! hit nor miss, so neither reaches this. + unsigned long long resolved; + //! Of those, the ones the cache did not hold. + unsigned long long misses; + + //! Clamped, because a host copy races a running fetch and the two are separate atomics: + //! a resolve can land between them and show a miss its selection has not been counted + //! for yet. A diagnostic must not be able to report a negative. + [[nodiscard]] unsigned long long hits() const noexcept + { + return resolved >= misses ? resolved - misses : 0; + } +}; + +//! One selection's resolve across a batch: turn each row's selection of unit ordinals into +//! cache entry indices, allocating entries for the units that are not resident. +//! +//! The scope is one selection, which is usually one layer: a layer's selection normally +//! depends on the previous layer's output, so they are not all known at the start of a +//! step and each gets its own call, letting the first layer's copies start as early as +//! possible. Where a model runs one indexer across several layers they are known together, +//! and those layers name one selection and resolve once between them. +//! +//! One result serves every sparse buffer of every layer in the selection, since they hold +//! the same units in mirrored cells and share one metadata row. Layouts are therefore +//! [row][...], with no buffer or layer dimension. Every array is device memory. +struct SparseResolveArgs +{ + //! [numRows] the arena slot each row holds, as Batch::sparseLeases() publishes it. A + //! row holding kNoSparseLease has no reservation -- free, or a member that has + //! suspended -- and is skipped. + SparseLeaseIndex const* leases; + //! The life cycle's arena, which every row's lease is addressed within. + SparseArenaView arena; + int32_t numRows; + //! Which of the life cycle's selections this call resolves: the metadata row every + //! sparse buffer of every layer naming that selection shares. + int32_t selectionSlot; + + //! [numRows][selectionWidth] unit ordinals, negative where the selection is padded. + //! Ordinals must be unique within a row; duplicates would be resolved as separate + //! misses and waste an entry each. + int32_t const* selected; + //! Selected units per row, uniform across requests so shapes stay static. + int32_t selectionWidth; + + //! [numRows] history length and capacity, in tokens, as Batch::watermarks() publishes + //! them; `unitTokens` divides both to units. Taking the source and the divisor rather + //! than the two derived counts keeps one truth, and spares the caller a conversion pass + //! over arrays it would otherwise have to allocate. + //! + //! Device-resident because both move every step: a captured graph must read the current + //! value rather than one baked in at capture. + //! + //! Ordinals below historyUnits are selectable history. Ordinals in + //! [historyUnits, capacityUnits) are in flight: the entry holds the only copy of their + //! data, so it is pinned against eviction and there is nothing to fetch for it. See + //! "Unflushed entries carry no state" in SPARSE_ATTENTION_DESIGN.md. + int32_t const* historyTokens; + int32_t const* capacityTokens; + //! tokens_per_block / unitsPerPage, a per-layer-group constant. + int32_t unitTokens; + + //! [numRows][selectionWidth] entry index per selection, kNoSparseSlot where padded or + //! where the row holds no reservation. One list serves every buffer of the layer. + int32_t* outSlot; + + //! [numRows][maxInFlightUnits] entry index per in-flight unit, in ordinal order from + //! historyUnits. Entries past the live count are left untouched. The model writes new + //! KV through these, so they are resolved before the selection takes any victim. + int32_t* inFlightSlot; + //! Array stride, bounding capacityUnits - historyUnits. Follows from the cache width, + //! so the shape stays static while the live count varies. + int32_t maxInFlightUnits; + + //! The history units that were not resident, and the entries given to them. One list, + //! since a call resolves one selection; every buffer of every layer naming that + //! selection copies from it, each into its own mirror of the entry. + //! + //! An in-flight unit that misses is allocated an entry but never appears here: its host + //! page is not written until the watermark passes it, so there is nothing to fetch. + SparseWorkList work; + //! [numRows] entries this row contributed. The copy does not read it -- it walks the + //! compacted list -- but it is what makes a row's hit rate observable. + int32_t* missCount; + + //! Cumulative fetch statistics, which every row adds its own share of. Null skips it. + //! + //! Counted in the resolve because only the resolve can count either half: a padded + //! selection entry is negative and an in-flight unit is not a fetch, neither of which + //! the host can see. Both halves come from one pass, so they agree on what they count. + //! + //! A row adds its own contribution rather than anything reading a total afterwards. The + //! step's counters are cleared once per publish, not per fetch, so a reader downstream + //! would count two fetches of one layer twice. + SparseHitMiss* hitMiss = nullptr; +}; + +//! Shared memory one CTA needs for `entriesPerSelection` entries and `selectionWidth` +//! selections. The probe index grows as next_pow2(2 * (selectionWidth + maxInFlightUnits)); +//! per-entry recency and eviction scratch grow linearly with entriesPerSelection. +[[nodiscard]] size_t sparseResolveSharedBytes( + int32_t entriesPerSelection, int32_t selectionWidth, int32_t maxInFlightUnits); + +//! Return the resolve's shared-memory requirement after checking it against the current +//! device's opt-in limit less the kernel's static shared memory. Used when constructing +//! a Batch and before launching the kernel; also opts the kernel in to that limit. +[[nodiscard]] size_t checkedSparseResolveSharedBytes( + int32_t entriesPerSelection, int32_t selectionWidth, int32_t maxInFlightUnits); + +//! Multiprocessors `device` has, which is the one part of the grid sizing below that costs +//! a driver call and so the one part worth holding on to. +[[nodiscard]] int32_t sparseCopyMultiprocessors(CUdevice device); + +//! CTAs to give a copy: enough to fill a machine of `smCount` multiprocessors, capped by +//! the most work items the list could ever hold, since beyond that the extra CTAs would +//! only exit. +//! +//! Fixed before a graph captures the copy, so it follows the machine rather than the step. +[[nodiscard]] int32_t sparseCopyGridCtas(int32_t workCapacity, int32_t smCount) noexcept; + +//! Which way a unit moves between its cold page and its cache entry. +enum class SparseCopyDirection : uint8_t +{ + //! Host page -> cache entry, filling an entry the resolve allocated for a miss. + kFetch = 0, + //! Cache entry -> host page, writing back a unit the watermark has passed. + kWriteBack = 1, +}; + +//! Vector type one pool's copy moves at a time, chosen from the layout's sizes alone. +//! +//! Sizes rather than addresses, so the width is a property of the layout and can live in a +//! descriptor built once. The two base pointers are checked against it on the host, where a +//! misaligned pool is a configuration error worth naming rather than a case to handle. +enum class SparseCopyWidth : int32_t +{ + kByte = 0, + kWord = 1, + kWord2 = 2, + kWord4 = 3, +}; + +//! A maximal set of consecutive buffers sharing a pool, as the buffer plan groups them. +//! +//! Host-side only: it is what the manager knows about a pool, before the Batch folds each +//! buffer's own slot and index-table entry into the per-buffer descriptors below. Buffers +//! are ordered (layer, pool), so a run never spans layers. +struct SparseCopyRun +{ + int32_t bufferBase; + int32_t numBuffers; + //! Base of this pool's region in the cell arena, before a buffer's slot is added. + CUdeviceptr deviceBase; + uint32_t bufferBytes; + int32_t cellsPerBuffer; + int32_t cellsPerLease; + uint32_t hostPageStride; + SelectionUnitLayout layout; + SparseCopyWidth width; +}; + +//! One buffer's share of a copy: everything a block needs that does not vary per work item. +//! +//! Device-resident and constant for the Batch's lifetime. Every field derives from the +//! buffer plan and the index tables, both fixed at construction, so these are built and +//! uploaded once -- which is what lets a captured graph replay a copy without any host +//! write, and what lets every buffer of every pool share one launch. +//! +//! Per buffer rather than per pool so a block resolves itself in a single load, and so the +//! two indirections a pool-level descriptor needs -- the buffer's slot within its pool and +//! its entry within the page table -- can be folded in here instead of being paid per work +//! item. +//! +//! Sixteen-byte aligned and exactly four of them wide, so the load vectorises; the strides +//! are 32-bit because a page, a lease and a cell are all far smaller than that, and only +//! their products with a page or lease index need the wider type. +struct alignas(16) SparseCopyDesc +{ + //! Base this buffer's converted page indices are relative to. Shared by the buffers of + //! one run: the converted index already carries each buffer's own offset within the + //! coalesced page, which is why the pool's slot size plays no part and the stride below + //! is one buffer. + CUdeviceptr hostBase; + //! Converted page indices, already at this buffer's entry, so a row's run of blocks is + //! `pageTable[row * pageTableStride ...]`. Buffers that merged onto one index table + //! share an entry and so share this pointer; buffers of one request do not. + int32_t const* pageTable; + //! This buffer's sub-region of its pool in the cell arena, which is the pool's base + //! plus its slot there. A lease's cells follow at `leaseIndex * leaseStride`. + CUdeviceptr deviceBase; + + uint32_t hostPageStride; + uint32_t pageTableStride; + //! Blocks a row's page table describes; an ordinal past them is outside this call. + uint32_t numBlocks; + //! Bytes one lease takes in this pool's region, which is its cells there times the + //! cell size. Multiplied by a lease index, so widened at the multiply. + uint32_t leaseStride; + //! Bytes of one cell, which is one buffer's worth of pages. + uint32_t cellBytes; + + SelectionUnitLayout layout; + + SparseCopyWidth width; +}; + +static_assert(sizeof(SparseCopyDesc) == 64, "A copy descriptor is meant to be four vector loads"); + +//! Move units between their cold pages and their cache entries. +//! +//! The default codec relocates buffer bytes verbatim, so a unit has the same internal +//! shape on both sides and one layout describes both; only the buffer's offset within a +//! page and the unit's index differ. See "Fast path: a layout value" in +//! SPARSE_ATTENTION_DESIGN.md. +//! +//! Both directions use the same addressing: the cold side is reached through the page +//! table by unit ordinal, the hot side through the cache page list by entry index. Only +//! which of the two is the source changes. +struct SparseUnitCopyArgs +{ + //! [numRows] the arena slot each row holds; a batched copy spans rows whose cells + //! belong to different requests and so are not contiguous. kNoSparseLease means the + //! row holds no reservation. + SparseLeaseIndex const* leases; + //! The life cycle's arena: where a lease's cells sit, and the per-buffer pool slots. + SparseArenaView arena; + int32_t numRows; + + //! One descriptor per BufferId of the life cycle, indexed by `blockIdx.y` directly. + SparseCopyDesc const* copyDesc; + + //! The buffers this call covers: `numBuffers` of them from `bufferBase`, which for a + //! fetch is one layer's and for a write-back is the life cycle's. Consecutive, because + //! buffers are grouped by layer and, within a layer, by pool. + int32_t bufferBase; + int32_t numBuffers; + + SparseCopyDirection direction; + + //! Selection units one page holds. A life cycle's property, so every pool shares it. + int32_t unitsPerPage; + + //! The work list: the resolve's misses when fetching, the flush list when writing back. + //! Compacted, so a warp takes the next entry rather than the next slot of a row that may + //! hold nothing. The pools of one layer share it, because they hold the same units. + SparseWorkList work; + + //! CTAs along the work dimension. A work item is one warp's worth -- a selection unit is + //! 288 bytes of NVFP4 data or 36 of scales, which is 18 or 36 lanes, so a block per item + //! would leave most of it idle -- and the warps walk the list with a grid-stride loop. + //! + //! Sized from the machine, not the work: the fetch is captured in a graph, so the grid + //! is fixed before anyone knows how many units a step will miss. The device-side total + //! is what the loop actually follows. + int32_t gridCtas; +}; + +//! Widest vector a pool's layout admits, from its sizes alone. A layout admitting none +//! falls to bytes, so there is never a tail to handle in the kernel. +[[nodiscard]] SparseCopyWidth sparseCopyWidth( + SelectionUnitLayout const& layout, uint64_t hostPageStride, uint64_t bufferBytes); + +//! Bytes one `width` moves, for the caller that has to check its bases against it. +[[nodiscard]] size_t sparseCopyWidthBytes(SparseCopyWidth width) noexcept; + +//! Launch the copies for every pool at once: the grid is `(gridCtas, numBuffers)`, where the +//! buffers of every pool form one flat range. +//! +//! One launch rather than one per pool, because nothing orders the pools against each other +//! -- their cells are disjoint and the converted table they share is read-only -- and a +//! single grid both runs them concurrently and lets every one of them carry the programmatic +//! dependency on the kernel that produced the work list. +//! +//! Enrolled in PDL when the environment enables it. The kernel resolves its pool, buffer and +//! list before waiting, so a block is resident and addressed by the time its producer drains. +void launchSparseUnitCopy(SparseUnitCopyArgs const& args, CUstream stream); + +//! Turn a resolve's entry indices into the index list an attention kernel can read. +//! +//! The entry index the resolve produced is private to the cache: it names a unit within a +//! layer's partition. What a consumer needs is a value in the expanded index space, +//! `(page * buffersPerSlot + cell) * unitsPerPage + unit`. +//! +//! The result is compact -- one entry per selection slot, in the selection's own order -- +//! which is the shape every sparse attention consumer takes: a token-sparse kernel reads it +//! in place of its page table, a block-sparse kernel beside one. +//! +//! The buffer dimension cannot be collapsed. A layer's buffers share the resolve but not +//! their cells, because the cache packs one buffer's entries densely across every cell of +//! the pages it reserved rather than leaving each at its nominal cell of a coalesced slot. +//! So one selection resolves to a different index per buffer. +struct SparseOutIndexArgs +{ + //! [numRows] the arena slot each row holds; a row holding kNoSparseLease is filled + //! with kBadPageIndex and nothing else, which is what a free row should look like. + SparseLeaseIndex const* leases; + //! The life cycle's arena, which turns a lease into the cells its indices name. + SparseArenaView arena; + int32_t numRows; + + //! [numRows][selectionWidth] the ordinals selected, and the entries the resolve gave + //! them. Negative in either array means the slot carried no selection. + int32_t const* selected; + int32_t const* outSlot; + int32_t selectionWidth; + + //! Row stride of `outSlot`, when it is wider than the part this call reads. Zero means + //! `selectionWidth`. The write-index path needs it: its slots come from the resolve's + //! in-flight array, whose row is the reservation's headroom, while only the units + //! actually in flight -- at most one -- are read out. + int32_t outSlotStride = 0; + + //! [numRows][outBuffers][selectionWidth] written in full: entry `k` is the index + //! holding `selected[row][k]` for the buffer, and kBadPageIndex where the selection is + //! padded or the row holds no reservation. Padding stays where the selection put it + //! rather than being packed out, so slot `k` of the output and slot `k` of the + //! selection describe the same unit. + //! + //! One call fills `numBuffers` of the row's `outBuffers` entries, starting at `outBase`, + //! because a call serves one pool while the caller asks for a whole layer. + int32_t* out; + int32_t outBuffers; + int32_t outBase; + + //! The buffers this call writes, as for the copy. + int32_t bufferBase; + int32_t numBuffers; + + int32_t cellsPerBuffer; + int32_t cellsPerLease; + int32_t unitsPerPage; +}; + +//! Launch the index scatter. One CTA per (row, buffer). +void launchSparseOutIndices(SparseOutIndexArgs const& args, CUstream stream); + +//! Collect the entries holding units the watermark has just passed, so they can be written +//! back to their host pages. +//! +//! Found by scanning the reverse map rather than probing: the units wanted are a contiguous +//! ordinal range, and the map has to be read in full either way. +struct SparseFlushListArgs +{ + //! [numRows] the arena slot each row holds. kNoSparseLease means the row holds no + //! reservation and is skipped, as in the fetch. + SparseLeaseIndex const* leases; + //! The life cycle's arena, which every row's lease is addressed within. + SparseArenaView arena; + int32_t numRows; + + //! [numRows] each: the unit ordinals that row's watermark moved across, as [from, to). + //! Device resident for the same reason the resolve's watermarks are. + int32_t const* fromUnits; + int32_t const* toUnits; + + //! Units to write back and the entries holding them: one list per layer, because the + //! copy that consumes it walks buffers and a buffer belongs to one layer. Batch-owned, + //! so one launch covering every row hands the copy one pointer -- the reservation in a + //! row's lease is its own, but this is the step's. + SparseWorkList work; + //! [numRows][numLayers] entries this (row, layer) contributed, for the same reason the + //! resolve keeps its miss count: the copy does not need it, an observer does. + int32_t* outCount; +}; + +//! Launch the flush-list scan. One CTA per (layer, row). +void launchSparseFlushList(SparseFlushListArgs const& args, CUstream stream); + +//! Launch the resolve. One CTA per row; enqueues and returns. +//! +//! Rejects shapes that are inconsistent or need more shared memory than the device allows. +void launchSparseResolve(SparseResolveArgs const& args, CUstream stream); + +} // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/stagingBuffer.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/stagingBuffer.cpp index e7414f403d44..89b9975f862d 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/stagingBuffer.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/stagingBuffer.cpp @@ -35,14 +35,7 @@ namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 namespace { -using MemoryOwner = std::variant; - -CudaUniqPtr allocateDeviceMemory(size_t size) -{ - CUdeviceptr ptr = 0; - cuCheck(cuMemAlloc(&ptr, size)); - return CudaUniqPtr{reinterpret_cast(ptr)}; -} +using MemoryOwner = std::variant>; MemoryOwner makeMemoryOwner(size_t size, StagingBufferMemory memory) { @@ -50,14 +43,15 @@ MemoryOwner makeMemoryOwner(size_t size, StagingBufferMemory memory) switch (memory) { case StagingBufferMemory::kPinnedHost: return MemoryOwner{std::in_place_type, size}; - case StagingBufferMemory::kDevice: return MemoryOwner{std::in_place_type, allocateDeviceMemory(size)}; + case StagingBufferMemory::kDevice: + return MemoryOwner{std::in_place_type>, allocateDeviceMemory(size)}; default: throw std::invalid_argument("StagingBufferManager received an invalid memory kind"); } } } // namespace -void DeviceDeleter::operator()(std::byte* ptr) const noexcept +void DeviceDeleter::operator()(void* ptr) const noexcept { if (ptr != nullptr) { @@ -174,7 +168,7 @@ MemAddress StagingBufferManager::baseAddress() const noexcept { return host->address(); } - return reinterpret_cast(std::get(mMemoryOwner).get()); + return reinterpret_cast(std::get>(mMemoryOwner).get()); } StagingBufferMemory StagingBufferManager::memory() const noexcept diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/stagingBuffer.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/stagingBuffer.h index 942f33964df8..72a0b11c464e 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/stagingBuffer.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/stagingBuffer.h @@ -41,10 +41,24 @@ enum class StagingBufferMemory struct DeviceDeleter { - void operator()(std::byte* ptr) const noexcept; + void operator()(void* ptr) const noexcept; }; -using CudaUniqPtr = std::unique_ptr; +template +using CudaUniqPtr = std::unique_ptr; + +template +CudaUniqPtr allocateDeviceMemory(size_t size) +{ + // Null for an empty request, which cuMemAlloc refuses: a caller whose size is a + // computed extent can then allocate unconditionally and hold null for "none". + CUdeviceptr ptr = 0; + if (size != 0) + { + cuCheck(cuMemAlloc(&ptr, sizeof(T) * size)); + } + return CudaUniqPtr{reinterpret_cast(ptr)}; +} struct StagingBufferRange; @@ -133,7 +147,7 @@ class StagingBufferManager void retire(StagingBufferRange* range, CachedCudaEvent readyEvent) noexcept; size_t mSize; - std::variant mMemoryOwner; + std::variant> mMemoryOwner; // Full spatial partition; retired fragments retain their previous temporal owner event. std::list mRanges; std::list::iterator mHead; diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/config.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/config.cpp index a5d17a7817e1..b4dcd77c153b 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/config.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/config.cpp @@ -154,12 +154,14 @@ StorageConfig createStorageConfig(KVCacheManagerConfig const& config) LifeCycleRegistry registry{config}; int tokensPerBlock = config.tokensPerBlock; - // Map: lifeCycleId → ((coalescingGroup, expanded buffer size) → list of BufferId). + // Map: lifeCycleId → ((coalescingGroup, expanded buffer size, sparsity) → list of BufferId). // NOTE: Python uses insertion-ordered dict/defaultdict here. The std::map // grouping below intentionally remains sorted, which can make observable // pool-group indices/layout order differ from Python. Later lookups use // StorageConfig mappings, so this is not expected to affect correctness. - using CoalescingKey = std::pair; + // Keyed on sparsity as well as group and size: a pool is uniformly sparse or uniformly + // dense, so one layer's buffers may land in both. + using CoalescingKey = std::tuple; std::map>> bufferGroups; std::unordered_map expansionMap; @@ -179,7 +181,7 @@ StorageConfig createStorageConfig(KVCacheManagerConfig const& config) BufferId bid{cfg.layerId, buf.role}; expansionMap[bid] = exp; - bufferGroups[lcId][{buf.coalescingGroup, expandedSize}].push_back(bid); + bufferGroups[lcId][{buf.coalescingGroup, expandedSize, buf.isSparse}].push_back(bid); } }, layer); @@ -195,28 +197,33 @@ StorageConfig createStorageConfig(KVCacheManagerConfig const& config) var.lifeCycleId = lcId; for (auto const& [key, bufIds] : coalescingMap) { + [[maybe_unused]] auto const& [coalescingGroup, sz, isSparse] = key; CoalescedBuffer cb; - cb.singleBufferSize = key.second; + cb.singleBufferSize = sz; + cb.isSparse = isSparse; cb.bufferIds = bufIds; var.coalescedBuffers.push_back(std::move(cb)); } - // Sort descending by size. - std::sort(var.coalescedBuffers.begin(), var.coalescedBuffers.end(), + // Sort descending by size. Stable, so pools of equal size keep the order their + // keys gave them -- by coalescing group, then sparsity -- which `size()` alone + // does not distinguish. + std::ranges::stable_sort(var.coalescedBuffers, [](CoalescedBuffer const& a, CoalescedBuffer const& b) { return a.size() > b.size(); }); slotGroups.push_back(std::move(var)); } - // Distinct sparsity and reuse policies need separate GPU pools and eviction queues. + // Distinct reuse policies need separate GPU pools and eviction queues. Keying on + // per-pool sparsity too keeps a shared pool uniformly host-backed or device-backed. struct PoolGroupKey { AttentionReusePolicy reusePolicy = AttentionReusePolicy::REQUIRED; std::vector slotSizes; - bool isSparse = false; + std::vector sparsity; bool operator<(PoolGroupKey const& other) const { - return std::tie(reusePolicy, slotSizes, isSparse) - < std::tie(other.reusePolicy, other.slotSizes, other.isSparse); + return std::tie(reusePolicy, slotSizes, sparsity) + < std::tie(other.reusePolicy, other.slotSizes, other.sparsity); } }; @@ -225,13 +232,20 @@ StorageConfig createStorageConfig(KVCacheManagerConfig const& config) { auto const sizes = sg.slotSizeList(); auto const* attn = std::get_if(®istry.getLifeCycle(sg.lifeCycleId)); - bool const isSparse = attn != nullptr && attn->isSparse; auto const policy = attn ? attn->reusePolicy : AttentionReusePolicy::REQUIRED; - poolGroups[{.reusePolicy = policy, .slotSizes = sizes.raw(), .isSparse = isSparse}].push_back(std::move(sg)); + TypedVec sparsity; + sparsity.reserve(sizes.size()); + for (PoolIndex pool{0}; pool < sizes.size(); ++pool) + { + sparsity.push_back(sg.isSparse(pool)); + } + poolGroups[{.reusePolicy = policy, .slotSizes = sizes.raw(), .sparsity = sparsity.raw()}].push_back( + std::move(sg)); } StorageConfig out; out.cacheTiers = TypedVec{config.cacheTiers}; + out.hostSparsePoolQuota = config.hostSparsePoolQuota; out.expansion = expansionMap; for (auto& [key, variants] : poolGroups) diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/config.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/config.h index ef7fee8ec1b6..8be1a9eccf6e 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/config.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/config.h @@ -88,6 +88,19 @@ struct BufferId } }; +//! Buffers a consumer must address from one table, under one base pointer. +//! +//! Some attention ops take a single block table covering K and V rather than one each. +//! Order defines the table's buffer dimension: a kernel reads buffer `i` at entry `i`. +//! +//! All buffers must share a life cycle and a page-index scale; one table cannot serve +//! buffers whose converted indices differ. A spec spanning either is refused when the +//! Batch is constructed. +struct IndexTableSpec +{ + std::vector buffers; +}; + struct BufferIdHash { size_t operator()(BufferId const& id) const noexcept @@ -108,6 +121,12 @@ struct CoalescedBuffer size_t singleBufferSize = 0; std::vector bufferIds; + //! Whether this pool's pages live in the sparse GPU cache: host memory, device-mapped, + //! and not charged to the GPU tier's quota. + //! + //! Per pool, not per life cycle -- one layer may hold both a sparse buffer and a dense one. + bool isSparse = false; + [[nodiscard]] size_t size() const noexcept { return singleBufferSize * bufferIds.size(); @@ -128,6 +147,11 @@ struct SlotDescVariant LifeCycleId lifeCycleId{0}; TypedVec coalescedBuffers; // sorted size desc + [[nodiscard]] bool isSparse(PoolIndex poolIdx) const + { + return coalescedBuffers.at(poolIdx).isSparse; + } + // Slot size for each pool in this group. [[nodiscard]] TypedVec slotSizeList() const { @@ -154,6 +178,13 @@ struct SlotDesc { return getUniformAttribute(variants, [](auto const& v) { return v.slotSizeList(); }); } + + //! Life cycles share a pool group only when their pools agree position-wise on + //! sparsity as well as size, so this is uniform across variants. + [[nodiscard]] bool isSparse(PoolIndex poolIdx) const + { + return getUniformAttribute(variants, [poolIdx](auto const& v) { return v.isSparse(poolIdx); }); + } }; // --------------------------------------------------------------------------- @@ -221,6 +252,9 @@ struct StorageConfig TypedVec cacheTiers; TypedVec slotDescList; + //! Host bytes the GPU tier's host-backed pools may take, or unset for unbounded. + std::optional hostSparsePoolQuota; + // Expansion factor per buffer (for heterogeneous tokens_per_block). std::unordered_map expansion; diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/core.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/core.cpp index 0e23097688fc..0f9ce55daf0e 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/core.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/core.cpp @@ -17,6 +17,8 @@ #include "kv_cache_manager_v2/storage/core.h" +#include "kv_cache_manager_v2/utils/funcGuard.h" + #include "tensorrt_llm/common/assert.h" #include #include @@ -26,8 +28,10 @@ #include #include #include +#include #include #include +#include #include #include @@ -39,13 +43,19 @@ size_t tierAllocationGranularity(CacheTier tier, size_t quota) switch (tier) { case CacheTier::GPU_MEM: - case CacheTier::HOST_MEM: { + // Device memory is mapped and unmapped in this unit and never released any other + // way, so a large tier wants a large one: the mappings, not the bytes, are what a + // bigger quota would otherwise multiply. constexpr size_t kPageSize = 2ULL << 20; size_t const ratio = quota / (kPageSize * 512); int const exponent = ratio == 0 ? 0 : std::min(4, static_cast(std::log2(ratio))); return kPageSize << exponent; } + // Host memory has no such constraint: madvise and mprotect take any page-aligned range, + // so the unit is a fill chunk rather than a quantum and one size serves every tier. The + // backing raises it to a page should a platform ever ask for more. + case CacheTier::HOST_MEM: return hostMemSizeGranularity(); case CacheTier::DISK: return kDiskPoolSizeGranularity; default: throw std::invalid_argument("Invalid cache tier"); } @@ -355,12 +365,6 @@ void GpuSlotPool::resize(SlotCount newNumSlots) mVirtMem.realloc(physSize * newPhysMem); } -SlotCount GpuSlotPool::extendByOnePhysMem() -{ - mVirtMem.extend(1); - return numSlots(); -} - Address GpuSlotPool::slotAddress(SlotId slot) const { TLLM_CHECK_DEBUG_WITH_INFO(isValidSlotId(slot), "GpuSlotPool::slotAddress: slot index out of bounds"); @@ -371,25 +375,18 @@ Address GpuSlotPool::slotAddress(SlotId slot) const // HostSlotPool // --------------------------------------------------------------------------- -//! The unit this pool's memory is registered and released in. -//! -//! A whole number of slots, because a slot must not straddle two units: each -//! unit is a separate registration, the driver rejects a copy that crosses one, -//! and every copy here is a whole slot. Nothing requires it to be a whole number -//! of pages -- the backing rounds its madvise and mprotect calls instead. -size_t hostCommitUnit(size_t tierUnit, size_t slotSize) -{ - return roundUp(std::max(tierUnit, slotSize), slotSize); -} - HostSlotPool::HostSlotPool(size_t slotSize, SlotCount numSlots, size_t vmSize, HostMemBackingOptions const& options) : SlotPoolBase(slotSize) - , mCommitUnit(hostCommitUnit(options.commitUnit, slotSize)) - // The pool does not wait for the fill: HostPoolGroup blocks per allocation - // on only the range it hands out, so serving can start before the whole - // tier is committed. + , mCommitUnit(hostMemSizeGranularity()) + // registerUnit = slotSize: every copy here is a whole slot and the driver rejects + // one that crosses a registration, so the backing keeps its registration bounds on + // slot multiples -- which leaves the commit unit free of the slot size. + // + // waitForFill = false: the owning pool group blocks per allocation on only the range + // it hands out, so serving can start before the whole tier is committed. , mHostMem(std::max(vmSize, alignedSize(numSlots)), alignedSize(numSlots), - HostMemBackingOptions{mCommitUnit, options.allowRemoteNumaFallback}, /*waitForFill=*/false) + HostMemBackingOptions{.registerUnit = slotSize, .allowRemoteNumaFallback = options.allowRemoteNumaFallback}, + /*waitForFill=*/false) { } @@ -533,12 +530,36 @@ SlotCount PoolGroupBase::numSlots() const noexcept Slot PoolGroupBase::allocate() { - return mSlotAllocator.allocate(); + Slot slot = mSlotAllocator.allocate(); + // A failed fill surfaces out of _waitSlotUsable, after the slot is already taken, so it + // is returned to the allocator before the error propagates. + auto rollback = FuncGuard([&]() { release(std::move(slot)); }); + _waitSlotUsable(slot.slotId()); + rollback.cancel(); + return slot; +} + +void PoolGroupBase::_waitSlotUsable(SlotId slot) const +{ + for (auto const& pool : mPools) + { + pool->waitSlotUsable(slot); + } } std::vector PoolGroupBase::allocateMultiple(SlotCount numSlots) { - return mSlotAllocator.allocateMultiple(numSlots); + if (numSlots == SlotCount{0}) + { + return {}; + } + std::vector slots = mSlotAllocator.allocateMultiple(numSlots); + auto rollback + = FuncGuard([&]() { std::ranges::for_each(slots, [this](Slot& slot) { release(std::move(slot)); }); }); + // Slot ids are not returned in order, so the highest one bounds them all. + _waitSlotUsable(std::ranges::max(slots | std::views::transform(&Slot::slotId))); + rollback.cancel(); + return slots; } void PoolGroupBase::release(Slot slot) @@ -594,10 +615,17 @@ TypedVec PoolGroupBase::slotSize() const // GpuPoolGroup // --------------------------------------------------------------------------- -GpuPoolGroup::GpuPoolGroup( - SlotCount numSlots, TypedVec const& slotSizeList, PooledPhysMemAllocator& physMemAllocator) +GpuPoolGroup::GpuPoolGroup(SlotCount numSlots, TypedVec const& slotSizeList, + PooledPhysMemAllocator& physMemAllocator, TypedVec sparsity, + HostMemBackingOptions const& hostOptions, std::optional sparseReservation) : PoolGroupBase(numSlots) { + if (sparsity.empty()) + { + sparsity = TypedVec(slotSizeList.size(), false); + } + TLLM_CHECK_WITH_INFO( + sparsity.size() == slotSizeList.size(), "GpuPoolGroup: one sparsity flag per pool is required"); size_t physMemSize = physMemAllocator.physMemSize(); // Query total GPU memory to size virtual address space (mirrors Python GpuPoolGroup). size_t totalGpuMem = 0; @@ -606,16 +634,43 @@ GpuPoolGroup::GpuPoolGroup( cuCheck(cuCtxGetDevice(&dev)); cuCheck(cuDeviceTotalMem(&totalGpuMem, dev)); } - // @TODO: We should replace maxSlotSize with sum. This should also be updated in Python. Will do it later. TLLM_CHECK_WITH_INFO(!slotSizeList.empty(), "GpuPoolGroup: slotSizeList must not be empty"); - size_t maxSlotSize = *std::max_element(slotSizeList.begin(), slotSizeList.end()); - for (size_t sz : slotSizeList) + // Every pool reserves address space for the most slots this group could ever hold: the + // count all of device memory would buy. A slot is one row spanning every pool, and only + // the device-backed pools draw on that memory, so they alone set the count. No quota can + // exceed device memory, so this reservation is a bound the allocator cannot reach. + size_t deviceSlotBytes = 0; + for (PoolIndex p{0}; p < slotSizeList.size(); ++p) { - // VA proportional to GPU memory, scaled by slot size ratio (mirrors Python). - // Compute ratio as double first to avoid size_t overflow. - double sizeRatio = static_cast(sz) / static_cast(maxSlotSize); - size_t vmSize = roundDown(static_cast(static_cast(totalGpuMem) * sizeRatio), physMemSize); - vmSize = std::max(vmSize, roundUp(slotCountToSizeT(numSlots) * sz, physMemSize)); + if (!sparsity[p]) + { + deviceSlotBytes += slotSizeList.at(p); + } + } + TLLM_CHECK_WITH_INFO(deviceSlotBytes > 0, + "GpuPoolGroup: a wholly sparse group has no device bytes per slot to size its reservation from"); + size_t const maxSlots = totalGpuMem / deviceSlotBytes; + for (PoolIndex p{0}; p < slotSizeList.size(); ++p) + { + size_t const sz = slotSizeList.at(p); + size_t const grain = sparsity[p] ? hostMemSizeGranularity() : physMemSize; + size_t reserved = maxSlots * sz; + if (sparsity[p] && sparseReservation.has_value()) + { + reserved = std::min(reserved, *sparseReservation); + } + size_t const vmSize = std::max(roundUp(reserved, grain), roundUp(slotCountToSizeT(numSlots) * sz, grain)); + if (sparsity[p]) + { + // A sparse pool's pages: host memory, device-mapped, so the sparse cache fills + // from them without a host round trip. The reservation is host address space, + // not device, but it is as hard a bound: HostMem keeps a fixed base and rejects + // a resize past it. maxSlots alone would keep it unreachable, but scaled by a + // sparse slot it can exceed the address space, so `sparseReservation` caps it + // and growth past the cap is refused. + mPools.push_back(std::make_unique(sz, numSlots, vmSize, hostOptions)); + continue; + } mPools.push_back(std::make_unique(sz, vmSize, physMemAllocator, numSlots)); } } @@ -636,64 +691,15 @@ HostPoolGroup::HostPoolGroup(SlotCount numSlots, TypedVec con for (size_t sz : slotSizeList) { double const sizeRatio = static_cast(sz) / static_cast(maxSlotSize); - auto const poolVmSize = static_cast(static_cast(groupVmSize) * sizeRatio); + // Rounded here because the share is scaled by a ratio and so lands anywhere, while + // HostMem takes only whole granules. Up rather than down, and the shares may then + // sum past groupVmSize: this is address space, which costs nothing until committed. + auto const poolVmSize + = roundUp(static_cast(static_cast(groupVmSize) * sizeRatio), hostMemSizeGranularity()); mPools.push_back(std::make_unique(sz, numSlots, poolVmSize, options)); } } -void HostPoolGroup::waitSlotUsable(SlotId slot) const -{ - for (auto const& pool : mPools) - { - static_cast(*pool).waitSlotUsable(slot); - } -} - -// A failed fill surfaces out of waitSlotUsable, after the slots are already -// taken, so they are returned to the allocator before the error propagates. -Slot HostPoolGroup::allocate() -{ - Slot slot = PoolGroupBase::allocate(); - try - { - waitSlotUsable(slot.slotId()); - } - catch (...) - { - release(std::move(slot)); - throw; - } - return slot; -} - -std::vector HostPoolGroup::allocateMultiple(SlotCount numSlots) -{ - std::vector slots = PoolGroupBase::allocateMultiple(numSlots); - if (slots.empty()) - { - return slots; - } - // Slot ids are not returned in order, so the highest one bounds them all. - SlotId highest = slots.front().slotId(); - for (auto const& slot : slots) - { - highest = std::max(highest, slot.slotId()); - } - try - { - waitSlotUsable(highest); - } - catch (...) - { - for (auto& slot : slots) - { - release(std::move(slot)); - } - throw; - } - return slots; -} - // --------------------------------------------------------------------------- // DiskPoolGroup // --------------------------------------------------------------------------- @@ -737,15 +743,24 @@ TypedVec CacheLevelStorage::slotAddress(PoolGroupIndex pgIdx // --------------------------------------------------------------------------- GpuCacheLevelStorage::GpuCacheLevelStorage(TypedVec const& slotDescList, - TypedVec const& slotCountList, PooledPhysMemAllocator& physMemAllocator) + TypedVec const& slotCountList, PooledPhysMemAllocator& physMemAllocator, + HostMemBackingOptions const& hostOptions, std::optional sparseReservation) : mPhysMemAllocator(physMemAllocator) { TLLM_CHECK_WITH_INFO(slotCountList.size() == slotDescList.size(), "GpuCacheLevelStorage: slotCountList and slotDescList must have the same length"); for (PoolGroupIndex pgIdx{0}; pgIdx < slotDescList.size(); ++pgIdx) { + auto const& slotDesc = slotDescList[pgIdx]; + auto const sizes = slotDesc.slotSizeList(); + TypedVec sparsity; + sparsity.reserve(sizes.size()); + for (PoolIndex pool{0}; pool < sizes.size(); ++pool) + { + sparsity.push_back(slotDesc.isSparse(pool)); + } mPoolGroups.push_back(std::make_unique( - slotCountList[pgIdx], slotDescList[pgIdx].slotSizeList(), mPhysMemAllocator)); + slotCountList[pgIdx], sizes, mPhysMemAllocator, std::move(sparsity), hostOptions, sparseReservation)); } } @@ -762,7 +777,7 @@ HostCacheLevelStorage::HostCacheLevelStorage(TypedVec // unit before they are compared, so the bound is quantized the same way: a // bound that sat between two multiples would reject the very quota the // pools had already reserved for. - , mMaxQuota(roundUp(maxQuota.value_or(hostTotalMemory()), options.commitUnit)) + , mMaxQuota(roundUp(maxQuota.value_or(hostTotalMemory()), hostMemSizeGranularity())) { TLLM_CHECK_WITH_INFO(slotCountList.size() == slotDescList.size(), "HostCacheLevelStorage: slotCountList and slotDescList must have the same length"); @@ -796,7 +811,8 @@ DiskCacheLevelStorage::DiskCacheLevelStorage(TypedVec std::unique_ptr createCacheLevelStorage(CacheTierConfig const& tierCfg, TypedVec const& slotDescList, TypedVec const& slotCountList, - PooledPhysMemAllocator* gpuPhysMemAllocator, HostMemBackingOptions const& options) + PooledPhysMemAllocator* gpuPhysMemAllocator, HostMemBackingOptions const& options, + std::optional sparseReservation) { TLLM_CHECK((cacheTierOf(tierCfg) == CacheTier::GPU_MEM) == (gpuPhysMemAllocator != nullptr)); return std::visit( @@ -805,7 +821,8 @@ std::unique_ptr createCacheLevelStorage(CacheTierConfig const using T = std::decay_t; if constexpr (std::is_same_v) { - return std::make_unique(slotDescList, slotCountList, *gpuPhysMemAllocator); + return std::make_unique( + slotDescList, slotCountList, *gpuPhysMemAllocator, options, sparseReservation); } else if constexpr (std::is_same_v) { @@ -826,21 +843,21 @@ std::unique_ptr createCacheLevelStorage(CacheTierConfig const // Mirrors Python CacheLevelStorage._grains_to_slots. // --------------------------------------------------------------------------- -std::pair CacheLevelStorage::grainsToSlots( - size_t pgGrains, TypedVec const& slotSizeList, size_t granularity) +std::pair CacheLevelStorage::grainsToSlots(size_t pgGrains, + TypedVec const& slotSizeList, std::vector const& which, size_t granularity) { TLLM_CHECK_WITH_INFO(granularity > 0, "Cache storage granularity must be positive"); - TypedVec minPoolGrains(slotSizeList.size()); - for (PoolIndex poolIdx{0}; poolIdx < slotSizeList.size(); ++poolIdx) + // With nothing to divide the grains between there is no slot count to find, and the + // search below would run past SlotCount's range looking for one. + TLLM_CHECK_WITH_INFO(!which.empty(), "grainsToSlots needs at least one pool to size against"); + // Indexed by real pool index; entries outside `which` stay zero and are never read. + TypedVec minPoolGrains(slotSizeList.size(), 0); + size_t minTotal = 0; + for (PoolIndex poolIdx : which) { TLLM_CHECK_WITH_INFO(slotSizeList[poolIdx] > 0, "Cache slot size must be positive"); minPoolGrains[poolIdx] = divUp(slotSizeList[poolIdx], granularity); - } - - size_t minTotal = 0; - for (auto g : minPoolGrains) - { - minTotal += g; + minTotal += minPoolGrains[poolIdx]; } if (pgGrains < minTotal) return {0, 0}; @@ -848,13 +865,9 @@ std::pair CacheLevelStorage::grainsToSlots( SlotCount numSlots{std::numeric_limits::max()}; size_t remainingPgGrains = pgGrains; - // Sort pools by slot size ascending. - std::vector poolOrder; - poolOrder.reserve(slotSizeList.stdSize()); - for (PoolIndex poolIdx{0}; poolIdx < slotSizeList.size(); ++poolIdx) - { - poolOrder.push_back(poolIdx); - } + // Sort the named pools by slot size ascending. Ties may fall either way: every term + // below is a function of the sizes alone. + std::vector poolOrder = which; std::sort(poolOrder.begin(), poolOrder.end(), [&](PoolIndex a, PoolIndex b) { return slotSizeList[a] < slotSizeList[b]; }); @@ -876,7 +889,7 @@ std::pair CacheLevelStorage::grainsToSlots( TLLM_CHECK_DEBUG(remainingPgGrains == 0); TLLM_CHECK_DEBUG(numSlots > 0); - auto slotsToGrains = [&](SlotCount slots) { return grainsForSlots(slots, slotSizeList, granularity); }; + auto slotsToGrains = [&](SlotCount slots) { return grainsForSlots(slots, slotSizeList, which, granularity); }; SlotCount lo = numSlots; SlotCount step = 1; SlotCount hi = lo + step; @@ -910,12 +923,12 @@ std::pair CacheLevelStorage::grainsToSlots( // Mirrors Python CacheLevelStorage._grains_for_slots. // --------------------------------------------------------------------------- -size_t CacheLevelStorage::grainsForSlots( - SlotCount numSlots, TypedVec const& slotSizeList, size_t granularity) +size_t CacheLevelStorage::grainsForSlots(SlotCount numSlots, TypedVec const& slotSizeList, + std::vector const& which, size_t granularity) { size_t total = 0; - for (auto s : slotSizeList) - total += divUp(slotCountToSizeT(numSlots) * s, granularity); + for (PoolIndex poolIdx : which) + total += divUp(slotCountToSizeT(numSlots) * slotSizeList[poolIdx], granularity); return total; } @@ -924,143 +937,725 @@ size_t CacheLevelStorage::grainsForSlots( // Mirrors Python CacheLevelStorage.ratio_to_slot_count_list. // --------------------------------------------------------------------------- +TypedVec> CacheLevelStorage::poolsPaidByQuota( + TypedVec> const& slotSizeLists, + std::optional const& secondBudget) +{ + TypedVec> ret; + ret.reserve(slotSizeLists.size()); + for (PoolGroupIndex pgIdx{0}; pgIdx < slotSizeLists.size(); ++pgIdx) + { + auto const numPools = slotSizeLists[pgIdx].size(); + TypedVec taken(numPools, false); + if (secondBudget.has_value()) + { + for (PoolIndex poolIdx : secondBudget->pools[pgIdx]) + { + TLLM_CHECK_WITH_INFO(poolIdx < numPools, + "A second budget names pool %d of group %d, which has only %d pools", poolIdx.value(), + pgIdx.value(), numPools.value()); + // Naming a pool twice would charge it twice and leave the quota paying for + // one fewer than it thinks. + TLLM_CHECK_WITH_INFO( + !taken[poolIdx], "A second budget names pool %d of group %d twice", poolIdx.value(), pgIdx.value()); + taken[poolIdx] = true; + } + } + std::vector paid; + paid.reserve(toSizeT(numPools)); + for (PoolIndex poolIdx{0}; poolIdx < numPools; ++poolIdx) + { + if (!taken[poolIdx]) + { + paid.push_back(poolIdx); + } + } + ret.push_back(std::move(paid)); + } + return ret; +} + +namespace +{ + +//! Bytes one pool group holds at `slots`, each pool rounded to its own grain. This is the +//! quantity the pool ratio speaks in: what the group actually occupies, device and host +//! together. +size_t allBytesForSlots( + TypedVec const& sizes, TypedVec const& grains, SlotCount slots) +{ + auto const n = slotCountToSizeT(slots); + size_t total = 0; + for (PoolIndex poolIdx{0}; poolIdx < sizes.size(); ++poolIdx) + { + total += roundUp(n * sizes[poolIdx], grains[poolIdx]); + } + return total; +} + +//! Steps a group may take to correct itself before only a budget may carry it further. +constexpr int kSelfCorrectionLimit = 8; + +//! Two levels within this of each other are the same level. Paired with the float32 ratio: +//! the band sits four orders of magnitude above double's own noise, so no comparison below +//! needs exact arithmetic. +constexpr double kRatioEps = 1.0 / 16777216.0; + +//! The largest rung at or below `budget` bytes, which is free-run closed by construction. +//! +//! Bracketed, not galloped: `A(n) >= n * S` with `S` the summed slot size, so no count above +//! `budget / S` can fit and bisection has an upper end without searching for one. +SlotCount rungAtOrBelow( + TypedVec const& sizes, TypedVec const& grains, size_t budget) +{ + size_t sumSizes = 0; + for (PoolIndex poolIdx{0}; poolIdx < sizes.size(); ++poolIdx) + { + sumSizes += sizes[poolIdx]; + } + TLLM_CHECK_WITH_INFO(sumSizes > 0, "A pool group must have a positive slot size"); + SlotCount lo{0}; + auto hi = slotCountValueFromSize(budget / sumSizes + 1); + if (allBytesForSlots(sizes, grains, hi) <= budget) + { + return hi; + } + while (lo + 1 < hi) + { + SlotCount const mid = lo + ((hi - lo) / 2); + if (allBytesForSlots(sizes, grains, mid) <= budget) + { + lo = mid; + } + else + { + hi = mid; + } + } + return lo; +} + +//! The most slots the grains already bought for `slots` will hold, in every pool at once. +//! +//! The counterpart of PoolGroupBase::getNumSlotsFromPools, and what makes a count +//! self-consistent: growing to it costs nothing in any medium, so it never disturbs a +//! budget. +//! +//! Every pool is in view, including ones the second budget pays for. While that budget is +//! still paying for a pool its grains are as real a cost as the quota's, so a run that held +//! one of the two fixed while the other moved would not be free. +SlotCount freeRun(TypedVec const& sizes, TypedVec const& grains, SlotCount slots) +{ + auto const n = slotCountToSizeT(slots); + SlotCount best{std::numeric_limits::max()}; + for (PoolIndex poolIdx{0}; poolIdx < sizes.size(); ++poolIdx) + { + auto const g = grains[poolIdx]; + auto const sz = sizes[poolIdx]; + best = std::min(best, slotCountValueFromSize(divUp(n * sz, g) * g / sz)); + } + return best; +} + +//! The most slots the grains the quota has already been charged for will hold. +//! +//! The same run over the pools the quota pays for. Once the second budget's level has +//! apportioned all of it, what is left over buys its pools at no cost anyone is accounting +//! for, so those pools no longer belong in a free run: leaving them in pins the count to +//! the bottom of the range over which the quota's own cost is flat, at a position fixed by +//! their width -- the one thing the quota does not measure -- while every slot up to the +//! top of that range is already bought and sits unused. +SlotCount freeRunOnQuota( + TypedVec const& sizes, std::vector const& paid, size_t granularity, SlotCount slots) +{ + auto const n = slotCountToSizeT(slots); + SlotCount best{std::numeric_limits::max()}; + for (PoolIndex poolIdx : paid) + { + auto const sz = sizes[poolIdx]; + best = std::min(best, slotCountValueFromSize(divUp(n * sz, granularity) * granularity / sz)); + } + return best; +} + +} // namespace + TypedVec CacheLevelStorage::ratioToSlotCountList(size_t quota, - TypedVec> const& sizeLists, + TypedVec> const& slotSizeLists, TypedVec const& ratioList, size_t granularity, - TypedVec const& minSlots) + TypedVec const& minSlots, std::optional const& secondBudget) { - PoolGroupIndex numPg = sizeLists.size(); + PoolGroupIndex const numPg = slotSizeLists.size(); TLLM_CHECK_DEBUG(ratioList.size() == numPg); - TLLM_CHECK_DEBUG_WITH_INFO(std::all_of(ratioList.begin(), ratioList.end(), [](auto x) { return x > 0; }), - "ratioToSlotCountList: all ratios must be positive"); - TLLM_CHECK_DEBUG(quota % granularity == 0); - size_t totalGrains = quota / granularity; - if (TLLM_UNLIKELY(gDebug)) + TLLM_CHECK_WITH_INFO(granularity > 0, "Cache storage granularity must be positive"); + + bool const hasSecond = secondBudget.has_value(); + bool const bounded = hasSecond && secondBudget->quota.has_value(); + size_t const grain2 = hasSecond ? secondBudget->granularity : granularity; + TLLM_CHECK_WITH_INFO(!hasSecond || grain2 > 0, "A second budget needs a positive granularity"); + + auto const pools = poolsPaidByQuota(slotSizeLists, secondBudget); + static std::vector const kNoPools; + auto secondPools = [&](PoolGroupIndex pgIdx) -> std::vector const& + { return hasSecond ? secondBudget->pools[pgIdx] : kNoPools; }; + + for (PoolGroupIndex pgIdx{0}; pgIdx < numPg; ++pgIdx) { - [[maybe_unused]] size_t minGrains = 0; - for (auto const& sizes : sizeLists) - minGrains += toSizeT(sizes.size()); - TLLM_CHECK_WITH_INFO(totalGrains >= minGrains, - "ratioToSlotCountList: insufficient total grains for at least 1 slot per pool group"); + TLLM_CHECK_WITH_INFO(ratioList[pgIdx] > 0.f, "ratioToSlotCountList: all ratios must be positive"); + TLLM_CHECK_WITH_INFO(!pools[pgIdx].empty(), + "ratioToSlotCountList: pool group %d has no pool the primary quota pays for", pgIdx.value()); } + // One grain per pool, settled once: a pool's grain follows its medium, so the budget a + // pool belongs to decides only which total it lands in, never how it rounds. + TypedVec> grains; + grains.reserve(numPg); + for (PoolGroupIndex pgIdx{0}; pgIdx < numPg; ++pgIdx) + { + TypedVec pgGrains(slotSizeLists[pgIdx].size(), granularity); + for (PoolIndex poolIdx : secondPools(pgIdx)) + { + pgGrains[poolIdx] = grain2; + } + grains.push_back(std::move(pgGrains)); + } + + auto quotaGrains = [&](PoolGroupIndex pgIdx, SlotCount slots) + { return grainsForSlots(slots, slotSizeLists[pgIdx], pools[pgIdx], granularity); }; + auto secondGrains = [&](PoolGroupIndex pgIdx, SlotCount slots) + { return grainsForSlots(slots, slotSizeLists[pgIdx], secondPools(pgIdx), grain2); }; + auto allBytes = [&](PoolGroupIndex pgIdx, SlotCount slots) + { return allBytesForSlots(slotSizeLists[pgIdx], grains[pgIdx], slots); }; + + size_t const budgetGrains = quota / granularity; + size_t const budget2Grains = bounded ? *secondBudget->quota / grain2 : 0; + + // The floors are not negotiable, so they come off the top before anything is shared. + // Snapping each to its free run costs no grain in any medium, and starting there is what + // makes "never decrease" enough to keep them. TypedVec slotCntList(numPg, 0); - size_t remainingGrains = totalGrains; - std::vector activePgs(toSizeT(numPg)); - std::iota(activePgs.begin(), activePgs.end(), PoolGroupIndex{0}); - - // Iteratively peel off constrained PGs until all active PGs are - // unconstrained: - // 1. Distribute remaining quota among active PGs by ratio. - // 2. Any PG with slots <= min_slots is constrained — pin it to - // min_slots and subtract its grains from the budget. - // 3. Repeat with the remaining PGs and re-normalized ratios. - // Each iteration removes at least one PG, so this terminates. - while (!activePgs.empty()) - { - // Distribute remainingGrains among active PGs by ratio. - size_t nActive = activePgs.size(); - std::vector activeRatio(nActive); - for (size_t i = 0; i < nActive; ++i) - activeRatio[i] = ratioList[activePgs[i]]; - - std::vector slotsForActive(nActive, 0); - std::vector grainsForActive(nActive, 0); - size_t budget = remainingGrains; - - // Sort indices by ratio ascending. - std::vector idxLst(nActive); - std::iota(idxLst.begin(), idxLst.end(), size_t{0}); - std::sort(idxLst.begin(), idxLst.end(), [&](size_t a, size_t b) { return activeRatio[a] < activeRatio[b]; }); - - for (size_t i = 0; i < idxLst.size(); ++i) + size_t floorQuota = 0; + size_t floorSecond = 0; + for (PoolGroupIndex pgIdx{0}; pgIdx < numPg; ++pgIdx) + { + slotCntList[pgIdx] = freeRun(slotSizeLists[pgIdx], grains[pgIdx], std::max(minSlots[pgIdx], SlotCount{0})); + floorQuota += quotaGrains(pgIdx, slotCntList[pgIdx]); + floorSecond += secondGrains(pgIdx, slotCntList[pgIdx]); + } + if (bounded) + { + TLLM_CHECK_WITH_INFO(floorSecond <= budget2Grains, + "Second memory budget of %zu bytes cannot hold the %zu the batches in `constraints` need. Raise the " + "budget or lower the constraints.", + *secondBudget->quota, floorSecond * grain2); + } + TLLM_CHECK_WITH_INFO(floorQuota <= budgetGrains, "Insufficient quota to satisfy min_slots constraints"); + + size_t remainingQuota = budgetGrains - floorQuota; + size_t remainingSecond = bounded ? budget2Grains - floorSecond : 0; + + // The floors are where every rule below starts, and none of them may fall back to it. + TypedVec const floorSlots = slotCntList; + + // Normalised once, order-independently: a float sum that depended on the order of the + // pool groups would carry that order into every comparison below, and the tie band + // magnifies a last-ulp difference into a different answer. + TypedVec ratio(numPg, 0.0); + { + std::vector values; + values.reserve(toSizeT(numPg)); + for (PoolGroupIndex pgIdx{0}; pgIdx < numPg; ++pgIdx) + { + values.push_back(static_cast(ratioList[pgIdx])); + } + std::sort(values.begin(), values.end()); + long double total = 0.0L; + for (double v : values) { - size_t idx = idxLst[i]; - double ratioSum = 0.0; - for (size_t j = i; j < idxLst.size(); ++j) - ratioSum += static_cast(activeRatio[idxLst[j]]); - double pct = (ratioSum > 0.0) ? static_cast(activeRatio[idx]) / ratioSum : 1.0; - auto [slots, used] = grainsToSlots(static_cast(std::nearbyint(static_cast(budget) * pct)), - sizeLists[activePgs[idx]], granularity); - slotsForActive[idx] = slots; - grainsForActive[idx] = used; - TLLM_CHECK_DEBUG(used <= budget); - budget -= used; + total += v; } + for (PoolGroupIndex pgIdx{0}; pgIdx < numPg; ++pgIdx) + { + ratio[pgIdx] = static_cast(static_cast(ratioList[pgIdx]) / total); + } + } - // Identify constrained PGs (slots <= min_slots). - std::vector constrained; - std::vector unconstrained; - for (size_t idx = 0; idx < nActive; ++idx) + // A level is a byte count: the all-bytes total a group of ratio 1 would be entitled to. + // Because a group's occupied bytes and its two grain costs are the same quantity in + // different units, this one scalar drives both budgets. + auto nearestRung = [&](PoolGroupIndex pgIdx, size_t level) + { + double const target = static_cast(level) * ratio[pgIdx]; + SlotCount const nLo = rungAtOrBelow(slotSizeLists[pgIdx], grains[pgIdx], static_cast(target)); + if (nLo < floorSlots[pgIdx]) + { + return floorSlots[pgIdx]; + } + SlotCount const nHi = freeRun(slotSizeLists[pgIdx], grains[pgIdx], nLo + 1); + // Of the two rungs bracketing the target take the nearer; an exact midpoint stays low. + return 2.0 * target > static_cast(allBytes(pgIdx, nLo) + allBytes(pgIdx, nHi)) ? nHi : nLo; + }; + + // Largest level whose spend fits, and whether the search ran out of headroom without + // the spend ever exceeding the residue -- which is how "this budget cannot bind" is + // reported. `spend` must be non-decreasing, which `nearestRung` is. + auto highestLevelFitting = [](auto&& spend, size_t residue) + { + if (spend(size_t{0}) > residue) + { + return std::pair{0, false}; + } + size_t lo = 0; + size_t hi = 1; + bool exceeded = false; + for (int i = 0; i < 64; ++i) { - PoolGroupIndex pgIdx = activePgs[idx]; - SlotCount const minSlotCount = minSlots[pgIdx]; - if (slotsForActive[idx] <= minSlotCount) - constrained.push_back(idx); + if (spend(hi) > residue) + { + exceeded = true; + break; + } + lo = hi; + if (hi > std::numeric_limits::max() / 2) + { + break; + } + hi *= 2; + } + if (!exceeded) + { + return std::pair{hi, true}; + } + while (lo + 1 < hi) + { + size_t const mid = lo + ((hi - lo) / 2); + if (spend(mid) <= residue) + { + lo = mid; + } else - unconstrained.push_back(idx); + { + hi = mid; + } + } + return std::pair{lo, false}; + }; + + // Among the levels paying for the same amount, the lowest. They differ only for a group + // whose cost in that budget is flat over a range of counts -- a small dense pool beside + // a large host-backed one -- and the extra level buys such a group nothing the budget + // pays for, while pushing it off its target ratio. + auto leastLevelReaching = [](auto&& spend, size_t target, size_t upper) + { + size_t lo = 0; + size_t hi = upper; + while (lo < hi) + { + size_t const mid = lo + ((hi - lo) / 2); + if (spend(mid) >= target) + { + hi = mid; + } + else + { + lo = mid + 1; + } } + return lo; + }; - if (constrained.empty()) + TypedVec onSecond(numPg, false); + bool anyOnSecond = false; + for (PoolGroupIndex pgIdx{0}; pgIdx < numPg; ++pgIdx) + { + onSecond[pgIdx] = !secondPools(pgIdx).empty(); + anyOnSecond = anyOnSecond || onSecond[pgIdx]; + } + + // One level for the whole host-backed class: equal levels are what an equal ratio means, + // so holding the class at one level is what keeps the ratio within it. + size_t level2 = 0; + bool hasLevel2 = false; + bool saturated2 = true; + if (bounded && anyOnSecond) + { + auto spendSecond = [&](size_t level) { - // All active PGs are unconstrained — accept their allocations. - for (size_t idx = 0; idx < nActive; ++idx) - slotCntList[activePgs[idx]] = slotsForActive[idx]; - break; + size_t total = 0; + for (PoolGroupIndex pgIdx{0}; pgIdx < numPg; ++pgIdx) + { + if (onSecond[pgIdx]) + { + total += secondGrains(pgIdx, nearestRung(pgIdx, level)) - secondGrains(pgIdx, floorSlots[pgIdx]); + } + } + return total; + }; + auto const found = highestLevelFitting(spendSecond, remainingSecond); + level2 = found.first; + saturated2 = found.second; + hasLevel2 = true; + } + + auto capped = [&](PoolGroupIndex pgIdx, size_t level) + { + if (hasLevel2 && onSecond[pgIdx] && level > level2) + { + level = level2; } + return nearestRung(pgIdx, level); + }; - // Pin constrained PGs to min_slots and subtract from budget. - for (size_t idx : constrained) + auto spendPrimary = [&](size_t level) + { + size_t total = 0; + for (PoolGroupIndex pgIdx{0}; pgIdx < numPg; ++pgIdx) { - PoolGroupIndex pgIdx = activePgs[idx]; - SlotCount const minSlotCount = minSlots[pgIdx]; - size_t minGrains = grainsForSlots(minSlotCount, sizeLists[pgIdx], granularity); - auto [slots, used] = grainsToSlots(minGrains, sizeLists[pgIdx], granularity); - slotCntList[pgIdx] = slots; - // remainingGrains is unsigned, so an over-subtraction would wrap past the - // `remainingGrains == 0` guard below instead of reporting an exhausted budget. - TLLM_CHECK_WITH_INFO(used <= remainingGrains, "Insufficient quota to satisfy min_slots constraints"); - remainingGrains -= used; + total += quotaGrains(pgIdx, capped(pgIdx, level)) - quotaGrains(pgIdx, floorSlots[pgIdx]); } + return total; + }; - if (unconstrained.empty()) + auto const foundPrimary = highestLevelFitting(spendPrimary, remainingQuota); + size_t level1 = foundPrimary.first; + if (!foundPrimary.second) + { + level1 = leastLevelReaching(spendPrimary, spendPrimary(level1), level1); + } + + // The secondary binds only where the primary would have carried the class higher. + bool const secondaryBinds = hasLevel2 && !saturated2 && level1 > level2; + + auto levelFor + = [&](PoolGroupIndex pgIdx) { return hasLevel2 && onSecond[pgIdx] && level1 > level2 ? level2 : level1; }; + + // The rung at or below the level, not the nearer one: which groups round up is a + // comparison between groups, and a per-group rounding settles it by whichever crossing + // the ratio's last bits happen to put first. The promotions below compare the crossings + // against each other instead, and so can see that two are tied. + for (PoolGroupIndex pgIdx{0}; pgIdx < numPg; ++pgIdx) + { + double const target = static_cast(levelFor(pgIdx)) * ratio[pgIdx]; + SlotCount want = rungAtOrBelow(slotSizeLists[pgIdx], grains[pgIdx], static_cast(target)); + want = std::max(want, floorSlots[pgIdx]); + remainingQuota -= quotaGrains(pgIdx, want) - quotaGrains(pgIdx, slotCntList[pgIdx]); + if (bounded) { - // All PGs are constrained — nothing left to redistribute. - break; + remainingSecond -= secondGrains(pgIdx, want) - secondGrains(pgIdx, slotCntList[pgIdx]); } + slotCntList[pgIdx] = want; + } - if (remainingGrains == 0) - throw std::runtime_error("Insufficient quota to satisfy min_slots constraints"); + // Finish the grains the quota has already been charged for -- see freeRunOnQuota. This + // runs once the promotions have stopped, because leftover is what the two levels did + // not apportion, and that is not known until they have stopped spending. + std::vector classA; + for (PoolGroupIndex pgIdx{0}; pgIdx < numPg; ++pgIdx) + { + if (onSecond[pgIdx]) + { + classA.push_back(pgIdx); + } + } - // Continue with unconstrained PGs only. - std::vector newActivePgs; - newActivePgs.reserve(unconstrained.size()); - for (size_t idx : unconstrained) - newActivePgs.push_back(activePgs[idx]); - activePgs = std::move(newActivePgs); + // Summed order-independently, like the ratio normaliser above: a float sum that + // depended on the order of the pool groups would carry that order into the level below. + double classShare = 0.0; + { + std::vector values; + values.reserve(classA.size()); + for (PoolGroupIndex pgIdx : classA) + { + values.push_back(ratio[pgIdx]); + } + std::sort(values.begin(), values.end()); + long double total = 0.0L; + for (double v : values) + { + total += v; + } + classShare = static_cast(total); } - // _g2s may under-count slots due to imperfect grain distribution - // across pools. Try bumping each PG's slot count while it still fits - // within the same grain budget. - for (PoolGroupIndex pgIdx{0}; pgIdx < sizeLists.size(); ++pgIdx) + TypedVec tops(numPg, 0); + auto closePaidGrains = [&]() { - size_t grainsNow = grainsForSlots(slotCntList[pgIdx], sizeLists[pgIdx], granularity); - while (grainsForSlots(slotCntList[pgIdx] + 1, sizeLists[pgIdx], granularity) <= grainsNow) - slotCntList[pgIdx] += 1; + bool wanted = false; + for (PoolGroupIndex pgIdx{0}; pgIdx < numPg; ++pgIdx) + { + tops[pgIdx] = freeRunOnQuota(slotSizeLists[pgIdx], pools[pgIdx], granularity, slotCntList[pgIdx]); + // A group the second budget names no pool of is already at its top: its free run + // is over the quota's pools alone. + TLLM_CHECK_DEBUG(onSecond[pgIdx] || tops[pgIdx] == slotCntList[pgIdx]); + wanted = wanted || tops[pgIdx] > slotCntList[pgIdx]; + } + if (!wanted) + { + return false; + } + size_t need = 0; + for (PoolGroupIndex pgIdx : classA) + { + need += secondGrains(pgIdx, tops[pgIdx]) - secondGrains(pgIdx, slotCntList[pgIdx]); + } + if (!bounded || need <= remainingSecond) + { + if (bounded) + { + remainingSecond -= need; + } + for (PoolGroupIndex pgIdx : classA) + { + slotCntList[pgIdx] = tops[pgIdx]; + } + return true; + } + + // More closures are wanted than the leftover can buy. One level for the whole + // host-backed class, Webster rounded: the rule level2 divided the budget by, + // applied to what it did not divide, so the ratio inside that class is what decides + // where the leftover goes. A level is a function of the groups' own data, so which + // of them it reaches cannot depend on the order they were passed in. + auto closureTarget = [&](PoolGroupIndex pgIdx, size_t level) + { + double const share = static_cast(level) * ratio[pgIdx] / classShare; + SlotCount const nLo = rungAtOrBelow(slotSizeLists[pgIdx], grains[pgIdx], static_cast(share)); + if (nLo >= tops[pgIdx]) + { + return tops[pgIdx]; + } + SlotCount const nHi = freeRun(slotSizeLists[pgIdx], grains[pgIdx], nLo + 1); + SlotCount const want + = 2.0 * share > static_cast(allBytes(pgIdx, nLo) + allBytes(pgIdx, nHi)) ? nHi : nLo; + return std::min(tops[pgIdx], std::max(slotCntList[pgIdx], want)); + }; + auto closureSpend = [&](size_t level) + { + size_t total = 0; + for (PoolGroupIndex pgIdx : classA) + { + total += secondGrains(pgIdx, closureTarget(pgIdx, level)) - secondGrains(pgIdx, slotCntList[pgIdx]); + } + return total; + }; + size_t const level3 = highestLevelFitting(closureSpend, remainingSecond).first; + bool moved = false; + for (PoolGroupIndex pgIdx : classA) + { + SlotCount const want = closureTarget(pgIdx, level3); + if (want > slotCntList[pgIdx]) + { + remainingSecond -= secondGrains(pgIdx, want) - secondGrains(pgIdx, slotCntList[pgIdx]); + slotCntList[pgIdx] = want; + moved = true; + } + } + return moved; + }; + + // Each group's next step becomes the nearer of its two rungs at its own crossing level. + // Taking the steps in ascending crossing order is rounding to the nearest rung done + // across the groups rather than one at a time; carrying on past the level is what spends + // the fraction of a step the level search had to leave behind. + TypedVec corrections(numPg, 0); + + //! Everything that distinguishes one pool group from another, in a fixed order. + //! + //! Two groups agreeing on all of it are interchangeable, so the index that settles them + //! last picks between relabelings of one answer rather than between two answers. + //! Anything that distinguishes two groups and is missing from here is a way for the + //! answer to depend on the order they were passed in. + //! + //! What a slot costs each budget leads: of two groups equally entitled to a step, the + //! one whose slots are cheaper leaves more of the residue for whoever comes next, which + //! is the tie-break's first rule carried from one step to the whole ladder. The slot + //! sizes cannot say this on their own -- two groups can hold the same sizes and differ + //! in which of them the second budget pays for, and so in what a slot costs the quota. + struct Identity + { + size_t quotaRate{0}; + size_t secondRate{0}; + std::vector sizes; + std::vector secondPools; + SlotCount floor{0}; + + bool operator<(Identity const& other) const + { + return std::tie(quotaRate, secondRate, sizes, secondPools, floor) + < std::tie(other.quotaRate, other.secondRate, other.sizes, other.secondPools, other.floor); + } + }; + + TypedVec identities; + identities.reserve(numPg); + for (PoolGroupIndex pgIdx{0}; pgIdx < numPg; ++pgIdx) + { + Identity ident; + for (PoolIndex poolIdx : pools[pgIdx]) + { + ident.quotaRate += slotSizeLists[pgIdx][poolIdx]; + } + for (PoolIndex poolIdx : secondPools(pgIdx)) + { + ident.secondRate += slotSizeLists[pgIdx][poolIdx]; + ident.secondPools.push_back(toSizeT(poolIdx)); + } + std::sort(ident.secondPools.begin(), ident.secondPools.end()); + ident.sizes.reserve(toSizeT(slotSizeLists[pgIdx].size())); + for (PoolIndex poolIdx{0}; poolIdx < slotSizeLists[pgIdx].size(); ++poolIdx) + { + ident.sizes.push_back(slotSizeLists[pgIdx][poolIdx]); + } + ident.floor = minSlots[pgIdx]; + identities.push_back(std::move(ident)); } - return slotCntList; -} + //! One step a group could take, and everything needed to choose between them. + struct Candidate + { + double crossing; + PoolGroupIndex pgIdx; + SlotCount next; + bool correcting; + size_t d1; + size_t d2; + }; + + // A cheaper step leaves more residue for whoever follows; at an equal level the larger + // ratio is the more entitled; the group's identity is the last canonical discriminator + // before the index. + auto cheaperStep = [&](Candidate const& lhs, Candidate const& rhs) + { + if (lhs.d1 != rhs.d1) + { + return lhs.d1 < rhs.d1; + } + if (lhs.d2 != rhs.d2) + { + return lhs.d2 < rhs.d2; + } + if (ratio[lhs.pgIdx] != ratio[rhs.pgIdx]) + { + return ratio[lhs.pgIdx] > ratio[rhs.pgIdx]; + } + if (identities[lhs.pgIdx] < identities[rhs.pgIdx]) + { + return true; + } + if (identities[rhs.pgIdx] < identities[lhs.pgIdx]) + { + return false; + } + return toSizeT(lhs.pgIdx) < toSizeT(rhs.pgIdx); + }; -// Instance convenience wrapper. -TypedVec CacheLevelStorage::computeSlotCountList( - TypedVec const& ratioList, TypedVec const& minSlots, - std::optional quota) const -{ - size_t q = quota.value_or(totalQuota()); - TLLM_CHECK_DEBUG(ratioList.size() == mPoolGroups.size()); - return ratioToSlotCountList(q, slotSizeLists(), ratioList, poolSizeGranularity(), minSlots); + std::vector cands; + cands.reserve(toSizeT(numPg)); + while (true) + { + size_t placed = 0; + for (PoolGroupIndex pgIdx{0}; pgIdx < numPg; ++pgIdx) + { + placed += allBytes(pgIdx, slotCntList[pgIdx]); + } + + cands.clear(); + for (PoolGroupIndex pgIdx{0}; pgIdx < numPg; ++pgIdx) + { + SlotCount const next = freeRun(slotSizeLists[pgIdx], grains[pgIdx], slotCntList[pgIdx] + 1); + size_t const d1 = quotaGrains(pgIdx, next) - quotaGrains(pgIdx, slotCntList[pgIdx]); + size_t const d2 = secondGrains(pgIdx, next) - secondGrains(pgIdx, slotCntList[pgIdx]); + if (d1 > remainingQuota || (bounded && d2 > remainingSecond)) + { + continue; + } + double const crossing = static_cast(allBytes(pgIdx, slotCntList[pgIdx]) + allBytes(pgIdx, next)) + / (2.0 * ratio[pgIdx]); + + // The level at which the rest of the allocation is exactly this group's share. + // Solving B = r * (B + rest) gives rest / (1 - r), a total that excludes the + // group's own bytes and so does not chase itself upward. A ratio at or above one + // leaves no rest to be consistent with; zero sends every step to the budget + // clause, which is the right rule for a lone group and is what stops the loop. + double const rest = static_cast(placed - allBytes(pgIdx, slotCntList[pgIdx])); + double cap = ratio[pgIdx] < 1.0 ? rest / (1.0 - ratio[pgIdx]) : 0.0; + if (secondaryBinds && onSecond[pgIdx]) + { + cap = std::min(cap, static_cast(level2)); + } + + bool correcting = true; + if (corrections[pgIdx] >= kSelfCorrectionLimit || crossing * (1.0 + kRatioEps) >= cap) + { + // Past its own share, so only the budget that sets this group's level may + // carry it further, and only out of that budget's own residue: a class-A + // group may not creep past its class level on device grains alone. A step + // neither bounded budget pays for moves the group off target for nothing. + if (secondaryBinds && onSecond[pgIdx]) + { + if (d2 < 1) + { + continue; + } + } + else if (d1 < 1) + { + continue; + } + correcting = false; + } + + cands.push_back(Candidate{crossing, pgIdx, next, correcting, d1, d2}); + } + if (cands.empty()) + { + // The two levels have stopped spending, so whatever the second budget still + // holds is leftover. Closing a group out of it can make that group's next step + // cost a quota grain, so the promotions and the closure alternate; two closures + // never run back to back, since the second has nothing left to reach. + if (closePaidGrains()) + { + continue; + } + break; + } + + // Everything within the ratio's own precision of the lowest crossing is tied, and + // the tie is settled by the groups' data, so no answer depends on where a group + // happens to sit in the list. + double lowest = cands.front().crossing; + for (auto const& cand : cands) + { + lowest = std::min(lowest, cand.crossing); + } + Candidate const* chosen = nullptr; + for (auto const& cand : cands) + { + if (cand.crossing > lowest * (1.0 + kRatioEps)) + { + continue; + } + if (chosen == nullptr || cheaperStep(cand, *chosen)) + { + chosen = &cand; + } + } + + if (chosen->correcting) + { + ++corrections[chosen->pgIdx]; + } + remainingQuota + -= quotaGrains(chosen->pgIdx, chosen->next) - quotaGrains(chosen->pgIdx, slotCntList[chosen->pgIdx]); + if (bounded) + { + remainingSecond + -= secondGrains(chosen->pgIdx, chosen->next) - secondGrains(chosen->pgIdx, slotCntList[chosen->pgIdx]); + } + slotCntList[chosen->pgIdx] = chosen->next; + } + + return slotCntList; } } // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/core.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/core.h index 5f4c5c8da375..b2f960d8daef 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/core.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/core.h @@ -29,6 +29,7 @@ #include #include #include +#include #include #include @@ -204,13 +205,23 @@ class SlotAllocator //! so this is just a rounding quantum. inline constexpr size_t kDiskPoolSizeGranularity = size_t{2} << 20; +//! The grain a hot-level pool of `medium` rounds its bytes to, where +//! `deviceGranularity` is the VMM allocation granularity device pools grow in. +//! +//! For the config stage, which sizes pools before any exists to be asked. Once storage is +//! built, ask the group: PoolGroupBase::poolSizeGranularity(CacheTier). The two agree +//! because a host-backed pool commits in hostMemSizeGranularity() whatever level it is on. +inline size_t hotPoolSizeGranularity(CacheTier medium, size_t deviceGranularity) noexcept +{ + return medium == CacheTier::HOST_MEM ? hostMemSizeGranularity() : deviceGranularity; +} + //! Unit a tier of the given size allocates and releases physical memory in, //! and therefore the quantum its pool sizes are rounded to. //! -//! For GPU and host it scales with the tier: 2 MiB for small tiers, rising to -//! 32 MiB, beyond which larger units measure no faster. A resize can give -//! memory back only in whole units, so this trades that precision against the -//! per-allocation cost. +//! For GPU it scales with the tier: 2 MiB for small tiers, rising to 32 MiB, beyond +//! which larger units measure no faster. A resize can give memory back only in whole +//! units, so this trades that precision against the per-allocation cost. //! //! This is the policy. What a given StorageManager actually uses for a tier can //! differ -- see StorageManager::tierCommitUnit -- because every GPU level @@ -250,6 +261,21 @@ class SlotPoolBase return mSlotSize * slotCountToSizeT(numSlots()); } + //! The unit this pool's backing is allocated in, and so the grain its byte total + //! rounds to. A property of the allocator behind the pool, not of the level it + //! sits in: two pools of one group can answer differently. + [[nodiscard]] virtual size_t sizeGranularity() const noexcept = 0; + + //! What this pool's pages are. A group asks its pools rather than holding the answer, + //! so a group of one tier may still hold a pool of another. + [[nodiscard]] virtual CacheTier medium() const noexcept = 0; + + //! Block until the backing for every slot up to `slot` exists. + //! + //! Only a progressively filled pool has anything to wait for; the rest are backed when + //! they are built, so this does nothing for them. + virtual void waitSlotUsable(SlotId) const {} + virtual void destroy() = 0; virtual void resize(SlotCount newNumSlots) = 0; virtual Address slotAddress(SlotId slot) const = 0; @@ -271,8 +297,16 @@ class GpuSlotPool : public SlotPoolBase void resize(SlotCount newNumSlots) override; Address slotAddress(SlotId slot) const override; - // Extend by exactly one physical memory chunk; returns new numSlots. - SlotCount extendByOnePhysMem(); + //! The VMM allocation granularity, which is what this pool's mapping grows in. + [[nodiscard]] size_t sizeGranularity() const noexcept override + { + return mVirtMem.physMemSize(); + } + + [[nodiscard]] CacheTier medium() const noexcept override + { + return CacheTier::GPU_MEM; + } static size_t computeNumPhysMem(size_t slotSize, SlotCount numSlots, size_t physMemSize) noexcept; static SlotCount computeNumSlots(size_t slotSize, size_t numPhysMem, size_t physMemSize) noexcept; @@ -290,6 +324,10 @@ class HostSlotPool : public SlotPoolBase //! vmSize is the address space reserved up front, and so the bound on every //! later resize. Reserving it all at once is what keeps the base address //! fixed, as it does for GpuSlotPool. + //! + //! The backing memory is committed on a worker thread, so a slot may exist + //! before it is backed. Closing that gap is the owning pool group's job: it + //! calls waitSlotUsable() before returning a slot from allocate(). HostSlotPool(size_t slotSize, SlotCount numSlots, size_t vmSize, HostMemBackingOptions const& options); SlotCount numSlots() const noexcept override; @@ -299,10 +337,20 @@ class HostSlotPool : public SlotPoolBase size_t alignedSize(SlotCount numSlots) const noexcept; - //! Blocks until `slot` is backed by physical memory. - //! - //! Rethrows the fill worker's exception if the fill failed. - void waitSlotUsable(SlotId slot) const; + //! The commit unit, which alignedSize() already rounds to. + [[nodiscard]] size_t sizeGranularity() const noexcept override + { + return mCommitUnit; + } + + [[nodiscard]] CacheTier medium() const noexcept override + { + return CacheTier::HOST_MEM; + } + + //! Blocks until `slot` is backed by physical memory, rethrowing the fill worker's + //! exception if the fill failed. + void waitSlotUsable(SlotId slot) const override; private: //! Declared before mHostMem: alignedSize() reads it while building mHostMem. @@ -325,6 +373,16 @@ class DiskSlotPool : public SlotPoolBase void resize(SlotCount newNumSlots) override; Address slotAddress(SlotId slot) const override; + [[nodiscard]] size_t sizeGranularity() const noexcept override + { + return kDiskPoolSizeGranularity; + } + + [[nodiscard]] CacheTier medium() const noexcept override + { + return CacheTier::DISK; + } + int fd() const noexcept { return mFd; @@ -349,6 +407,22 @@ class PoolGroupBase explicit PoolGroupBase(SlotCount numSlots); virtual ~PoolGroupBase(); + //! The tier this pool group belongs to, which is what names it to the levels above: + //! which quota pays for it, which budget it is rebalanced within. + //! + //! A property of the pool group rather than of its level, because groups of one level + //! can differ. This is not necessarily the medium of any given pool's pages -- see + //! poolMedium(). + [[nodiscard]] virtual CacheTier cacheTier() const noexcept = 0; + + //! What one pool's pages actually are. This is cacheTier() for every pool except a + //! host-backed pool of a GPU-tier group, where the group's tier would name device + //! memory for a host address. + [[nodiscard]] CacheTier poolMedium(PoolIndex poolIdx) const + { + return mPools.at(poolIdx)->medium(); + } + PoolIndex numPools() const noexcept { return mPools.size(); @@ -371,8 +445,14 @@ class PoolGroupBase return mSlotAllocator; } - virtual Slot allocate(); - virtual std::vector allocateMultiple(SlotCount numSlots); + //! Allocate, then block until the slots handed back are backed. + //! + //! The wait is the whole of a group's fill synchronization, and this is the single + //! point every later use of a slot passes through. Nothing downstream -- not + //! slotAddress(), not a device mapping read from it -- waits again. A pool with no + //! fill to wait on answers at once. + Slot allocate(); + std::vector allocateMultiple(SlotCount numSlots); void release(Slot slot); void destroy(); void resizePools(std::optional newNumSlots = std::nullopt); @@ -382,28 +462,80 @@ class PoolGroupBase // Sizes per pool. TypedVec slotSize() const; - // Total bytes across all pools for one slot. - size_t numBytes() const noexcept + //! The grain a pool of `medium` rounds its byte total to. Media of one group differ -- + //! a host-backed pool of a GPU group commits in host units, not in the device's VMM + //! granularity -- but pools sharing a medium do not: the grain follows the medium's + //! allocator, not the pool's slot size, which registration bounds absorb instead. + //! + //! Returns 0 if no pool of `medium` is here to ask. + [[nodiscard]] size_t poolSizeGranularity(CacheTier medium) const { - size_t total = 0; - for (auto const& pool : mPools) - total += pool->numBytes(); - return total; + auto ofMedium = std::views::iota(0, mPools.size().value()) + | std::views::filter([this, medium](int pool) { return poolMedium(PoolIndex{pool}) == medium; }); + if (std::ranges::empty(ofMedium)) + { + return 0; + } + return getUniformAttribute( + ofMedium, [this](int pool) { return mPools.at(PoolIndex{pool})->sizeGranularity(); }); + } + + //! The medium of this group's pools that are not of its tier, if it has any. At most + //! one: a group's pools are its tier's own or of a single other medium. + //! + //! This is the group's half of what SecondaryPoolBudget calls the second medium. Ask + //! it rather than comparing against cacheTier() at each site: a group of the host tier + //! has no second medium even though all its pools are host memory, and a test written + //! as "pools that are HOST_MEM" names every one of them. + [[nodiscard]] std::optional secondaryMedium() const + { + auto offTier = std::views::iota(0, mPools.size().value()) + | std::views::filter([this](int pool) { return poolMedium(PoolIndex{pool}) != cacheTier(); }); + if (std::ranges::empty(offTier)) + { + return std::nullopt; + } + return getUniformAttribute(offTier, [this](int pool) { return poolMedium(PoolIndex{pool}); }); } - // Total bytes across all pools, rounding each pool's bytes up to granularity. - // Mirrors Python total_quota: sum(round_up(p.num_bytes, granularity) for p in pg._pools). - size_t roundedNumBytes(size_t granularity) const noexcept + //! Bytes this group would hold at `slots`, each pool rounded to its own allocator's + //! grain. An absent `medium` counts every pool; a named one counts only its own. + //! + //! One slot count for every pool, though pools can report different counts: each + //! rounds its own allocation up to its own grain, so the excess is under one grain + //! and rounding returns the allocation whichever count is used. + [[nodiscard]] size_t roundedBytesWith(SlotCount slots, std::optional medium) const { size_t total = 0; for (auto const& pool : mPools) - total += roundUp(pool->numBytes(), granularity); + { + if (medium.has_value() && pool->medium() != *medium) + { + continue; + } + total += roundUp(slotCountToSizeT(slots) * pool->slotSize(), pool->sizeGranularity()); + } return total; } + //! Bytes this group holds now in pools whose pages are `medium`. Pools of any other + //! medium contribute nothing. + //! + //! Bytes of different media are not comparable -- they are paid for from different + //! budgets and rounded to different grains -- so every caller measuring them names the + //! one it means. roundedNumBytes(cacheTier()) is what this group's own tier pays for; + //! roundedNumBytes(HOST_MEM) on a GPU-tier group is what it holds in host pages besides. + [[nodiscard]] size_t roundedNumBytes(CacheTier medium) const + { + return roundedBytesWith(numSlots(), medium); + } + protected: // Mirrors Python _get_num_slots_from_pools: min(p.num_slots for p in pools). - SlotCount getNumSlotsFromPools() const noexcept; + [[nodiscard]] SlotCount getNumSlotsFromPools() const noexcept; + + //! Block until every pool backs the range ending just past `slot`. + void _waitSlotUsable(SlotId slot) const; SlotAllocator mSlotAllocator; TypedVec> mPools; @@ -416,13 +548,34 @@ class PoolGroupBase class GpuPoolGroup : public PoolGroupBase { public: - GpuPoolGroup( - SlotCount numSlots, TypedVec const& slotSizeList, PooledPhysMemAllocator& physMemAllocator); + //! `sparsity`, when given, marks the pools to allocate from host memory instead of + //! device memory. They hold a sparse buffer's pages: ordinary hot pages that are read + //! through the sparse GPU cache rather than directly, so they are device-mapped but do + //! not consume the tier's device budget. + //! + //! The group still reports GPU_MEM. Migration never reads the hot side's tier -- the + //! codec path keeps only the cold side's -- and host memory is device-accessible, so a + //! mixed group needs no second tier. + //! + //! `hostOptions` backs the sparse pools. One object for the whole level, so every + //! host-backed pool commits in the same unit. + //! + //! `sparseReservation`, when given, caps the host address space each sparse pool + //! reserves; a pool never reserves less than its initial slots. Unset reserves for the + //! most slots device memory could buy, which exceeds the address space once a sparse + //! slot is several hundred times the group's device bytes per slot. + GpuPoolGroup(SlotCount numSlots, TypedVec const& slotSizeList, + PooledPhysMemAllocator& physMemAllocator, TypedVec sparsity = {}, + HostMemBackingOptions const& hostOptions = {}, std::optional sparseReservation = std::nullopt); + + [[nodiscard]] CacheTier cacheTier() const noexcept override + { + return CacheTier::GPU_MEM; + } }; //! Host pools are committed progressively on a background thread, so a slot may -//! be handed out before its memory is backed. Both allocation entry points block -//! until the slots they return are usable. +//! be handed out before its memory is backed. PoolGroupBase::allocate waits for it. //! //! The wait is cheap in the steady state: fresh slot ids ascend from zero and //! recycled slots are reused first, so the Nth fresh allocation waits only for @@ -439,18 +592,50 @@ class HostPoolGroup : public PoolGroupBase HostPoolGroup(SlotCount numSlots, TypedVec const& slotSizeList, size_t groupVmSize, HostMemBackingOptions const& options = {}); - Slot allocate() override; - std::vector allocateMultiple(SlotCount numSlots) override; - -private: - //! Blocks until every pool backs the byte range ending just past `slot`. - void waitSlotUsable(SlotId slot) const; + [[nodiscard]] CacheTier cacheTier() const noexcept override + { + return CacheTier::HOST_MEM; + } }; class DiskPoolGroup : public PoolGroupBase { public: DiskPoolGroup(SlotCount numSlots, TypedVec const& slotSizeList, std::string const& directory); + + [[nodiscard]] CacheTier cacheTier() const noexcept override + { + return CacheTier::DISK; + } +}; + +//! A second medium's budget over a disjoint subset of each pool group's pools. +//! +//! A slot spans every pool of its group, so one slot count drives the bytes taken in every +//! pool -- but the pools of the GPU tier are of two media, each with its own total and its +//! own allocation grain. This names the second one. `pools[pg]` lists exactly the pools +//! `quota` pays for; the rest are what the primary quota pays for, derived by +//! poolsPaidByQuota(). The two are complements: no pool is named by both, and together +//! they are all of the group's pools. +//! +//! `pools[pg]` may be empty, which is how "this group costs the second medium nothing" is +//! said. Such a group takes no share of `quota` and can never be bound by it. +struct SecondaryPoolBudget +{ + //! Bytes of the second medium, for the whole level, or unset to leave it unbounded. + //! A budget, not a per-group maximum: how it divides between groups is what the + //! allocator decides. + //! + //! `pools` is a property of the pools rather than of any budget, so a level with a + //! second medium supplies this struct even when nothing bounds it; only `quota` says + //! whether anything does. + std::optional quota; + //! The second medium's allocation grain, which is a property of its allocator rather + //! than of the level. Host-backed pools round to HostMem::sizeGranularity(), not to the device + //! granularity, so charging them in device grains would over-count every one of them. + size_t granularity = 0; + //! The pools `quota` pays for, per pool group. + TypedVec> pools; }; // --------------------------------------------------------------------------- @@ -511,15 +696,55 @@ class CacheLevelStorage return std::get(mPoolGroups.at(pgIdx)->slotAddress(slotId).at(poolIdx)); } - size_t totalQuota() const noexcept + //! Bytes this level holds in pools whose pages are `medium`, each rounded to that + //! medium's grain. Every group of a level is of the level's tier, so the media + //! reachable here are this tier's own and, on a device level, the host pages its + //! host-backed pools hold. + //! + //! At cacheTier() this is what the level's quota pays for. Host-backed pools of a + //! device level fall outside it, which is the point: a rebalance given an inflated + //! figure would hand the level more slots than the device can hold. + [[nodiscard]] size_t roundedNumBytes(CacheTier medium) const { - size_t granularity = poolSizeGranularity(); size_t total = 0; for (auto const& pg : mPoolGroups) - total += pg->roundedNumBytes(granularity); + total += pg->roundedNumBytes(medium); return total; } + //! The medium of this level's pools that are not of its own tier, if any group has + //! one. At most one across the level: every group here is of the level's tier, and a + //! group has at most one second medium. + [[nodiscard]] std::optional secondaryMedium() const + { + for (auto const& pg : mPoolGroups) + { + if (auto const medium = pg->secondaryMedium()) + { + return medium; + } + } + return std::nullopt; + } + + //! The grain pools of `medium` round to at this level, or 0 when it has none. Groups + //! holding that medium are held to agree, as the pools within one are. + //! + //! Read off a pool, for every medium alike. A host pool commits in the unit its level + //! was configured with and a device pool in its allocator's, so asking the pool gives + //! the quantum the quota is expressed in as well as the second medium's grain -- + //! poolSizeGranularity(cacheTier()) is the former. + [[nodiscard]] size_t poolSizeGranularity(CacheTier medium) const + { + auto grainOf = [medium](auto const& pg) { return pg->poolSizeGranularity(medium); }; + auto withMedium = mPoolGroups | std::views::filter([&](auto const& pg) { return grainOf(pg) != 0; }); + if (std::ranges::empty(withMedium)) + { + return 0; + } + return getUniformAttribute(withMedium, grainOf); + } + // Returns numSlots() per pool group. TypedVec slotCountList() const { @@ -530,6 +755,33 @@ class CacheLevelStorage return ret; } + //! The pools whose pages are `medium`, as indices into slotSizeLists()[pg]. + //! + //! For "the pools this level's quota does not pay for", pass secondaryMedium() rather + //! than a literal: on a host level every pool is HOST_MEM, yet none of them is a + //! second medium. + //! + //! The complement is what the quota pays for, which the allocator derives rather than + //! taking: that way no pool can be charged twice or missed. + [[nodiscard]] TypedVec> poolsOfMedium(CacheTier medium) const + { + TypedVec> ret; + ret.reserve(mPoolGroups.size()); + for (auto const& pg : mPoolGroups) + { + std::vector matching; + for (PoolIndex pool{0}; pool < pg->slotSize().size(); ++pool) + { + if (pg->poolMedium(pool) == medium) + { + matching.push_back(pool); + } + } + ret.push_back(std::move(matching)); + } + return ret; + } + TypedVec> slotSizeLists() const { TypedVec> ret; @@ -546,14 +798,43 @@ class CacheLevelStorage return ret; } - // Current ratio list: proportion of bytes per pool group. - TypedVec ratioList() const + //! Bytes each pool group holds at `slotCounts`, every pool in its own allocator's grain. + //! + //! This is the measure the pool ratio speaks in, and the one the allocator divides + //! against. Measuring a host-backed pool in device grains instead -- they differ by up + //! to 8192x -- would put the two sides of a rebalance's comparison in different units. + [[nodiscard]] TypedVec allMediaBytes( + TypedVec const& slotCounts) const { + TypedVec ret(mPoolGroups.size(), 0); + for (PoolGroupIndex pgIdx{0}; pgIdx < mPoolGroups.size(); ++pgIdx) + { + ret[pgIdx] = mPoolGroups[pgIdx]->roundedBytesWith(slotCounts[pgIdx], std::nullopt); + } + return ret; + } + + //! Current ratio list: proportion of bytes per pool group. + //! + //! Counts every pool, host-backed ones included -- unlike the quota, which pays for + //! the dense ones alone. The asymmetry is deliberate: this ratio is projected onto cold + //! levels, where one pool holds dense and sparse buffers together, and it is what a + //! rebalance compares its target against. Both sides of that comparison have to measure + //! the same bytes or it never settles; see "The quota excludes sparse pools; the ratio + //! counts them" in SPARSE_ATTENTION_DESIGN.md. + [[nodiscard]] TypedVec ratioList() const + { + TypedVec counts(mPoolGroups.size(), 0); + for (PoolGroupIndex pgIdx{0}; pgIdx < mPoolGroups.size(); ++pgIdx) + { + counts[pgIdx] = mPoolGroups[pgIdx]->numSlots(); + } + auto const bytes = allMediaBytes(counts); TypedVec ret(mPoolGroups.size(), 0.f); float total = 0.f; for (PoolGroupIndex pgIdx{0}; pgIdx < mPoolGroups.size(); ++pgIdx) { - auto sz = static_cast(mPoolGroups[pgIdx]->numBytes()); + auto sz = static_cast(bytes[pgIdx]); ret[pgIdx] = sz; total += sz; } @@ -563,11 +844,6 @@ class CacheLevelStorage return ret; } - virtual size_t poolSizeGranularity() const noexcept - { - return kDiskPoolSizeGranularity; - } - //! Largest quota this level may be resized to, when it has one. //! //! Unset means unbounded here: GPU and disk levels are limited by the @@ -577,26 +853,49 @@ class CacheLevelStorage return std::nullopt; } - // Compute slot counts per pool group for a given ratio, min_slots, and optional quota. - // Instance convenience method — delegates to the static version below. - TypedVec computeSlotCountList(TypedVec const& ratioList, - TypedVec const& minSlots, std::optional quota = std::nullopt) const; - - // Static version: compute slot counts from quota, slot size lists, ratio, granularity, and min_slots. - // Mirrors Python CacheLevelStorage.ratio_to_slot_count_list. + //! Divide `quota` between pool groups by `ratioList`, honouring `minSlots` as a floor, + //! and `secondBudget` as a second medium's total over a disjoint set of the same pools. + //! + //! `slotSizeLists` holds every pool's size; the pools `quota` pays for are those + //! `secondBudget->pools[pg]` does not name. A slot spans all of them, so one count per + //! group settles both media at once. + //! + //! Both budgets and the floors are kept exactly; the ratio is the only soft bound. The + //! floors come off the top rather than out of anyone's share, so no division of the rest + //! can break one, and a group bound by either budget releases what it does not use of + //! the other. A group named by neither budget for a medium costs that medium nothing. + //! + //! The two budgets can only conflict at the floors, and that condition names neither the + //! ratio nor `quota`. It is therefore a property of the configuration: it is rejected the + //! first time this runs, and a rebalance -- which changes both of those and none of the + //! terms in the condition -- cannot begin to fail. + //! + //! With `secondBudget` unset, `quota` is the only bound and every pool is charged to it. + //! + //! @throws std::runtime_error when the floors alone exceed either budget. static TypedVec ratioToSlotCountList(size_t quota, TypedVec> const& slotSizeLists, TypedVec const& ratioList, size_t granularity, - TypedVec const& minSlots); + TypedVec const& minSlots, + std::optional const& secondBudget = std::nullopt); + + //! The pools `quota` pays for: every pool the second budget does not name. + //! + //! Derived rather than given, so the two budgets partition the group's pools by + //! construction -- a pool cannot be charged to both, nor left to neither. + static TypedVec> poolsPaidByQuota( + TypedVec> const& slotSizeLists, + std::optional const& secondBudget); - // Distribute grains among pools within a pool group. - // Returns {num_slots, grains_consumed}. - static std::pair grainsToSlots( - size_t pgGrains, TypedVec const& slotSizeList, size_t granularity); + // Distribute grains among the named pools of one pool group. + // Returns {num_slots, grains_consumed}. `which` must not be empty. + static std::pair grainsToSlots(size_t pgGrains, TypedVec const& slotSizeList, + std::vector const& which, size_t granularity); - // Compute minimum grains needed for numSlots in a pool group. - static size_t grainsForSlots( - SlotCount numSlots, TypedVec const& slotSizeList, size_t granularity); + // Minimum grains the named pools of one pool group need for numSlots. Zero for an empty + // `which`, which is what lets a group with no pool of a medium cost that medium nothing. + static size_t grainsForSlots(SlotCount numSlots, TypedVec const& slotSizeList, + std::vector const& which, size_t granularity); protected: TypedVec> mPoolGroups; @@ -605,19 +904,17 @@ class CacheLevelStorage class GpuCacheLevelStorage : public CacheLevelStorage { public: + //! Each SlotDesc declares per pool which pools are sparse, and so host-allocated. + //! `sparseReservation` is passed to every GpuPoolGroup. GpuCacheLevelStorage(TypedVec const& slotDescList, - TypedVec const& slotCountList, PooledPhysMemAllocator& physMemAllocator); + TypedVec const& slotCountList, PooledPhysMemAllocator& physMemAllocator, + HostMemBackingOptions const& hostOptions = {}, std::optional sparseReservation = std::nullopt); CacheTier cacheTier() const noexcept override { return CacheTier::GPU_MEM; } - size_t poolSizeGranularity() const noexcept override - { - return mPhysMemAllocator.physMemSize(); - } - private: PooledPhysMemAllocator& mPhysMemAllocator; }; @@ -625,8 +922,6 @@ class GpuCacheLevelStorage : public CacheLevelStorage class HostCacheLevelStorage : public CacheLevelStorage { public: - //! `commitUnit` is the granularity memory is allocated and released in; it - //! is also what pool sizes are quantized to. HostCacheLevelStorage(TypedVec const& slotDescList, TypedVec const& slotCountList, std::optional maxQuota = std::nullopt, HostMemBackingOptions const& options = {}); @@ -636,14 +931,9 @@ class HostCacheLevelStorage : public CacheLevelStorage return CacheTier::HOST_MEM; } - size_t poolSizeGranularity() const noexcept override - { - return mOptions.commitUnit; - } - //! Always set: an unset configured maxQuota resolves to the host memory the //! OS reports, and the pools reserve against that resolved value. Rounded - //! up to poolSizeGranularity(), which is the quantum quotas are expressed + //! up to poolSizeGranularity(cacheTier()), which is the quantum quotas are expressed //! in, so the configured bound is reachable. [[nodiscard]] std::optional maxQuota() const noexcept override { @@ -676,10 +966,10 @@ class DiskCacheLevelStorage : public CacheLevelStorage }; // Factory: create appropriate CacheLevelStorage for a given tier config. -//! `commitUnit` is the tier's allocation granularity, decided once from its -//! quota and fixed for the tier's lifetime. +// `sparseReservation` applies to a GPU tier only; see GpuPoolGroup. std::unique_ptr createCacheLevelStorage(CacheTierConfig const& tierCfg, TypedVec const& slotDescList, TypedVec const& slotCountList, - PooledPhysMemAllocator* gpuPhysMemAllocator = nullptr, HostMemBackingOptions const& options = {}); + PooledPhysMemAllocator* gpuPhysMemAllocator = nullptr, HostMemBackingOptions const& options = {}, + std::optional sparseReservation = std::nullopt); } // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.cpp index ccfad1c559e2..0735340b929f 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.cpp @@ -21,9 +21,9 @@ #include "kv_cache_manager_v2/common.h" #include "kv_cache_manager_v2/copyEngine.h" #include "kv_cache_manager_v2/exceptions.h" -#include "kv_cache_manager_v2/kvCache.h" #include "kv_cache_manager_v2/page.h" #include "kv_cache_manager_v2/stagingBuffer.h" +#include "kv_cache_manager_v2/utils/funcGuard.h" #include "kv_cache_manager_v2/utils/hostMem.h" #include "kv_cache_manager_v2/utils/math.h" #include "tensorrt_llm/common/logger.h" @@ -105,13 +105,14 @@ PoolGroupIndex LifeCyclePoolGroupMapping::numPoolGroups() const noexcept CacheLevelManager::CacheLevelManager(TypedVec const& lifeCycleGrouping, CacheLevel cl, CacheTierConfig const& tierConfig, TypedVec const& slotDescList, TypedVec const& slotCountList, PooledPhysMemAllocator* gpuPhysMemAllocator, - HostMemBackingOptions const& options) + HostMemBackingOptions const& options, std::optional sparseReservation) : cacheLevel(cl) , cacheTier(CacheTier(tierConfig.index())) , controller(lifeCycleGrouping, cl) { TLLM_CHECK((cacheTier == CacheTier::GPU_MEM) == (gpuPhysMemAllocator != nullptr)); - storage = createCacheLevelStorage(tierConfig, slotDescList, slotCountList, gpuPhysMemAllocator, options); + storage = createCacheLevelStorage( + tierConfig, slotDescList, slotCountList, gpuPhysMemAllocator, options, sparseReservation); } // --------------------------------------------------------------------------- @@ -126,6 +127,23 @@ bool isGpuAccessibleMemory(CacheTier tier) noexcept return tier == CacheTier::GPU_MEM || tier == CacheTier::HOST_MEM; } +// The host sparse pool quota an unset configuration gets: 80% of the host memory the host +// tiers' quotas leave. Host tiers configured past physical memory leave none. +size_t defaultHostSparsePoolQuota(TypedVec const& cacheTiers) +{ + size_t hostTierQuota = 0; + for (auto const& tier : cacheTiers) + { + if (cacheTierOf(tier) == CacheTier::HOST_MEM) + { + hostTierQuota += cacheTierQuota(tier); + } + } + size_t const total = hostTotalMemory(); + size_t const remaining = total > hostTierQuota ? total - hostTierQuota : 0; + return remaining / 5 * 4; +} + // Compute the slot-to-page-indices scale factors. // For each (lcId, poolIdx), scale = numBuffersInCoalescedSlot. // Python: _slot_to_page_indices[lc_id][pool_idx] = numBuffers @@ -194,8 +212,7 @@ bool StorageManager::submitColdPageCodec( TLLM_CHECK_DEBUG(chunkPages > 0); // The index vector is ephemeral. The helper captures it before returning while the device update remains // asynchronous. - detail::copyPageIndicesToDevice( - static_cast(device.address()), pageIndices + offset, chunkPages, stream); + copyToDeviceSyncSource(static_cast(device.address()), pageIndices + offset, chunkPages, stream); if (!submit(reinterpret_cast(device.address()), chunkPages, stream)) { @@ -258,12 +275,35 @@ StorageManager::StorageManager(LifeCycleRegistry const& lifeCycles, StorageConfi TLLM_CHECK_DEBUG_WITH_INFO( std::holds_alternative(config.cacheTiers[kHotLevel]), "First cache tier must be GPU"); - // Compute slot size lists for all pool groups. + // Every pool's size, and which of them are sparse. One array of sizes, so the two budgets + // cannot disagree about what a pool costs; the selector only says who pays. TypedVec> slotSizeLists; + TypedVec> sparsePools; slotSizeLists.reserve(hotSlotDescList.size()); + sparsePools.reserve(hotSlotDescList.size()); for (auto const& sd : hotSlotDescList) { - slotSizeLists.push_back(sd.slotSizeList()); + auto const all = sd.slotSizeList(); + std::vector sparse; + size_t dense = 0; + for (PoolIndex p{0}; p < all.size(); ++p) + { + if (sd.isSparse(p)) + { + sparse.push_back(p); + } + else + { + ++dense; + } + } + // Only the GPU tier needs this: sparsity exists at the hot level alone, so a cold + // level's groups size entirely from their own quota. + TLLM_CHECK_WITH_INFO(dense > 0, + "A pool group must hold at least one non-sparse pool; a wholly sparse group has no device " + "bytes from which to size its slots"); + slotSizeLists.push_back(all); + sparsePools.push_back(std::move(sparse)); } size_t const gpuQuota = cacheTierQuota(config.cacheTiers[kHotLevel]); @@ -274,7 +314,7 @@ StorageManager::StorageManager(LifeCycleRegistry const& lifeCycles, StorageConfi // Constraints are hot-level feasibility floors. They are scaled by // 1/maxUtilForResume because KvCache::resume rejects a pool group above that // utilization. Other levels need only the structural one-slot floor. - mMinSlots + mHotMinSlots = computePoolGroupMinSlotsFromConstraints(constraints, tokensPerBlock, mSwaScratchReuse, maxUtilForResume); // Derive hot-tier lifecycle byte weights. Cold initialization preserves the slot-count proportions implied by @@ -302,29 +342,38 @@ StorageManager::StorageManager(LifeCycleRegistry const& lifeCycles, StorageConfi } else if (typicalBatch.has_value()) { - lifeCycleRatio = ratioFromBatch(*typicalBatch, tokensPerBlock, mSwaScratchReuse, gpuGranularity); + lifeCycleRatio = ratioFromBatch(*typicalBatch, tokensPerBlock, mSwaScratchReuse, gpuGranularity, sparsePools); } else if (!constraints.empty()) { auto lifeCycleSlots = computeSlotsFromConstraints(constraints, tokensPerBlock, mSwaScratchReuse, maxUtilForResume); - lifeCycleRatio = normalizeToRatio(slotsToBytes(lifeCycleSlots, gpuGranularity)); + lifeCycleRatio = normalizeToRatio(slotsToBytes(lifeCycleSlots, gpuGranularity, sparsePools)); } else { // Fallback: average history length 2048. BatchDesc fallback; fallback.kvCaches.push_back(KVCacheDesc{2049, 2048}); - lifeCycleRatio = ratioFromBatch(fallback, tokensPerBlock, mSwaScratchReuse, gpuGranularity); + lifeCycleRatio = ratioFromBatch(fallback, tokensPerBlock, mSwaScratchReuse, gpuGranularity, sparsePools); } auto const hotRatio = toPoolGroupRatio(kHotLevel, lifeCycleRatio); mLevels.reserve(config.cacheTiers.size()); - auto gpuSlotCounts = computeSlotCountForLevel(config.cacheTiers[kHotLevel], slotSizeLists, hotRatio, mMinSlots); + // A sparse pool's hot pages are host memory. Its bytes are excluded from the GPU quota + // rather than sized from it: a slot spans every pool of its group, so a sparse pool's + // row count follows from the group's device-backed pools, which the quota does size. + // Those pages still need host memory, so an unset quota is filled in here and kept: it + // bounds the slot counts and each sparse pool's address-space reservation alike. + mHostSparsePoolQuota = config.hostSparsePoolQuota.has_value() ? *config.hostSparsePoolQuota + : defaultHostSparsePoolQuota(config.cacheTiers); + // From the descriptors: the level's storage does not exist yet. + auto gpuSlotCounts = computeSlotCountForLevel(config.cacheTiers[kHotLevel], slotSizeLists, hotRatio, mHotMinSlots, + gpuGranularity, _secondaryBudget(sparsePools, hostMemSizeGranularity(), mHostSparsePoolQuota)); mLevels.emplace_back(lifeCycleGrouping(kHotLevel), kHotLevel, config.cacheTiers[kHotLevel], slotDescList(kHotLevel), - gpuSlotCounts, mGpuPhysMemAllocator.get(), tierBackingOptions(config.cacheTiers[kHotLevel])); + gpuSlotCounts, mGpuPhysMemAllocator.get(), HostMemBackingOptions{}, mHostSparsePoolQuota); auto const gpuDescs = poolGroupDescs(); TLLM_CHECK_WITH_INFO( @@ -430,7 +479,7 @@ StorageManager::StorageManager(LifeCycleRegistry const& lifeCycles, StorageConfi { continue; } - size_t const granularity = tierCommitUnit(tierConfig); + size_t const granularity = tierCommitUnit(tierConfig, gpuGranularity); if (size_t const minQuota = minQuotaForLevel(coldSlotSizeLists, granularity, coldMinSlots); *maxQuota < minQuota) { @@ -444,12 +493,14 @@ StorageManager::StorageManager(LifeCycleRegistry const& lifeCycles, StorageConfi { TLLM_CHECK(appendLevelSlotDescList(coldSlotDescList) == level); auto const coldRatio = projectPoolGroupRatio(kHotLevel, level, lifeCycleRatio); - auto slotCounts - = computeSlotCountForLevel(config.cacheTiers[level], coldSlotSizeLists, coldRatio, coldMinSlots); + // No second budget: a cold group is a single pool holding whatever the codec packs + // into it, dense and sparse buffers alike, so its quota pays for all of it. + auto slotCounts = computeSlotCountForLevel( + config.cacheTiers[level], coldSlotSizeLists, coldRatio, coldMinSlots, gpuGranularity); auto* gpuPhysMemAllocator = cacheTierOf(config.cacheTiers[level]) == CacheTier::GPU_MEM ? mGpuPhysMemAllocator.get() : nullptr; mLevels.emplace_back(lifeCycleGrouping(level), level, config.cacheTiers[level], slotDescList(level), slotCounts, - gpuPhysMemAllocator, tierBackingOptions(config.cacheTiers[level])); + gpuPhysMemAllocator, tierBackingOptions(config.cacheTiers[level], gpuGranularity)); } for (CacheLevel level{0}; level < mLevels.size(); ++level) @@ -593,8 +644,6 @@ void StorageManager::submitMigrationBatch(CacheLevel dstLevel, CacheLevel srcLev TLLM_CHECK_DEBUG(pageIndices != nullptr && numPages > 0); PoolGroupIndex const srcPgIdx = getPoolGroupIndex(srcLevel, batchingLayerGroupId); PoolGroupIndex const dstPgIdx = getPoolGroupIndex(dstLevel, batchingLayerGroupId); - CacheTier const srcTier = cacheTier(srcLevel); - CacheTier const dstTier = cacheTier(dstLevel); bool const srcIsHot = srcLevel == kHotLevel; bool const dstIsHot = dstLevel == kHotLevel; auto& srcPoolGroup = poolGroup(srcLevel, srcPgIdx); @@ -619,14 +668,22 @@ void StorageManager::submitMigrationBatch(CacheLevel dstLevel, CacheLevel srcLev tasks.push_back({dstPoolGroup.slotAddress(toSlotId(pageIndices[index].dst)).at(poolIdx), srcPoolGroup.slotAddress(toSlotId(pageIndices[index].src)).at(poolIdx)}); } - mCopyEngine->transfer(dstTier, srcTier, srcSizes.at(poolIdx), tasks, stream); + // Per pool, not per group: a GPU-tier group holds device-backed pools beside + // host-backed sparse ones, and the copy engine picks its implementation from + // the medium. Taking the group's tier would run a device-to-device kernel over + // host pointers -- legal, since the host memory is registered, but tuned for a + // bandwidth it does not have. + mCopyEngine->transfer(poolMedium(dstLevel, dstPgIdx, poolIdx), poolMedium(srcLevel, srcPgIdx, poolIdx), + srcSizes.at(poolIdx), tasks, stream); } return; } CacheLevel const coldLevel = srcIsHot ? dstLevel : srcLevel; PoolGroupIndex const coldPgIdx = srcIsHot ? dstPgIdx : srcPgIdx; - CacheTier const coldTier = srcIsHot ? dstTier : srcTier; + // Only a GPU-tier pool is ever host-backed, so for a cold level cacheTier() and + // poolMedium() agree. + CacheTier const coldTier = cacheTier(coldLevel); size_t const coldPageBytes = slotSize(coldLevel, coldPgIdx).at(PoolIndex{0}); PageIndexLocation const pageIndexLocation = mPageIndexLocations.at(batchingLayerGroupId); auto encodeBatch = [this, batchingLayerGroupId, pageIndexLocation]( @@ -740,6 +797,21 @@ CacheTier StorageManager::cacheTier(CacheLevel level) const return mLevels.at(level).cacheTier; } +CacheTier StorageManager::poolMedium(CacheLevel level, PoolGroupIndex pgIdx, PoolIndex poolIdx) const +{ + return mLevels.at(level).storage->poolGroup(pgIdx).poolMedium(poolIdx); +} + +size_t StorageManager::getNumBytes(CacheLevel level, CacheTier medium) const +{ + return mLevels.at(level).storage->roundedNumBytes(medium); +} + +std::optional StorageManager::sparseBackingMedium() const +{ + return mLevels.at(kHotLevel).storage->secondaryMedium(); +} + void StorageManager::releaseSlot(LifeCycleId lc, CacheLevel level, Slot slot) { PoolGroupIndex pg = getPoolGroupIndex(level, lc); @@ -1395,196 +1467,6 @@ void StorageManager::batchedMigrate( } } -void StorageManager::offloadSparsePages(KvCache& requestingCache, std::vector> const& pages, - MigrationRecorder const& migrationRecorder, DropRecorder const& dropRecorder) -{ - if (requestingCache.storageManager() != this) - { - throw LogicError("Offload pages and requesting cache must belong to the same manager"); - } - _migrateLockedSparsePages( - requestingCache.cudaStream(), kSparseHistoryLevel, pages, &requestingCache, migrationRecorder, dropRecorder); -} - -void StorageManager::promoteSparsePages(CUstream stream, std::vector> const& pages, - MigrationRecorder const& migrationRecorder, DropRecorder const& dropRecorder) -{ - _migrateLockedSparsePages(stream, kHotLevel, pages, nullptr, migrationRecorder, dropRecorder); -} - -void StorageManager::_migrateLockedSparsePages(CUstream stream, CacheLevel dstLevel, - std::vector> const& pages, KvCache const* offloadingCache, - MigrationRecorder const& migrationRecorder, DropRecorder const& dropRecorder) -{ - CacheLevel const srcLevel = dstLevel == kHotLevel ? kSparseHistoryLevel : kHotLevel; - - struct MigrationBatch - { - std::vector> srcPages; - std::vector> srcPageLocks; - std::vector dstSlots; - std::vector srcDstPageIndices; - }; - - std::map batches; - std::set seen; - std::set ownerStreams; - for (auto const& page : pages) - { - if (!page || page->manager != this) - { - throw LogicError("Sparse migration pages must belong to the same manager"); - } - if (!seen.insert(page.get()).second) - { - continue; - } - auto holder = page->holder.lock(); - auto lock = holder ? holder->uniqLock.lock() : nullptr; - if (!lock) - { - throw LogicError("Sparse migration requires a locked page"); - } - if (offloadingCache) - { - lock->prepareSparseOffload(*offloadingCache); - } - else - { - lock->prepareSparsePromotion(); - } - if (page->cacheLevel == dstLevel) - { - continue; - } - auto& batch = batches[getMigrationBatchingLayerGroupId(dstLevel, srcLevel, page->lifeCycle)]; - batch.srcPages.push_back(page); - for (auto const& owner : lock->owners()) - { - ownerStreams.insert(owner.kvCache->cudaStream()); - } - batch.srcPageLocks.push_back(std::move(lock)); - } - if (batches.empty()) - { - return; - } - - TypedVec requirements(numPoolGroups(dstLevel), 0); - for (auto const& [layerGroup, batch] : batches) - { - requirements[getPoolGroupIndex(dstLevel, layerGroup)] += slotCountValueFromSize(batch.srcPages.size()); - } - prepareFreeSlots(dstLevel, requirements, migrationRecorder, dropRecorder); - auto releaseDestinations = FuncGuard( - [&]() - { - for (auto& [layerGroup, batch] : batches) - { - for (auto& slot : batch.dstSlots) - { - if (slot.hasValidSlot()) - { - releaseSlot(layerGroup, dstLevel, std::move(slot)); - } - } - } - }); - for (auto& [layerGroup, batch] : batches) - { - auto& pool = poolGroup(dstLevel, getPoolGroupIndex(dstLevel, layerGroup)); - batch.dstSlots = pool.allocateMultiple(slotCountValueFromSize(batch.srcPages.size())); - batch.srcDstPageIndices.reserve(batch.srcPages.size()); - for (size_t i = 0; i < batch.srcPages.size(); ++i) - { - batch.srcDstPageIndices.push_back({.dst = slotIdToPageIndexValue(batch.dstSlots[i].slotId()), - .src = slotIdToPageIndexValue(batch.srcPages[i]->slotId())}); - } - } - - auto const cudaStream = reinterpret_cast(stream); - std::vector ownerEvents; - ownerEvents.reserve(ownerStreams.size()); - for (auto const ownerStream : ownerStreams) - { - ownerEvents.emplace_back(reinterpret_cast(ownerStream)); - ownerEvents.back().waitInStream(cudaStream); - } - for (auto const& [layerGroup, batch] : batches) - { - for (size_t i = 0; i < batch.srcPages.size(); ++i) - { - batch.srcPages[i]->readyEvent.waitInStream(cudaStream); - batch.dstSlots[i].readyEvent.waitInStream(cudaStream); - for (auto const& event : batch.srcPageLocks[i]->finishEvents) - { - event.waitInStream(cudaStream); - } - } - } - - // Install the fence even when a codec rejects after enqueueing only part of a batch. - CachedCudaEvent completion = CachedCudaEvent::makeNull(); - auto fenceCopies = FuncGuard( - [&]() - { - completion = CachedCudaEvent(cudaStream); - for (auto& [layerGroup, batch] : batches) - { - for (size_t i = 0; i < batch.srcPages.size(); ++i) - { - batch.dstSlots[i].readyEvent = completion; - batch.srcPageLocks[i]->recordMigrationEvent(completion); - } - } - }); - for (auto const& [layerGroup, batch] : batches) - { - submitMigrationBatch( - dstLevel, srcLevel, layerGroup, batch.srcDstPageIndices.data(), batch.srcDstPageIndices.size(), stream); - } - fenceCopies.run(); - - // Subsequent readers on every owner's stream must observe the completed copy. - for (auto const ownerStream : ownerStreams) - { - completion.waitInStream(reinterpret_cast(ownerStream)); - } - for (auto const& [layerGroup, batch] : batches) - { - if (migrationRecorder) - { - migrationRecorder(batch.srcPages, batch.dstSlots, srcLevel, dstLevel); - } - } - for (auto& [layerGroup, batch] : batches) - { - for (size_t i = 0; i < batch.srcPages.size(); ++i) - { - Slot source = batch.srcPageLocks[i]->moveToCacheLevel(dstLevel, std::move(batch.dstSlots[i])); - releaseSlot(batch.srcPages[i]->lifeCycle, srcLevel, std::move(source)); - } - } - if (mEventSink) - { - for (auto const& [layerGroup, batch] : batches) - { - for (auto const& page : batch.srcPages) - { - if (page->isCommitted()) - { - auto const& committed = static_cast(*page); - auto const* block = committed.block; - if (block && !block->isOrphan() && block->holdsPage(committed)) - { - mEventSink->addCacheLevelUpdated(block->key, srcLevel, dstLevel, page->lifeCycle); - } - } - } - } - } -} - int64_t StorageManager::prefetch( CacheLevel dstLevel, TypedVec>>> const& pages) { @@ -1650,24 +1532,6 @@ int64_t StorageManager::prefetch( for (auto& [migrationPath, migrationPages] : migrationGroups) { CacheLevel const srcLevel = migrationPath.first; - if (dstLevel == kHotLevel && srcLevel == kSparseHistoryLevel) - { - std::vector> lockedPages; - for (auto const& page : migrationPages) - { - if (page->status() == PageStatus::LOCKED) - { - lockedPages.push_back(page); - } - } - if (!lockedPages.empty()) - { - TemporaryCudaStream stream({}); - auto scope = stream.enter(); - promoteSparsePages(stream.get(), lockedPages); - std::erase_if(migrationPages, [](auto const& page) { return page->cacheLevel == kHotLevel; }); - } - } _batchedMigrate(dstLevel, srcLevel, migrationPages, /*updateSrc=*/true); // Per batch, after it landed: the pages are already grouped by source level, so this costs // nothing per page and never credits a batch that did not run. @@ -1758,7 +1622,8 @@ TypedVec StorageManager::poolGroupDescs() const pools.reserve(slotSizeList.size()); for (PoolIndex poolIdx{0}; poolIdx < slotSizeList.size(); ++poolIdx) { - pools.push_back(PoolDesc{poolIdx, getMemPoolBaseAddress(pgIdx, poolIdx), slotSizeList.at(poolIdx)}); + pools.push_back(PoolDesc{poolIdx, getMemPoolBaseAddress(pgIdx, poolIdx), slotSizeList.at(poolIdx), + poolMedium(kHotLevel, pgIdx, poolIdx)}); } result.push_back(PoolGroupDesc{pgIdx, numSlots(pgIdx, kHotLevel), descList.at(pgIdx), std::move(pools)}); } @@ -1954,11 +1819,11 @@ void StorageManager::adjustCacheLevel(CacheLevel level, std::optional ne auto& lvlStorage = *mLevels.at(level).storage; auto oldNumSlots = lvlStorage.slotCountList(); size_t quota = newQuota.has_value() - ? roundUp(newQuota.value(), static_cast(lvlStorage.poolSizeGranularity())) - : lvlStorage.totalQuota(); - auto const minSlots - = level == kHotLevel ? mMinSlots : TypedVec(numPoolGroups(level), SlotCount{1}); - size_t minQuota = minQuotaForLevel(lvlStorage.slotSizeLists(), lvlStorage.poolSizeGranularity(), minSlots); + ? roundUp(newQuota.value(), lvlStorage.poolSizeGranularity(lvlStorage.cacheTier())) + : lvlStorage.roundedNumBytes(lvlStorage.cacheTier()); + // Chargeable, not total: a rebalance that charged host bytes to the device budget would + // under-size the level and evict live pages to reach it. + size_t minQuota = minQuotaForLevel(level); if (quota < minQuota) { throw std::invalid_argument("Quota " + std::to_string(quota) @@ -1972,7 +1837,7 @@ void StorageManager::adjustCacheLevel(CacheLevel level, std::optional ne throw std::invalid_argument("Quota " + std::to_string(quota) + " exceeds the reserved maximum of " + std::to_string(*maxQuota) + " for this cache level (raise the tier's max_quota)"); } - auto newNumSlots = lvlStorage.computeSlotCountList(ratioList, minSlots, quota); + auto newNumSlots = slotCountsForLevel(level, quota, ratioList); TLLM_CHECK_DEBUG(isLastLevel(level) || persistentPages == nullptr); @@ -2039,7 +1904,7 @@ TypedVec StorageManager::ratioFromLength( int const nonStaleShared = std::min(sharedRange.length() - intersect(stale, sharedRange).length(), nonStale); int const nonStaleBeam = nonStale - nonStaleShared; - numRequiredBlocks = std::max(nonStaleShared + effBeamWidth * nonStaleBeam, 1); + numRequiredBlocks = nonStaleShared + effBeamWidth * nonStaleBeam; } numBytes[lifeCycle] = static_cast(numRequiredBlocks) * slotBytes; } @@ -2082,7 +1947,8 @@ TypedVec StorageManager::projectPoolGroupRatio( } TypedVec StorageManager::ratioFromBatch(BatchDesc const& batch, int tokensPerBlock, - std::optional const& swaScratchReuse, size_t granularity) const + std::optional const& swaScratchReuse, size_t deviceGranularity, + TypedVec> const& sparsePools) const { auto numSlots = computeSlotsForBatch(batch, tokensPerBlock, swaScratchReuse); // A soft preference, not a feasibility floor: keep one optional endpoint @@ -2094,7 +1960,7 @@ TypedVec StorageManager::ratioFromBatch(BatchDesc const& bat numSlots[lc] += slotCountValueFromSize(batch.kvCaches.size()); } } - return normalizeToRatio(slotsToBytes(numSlots, granularity)); + return normalizeToRatio(slotsToBytes(numSlots, deviceGranularity, sparsePools)); } // --------------------------------------------------------------------------- @@ -2274,30 +2140,21 @@ TypedVec StorageManager::computePoolGroupSlotsForBatc // slotsToBytes // --------------------------------------------------------------------------- -TypedVec StorageManager::slotsToBytes( - TypedVec const& numSlots, size_t granularity) const +TypedVec StorageManager::slotsToBytes(TypedVec const& numSlots, + size_t deviceGranularity, TypedVec> const& sparsePools) const { TypedVec numBytes(numLifeCycles(), 0); for (LifeCycleId lifeCycle{0}; lifeCycle < numSlots.size(); ++lifeCycle) { auto const poolGroup = getPoolGroupIndex(kHotLevel, lifeCycle); - for (auto const poolSize : slotSize(kHotLevel, poolGroup)) - { - numBytes[lifeCycle] += roundUp(slotCountToSizeT(numSlots[lifeCycle]) * poolSize, granularity); - } - } - return numBytes; -} - -TypedVec StorageManager::slotsToBytes( - TypedVec const& numSlots, size_t granularity) const -{ - TypedVec numBytes(numPoolGroups(), 0); - for (PoolGroupIndex pgIdx{0}; pgIdx < numSlots.size(); ++pgIdx) - { - for (auto poolSize : slotSize(pgIdx)) + auto const& sparse = sparsePools.at(poolGroup); + auto const& sizes = slotSize(kHotLevel, poolGroup); + for (PoolIndex pool{0}; pool < sizes.size(); ++pool) { - numBytes[pgIdx] += roundUp(slotCountToSizeT(numSlots[pgIdx]) * poolSize, granularity); + auto const medium + = std::ranges::find(sparse, pool) == sparse.end() ? CacheTier::GPU_MEM : CacheTier::HOST_MEM; + numBytes[lifeCycle] += roundUp( + slotCountToSizeT(numSlots[lifeCycle]) * sizes[pool], hotPoolSizeGranularity(medium, deviceGranularity)); } } return numBytes; @@ -2309,29 +2166,60 @@ TypedVec StorageManager::slotsToBytes( TypedVec StorageManager::computeSlotCountForLevel(CacheTierConfig const& tierConfig, TypedVec> const& slotSizeLists, - TypedVec const& ratio, TypedVec const& minSlots) const + TypedVec const& ratio, TypedVec const& minSlots, + size_t gpuPhysMemSize, std::optional const& secondBudget) { size_t quota = cacheTierQuota(tierConfig); - size_t const granularity = tierCommitUnit(tierConfig); - quota = std::max(minQuotaForLevel(slotSizeLists, granularity, minSlots), roundUp(quota, granularity)); - return CacheLevelStorage::ratioToSlotCountList(quota, slotSizeLists, ratio, granularity, minSlots); + size_t const granularity = tierCommitUnit(tierConfig, gpuPhysMemSize); + quota = std::max(minQuotaForLevel(slotSizeLists, granularity, minSlots, secondBudget), roundUp(quota, granularity)); + return CacheLevelStorage::ratioToSlotCountList(quota, slotSizeLists, ratio, granularity, minSlots, secondBudget); } -size_t StorageManager::tierCommitUnit(CacheTierConfig const& tierConfig) const +//! Host-backed pools round to HostMem::sizeGranularity(), not to the device granularity, so the +//! budget is counted in that grain: charging them in device grains would over-count every +//! pool of every group by up to one device grain. +std::optional StorageManager::_secondaryBudget(CacheLevelStorage const& level) const +{ + auto const medium = level.secondaryMedium(); + if (!medium.has_value()) + { + // Nothing here sits outside what this level's quota pays for. + return std::nullopt; + } + return _secondaryBudget(level.poolsOfMedium(*medium), level.poolSizeGranularity(*medium), mHostSparsePoolQuota); +} + +std::optional StorageManager::_secondaryBudget( + TypedVec> secondaryPools, size_t granularity, std::optional quota) +{ + // Returned whenever the level has such a pool, bounded or not: the allocator needs to + // know which pools this level's quota does not pay for before it needs to know whether + // anything limits them. Only `quota` carries the latter. + if (std::ranges::none_of(secondaryPools, [](auto const& pools) { return !pools.empty(); })) + { + return std::nullopt; + } + SecondaryPoolBudget budget; + budget.quota = quota; + budget.granularity = granularity; + budget.pools = std::move(secondaryPools); + return budget; +} + +size_t StorageManager::tierCommitUnit(CacheTierConfig const& tierConfig, size_t gpuPhysMemSize) { // Every GPU level shares one allocator, whose unit was derived from the hot // level's quota. A cold GPU level must report that unit rather than one // recomputed from its own quota, which is smaller and would not match the // chunks the allocator actually hands out. return cacheTierOf(tierConfig) == CacheTier::GPU_MEM - ? mGpuPhysMemAllocator->physMemSize() + ? gpuPhysMemSize : tierAllocationGranularity(cacheTierOf(tierConfig), cacheTierQuota(tierConfig)); } -HostMemBackingOptions StorageManager::tierBackingOptions(CacheTierConfig const& tierConfig) const +HostMemBackingOptions StorageManager::tierBackingOptions(CacheTierConfig const& tierConfig, size_t gpuPhysMemSize) { HostMemBackingOptions options; - options.commitUnit = tierCommitUnit(tierConfig); if (auto const* host = std::get_if(&tierConfig)) { options.allowRemoteNumaFallback = host->allowRemoteNumaFallback; @@ -2343,15 +2231,39 @@ HostMemBackingOptions StorageManager::tierBackingOptions(CacheTierConfig const& // minQuotaForLevel // --------------------------------------------------------------------------- +TypedVec StorageManager::_minSlotsFor(CacheLevel level) const +{ + return level == kHotLevel ? mHotMinSlots : TypedVec(numPoolGroups(level), SlotCount{1}); +} + +size_t StorageManager::minQuotaForLevel(CacheLevel level) const +{ + auto const& storage = *mLevels.at(level).storage; + return minQuotaForLevel(storage.slotSizeLists(), storage.poolSizeGranularity(storage.cacheTier()), + _minSlotsFor(level), _secondaryBudget(storage)); +} + +TypedVec StorageManager::slotCountsForLevel( + CacheLevel level, size_t quota, TypedVec const& ratio) const +{ + auto const& storage = *mLevels.at(level).storage; + return CacheLevelStorage::ratioToSlotCountList(quota, storage.slotSizeLists(), ratio, + storage.poolSizeGranularity(storage.cacheTier()), _minSlotsFor(level), _secondaryBudget(storage)); +} + size_t StorageManager::minQuotaForLevel(TypedVec> const& slotSizeLists, - size_t granularity, TypedVec const& minSlots) const + size_t granularity, TypedVec const& minSlots, + std::optional const& secondBudget) { + auto const pools = CacheLevelStorage::poolsPaidByQuota(slotSizeLists, secondBudget); size_t total = 0; for (PoolGroupIndex pgIdx{0}; pgIdx < slotSizeLists.size(); ++pgIdx) { - for (auto slotSize : slotSizeLists[pgIdx]) + // Only the pools this quota pays for: a sparse pool's floor is charged to the host + // budget instead, and counting it here would inflate the device quota's own. + for (PoolIndex poolIdx : pools[pgIdx]) { - total += roundUp(slotCountToSizeT(minSlots[pgIdx]) * slotSize, granularity); + total += roundUp(slotCountToSizeT(minSlots[pgIdx]) * slotSizeLists[pgIdx][poolIdx], granularity); } } return total; @@ -2365,10 +2277,12 @@ TypedVec StorageManager::constrainPoolGroupRatio( TypedVec const& ratio) const { auto& gpuStorage = *mLevels[kHotLevel].storage; - size_t granularity = gpuStorage.poolSizeGranularity(); - auto slotCountList = gpuStorage.computeSlotCountList(ratio, mMinSlots); - auto numBytes = slotsToBytes(slotCountList, granularity); - return normalizeToRatio(numBytes); + auto slotCountList = slotCountsForLevel(kHotLevel, gpuStorage.roundedNumBytes(gpuStorage.cacheTier()), ratio); + // All-bytes, each medium in its own grain -- the same measure ratioList() reports, so + // the target this derives and the layout a rebalance then achieves are comparable. They + // have to be: _needAdjustment treats this as a prediction of the allocator's own output + // and re-fires until the two agree. + return normalizeToRatio(gpuStorage.allMediaBytes(slotCountList)); } } // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.h index 2b423590aeb7..fd4b8366fd10 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.h @@ -101,7 +101,7 @@ class CacheLevelManager CacheLevelManager(TypedVec const& lifeCycleGrouping, CacheLevel cacheLevel, CacheTierConfig const& tierConfig, TypedVec const& slotDescList, TypedVec const& slotCountList, PooledPhysMemAllocator* gpuPhysMemAllocator = nullptr, - HostMemBackingOptions const& options = {}); + HostMemBackingOptions const& options = {}, std::optional sparseReservation = std::nullopt); CacheLevel cacheLevel; CacheTier cacheTier; @@ -147,6 +147,8 @@ class StorageManager : public std::enable_shared_from_this TypedVec const& numSlotsPerLc, MigrationRecorder const& migrationRecorder = {}, DropRecorder const& dropRecorder = {}); + // A sparse life cycle allocates here too: its pages are host memory mapped into a + // hot-level pool, not a separate level. TypedVec> newGpuSlots(TypedVec const& numSlotsPerLc, MigrationRecorder const& migrationRecorder = {}, DropRecorder const& dropRecorder = {}); @@ -197,18 +199,9 @@ class StorageManager : public std::enable_shared_from_this void batchedMigrate( CacheLevel dstLevel, std::vector> const& pages, MigrationRecorder const& migrationRecorder); - //! Demote complete locked sparse pages on the requesting owner's stream, updating every live owner. - //! Caller holds the manager's exclusive lock. Allocation/copy failure preserves GPU ownership. - void offloadSparsePages(KvCache& requestingCache, std::vector> const& pages, - MigrationRecorder const& migrationRecorder = {}, DropRecorder const& dropRecorder = {}); - - //! Promote locked sparse host pages into one shared GPU allocation per page, updating all owners. - //! Caller holds the manager's exclusive lock. Allocation/copy failure preserves host ownership. - void promoteSparsePages(CUstream stream, std::vector> const& pages, - MigrationRecorder const& migrationRecorder = {}, DropRecorder const& dropRecorder = {}); - // Best-effort migration of grouped pages to a destination cache level. Returns how many pages - // it moved off the disk tier. If a later batch fails, earlier batches may already have migrated. + // it moved off the disk tier, counted per migrated batch rather than per page. A throw reports + // nothing, which in practice means slot preparation failed before anything moved. int64_t prefetch( CacheLevel dstLevel, TypedVec>>> const& pages); @@ -315,6 +308,19 @@ class StorageManager : public std::enable_shared_from_this // Cache tier for a given level. CacheTier cacheTier(CacheLevel level) const; + //! The medium of one pool's pages. Prefer this to cacheTier(level): the two agree for + //! every pool but a host-backed one of a GPU-tier group, where the level's tier names + //! device memory for a host address. + [[nodiscard]] CacheTier poolMedium(CacheLevel level, PoolGroupIndex pgIdx, PoolIndex poolIdx) const; + + //! Bytes `level` holds in pools whose pages are `medium`, each rounded to that + //! medium's grain. Zero when the level has no pool of it. + [[nodiscard]] size_t getNumBytes(CacheLevel level, CacheTier medium) const; + + //! The medium backing sparse buffers, or nullopt when no life cycle is sparse. + //! The hot level's second medium: host pages the GPU reads through a device mapping. + [[nodiscard]] std::optional sparseBackingMedium() const; + // NOTE: Python's get_statistics(level) returns a list over all pool groups. // C++ takes a single pgIdx for flexibility; the nanobind wrapper loops over // all pool groups to match Python's signature. @@ -323,7 +329,7 @@ class StorageManager : public std::enable_shared_from_this float getOverallUtilization(CacheLevel level = kHotLevel) const; // Pool-group slot count (number of pages). - SlotCount numSlots(PoolGroupIndex pgIdx, CacheLevel level = kHotLevel) const; + [[nodiscard]] SlotCount numSlots(PoolGroupIndex pgIdx, CacheLevel level = kHotLevel) const; // Layer-to-lifecycle mapping (for KvCacheManager queries). std::unordered_map const& layerToLifeCycleIds() const noexcept @@ -355,6 +361,8 @@ class StorageManager : public std::enable_shared_from_this BatchDesc const& batch, int tokensPerBlock, std::optional const& swaScratchReuse) const; // Constraint floors apply only to the hot level. Lifecycle demands remain available only as ratio weights. + // + // Both read the hot pool-group mapping and slot descs, so they run after those are built. TypedVec computePoolGroupMinSlotsFromConstraints( std::vector const& constraints, int tokensPerBlock, std::optional const& swaScratchReuse, float maxUtilForResume = 1.0f) const; @@ -362,29 +370,67 @@ class StorageManager : public std::enable_shared_from_this int tokensPerBlock, std::optional const& swaScratchReuse, float maxUtilForResume = 1.0f) const; TypedVec ratioFromBatch(BatchDesc const& batch, int tokensPerBlock, - std::optional const& swaScratchReuse, size_t granularity) const; + std::optional const& swaScratchReuse, size_t deviceGranularity, + TypedVec> const& sparsePools) const; // Project source-level lifecycle byte weights onto destination pool groups while preserving the implied // lifecycle slot-count proportions. TypedVec projectPoolGroupRatio( CacheLevel srcLevel, CacheLevel dstLevel, TypedVec const& srcLifeCycleRatio) const; - TypedVec slotsToBytes( - TypedVec const& numSlots, size_t granularity) const; - TypedVec slotsToBytes( - TypedVec const& numSlots, size_t granularity) const; + //! Bytes a hot-level slot count costs, each pool rounded by hotPoolSizeGranularity() to the + //! grain of the medium backing it -- the same measure as CacheLevelStorage::allMediaBytes(), + //! which the resulting ratio is compared against. The storage does not exist yet here, so + //! `sparsePools` names each group's sparse pools in place of PoolGroupBase::poolMedium(). + TypedVec slotsToBytes(TypedVec const& numSlots, + size_t deviceGranularity, TypedVec> const& sparsePools) const; + //! Allocation granularity for a tier, decided once from its configured //! quota. Fixed for the tier's lifetime, so every size derived from it stays //! consistent with what the pools were built for. - [[nodiscard]] size_t tierCommitUnit(CacheTierConfig const& tierConfig) const; + //! + //! `gpuPhysMemSize` is the shared GPU allocator's unit, which every GPU level reports. + [[nodiscard]] static size_t tierCommitUnit(CacheTierConfig const& tierConfig, size_t gpuPhysMemSize); //! Backing options for a tier: its fixed allocation unit together with the - //! page-placement policy the tier's config asks for. - [[nodiscard]] HostMemBackingOptions tierBackingOptions(CacheTierConfig const& tierConfig) const; - - TypedVec computeSlotCountForLevel(CacheTierConfig const& tierConfig, + //! page-placement policy the tier's config asks for. `gpuPhysMemSize` is the shared + //! GPU allocator's unit, as tierCommitUnit takes it. + [[nodiscard]] static HostMemBackingOptions tierBackingOptions( + CacheTierConfig const& tierConfig, size_t gpuPhysMemSize); + + //! `slotSizeLists` holds every pool's size. The level's quota pays for all of them bar + //! the ones `secondBudget` names, which are host memory charged to it instead. + [[nodiscard]] static TypedVec computeSlotCountForLevel(CacheTierConfig const& tierConfig, TypedVec> const& slotSizeLists, - TypedVec const& ratio, TypedVec const& minSlots) const; - size_t minQuotaForLevel(TypedVec> const& slotSizeLists, - size_t granularity, TypedVec const& minSlots) const; + TypedVec const& ratio, TypedVec const& minSlots, + size_t gpuPhysMemSize, std::optional const& secondBudget = std::nullopt); + //! For the config stage, which sizes a level before it exists: every fact it needs is + //! named in the call rather than read off the manager. + [[nodiscard]] static size_t minQuotaForLevel( + TypedVec> const& slotSizeLists, size_t granularity, + TypedVec const& minSlots, + std::optional const& secondBudget = std::nullopt); + + //! The same floor for a level that exists, reading its sizes, grain, second budget and + //! slot floors off it. + [[nodiscard]] size_t minQuotaForLevel(CacheLevel level) const; + + //! Slot counts `quota` buys at `ratio`, for a level that exists. + [[nodiscard]] TypedVec slotCountsForLevel( + CacheLevel level, size_t quota, TypedVec const& ratio) const; + + //! Slot floors for a level: the hot level's feasibility constraints, and elsewhere the + //! structural one slot per pool group. + [[nodiscard]] TypedVec _minSlotsFor(CacheLevel level) const; + + //! The second medium's budget over `secondaryPools`, as the allocator takes it, or nullopt + //! when no pool is named. `granularity` is the grain those pools round to; `quota` is the + //! configured bound, if any. + [[nodiscard]] static std::optional _secondaryBudget( + TypedVec> secondaryPools, size_t granularity, + std::optional quota); + + //! The second medium's budget for a level that already exists, reading the pools and + //! their grain off it. + [[nodiscard]] std::optional _secondaryBudget(CacheLevelStorage const& level) const; // Internal helpers. // Register the slot descriptors for the next cache level and derive its slot-size table. @@ -404,8 +450,6 @@ class StorageManager : public std::enable_shared_from_this std::optional> _batchedMigrate(CacheLevel dstLevel, CacheLevel srcLevel, std::vector> const& srcPages, bool updateSrc, MigrationRecorder const& migrationRecorder = {}, bool defrag = false); - void _migrateLockedSparsePages(CUstream stream, CacheLevel dstLevel, std::vector> const& pages, - KvCache const* offloadingCache, MigrationRecorder const& migrationRecorder, DropRecorder const& dropRecorder); [[nodiscard]] LayerGroupId getMigrationBatchingLayerGroupId( CacheLevel dstLevel, CacheLevel srcLevel, LifeCycleId lifeCycle) const; void submitMigrationBatch(CacheLevel dstLevel, CacheLevel srcLevel, LayerGroupId batchingLayerGroupId, @@ -470,7 +514,15 @@ class StorageManager : public std::enable_shared_from_this TypedVec> mSlotDescLists; // Slot sizes per (level, pool group), built from mSlotDescLists. TypedVec>> mSlotSizes; - TypedVec mMinSlots; + //! Hot-level slot floors, from the batch constraints and scaled by 1/maxUtilForResume. + TypedVec mHotMinSlots; + + //! Host memory the hot level's host-backed pools may take: the configured quota, or the + //! default construction fills in when it is unset. Kept because a rebalance moves the + //! slot counts those pools follow, so the bound has to be reapplied rather than assumed + //! still satisfied. + size_t mHostSparsePoolQuota = 0; + // All GPU cache levels borrow this allocator. It must outlive mLevels. std::unique_ptr mGpuPhysMemAllocator; RetentionExpiryQueue mRetentionExpiries; diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/generator.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/generator.h new file mode 100644 index 000000000000..f7cafa8d143d --- /dev/null +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/generator.h @@ -0,0 +1,348 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/common/assert.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 +{ + +// --------------------------------------------------------------------------- +// Generator — a coroutine that yields values, optionally receiving one back. +// +// `Generator` is the familiar lazy sequence, close enough to C++23's +// std::generator to be replaced by it later: nothing runs until the first +// iteration, and iterating resumes the body to its next `co_yield`. +// +// Generator evens(int n) { for (int i = 0; i < n; i += 2) co_yield i; } +// for (int v : evens(10)) { ... } +// +// `Generator` additionally carries a value back in, because `co_yield v` +// is `co_await promise.yield_value(v)` and so has a value of its own: +// +// Generator layout() { +// int32_t* base = co_yield elements; // yields a count, resumes with memory +// ... +// } +// +// That direction is what lets a caller collect what several bodies need, satisfy +// them all at once, and let each continue -- which a plain sequence of calls +// cannot express, since the first would have to be satisfied before the last had +// been asked. Each body stays one function, so what its first half computed is +// still in scope for its second: the values live in the coroutine frame. +// +// `Generator` additionally carries a result out of `co_return`, so a +// body that was a function returning Return stays one: +// +// Generator resize(...) { +// for (lc ...) { auto slots = co_yield needed[lc]; ... } +// co_return true; +// } +// +// A body may yield any number of times; each suspension is a separate request, and +// value()/send() service them one at a time. That is what lets a caller drive several +// bodies round by round -- collecting every body's request for one life cycle, meeting +// them together, and only then letting each continue -- which is how allocations across +// bodies can be coalesced, or what one body released reused by another. A driver that +// expects a single round should check done() afterwards rather than assume it. +// +// A body's parameters are copied into its frame, so they are safe to hold across a +// suspension. A lambda's captures are not: the frame only refers to the closure, so an +// immediately-invoked capturing lambda leaves the frame pointing at a destroyed object. +// Write such a body as a function taking what it needs. +// +// Not std::generator: no recursive `elements_of`, no allocator customization, and +// the iterator is single-pass. Move-only, and an exception from the body is +// rethrown by whichever call resumed it. +// --------------------------------------------------------------------------- +namespace detail +{ + +//! Supplies the promise with return_value or return_void, which a promise may have one +//! of but never both. Both clear the pending value, so a finished body is told apart +//! from one suspended on a request by asking for that value. +template +struct GeneratorReturn +{ + void return_value(Return result) + { + auto* self = static_cast(this); + self->mValue.reset(); + mReturn = std::move(result); + } + + std::optional mReturn; +}; + +template +struct GeneratorReturn +{ + void return_void() noexcept + { + static_cast(this)->mValue.reset(); + } +}; + +} // namespace detail + +template +class Generator +{ + static constexpr bool kSends = !std::is_void_v; + static constexpr bool kReturns = !std::is_void_v; + +public: + struct promise_type : detail::GeneratorReturn + { + Generator get_return_object() + { + return Generator{std::coroutine_handle::from_promise(*this)}; + } + + //! Lazy, as std::generator is: the body runs on the first pull, not on the call. + [[nodiscard]] std::suspend_always initial_suspend() const noexcept + { + return {}; + } + + //! Suspends at the end so the frame outlives the body, which is what lets a + //! caller read a final exception out of the promise. + [[nodiscard]] std::suspend_always final_suspend() const noexcept + { + return {}; + } + + void unhandled_exception() noexcept + { + mError = std::current_exception(); + } + + auto yield_value(T value) noexcept + { + mValue = std::move(value); + if constexpr (kSends) + { + struct Resumption + { + promise_type* promise; + + [[nodiscard]] bool await_ready() const noexcept + { + return false; + } + + void await_suspend(std::coroutine_handle<>) const noexcept {} + + [[nodiscard]] Send await_resume() const + { + // Resumed without a value: the body would read whatever this + // returned and carry on as if it were real. + TLLM_CHECK_WITH_INFO( + promise->mSent.has_value(), "A two-way generator was resumed without being sent a value"); + return std::move(*promise->mSent); + } + }; + + return Resumption{this}; + } + else + { + return std::suspend_always{}; + } + } + + std::optional mValue; + std::exception_ptr mError; + // Only the two-way form carries one; std::monostate keeps the one-way form from + // paying for it. + std::conditional_t>, std::monostate> mSent; + }; + + explicit Generator(std::coroutine_handle handle) + : mHandle(handle) + { + } + + Generator(Generator&& other) noexcept + : mHandle(std::exchange(other.mHandle, {})) + , mStarted(std::exchange(other.mStarted, false)) + { + } + + Generator(Generator const&) = delete; + Generator& operator=(Generator const&) = delete; + Generator& operator=(Generator&&) = delete; + + //! Destroys the frame, running the destructors of a suspended body's locals. + ~Generator() + { + if (mHandle) + { + mHandle.destroy(); + } + } + + //! Whether the body has run to completion. + [[nodiscard]] bool done() const + { + _rethrow(); + return !mHandle || mHandle.done(); + } + + //! The value the body is suspended on, priming it if it has not started. + //! Empty once the body has returned. + [[nodiscard]] std::optional const& value() + { + _prime(); + return mHandle.promise().mValue; + } + + //! What the body returned, once it has finished. Empty until then. + //! + //! A template so that the signature is formed only where it is called. A + //! requires-clause gates overload resolution, not the declaration, so a void Return + //! would otherwise reach std::optional here and make the class ill-formed. + template + [[nodiscard]] std::optional const& result() const requires(kReturns) + { + _rethrow(); + return mHandle.promise().mReturn; + } + + //! Resume a one-way body. Available only when nothing is sent back. + void next() requires(!kSends) + { + _prime(); + _step(); + } + + //! Hand the body a value and resume it. The `co_yield` it was suspended on + //! evaluates to `sent`. + //! + //! A template for the same reason as result(), with S kept out of deduction so the + //! parameter stays exactly Send rather than whatever the caller passed. + template + void send(std::type_identity_t sent) requires(kSends) + { + _prime(); + if (mHandle.done()) + { + return; + } + mHandle.promise().mSent = std::move(sent); + _step(); + } + + // ---- one-way iteration, so `for (auto v : gen)` works ------------------- + + class Iterator + { + public: + using iterator_category = std::input_iterator_tag; + using value_type = T; + using difference_type = std::ptrdiff_t; + + explicit Iterator(Generator* owner) noexcept + : mOwner(owner) + { + } + + T const& operator*() const + { + return *mOwner->mHandle.promise().mValue; + } + + Iterator& operator++() + { + mOwner->_step(); + if (mOwner->mHandle.done()) + { + mOwner = nullptr; + } + return *this; + } + + void operator++(int) + { + ++*this; + } + + bool operator==(std::default_sentinel_t) const noexcept + { + return mOwner == nullptr; + } + + private: + Generator* mOwner; + }; + + [[nodiscard]] Iterator begin() requires(!kSends) + { + _prime(); + return Iterator{mHandle.done() ? nullptr : this}; + } + + [[nodiscard]] std::default_sentinel_t end() const noexcept requires(!kSends) + { + return {}; + } + +private: + //! Run to the first yield. Idempotent, so a caller need not track whether the body + //! has started. + void _prime() + { + if (mHandle && !mStarted) + { + mStarted = true; + mHandle.resume(); + _rethrow(); + } + } + + void _step() + { + if (mHandle && !mHandle.done()) + { + mHandle.resume(); + _rethrow(); + } + } + + void _rethrow() const + { + if (mHandle && mHandle.promise().mError) + { + std::rethrow_exception(mHandle.promise().mError); + } + } + + std::coroutine_handle mHandle; + bool mStarted = false; +}; + +} // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/hostMem.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/hostMem.cpp index 75e533250c4f..532b86d43a24 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/hostMem.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/hostMem.cpp @@ -166,30 +166,19 @@ constexpr int kInitialFillHelpers = 7; } // namespace -size_t HostMem::chunkCount() const noexcept -{ - return mChunkSize == 0 ? 0 : divUp(mSize, mChunkSize); -} - -//! Commits [beginOffset, mSize) in chunk-sized steps, publishing the watermark -//! after each one. -//! -//! `beginOffset` is a byte offset rather than a chunk index because a grow -//! resumes at the old size, which need not fall on a chunk boundary. The step -//! that straddles the boundary is committed as a short range, so no byte is -//! ever committed twice. +//! Steps the backing to its target, publishing the watermark after each step. //! //! With `helperCount` > 0 the range is opened up front and that many threads -//! populate it chunk by chunk ahead of the commits. Populating in parallel -//! slows everything else in the process, so only the initial fill, which runs -//! before serving, uses helpers. Until it ends, a reader that skips -//! waitUsable() reads zeros past the watermark instead of faulting. +//! populate it step by step ahead of the commits. Populating in parallel slows +//! everything else in the process, so only the initial fill, which runs before +//! serving, uses helpers. Until it ends, a reader that skips waitUsable() reads +//! zeros past the watermark instead of faulting. //! -//! A step is populated by whichever thread claims it first, and one a helper -//! has claimed is committed only once that helper is done with it; a helper -//! that fails hands its step back. Where MADV_POPULATE_WRITE is unavailable -//! populating means zeroing, which must not reach a step already handed out. -void HostMem::commitRange(size_t beginOffset, int helperCount) +//! A step is populated by whichever thread claims it first, and one a helper has +//! claimed is committed only once that helper is done with it; a helper that +//! fails hands its step back. Where MADV_POPULATE_WRITE is unavailable populating +//! means zeroing, which must not reach a step already handed out. +void HostMem::fillToTarget(int helperCount) { enum StepState : uint8_t { @@ -198,14 +187,19 @@ void HostMem::commitRange(size_t beginOffset, int helperCount) kPopulated }; - bool const useHelpers = helperCount > 0 && mSize - beginOffset > mChunkSize; - size_t const firstChunk = beginOffset / mChunkSize; + size_t const begin = mBacking.committedSize(); + size_t const target = mBacking.targetSize(); + size_t const chunk = HostMemBacking::commitChunkSize(); + // Only a grow is populated ahead: a shrink releases pages rather than taking + // them, and a single-step range has nothing to overlap with. + bool const useHelpers = helperCount > 0 && target > begin && target - begin > chunk; + size_t const firstChunk = begin / chunk; // One entry per step, indexed by its chunk relative to firstChunk. - std::vector> stepStates(useHelpers ? chunkCount() - firstChunk : 0); + std::vector> stepStates(useHelpers ? divUp(target, chunk) - firstChunk : 0); auto const stepRange = [&](size_t step) { - size_t const chunk = firstChunk + step; - return std::pair{std::max(chunk * mChunkSize, beginOffset), std::min((chunk + 1) * mChunkSize, mSize)}; + size_t const chunkIndex = firstChunk + step; + return std::pair{std::max(chunkIndex * chunk, begin), std::min((chunkIndex + 1) * chunk, target)}; }; std::atomic nextHelperStep{0}; // Declared after what they use: destruction stops and joins them on every @@ -213,7 +207,7 @@ void HostMem::commitRange(size_t beginOffset, int helperCount) std::vector helpers; if (useHelpers) { - mBacking.open(beginOffset, mSize - beginOffset); + mBacking.open(begin, target - begin); try { for (int i = 0; i < helperCount; ++i) @@ -234,14 +228,14 @@ void HostMem::commitRange(size_t beginOffset, int helperCount) { continue; } - auto const [begin, end] = stepRange(step); + auto const [stepBegin, stepEnd] = stepRange(step); try { - mBacking.populate(begin, end - begin); + mBacking.populate(stepBegin, stepEnd - stepBegin); } catch (...) { - // Handed back for commit() to populate. + // Handed back for the committing thread to populate. state.store(kPending, std::memory_order_release); state.notify_all(); throw; @@ -252,19 +246,19 @@ void HostMem::commitRange(size_t beginOffset, int helperCount) } catch (...) { - // commit() populates whatever a helper leaves and - // reports any failure itself. + // The committing thread populates whatever a helper + // leaves, and reports any failure itself. } }); } } catch (std::system_error const&) { - // Out of threads: those running and commit() cover the range. + // Out of threads: those running and the committing thread cover the range. } } - // Whether a helper populated the step; if not, the caller now owns it. + // Whether a helper populated the step; if not, this thread now owns it. auto const claimOrAwaitStep = [&](size_t step) { auto& state = stepStates[step]; @@ -283,25 +277,32 @@ void HostMem::commitRange(size_t beginOffset, int helperCount) } }; - size_t offset = beginOffset; - while (offset < mSize) + while (mBacking.committedSize() != mBacking.targetSize()) { - size_t const end = std::min(roundUp(offset + 1, mChunkSize), mSize); - bool const populated = useHelpers && claimOrAwaitStep(offset / mChunkSize - firstChunk); - mBacking.commit(offset, end - offset, populated); - offset = end; - // A chunk counts as usable only once it is committed to its end, so a - // partially committed trailing chunk does not advance the watermark. - mFilledChunks.store( - static_cast(offset == mSize ? chunkCount() : offset / mChunkSize), std::memory_order_release); - mFilledChunks.notify_all(); + bool const populated = useHelpers && claimOrAwaitStep(mBacking.committedSize() / chunk - firstChunk); + // Clamped because a shrink's steps report a registered end that is still + // above the target; the watermark must never run ahead of it. + mUsable.store(std::min(mBacking.resizeStep(populated), mBacking.targetSize()), std::memory_order_release); + mUsable.notify_all(); } + // The target rather than the registered end, which trails it by whatever + // does not fill a whole register unit. No consumer can use a remainder + // smaller than that, and publishing the shorter value would hang a wait for + // the whole range. + // + // Checked rather than assumed: a larger gap publishes memory the GPU cannot + // reach, and nothing notices until a copy fails far from here. + TLLM_CHECK_WITH_INFO(mBacking.targetSize() - mBacking.registeredSize() < mBacking.registerUnit(), + "HostMem: %zu of %zu bytes unregistered at the target, which is a whole register unit of %zu or more", + mBacking.targetSize() - mBacking.registeredSize(), mBacking.targetSize(), mBacking.registerUnit()); + mUsable.store(mBacking.targetSize(), std::memory_order_release); + mUsable.notify_all(); } -void HostMem::startFill(size_t beginOffset, int helperCount) +void HostMem::startFill(int helperCount) { mFillWorker = std::thread( - [this, beginOffset, helperCount] + [this, helperCount] { // A new thread starts on device 0 whatever the creator selected, and // registration needs a context on the right one. Selecting it here @@ -317,22 +318,22 @@ void HostMem::startFill(size_t beginOffset, int helperCount) cudaStreamCaptureMode captureMode = cudaStreamCaptureModeRelaxed; TLLM_CUDA_CHECK(cudaThreadExchangeStreamCaptureMode(&captureMode)); runtime::TopologyDetector::getInstance().bindThreadByCurrentGpu(); - commitRange(beginOffset, helperCount); + fillToTarget(helperCount); } catch (...) { - if (helperCount > 0 && mSize - beginOffset > mChunkSize) + if (helperCount > 0) { - // Helpers may have populated past the failed commit, which - // released only its own step. - size_t const committed = mFilledChunks.load(std::memory_order_relaxed) * mChunkSize; - mBacking.decommit(committed, mSize - committed); + // Helpers may have populated past the failed step, which + // released only its own range. + size_t const committed = mBacking.committedSize(); + mBacking.discardUncommitted(committed, mBacking.targetSize() - committed); } // Stored before the sentinel: a waiter that observes the // sentinel through an acquire load is guaranteed to see it. mFillFailure = std::current_exception(); - mFilledChunks.store(kFillFailed, std::memory_order_release); - mFilledChunks.notify_all(); + mUsable.store(kFillFailed, std::memory_order_release); + mUsable.notify_all(); } }); } @@ -347,18 +348,17 @@ void HostMem::joinFill() noexcept void HostMem::waitUsable(size_t endOffset) const { - if (endOffset == 0 || mChunkSize == 0) + if (endOffset == 0) { return; } TLLM_CHECK_DEBUG(endOffset <= mSize); - auto const needed = static_cast(divUp(endOffset, mChunkSize)); - uint32_t current = mFilledChunks.load(std::memory_order_acquire); - while (current < needed) + size_t current = mUsable.load(std::memory_order_acquire); + while (current < endOffset) { - mFilledChunks.wait(current, std::memory_order_acquire); - current = mFilledChunks.load(std::memory_order_acquire); + mUsable.wait(current, std::memory_order_acquire); + current = mUsable.load(std::memory_order_acquire); } if (current == kFillFailed) { @@ -367,7 +367,7 @@ void HostMem::waitUsable(size_t endOffset) const } HostMem::HostMem(size_t vmSize, size_t initialSize, HostMemBackingOptions const& options, bool waitForFill) - : mBacking(options) + : mBacking(vmSize, options) , mWaitForFill(waitForFill) { TLLM_CUDA_CHECK(cudaGetDevice(&mDevice)); @@ -380,16 +380,14 @@ HostMem::HostMem(size_t vmSize, size_t initialSize, HostMemBackingOptions const& size_t const granularity = mBacking.commitGranularity(); TLLM_CHECK_WITH_INFO(initialSize % granularity == 0, "HostMem size %zu is not a multiple of the commit granularity %zu", initialSize, granularity); - mChunkSize = roundUp(kFillChunkSize, granularity); try { - // The reservation is rounded up so that the largest accepted resize is - // itself a valid, aligned size. - mVmSize = roundUp(vmSize, granularity); + mVmSize = mBacking.reservedSize(); mSize = initialSize; - mAddr = mBacking.reserve(mVmSize); - startFill(0, kInitialFillHelpers); + mAddr = mBacking.address(); + mBacking.setTargetSize(mSize); + startFill(kInitialFillHelpers); if (mWaitForFill) { waitUsable(mSize); @@ -420,13 +418,24 @@ HostMem::~HostMem() void HostMem::destroy() { joinFill(); + // Tearing down with a waiter outstanding is a caller error; the sentinel makes it one + // that reports itself rather than blocking forever on a watermark that never rises. + // + // Set before the sentinel and kept afterwards: waitUsable() reads it only once it has + // observed the sentinel through an acquire load, and rethrowing a null is undefined. A + // fill that already failed keeps its own exception, which says more than this one. + if (mFillFailure == nullptr) + { + mFillFailure = std::make_exception_ptr(LogicError("HostMem was destroyed while a wait on it was outstanding")); + } + // Published before the mapping goes: a watermark left standing names memory that is + // already gone, and a reader taking it at its word faults. + mUsable.store(kFillFailed, std::memory_order_release); + mUsable.notify_all(); mBacking.destroy(); mAddr = 0; mSize = 0; mVmSize = 0; - mChunkSize = 0; - mFillFailure = nullptr; - mFilledChunks.store(0, std::memory_order_release); } void HostMem::resize(size_t newSize) @@ -436,7 +445,7 @@ void HostMem::resize(size_t newSize) // The fill must be finished before the mapping changes underneath it, and a // failed fill has to surface here rather than be resized over. joinFill(); - if (mFilledChunks.load(std::memory_order_acquire) == kFillFailed) + if (mUsable.load(std::memory_order_acquire) == kFillFailed) { std::rethrow_exception(mFillFailure); } @@ -451,29 +460,31 @@ void HostMem::resize(size_t newSize) TLLM_CHECK_WITH_INFO(newSize <= mVmSize, "HostMem resize to %zu exceeds the reserved maximum of %zu; raise the tier's maxQuota", newSize, mVmSize); + mBacking.setTargetSize(newSize); if (newSize < mSize) { // Shrink within the existing reservation: only the tail is released and - // the base address is unchanged. The release is exact, because newSize - // is a multiple of the unit the backing allocates in. - mBacking.decommit(newSize, mSize - newSize); + // the base address is unchanged. Done on this thread, because the caller + // has already had to stop every consumer for the range to be releasable + // at all, so there is nothing for a worker to overlap with. mSize = newSize; - mFilledChunks.store(static_cast(chunkCount()), std::memory_order_release); - mFilledChunks.notify_all(); + // Lowered before anything is released, not after, so no waiter is ever + // holding a watermark that covers memory already gone. + mUsable.store(newSize, std::memory_order_release); + // A shrink releases pages rather than taking them, so there is nothing + // for a helper to populate ahead of. + fillToTarget(0); return; } // Grow: the address space is already reserved, so this only commits the new // tail. Everything below keeps its contents and its address. - size_t const oldSize = mSize; mSize = newSize; - // The watermark counts whole chunks, and the old size rounded up to one. - // Growing makes that rounding wrong -- it would claim the new tail is - // already usable -- so it is rebased to the chunks genuinely complete. - mFilledChunks.store(static_cast(oldSize / mChunkSize), std::memory_order_release); try { - startFill(oldSize, 0); + // No helpers once serving has begun: populating in parallel slows every + // other thread in the process. + startFill(0); } catch (...) { @@ -481,8 +492,8 @@ void HostMem::resize(size_t newSize) // wait past the old size would block on a watermark that no worker is // left to advance. mFillFailure = std::current_exception(); - mFilledChunks.store(kFillFailed, std::memory_order_release); - mFilledChunks.notify_all(); + mUsable.store(kFillFailed, std::memory_order_release); + mUsable.notify_all(); throw; } if (mWaitForFill) diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/hostMem.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/hostMem.h index 0288d34537d8..f8354fa6a343 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/hostMem.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/hostMem.h @@ -44,51 +44,48 @@ namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 class HostMem { public: - //! Required alignment of every size passed to HostMem. Only the single-size - //! constructor rounds up on the caller's behalf; elsewhere a size that is - //! not a multiple of this is a bug and is rejected. + //! What every size passed to HostMem must be a multiple of. Only the single-size + //! constructor rounds up on the caller's behalf; elsewhere a size that is not a + //! multiple of this is a bug and is rejected. //! - //! The base is only page-aligned: registration takes any address, and the - //! mprotect and madvise calls round to pages themselves. - static constexpr size_t kAlignment = kHostMemAlignment; - - //! Bytes committed per watermark advance. Large enough that the commit cost - //! dominates the notify, small enough that a waiter near the start of the - //! range is released early. - //! - //! Unrelated to the unit memory is allocated and released in: this only sets - //! how often the fill publishes progress. See HostMemBacking. - static constexpr size_t kFillChunkSize = 512ULL << 20; // 512 MB + //! A granularity of sizes and offsets, not an alignment of addresses: the base + //! comes from mmap and is only page-aligned. + [[nodiscard]] static size_t sizeGranularity() noexcept + { + return hostMemSizeGranularity(); + } - //! Reserves `vmSize` of address space and commits `initialSize` of it. + //! Reserves `vmSize` of address space and commits `initialSize` of it. Both must be + //! multiples of sizeGranularity(); neither is rounded on the caller's behalf, since + //! a reservation quietly made larger is one the caller does not know it can resize + //! into, and resize() rejects the same way. //! //! The base address is fixed for the lifetime of the object: resize() only //! moves the committed boundary within the reservation, and a resize past //! `vmSize` is rejected rather than served by relocating. Reserving costs //! address space, not memory. //! - //! `options.commitUnit` is the granularity memory is allocated and released - //! in, and therefore the alignment every size must satisfy. It is - //! independent of kFillChunkSize, which only paces progress reporting. - //! //! The fill always runs on a worker thread. `waitForFill` selects whether //! the caller is exposed to that: when true, the constructor and resize() //! return only once the whole range is committed, which is what an owner //! that hands out raw pointers without any further bookkeeping needs. //! - //! Pass false only when every reader of this range calls waitUsable() with - //! the offset it is about to touch. Reaching past the watermark faults, so - //! a reader that skips it fails loudly rather than reading zeros. + //! Pass false only when the range is covered by a waitUsable() somewhere + //! above. One wait where the memory is handed out is what that takes, not a + //! wait per read: once a byte is backed it stays backed, so an owner that + //! waits before publishing a sub-range has already covered every later read + //! of it. Reaching past the watermark faults, so a range that nothing waits + //! for fails loudly rather than reading zeros. HostMem(size_t vmSize, size_t initialSize, HostMemBackingOptions const& options = {}, bool waitForFill = true); //! Reserves exactly as much as it commits, for an allocation that never //! resizes. Takes no waitForFill: a caller that cannot resize has no reason //! to observe a partially filled range. //! - //! `size` may be any byte count: it is rounded up to kAlignment, and size() + //! `size` may be any byte count: it is rounded up to sizeGranularity(), and size() //! reports the rounded value. explicit HostMem(size_t size) - : HostMem(roundUp(size, kAlignment), roundUp(size, kAlignment), HostMemBackingOptions{}, + : HostMem(roundUp(size, sizeGranularity()), roundUp(size, sizeGranularity()), HostMemBackingOptions{}, /*waitForFill=*/true) { } @@ -133,28 +130,23 @@ class HostMem private: //! Sentinel stored in the watermark when the worker failed. Distinct from - //! every real chunk count, so a waiter cannot mistake it for progress. - static constexpr uint32_t kFillFailed = UINT32_MAX; + //! every reachable byte count, so a waiter cannot mistake it for progress. + static constexpr size_t kFillFailed = SIZE_MAX; - void commitRange(size_t beginOffset, int helperCount); - void startFill(size_t beginOffset, int helperCount); + void fillToTarget(int helperCount); + void startFill(int helperCount); void joinFill() noexcept; - [[nodiscard]] size_t chunkCount() const noexcept; HostMemBacking mBacking; MemAddress mAddr = 0; size_t mSize = 0; size_t mVmSize = 0; - - //! Bytes committed between watermark advances. A multiple of the backing's - //! commit granularity, so a fill step never lands inside an allocation. - size_t mChunkSize = 0; bool mWaitForFill = true; - //! Read on every allocation, written once per committed chunk. Kept on its - //! own cache line so publishing progress does not invalidate the line that - //! readers are spinning on. - alignas(64) mutable std::atomic mFilledChunks{0}; + //! Bytes a consumer may touch, from the base. Read on every allocation, + //! written once per fill step. Kept on its own cache line so publishing + //! progress does not invalidate the line that readers are spinning on. + alignas(64) mutable std::atomic mUsable{0}; //! Written before the failure sentinel is stored, and only read after that //! sentinel has been observed, so no lock is needed. diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/hostMemBacking.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/hostMemBacking.cpp index fac0d808912b..26c9188a71f4 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/hostMemBacking.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/hostMemBacking.cpp @@ -52,6 +52,29 @@ size_t commitAlignment() return static_cast(::sysconf(_SC_PAGESIZE)); } +} // namespace + +size_t hostMemSizeGranularity() noexcept +{ + // A page is the only hard bound; 2 MiB is the floor because that is the huge page the + // backing asks for, and a size quantized below it would straddle one for no gain. + static size_t const unit = std::max(size_t{2} << 20, commitAlignment()); + return unit; +} + +namespace +{ + +//! Bytes one step populates. Large enough that the call overhead disappears against the +//! faulting, small enough that a caller publishing progress between steps releases an +//! early waiter early. +constexpr size_t kCommitChunkSize = 32ULL << 20; + +//! Bytes one registration covers, before the register unit rounds it down. Near the commit +//! chunk so the two advance together, and small enough that the tail a shrink has to split +//! is cheap to put back. +constexpr size_t kRegisterChunkTarget = 32ULL << 20; + //! The pages [addr, addr + size) touches, widened to whole pages. //! //! For calls that must cover at least the range: populating or opening a few @@ -160,45 +183,41 @@ void bindRangeToGpuNumaNode(MemAddress addr, size_t size, bool allowRemoteFallba } // namespace -HostMemBacking::HostMemBacking(HostMemBackingOptions const& options) - : mUnit(std::max(options.commitUnit, commitAlignment())) +HostMemBacking::HostMemBacking(size_t vmSize, HostMemBackingOptions const& options) + : mUnit(hostMemSizeGranularity()) + , mRegUnit(std::max(options.registerUnit, 1)) + // Whole register units, so no bound falls inside one. + , mRegChunk(std::max(kRegisterChunkTarget / mRegUnit * mRegUnit, mRegUnit)) , mAllowRemoteNumaFallback(options.allowRemoteNumaFallback) { -} - -HostMemBacking::~HostMemBacking() -{ - HostMemBacking::destroy(); -} - -MemAddress HostMemBacking::reserve(size_t size) -{ - TLLM_CHECK_WITH_INFO(mAddr == 0, "HostMem backing reserved twice"); - TLLM_CHECK_DEBUG(size % mUnit == 0); + if (vmSize == 0) + { + return; + } + // The caller rounds: a reservation silently made larger is one the caller does not + // know it can resize into. + TLLM_CHECK_WITH_INFO( + vmSize % mUnit == 0, "HostMemBacking reservation %zu is not a multiple of the granularity %zu", vmSize, mUnit); - // Captured here rather than in commit(): this runs on the thread that owns - // the context, while the fill worker only sets the device. cuCheck(cuCtxGetCurrent(reinterpret_cast(&mContext))); - TLLM_CHECK_WITH_INFO(mContext != nullptr, "HostMemBacking::reserve needs a current CUDA context"); - - // No alignment is asked of the mapping. Unit boundaries line up with slot - // boundaries because slot 0 sits at the base and a unit is a whole number of - // slots, wherever the base happens to land; registration needs no alignment, - // and the madvise and mprotect calls round to pages themselves. The mapped - // length is page-rounded so the guard below covers whole pages. - mMapped = roundUp(size, commitAlignment()); + TLLM_CHECK_WITH_INFO(mContext != nullptr, "HostMemBacking needs a current CUDA context"); + + // No alignment is asked of the mapping: registration needs none, and the + // madvise and mprotect calls round to pages themselves. The mapped length is + // page-rounded so the guard below covers whole pages. + mMapped = roundUp(vmSize, commitAlignment()); mAddr = hostMmap(mMapped); - mSize = size; + mSize = vmSize; //! Applied once over the whole mapping. Per-chunk page-mode advice would //! change VMA flags and split the mapping into one VMA per chunk. hostMadvisePageMode(mAddr, mSize, hostUseThp()); bindRangeToGpuNumaNode(mAddr, mSize, mAllowRemoteNumaFallback); - // Nothing is reachable until commit() opens it, so a read that runs ahead - // of the fill faults rather than returning zeros from an untouched page. - // Costs one VMA and is independent of the size: with no resident pages - // there are no page tables to rewrite. + // Nothing is reachable until a step opens it, so a read that runs ahead of + // the fill faults rather than returning zeros from an untouched page. Costs + // one VMA and is independent of the size: with no resident pages there are + // no page tables to rewrite. auto const [guardAddr, guardSize] = pagesCovering(mAddr, mSize); if (::mprotect(reinterpret_cast(guardAddr), guardSize, PROT_NONE) != 0) { @@ -208,49 +227,87 @@ MemAddress HostMemBacking::reserve(size_t size) mMapped = 0; throw HostOOMError(std::string("mprotect(PROT_NONE) failed: ") + std::strerror(errno)); } - return mAddr; } -void HostMemBacking::commit(size_t offset, size_t size, bool populated) +HostMemBacking::~HostMemBacking() { - TLLM_CHECK_DEBUG(mAddr != 0 && offset + size <= mSize); - TLLM_CHECK_DEBUG(offset % mUnit == 0 && size % mUnit == 0); + HostMemBacking::destroy(); +} + +void HostMemBacking::setTargetSize(size_t size) +{ + TLLM_CHECK_WITH_INFO(mAddr != 0, "HostMemBacking::setTargetSize on a backing that reserved nothing"); + TLLM_CHECK_WITH_INFO(size <= mSize, "HostMemBacking target %zu exceeds the reservation of %zu", size, mSize); + TLLM_CHECK_WITH_INFO( + size % mUnit == 0, "HostMemBacking target %zu is not a multiple of the granularity %zu", size, mUnit); + TLLM_CHECK_WITH_INFO(mCommitted == mTarget, "HostMemBacking: a new target with %zu of %zu bytes still to commit", + mTarget - mCommitted, mTarget); + mTarget = size; +} +size_t HostMemBacking::resizeStep(bool populated) +{ + if (mCommitted == mTarget) + { + return mRegisteredEnd; + } // cuMemHostRegister needs a current context; the fill worker has only set // the device, which does not give it one. cuCheck(cuCtxSetCurrent(static_cast(mContext))); + if (mCommitted < mTarget) + { + // Rounding up from one past the end lands on the next chunk boundary + // whether or not the previous step ended on one, so a step that resumed + // a partial chunk completes it rather than carrying the offset forward. + commitUpTo(std::min(roundUp(mCommitted + 1, kCommitChunkSize), mTarget), populated); + } + else + { + decommitFrom(std::max(mTarget, (mCommitted - 1) / kCommitChunkSize * kCommitChunkSize)); + } + return mRegisteredEnd; +} + +size_t HostMemBacking::commitChunkSize() noexcept +{ + return kCommitChunkSize; +} + +size_t HostMemBacking::nextStepEnd() const noexcept +{ + return mCommitted < mTarget ? std::min(roundUp(mCommitted + 1, kCommitChunkSize), mTarget) : mCommitted; +} + +void HostMemBacking::resizeTo(size_t size) +{ + setTargetSize(size); + while (mCommitted != mTarget) + { + resizeStep(); + } +} + +void HostMemBacking::commitUpTo(size_t end, bool populated) +{ + // With `populated` the caller has already opened and populated this range; + // the step only has to register it. if (!populated) { - open(offset, size); - populate(offset, size); + open(mCommitted, end - mCommitted); + populate(mCommitted, end - mCommitted); } - size_t committed = 0; + size_t const begin = mCommitted; + mCommitted = end; try { - for (; committed < size; committed += mUnit) - { - MemAddress const unitAddr = mAddr + offset + committed; - CUresult const registered = cuMemHostRegister( - reinterpret_cast(unitAddr), mUnit, CU_MEMHOSTREGISTER_PORTABLE | CU_MEMHOSTREGISTER_DEVICEMAP); - if (registered == CUDA_ERROR_OUT_OF_MEMORY) - { - throw HostOOMError("cuMemHostRegister failed: out of memory"); - } - cuCheck(registered); - mRegistered.emplace(offset + committed, mUnit); - } + registerUpTo(end, /*allowPartialChunk=*/end == mTarget); } catch (...) { - // The caller treats a failed commit as having taken nothing. - decommit(offset, committed); - auto const [closeAddr, closeSize] = pagesWithin(mAddr + offset + committed, size - committed); - if (closeSize != 0 && ::mprotect(reinterpret_cast(closeAddr), closeSize, PROT_NONE) != 0) - { - TLLM_LOG_WARNING("KVCM2 host tier: mprotect(PROT_NONE) failed while unwinding a commit, errno %d", errno); - } + // A failed step must take nothing, so what it populated goes back. + decommitFrom(begin); throw; } } @@ -284,37 +341,118 @@ void HostMemBacking::populate(size_t offset, size_t size) const } } -void HostMemBacking::unregisterRange(size_t offset, size_t size) noexcept +void HostMemBacking::registerUpTo(size_t end, bool allowPartialChunk) { - auto unit = mRegistered.lower_bound(offset); - while (unit != mRegistered.end() && unit->first < offset + size) + size_t const target = end / mRegUnit * mRegUnit; + while (mRegisteredEnd < target) { - CUresult const result = cuMemHostUnregister(reinterpret_cast(mAddr + unit->first)); + // Rounding up from one past the start stops at the next chunk boundary, so + // no registration ever spans one. + size_t const chunkEnd = std::min(roundUp(mRegisteredEnd + 1, mRegChunk), target); + if (!allowPartialChunk && chunkEnd % mRegChunk != 0) + { + return; + } + CUresult const registered = cuMemHostRegister(reinterpret_cast(mAddr + mRegisteredEnd), + chunkEnd - mRegisteredEnd, CU_MEMHOSTREGISTER_PORTABLE | CU_MEMHOSTREGISTER_DEVICEMAP); + if (registered == CUDA_ERROR_OUT_OF_MEMORY) + { + throw HostOOMError("cuMemHostRegister failed: out of memory"); + } + cuCheck(registered); + recordRegistrationStart(mRegisteredEnd); + mRegisteredEnd = chunkEnd; + } +} + +void HostMemBacking::recordRegistrationStart(size_t start) +{ + // A start on a chunk boundary is the one innermostRegistrationStart() derives, + // so remembering it would only duplicate the arithmetic. + if (start % mRegChunk != 0) + { + mUnalignedRegStarts.push(start); + } +} + +size_t HostMemBacking::innermostRegistrationStart() const noexcept +{ + TLLM_CHECK_DEBUG(mRegisteredEnd > 0); + // A registration never reaches past the chunk boundary above where it began, + // so the boundary below its last byte is where it began -- unless it began + // off a boundary, which is the case the stack exists to record. + size_t const onChunkBoundary = (mRegisteredEnd - 1) / mRegChunk * mRegChunk; + return mUnalignedRegStarts.empty() ? onChunkBoundary : std::max(mUnalignedRegStarts.top(), onChunkBoundary); +} + +size_t HostMemBacking::dropInnermostRegistration() noexcept +{ + size_t const start = innermostRegistrationStart(); + unregisterAt(start); + if (!mUnalignedRegStarts.empty() && mUnalignedRegStarts.top() == start) + { + mUnalignedRegStarts.pop(); + } + // Registrations are contiguous, so the one below ends where this began. + mRegisteredEnd = start; + return start; +} + +void HostMemBacking::unregisterAt(size_t offset) noexcept +{ + CUresult const result = cuMemHostUnregister(reinterpret_cast(mAddr + offset)); + if (result != CUDA_SUCCESS) + { + char const* name = nullptr; + cuGetErrorName(result, &name); + // The driver and this object disagree about what is registered, which the + // pages are about to be freed on the assumption of. + Poison::set("HostMemBacking::unregisterAt", name != nullptr ? name : "cuMemHostUnregister failed"); + } +} + +void HostMemBacking::unregisterFrom(size_t end) noexcept +{ + size_t const target = end / mRegUnit * mRegUnit; + while (mRegisteredEnd > target) + { + size_t const start = dropInnermostRegistration(); + if (start >= target) + { + continue; + } + // That one straddled `target`. A registration cannot be shortened, so it + // had to go whole and its surviving prefix is registered again. Leaving + // it would pin pages madvise is about to discard. + CUresult const result = cuMemHostRegister(reinterpret_cast(mAddr + start), target - start, + CU_MEMHOSTREGISTER_PORTABLE | CU_MEMHOSTREGISTER_DEVICEMAP); if (result != CUDA_SUCCESS) { char const* name = nullptr; cuGetErrorName(result, &name); - // The driver and this object disagree about what is registered, - // which the pages below are about to be freed on the assumption of. - Poison::set("HostMemBacking::unregisterRange", name != nullptr ? name : "cuMemHostUnregister failed"); + // The prefix stays committed but unregistered, so a copy over it would + // fail rather than read the wrong pages. Nothing can repair that here. + Poison::set("HostMemBacking::unregisterFrom", name != nullptr ? name : "cuMemHostRegister failed"); + return; } - unit = mRegistered.erase(unit); + recordRegistrationStart(start); + mRegisteredEnd = target; + return; } } -void HostMemBacking::decommit(size_t offset, size_t size) noexcept +void HostMemBacking::decommitFrom(size_t end) noexcept { - if (mAddr == 0 || size == 0) + if (mAddr == 0 || end >= mCommitted) { return; } - TLLM_CHECK_DEBUG(offset % mUnit == 0 && size % mUnit == 0); - // Order matters and nothing enforces it: registration does not hold the // pages, so discarding them first destroys memory the GPU still has // registered, with every call reporting success. - unregisterRange(offset, size); - auto const [dropAddr, dropSize] = pagesWithin(mAddr + offset, size); + unregisterFrom(end); + auto const [dropAddr, dropSize] = pagesWithin(mAddr + end, mCommitted - end); + mCommitted = end; if (dropSize == 0) { return; @@ -333,6 +471,31 @@ void HostMemBacking::decommit(size_t offset, size_t size) noexcept } } +void HostMemBacking::discardUncommitted(size_t offset, size_t size) noexcept +{ + if (mAddr == 0 || size == 0) + { + return; + } + // Never registered -- registration happens in the commit -- so unlike + // decommitFrom() this only has pages to drop. + TLLM_CHECK_DEBUG(offset >= mCommitted); + auto const [dropAddr, dropSize] = pagesWithin(mAddr + offset, size); + if (dropSize == 0) + { + return; + } + if (::madvise(reinterpret_cast(dropAddr), dropSize, MADV_DONTNEED) != 0) + { + TLLM_LOG_WARNING("KVCM2 host tier: madvise(MADV_DONTNEED) failed with errno %d", errno); + } + // Back to unreachable, undoing the open() that let it be populated. + if (::mprotect(reinterpret_cast(dropAddr), dropSize, PROT_NONE) != 0) + { + TLLM_LOG_WARNING("KVCM2 host tier: mprotect(PROT_NONE) failed with errno %d", errno); + } +} + void HostMemBacking::destroy() noexcept { if (mAddr == 0) @@ -346,13 +509,16 @@ void HostMemBacking::destroy() noexcept KVCM2_POISON_ON_EXCEPT( [&] { - unregisterRange(0, mSize); + unregisterFrom(0); hostMunmap(mAddr, mMapped); }); - mRegistered.clear(); + mUnalignedRegStarts = {}; + mRegisteredEnd = 0; mAddr = 0; mSize = 0; mMapped = 0; + mTarget = 0; + mCommitted = 0; } } // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/hostMemBacking.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/hostMemBacking.h index da8cbc9fb0d5..5b06fa9c3845 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/hostMemBacking.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/hostMemBacking.h @@ -21,29 +21,30 @@ #include #include -#include #include +#include namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 { -//! Smallest unit the backing allocates in, and the alignment every size it is -//! given must satisfy. -inline constexpr size_t kHostMemAlignment = 2ULL << 20; +//! What every size a caller asks of the host tier must be a multiple of. +//! +//! A granularity, not an address alignment: the reservation's base comes from mmap and is +//! only page-aligned, so this constrains sizes and offsets within it rather than addresses. +//! +//! One value for the process, settled from the page size on first use. +[[nodiscard]] size_t hostMemSizeGranularity() noexcept; struct HostMemBackingOptions { - //! Unit memory is registered and released in. Also the alignment every - //! commit and decommit bound must satisfy, and therefore how finely a resize - //! can give memory back. + //! What every registration bound is a multiple of, so no registration ever + //! splits a range of this size. The registered prefix therefore trails the + //! committed one by up to this much. //! - //! Each unit is a separate registration, so it must divide whatever ranges - //! are copied in one call: the driver rejects a copy that straddles two - //! registrations. It need not be a multiple of the page size -- registration - //! accepts any alignment, and the madvise and mprotect calls round to pages - //! themselves, outwards when they must cover the range and inwards when they - //! must not reach past it. - size_t commitUnit = kHostMemAlignment; + //! A caller sets it to the largest range it hands to one cudaMemcpy, which + //! the driver rejects if it straddles two registrations. 1 imposes no + //! constraint, which is what a caller that never copies wants. + size_t registerUnit = 1; //! Whether a page may come from a NUMA node other than the one the GPU //! attaches to, when the local node cannot satisfy it. @@ -57,13 +58,13 @@ struct HostMemBackingOptions bool allowRemoteNumaFallback = true; }; -//! Physical backing for the host tier: reserve an address range once, then -//! commit it a unit at a time. +//! Physical backing for the host tier: an address range reserved on +//! construction, whose committed prefix then moves towards a target a chunk at +//! a time. //! -//! The base address returned by reserve() is fixed for the lifetime of the -//! object. Growing and shrinking touch only the delta, so a resize costs -//! O(delta) rather than O(total) and never invalidates a pointer the allocator -//! above has handed out. +//! The base address is fixed for the lifetime of the object. Growing and +//! shrinking touch only the delta, so a resize costs O(delta) rather than +//! O(total) and never invalidates a pointer the allocator above has handed out. //! //! Pages are populated through the kernel's fault path (MADV_POPULATE_WRITE) //! and then handed to the driver with cuMemHostRegister. Allocating and zeroing @@ -72,14 +73,27 @@ struct HostMemBackingOptions //! an 8-GPU node collapsed per-process fill throughput by roughly 5x and blocked //! unrelated CUDA calls on the same process for seconds at a time. //! -//! Registration happens in units rather than over the whole range so that a -//! shrink can unregister exactly what it releases; a registration can only be -//! dropped whole. Each unit stays far below the 2 GiB-per-call pinning limit on -//! Linux 6.11. +//! Registration happens in chunks rather than over the whole range so that a +//! shrink can unregister close to exactly what it releases; a registration can +//! only be dropped whole. Each chunk stays far below the 2 GiB-per-call pinning +//! limit on Linux 6.11. +//! +//! A step stops on a chunk boundary and leaves the remainder to its successor, +//! so the range carries one registration per chunk rather than one per step. class HostMemBacking { public: - explicit HostMemBacking(HostMemBackingOptions const& options); + //! Reserves `vmSize` bytes of address space, page-aligned as mmap returns it + //! and no coarser, and takes no memory for it. No page is resident and none + //! is reachable: the whole range is PROT_NONE until a step opens it, so a + //! read that runs ahead of the fill faults instead of returning zeros. + //! + //! `vmSize` must be a multiple of hostMemSizeGranularity(). Zero reserves + //! nothing, for an owner that has no range to back. + //! + //! Needs a current CUDA context, which is captured here: registration needs + //! one and the fill worker only sets the device, which does not give it one. + explicit HostMemBacking(size_t vmSize, HostMemBackingOptions const& options = {}); ~HostMemBacking(); HostMemBacking(HostMemBacking const&) = delete; @@ -87,67 +101,166 @@ class HostMemBacking HostMemBacking(HostMemBacking&&) = delete; HostMemBacking& operator=(HostMemBacking&&) = delete; - //! Reserves `size` bytes of address space, page-aligned as mmap returns it - //! and no coarser. No page is resident and none is reachable: the whole - //! range is PROT_NONE until commit() opens it, so a read that runs ahead of - //! the fill faults instead of returning zeros. - //! - //! Must be called exactly once, before any commit. - [[nodiscard]] MemAddress reserve(size_t size); + //! Base of the reservation, fixed for the lifetime of the object. Zero once + //! destroy() has run, and for a backing reserving nothing. + [[nodiscard]] MemAddress address() const noexcept + { + return mAddr; + } - //! Makes [offset, offset + size) resident and reachable by both the CPU and - //! the GPU. Both bounds must be multiples of commitGranularity(). - //! - //! Throws HostOOMError when the system cannot back the range; that is the - //! expected failure and callers are required to handle it. - //! - //! With `populated` true the caller has already run open() and populate() - //! on the range, and commit() only registers it. - void commit(size_t offset, size_t size, bool populated = false); + //! Bytes of address space held, i.e. the largest target this accepts. + [[nodiscard]] size_t reservedSize() const noexcept + { + return mSize; + } - //! The first two steps of commit(): open() makes the range reachable and - //! populate() allocates its pages. Neither unwinds on failure; decommit() + //! The first two steps a commit takes: open() makes the range reachable and + //! populate() allocates its pages. Neither unwinds on failure; decommitFrom() //! releases what they took. //! //! open() changes the mapping, which waits for every populate() in flight //! on any thread: open a range once, then populate it in parallel. Both //! are safe to call while another thread commits a disjoint range. + //! + //! Together they let a caller populate a step ahead of resizeStep(), which it + //! then passes `populated` so the step only registers what is already there. void open(size_t offset, size_t size); void populate(size_t offset, size_t size) const; - //! Releases the physical backing of [offset, offset + size), which stays - //! reserved and may be committed again. Both bounds must be multiples of - //! commitGranularity(). - void decommit(size_t offset, size_t size) noexcept; + //! Undoes open() and populate() over a range that was never committed, which + //! is what a fill that populated ahead of its commits leaves behind when it + //! fails. Only valid at or above committedSize(); committed bytes are released + //! by lowering the target instead. + void discardUncommitted(size_t offset, size_t size) noexcept; + + //! Declares how much of the reservation should end up committed. Must be a + //! multiple of commitGranularity() and no larger than the reservation. + //! + //! Takes no memory by itself; resizeStep() does the work. The previous target + //! must have been reached: the caller finishes or abandons an in-flight fill first. + void setTargetSize(size_t size); + + //! Moves committedSize() one chunk towards the target and returns + //! registeredSize(). + //! + //! A no-op once committedSize() == targetSize(), which is how a loop over it + //! terminates. + //! + //! Throws HostOOMError when the system cannot back the chunk; that is the + //! expected failure and callers are required to handle it. The step then + //! takes nothing: committedSize() is left where it was found. + //! + //! Pass `populated` when the range up to nextStepEnd() has already been + //! opened and populated, which lets the step skip straight to registering it. + size_t resizeStep(bool populated = false); + + //! Where the next growing step will stop, so a caller can populate that range + //! before asking for it. Equals committedSize() when there is no step to take. + [[nodiscard]] size_t nextStepEnd() const noexcept; + + //! Bytes a growing step advances by. Steps land on multiples of it, so a caller + //! populating ahead can derive every step boundary without taking one. + [[nodiscard]] static size_t commitChunkSize() noexcept; + + //! Sets the target and steps until it is reached, for a caller with nothing + //! to do in between. + void resizeTo(size_t size); //! Releases everything, including the reservation. Idempotent. void destroy() noexcept; - //! The unit memory is populated, registered and released in. + //! The granularity of every size this object is asked for. [[nodiscard]] size_t commitGranularity() const noexcept { return mUnit; } + [[nodiscard]] size_t targetSize() const noexcept + { + return mTarget; + } + + //! Bytes resident and reachable, from the base. Moves one chunk per step. + [[nodiscard]] size_t committedSize() const noexcept + { + return mCommitted; + } + + //! Bytes the GPU may be copied to or from, from the base. Never exceeds + //! committedSize(), and trails it by less than one register unit once the + //! target is reached. + [[nodiscard]] size_t registeredSize() const noexcept + { + return mRegisteredEnd; + } + + //! What registration bounds are multiples of, as the options asked, or 1 when + //! they imposed no constraint. + [[nodiscard]] size_t registerUnit() const noexcept + { + return mRegUnit; + } + private: - void unregisterRange(size_t offset, size_t size) noexcept; + //! Populate up to `end`, which must be at or above mCommitted. + void commitUpTo(size_t end, bool populated = false); + + //! Release from `end`, which must be at or below mCommitted. + void decommitFrom(size_t end) noexcept; + + //! Extend registration towards the last whole register unit at or below `end`. + //! Unless `allowPartialChunk`, stops at the last whole chunk boundary so the + //! remainder joins the next chunk rather than becoming a registration of its + //! own. + void registerUpTo(size_t end, bool allowPartialChunk); + + //! Drop registration above the last whole register unit at or below `end`. + void unregisterFrom(size_t end) noexcept; + + //! Remember `start` if it is not where arithmetic would look for it. + void recordRegistrationStart(size_t start); + + //! Where the registration ending at mRegisteredEnd began. Requires one to exist. + [[nodiscard]] size_t innermostRegistrationStart() const noexcept; + + //! Unregister that one and pull mRegisteredEnd back to its start, which is + //! returned. + size_t dropInnermostRegistration() noexcept; + + //! cuMemHostUnregister takes only a base address; the size is the driver's. + void unregisterAt(size_t offset) noexcept; size_t mUnit = 0; + size_t mRegUnit = 1; + //! Bytes per registration: the most whole register units that fit in the target + //! chunk size, or one register unit when that is the larger. + size_t mRegChunk = 0; bool mAllowRemoteNumaFallback = true; MemAddress mAddr = 0; + size_t mSize = 0; - //! Bytes actually mapped: mSize rounded up to a page, since a unit is a - //! whole number of slots and need not end on one. + size_t mTarget = 0; + size_t mCommitted = 0; + + //! Bytes actually mapped: mSize rounded up to a page. size_t mMapped = 0; - //! Offsets of the units currently registered. A unit is registered exactly - //! when it is committed, so this is also what destroy() must undo. - std::map mRegistered; + //! Innermost last. Registrations are contiguous and grow and shrink only at + //! the end, so a stack is the whole structure. + //! + //! Only the unaligned ones, because the rest are arithmetic: a step registers + //! whole chunks from a boundary, so the start of the registration covering a + //! byte is normally the chunk boundary below it. One case breaks that -- a + //! fill resumed at a higher target begins where the previous target's partial + //! tail ended -- and those are what this holds, one per resumed fill rather + //! than one per chunk of the tier. + std::stack mUnalignedRegStarts; + + //! End of the registered prefix, always a whole number of register units. + size_t mRegisteredEnd = 0; - //! Captured in reserve(), which runs on the thread that owns the context. - //! cuMemHostRegister needs a current context and the fill worker only sets - //! the device, which leaves it without one. + //! The context cuMemHostRegister runs under; the fill worker only sets the device. void* mContext = nullptr; }; diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/math.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/math.h index c946a6810607..c36b96e0aad4 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/math.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/math.h @@ -25,6 +25,7 @@ #include #include #include +#include #include #include #include @@ -80,13 +81,17 @@ template // Extract an attribute from the first element and assert (debug) all elements agree. // Mirrors Python's get_uniform_attribute(_utils.py:202). +// +// Takes the range by forwarding reference so an adapted view may be passed: a +// filter_view's begin() is not const-qualified, so a const reference to one cannot +// be iterated. template -[[nodiscard]] auto getUniformAttribute(Range const& range, Func&& func) -> decltype(func(*range.begin())) +[[nodiscard]] auto getUniformAttribute(Range&& range, Func&& func) -> decltype(func(*std::ranges::begin(range))) { - auto it = range.begin(); - TLLM_CHECK_DEBUG(it != range.end()); + auto it = std::ranges::begin(range); + TLLM_CHECK_DEBUG(it != std::ranges::end(range)); auto result = func(*it); - TLLM_CHECK_DEBUG(std::all_of(range.begin(), range.end(), [&](auto const& item) { return func(item) == result; })); + TLLM_CHECK_DEBUG(std::ranges::all_of(range, [&](auto const& item) { return func(item) == result; })); return result; } diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/sharedPtr.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/sharedPtr.h index 37b1d9c5be61..9db79737e2f0 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/sharedPtr.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/sharedPtr.h @@ -560,7 +560,7 @@ SharedPtr makeShared(Args&&... args) // Wire up EnableSharedFromThis if T (or a base) derives from it. if constexpr (detail::hasEnable) { - result.mPtr->detail::template EnableBase::mWeakThis = result; + static_cast&>(*result.mPtr).mWeakThis = result; } return result; diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/staticHostMem.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/staticHostMem.cpp index b619f5a257cf..7d40ca980417 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/staticHostMem.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/staticHostMem.cpp @@ -42,7 +42,10 @@ void StaticHostMem::Deleter::operator()(void* ptr) const noexcept StaticHostMem::StaticHostMem(size_t size) { - TLLM_CHECK_WITH_INFO(size > 0, "StaticHostMem size must be non-zero"); + if (size == 0) + { + return; + } void* ptr = nullptr; cuCheck(cuMemHostAlloc(&ptr, size, CU_MEMHOSTALLOC_PORTABLE | CU_MEMHOSTALLOC_DEVICEMAP)); mPtr.reset(ptr); diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/staticHostMem.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/staticHostMem.h index f2a72c7814cd..85d2fc6876cd 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/staticHostMem.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/staticHostMem.h @@ -40,10 +40,20 @@ namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 class StaticHostMem { public: + //! Holds nothing. Move-assign an allocated one over it to fill it. + StaticHostMem() = default; + + //! A size of zero holds nothing, as the default constructor does. + //! //! Throws CuError when the allocation fails, which is also how an oversized //! request is reported -- there is no internal chunking to fall back on. explicit StaticHostMem(size_t size); + [[nodiscard]] bool empty() const noexcept + { + return mSize == 0; + } + [[nodiscard]] MemAddress address() const noexcept { return reinterpret_cast(mPtr.get()); diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/typedIndex.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/typedIndex.h index bc50645d4941..0f933dc3b31f 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/typedIndex.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/typedIndex.h @@ -18,6 +18,7 @@ #pragma once #include "tensorrt_llm/common/assert.h" +#include #include #include #include @@ -167,6 +168,45 @@ constexpr T operator-(StrongIndex lhs, StrongIndex(lhs.value() - rhs.value()); } +namespace detail +{ +//! A StrongIndex is its value and nothing else: one member, no padding, no vtable. Code that +//! stages indices in a buffer a kernel reads, or hands them to a copy that moves bytes, rests +//! on this; checked here so none of it has to say so again. +//! +//! Spelled in the native types rather than the fixed-width aliases: int64_t names one of +//! long and long long, so those would leave the other unchecked. Every type is_integral +//! admits is here, which is wider than the arithmetic ones -- char is a third type beside +//! its signed and unsigned spellings, and bool and the character types are integral too. +template +using IndexLayoutProbe = StrongIndex; + +#define KVCM2_CHECK_INDEX_LAYOUT(T) \ + static_assert(sizeof(IndexLayoutProbe) == sizeof(T)); \ + static_assert(alignof(IndexLayoutProbe) == alignof(T)); \ + static_assert(std::is_trivially_copyable_v>); \ + static_assert(std::is_standard_layout_v>) + +KVCM2_CHECK_INDEX_LAYOUT(bool); +KVCM2_CHECK_INDEX_LAYOUT(char); +KVCM2_CHECK_INDEX_LAYOUT(signed char); +KVCM2_CHECK_INDEX_LAYOUT(short); +KVCM2_CHECK_INDEX_LAYOUT(int); +KVCM2_CHECK_INDEX_LAYOUT(long); +KVCM2_CHECK_INDEX_LAYOUT(long long); +KVCM2_CHECK_INDEX_LAYOUT(unsigned char); +KVCM2_CHECK_INDEX_LAYOUT(unsigned short); +KVCM2_CHECK_INDEX_LAYOUT(unsigned int); +KVCM2_CHECK_INDEX_LAYOUT(unsigned long); +KVCM2_CHECK_INDEX_LAYOUT(unsigned long long); +KVCM2_CHECK_INDEX_LAYOUT(wchar_t); +KVCM2_CHECK_INDEX_LAYOUT(char8_t); +KVCM2_CHECK_INDEX_LAYOUT(char16_t); +KVCM2_CHECK_INDEX_LAYOUT(char32_t); + +#undef KVCM2_CHECK_INDEX_LAYOUT +} // namespace detail + template [[nodiscard]] std::size_t toSizeT(StrongIndex index) noexcept { @@ -186,6 +226,11 @@ class TypedVec using ContainerType = std::vector; using iterator = typename ContainerType::iterator; using const_iterator = typename ContainerType::const_iterator; + //! The container's own reference types rather than `T&`/`T const&`, so that + //! TypedVec works: std::vector hands out a proxy, not a `bool&`. + //! For every other T these are exactly `T&` and `T const&`. + using reference = typename ContainerType::reference; + using const_reference = typename ContainerType::const_reference; TypedVec() = default; @@ -209,22 +254,22 @@ class TypedVec { } - [[nodiscard]] T& operator[](Index index) noexcept + [[nodiscard]] reference operator[](Index index) noexcept { return mData[toSizeT(index)]; } - [[nodiscard]] T const& operator[](Index index) const noexcept + [[nodiscard]] const_reference operator[](Index index) const noexcept { return mData[toSizeT(index)]; } - [[nodiscard]] T& at(Index index) + [[nodiscard]] reference at(Index index) { return mData.at(toSizeT(index)); } - [[nodiscard]] T const& at(Index index) const + [[nodiscard]] const_reference at(Index index) const { return mData.at(toSizeT(index)); } @@ -275,7 +320,7 @@ class TypedVec } template - T& emplace_back(Args&&... args) + reference emplace_back(Args&&... args) { return mData.emplace_back(std::forward(args)...); } @@ -285,22 +330,22 @@ class TypedVec mData.pop_back(); } - [[nodiscard]] T& front() + [[nodiscard]] reference front() { return mData.front(); } - [[nodiscard]] T const& front() const + [[nodiscard]] const_reference front() const { return mData.front(); } - [[nodiscard]] T& back() + [[nodiscard]] reference back() { return mData.back(); } - [[nodiscard]] T const& back() const + [[nodiscard]] const_reference back() const { return mData.back(); } @@ -359,6 +404,92 @@ class TypedVec ContainerType mData; }; +//! Index of an EnumMap: a StrongIndex as elsewhere, or a plain enum whose enumerators run +//! from zero. Both name a position; neither is an integer a caller can pass by accident. +template +[[nodiscard]] constexpr std::size_t toMapSlot(Index index) noexcept +{ + if constexpr (std::is_enum_v) + { + return static_cast(index); + } + else + { + return toSizeT(index); + } +} + +// --------------------------------------------------------------------------- +// EnumMap -- one value per enumerator, in an array rather than an allocation. +// +// For a set whose size is known when the type is written: the slots are fixed, so a lookup +// is a load at a constant displacement where the value it replaces was arithmetic over +// members the compiler cannot fold. TypedVec is the one to reach for when the count is a +// runtime answer. +// --------------------------------------------------------------------------- +template +class EnumMap +{ +public: + using IndexType = Index; + using ValueType = T; + + constexpr EnumMap() = default; + + constexpr explicit EnumMap(std::array const& values) + : mData(values) + { + } + + [[nodiscard]] constexpr T& operator[](Index index) noexcept + { + return mData[toMapSlot(index)]; + } + + [[nodiscard]] constexpr T const& operator[](Index index) const noexcept + { + return mData[toMapSlot(index)]; + } + + [[nodiscard]] constexpr T& at(Index index) + { + return mData.at(toMapSlot(index)); + } + + [[nodiscard]] constexpr T const& at(Index index) const + { + return mData.at(toMapSlot(index)); + } + + [[nodiscard]] static constexpr std::size_t size() noexcept + { + return Count; + } + + [[nodiscard]] constexpr typename std::array::iterator begin() noexcept + { + return mData.begin(); + } + + [[nodiscard]] constexpr typename std::array::iterator end() noexcept + { + return mData.end(); + } + + [[nodiscard]] constexpr typename std::array::const_iterator begin() const noexcept + { + return mData.begin(); + } + + [[nodiscard]] constexpr typename std::array::const_iterator end() const noexcept + { + return mData.end(); + } + +private: + std::array mData{}; +}; + } // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 namespace std diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp index c503e0250648..27f355afafbd 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp @@ -69,14 +69,6 @@ namespace tensorrt_llm::nanobind::batch_manager namespace { -//! A DLPack view keeps the batch's device allocation alive without importing torch. -struct BatchDeviceArray -{ - std::shared_ptr batch; - kv::LayerGroupId group; - bool counts; -}; - // Exposed via introspection sub-module for tests. class TestPaddingColdPageCodec final : public kv::IKvCacheColdPageCodec { @@ -315,6 +307,43 @@ static kv::TypedVec> typ return result; } +//! Pool selectors, defaulting to every pool of each group when not given -- which is what +//! a single-budget caller means, and keeps the simple form of these helpers simple. +static kv::TypedVec> typedPoolSelectors( + std::optional>> const& pools, + kv::TypedVec> const& sizes) +{ + kv::TypedVec> result; + if (!pools.has_value()) + { + for (auto const& group : sizes) + { + std::vector all; + for (kv::PoolIndex p{0}; p < group.size(); ++p) + { + all.push_back(p); + } + result.push_back(std::move(all)); + } + return result; + } + for (auto const& group : *pools) + { + std::vector picked; + picked.reserve(group.size()); + for (int32_t p : group) + { + if (p < 0) + { + throw std::invalid_argument("pool indices must be non-negative"); + } + picked.push_back(kv::PoolIndex{p}); + } + result.push_back(std::move(picked)); + } + return result; +} + static kv::TypedVec typedSlotCounts(std::vector const& slotCounts) { kv::TypedVec result; @@ -353,6 +382,50 @@ static kv::KvCache::PriorityCb castPriorityCallback(kv::KvCacheManager const& ma }; } +//! Wrap device memory the caller does not own as a DLPack array, so torch.from_dlpack and +//! friends can consume it. `owner` is kept alive for the array's lifetime, which is what +//! stops the underlying allocation vanishing beneath it. +using DeviceArray1D = nb::ndarray, nb::c_contig, nb::device::cuda>; +using DeviceArray2D = nb::ndarray, nb::c_contig, nb::device::cuda>; +using DeviceArray3D = nb::ndarray, nb::c_contig, nb::device::cuda>; + +static int currentCudaDevice() +{ + CUdevice device = 0; + kv::cuCheck(cuCtxGetDevice(&device)); + return static_cast(device); +} + +static DeviceArray1D makeDeviceArray1D(CUdeviceptr data, int32_t count, nb::handle owner) +{ + return DeviceArray1D{reinterpret_cast(data), {static_cast(count)}, owner, {}, nb::dtype(), + nb::device::cuda::value, currentCudaDevice()}; +} + +static DeviceArray2D makeDeviceArray2D(CUdeviceptr data, int32_t rows, int32_t cols, nb::handle owner) +{ + return DeviceArray2D{reinterpret_cast(data), {static_cast(rows), static_cast(cols)}, owner, + {}, nb::dtype(), nb::device::cuda::value, currentCudaDevice()}; +} + +static DeviceArray3D makeDeviceArray3D(CUdeviceptr data, int32_t d0, int32_t d1, int32_t d2, nb::handle owner) +{ + return DeviceArray3D{reinterpret_cast(data), + {static_cast(d0), static_cast(d1), static_cast(d2)}, owner, {}, nb::dtype(), + nb::device::cuda::value, currentCudaDevice()}; +} + +//! The host counterpart, for staging a consumer reads rather than copies. +//! Read-only: the staging block behind it is the members' to write, through their page +//! locks, and a publish may be reading it. +using HostArray2D = nb::ndarray, nb::c_contig, nb::device::cpu>; + +static HostArray2D makeHostArray2D(int32_t const* data, int32_t rows, int32_t cols, nb::handle owner) +{ + return HostArray2D{data, {static_cast(rows), static_cast(cols)}, owner, {}, nb::dtype(), + nb::device::cpu::value, 0}; +} + static nb::tuple bufferIdTuple(kv::BufferId const& self) { return nb::make_tuple(self.layerId, self.role); @@ -1396,28 +1469,20 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) .value("SHARED", kv::PageIndexMode::SHARED) .value("PER_LAYER", kv::PageIndexMode::PER_LAYER); - nb::class_(m, "ScratchDesc") - .def_ro("range", &kv::ScratchDesc::range) - .def_ro("slot_ids", &kv::ScratchDesc::slotIds) - .def("__bool__", [](kv::ScratchDesc const& self) { return static_cast(self); }); - nb::enum_(m, "AttentionReusePolicy") .value("REQUIRED", kv::AttentionReusePolicy::REQUIRED) .value("PRIVATE", kv::AttentionReusePolicy::PRIVATE) .value("OPTIONAL", kv::AttentionReusePolicy::OPTIONAL); nb::class_(m, "AttnLifeCycle") - .def(nb::init, int, bool, kv::AttentionReusePolicy>(), nb::arg("window_size").none(), - nb::arg("num_sink_blocks"), nb::arg("is_sparse") = false, - nb::arg("reuse_policy") = kv::AttentionReusePolicy::REQUIRED) + .def(nb::init, int, kv::AttentionReusePolicy>(), nb::arg("window_size").none(), + nb::arg("num_sink_blocks"), nb::arg("reuse_policy") = kv::AttentionReusePolicy::REQUIRED) // Sink tokens round up to whole blocks. Bound rather than repeated in Python so the // connector's view of a life cycle is built by the same code as the allocator's. .def_static("make", &kv::AttnLifeCycle::make, nb::arg("window_size").none(), nb::arg("num_sink_tokens").none(), - nb::arg("tokens_per_block"), nb::arg("is_sparse") = false, - nb::arg("reuse_policy") = kv::AttentionReusePolicy::REQUIRED) + nb::arg("tokens_per_block"), nb::arg("reuse_policy") = kv::AttentionReusePolicy::REQUIRED) .def_prop_ro("window_size", [](kv::AttnLifeCycle const& self) { return self.windowSize; }) .def_ro("num_sink_blocks", &kv::AttnLifeCycle::numSinkBlocks) - .def_ro("is_sparse", &kv::AttnLifeCycle::isSparse) .def_ro("reuse_policy", &kv::AttnLifeCycle::reusePolicy) .def("get_stale_range", &kv::AttnLifeCycle::getStaleRange, nb::arg("history_length"), nb::arg("tokens_per_block")) @@ -1563,7 +1628,12 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) nb::class_(m, "PoolDesc") .def_prop_ro("pool_index", [](kv::PoolDesc const& self) { return self.poolIndex.value(); }) .def_ro("base_address", &kv::PoolDesc::baseAddress) - .def_ro("slot_bytes", &kv::PoolDesc::slotBytes); + .def_ro("slot_bytes", &kv::PoolDesc::slotBytes) + .def_ro("storage_medium", &kv::PoolDesc::storageMedium, + "What this pool's pages actually are. The same as the group's tier except for a sparse " + "pool of the GPU tier, which is host-backed while the group still reports GPU_MEM. Ask " + "per pool, not per group, when the medium of an address matters -- registering it for " + "transfer, or picking a copy implementation."); nb::class_(m, "PoolGroupDesc") .def_prop_ro("pool_group_index", [](kv::PoolGroupDesc const& self) { return self.poolGroupIndex.value(); }) @@ -1615,21 +1685,36 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) .def_prop_ro("tier", &kv::DiskCacheTierConfig::tier) .def("validate", &kv::DiskCacheTierConfig::validate) DEF_COPY(kv::DiskCacheTierConfig); + nb::class_(m, "SparseCacheConfig") + .def(nb::init>(), nb::arg("num_selected_blocks"), nb::arg("max_cached_units"), + nb::arg("unit_tokens"), nb::arg("selection_id").none() = nb::none()) + .def_rw("num_selected_blocks", &kv::SparseCacheConfig::numSelectedBlocks) + .def_rw("max_cached_units", &kv::SparseCacheConfig::maxCachedUnits) + .def_rw("unit_tokens", &kv::SparseCacheConfig::unitTokens) + .def_rw("selection_id", &kv::SparseCacheConfig::selectionId, + "Layers whose selection is one selection share this id, and with it one resident " + "set and one resolve. None means the layer selects for itself: it is read as the " + "layer's own id, so a layer that names nothing never shares, and no explicit id " + "may name a layer that named nothing.") DEF_COPY(kv::SparseCacheConfig); + nb::class_(m, "BufferConfig") - .def(nb::init, bool, int>(), nb::arg("role"), nb::arg("size"), + .def(nb::init, bool, int, int>(), nb::arg("role"), nb::arg("size"), nb::arg("tokens_per_block_override") = std::nullopt, nb::arg("is_sparse") = false, - nb::arg("coalescing_group") = 0) + nb::arg("coalescing_group") = 0, nb::arg("sparse_unit_chunks") = 1) .def_rw("role", &kv::BufferConfig::role) .def_rw("size", &kv::BufferConfig::size) .def_rw("tokens_per_block_override", &kv::BufferConfig::tokensPerBlockOverride) .def_rw("is_sparse", &kv::BufferConfig::isSparse) - .def_rw("coalescing_group", &kv::BufferConfig::coalescingGroup) DEF_COPY(kv::BufferConfig); + .def_rw("coalescing_group", &kv::BufferConfig::coalescingGroup) + .def_rw("sparse_unit_chunks", &kv::BufferConfig::sparseUnitChunks) DEF_COPY(kv::BufferConfig); nb::class_(m, "AttentionLayerConfig") .def(nb::init, std::optional, std::optional, - kv::AttentionReusePolicy>(), + kv::AttentionReusePolicy, std::optional>(), nb::arg("layer_id"), nb::arg("buffers"), nb::arg("sliding_window_size") = std::nullopt, - nb::arg("num_sink_tokens") = std::nullopt, nb::arg("reuse_policy") = kv::AttentionReusePolicy::REQUIRED) + nb::arg("num_sink_tokens") = std::nullopt, nb::arg("reuse_policy") = kv::AttentionReusePolicy::REQUIRED, + nb::arg("sparse_cache") = std::nullopt) + .def_rw("sparse_cache", &kv::AttentionLayerConfig::sparseCache) .def_rw("layer_id", &kv::AttentionLayerConfig::layerId) .def_rw("buffers", &kv::AttentionLayerConfig::buffers) .def_rw("sliding_window_size", &kv::AttentionLayerConfig::slidingWindowSize) @@ -1753,7 +1838,7 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) std::optional> initialPoolRatio, std::optional swaScratchReuse, bool commitMinSnapshot, bool enableStats, bool textOnly, bool enablePartialCommit, kv::PoolRebalanceConfig poolRebalance, - kv::Priority coldTierOffloadMinPriority) + kv::Priority coldTierOffloadMinPriority, std::optional hostSparsePoolQuota) { new (cfg) kv::KVCacheManagerConfig(); cfg->tokensPerBlock = tokensPerBlock; @@ -1779,6 +1864,7 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) cfg->enablePartialCommit = enablePartialCommit; cfg->poolRebalance = std::move(poolRebalance); cfg->coldTierOffloadMinPriority = coldTierOffloadMinPriority; + cfg->hostSparsePoolQuota = hostSparsePoolQuota; // Mirror Python's __post_init__: validate at construction. Config-integrity // failures raise AssertionError (translated below). cfg->validate(); @@ -1790,7 +1876,8 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) nb::arg("swa_scratch_reuse").none() = std::nullopt, nb::arg("commit_min_snapshot") = false, nb::arg("enable_stats") = true, nb::arg("text_only") = false, nb::arg("enable_partial_commit") = true, nb::arg("pool_rebalance") = kv::PoolRebalanceConfig{}, - nb::arg("cold_tier_offload_min_priority") = tensorrt_llm::batch_manager::kDefaultColdTierOffloadMinPriority) + nb::arg("cold_tier_offload_min_priority") = tensorrt_llm::batch_manager::kDefaultColdTierOffloadMinPriority, + nb::arg("host_sparse_pool_quota").none() = std::nullopt) .def_rw("tokens_per_block", &kv::KVCacheManagerConfig::tokensPerBlock) .def_rw("cache_tiers", &kv::KVCacheManagerConfig::cacheTiers) .def_rw("layers", &kv::KVCacheManagerConfig::layers) @@ -1801,6 +1888,17 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) .def_rw("cold_tier_offload_min_priority", &kv::KVCacheManagerConfig::coldTierOffloadMinPriority) .def_rw("typical_step", &kv::KVCacheManagerConfig::typicalStep) .def_rw("constraints", &kv::KVCacheManagerConfig::constraints) + .def_rw("host_sparse_pool_quota", &kv::KVCacheManagerConfig::hostSparsePoolQuota, + "Host bytes the GPU tier's host-backed pools may take. None has the storage manager set it " + "at construction to 80% of the host memory the host tiers' quotas leave, and keep that value.\n\n" + "A sparse buffer's pages are host memory, so the device quota does not pay for them and " + "nothing else bounds them: a slot spans every pool of its group, so their row count " + "follows from the device-backed pools the quota does size. This is the bound they " + "otherwise lack. It only lowers a slot count, and pool groups holding such pools share " + "it by the same ratio that divides the device quota.\n\n" + "Kept exactly, as is the feasibility floor `constraints` implies; the ratio is what gives " + "way between them. The two can only conflict if the floors alone exceed this budget, " + "which construction rejects.") .def_rw("initial_pool_ratio", &kv::KVCacheManagerConfig::initialPoolRatio, "One positive, normalized cache-tier quota weight per layer group.") .def_rw("swa_scratch_reuse", &kv::KVCacheManagerConfig::swaScratchReuse) @@ -1817,111 +1915,6 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) nb::class_(m, "PlannedDropHandle") .def("drop", &kv::PlannedDropHandle::drop, nb::call_guard()); - nb::class_(m, "PageStorageSnapshot") - .def_prop_ro("version", &kv::PageStorageSnapshot::version) - .def_prop_ro("row", &kv::PageStorageSnapshot::row) - .def_prop_ro("base_page_indices", &kv::PageStorageSnapshot::basePageIndices) - .def_prop_ro("cache_levels", - [](kv::PageStorageSnapshot const& self) - { - std::vector> levels; - levels.reserve(self.cacheLevels().size()); - for (auto const& level : self.cacheLevels()) - levels.push_back(level ? std::optional{level->value()} : std::nullopt); - return levels; - }) - .def_prop_ro("eligible_history_blocks", &kv::PageStorageSnapshot::eligibleHistoryBlocks) - .def("wait_ready", &kv::PageStorageSnapshot::waitReady, nb::arg("cuda_stream"), - nb::call_guard()); - - nb::class_(m, "BatchDeviceArray") - .def("__dlpack_device__", - [](BatchDeviceArray const& self) - { return std::make_pair(nb::device::cuda::value, self.batch->deviceId()); }) - .def( - "__dlpack__", - [](BatchDeviceArray const& self, std::optional stream, nb::kwargs kwargs) - { - if (kwargs.contains("copy") && !kwargs["copy"].is_none() && nb::cast(kwargs["copy"])) - { - throw std::invalid_argument("Batch arrays support only zero-copy DLPack export"); - } - if (kwargs.contains("dl_device") && !kwargs["dl_device"].is_none() - && nb::cast>(kwargs["dl_device"]) - != std::make_pair(nb::device::cuda::value, self.batch->deviceId())) - { - throw std::invalid_argument("Batch arrays cannot be exported to a different device"); - } - // DLPack 1/2 denote the legacy/per-thread default CUDA streams. - auto cudaStream = stream.value_or(1); - if (cudaStream == 0 || cudaStream < -1) - { - throw std::invalid_argument("Invalid DLPack CUDA stream"); - } - if (cudaStream != -1) - { - auto const consumerStream = cudaStream == 1 ? CU_STREAM_LEGACY - : cudaStream == 2 ? CU_STREAM_PER_THREAD - : reinterpret_cast(cudaStream); - nb::gil_scoped_release release; - self.batch->waitReady(reinterpret_cast(consumerStream)); - } - else - { - nb::gil_scoped_release release; - if (!self.batch->dirtyRows().empty()) - { - throw kv::LogicError("Publish Batch metadata before exporting it"); - } - } - std::vector shape{ - static_cast(self.batch->maxRows()), static_cast(self.batch->maxBeamWidth())}; - auto address - = self.counts ? self.batch->numBlocksAddress(self.group) : self.batch->pageTableAddress(self.group); - if (!self.counts) - { - shape.push_back(self.batch->maxBlocks()); - } - return nb::ndarray(reinterpret_cast(address), shape.size(), - shape.data(), nb::cast(self.batch), nullptr, nb::dtype(), nb::device::cuda::value, - self.batch->deviceId()); - }, - nb::arg("stream").none() = nb::none(), nb::arg("kwargs")); - - nb::class_(m, "Batch") - .def(nb::init, int, int, int>(), nb::arg("manager"), nb::arg("max_rows"), - nb::arg("max_blocks"), nb::arg("max_beam_width") = 1, nb::call_guard()) - .def("add", &kv::Batch::add, nb::arg("kv_cache"), nb::arg("row").none() = std::nullopt, - nb::call_guard()) - .def("remove", &kv::Batch::remove, nb::arg("kv_cache"), nb::call_guard()) - .def("close", &kv::Batch::close, nb::call_guard()) - .def("publish", &kv::Batch::publish, nb::arg("cuda_stream"), nb::call_guard()) - .def("wait_ready", &kv::Batch::waitReady, nb::arg("cuda_stream"), nb::call_guard()) - .def("record_read", &kv::Batch::recordRead, nb::arg("cuda_stream"), nb::call_guard()) - .def("resize", &kv::Batch::resize, nb::arg("capacities"), nb::arg("history_lengths"), nb::arg("cuda_stream"), - nb::call_guard()) - .def_prop_ro("dirty_rows", &kv::Batch::dirtyRows, nb::call_guard()) - .def_prop_ro("max_rows", &kv::Batch::maxRows) - .def_prop_ro("max_blocks", &kv::Batch::maxBlocks) - .def_prop_ro("max_beam_width", &kv::Batch::maxBeamWidth) - .def_prop_ro("num_layer_groups", &kv::Batch::numLayerGroups) - .def( - "page_table", - [](std::shared_ptr self, int group) - { - self->pageTableAddress(kv::LayerGroupId{group}); - return BatchDeviceArray{std::move(self), kv::LayerGroupId{group}, false}; - }, - nb::arg("layer_group_id")) - .def( - "num_blocks", - [](std::shared_ptr self, int group) - { - self->numBlocksAddress(kv::LayerGroupId{group}); - return BatchDeviceArray{std::move(self), kv::LayerGroupId{group}, true}; - }, - nb::arg("layer_group_id")); - nb::class_(m, "ReuseGroupStatus") .def_prop_ro("group_id", [](kv::ReuseGroupStatus const& s) { return s.groupId.value(); }) .def_ro("policy", &kv::ReuseGroupStatus::policy) @@ -1933,8 +1926,7 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) nb::class_(m, "_KVCache") .def( "resume", - [](kv::KvCache& self, nb::object stream, std::optional isDecoding, - std::optional> groups) + [](kv::KvCache& self, nb::object stream, std::optional> groups) { std::optional optStream; if (!stream.is_none()) @@ -1949,31 +1941,14 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) } } nb::gil_scoped_release rel; - return self.resume(optStream, isDecoding, std::move(selected)); + return self.resume(optStream, std::move(selected)); }, - nb::arg("cuda_stream") = nb::none(), nb::arg("is_decoding") = nb::none(), - nb::arg("optional_reuse_groups") = nb::none()) - .def("enter_decode", &kv::KvCache::enterDecode, nb::call_guard()) - .def_prop_ro("is_decoding", &kv::KvCache::isDecoding) - .def_prop_ro("page_storage_version", &kv::KvCache::pageStorageVersion, nb::call_guard()) - .def_prop_ro("page_storage_dirty", &kv::KvCache::pageStorageDirty, nb::call_guard()) - .def_prop_ro("page_storage_row", &kv::KvCache::pageStorageRow, nb::call_guard()) - .def("bind_page_storage_row", &kv::KvCache::bindPageStorageRow, nb::arg("row").none(), - nb::call_guard()) - .def("acknowledge_page_storage", &kv::KvCache::acknowledgePageStorage, nb::arg("version"), - nb::call_guard()) + nb::arg("cuda_stream") = nb::none(), nb::arg("optional_reuse_groups") = nb::none()) .def( "get_page_priorities", [](kv::KvCache const& self, int layerGroupId, int beamIndex) { return self.getPagePriorities(kv::LayerGroupId{layerGroupId}, kv::BeamIndex{beamIndex}); }, nb::arg("layer_group_id"), nb::arg("beam_index") = 0, nb::call_guard()) - .def( - "get_page_storage_snapshot", - [](kv::KvCache const& self, int layerGroupId, int beamIdx) - { return self.getPageStorageSnapshot(kv::LayerGroupId{layerGroupId}, kv::BeamIndex{beamIdx}); }, - nb::arg("layer_group_id"), nb::arg("beam_id") = 0, nb::call_guard()) - .def("record_page_storage_read", &kv::KvCache::recordPageStorageRead, nb::arg("cuda_stream"), - nb::call_guard()) .def("suspend", &kv::KvCache::suspend, nb::call_guard()) .def( "prefetch", [](kv::KvCache& self, int target) { return self.prefetch(kv::CacheLevel{target}); }, @@ -1993,6 +1968,14 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) { return self.resize(capacity, historyLength); }, nb::arg("capacity") = std::nullopt, nb::arg("history_length") = std::nullopt, nb::call_guard()) + .def_prop_ro( + "max_input_length", [](kv::KvCache const& self) { return self.maxInputLength(); }, + nb::call_guard(), + "Bound on capacity - history_length, or None without a sparse fetch path.\n\n" + "One token, and read-only: a sparse request contributes exactly one query token " + "per step, because fetch_sparse_pages resolves one selection per request while an " + "attention kernel indexes its selection per query token. Prefill and speculative " + "decoding give a request several and are out of scope; this is the guard.") .def( "commit", [](kv::KvCache& self, nb::object acceptedInputTokens, nb::object beamSearchIndices, bool isEnd) @@ -2046,10 +2029,20 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) { return self.getSsmBlockBaseIndex(kv::LayerGroupId{layerGroupId}, kv::BeamIndex{beamId}); }, nb::arg("layer_group_id"), nb::arg("beam_id") = 0, nb::call_guard()) .def( - "get_scratch_desc", + "get_swa_scratch_range", + [](kv::KvCache const& self, int layerGroupId) + { return self.getSwaScratchRange(kv::LayerGroupId{layerGroupId}); }, + nb::arg("layer_group_id"), nb::call_guard(), + "Block ordinal range this life cycle serves from scratch slots; empty when none of it does.") + .def( + "get_scratch_slots", [](kv::KvCache const& self, int layerGroupId) - { return self.getScratchDesc(kv::LayerGroupId{layerGroupId}); }, - nb::arg("layer_group_id"), nb::call_guard()) + { + auto const slots = self.getScratchSlots(kv::LayerGroupId{layerGroupId}); + return std::vector(slots.begin(), slots.end()); + }, + nb::arg("layer_group_id"), nb::call_guard(), + "Page indices of this life cycle's scratch slots, in slot order.") .def_prop_ro("has_scratch_slots", &kv::KvCache::hasScratchSlots) .def_prop_rw("enable_swa_scratch_reuse", &kv::KvCache::isSwaScratchReuseEnabled, [](kv::KvCache& self, bool enable) @@ -2063,10 +2056,24 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) nb::gil_scoped_release release; self.setTextOnly(textOnly); }) - .def("supports_index_mode", &kv::KvCache::supportsIndexMode, nb::arg("mode")) + .def( + "supports_index_mode", + [](kv::KvCache const& self, kv::PageIndexMode mode, int layerGroupId) + { return self.supportsIndexMode(mode, kv::LayerGroupId{layerGroupId}); }, + nb::arg("mode"), nb::arg("layer_group_id"), "Per life cycle, which is where the scratch slots live.") .def_prop_ro("status", [](kv::KvCache const& kvc) { return kvc.status(); }) .def_prop_ro("is_active", &kv::KvCache::isActive) .def_prop_ro("num_blocks", [](kv::KvCache const& self) { return self.numBlocks().value(); }) + .def_prop_ro( + "batch_row", + [](kv::KvCache const& self) -> std::optional + { + auto const row = self.batchRow(); + return row.has_value() ? std::optional{row->value()} : std::nullopt; + }, + "Row this cache occupies in its Batch; None when it belongs to none.\n\n" + "Also the answer to who owns its page index buffers: a member's are the Batch's " + "rows, so set_base_page_index_buf refuses and suspend() leaves them bound.") .def_prop_ro("num_committed_blocks", &kv::KvCache::numCommittedBlocks) .def_prop_ro("num_committed_tokens", &kv::KvCache::numCommittedTokens) .def_prop_ro("cached_tokens_by_level", [](kv::KvCache const& self) { return self.cachedTokensByLevel().raw(); }) @@ -2133,7 +2140,6 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) kv::LayerGroupId const typedLayerGroupId{layerGroupId}; if (bufObj.is_none()) { - nb::gil_scoped_release release; self.setBasePageIndexBuf(typedBeamIdx, typedLayerGroupId, nullptr, 0); return; } @@ -2153,7 +2159,6 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) } cleanup{&view}; if (std::string(view.format) != "i" || view.ndim != 1) throw std::invalid_argument("set_base_page_index_buf: buffer must be 1-D int32 ('i')"); - nb::gil_scoped_release release; self.setBasePageIndexBuf(typedBeamIdx, typedLayerGroupId, static_cast(view.buf), static_cast(view.len / sizeof(int32_t))); }, @@ -2403,6 +2408,101 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) .def_ro("node_page_counts", &kv::HostMemPlacement::nodePageCounts); mIntrospection.def("probe_host_mem_placement", &kv::KvCacheIntrospection::probeHostMemPlacement, nb::arg("size"), nb::arg("allow_remote_numa_fallback"), nb::call_guard()); + mIntrospection.def("sparse_cache_resident_ordinals", &kv::KvCacheIntrospection::sparseCacheResidentOrdinals, + nb::arg("kv_cache"), nb::arg("lc_id"), nb::call_guard()); + mIntrospection.def("sparse_cache_base_cell", &kv::KvCacheIntrospection::sparseCacheBaseCell, nb::arg("kv_cache"), + nb::arg("lc_id"), nb::arg("pool") = 0, nb::call_guard()); + mIntrospection.def("sparse_flush_counts", &kv::KvCacheIntrospection::sparseFlushCounts, nb::arg("kv_cache"), + nb::arg("lc_id"), nb::call_guard()); + mIntrospection.def("batch_page_table", &kv::KvCacheIntrospection::batchPageTable, nb::arg("batch"), + nb::arg("layer_group_id"), nb::call_guard()); + + mIntrospection.def( + "batch_index_table_slot", + [](kv::Batch const& batch, int layerId, kv::DataRole const& role) + { + auto const [table, entry] = kv::KvCacheIntrospection::batchIndexTableSlot(batch, layerId, role); + return nb::make_tuple(table, entry); + }, + nb::arg("batch"), nb::arg("layer_id"), nb::arg("role"), + "(table, entry) serving this buffer. A consumer names a table by the spec it asked " + "for -- index_table_id -- and slices it at the buffer's position in that spec; this " + "is the same answer reached from the buffer."); + mIntrospection.def("batch_index_table_base", &kv::KvCacheIntrospection::batchIndexTableBase, nb::arg("batch"), + nb::arg("layer_id"), nb::arg("role"), + "What a converted index for this buffer is added to. Buffers sharing a pool share it, " + "their own offsets being already in the entries."); + mIntrospection.def("batch_index_table", &kv::KvCacheIntrospection::batchIndexTable, nb::arg("batch"), + nb::arg("table_id"), + "One converted index table, row-major [max_batch][entries][blocks_per_seq]. Reads the " + "device copy, so it shows what the last publish() converted. Synchronises."); + mIntrospection.def( + "batch_gpu_coalesced_page_table", + [](nb::object batchObj, int layerGroupId) + { + auto const& batch = nb::cast(batchObj); + auto const view = kv::KvCacheIntrospection::batchGpuCoalescedPageTable(batch, layerGroupId); + // The Batch owns the memory; batchObj keeps it alive for the array's lifetime. + return makeDeviceArray2D(view.data, view.numRows, view.blocksPerRow, batchObj); + }, + nb::arg("batch"), nb::arg("layer_group_id"), + "int32[max_batch, max_blocks_per_seq] coalesced page indices for one layer group, as a " + "DLPack array on the device.\n\n" + "A view, not a copy, and so unlike its neighbours it neither synchronises nor returns a " + "plain list. The device copy exists for the index conversion to read, and the kernels a " + "consumer writes read the converted table instead: Batch.coalesced_page_table() is where " + "the indices themselves are read.\n\n" + "It shows the staged step rather than the published one, and it is live. Read it " + "between wait_published() and the next publish(): before the first the upload filling " + "it may not have landed, and once the second is enqueued the upload stream is free to " + "overwrite it."); + mIntrospection.def("batch_sparse_entries_per_row", &kv::KvCacheIntrospection::batchSparseEntriesPerRow, + nb::arg("batch"), nb::arg("layer_group_id"), nb::call_guard()); + mIntrospection.def( + "batch_watermarks", + [](kv::Batch const& batch) + { + auto r = kv::KvCacheIntrospection::batchWatermarks(batch); + return std::make_tuple(std::move(r.first), std::move(r.second)); + }, + nb::arg("batch")); + mIntrospection.def( + "sparse_resolve", + [](std::vector const& residentOrdinal, std::vector const& lruOrder, int numLayers, + int entriesPerLayer, std::vector const& selected, int selectionWidth, + std::vector const& historyUnits, std::vector const& capacityUnits, int maxInFlightUnits, + std::vector const& inactiveRows) + { + auto r = kv::KvCacheIntrospection::sparseResolve(residentOrdinal, lruOrder, numLayers, entriesPerLayer, + selected, selectionWidth, historyUnits, capacityUnits, maxInFlightUnits, inactiveRows); + // No gil_scoped_release call guard: the tuple below is built from the lambda body, + // which the guard would have left running without the GIL. + return std::make_tuple(std::move(r.residentOrdinal), std::move(r.lruOrder), std::move(r.outSlot), + std::move(r.missRow), std::move(r.missOrdinal), std::move(r.missSlot), std::move(r.missTotal), + std::move(r.missCount), std::move(r.inFlightSlot)); + }, + nb::arg("resident_ordinal"), nb::arg("lru_order"), nb::arg("num_layers"), nb::arg("entries_per_layer"), + nb::arg("selected"), nb::arg("selection_width"), nb::arg("history_units"), nb::arg("capacity_units"), + nb::arg("max_in_flight_units"), nb::arg("inactive_rows")); + mIntrospection.def("sparse_unit_copy", &kv::KvCacheIntrospection::sparseUnitCopy, nb::arg("host_pool"), + nb::arg("host_page_stride"), nb::arg("page_table"), nb::arg("device_arena"), nb::arg("base_cell"), + nb::arg("num_blocks"), nb::arg("cells_per_buffer"), nb::arg("buffer_bytes"), nb::arg("num_chunks"), + nb::arg("chunk_bytes"), nb::arg("chunk_stride"), nb::arg("unit_stride"), nb::arg("expansion"), + nb::arg("miss_ordinal"), nb::arg("miss_slot"), nb::arg("miss_count"), nb::arg("selection_width"), + nb::arg("buffer_layer"), nb::arg("write_back") = false, nb::arg("grid_ctas") = 1, + nb::call_guard()); + mIntrospection.def( + "sparse_flush_list", + [](std::vector const& residentOrdinal, int numLayers, int entriesPerLayer, int fromUnits, int toUnits, + int maxFlushUnits) + { + auto r = kv::KvCacheIntrospection::sparseFlushList( + residentOrdinal, numLayers, entriesPerLayer, fromUnits, toUnits, maxFlushUnits); + return std::make_tuple( + std::move(r.row), std::move(r.ordinal), std::move(r.slot), std::move(r.total), std::move(r.count)); + }, + nb::arg("resident_ordinal"), nb::arg("num_layers"), nb::arg("entries_per_layer"), nb::arg("from_units"), + nb::arg("to_units"), nb::arg("max_flush_units")); mIntrospection.def( "is_commit_allowed", [](kv::KvCache const& kvCache) { return kvCache.commitState() == kv::KvCache::CommitState::ALLOWED; }, @@ -2548,33 +2648,61 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) nb::arg("manager"), nb::arg("cache_level") = kv::kHotLevel.value(), nb::call_guard()); mIntrospection.def( "grains_for_slots", - [](kv::SlotCount numSlots, std::vector const& slotSizeList, size_t granularity) + [](kv::SlotCount numSlots, std::vector const& slotSizeList, size_t granularity, + std::optional> const& pools) { if (numSlots < 0) { throw std::invalid_argument("num_slots must be non-negative"); } - return kv::CacheLevelStorage::grainsForSlots(numSlots, typedPoolSizeList(slotSizeList), granularity); + auto const sizes = typedPoolSizeList(slotSizeList); + auto const which = typedPoolSelectors( + pools.has_value() ? std::optional{std::vector>{*pools}} : std::nullopt, + kv::TypedVec>{{sizes}}); + return kv::CacheLevelStorage::grainsForSlots(numSlots, sizes, which[kv::PoolGroupIndex{0}], granularity); }, - nb::arg("num_slots"), nb::arg("slot_size_list"), nb::arg("granularity"), + nb::arg("num_slots"), nb::arg("slot_size_list"), nb::arg("granularity"), nb::arg("pools").none() = std::nullopt, nb::call_guard()); mIntrospection.def( "grains_to_slots", - [](size_t pgGrains, std::vector const& slotSizeList, size_t granularity) + [](size_t pgGrains, std::vector const& slotSizeList, size_t granularity, + std::optional> const& pools) { + auto const sizes = typedPoolSizeList(slotSizeList); + auto const which = typedPoolSelectors( + pools.has_value() ? std::optional{std::vector>{*pools}} : std::nullopt, + kv::TypedVec>{{sizes}}); auto [slots, used] - = kv::CacheLevelStorage::grainsToSlots(pgGrains, typedPoolSizeList(slotSizeList), granularity); + = kv::CacheLevelStorage::grainsToSlots(pgGrains, sizes, which[kv::PoolGroupIndex{0}], granularity); return std::make_tuple(slots, used); }, - nb::arg("pg_grains"), nb::arg("slot_size_list"), nb::arg("granularity"), + nb::arg("pg_grains"), nb::arg("slot_size_list"), nb::arg("granularity"), nb::arg("pools").none() = std::nullopt, nb::call_guard()); mIntrospection.def( "ratio_to_slot_count_list", [](size_t quota, std::vector> const& slotSizeLists, std::vector const& ratioList, - size_t granularity, std::vector const& minSlots) + size_t granularity, std::vector const& minSlots, std::optional secondQuota, + std::optional secondGranularity, + std::optional>> const& secondPools) { - auto slotCountList = kv::CacheLevelStorage::ratioToSlotCountList(quota, typedSlotSizeLists(slotSizeLists), - kv::TypedVec{ratioList}, granularity, typedSlotCounts(minSlots)); + auto const sizes = typedSlotSizeLists(slotSizeLists); + std::optional second; + if (secondQuota.has_value() || secondPools.has_value()) + { + // Omitting the pools means every pool, which would charge the whole layout to + // the second budget and leave `quota` paying for nothing. That is never what a + // caller naming a second quota wants, so it has to say which pools it covers. + TLLM_CHECK_WITH_INFO(secondPools.has_value(), + "second_quota needs second_pools: without it every pool is charged to the " + "second budget and the primary quota pays for nothing"); + kv::SecondaryPoolBudget budget; + budget.quota = secondQuota; + budget.granularity = secondGranularity.value_or(granularity); + budget.pools = typedPoolSelectors(secondPools, sizes); + second = std::move(budget); + } + auto slotCountList = kv::CacheLevelStorage::ratioToSlotCountList(quota, sizes, + kv::TypedVec{ratioList}, granularity, typedSlotCounts(minSlots), second); std::vector result; result.reserve(slotCountList.stdSize()); for (kv::SlotCount slotCount : slotCountList) @@ -2584,7 +2712,14 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) return result; }, nb::arg("quota"), nb::arg("slot_size_lists"), nb::arg("ratio_list"), nb::arg("granularity"), - nb::arg("min_slots"), nb::call_guard()); + nb::arg("min_slots"), nb::arg("second_quota").none() = std::nullopt, + nb::arg("second_granularity").none() = std::nullopt, nb::arg("second_pools").none() = std::nullopt, + nb::call_guard()); + + mIntrospection.def("tier_allocation_granularity", &kv::tierAllocationGranularity, nb::arg("tier"), nb::arg("quota"), + "The unit a tier of this quota allocates its pools in. Reached through _introspection " + "because it carries no stability promise: it exists so a test can size against the rule " + "the native tier actually applies rather than restating it."); // CUDA virtual-memory primitives, reached through _introspection because they carry no // stability promise: the native disaggregated bounce buffer reserves one contiguous fabric @@ -2608,6 +2743,503 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) "native codec factory exposes an owning IKvCacheColdPageCodec object for transfer into KVCacheManager. Any " "KVCacheManager construction attempt consumes an explicitly supplied codec, including an attempt that fails."); + // ---- Batch ------------------------------------------------------------- + nb::class_(m, "IndexTableSpec") + .def(nb::init<>()) + .def( + "__init__", + [](kv::IndexTableSpec* self, std::vector buffers) + { new (self) kv::IndexTableSpec{std::move(buffers)}; }, + nb::arg("buffers"), + "Buffers a consumer must address from one table, under one base pointer.\n\n" + "States a kernel's requirement, not a preference: some attention ops take a single " + "block table covering K and V rather than one each. Order defines the buffer " + "dimension, so a kernel reads buffer i of this request at entry i. Every buffer must " + "share a life cycle and a page index scale, or the request is refused.") + .def_rw("buffers", &kv::IndexTableSpec::buffers); + + nb::class_(m, "Batch") + .def( + "__init__", + [](kv::Batch* self, kv::KvCacheManager& manager, intptr_t stream, + std::vector indexTables, std::optional maxBatchSize, + std::optional maxBlocksPerSeq, std::optional mode, + std::optional maxTotalInputTokens) + { + new (self) kv::Batch(manager, reinterpret_cast(stream), std::move(indexTables), maxBatchSize, + maxBlocksPerSeq, mode, maxTotalInputTokens); + }, + nb::arg("manager"), nb::arg("stream"), nb::arg("index_tables"), nb::arg("max_batch_size") = nb::none(), + nb::arg("max_blocks_per_seq") = nb::none(), nb::arg("mode") = nb::none(), + nb::arg("max_total_input_tokens") = nb::none(), nb::keep_alive<1, 2>(), + "The KV caches stepped together, and their device-side page tables.\n\n" + "index_tables names what a consumer needs addressable together, and is required: pass [] " + "to convert no indices, which is a choice about what the consumer reads rather than an " + "absence the Batch could fill in.\n\n" + "max_batch_size and max_blocks_per_seq bound device arrays whose shape must not change " + "between steps, so neither can be grown later. Left unset they are derived from the " + "manager's constraints, which fails when none were declared.\n\n" + "max_total_input_tokens is the tokens one step may ingest across every row. It sizes the " + "write-back's work list; left unset it assumes a decode step, where each row advances by " + "one token and so crosses at most one unit boundary.\n\n" + "mode applies to every life cycle. Left unset, each takes the smaller table: SHARED where " + "KVCacheManager.supports_index_mode says it is always available, PER_LAYER otherwise. Set, " + "it overrides that answer rather than being checked against it, so SHARED may be named " + "where the manager returns None. The manager answers for every request a life cycle could " + "carry, so a sliding window leaves it undecided; naming SHARED claims that no member of " + "this batch will hold scratch slots there, which a decode-only batch can know because " + "scratch covers the blocks a prefill chunk adds. That is what lets prefill and decode run " + "as two batches over one manager with SWA scratch enabled, the decode one still taking " + "SHARED for the sliding life cycles. Every publish() validates the claim against each " + "member's own KvCache.supports_index_mode and raises if one holds scratch.\n\n" + "The Batch stages the host page indices itself; host_page_table() exposes them for a " + "consumer that needs the same numbers, which leaves one buffer with one writer and a " + "member's binding alive for the whole decode.") + .def( + "add_kv_cache", + [](kv::Batch& self, kv::KvCache& kvCache, std::optional row) + { + std::optional seat; + if (row.has_value()) + { + seat = kv::BatchRow{*row}; + } + return self.addKvCache(kvCache, seat).value(); + }, + nb::arg("kv_cache"), nb::arg("row") = nb::none(), nb::call_guard(), + "Seat a cache in `row`, or in the lowest free row when none is given, returning the " + "row. Raises when the row is taken or out of range, the batch is full, or the cache " + "already belongs to a batch.\n\n" + "The batch holds a reference while the cache is seated, so a caller may keep only the " + "batch; remove_kv_cache hands it back.\n\n" + "Seat a request only once it is ready to be stepped. Seating takes the sparse arena " + "lease, so a request that holds pages without being stepped -- a disaggregated one " + "still receiving its KV, or one sitting out a step -- costs the arena nothing while it " + "stays out, and the arena is sized for the rows rather than for every request the " + "scheduler admits because of that. Nothing checks it: an early seat shows up as arena " + "pressure, not as an error.") + .def( + "remove_kv_cache", [](kv::Batch& self, kv::KvCache& kvCache) { return self.removeKvCache(kvCache); }, + nb::arg("kv_cache"), nb::call_guard(), + "Free the cache's row, leaving a hole for reuse, and return the cache. Raises when it " + "is not a member.\n\n" + "Returns the arena lease, which discards whatever the GPU cache held: a lease comes " + "back with every entry free. Call flush_sparse_units() first for anything worth " + "keeping.") + .def( + "remove_row", [](kv::Batch& self, int row) { return self.removeKvCache(kv::BatchRow{row}); }, + nb::arg("row"), nb::call_guard(), + "The same, by the row add_kv_cache returned. None when nothing is seated there, so a " + "caller sweeping rows need not ask twice about each one.") + .def( + "move", + [](kv::Batch& self, int dstRow, int srcRow) { self.move(kv::BatchRow{dstRow}, kv::BatchRow{srcRow}); }, + nb::arg("dst_row"), nb::arg("src_row"), nb::call_guard(), + "Move the member at src_row to the empty dst_row, leaving src_row free.\n\n" + "Membership does not lapse: the sparse cache lease, its contents and the stream " + "binding all survive, which is what separates this from remove_row() followed by " + "add_kv_cache() - that pair hands the lease back with every entry free. It also " + "writes each life cycle's index-table row once instead of twice.\n\n" + "Raises when either row is out of range, src_row holds nothing, or dst_row is " + "occupied. Waits for the previous staging upload. Device tables reach the new row on " + "the next publish(), so a step may resize, move and publish once rather than " + "publishing per operation.") + .def( + "index_mode", + [](kv::Batch const& self, int layerGroupId) { return self.indexMode(kv::LayerGroupId{layerGroupId}); }, + nb::arg("layer_group_id"), + "Page index mode of one life cycle, as chosen at construction.\n\n" + "Unset at construction, each life cycle takes the smaller table: SHARED where the " + "manager reports it always available, PER_LAYER otherwise.") + .def_prop_ro("num_index_tables", [](kv::Batch const& self) { return self.numIndexTables().value(); }) + .def( + "index_table", + [](kv::Batch const& self, int tableId) { return self.indexTable(kv::IndexTableId{tableId}); }, + nb::arg("table_id"), + "Every request this table satisfies. A request's buffer order is its own mapping " + "into the buffer dimension, so the first buffer of each sits at entry zero.") + .def( + "index_table_id", + [](kv::Batch const& self, kv::IndexTableSpec const& spec) { return self.indexTableOf(spec).value(); }, + nb::arg("spec"), + "Which table serves this spec, which must be one the Batch was built with.\n\n" + "The query to prefer: a caller that asked for a table by naming its buffers gets it " + "back without going through a buffer, and equivalent specs resolve to the same id -- " + "a layer's data and its block scales name one table between them. The spec's buffers " + "occupy entries [0, len(spec.buffers)) in its own order, so an entry needs no " + "separate lookup, and converted_index_table(id)[:, entry, :] is a buffer's plane.") + .def( + "pool_spec_for", + [](kv::Batch const& self, kv::IndexTableSpec const& spec) + { + auto const pool = self.poolSpecFor(spec); + return nb::make_tuple(pool.base, pool.pageStride); + }, + nb::arg("spec"), + "(base, page_stride) for this spec's converted indices: a page index p of this " + "spec addresses base + p * page_stride.\n\n" + "Keyed by the spec because the base is. Specs whose entries coincide are merged " + "onto one table, and reading that one set of entries through different bases is " + "what keeps a quantised layer's data and its block scales apart -- so a table has " + "as many bases as it serves specs, and only the spec picks one.") + .def( + "converted_index_table", + [](nb::object self, int tableId) + { + auto const& batch = nb::cast(self); + auto const view = batch.convertedIndexTable(kv::IndexTableId{tableId}); + // The Batch owns the memory; self keeps it alive for the array's lifetime. + return makeDeviceArray3D(view.data, view.numRows, view.numBuffers, view.blocksPerRow, self); + }, + nb::arg("table_id"), + "int32[max_batch, num_buffers, max_blocks_per_seq] converted page indices for one " + "table, as a DLPack array on the device. A buffer reads the entry index_table_slot " + "gives it; a row's entries are adjacent, which is the shape an attention kernel " + "takes a key and value page list in.\n\n" + "The address is stable for the Batch's lifetime, so a captured graph may bake it; " + "the array keeps the Batch alive meanwhile.") + .def( + "resize", + [](kv::Batch& self, std::vector> const& capacities, + std::vector> const& historyLengths) + { + nb::gil_scoped_release rel; + return self.resize(capacities, historyLengths); + }, + nb::arg("capacities"), nb::arg("history_lengths"), + "Resize every member, then publish once.\n\n" + "The sanctioned way to advance a batched step: it waits on the previous publish " + "before touching anyone, so the staging hazard is closed without the caller " + "remembering wait_published(). One result per row, in row order, 0 for a free " + "row; a member that fails rolls itself back, so a batch can partially fail and the " + "table published afterwards describes the post-rollback state.\n\n" + "Rejects, before touching any member, a history length that falls or that outruns the " + "capacity its row already had - capacity this call grants makes room for tokens written " + "after it. A row whose growth fails still takes its history advance, which needs no " + "slots and lets the units below it be written back; its result is still 0, since the " + "row is not at the state that was asked for.") + .def( + "flush_sparse_units", + [](kv::Batch& self, std::vector const& newHistoryLengths) + { + nb::gil_scoped_release rel; + self.flushSparseUnits(newHistoryLengths); + }, + nb::arg("new_history_lengths"), + "Write back the units each row's watermark is about to pass, given the history you are about " + "to set. One per row; a row may repeat its current history, which asks for nothing.\n\n" + "Call this before advancing history, not after: a block must be in its host page before it can " + "be selected, and resize() refuses a history this has not covered. Bounded by the capacity each " + "row started from, so a resumed request moving history across a whole prompt writes back " + "nothing. Reads the index table as of the last publish.") + .def( + "prepare_sparse_write_indices", + [](kv::Batch& self, std::vector const& buffers, DeviceArray2D const& out, nb::object stream) + { + if (static_cast(out.shape(0)) != self.maxBatchSize()) + { + throw std::invalid_argument( + "prepare_sparse_write_indices: out must carry one row per batch row, not one per " + "query token. A request has at most one unit in flight; prefill and speculative " + "decoding give it several and are out of scope"); + } + if (out.shape(1) != buffers.size()) + { + throw std::invalid_argument("prepare_sparse_write_indices: out's last extent must be len(buffers)"); + } + auto const handle = stream.is_none() + ? std::optional{} + : std::optional{reinterpret_cast(nb::cast(stream))}; + auto* outPtr = out.data(); + nb::gil_scoped_release rel; + self.prepareSparseWriteIndices(buffers, outPtr, handle); + }, + nb::arg("buffers"), nb::arg("out").noconvert(), nb::arg("stream") = nb::none(), + "Give each request's in-flight unit a cell, and report where to write it.\n\n" + "The unit between the two watermarks is the one the model is about to produce. It has no " + "host copy yet -- a unit reaches its page only when resize() carries the watermark past " + "it -- so nothing is fetched: the call assigns a cell and leaves its bytes to the model.\n\n" + "Separate from fetch_sparse_pages because the orders differ. A fetch needs the selection, " + "which the indexer produces from the previous layer; a write needs only the watermarks, so " + "it can run first, which is the order the model requires.\n\n" + "Idempotent while the watermarks and the reservation hold: a later call, including the " + "fetch's own resolve, finds the entry and reuses it. The entry is pinned while its unit is " + "in flight, so no fetch can evict it.\n\n" + "buffers is one layer's sparse buffers, exactly as for fetch_sparse_pages. out is " + "int32[batch, len(buffers)], written in full, in the same expanded index space " + "fetch_sparse_pages produces and relative to the same get_sparse_cache_base_address " + "region: BAD_PAGE_INDEX where the row holds no reservation or nothing in flight.\n\n" + "Device-driven and stream-ordered: enqueues and returns, and every shape is static, so a " + "CUDA graph can capture it.") + .def( + "fetch_sparse_pages", + [](kv::Batch& self, std::vector const& buffers, DeviceArray2D const& selected, + DeviceArray3D const& out, nb::object stream) + { + if (selected.shape(0) != out.shape(0)) + { + throw std::invalid_argument("fetch_sparse_pages: selected and out must agree on the batch size"); + } + if (static_cast(selected.shape(0)) != self.maxBatchSize()) + { + throw std::invalid_argument( + "fetch_sparse_pages: selected must carry one row per batch row, not one per query " + "token. A request contributes a single selection per step; prefill and speculative " + "decoding give it several and are out of scope"); + } + if (out.shape(1) != buffers.size()) + { + throw std::invalid_argument("fetch_sparse_pages: out's middle extent must be len(buffers)"); + } + if (out.shape(2) != selected.shape(1)) + { + throw std::invalid_argument( + "fetch_sparse_pages: out's last extent must be the selection width; the result is compact, " + "one index per selection slot, not a table indexed by ordinal"); + } + auto const handle = stream.is_none() + ? std::optional{} + : std::optional{reinterpret_cast(nb::cast(stream))}; + auto* selectedPtr = selected.data(); + auto* outPtr = out.data(); + auto const width = static_cast(selected.shape(1)); + nb::gil_scoped_release rel; + self.fetchSparsePages(buffers, selectedPtr, width, outPtr, handle); + }, + nb::arg("buffers"), nb::arg("selected"), nb::arg("out").noconvert(), nb::arg("stream") = nb::none(), + "Resolve a batch's selections against each buffer's per-request history cache.\n\n" + "buffers must be exactly one layer's sparse buffers, in the cache's order. One layer, " + "not one layer group, because a layer's selection is normally computed from the " + "previous layer's output. The whole layer, not a subset, because the resolve marks " + "entries resident on behalf of every buffer at once.\n\n" + "Only for a layer whose selection is its own. Layers sharing a selection_id know their " + "selection together and resolve it once between them: call resolve_sparse_selection " + "once and fetch_resolved_pages per layer. This rejects them rather than letting the " + "second resolve count the first's admissions as hits and leave its cells unfilled.\n\n" + "selected is int32[batch, topk] unit ordinals, one row per request rather than per " + "query token, negative where padded, and must name only units below the history " + "watermark. out is int32[batch, len(buffers), topk], " + "written in full in the expanded index space: entry k is the index holding " + "selected[b][k], or BAD_PAGE_INDEX where that slot carries no selection. Compact, in " + "the selection's own order, because that is the shape sparse attention kernels take.\n\n" + "Indices are relative to get_sparse_cache_base_address for the buffer, not to " + "get_mem_pool_base_address and not to a cache level: the pool names the buffer's " + "pages, which are host memory for a sparse life cycle, while the cells a kernel " + "reads live in the arena.\n\n" + "Device-driven and stream-ordered: enqueues and returns, and every shape is static, " + "so a CUDA graph can capture it.\n\n" + "stream defaults to the batch's own, which is the only correct choice while layers are " + "fetched in order, since a layer's selection comes from the one before it. Name a stream " + "only to run ahead of that, and then order it against the batch's stream yourself: the " + "resolve's scratch and counters are batch state the step clears on that stream.") + .def( + "resolve_sparse_selection", + [](kv::Batch& self, std::vector const& buffers, DeviceArray2D const& selected, + nb::object stream) + { + if (static_cast(selected.shape(0)) != self.maxBatchSize()) + { + throw std::invalid_argument( + "resolve_sparse_selection: selected must carry one row per batch row, not one per " + "query token"); + } + auto const handle = stream.is_none() + ? std::optional{} + : std::optional{reinterpret_cast(nb::cast(stream))}; + auto* selectedPtr = selected.data(); + auto const width = static_cast(selected.shape(1)); + nb::gil_scoped_release rel; + self.resolveSparseSelection(buffers, selectedPtr, width, handle); + }, + nb::arg("buffers"), nb::arg("selected"), nb::arg("stream") = nb::none(), + "Resolve one selection, admitting what it names and listing what has to be copied.\n\n" + "The first half of fetch_sparse_pages, split out for the layers that share a " + "selection_id: they resolve once between them and then fetch a layer at a time. " + "buffers names any one member layer -- the residency and the miss list belong to the " + "selection, not to the layer that happened to name it.\n\n" + "The miss list it leaves behind is read by every fetch_resolved_pages that follows, " + "and is overwritten by the next resolve of the same selection.") + .def( + "fetch_resolved_pages", + [](kv::Batch& self, std::vector const& buffers, DeviceArray2D const& selected, + DeviceArray3D const& out, nb::object stream) + { + if (selected.shape(0) != out.shape(0)) + { + throw std::invalid_argument("fetch_resolved_pages: selected and out must agree on the batch size"); + } + if (static_cast(selected.shape(0)) != self.maxBatchSize()) + { + throw std::invalid_argument( + "fetch_resolved_pages: selected must carry one row per batch row, not one per " + "query token"); + } + if (out.shape(1) != buffers.size()) + { + throw std::invalid_argument("fetch_resolved_pages: out's middle extent must be len(buffers)"); + } + if (out.shape(2) != selected.shape(1)) + { + throw std::invalid_argument("fetch_resolved_pages: out's last extent must be the selection width"); + } + auto const handle = stream.is_none() + ? std::optional{} + : std::optional{reinterpret_cast(nb::cast(stream))}; + auto* selectedPtr = selected.data(); + auto* outPtr = out.data(); + auto const width = static_cast(selected.shape(1)); + nb::gil_scoped_release rel; + self.fetchResolvedPages(buffers, selectedPtr, width, outPtr, handle); + }, + nb::arg("buffers"), nb::arg("selected"), nb::arg("out").noconvert(), nb::arg("stream") = nb::none(), + "Copy one layer's share of a resolved selection into its cells, and write its indices.\n\n" + "The second half of fetch_sparse_pages. selected must be the array the matching " + "resolve_sparse_selection read, since the indices written are those units'; out is as " + "fetch_sparse_pages describes it, for this layer's buffers alone.") + .def( + "flushed_history", [](kv::Batch const& self, int row) { return self.flushedHistory(kv::BatchRow{row}); }, + nb::arg("row"), + "The history a row has been written back to. A row that has never needed a flush reads its own " + "history, so a caller may always compare against it.") + .def( + "publish", + [](kv::Batch& self) + { + nb::gil_scoped_release rel; + self.publish(); + }, + "Copy the staged page indices to the device, enqueuing on the Batch's stream. What this adds " + "over the staging rows, which members keep current themselves, is a stream-ordered " + "snapshot: without it the host would be mutating a table a previously enqueued fetch " + "is still reading.\n\n" + "The copy is asynchronous and reads pinned memory, so it is still in flight on return. " + "Call wait_published() before any later host write to a member's page indices - a " + "resize, a commit, a beam fork, an SWA recycle - or that write races the copy.") + .def( + "wait_published", + [](kv::Batch& self) + { + nb::gil_scoped_release rel; + self.waitPublished(); + }, + "Block until the last publish() has finished reading the staging rows. Cheap once " + "complete and a no-op when nothing is in flight, so it is safe at the top of a step.") + .def( + "coalesced_page_table", + [](nb::object self, int layerGroupId) + { + auto const& batch = nb::cast(self); + auto const view = batch.coalescedPageTable(kv::LayerGroupId{layerGroupId}); + // The Batch owns the memory; self keeps it alive for the array's lifetime. + return makeHostArray2D(view.data, view.numRows, view.blocksPerRow, self); + }, + nb::arg("layer_group_id"), + "int32[max_batch, max_blocks_per_seq] coalesced page indices for one layer group, " + "as a read-only DLPack array on the host. Pinned, and written by the members " + "themselves, so a consumer needing these numbers reads them here rather than keeping " + "a buffer of its own. Writing here would race a publish already reading the same " + "memory, and would be lost in any case: a resize() rewrites only what it changes. " + "An entry names a coalesced slot; the converted tables name one buffer inside " + "one.") + .def( + "history_lengths", + [](nb::object self) + { + auto const& batch = nb::cast(self); + auto const view = batch.watermarks(); + return makeDeviceArray1D(view.historyLength, view.numRows, self); + }, + "int32[max_batch] per-row history length, in tokens.\n\n" + "Published by the same transfer as the page tables, because it bounds which entries of " + "a row are valid: pairing one step's table with another step's lengths would read past " + "the end of a reused row. num_blocks is history_length // tokens_per_block; the resolve " + "works in units and divides by unit_tokens instead. Publishing the source rather than " + "the derived counts keeps one truth.") + .def( + "capacities", + [](nb::object self) + { + auto const& batch = nb::cast(self); + auto const view = batch.watermarks(); + return makeDeviceArray1D(view.capacity, view.numRows, self); + }, + "int32[max_batch] per-row capacity, in tokens. Published with history_lengths().") + + .def( + "row_of", + [](kv::Batch const& self, kv::KvCache const& kvCache) -> std::optional + { + auto const row = self.rowOf(kvCache); + return row.has_value() ? std::optional{row->value()} : std::nullopt; + }, + nb::arg("kv_cache"), "Row of a member; None when this cache is not one.") + .def( + "__getitem__", [](kv::Batch const& self, int row) { return self[kv::BatchRow{row}]; }, nb::arg("row"), + "The member seated in `row`, or None when that row is free. The inverse of " + "row_of(). Raises IndexError when the row is beyond max_batch_size, so None " + "always means an empty row rather than a bad index.") + .def( + "peek_hit_miss_counter", + [](kv::Batch const& self) + { + kv::SparseHitMiss stats; + { + // Copies device counters and synchronises, which is not work to hold the + // GIL across; make_tuple below needs it again. + nb::gil_scoped_release rel; + stats = self.peekHitMissCounter(); + } + return nb::make_tuple(stats.hits(), stats.misses); + }, + "Cumulative (hits, misses) in selection units since the Batch was built or the " + "counters were reset.\n\n" + "A hit is a selected unit the GPU cache already held; a miss is one the fetch copied " + "from its host page. Padded selection entries and in-flight units are neither, so they " + "are in neither figure.\n\n" + "Reads on a temporary stream and syncs only that, so the pair may be a few launches " + "stale -- which does not move a ratio. Synchronise the device first for an exact one.") + .def("reset_hit_miss_counter", &kv::Batch::resetHitMissCounter, + "Zero the cumulative pair, on the Batch's stream.") + .def_prop_ro("sparse_arena_bytes", &kv::Batch::sparseArenaBytes, + "Device bytes every life cycle's sparse arena holds, 0 without sparse offload.\n\n" + "The arenas sit outside the cache tiers' quota, so a caller budgeting device memory " + "accounts for them separately and otherwise has to predict the layout to do it. This " + "is what was allocated, which makes that prediction checkable.") + .def_prop_ro("num_kv_caches", &kv::Batch::numKvCaches, "Members seated, which is the rows occupied.") + .def_prop_ro("is_contiguous", &kv::Batch::isContiguous, + "Whether the occupied rows are exactly [0, num_kv_caches).\n\n" + "A Batch does not require it: removing a member leaves its row free and the next join " + "takes the lowest free one, so a hole is legal and transient. An engine that maps rows " + "onto its own request slots usually does require it, and asking here beats keeping a " + "mirror of the membership to check against.") + .def( + "sparse_out_slots", + [](kv::Batch const& self, std::vector const& buffers) + { return self.sparseOutSlots(buffers); }, + nb::arg("buffers"), + "Where each buffer reads its indices in fetch_sparse_pages' output.\n\n" + "Buffers whose converted indices coincide share a slice: a layer's data and scale pools " + "mirror, so one slice serves both and the caller applies the base of whichever pool it " + "reads. A scale shared by key and value is its own index table and keeps a slice of its " + "own. One entry per buffer, in the order given; out's middle extent is one more than the " + "largest.") + .def_prop_rw( + "cuda_stream", [](kv::Batch const& self) { return reinterpret_cast(self.cudaStream()); }, + [](kv::Batch& self, intptr_t stream) { self.setCudaStream(reinterpret_cast(stream)); }, + "The stream the Batch and every member run on. Assigning moves the Batch and every " + "member onto the new stream, which first waits on what the old one has queued.") + .def_prop_ro("max_batch_size", &kv::Batch::maxBatchSize) + .def("get_sparse_cache_base_address", &kv::Batch::getSparseCacheBaseAddress, nb::arg("layer_id"), + nb::arg("data_role"), nb::call_guard(), + "Base of the sparse cache arena backing this buffer, for a consumer reading units the " + "fetch resolved.\n\n" + "The arena is the Batch's: it is sized by max_batch_size, one lease per row, so the " + "address is only meaningful for caches seated in this Batch.\n\n" + "Distinct from KVCacheManager.get_mem_pool_base_address, which names the buffer's pages " + "for every buffer alike - what the transceiver and the KV transfer path address. The " + "arena is a second allocation holding the only GPU-resident copy a kernel can read.") + .def_prop_ro("max_blocks_per_seq", &kv::Batch::maxBlocksPerSeq); + // ---- KvCacheManager ---------------------------------------------------- nb::class_(m, "KVCacheManager") .def( @@ -2669,7 +3301,6 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) }); }, nb::arg("reuse_scope"), nb::arg("input_tokens"), nb::arg("target")) - .def("is_sparse", &kv::KvCacheManager::isSparse, nb::arg("layer_id"), nb::arg("data_role")) .def( "clear_reusable_blocks", &kv::KvCacheManager::clearReusableBlocks, nb::call_guard()) .def( @@ -2755,9 +3386,34 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) .def("get_mem_pool_base_address", &kv::KvCacheManager::getMemPoolBaseAddress, nb::arg("layer_id"), nb::arg("data_role"), nb::arg("index_mode") = std::nullopt, nb::call_guard()) .def("get_page_stride", &kv::KvCacheManager::getPageStride, nb::arg("layer_id"), nb::arg("data_role")) + .def("sparse_index_table_specs", &kv::KvCacheManager::sparseIndexTableSpecs, + "The index table specs every sparse life cycle needs, one per (layer, pool) run.\n\n" + "A Batch still has to be given these: it refuses a sparse buffer it was told nothing about, " + "because a fetch with no converted indices would read the wrong pages rather than fail. Their " + "shape is not a choice, so this saves a caller re-deriving the pool layout to restate it.") .def("get_page_index_scale", &kv::KvCacheManager::getPageIndexScale, nb::arg("layer_id"), nb::arg("data_role")) .def("get_page_index_upper_bound", &kv::KvCacheManager::getPageIndexUpperBound, nb::arg("layer_id"), nb::arg("data_role"), nb::call_guard()) + .def("is_sparse", &kv::KvCacheManager::isSparse, nb::arg("layer_id"), nb::arg("data_role"), + nb::call_guard(), + "Whether this buffer was declared BufferConfig.is_sparse.\n\n" + "Every buffer locks its pages at the hot level, sparse or not. What a sparse buffer " + "changes is the medium: its pages are host memory, device-mapped, and are read through " + "the per-request GPU cache rather than directly.") + .def( + "get_num_bytes", + [](kv::KvCacheManager const& self, int cacheLevel, kv::CacheTier medium) + { return self.getNumBytes(kv::CacheLevel{cacheLevel}, medium); }, + nb::arg("level"), nb::arg("medium"), nb::call_guard(), + "Bytes this level holds in pools whose pages are `medium`.\n\n" + "A level's quota pays only for its own tier, so bytes of a second medium appear in no " + "other total a caller can read and get_quota(level) excludes them. Pair this with " + "sparse_backing_medium() to count what sparse buffers are pinning.") + .def("sparse_backing_medium", &kv::KvCacheManager::sparseBackingMedium, + nb::call_guard(), + "The medium a sparse buffer's pages are, or None when no life cycle is sparse.\n\n" + "The hot level's second medium: host memory the GPU reads through a device mapping, " + "which is what sparse loading pulls into the per-request GPU cache.") .def( "resize", [](kv::KvCacheManager& self, int cacheLevel, size_t quota, bool bestEfforts) @@ -2918,16 +3574,24 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) .def_prop_ro("enable_partial_match", &kv::KvCacheManager::enablePartialMatch) .def_prop_ro("enable_partial_commit", &kv::KvCacheManager::enablePartialCommit) .def_prop_ro("enable_swa_scratch_reuse", &kv::KvCacheManager::isSwaScratchReuseEnabled) + .def("max_scratch_slots", &kv::KvCacheManager::maxScratchSlots, nb::arg("max_blocks_per_seq"), + "The most scratch slots one request can need, over every life cycle and buffer, given the " + "widest page table a row will get. Zero where nothing can hold scratch.\n\n" + "Sizes scratch staging. A caller that stages scratch slots itself asks this rather than " + "deriving it: which blocks can be stale turns on the window, the sinks and the rewind " + "margin, and a second copy of that reasoning is a second thing to keep right.") .def( "supports_index_mode", - [](kv::KvCacheManager const& self, kv::PageIndexMode mode) -> nb::object + [](kv::KvCacheManager const& self, kv::PageIndexMode mode, int layerGroupId) -> nb::object { - auto result = self.supportsIndexMode(mode); + auto result = self.supportsIndexMode(mode, kv::LayerGroupId{layerGroupId}); if (!result.has_value()) return nb::none(); return nb::cast(*result); }, - nb::arg("mode")) + nb::arg("mode"), nb::arg("layer_group_id"), + "None means it depends on the request. The answer is definitive for every life cycle " + "that cannot hold scratch slots.") .def_prop_ro("commit_min_snapshot", &kv::KvCacheManager::commitMinSnapshot) .def_prop_ro("num_layers", &kv::KvCacheManager::numLayers) .def_prop_ro("layer_ids", &kv::KvCacheManager::layerIds) @@ -2941,6 +3605,31 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) "get_layer_group_id", [](kv::KvCacheManager const& self, kv::LayerId layerId) { return self.getLayerGroupId(layerId).value(); }, nb::arg("layer_id")) + .def( + "get_sparse_cache_max_units", + [](kv::KvCacheManager const& self, kv::LayerId layerId) + { return self.sparseCacheMaxUnits(self.getLayerGroupId(layerId)); }, + nb::arg("layer_id"), + "Cells a request's sparse cache holds per buffer of this layer, in selection units, " + "counting both the cached history and everything in flight. The lease follows from " + "this alone. Only one unit is ever in flight, so the surplus above the selection " + "width is retention rather than headroom.") + .def( + "get_sparse_cache_units_per_page", + [](kv::KvCacheManager const& self, kv::LayerId layerId) + { return self.sparseCacheUnitsPerPage(self.getLayerGroupId(layerId)); }, + nb::arg("layer_id"), + "Selection units one page of this layer's buffers holds; one when a unit is a whole page.\n\n" + "Distinct from PageIndexConverter.expansion, which is how much a buffer's page index space " + "is stretched by tokens_per_block_override. A sparse buffer does not stretch its index " + "space, so that stays 1 while this counts the cache's own granularity.") + .def( + "get_sparse_cache_buffer_count", + [](kv::KvCacheManager const& self, kv::LayerId layerId) + { return self.sparseCacheBufferCount(self.getLayerGroupId(layerId)); }, + nb::arg("layer_id"), + "Sparse buffers of this layer's life cycle, one per (layer, role), each with its " + "own cells. A lease is this many runs of get_sparse_cache_max_units cells.") .def("get_page_index_converter", &kv::KvCacheManager::getPageIndexConverter, nb::arg("layer_id"), nb::arg("data_role")) .def( @@ -2971,18 +3660,15 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) .def( "__call__", [](kv::PageIndexConverter const& self, std::vector const& baseIndices, - std::optional indexMode, nb::object scratchObj) -> std::vector + std::optional indexMode, + std::optional> scratchRange, + std::vector const& scratchSlots) -> std::vector { - kv::ScratchDesc const* scratch = nullptr; - std::optional scratchHolder; - if (!scratchObj.is_none()) - { - scratchHolder = nb::cast(scratchObj); - scratch = &*scratchHolder; - } - return self(baseIndices, indexMode, scratch); + return self(baseIndices, indexMode, scratchRange.value_or(kv::HalfOpenRange{}), + kv::Span{scratchSlots.data(), static_cast(scratchSlots.size())}); }, - nb::arg("base_indices"), nb::arg("index_mode") = nb::none(), nb::arg("scratch") = nb::none()); + nb::arg("base_indices"), nb::arg("index_mode") = nb::none(), nb::arg("scratch_range") = nb::none(), + nb::arg("scratch_slots") = std::vector{}); m.def( "gen_multimodal_cache_key_tokens", diff --git a/cpp/tests/unit_tests/batch_manager/CMakeLists.txt b/cpp/tests/unit_tests/batch_manager/CMakeLists.txt index 0f2879b1dd77..1f2b1c77bad7 100644 --- a/cpp/tests/unit_tests/batch_manager/CMakeLists.txt +++ b/cpp/tests/unit_tests/batch_manager/CMakeLists.txt @@ -288,6 +288,18 @@ add_tllm_gtest( LABELS "unit;batch-manager;gpu") +add_tllm_gtest( + kvCacheManagerV2GeneratorTest + OWNER + batch-manager + SOURCES + kvCacheManagerV2GeneratorTest.cpp + LINK_LIBRARIES + tllm::kv_cache + $ + LABELS + "unit;batch-manager;host") + add_tllm_gtest( kvCacheManagerV2DigestPoolTest OWNER diff --git a/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2ColdPageTest.cpp b/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2ColdPageTest.cpp index 011f8b9dec98..496d2c413acf 100644 --- a/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2ColdPageTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2ColdPageTest.cpp @@ -16,14 +16,11 @@ */ #include "kvCacheManagerV2TestUtils.h" -#include "tensorrt_llm/batch_manager/kv_cache_manager_v2/batch.h" #include "tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.h" #include "tensorrt_llm/batch_manager/kv_cache_manager_v2/config.h" -#include "tensorrt_llm/batch_manager/kv_cache_manager_v2/eventManager.h" #include "tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.h" #include "tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.h" #include "tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.h" -#include "tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/cudaEvent.h" #include "tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/funcGuard.h" #include "tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/hostMem.h" #include "tensorrt_llm/common/tllmException.h" @@ -34,8 +31,8 @@ #include #include #include -#include #include +#include #include #include #include @@ -242,14 +239,14 @@ TEST_P(OptionalDiskCheckpointTest, PayloadRestoreRetryAndIndependentEviction) ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); auto reservations = storage.newSlotsForPoolGroup(kHotLevel, hotGroup, storage.getStatistics(kHotLevel, hotGroup).free); - EXPECT_FALSE(reader->resume(stream, std::nullopt, std::vector{encoder, peer})); + EXPECT_FALSE(reader->resume(stream, std::vector{encoder, peer})); EXPECT_FALSE(reader->isActive()); EXPECT_EQ(reader->numCommittedTokens(), 8); for (auto& slot : reservations) { storage.releaseSlot(encoder, kHotLevel, std::move(slot)); } - ASSERT_TRUE(reader->resume(stream, std::nullopt, std::vector{encoder, peer})); + ASSERT_TRUE(reader->resume(stream, std::vector{encoder, peer})); ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); EXPECT_EQ(block->getPage(encoder), encoderPage); EXPECT_EQ(block->getPage(peer), peerPage); @@ -308,12 +305,38 @@ KVCacheManagerConfig makeGpuTieredConfig() return config; } +//! One sparse attention layer, over a hot tier and a cold tier. +//! +//! The hot pool group is host-backed because the layer is sparse; the cold tier is an +//! ordinary codec target, which is the whole point of the arrangement. +KVCacheManagerConfig makeSparseTieredConfig() +{ + KVCacheManagerConfig config; + config.tokensPerBlock = 4; + config.cacheTiers.emplace_back(GpuCacheTierConfig{4 << 20}); + config.cacheTiers.emplace_back(HostCacheTierConfig{4 << 20}); + AttentionLayerConfig layer; + layer.layerId = 0; + BufferConfig buffer{"key", 2 << 20, std::nullopt}; + buffer.isSparse = true; + layer.buffers.push_back(std::move(buffer)); + layer.sparseCache = SparseCacheConfig{/*numSelectedBlocks=*/1, /*maxCachedUnits=*/2, /*unitTokens=*/4}; + // A pool group sizes its slots from its device-backed pools, so a sparse layer carries + // one -- which is also the shape a model uses. + layer.buffers.push_back(BufferConfig{"index", 256, std::nullopt}); + config.layers.emplace_back(std::move(layer)); + BatchDesc constraint; + constraint.kvCaches.assign(2, KVCacheDesc{/*capacity=*/16, /*historyLength=*/0}); + config.constraints.push_back(std::move(constraint)); + return config; +} + KVCacheManagerConfig makeSplitColdGroupingConfig() { KVCacheManagerConfig config; config.tokensPerBlock = 4; config.cacheTiers.emplace_back(GpuCacheTierConfig{4 << 20}); - // The host tier is distributed in units of HostMem::kAlignment, so a quota + // The host tier is distributed in units of HostMem::sizeGranularity(), so a quota // of only a few multiples of it cannot express a pool ratio: every group // rounds to the same single unit. Sized here to leave enough units for the // ratios these tests assert on. @@ -335,7 +358,7 @@ KVCacheManagerConfig makeSplitColdGroupingConfig() //! A config whose cold tier holds exactly one page. //! -//! The host tier is allocated in whole multiples of HostMem::kAlignment, so a +//! The host tier is allocated in whole multiples of HostMem::sizeGranularity(), so a //! pool contains a single slot only when one page fills that unit. Tests that //! need two pages to contend for one cold slot size the page accordingly rather //! than shrinking the quota, which cannot go below one unit. @@ -344,12 +367,12 @@ KVCacheManagerConfig makeSingleColdSlotConfig() KVCacheManagerConfig config = makeSplitColdGroupingConfig(); for (auto& layer : config.layers) { - std::get(layer).buffers[0].size = HostMem::kAlignment; + std::get(layer).buffers[0].size = HostMem::sizeGranularity(); } // The hot tier holds pages of the same size, so it is sized in the same // unit to leave room for the slots these tests allocate up front. - config.cacheTiers[0] = GpuCacheTierConfig{8 * HostMem::kAlignment}; - config.cacheTiers[1] = HostCacheTierConfig{HostMem::kAlignment}; + config.cacheTiers[0] = GpuCacheTierConfig{8 * HostMem::sizeGranularity()}; + config.cacheTiers[1] = HostCacheTierConfig{HostMem::sizeGranularity()}; return config; } @@ -504,20 +527,15 @@ class MixedIndexLocationColdPageCodec final : public IKvCacheColdPageCodec } }; -class AsyncRejectingColdPageCodec final : public IKvCacheColdPageCodec +//! Records which life cycles were encoded and decoded, without interfering. +//! +//! Compression lives behind this interface, so "can sparse history be compressed" reduces +//! to "does a sparse life cycle's page reach the codec at all". Before sparse history moved +//! to the hot tier it could not: the codec runs on the hot/cold boundary, and that history +//! was locked on the cold side and never crossed. +class RecordingColdPageCodec final : public IKvCacheColdPageCodec { public: - enum class Operation - { - kEncode, - kDecode, - }; - - explicit AsyncRejectingColdPageCodec(Operation operation) - : mOperation(operation) - { - } - bool configure(PoolGroupDesc const*, PoolGroupIndex) noexcept override { return true; @@ -533,120 +551,77 @@ class AsyncRejectingColdPageCodec final : public IKvCacheColdPageCodec return PageIndexLocation::kHost; } - bool encode(LayerGroupId, void*, PageIndexPair const*, size_t, cudaStream_t stream) noexcept override + bool encode(LayerGroupId lg, void*, PageIndexPair const*, size_t pages, cudaStream_t) noexcept override { - return mOperation == Operation::kEncode ? reject(stream) : true; + mEncoded[lg.value()] += pages; + return true; } - bool decode(LayerGroupId, void const*, PageIndexPair const*, size_t, cudaStream_t stream) noexcept override + bool decode(LayerGroupId lg, void const*, PageIndexPair const*, size_t pages, cudaStream_t) noexcept override { - return mOperation == Operation::kDecode ? reject(stream) : true; + mDecoded[lg.value()] += pages; + return true; } - bool launched() const noexcept + size_t encoded(LayerGroupId lg) const { - return mLaunched.load(std::memory_order_acquire); + auto const it = mEncoded.find(lg.value()); + return it == mEncoded.end() ? 0 : it->second; } - void release() noexcept + size_t decoded(LayerGroupId lg) const { - mRelease.store(true, std::memory_order_release); + auto const it = mDecoded.find(lg.value()); + return it == mDecoded.end() ? 0 : it->second; } private: - static void CUDART_CB waitForRelease(void* data) - { - auto& codec = *static_cast(data); - while (!codec.mRelease.load(std::memory_order_acquire)) - { - std::this_thread::yield(); - } - } - - bool reject(cudaStream_t stream) noexcept - { - bool const launched = cudaLaunchHostFunc(stream, waitForRelease, this) == cudaSuccess; - mLaunched.store(launched, std::memory_order_release); - return false; - } - - Operation mOperation; - std::atomic mLaunched{false}; - std::atomic mRelease{false}; + std::map mEncoded; + std::map mDecoded; }; -class ObservingColdPageCodec final : public IKvCacheColdPageCodec +class AsyncRejectingColdPageCodec final : public IKvCacheColdPageCodec { public: - bool configure(PoolGroupDesc const* descriptors, PoolGroupIndex count) noexcept override + enum class Operation { - return mCodec->configure(descriptors, count); - } + kEncode, + kDecode, + }; - size_t queryColdPageBytes(LayerGroupId layerGroup) const noexcept override + explicit AsyncRejectingColdPageCodec(Operation operation) + : mOperation(operation) { - return mCodec->queryColdPageBytes(layerGroup); } - LayerGroupId getBatchingLayerGroupId(LayerGroupId layerGroup) const noexcept override + bool configure(PoolGroupDesc const*, PoolGroupIndex) noexcept override { - return mCodec->getBatchingLayerGroupId(layerGroup); + return true; } - PageIndexLocation queryPageIndexLocation(LayerGroupId layerGroup) const noexcept override + size_t queryColdPageBytes(LayerGroupId) const noexcept override { - return mCodec->queryPageIndexLocation(layerGroup); + return 2 << 20; } - bool encode(LayerGroupId layerGroup, void* destination, PageIndexPair const* indices, size_t count, - cudaStream_t stream) noexcept override + PageIndexLocation queryPageIndexLocation(LayerGroupId) const noexcept override { - ++encodeCalls; - encodedPages += count; - encodeStream = stream; - bool const submitted = mCodec->encode(layerGroup, destination, indices, count, stream); - return submitted && encodeCalls != rejectEncodeCall; + return PageIndexLocation::kHost; } - bool decode(LayerGroupId layerGroup, void const* source, PageIndexPair const* indices, size_t count, - cudaStream_t stream) noexcept override + bool encode(LayerGroupId, void*, PageIndexPair const*, size_t, cudaStream_t stream) noexcept override { - ++decodeCalls; - decodedPages += count; - bool const submitted = mCodec->decode(layerGroup, source, indices, count, stream); - return submitted && decodeCalls != rejectDecodeCall; + return mOperation == Operation::kEncode ? reject(stream) : true; } - size_t encodeCalls = 0; - size_t encodedPages = 0; - size_t rejectEncodeCall = 0; - size_t decodeCalls = 0; - size_t decodedPages = 0; - size_t rejectDecodeCall = 0; - cudaStream_t encodeStream{}; - -private: - std::unique_ptr mCodec = createDefaultKvCacheColdPageCodec(); -}; - -class StreamGate -{ -public: - ~StreamGate() + bool decode(LayerGroupId, void const*, PageIndexPair const*, size_t, cudaStream_t stream) noexcept override { - release(); - if (mStream) - { - cudaStreamSynchronize(mStream); - } + return mOperation == Operation::kDecode ? reject(stream) : true; } - cudaError_t enqueue(cudaStream_t stream) + bool launched() const noexcept { - // Event merging must not create streams while a gate is held: stream creation can wait for host callbacks. - CudaStreamPool::instance(); - mStream = stream; - return cudaLaunchHostFunc(stream, wait, this); + return mLaunched.load(std::memory_order_acquire); } void release() noexcept @@ -655,17 +630,25 @@ class StreamGate } private: - static void CUDART_CB wait(void* data) + static void CUDART_CB waitForRelease(void* data) { - auto& gate = *static_cast(data); - while (!gate.mRelease.load(std::memory_order_acquire)) + auto& codec = *static_cast(data); + while (!codec.mRelease.load(std::memory_order_acquire)) { std::this_thread::yield(); } } + bool reject(cudaStream_t stream) noexcept + { + bool const launched = cudaLaunchHostFunc(stream, waitForRelease, this) == cudaSuccess; + mLaunched.store(launched, std::memory_order_release); + return false; + } + + Operation mOperation; + std::atomic mLaunched{false}; std::atomic mRelease{false}; - cudaStream_t mStream{}; }; SharedPtr makeCommittedPage(KvCacheManager& manager, StorageManager& storage, CacheLevel level, @@ -1105,9 +1088,9 @@ TEST(KvCacheManagerV2ColdPageTest, MaxQuotaBelowTheStructuralMinimumIsRejected) { ASSERT_EQ(cudaSetDevice(0), cudaSuccess); // Two cold pool groups, from the two distinct sliding windows. - constexpr size_t kOneUnit = HostMem::kAlignment; + size_t const kOneUnit = HostMem::sizeGranularity(); - auto makeConfig = [](std::optional maxQuota) + auto makeConfig = [kOneUnit](std::optional maxQuota) { auto config = makeSplitColdGroupingConfig(); config.cacheTiers[1] = HostCacheTierConfig{kOneUnit, maxQuota}; @@ -1156,7 +1139,7 @@ TEST(KvCacheManagerV2ColdPageTest, RecursiveFallenPageMergeResortsByPriority) { ASSERT_EQ(cudaSetDevice(0), cudaSuccess); auto config = makeSingleColdSlotConfig(); - config.cacheTiers.emplace_back(HostCacheTierConfig{HostMem::kAlignment}); + config.cacheTiers.emplace_back(HostCacheTierConfig{HostMem::sizeGranularity()}); auto manager = std::make_shared(std::move(config)); auto& storage = manager->storage(); CacheLevel const firstColdLevel{1}; @@ -1185,2592 +1168,50 @@ TEST(KvCacheManagerV2ColdPageTest, RecursiveFallenPageMergeResortsByPriority) EXPECT_EQ(lowestPriorityPage->cacheLevel, kHotLevel); } -class KvCacheManagerV2PageLockTest : public ::testing::Test -{ -protected: - void SetUp() override - { - ASSERT_EQ(cudaSetDevice(0), cudaSuccess); - ASSERT_EQ(cudaStreamCreateWithFlags(&mStream, cudaStreamNonBlocking), cudaSuccess); - } - - void TearDown() override - { - EXPECT_EQ(cudaStreamDestroy(mStream), cudaSuccess); - } - - static KVCacheManagerConfig sparseConfig() - { - auto config = makeTieredConfig(); - std::get(config.layers.front()).buffers.front().isSparse = true; - return config; - } - - static SharedPtr seedPrefix( - KvCacheManager& manager, CacheLevel level, LifeCycleId lc = LifeCycleId{0}) - { - auto& storage = manager.storage(); - TypedVec counts(storage.numLifeCycles(), 0); - counts[lc] = 1; - auto slots = storage.newSlots(level, counts); - return makeCommittedPage(manager, storage, level, slots[lc].front(), lc); - } - - static SharedPtr pageAt(KvCache const& cache, int ordinal = 0, LifeCycleId lc = LifeCycleId{0}) - { - return blockPageGetPage(cache.blocks().at(BlockOrdinal{ordinal}).pages.at(kDefaultBeamIndex).at(lc)); - } - - CUstream stream() const - { - return reinterpret_cast(mStream); - } - - TokenSpan tokens(int length = 4) const - { - return {mTokens.data(), length}; - } - - std::vector const mTokens{TokenIdExt{0}, TokenIdExt{1}, TokenIdExt{2}, TokenIdExt{3}}; - cudaStream_t mStream{}; -}; - -class KvCacheManagerV2SparseOffloadTest : public KvCacheManagerV2PageLockTest -{ -}; - -class KvCacheManagerV2PageStorageTest : public KvCacheManagerV2PageLockTest -{ -}; - -class KvCacheManagerV2BatchTest : public KvCacheManagerV2PageLockTest -{ -protected: - CudaStream batchStream() const - { - return reinterpret_cast(stream()); - } - - std::vector read(MemAddress address, size_t size) - { - cuCheck(cuStreamSynchronize(stream())); - std::vector result(size); - cuCheck(cuMemcpyDtoH(result.data(), address, size * sizeof(int32_t))); - return result; - } -}; - -TEST_F(KvCacheManagerV2BatchTest, PublishesRawMixedTierRowsAndSkipsUnchangedRows) -{ - auto config = makeSplitColdGroupingConfig(); - auto& sparse = std::get(config.layers.front()); - sparse.buffers.front().isSparse = true; - sparse.buffers.push_back({.role = "value", .size = 2048, .tokensPerBlockOverride = 2, .isSparse = true}); - auto coalesced = sparse; - coalesced.layerId = 2; - config.layers.push_back(std::move(coalesced)); - auto manager = std::make_shared(std::move(config)); - Batch batch(manager, 3, 4); - auto cache = manager->createKvCache(); - auto closeCache = FuncGuard([&]() { cache->close(); }); - ASSERT_TRUE(cache->resume(stream())); - ASSERT_TRUE(cache->resize(12, 6)); - EXPECT_EQ(batch.add(*cache, 2), 2); - auto const sparseGroup = manager->getLayerGroupId(0); - auto const denseGroup = manager->getLayerGroupId(1); - auto const address = batch.pageTableAddress(sparseGroup); - EXPECT_EQ(batch.publish(batchStream()), (std::vector{0, 1, 2})); - EXPECT_EQ(read(batch.numBlocksAddress(sparseGroup), 3), (std::vector{0, 0, 0})); - EXPECT_EQ(batch.publish(batchStream()), (std::vector{})); - EXPECT_FALSE(cache->pageStorageDirty()); - - ASSERT_TRUE(cache->enterDecode()); - EXPECT_EQ(batch.dirtyRows(), (std::vector{2})); - EXPECT_THROW(batch.waitReady(batchStream()), LogicError); - EXPECT_EQ(batch.publish(batchStream()), (std::vector{2})); - for (auto const group : {sparseGroup, denseGroup}) - { - auto expected = std::vector(12, kBadPageIndex.value()); - auto const snapshot = cache->getPageStorageSnapshot(group); - std::copy(snapshot.basePageIndices().begin(), snapshot.basePageIndices().end(), expected.begin() + 8); - EXPECT_EQ(read(batch.pageTableAddress(group), 12), expected); - } - EXPECT_EQ(read(batch.numBlocksAddress(sparseGroup), 3), (std::vector{0, 0, 1})); - EXPECT_EQ(read(batch.numBlocksAddress(denseGroup), 3), (std::vector{0, 0, 0})); - EXPECT_EQ(batch.pageTableAddress(sparseGroup), address); - - cache->suspend(); - EXPECT_EQ(batch.publish(batchStream()), (std::vector{2})); - EXPECT_EQ(read(address, 12), (std::vector(12, kBadPageIndex.value()))); - EXPECT_EQ(read(batch.numBlocksAddress(sparseGroup), 3), (std::vector{0, 0, 0})); - ASSERT_TRUE(cache->resume()); - EXPECT_EQ(batch.publish(batchStream()), (std::vector{2})); - EXPECT_EQ(read(batch.numBlocksAddress(sparseGroup), 3).back(), 1); -} - -TEST_F(KvCacheManagerV2BatchTest, PublishesIncrementalOffloadAndClearsReusedRows) -{ - auto config = makeSplitColdGroupingConfig(); - auto& sparse = std::get(config.layers.front()); - sparse.buffers.front().isSparse = true; - sparse.buffers.push_back({.role = "value", .size = 2048, .tokensPerBlockOverride = 2, .isSparse = true}); - auto coalesced = sparse; - coalesced.layerId = 2; - config.layers.push_back(std::move(coalesced)); - auto codec = std::make_unique(); - auto* observer = codec.get(); - auto manager = std::make_shared(std::move(config), nullptr, std::move(codec)); - auto const apiLock = manager->lockExclusive(); - auto& storage = manager->storage(); - auto const sparseGroup = manager->getLayerGroupId(0); - auto const denseGroup = manager->getLayerGroupId(1); - auto const gpuPool = storage.getPoolGroupIndex(kHotLevel, sparseGroup); - auto const hostPool = storage.getPoolGroupIndex(kSparseHistoryLevel, sparseGroup); - auto const& sizes = storage.slotSize(kHotLevel, gpuPool); - ASSERT_EQ(sizes.size(), PoolIndex{1}); - size_t const bytes = sizes[PoolIndex{0}]; - Batch batch(manager, 2, 4); - auto const tableAddress = batch.pageTableAddress(sparseGroup); - auto cache = manager->createKvCache(); - auto closeCache = FuncGuard([&]() { cache->close(); }); - ASSERT_TRUE(cache->resume(stream())); - ASSERT_TRUE(cache->resize(12, 6)); - for (int ordinal = 0; ordinal < 3; ++ordinal) - { - auto const page = pageAt(*cache, ordinal, sparseGroup); - auto const address - = std::get(storage.slotAddress(kHotLevel, gpuPool, page->slotId(), PoolIndex{0})); - ASSERT_EQ(cudaMemsetAsync(reinterpret_cast(address), 0x40 + ordinal, bytes, mStream), cudaSuccess); - } - batch.add(*cache, 1); - auto checkPublished = [&](int eligible) - { - for (auto const group : {sparseGroup, denseGroup}) - { - auto expected = std::vector(8, kBadPageIndex.value()); - for (int ordinal = 0; ordinal < 3; ++ordinal) - { - auto const page = pageAt(*cache, ordinal, group); - expected[4 + ordinal] = slotIdToPageIndexValue(page->slotId()); - EXPECT_EQ( - page->cacheLevel, group == sparseGroup && ordinal < eligible ? kSparseHistoryLevel : kHotLevel); - } - EXPECT_EQ(read(batch.pageTableAddress(group), 8), expected); - EXPECT_EQ( - read(batch.numBlocksAddress(group), 2), (std::vector{0, group == sparseGroup ? eligible : 0})); - } - for (int ordinal = 0; ordinal < eligible; ++ordinal) - { - auto const page = pageAt(*cache, ordinal, sparseGroup); - auto const address = std::get( - storage.slotAddress(kSparseHistoryLevel, hostPool, page->slotId(), PoolIndex{0})); - auto const* data = reinterpret_cast(address); - EXPECT_TRUE(std::all_of(data, data + bytes, [ordinal](uint8_t value) { return value == 0x40 + ordinal; })); - } - }; - batch.publish(batchStream()); - checkPublished(0); - auto const freeGpuBefore = manager->getStorageStatistics(kHotLevel)[gpuPool].free; - ASSERT_TRUE(cache->enterDecode()); - batch.publish(batchStream()); - checkPublished(1); - auto const firstHostSlot = pageAt(*cache, 0, sparseGroup)->slotId(); - auto const version = cache->pageStorageVersion(); - EXPECT_EQ(batch.resize({std::nullopt, 12}, {std::nullopt, 7}, batchStream()), - (std::vector>{std::nullopt, true})); - EXPECT_EQ(cache->pageStorageVersion(), version); - EXPECT_EQ(observer->encodedPages, 1); - checkPublished(1); - EXPECT_EQ(batch.resize({std::nullopt, 12}, {std::nullopt, 8}, batchStream()), - (std::vector>{std::nullopt, true})); - checkPublished(2); - EXPECT_EQ(observer->encodedPages, 2); - EXPECT_EQ(pageAt(*cache, 0, sparseGroup)->slotId(), firstHostSlot); - EXPECT_EQ(manager->getStorageStatistics(kHotLevel)[gpuPool].free, freeGpuBefore + 2); - EXPECT_EQ(batch.publish(batchStream()), (std::vector{})); - - cache->close(); - auto replacement = manager->createKvCache(); - auto closeReplacement = FuncGuard([&]() { replacement->close(); }); - ASSERT_TRUE(replacement->resume(stream())); - ASSERT_TRUE(replacement->resize(4, 0)); - batch.add(*replacement, 1); - EXPECT_EQ(batch.publish(batchStream()), (std::vector{1})); - auto expected = std::vector(8, kBadPageIndex.value()); - expected[4] = slotIdToPageIndexValue(pageAt(*replacement, 0, sparseGroup)->slotId()); - EXPECT_EQ(read(tableAddress, 8), expected); - EXPECT_EQ(read(batch.numBlocksAddress(sparseGroup), 2), (std::vector{0, 0})); - EXPECT_EQ(batch.pageTableAddress(sparseGroup), tableAddress); -} - -TEST_F(KvCacheManagerV2BatchTest, MembershipCloseAndPublicationFailuresPreserveRows) -{ - auto manager = std::make_shared(sparseConfig()); - Batch batch(manager, 3, 1); - Batch other(manager, 3, 1); - auto first = manager->createKvCache(); - auto second = manager->createKvCache(); - auto closeCaches = FuncGuard( - [&]() - { - first->close(); - if (second) - { - second->close(); - } - }); - EXPECT_EQ(batch.add(*first, 2), 2); - EXPECT_EQ(batch.add(*first), 2); - EXPECT_EQ(batch.add(*second), 0); - EXPECT_THROW(other.add(*first), LogicError); - EXPECT_THROW(first->bindPageStorageRow(1), LogicError); - EXPECT_THROW(batch.add(*second, 2), LogicError); - batch.publish(batchStream()); - first->close(); - EXPECT_EQ(batch.dirtyRows(), (std::vector{2})); - EXPECT_EQ(second->pageStorageRow(), 0); - batch.remove(*second); - EXPECT_EQ(other.add(*second, 1), 1); - other.close(); - EXPECT_FALSE(second->pageStorageRow().has_value()); - EXPECT_FALSE(second->isClosed()); - EXPECT_EQ(batch.add(*second, 1), 1); - ASSERT_TRUE(second->resume(stream())); - ASSERT_TRUE(second->resize(8, 0)); - EXPECT_THROW(batch.publish(batchStream()), std::invalid_argument); - EXPECT_THROW(batch.waitReady(batchStream()), LogicError); - EXPECT_TRUE(second->pageStorageDirty()); - EXPECT_FALSE(batch.dirtyRows().empty()); - ASSERT_TRUE(second->resize(4, 0)); - batch.publish(batchStream()); - EXPECT_FALSE(second->pageStorageDirty()); - second.reset(); - EXPECT_EQ(batch.dirtyRows(), (std::vector{1})); - batch.publish(batchStream()); - EXPECT_EQ(read(batch.pageTableAddress(LifeCycleId{0}), 3), (std::vector(3, -1))); - closeCaches.cancel(); -} - -TEST_F(KvCacheManagerV2BatchTest, BatchedResizePublishesFinalStatesAfterPartialOom) -{ - auto config = sparseConfig(); - config.cacheTiers[0] = GpuCacheTierConfig{8 << 20}; - auto manager = std::make_shared(std::move(config)); - Batch batch(manager, 3, 4); - auto first = manager->createKvCache(); - auto second = manager->createKvCache(); - auto closeCaches = FuncGuard( - [&]() - { - first->close(); - second->close(); - }); - ASSERT_TRUE(first->resume(stream())); - ASSERT_TRUE(second->resume(stream())); - ASSERT_TRUE(first->resize(4, 0)); - ASSERT_TRUE(second->resize(4, 0)); - batch.add(*first, 0); - batch.add(*second, 2); - auto const results = batch.resize({8, std::nullopt, 12}, {0, std::nullopt, 0}, batchStream()); - EXPECT_EQ(results, (std::vector>{true, std::nullopt, false})); - EXPECT_EQ(first->capacity(), 8); - EXPECT_EQ(second->capacity(), 4); - EXPECT_TRUE(batch.dirtyRows().empty()); - auto expected = std::vector(12, -1); - auto const firstIndices = first->getBasePageIndices(LifeCycleId{0}); - auto const secondIndices = second->getBasePageIndices(LifeCycleId{0}); - std::copy(firstIndices.data(), firstIndices.data() + 2, expected.begin()); - expected[8] = secondIndices[0]; - EXPECT_EQ(read(batch.pageTableAddress(LifeCycleId{0}), 12), expected); -} - -TEST_F(KvCacheManagerV2BatchTest, SharedOffloadInvalidatesEveryOwner) -{ - auto manager = std::make_shared(sparseConfig()); - auto const apiLock = manager->lockExclusive(); - auto page = seedPrefix(*manager, kHotLevel); - auto first = manager->createKvCache({}, tokens()); - auto second = manager->createKvCache({}, tokens()); - auto closeCaches = FuncGuard( - [&]() - { - first->close(); - second->close(); - }); - ASSERT_TRUE(first->resume(stream())); - ASSERT_TRUE(second->resume(stream())); - Batch batch(manager, 2, 1); - batch.add(*first); - batch.add(*second); - batch.publish(batchStream()); - first->offloadSparsePages({page}); - EXPECT_EQ(batch.dirtyRows(), (std::vector{0, 1})); - ASSERT_TRUE(first->enterDecode()); - ASSERT_TRUE(second->enterDecode()); - EXPECT_EQ(batch.publish(batchStream()), (std::vector{0, 1})); - EXPECT_EQ(read(batch.pageTableAddress(LifeCycleId{0}), 2), - (std::vector(2, slotIdToPageIndexValue(page->slotId())))); - EXPECT_EQ(read(batch.numBlocksAddress(LifeCycleId{0}), 2), (std::vector{1, 1})); -} - -TEST_F(KvCacheManagerV2BatchTest, DeferredHistoryGapClosesWithoutAdvancingWatermark) -{ - for (bool const closeOwner : {false, true}) - { - SCOPED_TRACE(closeOwner); - auto config = sparseConfig(); - config.cacheTiers[0] = GpuCacheTierConfig{8 << 20}; - auto codec = std::make_unique(); - auto* observer = codec.get(); - auto manager = std::make_shared(std::move(config), nullptr, std::move(codec)); - auto const apiLock = manager->lockExclusive(); - auto page = seedPrefix(*manager, kHotLevel); - auto decoder = manager->createKvCache({}, tokens()); - auto prefill = manager->createKvCache({}, tokens()); - auto closeCaches = FuncGuard( - [&]() - { - decoder->close(); - prefill->close(); - }); - ASSERT_TRUE(decoder->resume(stream())); - ASSERT_TRUE(prefill->resume(stream())); - ASSERT_TRUE(decoder->resize(12, 8)); - Batch batch(manager, 1, 3); - batch.add(*decoder); - auto const group = manager->getLayerGroupId(0); - ASSERT_TRUE(decoder->enterDecode()); - EXPECT_EQ(page->cacheLevel, kHotLevel); - EXPECT_EQ(pageAt(*decoder, 1)->cacheLevel, kSparseHistoryLevel); - EXPECT_EQ(pageAt(*decoder, 2)->cacheLevel, kHotLevel); - auto const snapshot = decoder->getPageStorageSnapshot(group); - EXPECT_EQ(snapshot.eligibleHistoryBlocks(), 0); - EXPECT_EQ(snapshot.cacheLevels(), - (std::vector>{kHotLevel, kSparseHistoryLevel, kHotLevel})); - EXPECT_EQ(batch.publish(batchStream()), (std::vector{0})); - EXPECT_EQ(read(batch.pageTableAddress(group), 3), snapshot.basePageIndices()); - EXPECT_EQ(read(batch.numBlocksAddress(group), 1), (std::vector{0})); - EXPECT_EQ(observer->encodedPages, 1); - auto const hostSlot = pageAt(*decoder, 1)->slotId(); - auto const version = decoder->pageStorageVersion(); - - if (closeOwner) - { - prefill->close(); - } - else - { - prefill->suspend(); - } - EXPECT_EQ(page->cacheLevel, kHotLevel); - EXPECT_EQ(decoder->pageStorageVersion(), version); - EXPECT_TRUE(batch.dirtyRows().empty()); - EXPECT_EQ(batch.publish(batchStream()), (std::vector{0})); - EXPECT_EQ(decoder->historyLength(), 8); - EXPECT_GT(decoder->pageStorageVersion(), version); - EXPECT_EQ(page->cacheLevel, kSparseHistoryLevel); - EXPECT_EQ(pageAt(*decoder, 1)->slotId(), hostSlot); - EXPECT_EQ(observer->encodedPages, 2); - auto const finalSnapshot = decoder->getPageStorageSnapshot(group); - EXPECT_EQ(finalSnapshot.eligibleHistoryBlocks(), 2); - EXPECT_EQ(read(batch.numBlocksAddress(group), 1), (std::vector{2})); - EXPECT_EQ(read(batch.pageTableAddress(group), 3), finalSnapshot.basePageIndices()); - EXPECT_EQ(batch.publish(batchStream()), (std::vector{})); - EXPECT_EQ(observer->encodeCalls, 2); - } -} - -TEST_F(KvCacheManagerV2BatchTest, DeferredOffloadRetainsWorkAfterHostOomAndCodecRejection) +TEST(KvCacheManagerV2ColdPageTest, SparseHistoryPassesThroughTheCodec) { - auto codec = std::make_unique(); - auto* observer = codec.get(); - auto manager = std::make_shared(sparseConfig(), nullptr, std::move(codec)); - auto const apiLock = manager->lockExclusive(); + ASSERT_EQ(cudaSetDevice(0), cudaSuccess); + auto codec = std::make_unique(); + auto* codecPtr = codec.get(); + auto manager = std::make_shared(makeSparseTieredConfig(), nullptr, std::move(codec)); auto& storage = manager->storage(); - auto page = seedPrefix(*manager, kHotLevel); - auto decoder = manager->createKvCache({}, tokens()); - auto prefill = manager->createKvCache({}, tokens()); - auto closeCaches = FuncGuard( - [&]() - { - decoder->close(); - prefill->close(); - }); - ASSERT_TRUE(decoder->resume(stream())); - ASSERT_TRUE(prefill->resume(stream())); - auto const freeHost = storage.getStatistics(kSparseHistoryLevel).free; - auto blockers = storage.newSlots(kSparseHistoryLevel, TypedVec{freeHost}); - auto releaseBlockers = FuncGuard( - [&]() - { - for (auto& slot : blockers[LifeCycleId{0}]) - { - storage.releaseSlot(LifeCycleId{0}, kSparseHistoryLevel, std::move(slot)); - } - }); - ASSERT_TRUE(decoder->enterDecode()); - Batch batch(manager, 1, 1); - batch.add(*decoder); - auto const group = manager->getLayerGroupId(0); - batch.publish(batchStream()); - auto const gpuSlot = page->slotId(); - prefill->close(); - EXPECT_THROW(batch.publish(batchStream()), OutOfPagesError); - EXPECT_EQ(observer->encodeCalls, 0); - EXPECT_EQ(page->slotId(), gpuSlot); - EXPECT_EQ(page->cacheLevel, kHotLevel); - EXPECT_EQ(read(batch.numBlocksAddress(group), 1), (std::vector{0})); - releaseBlockers.run(); - - observer->rejectEncodeCall = 1; - EXPECT_THROW(batch.publish(batchStream()), TllmException); - EXPECT_EQ(page->cacheLevel, kHotLevel); - EXPECT_EQ(page->slotId(), gpuSlot); - EXPECT_TRUE(decoder->isDecoding()); - EXPECT_EQ(decoder->historyLength(), 4); - EXPECT_EQ(decoder->getPageStorageSnapshot(group).eligibleHistoryBlocks(), 0); - EXPECT_EQ(read(batch.pageTableAddress(group), 1), (std::vector{slotIdToPageIndexValue(gpuSlot)})); - EXPECT_EQ(read(batch.numBlocksAddress(group), 1), (std::vector{0})); - observer->rejectEncodeCall = 0; - EXPECT_EQ(batch.publish(batchStream()), (std::vector{0})); - EXPECT_EQ(page->cacheLevel, kSparseHistoryLevel); - EXPECT_EQ(read(batch.numBlocksAddress(group), 1), (std::vector{1})); - EXPECT_EQ(storage.getStatistics(kSparseHistoryLevel).free, freeHost - 1); -} - -TEST_F(KvCacheManagerV2BatchTest, RetainsStagingUntilUploadAndOrdersTableReuseAfterReaders) -{ - auto config = sparseConfig(); - std::get(config.layers.front()).buffers.front().size = 4096; - auto manager = std::make_shared(std::move(config)); - Batch batch(manager, 1, 2); - auto cache = manager->createKvCache(); - auto closeCache = FuncGuard([&]() { cache->close(); }); - ASSERT_TRUE(cache->resume(stream())); - ASSERT_TRUE(cache->resize(4, 0)); - batch.add(*cache); - batch.publish(batchStream()); - auto const expected = read(batch.pageTableAddress(LifeCycleId{0}), 2); - HostMem readback(2 * sizeof(int32_t)); - cudaStream_t readerStream{}; - ASSERT_EQ(cudaStreamCreateWithFlags(&readerStream, cudaStreamNonBlocking), cudaSuccess); - auto destroyReader = FuncGuard([&]() { cudaStreamDestroy(readerStream); }); - StreamGate uploadGate; - StreamGate readerGate; - auto releaseGates = FuncGuard( - [&]() - { - uploadGate.release(); - readerGate.release(); - }); - ASSERT_EQ(uploadGate.enqueue(mStream), cudaSuccess); - EXPECT_THROW(cache->bindPageStorageRow(std::nullopt), LogicError); - cache->setBasePageIndexBuf(kDefaultBeamIndex, LifeCycleId{0}, nullptr, 0); - // Commit changes lock ownership/readiness even when the numeric slot stays unchanged. - cache->commit(tokens()); - batch.publish(batchStream()); - batch.waitReady(reinterpret_cast(readerStream)); - ASSERT_EQ(readerGate.enqueue(readerStream), cudaSuccess); - ASSERT_EQ(cudaMemcpyAsync(reinterpret_cast(readback.address()), - reinterpret_cast(batch.pageTableAddress(LifeCycleId{0})), 2 * sizeof(int32_t), - cudaMemcpyDeviceToHost, readerStream), - cudaSuccess); - batch.recordRead(reinterpret_cast(readerStream)); - cache->close(); - batch.publish(batchStream()); - CachedCudaEvent cleared(batchStream()); - EXPECT_FALSE(cleared.queryComplete()); - uploadGate.release(); - EXPECT_FALSE(cleared.queryComplete()); - readerGate.release(); - cleared.synchronize(); - auto const* oldIndices = reinterpret_cast(readback.address()); - EXPECT_EQ(std::vector(oldIndices, oldIndices + 2), expected); - EXPECT_EQ(read(batch.pageTableAddress(LifeCycleId{0}), 2), (std::vector{-1, -1})); -} + LifeCycleId const lifeCycle{0}; + LayerGroupId const layerGroup{lifeCycle.value()}; + CacheLevel const coldLevel{1}; + TypedVec oneSlot(LifeCycleId{1}, 1); -TEST_F(KvCacheManagerV2BatchTest, DeviceAddressesSurviveGraphReplayAcrossPublication) -{ - auto manager = std::make_shared(sparseConfig()); - Batch batch(manager, 1, 2); - auto cache = manager->createKvCache(); - auto closeCache = FuncGuard([&]() { cache->close(); }); - ASSERT_TRUE(cache->resume(stream())); - ASSERT_TRUE(cache->resize(4, 0)); - batch.add(*cache); - batch.publish(batchStream()); - batch.waitReady(batchStream()); - auto const table = batch.pageTableAddress(LifeCycleId{0}); - void* output = nullptr; - ASSERT_EQ(cudaMalloc(&output, 2 * sizeof(int32_t)), cudaSuccess); - auto freeOutput = FuncGuard([&]() { cudaFree(output); }); - cudaGraph_t graph{}; - cudaGraphExec_t executable{}; - auto destroyGraph = FuncGuard( - [&]() - { - if (executable) - { - cudaGraphExecDestroy(executable); - } - if (graph) - { - cudaGraphDestroy(graph); - } - }); - ASSERT_EQ(cudaStreamBeginCapture(mStream, cudaStreamCaptureModeThreadLocal), cudaSuccess); - EXPECT_THROW(batch.publish(batchStream()), LogicError); - ASSERT_EQ(cudaMemcpyAsync( - output, reinterpret_cast(table), 2 * sizeof(int32_t), cudaMemcpyDeviceToDevice, mStream), - cudaSuccess); - ASSERT_EQ(cudaStreamEndCapture(mStream, &graph), cudaSuccess); - ASSERT_EQ(cudaGraphInstantiateWithFlags(&executable, graph, 0), cudaSuccess); - for (bool active : {true, false, true}) + // The arrangement this test is about: the sparse buffer's *hot* pages are host memory, + // sitting in a GPU-tier group beside a device-backed pool. Without this the test would + // pass for a dense life cycle too, and prove nothing. + PoolGroupIndex const hotPg = storage.getPoolGroupIndex(kHotLevel, lifeCycle); + ASSERT_EQ(storage.cacheTier(kHotLevel), CacheTier::GPU_MEM); + auto const hotPools = storage.slotSize(kHotLevel, hotPg); + ASSERT_EQ(hotPools.size(), PoolIndex{2}) << "a sparse pool beside a device-backed one"; + int sparsePools = 0; + for (PoolIndex pool{0}; pool < hotPools.size(); ++pool) { - if (active && !cache->isActive()) - { - ASSERT_TRUE(cache->resume()); - } - if (!active) - { - cache->suspend(); - } - batch.publish(batchStream()); - batch.waitReady(batchStream()); - ASSERT_EQ(cudaGraphLaunch(executable, mStream), cudaSuccess); - batch.recordRead(batchStream()); - auto expected = std::vector{-1, -1}; - if (active) - { - expected[0] = cache->getBasePageIndices(LifeCycleId{0})[0]; - } - EXPECT_EQ(read(reinterpret_cast(output), 2), expected); - EXPECT_EQ(batch.pageTableAddress(LifeCycleId{0}), table); - } -} - -TEST_F(KvCacheManagerV2BatchTest, PublicationWaitsForOffloadAndReaderFencesProtectHostSlots) -{ - for (bool commit : {false, true}) - { - SCOPED_TRACE(commit); - auto manager = std::make_shared(sparseConfig()); - Batch batch(manager, 1, 1); - auto const apiLock = manager->lockExclusive(); - auto& storage = manager->storage(); - auto const lc = LifeCycleId{0}; - auto const gpuPool = storage.getPoolGroupIndex(kHotLevel, lc); - auto const hostPool = storage.getPoolGroupIndex(kSparseHistoryLevel, lc); - size_t const bytes = storage.slotSize(kHotLevel, gpuPool)[PoolIndex{0}]; - cudaStream_t readerStream{}; - ASSERT_EQ(cudaStreamCreateWithFlags(&readerStream, cudaStreamNonBlocking), cudaSuccess); - auto destroyReader = FuncGuard([&]() { cudaStreamDestroy(readerStream); }); - void* gpuReadback = nullptr; - void* hostReadback = nullptr; - ASSERT_EQ(cudaMalloc(&gpuReadback, bytes), cudaSuccess); - auto freeGpuReadback = FuncGuard([&]() { cudaFree(gpuReadback); }); - ASSERT_EQ(cudaMallocHost(&hostReadback, bytes), cudaSuccess); - auto freeHostReadback = FuncGuard([&]() { cudaFreeHost(hostReadback); }); - auto cache = manager->createKvCache(); - auto closeCache = FuncGuard([&]() { cache->close(); }); - ASSERT_TRUE(cache->resume(stream())); - ASSERT_TRUE(cache->resize(4, 4)); - auto const gpuSlot = pageAt(*cache)->slotId(); - auto const gpuAddress = std::get(storage.slotAddress(kHotLevel, gpuPool, gpuSlot, PoolIndex{0})); - constexpr uint8_t kPattern = 0xD3; - ASSERT_EQ(cudaMemsetAsync(reinterpret_cast(gpuAddress), kPattern, bytes, mStream), cudaSuccess); - auto blocker = storage.newSlots(kSparseHistoryLevel, TypedVec{1}); - auto releaseBlocker - = FuncGuard([&]() { storage.releaseSlot(lc, kSparseHistoryLevel, std::move(blocker[lc].front())); }); - // Warm transfer staging before deliberately delaying the copy. - storage.copySlotData(lc, kSparseHistoryLevel, kHotLevel, blocker[lc].front().slotId(), gpuSlot, stream()); - batch.add(*cache); - batch.publish(batchStream()); - ASSERT_EQ(cudaStreamSynchronize(mStream), cudaSuccess); - StreamGate copyGate; - StreamGate readerGate; - auto releaseGates = FuncGuard( - [&]() - { - copyGate.release(); - readerGate.release(); - }); - ASSERT_EQ(copyGate.enqueue(mStream), cudaSuccess); - ASSERT_TRUE(cache->enterDecode()); - auto const snapshot = cache->getPageStorageSnapshot(lc); - ASSERT_EQ(snapshot.eligibleHistoryBlocks(), 1); - ASSERT_EQ(snapshot.readyEvents().size(), 1); - EXPECT_FALSE(snapshot.readyEvents().front().queryComplete()); - batch.publish(reinterpret_cast(readerStream)); - batch.waitReady(reinterpret_cast(readerStream)); - CachedCudaEvent published(reinterpret_cast(readerStream)); - EXPECT_FALSE(published.queryComplete()); - ASSERT_EQ(readerGate.enqueue(readerStream), cudaSuccess); - auto const hostSlot = pageAt(*cache)->slotId(); - auto const hostAddress - = std::get(storage.slotAddress(kSparseHistoryLevel, hostPool, hostSlot, PoolIndex{0})); - ASSERT_EQ(cudaMemcpyAsync(gpuReadback, reinterpret_cast(hostAddress), bytes, - cudaMemcpyHostToDevice, readerStream), - cudaSuccess); - ASSERT_EQ(cudaMemcpyAsync(hostReadback, gpuReadback, bytes, cudaMemcpyDeviceToHost, readerStream), cudaSuccess); - batch.recordRead(reinterpret_cast(readerStream)); - if (commit) - cache->commit(tokens()); - cache->close(); - batch.publish(batchStream()); - manager->clearReusableBlocks(); - auto recycled = storage.newSlots(kSparseHistoryLevel, TypedVec{1}); - auto releaseRecycled - = FuncGuard([&]() { storage.releaseSlot(lc, kSparseHistoryLevel, std::move(recycled[lc].front())); }); - EXPECT_EQ(recycled[lc].front().slotId(), hostSlot); - EXPECT_FALSE(recycled[lc].front().queryReady()); - copyGate.release(); - published.synchronize(); - EXPECT_FALSE(recycled[lc].front().queryReady()); - readerGate.release(); - recycled[lc].front().readyEvent.synchronize(); - EXPECT_EQ(read(batch.pageTableAddress(lc), 1), (std::vector{-1})); - EXPECT_EQ(read(batch.numBlocksAddress(lc), 1), (std::vector{0})); - auto const* readBytes = static_cast(hostReadback); - EXPECT_TRUE(std::all_of(readBytes, readBytes + bytes, [](uint8_t v) { return v == kPattern; })); + // A pool whose medium differs from its level's tier is the host-backed one. + sparsePools += storage.poolMedium(kHotLevel, hotPg, pool) != storage.cacheTier(kHotLevel) ? 1 : 0; } -} - -TEST_F(KvCacheManagerV2PageStorageTest, QueriesSparseBuffersAndRejectsUnknownBuffers) -{ - auto config = makeSplitColdGroupingConfig(); - auto& sparse = std::get(config.layers[0]); - sparse.buffers.front().isSparse = true; - sparse.buffers.push_back({.role = "value", .size = 2048, .tokensPerBlockOverride = 2, .isSparse = true}); - config.layers.emplace_back(SsmLayerConfig{.layerId = 2, .buffers = {{"state", 4096}}}); - config.commitMinSnapshot = true; - auto manager = std::make_shared(std::move(config)); - EXPECT_TRUE(manager->isSparse(0, "key")); - EXPECT_TRUE(manager->isSparse(0, "value")); - EXPECT_FALSE(manager->isSparse(1, "key")); - EXPECT_FALSE(manager->isSparse(2, "state")); - EXPECT_THROW(manager->isSparse(0, "missing"), std::out_of_range); - EXPECT_THROW(manager->isSparse(3, "key"), std::out_of_range); -} - -TEST_F(KvCacheManagerV2PageStorageTest, SnapshotsRawMixedTierIndicesAndDecodeEligibility) -{ - auto config = makeSplitColdGroupingConfig(); - auto& sparse = std::get(config.layers[0]); - sparse.buffers.front().isSparse = true; - sparse.buffers.push_back({.role = "value", .size = 2048, .tokensPerBlockOverride = 2, .isSparse = true}); - auto coalesced = sparse; - coalesced.layerId = 2; - config.layers.emplace_back(std::move(coalesced)); - auto manager = std::make_shared(std::move(config)); - auto const apiLock = manager->lockExclusive(); - auto const sparseGroup = manager->getLayerGroupId(0); - auto const denseGroup = manager->getLayerGroupId(1); - EXPECT_EQ(manager->getLayerGroupId(2), sparseGroup); - EXPECT_GT(manager->getPageIndexScale(0, "value"), 1); - auto cache = manager->createKvCache(); - std::vector externalIndices(8, kBadPageIndex.value()); - auto closeCache = FuncGuard([&]() { cache->close(); }); - ASSERT_TRUE(cache->resume(stream())); - ASSERT_TRUE(cache->resize(12, 6)); - cache->setBasePageIndexBuf(kDefaultBeamIndex, sparseGroup, externalIndices.data(), externalIndices.size()); - auto const prefill = cache->getPageStorageSnapshot(sparseGroup); - EXPECT_EQ(prefill.eligibleHistoryBlocks(), 0); - EXPECT_EQ(prefill.cacheLevels(), (std::vector>(3, kHotLevel))); - - ASSERT_TRUE(cache->enterDecode()); - auto const decode = cache->getPageStorageSnapshot(sparseGroup); - EXPECT_EQ(decode.eligibleHistoryBlocks(), 1); - ASSERT_EQ(decode.basePageIndices().size(), 3); - EXPECT_EQ(decode.basePageIndices(), (std::vector(externalIndices.begin(), externalIndices.begin() + 3))); - EXPECT_EQ( - decode.cacheLevels(), (std::vector>{kSparseHistoryLevel, kHotLevel, kHotLevel})); - EXPECT_EQ(cache->getPageStorageSnapshot(denseGroup).eligibleHistoryBlocks(), 0); - EXPECT_EQ(prefill.cacheLevels()[0], kHotLevel); - EXPECT_GT(decode.version(), prefill.version()); - for (int ord = 0; ord < 3; ++ord) - EXPECT_EQ(decode.basePageIndices()[ord], slotIdToPageIndexValue(pageAt(*cache, ord, sparseGroup)->slotId())); - - cache->suspend(); - auto const suspended = cache->getPageStorageSnapshot(sparseGroup); - EXPECT_EQ(suspended.eligibleHistoryBlocks(), 0); - EXPECT_EQ(suspended.basePageIndices(), (std::vector(3, kBadPageIndex.value()))); - EXPECT_EQ(suspended.cacheLevels(), (std::vector>(3, std::nullopt))); - EXPECT_TRUE(suspended.readyEvents().empty()); - ASSERT_TRUE(cache->resume()); - EXPECT_EQ(cache->getPageStorageSnapshot(sparseGroup).eligibleHistoryBlocks(), 1); - ASSERT_TRUE(cache->resize(12, 8)); - EXPECT_EQ(cache->getPageStorageSnapshot(sparseGroup).eligibleHistoryBlocks(), 2); - EXPECT_EQ(pageAt(*cache, 2, sparseGroup)->cacheLevel, kHotLevel); -} + ASSERT_EQ(sparsePools, 1) << "exactly one of the group's pools holds host-backed pages"; -TEST_F(KvCacheManagerV2PageStorageTest, SnapshotsSkipSharedPromptBeamRows) -{ - auto config = sparseConfig(); - config.enablePartialCommit = false; - std::get(config.layers.front()).buffers.front().size = 4096; - auto manager = std::make_shared(std::move(config)); - auto const apiLock = manager->lockExclusive(); - auto cache = manager->createKvCache({}, {}, std::nullopt, {}, /*expectedPromptLength=*/8); - auto closeCache = FuncGuard([&]() { cache->close(); }); - ASSERT_TRUE(cache->resume(stream())); - ASSERT_TRUE(cache->resize(12, 8)); - cache->setBeamWidth(BeamIndex{3}); - - for (BeamIndex bi{0}; bi < cache->beamWidth(); ++bi) - { - auto const snapshot = cache->getPageStorageSnapshot(LifeCycleId{0}, bi); - ASSERT_EQ(snapshot.basePageIndices().size(), 3U); - ASSERT_EQ(snapshot.cacheLevels().size(), 3U); - EXPECT_EQ(snapshot.eligibleHistoryBlocks(), 0); - EXPECT_EQ(snapshot.readyEvents().size(), bi == kDefaultBeamIndex ? 3U : 1U); - for (int ordinal = 0; ordinal < 3; ++ordinal) - { - if (bi != kDefaultBeamIndex && ordinal < 2) - { - EXPECT_EQ(snapshot.basePageIndices()[ordinal], kBadPageIndex.value()); - EXPECT_EQ(snapshot.cacheLevels()[ordinal], std::nullopt); - } - else - { - auto const& page = blockPageGetPage(cache->blocks()[BlockOrdinal{ordinal}].pages[bi][LifeCycleId{0}]); - ASSERT_NE(page, nullptr); - EXPECT_EQ(snapshot.basePageIndices()[ordinal], slotIdToPageIndexValue(page->slotId())); - EXPECT_EQ(snapshot.cacheLevels()[ordinal], kHotLevel); - } - } - } -} + // A hot page of the sparse life cycle. Host-backed, but an ordinary hot page: locked at + // the hot level and migratable like any other. + auto hotSlots = storage.newSlots(kHotLevel, oneSlot); + auto page = makeCommittedPage(*manager, storage, kHotLevel, hotSlots[lifeCycle].front()); + storage.excludeFromEviction(*page); + EXPECT_EQ(codecPtr->encoded(layerGroup), 0U); -TEST_F(KvCacheManagerV2PageStorageTest, DirtyAcknowledgmentTracksBindingsCommitAndRequestLifetime) -{ - auto manager = std::make_shared(sparseConfig()); - auto cache = manager->createKvCache(); - std::vector externalIndices(2, kBadPageIndex.value()); - auto closeCache = FuncGuard([&]() { cache->close(); }); - EXPECT_TRUE(cache->pageStorageDirty()); - EXPECT_FALSE(cache->pageStorageRow().has_value()); - EXPECT_THROW(cache->bindPageStorageRow(-1), LogicError); - cache->bindPageStorageRow(7); - ASSERT_TRUE(cache->resume(stream())); - ASSERT_TRUE(cache->resize(8, 4)); - auto const prefill = cache->getPageStorageSnapshot(LifeCycleId{0}); - EXPECT_EQ(prefill.row(), 7); - ASSERT_TRUE(cache->acknowledgePageStorage(prefill.version())); - EXPECT_FALSE(cache->pageStorageDirty()); - ASSERT_TRUE(cache->resize(8, 5)); - EXPECT_FALSE(cache->pageStorageDirty()); - ASSERT_TRUE(cache->enterDecode()); - EXPECT_TRUE(cache->pageStorageDirty()); - EXPECT_FALSE(cache->acknowledgePageStorage(prefill.version())); - auto const decode = cache->getPageStorageSnapshot(LifeCycleId{0}); - ASSERT_TRUE(cache->acknowledgePageStorage(decode.version())); - ASSERT_TRUE(cache->enterDecode()); - EXPECT_FALSE(cache->pageStorageDirty()); - - cache->bindPageStorageRow(7); - EXPECT_TRUE(cache->pageStorageDirty()); - EXPECT_FALSE(cache->acknowledgePageStorage(decode.version())); - cache->bindPageStorageRow(9); - cache->setBasePageIndexBuf(kDefaultBeamIndex, LifeCycleId{0}, externalIndices.data(), externalIndices.size()); - auto const rebound = cache->getPageStorageSnapshot(LifeCycleId{0}); - EXPECT_EQ(rebound.row(), 9); - EXPECT_EQ(rebound.basePageIndices(), decode.basePageIndices()); - ASSERT_TRUE(cache->acknowledgePageStorage(rebound.version())); - cache->setBasePageIndexBuf(kDefaultBeamIndex, LifeCycleId{0}, nullptr, 0); - EXPECT_TRUE(cache->pageStorageDirty()); - - auto const beforeCommit = cache->getPageStorageSnapshot(LifeCycleId{0}); - ASSERT_TRUE(cache->acknowledgePageStorage(beforeCommit.version())); - cache->commit(tokens()); - EXPECT_TRUE(cache->pageStorageDirty()); - auto const committed = cache->getPageStorageSnapshot(LifeCycleId{0}); - EXPECT_EQ(committed.basePageIndices(), beforeCommit.basePageIndices()); - EXPECT_EQ(committed.cacheLevels(), beforeCommit.cacheLevels()); - ASSERT_TRUE(cache->acknowledgePageStorage(committed.version())); - cache->suspend(); - EXPECT_TRUE(cache->pageStorageDirty()); - EXPECT_EQ(cache->pageStorageRow(), 9); - ASSERT_TRUE(cache->resume()); - EXPECT_EQ(cache->getPageStorageSnapshot(LifeCycleId{0}).eligibleHistoryBlocks(), 1); - cache->close(); - EXPECT_TRUE(cache->pageStorageDirty()); - EXPECT_FALSE(cache->pageStorageRow().has_value()); - auto const closed = cache->getPageStorageSnapshot(LifeCycleId{0}); - EXPECT_TRUE(closed.basePageIndices().empty()); - EXPECT_EQ(closed.eligibleHistoryBlocks(), 0); - EXPECT_THROW(cache->bindPageStorageRow(9), LogicError); - EXPECT_THROW(cache->recordPageStorageRead(reinterpret_cast(mStream)), LogicError); - - auto reused = manager->createKvCache({}, tokens()); - auto closeReused = FuncGuard([&]() { reused->close(); }); - EXPECT_TRUE(reused->pageStorageDirty()); - EXPECT_FALSE(reused->pageStorageRow().has_value()); - EXPECT_EQ(reused->getPageStorageSnapshot(LifeCycleId{0}).basePageIndices(), (std::vector{-1})); - ASSERT_TRUE(reused->resume(stream(), true)); - EXPECT_EQ(reused->getPageStorageSnapshot(LifeCycleId{0}).eligibleHistoryBlocks(), 1); -} + // Out of the hot tier: crossing the boundary is what invokes the codec. + storage.batchedMigrate(coldLevel, {page}, {}); + EXPECT_EQ(page->cacheLevel, coldLevel); + EXPECT_EQ(codecPtr->encoded(layerGroup), 1U) << "sparse history must reach the codec on demotion"; -TEST_F(KvCacheManagerV2PageStorageTest, ReadinessAndReaderFencesSurviveCommitAndClose) -{ - for (bool commit : {false, true}) - { - SCOPED_TRACE(commit); - auto manager = std::make_shared(sparseConfig()); - auto const apiLock = manager->lockExclusive(); - auto& storage = manager->storage(); - auto const lc = LifeCycleId{0}; - auto const gpuPool = storage.getPoolGroupIndex(kHotLevel, lc); - auto const hostPool = storage.getPoolGroupIndex(kSparseHistoryLevel, lc); - size_t const bytes = storage.slotSize(kHotLevel, gpuPool)[PoolIndex{0}]; - cudaStream_t readerStream{}; - ASSERT_EQ(cudaStreamCreateWithFlags(&readerStream, cudaStreamNonBlocking), cudaSuccess); - auto destroyReader = FuncGuard([&]() { cudaStreamDestroy(readerStream); }); - void* gpuReadback = nullptr; - void* hostReadback = nullptr; - ASSERT_EQ(cudaMalloc(&gpuReadback, bytes), cudaSuccess); - auto freeGpuReadback = FuncGuard([&]() { cudaFree(gpuReadback); }); - ASSERT_EQ(cudaMallocHost(&hostReadback, bytes), cudaSuccess); - auto freeHostReadback = FuncGuard([&]() { cudaFreeHost(hostReadback); }); - auto cache = manager->createKvCache(); - auto closeCache = FuncGuard([&]() { cache->close(); }); - ASSERT_TRUE(cache->resume(stream())); - ASSERT_TRUE(cache->resize(4, 4)); - auto const gpuSlot = pageAt(*cache)->slotId(); - auto const gpuAddress = std::get(storage.slotAddress(kHotLevel, gpuPool, gpuSlot, PoolIndex{0})); - constexpr uint8_t kPattern = 0xD3; - ASSERT_EQ(cudaMemsetAsync(reinterpret_cast(gpuAddress), kPattern, bytes, mStream), cudaSuccess); - auto blocker = storage.newSlots(kSparseHistoryLevel, TypedVec{1}); - auto releaseBlocker - = FuncGuard([&]() { storage.releaseSlot(lc, kSparseHistoryLevel, std::move(blocker[lc].front())); }); - // Warm transfer staging before deliberately delaying the copy. - storage.copySlotData(lc, kSparseHistoryLevel, kHotLevel, blocker[lc].front().slotId(), gpuSlot, stream()); - ASSERT_EQ(cudaStreamSynchronize(mStream), cudaSuccess); - StreamGate copyGate; - StreamGate readerGate; - auto releaseGates = FuncGuard( - [&]() - { - copyGate.release(); - readerGate.release(); - }); - ASSERT_EQ(copyGate.enqueue(mStream), cudaSuccess); - ASSERT_TRUE(cache->enterDecode()); - auto const snapshot = cache->getPageStorageSnapshot(lc); - ASSERT_EQ(snapshot.eligibleHistoryBlocks(), 1); - ASSERT_EQ(snapshot.readyEvents().size(), 1); - EXPECT_FALSE(snapshot.readyEvents().front().queryComplete()); - snapshot.waitReady(reinterpret_cast(readerStream)); - ASSERT_EQ(readerGate.enqueue(readerStream), cudaSuccess); - auto const hostSlot = pageAt(*cache)->slotId(); - auto const hostAddress - = std::get(storage.slotAddress(kSparseHistoryLevel, hostPool, hostSlot, PoolIndex{0})); - ASSERT_EQ(cudaMemcpyAsync(gpuReadback, reinterpret_cast(hostAddress), bytes, - cudaMemcpyHostToDevice, readerStream), - cudaSuccess); - ASSERT_EQ(cudaMemcpyAsync(hostReadback, gpuReadback, bytes, cudaMemcpyDeviceToHost, readerStream), cudaSuccess); - cache->recordPageStorageRead(reinterpret_cast(readerStream)); - if (commit) - cache->commit(tokens()); - cache->close(); - manager->clearReusableBlocks(); - auto recycled = storage.newSlots(kSparseHistoryLevel, TypedVec{1}); - auto releaseRecycled - = FuncGuard([&]() { storage.releaseSlot(lc, kSparseHistoryLevel, std::move(recycled[lc].front())); }); - EXPECT_EQ(recycled[lc].front().slotId(), hostSlot); - EXPECT_FALSE(recycled[lc].front().queryReady()); - copyGate.release(); - snapshot.readyEvents().front().synchronize(); - EXPECT_FALSE(recycled[lc].front().queryReady()); - readerGate.release(); - recycled[lc].front().readyEvent.synchronize(); - auto const* readBytes = static_cast(hostReadback); - EXPECT_TRUE(std::all_of(readBytes, readBytes + bytes, [](uint8_t v) { return v == kPattern; })); - } -} - -TEST_F(KvCacheManagerV2SparseOffloadTest, BatchesCompleteCoalescedPagesAndCountsPhysicalCopies) -{ - auto config = makeSplitColdGroupingConfig(); - config.enableStats = true; - for (auto& layerConfig : config.layers) - { - auto& layer = std::get(layerConfig); - layer.buffers.front().isSparse = true; - layer.buffers.push_back({.role = "value", .size = 2048, .tokensPerBlockOverride = 2, .isSparse = true}); - layer.buffers.push_back({.role = "scale", .size = 128, .isSparse = true}); - } - auto codec = std::make_unique(); - auto* observer = codec.get(); - auto manager = std::make_shared(std::move(config), nullptr, std::move(codec)); - auto const apiLock = manager->lockExclusive(); - auto& storage = manager->storage(); - ASSERT_EQ(storage.numLifeCycles(), LifeCycleId{2}); - ASSERT_EQ(storage.numPoolGroups(kHotLevel), PoolGroupIndex{1}); - ASSERT_GT(storage.numPools(PoolGroupIndex{0}), PoolIndex{1}); - auto cache = manager->createKvCache(); - auto closeCache = FuncGuard([&]() { cache->close(); }); - ASSERT_TRUE(cache->resume(stream())); - ASSERT_TRUE(cache->resize(12, 8)); - - std::vector> pages; - std::vector> expected; - for (int ordinal = 0; ordinal < 2; ++ordinal) - { - for (LifeCycleId lc{0}; lc < storage.numLifeCycles(); ++lc) - { - auto page = pageAt(*cache, ordinal, lc); - auto const pg = storage.getPoolGroupIndex(kHotLevel, lc); - auto const& sizes = storage.slotSize(kHotLevel, pg); - std::vector bytes; - for (PoolIndex pool{0}; pool < sizes.size(); ++pool) - { - auto const pattern = static_cast(17 * pages.size() + pool.value() + 1); - auto const address = std::get(storage.slotAddress(kHotLevel, pg, page->slotId(), pool)); - ASSERT_EQ( - cudaMemsetAsync(reinterpret_cast(address), pattern, sizes[pool], mStream), cudaSuccess); - bytes.insert(bytes.end(), sizes[pool], pattern); - } - pages.push_back(std::move(page)); - expected.push_back(std::move(bytes)); - } - } - auto const gpuFree = storage.getStatistics(kHotLevel).free; - auto const hostFree = storage.getStatistics(kSparseHistoryLevel).free; - auto targets = pages; - targets.push_back(pages.front()); - auto const versionBeforeOffload = cache->pageStorageVersion(); - cache->offloadSparsePages(targets); - EXPECT_EQ(observer->encodeCalls, 1); - EXPECT_EQ(observer->encodedPages, pages.size()); - EXPECT_EQ(observer->encodeStream, mStream); - EXPECT_EQ(storage.getStatistics(kHotLevel).free, gpuFree + pages.size()); - EXPECT_EQ(storage.getStatistics(kSparseHistoryLevel).free, hostFree - pages.size()); - EXPECT_GT(cache->pageStorageVersion(), versionBeforeOffload); - - for (size_t i = 0; i < pages.size(); ++i) - { - auto const& page = pages[i]; - EXPECT_EQ(page->cacheLevel, kSparseHistoryLevel); - EXPECT_EQ(page->status(), PageStatus::LOCKED); - EXPECT_FALSE(page->scheduledForEviction()); - page->readyEvent.synchronize(); - auto const pg = storage.getPoolGroupIndex(kSparseHistoryLevel, page->lifeCycle); - auto const address - = std::get(storage.slotAddress(kSparseHistoryLevel, pg, page->slotId(), PoolIndex{0})); - EXPECT_EQ(std::memcmp(reinterpret_cast(address), expected[i].data(), expected[i].size()), 0); - } - for (LifeCycleId lc{0}; lc < storage.numLifeCycles(); ++lc) - { - EXPECT_EQ(pageAt(*cache, 2, lc)->cacheLevel, kHotLevel); - auto const indices = cache->getBasePageIndices(lc); - for (int ordinal = 0; ordinal < 3; ++ordinal) - { - EXPECT_EQ(indices[ordinal], slotIdToPageIndexValue(pageAt(*cache, ordinal, lc)->slotId())); - } - } - auto const stats = manager->getAndResetIterationStats(); - for (LifeCycleId lc{0}; lc < storage.numLifeCycles(); ++lc) - { - EXPECT_EQ(stats.at(lc).iterOffloadBlocks, 2); - EXPECT_EQ(stats.at(lc).iterOffloadBytes, 2 * expected.front().size()); - } - auto const versionAfterOffload = cache->pageStorageVersion(); - cache->offloadSparsePages(targets); - EXPECT_EQ(observer->encodeCalls, 1); - EXPECT_EQ(cache->pageStorageVersion(), versionAfterOffload); - EXPECT_TRUE(manager->getAndResetIterationStats().empty()); -} - -TEST_F(KvCacheManagerV2SparseOffloadTest, SharedOwnersPublishHostIndicesAndKeepHistoryPinned) -{ - auto manager = std::make_shared(sparseConfig()); - auto const apiLock = manager->lockExclusive(); - auto& storage = manager->storage(); - auto page = seedPrefix(*manager, kHotLevel); - auto first = manager->createKvCache({}, tokens()); - auto second = manager->createKvCache({}, tokens()); - std::vector externalIndices(1, kBadPageIndex.value()); - auto closeCaches = FuncGuard( - [&]() - { - first->close(); - second->close(); - }); - ASSERT_TRUE(first->resume(stream())); - ASSERT_TRUE(second->resume(stream())); - second->setBasePageIndexBuf(kDefaultBeamIndex, LifeCycleId{0}, externalIndices.data(), externalIndices.size()); - auto const gpuSlot = page->slotId(); - auto hostBlockers = storage.newSlots(kSparseHistoryLevel, TypedVec{1}); - auto releaseBlockers = FuncGuard([&]() - { storage.releaseSlot(LifeCycleId{0}, kSparseHistoryLevel, std::move(hostBlockers[LifeCycleId{0}].front())); }); - - EXPECT_THROW(storage.batchedMigrate(kSparseHistoryLevel, {page}, {}), LogicError); - auto const firstVersion = first->pageStorageVersion(); - auto const secondVersion = second->pageStorageVersion(); - ASSERT_TRUE(first->acknowledgePageStorage(firstVersion)); - ASSERT_TRUE(second->acknowledgePageStorage(secondVersion)); - first->offloadSparsePages({page, page}); - EXPECT_NE(page->slotId(), gpuSlot); - EXPECT_EQ(first->getBasePageIndices(LifeCycleId{0})[0], slotIdToPageIndexValue(page->slotId())); - EXPECT_EQ(externalIndices[0], slotIdToPageIndexValue(page->slotId())); - EXPECT_GT(first->pageStorageVersion(), firstVersion); - EXPECT_GT(second->pageStorageVersion(), secondVersion); - EXPECT_TRUE(first->pageStorageDirty()); - EXPECT_TRUE(second->pageStorageDirty()); - auto const snapshot = second->getPageStorageSnapshot(LifeCycleId{0}); - EXPECT_EQ(snapshot.basePageIndices(), externalIndices); - EXPECT_EQ(snapshot.cacheLevels()[0], kSparseHistoryLevel); - EXPECT_EQ(snapshot.eligibleHistoryBlocks(), 0); - EXPECT_EQ(storage.getStatistics(kHotLevel).free, storage.getStatistics(kHotLevel).total); - EXPECT_FALSE(storage.isEvictable(*page)); - EXPECT_THROW(storage.batchedMigrate(kHotLevel, {page}, {}), LogicError); - ASSERT_TRUE(first->enterDecode()); - ASSERT_TRUE(second->enterDecode()); - first->suspend(); - EXPECT_EQ(page->status(), PageStatus::LOCKED); - ASSERT_TRUE(first->resume()); - EXPECT_EQ(pageAt(*first), page); - second->close(); - EXPECT_EQ(externalIndices[0], kBadPageIndex.value()); - EXPECT_EQ(page->status(), PageStatus::LOCKED); - EXPECT_EQ(page->cacheLevel, kSparseHistoryLevel); -} - -TEST_F(KvCacheManagerV2SparseOffloadTest, WaitsForLiveAndFinishedReadersBeforeRecyclingGpuSlot) -{ - for (auto const [deferred, finishReader] : - {std::pair{false, false}, std::pair{false, true}, std::pair{true, false}, std::pair{true, true}}) - { - SCOPED_TRACE(deferred); - SCOPED_TRACE(finishReader); - auto manager = std::make_shared(sparseConfig()); - auto const apiLock = manager->lockExclusive(); - auto& storage = manager->storage(); - cudaStream_t readerStream{}; - ASSERT_EQ(cudaStreamCreateWithFlags(&readerStream, cudaStreamNonBlocking), cudaSuccess); - auto destroyReaderStream = FuncGuard([&]() { cudaStreamDestroy(readerStream); }); - auto page = seedPrefix(*manager, kHotLevel); - auto first = manager->createKvCache({}, tokens()); - auto second = manager->createKvCache({}, tokens()); - auto closeCaches = FuncGuard( - [&]() - { - first->close(); - second->close(); - }); - ASSERT_TRUE(first->resume(stream())); - ASSERT_TRUE(second->resume(reinterpret_cast(readerStream))); - auto const lc = page->lifeCycle; - auto const pg = storage.getPoolGroupIndex(kHotLevel, lc); - size_t const bytes = storage.slotSize(kHotLevel, pg)[PoolIndex{0}]; - auto const gpuSlot = page->slotId(); - auto const gpuAddress = std::get(storage.slotAddress(kHotLevel, pg, gpuSlot, PoolIndex{0})); - constexpr uint8_t kPattern = 0xA6; - ASSERT_EQ(cudaMemsetAsync(reinterpret_cast(gpuAddress), kPattern, bytes, mStream), cudaSuccess); - auto gpuBlocker = storage.newGpuSlots(TypedVec{1}); - auto hostScratch = storage.newSlots(kSparseHistoryLevel, TypedVec{1}); - auto releaseSlots = FuncGuard( - [&]() - { - storage.releaseSlot(lc, kHotLevel, std::move(gpuBlocker[lc].front())); - storage.releaseSlot(lc, kSparseHistoryLevel, std::move(hostScratch[lc].front())); - }); - // Warm the codec's descriptor/index staging before deliberately blocking a stream. - storage.copySlotData(lc, kSparseHistoryLevel, kHotLevel, hostScratch[lc].front().slotId(), gpuSlot, stream()); - ASSERT_EQ(cudaStreamSynchronize(mStream), cudaSuccess); - auto const readback = std::get(storage.slotAddress(kSparseHistoryLevel, - storage.getPoolGroupIndex(kSparseHistoryLevel, lc), hostScratch[lc].front().slotId(), PoolIndex{0})); - if (deferred) - { - ASSERT_TRUE(first->enterDecode()); - EXPECT_EQ(page->cacheLevel, kHotLevel); - } - StreamGate gate; - ASSERT_EQ(gate.enqueue(readerStream), cudaSuccess); - ASSERT_EQ(cudaMemcpyAsync(reinterpret_cast(readback), reinterpret_cast(gpuAddress), bytes, - cudaMemcpyDeviceToHost, readerStream), - cudaSuccess); - if (finishReader) - { - second->suspend(); - } - - if (deferred) - { - ASSERT_TRUE(finishReader ? first->resize(4, 4) : second->enterDecode()); - EXPECT_EQ(first->historyLength(), 4); - } - else - { - first->offloadSparsePages({page}); - } - EXPECT_FALSE(page->queryReady()); - EXPECT_EQ(page->cacheLevel, kSparseHistoryLevel); - auto recycled = storage.newGpuSlots(TypedVec{1}); - auto releaseRecycled - = FuncGuard([&]() { storage.releaseSlot(lc, kHotLevel, std::move(recycled[lc].front())); }); - EXPECT_EQ(recycled[lc].front().slotId(), gpuSlot); - EXPECT_FALSE(recycled[lc].front().queryReady()); - recycled[lc].front().readyEvent.waitInStream(reinterpret_cast(mStream)); - ASSERT_EQ(cudaMemsetAsync(reinterpret_cast(gpuAddress), 0, bytes, mStream), cudaSuccess); - gate.release(); - ASSERT_EQ(cudaStreamSynchronize(mStream), cudaSuccess); - auto const* readBytes = reinterpret_cast(readback); - EXPECT_TRUE(std::all_of(readBytes, readBytes + bytes, [](uint8_t value) { return value == kPattern; })); - page->readyEvent.synchronize(); - auto const hostAddress = std::get(storage.slotAddress( - kSparseHistoryLevel, storage.getPoolGroupIndex(kSparseHistoryLevel, lc), page->slotId(), PoolIndex{0})); - auto const* hostBytes = reinterpret_cast(hostAddress); - EXPECT_TRUE(std::all_of(hostBytes, hostBytes + bytes, [](uint8_t value) { return value == kPattern; })); - } -} - -TEST_F(KvCacheManagerV2SparseOffloadTest, HostOomLeavesEntireBatchOnGpu) -{ - auto config = sparseConfig(); - config.cacheTiers[1] = HostCacheTierConfig{2 << 20}; - auto codec = std::make_unique(); - auto* observer = codec.get(); - auto manager = std::make_shared(std::move(config), nullptr, std::move(codec)); - auto const apiLock = manager->lockExclusive(); - auto& storage = manager->storage(); - auto cache = manager->createKvCache(); - auto closeCache = FuncGuard([&]() { cache->close(); }); - ASSERT_TRUE(cache->resume(stream())); - ASSERT_TRUE(cache->resize(8, 8)); - auto first = pageAt(*cache); - auto second = pageAt(*cache, 1); - auto const firstSlot = first->slotId(); - auto const secondSlot = second->slotId(); - auto const version = cache->pageStorageVersion(); - EXPECT_THROW(cache->offloadSparsePages({first, second}), OutOfPagesError); - EXPECT_EQ(observer->encodeCalls, 0); - EXPECT_EQ(first->cacheLevel, kHotLevel); - EXPECT_EQ(second->cacheLevel, kHotLevel); - EXPECT_EQ(first->slotId(), firstSlot); - EXPECT_EQ(second->slotId(), secondSlot); - EXPECT_EQ(first->status(), PageStatus::LOCKED); - EXPECT_EQ(second->status(), PageStatus::LOCKED); - EXPECT_EQ(cache->pageStorageVersion(), version); - EXPECT_EQ(storage.getStatistics(kSparseHistoryLevel).free, 1); - EXPECT_EQ(storage.getStatistics(kHotLevel).free, 0); - EXPECT_EQ(cache->getBasePageIndices(LifeCycleId{0})[0], slotIdToPageIndexValue(firstSlot)); - EXPECT_EQ(cache->getBasePageIndices(LifeCycleId{0})[1], slotIdToPageIndexValue(secondSlot)); -} - -TEST_F(KvCacheManagerV2SparseOffloadTest, AsynchronousRejectionFencesBothSlotsWithoutPublishingHostIndices) -{ - auto config = sparseConfig(); - config.enableStats = true; - auto codec = std::make_unique(AsyncRejectingColdPageCodec::Operation::kEncode); - auto* rejecting = codec.get(); - auto manager = std::make_shared(std::move(config), nullptr, std::move(codec)); - auto const apiLock = manager->lockExclusive(); - auto& storage = manager->storage(); - auto cache = manager->createKvCache(); - auto closeCache = FuncGuard([&]() { cache->close(); }); - ASSERT_TRUE(cache->resume(stream())); - ASSERT_TRUE(cache->resize(4, 4)); - auto page = pageAt(*cache); - auto const gpuSlot = page->slotId(); - auto const lc = page->lifeCycle; - auto blocker = storage.newSlots(kSparseHistoryLevel, TypedVec{1}); - auto releaseBlocker - = FuncGuard([&]() { storage.releaseSlot(lc, kSparseHistoryLevel, std::move(blocker[lc].front())); }); - auto releaseCodec = FuncGuard([&]() { rejecting->release(); }); - auto const version = cache->pageStorageVersion(); - EXPECT_THROW(cache->offloadSparsePages({page}), TllmException); - ASSERT_TRUE(rejecting->launched()); - EXPECT_EQ(page->cacheLevel, kHotLevel); - EXPECT_EQ(page->slotId(), gpuSlot); - EXPECT_EQ(page->status(), PageStatus::LOCKED); - EXPECT_FALSE(page->queryReady()); - EXPECT_GT(cache->pageStorageVersion(), version); - EXPECT_EQ(cache->getBasePageIndices(lc)[0], slotIdToPageIndexValue(gpuSlot)); - auto recycled = storage.newSlots(kSparseHistoryLevel, TypedVec{1}); - auto releaseRecycled - = FuncGuard([&]() { storage.releaseSlot(lc, kSparseHistoryLevel, std::move(recycled[lc].front())); }); - EXPECT_FALSE(recycled[lc].front().queryReady()); - EXPECT_TRUE(manager->getAndResetIterationStats().empty()); - rejecting->release(); - recycled[lc].front().readyEvent.synchronize(); -} - -TEST_F(KvCacheManagerV2SparseOffloadTest, LaterCodecBatchFailurePreservesAllSourcePages) -{ - auto config = makeSplitColdGroupingConfig(); - for (auto& layer : config.layers) - { - std::get(layer).buffers.front().isSparse = true; - } - std::get(config.layers[1]).buffers.front().size *= 2; - auto codec = std::make_unique(); - auto* observer = codec.get(); - observer->rejectEncodeCall = 2; - auto manager = std::make_shared(std::move(config), nullptr, std::move(codec)); - auto const apiLock = manager->lockExclusive(); - auto& storage = manager->storage(); - ASSERT_EQ(storage.numPoolGroups(kHotLevel), PoolGroupIndex{2}); - auto cache = manager->createKvCache(); - auto closeCache = FuncGuard([&]() { cache->close(); }); - ASSERT_TRUE(cache->resume(stream())); - ASSERT_TRUE(cache->resize(4, 4)); - auto first = pageAt(*cache, 0, LifeCycleId{0}); - auto second = pageAt(*cache, 0, LifeCycleId{1}); - auto const firstSlot = first->slotId(); - auto const secondSlot = second->slotId(); - auto const version = cache->pageStorageVersion(); - EXPECT_THROW(cache->offloadSparsePages({first, second}), TllmException); - EXPECT_EQ(observer->encodeCalls, 2); - EXPECT_EQ(first->cacheLevel, kHotLevel); - EXPECT_EQ(second->cacheLevel, kHotLevel); - EXPECT_EQ(first->slotId(), firstSlot); - EXPECT_EQ(second->slotId(), secondSlot); - EXPECT_GT(cache->pageStorageVersion(), version); - EXPECT_EQ(storage.getStatistics(kSparseHistoryLevel).free, storage.getStatistics(kSparseHistoryLevel).total); - observer->rejectEncodeCall = 0; - cache->offloadSparsePages({first, second}); - EXPECT_EQ(first->cacheLevel, kSparseHistoryLevel); - EXPECT_EQ(second->cacheLevel, kSparseHistoryLevel); -} - -TEST_F(KvCacheManagerV2SparseOffloadTest, RejectsWritablePartialAndDensePagesBeforeAnyCopy) -{ - for (bool const sparse : {false, true}) - { - for (int const history : {0, 2, 4}) - { - SCOPED_TRACE(sparse); - SCOPED_TRACE(history); - auto config = sparse ? sparseConfig() : makeTieredConfig(); - auto codec = std::make_unique(); - auto* observer = codec.get(); - auto manager = std::make_shared(std::move(config), nullptr, std::move(codec)); - auto const apiLock = manager->lockExclusive(); - auto cache = manager->createKvCache(); - auto closeCache = FuncGuard([&]() { cache->close(); }); - ASSERT_TRUE(cache->resume(stream())); - ASSERT_TRUE(cache->resize(8, history)); - auto first = pageAt(*cache); - auto input = pageAt(*cache, 1); - auto const version = cache->pageStorageVersion(); - EXPECT_THROW(cache->offloadSparsePages({first, input}), LogicError); - EXPECT_EQ(observer->encodeCalls, 0); - EXPECT_EQ(first->cacheLevel, kHotLevel); - EXPECT_EQ(input->cacheLevel, kHotLevel); - EXPECT_EQ(cache->pageStorageVersion(), version); - } - } -} - -TEST_F(KvCacheManagerV2SparseOffloadTest, RejectsPartialCommittedPagesAndNonOwners) -{ - auto config = sparseConfig(); - config.commitMinSnapshot = true; - auto manager = std::make_shared(std::move(config)); - auto const apiLock = manager->lockExclusive(); - auto cache = manager->createKvCache(); - auto other = manager->createKvCache(); - auto closeCaches = FuncGuard( - [&]() - { - cache->close(); - other->close(); - }); - EXPECT_THROW(cache->offloadSparsePages({}), LogicError); - ASSERT_TRUE(cache->resume(stream())); - ASSERT_TRUE(other->resume(stream())); - ASSERT_TRUE(cache->resize(4, 2)); - cache->commit(tokens(2), /*isEnd=*/true); - auto partial = pageAt(*cache); - ASSERT_TRUE(partial->isCommitted()); - EXPECT_THROW(cache->offloadSparsePages({partial}), LogicError); - EXPECT_EQ(partial->cacheLevel, kHotLevel); - ASSERT_TRUE(other->resize(4, 4)); - auto complete = pageAt(*other); - EXPECT_THROW(cache->offloadSparsePages({complete}), LogicError); - EXPECT_EQ(complete->cacheLevel, kHotLevel); - EXPECT_THROW(other->offloadSparsePages({nullptr}), LogicError); -} - -TEST_F(KvCacheManagerV2SparseOffloadTest, EmitsCommittedTierChangeEvenWhenSlotIndexIsUnchanged) -{ - auto events = std::make_shared(128); - auto manager = std::make_shared(sparseConfig(), events); - auto const apiLock = manager->lockExclusive(); - auto cache = manager->createKvCache(); - auto closeCache = FuncGuard([&]() { cache->close(); }); - ASSERT_TRUE(cache->resume(stream())); - ASSERT_TRUE(cache->resize(4, 4)); - cache->commit(tokens()); - auto page = pageAt(*cache); - ASSERT_TRUE(page->isCommitted()); - auto const originalIndex = cache->getBasePageIndices(LifeCycleId{0})[0]; - events->flushIterationEvents(); - events->getLatestEvents(/*timeoutMs=*/0); - - auto const version = cache->pageStorageVersion(); - cache->offloadSparsePages({page, page}); - EXPECT_EQ(cache->getBasePageIndices(LifeCycleId{0})[0], originalIndex); - EXPECT_GT(cache->pageStorageVersion(), version); - events->flushIterationEvents(); - auto const updates = events->getLatestEvents(/*timeoutMs=*/0); - ASSERT_EQ(updates.size(), 1); - EXPECT_EQ(updates.front().layerGroupId, 0); - auto const* data = std::get_if(&updates.front().data); - ASSERT_NE(data, nullptr); - ASSERT_TRUE(data->cacheLevel.has_value()); - EXPECT_EQ(data->cacheLevel->oldValue, kHotLevel.value()); - EXPECT_EQ(data->cacheLevel->newValue, kSparseHistoryLevel.value()); -} - -TEST_F(KvCacheManagerV2SparseOffloadTest, DemotedUncommittedHistoryCanCommitAndBeReused) -{ - auto manager = std::make_shared(sparseConfig()); - auto const apiLock = manager->lockExclusive(); - auto cache = manager->createKvCache(); - auto closeCache = FuncGuard([&]() { cache->close(); }); - ASSERT_TRUE(cache->resume(stream())); - ASSERT_TRUE(cache->resize(4, 4)); - cache->offloadSparsePages({pageAt(*cache)}); - ASSERT_TRUE(cache->enterDecode()); - auto const hostSlot = pageAt(*cache)->slotId(); - cache->commit(tokens()); - auto committed = pageAt(*cache); - ASSERT_TRUE(committed->isCommitted()); - EXPECT_EQ(committed->cacheLevel, kSparseHistoryLevel); - EXPECT_EQ(committed->slotId(), hostSlot); - auto reused = manager->createKvCache({}, tokens()); - auto closeReused = FuncGuard([&]() { reused->close(); }); - ASSERT_TRUE(reused->resume(stream(), true)); - EXPECT_EQ(pageAt(*reused), committed); - cache->close(); - EXPECT_EQ(committed->status(), PageStatus::LOCKED); - EXPECT_EQ(reused->getBasePageIndices(LifeCycleId{0})[0], slotIdToPageIndexValue(hostSlot)); -} - -class KvCacheManagerV2DecodeOffloadTest : public KvCacheManagerV2PageLockTest -{ -}; - -TEST_F(KvCacheManagerV2DecodeOffloadTest, PrefillRetainsSparseSwaHistoryAndRestoresGpuStorage) -{ - auto config = sparseConfig(); - config.swaScratchReuse = SwaScratchReuseConfig{}; - auto& layer = std::get(config.layers.front()); - layer.slidingWindowSize = 4; - layer.buffers.front().size = 4096; - auto manager = std::make_shared(std::move(config)); - auto const apiLock = manager->lockExclusive(); - auto& storage = manager->storage(); - auto cache = manager->createKvCache(); - auto closeCache = FuncGuard([&]() { cache->close(); }); - ASSERT_TRUE(cache->resume(stream())); - cache->setEnableSwaScratchReuse(true); - ASSERT_TRUE(cache->resize(12, 0)); - EXPECT_FALSE(cache->hasScratchSlots()); - ASSERT_TRUE(cache->resize(12, 8)); - EXPECT_FALSE(cache->isDecoding()); - EXPECT_EQ(cache->getPageStorageSnapshot(LifeCycleId{0}).eligibleHistoryBlocks(), 0); - for (int ord = 0; ord < 3; ++ord) - { - ASSERT_NE(pageAt(*cache, ord), nullptr); - EXPECT_EQ(pageAt(*cache, ord)->cacheLevel, kHotLevel); - } - - for (bool prefetch : {false, true}) - { - cache->suspend(); - storage.forceEvict(kHotLevel, TypedVec{3}); - for (int ord = 0; ord < 3; ++ord) - EXPECT_EQ(pageAt(*cache, ord)->cacheLevel, kSparseHistoryLevel); - if (prefetch) - ASSERT_TRUE(cache->prefetch(kHotLevel)); - ASSERT_TRUE(cache->resume()); - for (int ord = 0; ord < 3; ++ord) - EXPECT_EQ(pageAt(*cache, ord)->cacheLevel, kHotLevel); - } - ASSERT_TRUE(cache->resize(12, 12)); - EXPECT_EQ(cache->getPageStorageSnapshot(LifeCycleId{0}).eligibleHistoryBlocks(), 0); - ASSERT_TRUE(cache->enterDecode()); - for (int ord = 0; ord < 3; ++ord) - { - EXPECT_EQ(pageAt(*cache, ord)->cacheLevel, kSparseHistoryLevel); - EXPECT_EQ(pageAt(*cache, ord)->status(), PageStatus::LOCKED); - } -} - -TEST_F(KvCacheManagerV2DecodeOffloadTest, EntryScansUnchangedWatermarkAndResumeKeepsHostHistory) -{ - auto codec = std::make_unique(); - auto* observer = codec.get(); - auto manager = std::make_shared(sparseConfig(), nullptr, std::move(codec)); - auto const apiLock = manager->lockExclusive(); - auto cache = manager->createKvCache(); - auto closeCache = FuncGuard([&]() { cache->close(); }); - ASSERT_TRUE(cache->resume(stream())); - ASSERT_TRUE(cache->resize(8, 8)); - EXPECT_EQ(observer->encodeCalls, 0); - auto const prefillVersion = cache->pageStorageVersion(); - ASSERT_TRUE(cache->enterDecode()); - EXPECT_TRUE(cache->isDecoding()); - EXPECT_EQ(cache->historyLength(), 8); - EXPECT_EQ(observer->encodeCalls, 1); - EXPECT_EQ(observer->encodedPages, 2); - EXPECT_GT(cache->pageStorageVersion(), prefillVersion); - auto const decodeVersion = cache->pageStorageVersion(); - ASSERT_TRUE(cache->enterDecode()); - EXPECT_EQ(observer->encodeCalls, 1); - EXPECT_EQ(cache->pageStorageVersion(), decodeVersion); - EXPECT_THROW(cache->resize(8, 4), std::invalid_argument); - cache->suspend(); - EXPECT_THROW(cache->resume(std::nullopt, false), std::invalid_argument); - ASSERT_TRUE(cache->resume()); - EXPECT_EQ(observer->encodeCalls, 1); - EXPECT_GT(cache->pageStorageVersion(), decodeVersion); - for (int ord = 0; ord < 2; ++ord) - { - auto const page = pageAt(*cache, ord); - EXPECT_EQ(page->cacheLevel, kSparseHistoryLevel); - EXPECT_EQ(cache->getBasePageIndices(LifeCycleId{0})[ord], slotIdToPageIndexValue(page->slotId())); - } -} - -TEST_F(KvCacheManagerV2DecodeOffloadTest, BeamExpansionOffloadsOnlyExistingHistoryRows) -{ - auto config = sparseConfig(); - config.enablePartialCommit = false; - std::get(config.layers.front()).buffers.front().size = 4096; - auto manager = std::make_shared(std::move(config)); - auto const apiLock = manager->lockExclusive(); - auto cache = manager->createKvCache({}, {}, std::nullopt, {}, /*expectedPromptLength=*/8); - auto closeCache = FuncGuard([&]() { cache->close(); }); - ASSERT_TRUE(cache->resume(stream())); - ASSERT_TRUE(cache->resize(12, 8)); - cache->setBeamWidth(BeamIndex{3}); - ASSERT_EQ(cache->blocks()[BlockOrdinal{0}].pages.size(), BeamIndex{1}); - ASSERT_EQ(cache->blocks()[BlockOrdinal{1}].pages.size(), BeamIndex{1}); - ASSERT_EQ(cache->blocks()[BlockOrdinal{2}].pages.size(), BeamIndex{3}); - - ASSERT_TRUE(cache->enterDecode()); - EXPECT_TRUE(cache->isDecoding()); - for (int historyLength : {8, 12}) - { - ASSERT_TRUE(cache->resize(16, historyLength)); - for (BlockOrdinal ordinal{0}; ordinal < cache->numBlocks(); ++ordinal) - { - auto const& pages = cache->blocks()[ordinal].pages; - for (BeamIndex bi{0}; bi < pages.size(); ++bi) - { - auto const& page = blockPageGetPage(pages[bi][LifeCycleId{0}]); - ASSERT_NE(page, nullptr); - EXPECT_EQ(page->cacheLevel, - ordinal < BlockOrdinal{historyLength / cache->tokensPerBlock()} ? kSparseHistoryLevel : kHotLevel); - EXPECT_EQ(page->status(), PageStatus::LOCKED); - EXPECT_EQ(cache->getBasePageIndices(LifeCycleId{0}, bi)[ordinal.value()], - slotIdToPageIndexValue(page->slotId())); - } - } - EXPECT_EQ(cache->getPageStorageSnapshot(LifeCycleId{0}).eligibleHistoryBlocks(), - historyLength / cache->tokensPerBlock()); - EXPECT_EQ(cache->getPageStorageSnapshot(LifeCycleId{0}, BeamIndex{1}).eligibleHistoryBlocks(), 0); - } -} - -TEST_F(KvCacheManagerV2DecodeOffloadTest, SparseHistoryIsNotMarkedForTurnEndDrop) -{ - auto config = sparseConfig(); - std::get(config.layers.front()).slidingWindowSize = 4; - auto manager = std::make_shared(std::move(config)); - auto const apiLock = manager->lockExclusive(); - auto page = seedPrefix(*manager, kHotLevel); - auto cache = manager->createKvCache({}, tokens()); - auto closeCache = FuncGuard([&]() { cache->close(); }); - ASSERT_TRUE(cache->resume(stream(), true)); - cache->stopCommitting(); - auto dropPlan = cache->planCommittedBlockDrop(); - EXPECT_EQ(page->plannedDropCount, 0); -} - -TEST_F(KvCacheManagerV2DecodeOffloadTest, CachedPrefillRestoresGpuOnResumePrefetchAndRebase) -{ - for (int path : {0, 1, 2}) - { - SCOPED_TRACE(path); - auto manager = std::make_shared(sparseConfig()); - auto const apiLock = manager->lockExclusive(); - auto page = seedPrefix(*manager, kSparseHistoryLevel); - auto cache = path == 2 ? manager->createKvCache() : manager->createKvCache({}, tokens()); - auto closeCache = FuncGuard([&]() { cache->close(); }); - if (path == 1) - { - ASSERT_TRUE(cache->prefetch(kHotLevel)); - EXPECT_EQ(page->cacheLevel, kHotLevel); - } - ASSERT_TRUE(cache->resume(stream())); - if (path == 2) - { - ASSERT_TRUE(cache->resize(4, 4)); - cache->commit(tokens()); - } - EXPECT_EQ(pageAt(*cache), page); - EXPECT_EQ(page->cacheLevel, kHotLevel); - EXPECT_EQ(cache->historyLength(), 4); - EXPECT_FALSE(cache->isDecoding()); - EXPECT_EQ(cache->getPageStorageSnapshot(LifeCycleId{0}).eligibleHistoryBlocks(), 0); - EXPECT_EQ(cache->getBasePageIndices(LifeCycleId{0})[0], slotIdToPageIndexValue(page->slotId())); - } -} - -TEST_F(KvCacheManagerV2DecodeOffloadTest, MixedDenseSwaLocksAreRestoredAfterHostOom) -{ - auto config = sparseConfig(); - config.cacheTiers[0] = GpuCacheTierConfig{8 << 20}; - auto& sparse = std::get(config.layers.front()); - sparse.slidingWindowSize = 4; - auto dense = sparse; - dense.layerId = 1; - dense.buffers.front().isSparse = false; - config.layers.emplace_back(std::move(dense)); - auto manager = std::make_shared(std::move(config)); - auto const apiLock = manager->lockExclusive(); - auto& storage = manager->storage(); - auto cache = manager->createKvCache(); - auto closeCache = FuncGuard([&]() { cache->close(); }); - ASSERT_TRUE(cache->resume(stream())); - ASSERT_TRUE(cache->resize(8, 4)); - auto const densePage = pageAt(*cache, 0, LifeCycleId{1}); - auto const denseSlot = densePage->slotId(); - ASSERT_TRUE(cache->enterDecode()); - EXPECT_EQ(densePage->cacheLevel, kHotLevel); - EXPECT_EQ(pageAt(*cache)->cacheLevel, kSparseHistoryLevel); - EXPECT_EQ(pageAt(*cache, 1)->cacheLevel, kHotLevel); - auto const version = cache->pageStorageVersion(); - auto const before = cache->getPageStorageSnapshot(LifeCycleId{0}); - ASSERT_TRUE(cache->acknowledgePageStorage(version)); - auto const freeHost = manager->getStorageStatistics(kSparseHistoryLevel) - .at(storage.getPoolGroupIndex(kSparseHistoryLevel, LifeCycleId{0})) - .free; - auto blockers = storage.newSlots(kSparseHistoryLevel, TypedVec{freeHost, 0}); - auto releaseBlockers = FuncGuard( - [&]() - { - for (auto& slot : blockers[LifeCycleId{0}]) - storage.releaseSlot(LifeCycleId{0}, kSparseHistoryLevel, std::move(slot)); - }); - EXPECT_FALSE(cache->resize(8, 8)); - EXPECT_EQ(cache->historyLength(), 4); - EXPECT_GT(cache->pageStorageVersion(), version); - EXPECT_TRUE(cache->pageStorageDirty()); - EXPECT_FALSE(cache->acknowledgePageStorage(version)); - auto const after = cache->getPageStorageSnapshot(LifeCycleId{0}); - EXPECT_EQ(after.basePageIndices(), before.basePageIndices()); - EXPECT_EQ(after.cacheLevels(), before.cacheLevels()); - EXPECT_EQ(after.eligibleHistoryBlocks(), before.eligibleHistoryBlocks()); - EXPECT_EQ(pageAt(*cache, 0, LifeCycleId{1}), densePage); - EXPECT_EQ(densePage->status(), PageStatus::LOCKED); - EXPECT_EQ(densePage->cacheLevel, kHotLevel); - EXPECT_EQ(densePage->slotId(), denseSlot); - EXPECT_EQ(cache->getBasePageIndices(LifeCycleId{1})[0], slotIdToPageIndexValue(denseSlot)); - EXPECT_EQ(pageAt(*cache, 1)->cacheLevel, kHotLevel); -} - -TEST_F(KvCacheManagerV2DecodeOffloadTest, HistoryUpdatesOffloadOnlyNewFullPages) -{ - auto codec = std::make_unique(); - auto* observer = codec.get(); - auto manager = std::make_shared(sparseConfig(), nullptr, std::move(codec)); - auto const apiLock = manager->lockExclusive(); - auto cache = manager->createKvCache(); - auto closeCache = FuncGuard([&]() { cache->close(); }); - ASSERT_TRUE(cache->resume(stream())); - ASSERT_TRUE(cache->resize(8, 4)); - ASSERT_TRUE(cache->enterDecode()); - auto const version = cache->pageStorageVersion(); - auto const firstHostSlot = pageAt(*cache)->slotId(); - ASSERT_TRUE(cache->resize(8, 7)); - EXPECT_EQ(observer->encodedPages, 1); - EXPECT_EQ(cache->pageStorageVersion(), version); - EXPECT_EQ(pageAt(*cache, 1)->cacheLevel, kHotLevel); - ASSERT_TRUE(cache->resize(8, 8)); - EXPECT_EQ(observer->encodeCalls, 2); - EXPECT_EQ(observer->encodedPages, 2); - EXPECT_GT(cache->pageStorageVersion(), version); - EXPECT_EQ(pageAt(*cache)->slotId(), firstHostSlot); - EXPECT_EQ(pageAt(*cache, 1)->cacheLevel, kSparseHistoryLevel); - ASSERT_TRUE(cache->resize(12)); - ASSERT_TRUE(cache->resize(12, 9)); - EXPECT_EQ(pageAt(*cache, 2)->cacheLevel, kHotLevel); - EXPECT_EQ(observer->encodedPages, 2); -} - -TEST_F(KvCacheManagerV2BatchTest, PrefillPromotesOneSharedPageAndReleaseRetriesOffload) -{ - for (int path : {0, 1, 2, 3}) - { - SCOPED_TRACE(path); - auto codec = std::make_unique(); - auto* observer = codec.get(); - auto manager = std::make_shared(sparseConfig(), nullptr, std::move(codec)); - auto const apiLock = manager->lockExclusive(); - auto& storage = manager->storage(); - auto page = seedPrefix(*manager, kHotLevel); - auto first = manager->createKvCache({}, tokens()); - auto second = manager->createKvCache({}, tokens()); - std::shared_ptr prefill; - Batch batch(manager, 3, 1); - auto closeCaches = FuncGuard( - [&]() - { - if (prefill) - { - prefill->close(); - } - first->close(); - second->close(); - }); - ASSERT_TRUE(first->resume(stream())); - ASSERT_TRUE(second->resume(stream())); - if (path == 1) - { - prefill = manager->createKvCache({}, tokens()); - ASSERT_TRUE(prefill->resume(stream())); - prefill->suspend(); - } - auto const lc = page->lifeCycle; - auto const pg = storage.getPoolGroupIndex(kHotLevel, lc); - size_t const bytes = storage.slotSize(kHotLevel, pg)[PoolIndex{0}]; - auto const originalAddress - = std::get(storage.slotAddress(kHotLevel, pg, page->slotId(), PoolIndex{0})); - ASSERT_EQ(cudaMemsetAsync(reinterpret_cast(originalAddress), 0x3C, bytes, mStream), cudaSuccess); - ASSERT_TRUE(first->enterDecode()); - ASSERT_TRUE(second->enterDecode()); - ASSERT_EQ(page->cacheLevel, kSparseHistoryLevel); - batch.add(*first, 0); - batch.add(*second, 1); - batch.publish(batchStream()); - auto const firstVersion = first->pageStorageVersion(); - auto const secondVersion = second->pageStorageVersion(); - - if (!prefill) - { - prefill = path == 2 ? manager->createKvCache() : manager->createKvCache({}, tokens()); - } - if (path == 3) - { - ASSERT_TRUE(prefill->prefetch(kHotLevel)); - } - ASSERT_TRUE(prefill->resume(stream())); - if (path == 2) - { - ASSERT_TRUE(prefill->resize(4, 4)); - prefill->commit(tokens()); - } - EXPECT_EQ(page->cacheLevel, kHotLevel); - EXPECT_EQ(observer->decodeCalls, 1); - EXPECT_EQ(observer->decodedPages, 1); - EXPECT_GT(first->pageStorageVersion(), firstVersion); - EXPECT_GT(second->pageStorageVersion(), secondVersion); - for (auto const& cache : {first, second, prefill}) - { - EXPECT_EQ(pageAt(*cache), page); - EXPECT_EQ(cache->getBasePageIndices(lc)[0], slotIdToPageIndexValue(page->slotId())); - } - EXPECT_EQ(storage.getStatistics(kHotLevel).free, storage.getStatistics(kHotLevel).total - 1); - EXPECT_EQ(storage.getStatistics(kSparseHistoryLevel).free, storage.getStatistics(kSparseHistoryLevel).total); - batch.add(*prefill, 2); - EXPECT_EQ(batch.publish(batchStream()), (std::vector{0, 1, 2})); - EXPECT_EQ( - read(batch.pageTableAddress(lc), 3), (std::vector(3, slotIdToPageIndexValue(page->slotId())))); - EXPECT_EQ(read(batch.numBlocksAddress(lc), 3), (std::vector{0, 0, 0})); - auto const address = std::get(storage.slotAddress(kHotLevel, pg, page->slotId(), PoolIndex{0})); - auto const data = read(address, bytes / sizeof(int32_t)); - EXPECT_TRUE(std::all_of(data.begin(), data.end(), [](int32_t value) { return value == 0x3C3C3C3C; })); - - prefill->close(); - batch.publish(batchStream()); - EXPECT_EQ(page->cacheLevel, kSparseHistoryLevel); - EXPECT_EQ(observer->encodeCalls, 2); - EXPECT_EQ(storage.getStatistics(kHotLevel).free, storage.getStatistics(kHotLevel).total); - EXPECT_EQ(read(batch.numBlocksAddress(lc), 3), (std::vector{1, 1, 0})); - } -} - -TEST_F(KvCacheManagerV2DecodeOffloadTest, SharedPrefillOwnerDefersDemotionAcrossDecodeResume) -{ - auto manager = std::make_shared(sparseConfig()); - auto const apiLock = manager->lockExclusive(); - auto page = seedPrefix(*manager, kHotLevel); - auto first = manager->createKvCache({}, tokens()); - auto second = manager->createKvCache({}, tokens()); - auto closeCaches = FuncGuard( - [&]() - { - first->close(); - second->close(); - }); - ASSERT_TRUE(first->resume(stream())); - ASSERT_TRUE(second->resume(stream())); - auto const version = first->pageStorageVersion(); - ASSERT_TRUE(first->enterDecode()); - EXPECT_TRUE(first->isDecoding()); - EXPECT_GT(first->pageStorageVersion(), version); - EXPECT_EQ(page->cacheLevel, kHotLevel); - EXPECT_EQ(first->getPageStorageSnapshot(LifeCycleId{0}).eligibleHistoryBlocks(), 0); - second->suspend(); - ASSERT_TRUE(first->enterDecode()); - ASSERT_TRUE(second->resume()); - EXPECT_TRUE(second->isActive()); - EXPECT_FALSE(second->isDecoding()); - EXPECT_EQ(page->cacheLevel, kHotLevel); - EXPECT_EQ(pageAt(*second), page); - first->suspend(); - EXPECT_EQ(page->cacheLevel, kHotLevel); - ASSERT_TRUE(first->resume()); - EXPECT_TRUE(first->isActive()); - EXPECT_TRUE(first->isDecoding()); - EXPECT_EQ(page->cacheLevel, kHotLevel); - second->close(); - ASSERT_TRUE(first->enterDecode()); - EXPECT_EQ(first->historyLength(), 4); - EXPECT_EQ(page->cacheLevel, kSparseHistoryLevel); -} - -TEST_F(KvCacheManagerV2DecodeOffloadTest, PromotionGpuOomPreservesHostPageAndAllowsRetry) -{ - auto config = sparseConfig(); - config.maxUtilForResume = 1.0f; - auto codec = std::make_unique(); - auto* observer = codec.get(); - auto manager = std::make_shared(std::move(config), nullptr, std::move(codec)); - auto const apiLock = manager->lockExclusive(); - auto& storage = manager->storage(); - auto page = seedPrefix(*manager, kHotLevel); - auto decoder = manager->createKvCache({}, tokens()); - auto prefill = manager->createKvCache({}, tokens()); - auto closeCaches = FuncGuard( - [&]() - { - prefill->close(); - decoder->close(); - }); - ASSERT_TRUE(decoder->resume(stream())); - ASSERT_TRUE(decoder->enterDecode()); - auto const hostSlot = page->slotId(); - auto const version = decoder->pageStorageVersion(); - auto const lc = page->lifeCycle; - auto blockers = storage.newGpuSlots(TypedVec{storage.getStatistics(kHotLevel).free}); - auto releaseBlockers = FuncGuard( - [&]() - { - for (auto& slot : blockers[lc]) - { - storage.releaseSlot(lc, kHotLevel, std::move(slot)); - } - }); - EXPECT_FALSE(prefill->resume(stream())); - EXPECT_FALSE(prefill->isActive()); - EXPECT_TRUE(decoder->isActive()); - EXPECT_EQ(page->cacheLevel, kSparseHistoryLevel); - EXPECT_EQ(page->slotId(), hostSlot); - EXPECT_EQ(decoder->pageStorageVersion(), version); - EXPECT_EQ(decoder->getBasePageIndices(lc)[0], slotIdToPageIndexValue(hostSlot)); - EXPECT_EQ(observer->decodeCalls, 0); - releaseBlockers.run(); - ASSERT_TRUE(prefill->resume()); - EXPECT_EQ(pageAt(*prefill), page); - EXPECT_EQ(page->cacheLevel, kHotLevel); - EXPECT_EQ(observer->decodedPages, 1); -} - -TEST_F(KvCacheManagerV2DecodeOffloadTest, PromotionCodecBatchFailurePreservesAllHostPages) -{ - auto config = makeSplitColdGroupingConfig(); - config.enableStats = true; - for (auto& layer : config.layers) - { - std::get(layer).buffers.front().isSparse = true; - } - std::get(config.layers[1]).buffers.front().size *= 2; - auto codec = std::make_unique(); - auto* observer = codec.get(); - auto manager = std::make_shared(std::move(config), nullptr, std::move(codec)); - auto const apiLock = manager->lockExclusive(); - auto& storage = manager->storage(); - auto decoder = manager->createKvCache(); - std::shared_ptr prefill; - auto closeCaches = FuncGuard( - [&]() - { - if (prefill) - { - prefill->close(); - } - decoder->close(); - }); - ASSERT_TRUE(decoder->resume(stream())); - ASSERT_TRUE(decoder->resize(4, 4)); - decoder->commit(tokens()); - ASSERT_TRUE(decoder->enterDecode()); - prefill = manager->createKvCache({}, tokens()); - auto const first = pageAt(*decoder, 0, LifeCycleId{0}); - auto const second = pageAt(*decoder, 0, LifeCycleId{1}); - auto const firstSlot = first->slotId(); - auto const secondSlot = second->slotId(); - auto const version = decoder->pageStorageVersion(); - manager->getAndResetIterationStats(); - observer->rejectDecodeCall = 2; - EXPECT_THROW(prefill->resume(stream()), TllmException); - EXPECT_FALSE(prefill->isActive()); - EXPECT_EQ(observer->decodeCalls, 2); - EXPECT_EQ(first->cacheLevel, kSparseHistoryLevel); - EXPECT_EQ(second->cacheLevel, kSparseHistoryLevel); - EXPECT_EQ(first->slotId(), firstSlot); - EXPECT_EQ(second->slotId(), secondSlot); - EXPECT_EQ(decoder->getBasePageIndices(LifeCycleId{0})[0], slotIdToPageIndexValue(firstSlot)); - EXPECT_EQ(decoder->getBasePageIndices(LifeCycleId{1})[0], slotIdToPageIndexValue(secondSlot)); - EXPECT_GT(decoder->pageStorageVersion(), version); - EXPECT_EQ(storage.getStatistics(kHotLevel).free, storage.getStatistics(kHotLevel).total); - EXPECT_TRUE(manager->getAndResetIterationStats().empty()); - observer->rejectDecodeCall = 0; - ASSERT_TRUE(prefill->resume()); - EXPECT_EQ(pageAt(*prefill, 0, LifeCycleId{0}), first); - EXPECT_EQ(pageAt(*prefill, 0, LifeCycleId{1}), second); - EXPECT_EQ(first->cacheLevel, kHotLevel); - EXPECT_EQ(second->cacheLevel, kHotLevel); - EXPECT_EQ(storage.getStatistics(kSparseHistoryLevel).free, storage.getStatistics(kSparseHistoryLevel).total); - auto const stats = manager->getAndResetIterationStats(); - EXPECT_EQ(stats.at(LifeCycleId{0}).iterOnboardBlocks, 1); - EXPECT_EQ(stats.at(LifeCycleId{1}).iterOnboardBlocks, 1); -} - -TEST_F(KvCacheManagerV2DecodeOffloadTest, PromotionRejectionFencesSourceAndRecycledDestination) -{ - auto codec = std::make_unique(AsyncRejectingColdPageCodec::Operation::kDecode); - auto* rejecting = codec.get(); - auto manager = std::make_shared(sparseConfig(), nullptr, std::move(codec)); - auto const apiLock = manager->lockExclusive(); - auto& storage = manager->storage(); - auto page = seedPrefix(*manager, kHotLevel); - auto decoder = manager->createKvCache({}, tokens()); - auto prefill = manager->createKvCache({}, tokens()); - auto closeCaches = FuncGuard( - [&]() - { - prefill->close(); - decoder->close(); - }); - ASSERT_TRUE(decoder->resume(stream())); - ASSERT_TRUE(decoder->enterDecode()); - auto const lc = page->lifeCycle; - auto const hostSlot = page->slotId(); - auto blocker = storage.newGpuSlots(TypedVec{1}); - auto releaseBlocker = FuncGuard([&]() { storage.releaseSlot(lc, kHotLevel, std::move(blocker[lc].front())); }); - auto releaseCodec = FuncGuard([&]() { rejecting->release(); }); - EXPECT_THROW(prefill->resume(stream()), TllmException); - ASSERT_TRUE(rejecting->launched()); - EXPECT_FALSE(prefill->isActive()); - EXPECT_EQ(page->cacheLevel, kSparseHistoryLevel); - EXPECT_EQ(page->slotId(), hostSlot); - EXPECT_FALSE(page->queryReady()); - EXPECT_EQ(decoder->getBasePageIndices(lc)[0], slotIdToPageIndexValue(hostSlot)); - auto recycled = storage.newGpuSlots(TypedVec{1}); - auto releaseRecycled = FuncGuard([&]() { storage.releaseSlot(lc, kHotLevel, std::move(recycled[lc].front())); }); - EXPECT_FALSE(recycled[lc].front().queryReady()); - rejecting->release(); - recycled[lc].front().readyEvent.synchronize(); -} - -TEST_F(KvCacheManagerV2DecodeOffloadTest, PromotionWaitsForLiveAndFinishedReadersBeforeRecyclingHostSlot) -{ - for (bool const finishReader : {false, true}) - { - SCOPED_TRACE(finishReader); - auto manager = std::make_shared(sparseConfig()); - auto const apiLock = manager->lockExclusive(); - auto& storage = manager->storage(); - cudaStream_t readerStream{}; - ASSERT_EQ(cudaStreamCreateWithFlags(&readerStream, cudaStreamNonBlocking), cudaSuccess); - auto destroyReaderStream = FuncGuard([&]() { cudaStreamDestroy(readerStream); }); - auto page = seedPrefix(*manager, kHotLevel); - auto first = manager->createKvCache({}, tokens()); - auto second = manager->createKvCache({}, tokens()); - auto prefill = manager->createKvCache({}, tokens()); - auto closeCaches = FuncGuard( - [&]() - { - prefill->close(); - first->close(); - second->close(); - }); - ASSERT_TRUE(first->resume(stream())); - ASSERT_TRUE(second->resume(reinterpret_cast(readerStream))); - auto const lc = page->lifeCycle; - auto const gpuPool = storage.getPoolGroupIndex(kHotLevel, lc); - auto const hostPool = storage.getPoolGroupIndex(kSparseHistoryLevel, lc); - size_t const bytes = storage.slotSize(kHotLevel, gpuPool)[PoolIndex{0}]; - auto const originalAddress - = std::get(storage.slotAddress(kHotLevel, gpuPool, page->slotId(), PoolIndex{0})); - ASSERT_EQ(cudaMemsetAsync(reinterpret_cast(originalAddress), 0x3C, bytes, mStream), cudaSuccess); - ASSERT_TRUE(first->enterDecode()); - ASSERT_TRUE(second->enterDecode()); - auto const hostSlot = page->slotId(); - auto const hostAddress - = std::get(storage.slotAddress(kSparseHistoryLevel, hostPool, hostSlot, PoolIndex{0})); - auto readback = storage.newGpuSlots(TypedVec{1}); - auto hostBlocker = storage.newSlots(kSparseHistoryLevel, TypedVec{1}); - auto releaseSlots = FuncGuard( - [&]() - { - storage.releaseSlot(lc, kHotLevel, std::move(readback[lc].front())); - storage.releaseSlot(lc, kSparseHistoryLevel, std::move(hostBlocker[lc].front())); - }); - auto const readbackAddress = std::get( - storage.slotAddress(kHotLevel, gpuPool, readback[lc].front().slotId(), PoolIndex{0})); - // Warm the H2D codec before holding a CUDA callback. - storage.copySlotData(lc, kHotLevel, kSparseHistoryLevel, readback[lc].front().slotId(), hostSlot, stream()); - ASSERT_EQ(cudaStreamSynchronize(mStream), cudaSuccess); - std::pair overwrite{hostAddress, bytes}; - StreamGate gate; - ASSERT_EQ(gate.enqueue(readerStream), cudaSuccess); - ASSERT_EQ(cudaMemcpyAsync(reinterpret_cast(readbackAddress), reinterpret_cast(hostAddress), - bytes, cudaMemcpyHostToDevice, readerStream), - cudaSuccess); - if (finishReader) - { - second->suspend(); - } - ASSERT_TRUE(prefill->resume(stream())); - EXPECT_EQ(page->cacheLevel, kHotLevel); - EXPECT_FALSE(page->queryReady()); - EXPECT_EQ(pageAt(*prefill), page); - auto recycled = storage.newSlots(kSparseHistoryLevel, TypedVec{1}); - auto releaseRecycled - = FuncGuard([&]() { storage.releaseSlot(lc, kSparseHistoryLevel, std::move(recycled[lc].front())); }); - EXPECT_EQ(recycled[lc].front().slotId(), hostSlot); - EXPECT_FALSE(recycled[lc].front().queryReady()); - recycled[lc].front().readyEvent.waitInStream(reinterpret_cast(mStream)); - ASSERT_EQ(cudaLaunchHostFunc( - mStream, - [](void* data) - { - auto const& [address, size] = *static_cast const*>(data); - std::memset(reinterpret_cast(address), 0, size); - }, - &overwrite), - cudaSuccess); - gate.release(); - ASSERT_EQ(cudaStreamSynchronize(mStream), cudaSuccess); - auto const promotedAddress - = std::get(storage.slotAddress(kHotLevel, gpuPool, page->slotId(), PoolIndex{0})); - for (auto const address : {readbackAddress, promotedAddress}) - { - std::vector data(bytes); - cuCheck(cuMemcpyDtoH(data.data(), address, bytes)); - EXPECT_TRUE(std::all_of(data.begin(), data.end(), [](uint8_t value) { return value == 0x3C; })); - } - } -} - -TEST_F(KvCacheManagerV2DecodeOffloadTest, SharedOwnersEnterDecodeWithoutMutuallyBlocking) -{ - auto config = sparseConfig(); - config.enableStats = true; - auto codec = std::make_unique(); - auto* observer = codec.get(); - auto manager = std::make_shared(std::move(config), nullptr, std::move(codec)); - auto const apiLock = manager->lockExclusive(); - auto page = seedPrefix(*manager, kHotLevel); - auto first = manager->createKvCache({}, tokens()); - auto second = manager->createKvCache({}, tokens()); - auto third = manager->createKvCache({}, tokens()); - auto closeCaches = FuncGuard( - [&]() - { - first->close(); - second->close(); - third->close(); - }); - for (auto const& cache : {first, second, third}) - { - ASSERT_TRUE(cache->resume(stream())); - EXPECT_EQ(pageAt(*cache), page); - } - auto const gpuSlot = page->slotId(); - auto const hostFree = manager->storage().getStatistics(kSparseHistoryLevel).free; - ASSERT_TRUE(first->enterDecode()); - ASSERT_TRUE(second->enterDecode()); - auto const version = first->pageStorageVersion(); - ASSERT_TRUE(third->resize(4, 4)); - ASSERT_TRUE(first->enterDecode()); - ASSERT_TRUE(second->resize(4, 4)); - EXPECT_EQ(first->pageStorageVersion(), version); - EXPECT_EQ(page->cacheLevel, kHotLevel); - EXPECT_EQ(page->slotId(), gpuSlot); - EXPECT_EQ(observer->encodeCalls, 0); - EXPECT_EQ(manager->storage().getStatistics(kSparseHistoryLevel).free, hostFree); - EXPECT_EQ(first->getPageStorageSnapshot(LifeCycleId{0}).eligibleHistoryBlocks(), 0); - - ASSERT_TRUE(third->enterDecode()); - EXPECT_EQ(observer->encodedPages, 1); - EXPECT_EQ(page->cacheLevel, kSparseHistoryLevel); - for (auto const& cache : {first, second, third}) - { - EXPECT_EQ(cache->historyLength(), 4); - auto const snapshot = cache->getPageStorageSnapshot(LifeCycleId{0}); - EXPECT_EQ(snapshot.eligibleHistoryBlocks(), 1); - EXPECT_EQ(snapshot.basePageIndices()[0], slotIdToPageIndexValue(page->slotId())); - ASSERT_TRUE(cache->enterDecode()); - } - EXPECT_EQ(observer->encodeCalls, 1); - EXPECT_EQ(manager->getAndResetIterationStats().at(LifeCycleId{0}).iterOffloadBlocks, 1); -} - -TEST_F(KvCacheManagerV2DecodeOffloadTest, DeferredOwnerCanCloseBeforePrefillOwners) -{ - for (bool const closeAll : {false, true}) - { - SCOPED_TRACE(closeAll); - auto codec = std::make_unique(); - auto* observer = codec.get(); - auto manager = std::make_shared(sparseConfig(), nullptr, std::move(codec)); - auto const apiLock = manager->lockExclusive(); - auto page = seedPrefix(*manager, kHotLevel); - auto decoder = manager->createKvCache({}, tokens()); - auto prefill = manager->createKvCache({}, tokens()); - auto closeCaches = FuncGuard( - [&]() - { - decoder->close(); - prefill->close(); - }); - ASSERT_TRUE(decoder->resume(stream())); - ASSERT_TRUE(prefill->resume(stream())); - ASSERT_TRUE(decoder->enterDecode()); - decoder->close(); - if (closeAll) - { - prefill->close(); - EXPECT_EQ(observer->encodeCalls, 0); - prefill = manager->createKvCache({}, tokens()); - ASSERT_TRUE(prefill->resume(stream())); - } - ASSERT_TRUE(prefill->enterDecode()); - EXPECT_EQ(page->cacheLevel, kSparseHistoryLevel); - EXPECT_EQ(observer->encodedPages, 1); - } -} - -TEST_F(KvCacheManagerV2DecodeOffloadTest, HostOomDoesNotAdmitDecodeAndCanRetry) -{ - for (int admission : {0, 1, 2}) - { - SCOPED_TRACE(admission); - bool const suspended = admission != 0; - bool const firstResume = admission == 2; - int const history = firstResume ? 4 : 8; - auto manager = std::make_shared(sparseConfig()); - auto const apiLock = manager->lockExclusive(); - auto& storage = manager->storage(); - auto prefix = firstResume ? seedPrefix(*manager, kHotLevel) : nullptr; - auto cache = firstResume ? manager->createKvCache({}, tokens()) : manager->createKvCache(); - auto closeCache = FuncGuard([&]() { cache->close(); }); - if (!firstResume) - { - ASSERT_TRUE(cache->resume(stream())); - ASSERT_TRUE(cache->resize(history, history)); - if (suspended) - cache->suspend(); - } - manager->getAndResetIterationSuspendResumeStats(); - auto blockers = storage.newSlots(kSparseHistoryLevel, TypedVec{firstResume ? 2 : 1}); - auto releaseBlockers = FuncGuard( - [&]() - { - for (auto& slot : blockers[LifeCycleId{0}]) - storage.releaseSlot(LifeCycleId{0}, kSparseHistoryLevel, std::move(slot)); - }); - EXPECT_FALSE(suspended ? cache->resume(stream(), true) : cache->enterDecode()); - EXPECT_EQ(cache->isActive(), !suspended); - EXPECT_FALSE(cache->isDecoding()); - EXPECT_EQ(cache->historyLength(), history); - EXPECT_EQ(cache->capacity(), history); - EXPECT_EQ(cache->getPageStorageSnapshot(LifeCycleId{0}).eligibleHistoryBlocks(), 0); - EXPECT_EQ(manager->getAndResetIterationSuspendResumeStats(), (std::pair{0, 0})); - for (int ord = 0; ord < history / 4; ++ord) - { - auto const page = pageAt(*cache, ord); - EXPECT_EQ(page->cacheLevel, kHotLevel); - EXPECT_EQ(cache->getBasePageIndices(LifeCycleId{0})[ord], - suspended ? kBadPageIndex.value() : slotIdToPageIndexValue(page->slotId())); - } - releaseBlockers.run(); - ASSERT_TRUE(suspended ? cache->resume(std::nullopt, true) : cache->enterDecode()); - EXPECT_TRUE(cache->isDecoding()); - for (int ord = 0; ord < history / 4; ++ord) - EXPECT_EQ(pageAt(*cache, ord)->cacheLevel, kSparseHistoryLevel); - EXPECT_EQ(manager->getAndResetIterationSuspendResumeStats().second, admission == 1 ? 1 : 0); - } -} - -TEST_F(KvCacheManagerV2DecodeOffloadTest, HistoryOomRollsBackShortcutGrowthAndShrink) -{ - for (auto const [oldCapacity, newCapacity] : {std::pair{8, 8}, std::pair{8, 12}, std::pair{12, 8}}) - { - SCOPED_TRACE(newCapacity); - SCOPED_TRACE(oldCapacity); - auto manager = std::make_shared(sparseConfig()); - auto const apiLock = manager->lockExclusive(); - auto& storage = manager->storage(); - auto cache = manager->createKvCache(); - auto closeCache = FuncGuard([&]() { cache->close(); }); - ASSERT_TRUE(cache->resume(stream())); - ASSERT_TRUE(cache->resize(8, 4)); - ASSERT_TRUE(cache->enterDecode()); - ASSERT_TRUE(cache->resize(oldCapacity)); - auto const page = pageAt(*cache, 1); - auto const slotId = page->slotId(); - auto const version = cache->pageStorageVersion(); - auto const gpuFree = storage.getStatistics(kHotLevel).free; - auto blockers = storage.newSlots(kSparseHistoryLevel, TypedVec{1}); - auto releaseBlockers = FuncGuard( - [&]() - { - for (auto& slot : blockers[LifeCycleId{0}]) - storage.releaseSlot(LifeCycleId{0}, kSparseHistoryLevel, std::move(slot)); - }); - EXPECT_FALSE(cache->resize(newCapacity, 8)); - EXPECT_EQ(cache->capacity(), oldCapacity); - EXPECT_EQ(cache->historyLength(), 4); - EXPECT_TRUE(cache->isDecoding()); - EXPECT_EQ(cache->pageStorageVersion(), version); - EXPECT_EQ(pageAt(*cache, 1), page); - EXPECT_EQ(page->slotId(), slotId); - EXPECT_EQ(page->cacheLevel, kHotLevel); - EXPECT_EQ(cache->getBasePageIndices(LifeCycleId{0})[1], slotIdToPageIndexValue(slotId)); - EXPECT_EQ(storage.getStatistics(kHotLevel).free, gpuFree); - if (oldCapacity == 12) - EXPECT_EQ(pageAt(*cache, 2)->status(), PageStatus::LOCKED); - storage.releaseSlot(LifeCycleId{0}, kSparseHistoryLevel, std::move(blockers[LifeCycleId{0}].back())); - blockers[LifeCycleId{0}].clear(); - ASSERT_TRUE(cache->resize(newCapacity, 8)); - EXPECT_EQ(cache->capacity(), newCapacity); - EXPECT_EQ(cache->historyLength(), 8); - EXPECT_EQ(page->cacheLevel, kSparseHistoryLevel); - EXPECT_GT(cache->pageStorageVersion(), version); - } -} - -TEST_F(KvCacheManagerV2DecodeOffloadTest, CodecRejectionDoesNotPublishEntryOrHistoryUpdate) -{ - for (bool entry : {false, true}) - { - SCOPED_TRACE(entry); - auto codec = std::make_unique(); - auto* observer = codec.get(); - auto manager = std::make_shared(sparseConfig(), nullptr, std::move(codec)); - auto const apiLock = manager->lockExclusive(); - auto cache = manager->createKvCache(); - auto closeCache = FuncGuard([&]() { cache->close(); }); - ASSERT_TRUE(cache->resume(stream())); - int const history = entry ? 8 : 4; - ASSERT_TRUE(cache->resize(8, history)); - if (!entry) - ASSERT_TRUE(cache->enterDecode()); - auto const version = cache->pageStorageVersion(); - auto const before = cache->getPageStorageSnapshot(LifeCycleId{0}); - auto const page = pageAt(*cache, 1); - auto const gpuSlot = page->slotId(); - observer->rejectEncodeCall = observer->encodeCalls + 1; - EXPECT_THROW(entry ? cache->enterDecode() : cache->resize(12, 8), TllmException); - EXPECT_EQ(cache->isDecoding(), !entry); - EXPECT_EQ(cache->capacity(), 8); - EXPECT_EQ(cache->historyLength(), history); - EXPECT_GT(cache->pageStorageVersion(), version); - auto const after = cache->getPageStorageSnapshot(LifeCycleId{0}); - EXPECT_EQ(after.basePageIndices(), before.basePageIndices()); - EXPECT_EQ(after.cacheLevels(), before.cacheLevels()); - EXPECT_EQ(after.eligibleHistoryBlocks(), before.eligibleHistoryBlocks()); - EXPECT_EQ(page->cacheLevel, kHotLevel); - EXPECT_EQ(page->slotId(), gpuSlot); - EXPECT_EQ(cache->getBasePageIndices(LifeCycleId{0})[1], slotIdToPageIndexValue(gpuSlot)); - observer->rejectEncodeCall = 0; - ASSERT_TRUE(entry ? cache->enterDecode() : cache->resize(12, 8)); - EXPECT_TRUE(cache->isDecoding()); - EXPECT_EQ(page->cacheLevel, kSparseHistoryLevel); - } -} - -TEST_F(KvCacheManagerV2DecodeOffloadTest, RebaseOffloadsOlderGpuHistoryAtUnchangedWatermark) -{ - auto manager = std::make_shared(sparseConfig()); - auto const apiLock = manager->lockExclusive(); - auto page = seedPrefix(*manager, kHotLevel); - auto cache = manager->createKvCache(); - auto closeCache = FuncGuard([&]() { cache->close(); }); - ASSERT_TRUE(cache->resume(stream())); - ASSERT_TRUE(cache->resize(4, 4)); - ASSERT_TRUE(cache->enterDecode()); - auto const version = cache->pageStorageVersion(); - cache->commit(tokens()); - EXPECT_EQ(pageAt(*cache), page); - EXPECT_EQ(page->cacheLevel, kSparseHistoryLevel); - EXPECT_EQ(cache->historyLength(), 4); - EXPECT_GT(cache->pageStorageVersion(), version); -} - -TEST_F(KvCacheManagerV2DecodeOffloadTest, PrefillRebasePromotesSharedHostPage) -{ - auto manager = std::make_shared(sparseConfig()); - auto const apiLock = manager->lockExclusive(); - auto page = seedPrefix(*manager, kSparseHistoryLevel); - auto decoder = manager->createKvCache({}, tokens()); - auto prefill = manager->createKvCache(); - auto closeCaches = FuncGuard( - [&]() - { - decoder->close(); - prefill->close(); - }); - ASSERT_TRUE(decoder->resume(stream(), true)); - ASSERT_TRUE(prefill->resume(stream())); - ASSERT_TRUE(prefill->resize(4, 4)); - EXPECT_NO_THROW(prefill->commit(tokens())); - EXPECT_FALSE(prefill->isDecoding()); - EXPECT_EQ(prefill->numCommittedTokens(), 4); - EXPECT_EQ(pageAt(*prefill), page); - EXPECT_EQ(page->cacheLevel, kHotLevel); - EXPECT_EQ(page->status(), PageStatus::LOCKED); - EXPECT_EQ(prefill->getBasePageIndices(LifeCycleId{0})[0], slotIdToPageIndexValue(page->slotId())); - EXPECT_EQ(decoder->getBasePageIndices(LifeCycleId{0})[0], slotIdToPageIndexValue(page->slotId())); - EXPECT_EQ(manager->storage().getStatistics(kHotLevel).free, manager->storage().getStatistics(kHotLevel).total - 1); - EXPECT_EQ(manager->storage().getStatistics(kSparseHistoryLevel).free, - manager->storage().getStatistics(kSparseHistoryLevel).total); - EXPECT_NO_THROW(prefill->close()); - ASSERT_TRUE(decoder->enterDecode()); - EXPECT_EQ(page->cacheLevel, kSparseHistoryLevel); -} - -TEST_F(KvCacheManagerV2DecodeOffloadTest, RebaseOomPreservesMissingLifecyclePagesAndCanRetryCommit) -{ - auto config = makeSplitColdGroupingConfig(); - for (auto& layer : config.layers) - std::get(layer).buffers.front().isSparse = true; - auto manager = std::make_shared(std::move(config)); - auto const apiLock = manager->lockExclusive(); - auto& storage = manager->storage(); - auto existing = seedPrefix(*manager, kHotLevel, LifeCycleId{1}); - auto cache = manager->createKvCache(); - auto closeCache = FuncGuard([&]() { cache->close(); }); - ASSERT_TRUE(cache->resume(stream())); - ASSERT_TRUE(cache->resize(4, 4)); - ASSERT_TRUE(cache->enterDecode()); - auto first = pageAt(*cache, 0, LifeCycleId{0}); - auto second = pageAt(*cache, 0, LifeCycleId{1}); - auto const version = cache->pageStorageVersion(); - auto const before = cache->getPageStorageSnapshot(LifeCycleId{0}); - ASSERT_TRUE(cache->acknowledgePageStorage(version)); - auto const pool = storage.getPoolGroupIndex(kSparseHistoryLevel, LifeCycleId{1}); - auto const freeHost = manager->getStorageStatistics(kSparseHistoryLevel).at(pool).free; - auto blockers = storage.newSlots(kSparseHistoryLevel, TypedVec{0, freeHost}); - auto releaseBlockers = FuncGuard( - [&]() - { - for (auto& slot : blockers[LifeCycleId{1}]) - storage.releaseSlot(LifeCycleId{1}, kSparseHistoryLevel, std::move(slot)); - }); - EXPECT_THROW(cache->commit(tokens()), OutOfPagesError); - EXPECT_EQ(cache->numCommittedTokens(), 0); - EXPECT_EQ(cache->numCommittedBlocks(), 0); - EXPECT_EQ(cache->historyLength(), 4); - EXPECT_GT(cache->pageStorageVersion(), version); - EXPECT_TRUE(cache->pageStorageDirty()); - EXPECT_FALSE(cache->acknowledgePageStorage(version)); - auto const after = cache->getPageStorageSnapshot(LifeCycleId{0}); - EXPECT_EQ(after.basePageIndices(), before.basePageIndices()); - EXPECT_EQ(after.cacheLevels(), before.cacheLevels()); - EXPECT_EQ(after.eligibleHistoryBlocks(), before.eligibleHistoryBlocks()); - EXPECT_EQ(pageAt(*cache, 0, LifeCycleId{0}), first); - EXPECT_EQ(pageAt(*cache, 0, LifeCycleId{1}), second); - for (auto const& page : {first, second}) - { - EXPECT_FALSE(page->isCommitted()); - EXPECT_EQ(page->cacheLevel, kSparseHistoryLevel); - EXPECT_EQ(page->status(), PageStatus::LOCKED); - EXPECT_EQ(cache->getBasePageIndices(page->lifeCycle)[0], slotIdToPageIndexValue(page->slotId())); - } - EXPECT_EQ(existing->cacheLevel, kHotLevel); - releaseBlockers.run(); - cache->commit(tokens()); - EXPECT_EQ(cache->numCommittedTokens(), 4); - EXPECT_TRUE(pageAt(*cache, 0, LifeCycleId{0})->isCommitted()); - EXPECT_EQ(pageAt(*cache, 0, LifeCycleId{1}), existing); - EXPECT_EQ(existing->cacheLevel, kSparseHistoryLevel); -} - -TEST_F(KvCacheManagerV2PageLockTest, SparseHostPrefixStaysPinnedAcrossReuseAndResume) -{ - auto manager = std::make_shared(sparseConfig()); - auto const apiLock = manager->lockExclusive(); - auto& storage = manager->storage(); - auto page = seedPrefix(*manager, kSparseHistoryLevel); - SlotId const hostSlot = page->slotId(); - auto cache = manager->createKvCache({}, tokens()); - auto closeCache = FuncGuard([&]() { cache->close(); }); - ASSERT_TRUE(cache->prefetch(kSparseHistoryLevel)); - ASSERT_TRUE(cache->resume(stream(), true)); - EXPECT_EQ(pageAt(*cache), page); - EXPECT_EQ(page->cacheLevel, kSparseHistoryLevel); - EXPECT_EQ(page->slotId(), hostSlot); - EXPECT_EQ(page->status(), PageStatus::LOCKED); - EXPECT_FALSE(page->scheduledForEviction()); - EXPECT_FALSE(storage.isEvictable(*page)); - EXPECT_THROW(storage.batchedMigrate(kHotLevel, {page}, {}), LogicError); - - auto second = manager->createKvCache({}, tokens()); - auto closeSecond = FuncGuard([&]() { second->close(); }); - ASSERT_TRUE(second->prefetch(kSparseHistoryLevel)); - ASSERT_TRUE(second->resume(stream(), true)); - EXPECT_EQ(pageAt(*second), page); - cache->suspend(); - EXPECT_EQ(page->status(), PageStatus::LOCKED); - second->close(); - EXPECT_EQ(page->status(), PageStatus::HELD); - ASSERT_TRUE(cache->prefetch(kHotLevel)); - ASSERT_TRUE(cache->resume()); - EXPECT_EQ(page->cacheLevel, kSparseHistoryLevel); - EXPECT_EQ(cache->getBasePageIndices(LifeCycleId{0})[0], slotIdToPageIndexValue(hostSlot)); - cache->close(); - EXPECT_EQ(page->status(), PageStatus::DROPPABLE); - EXPECT_TRUE(page->scheduledForEviction()); - - // Destruction returns the slot to the host pool, leaving the GPU pool untouched. - storage.excludeFromEviction(*page); - page.reset(); - for (CacheLevel level : {kHotLevel, kSparseHistoryLevel}) - { - auto const stats = manager->getStorageStatistics(level).at(PoolGroupIndex{0}); - EXPECT_EQ(stats.free, stats.total); - } -} - -TEST_F(KvCacheManagerV2PageLockTest, SparseGpuPrefixStaysOnGpuAcrossReuseAndResume) -{ - auto manager = std::make_shared(sparseConfig()); - auto const apiLock = manager->lockExclusive(); - auto page = seedPrefix(*manager, kHotLevel); - SlotId const gpuSlot = page->slotId(); - EXPECT_EQ(page->queryLockLevel(), kHotLevel); - auto cache = manager->createKvCache({}, tokens()); - auto closeCache = FuncGuard([&]() { cache->close(); }); - ASSERT_TRUE(cache->prefetch(kHotLevel)); - ASSERT_TRUE(cache->resume(stream())); - EXPECT_EQ(pageAt(*cache), page); - EXPECT_EQ(page->cacheLevel, kHotLevel); - EXPECT_EQ(page->slotId(), gpuSlot); - EXPECT_EQ(page->status(), PageStatus::LOCKED); - - cache->suspend(); - ASSERT_TRUE(cache->resume()); + // And back, which is the decode half. + storage.batchedMigrate(kHotLevel, {page}, {}); EXPECT_EQ(page->cacheLevel, kHotLevel); - EXPECT_EQ(page->slotId(), gpuSlot); - EXPECT_EQ(page->status(), PageStatus::LOCKED); - auto const hostStats = manager->getStorageStatistics(kSparseHistoryLevel).at(PoolGroupIndex{0}); - EXPECT_EQ(hostStats.free, hostStats.total); -} - -TEST_F(KvCacheManagerV2PageLockTest, DensePrefixStillRequiresGpuLock) -{ - auto manager = std::make_shared(makeTieredConfig()); - auto const apiLock = manager->lockExclusive(); - auto page = seedPrefix(*manager, kSparseHistoryLevel); - EXPECT_EQ(page->queryLockLevel(), kHotLevel); - auto holder = page->hold(); - EXPECT_THROW(makeShared(holder), LogicError); - auto cache = manager->createKvCache({}, tokens()); - auto closeCache = FuncGuard([&]() { cache->close(); }); - ASSERT_TRUE(cache->resume(stream())); - EXPECT_EQ(page->cacheLevel, kHotLevel); - EXPECT_EQ(page->status(), PageStatus::LOCKED); -} - -TEST_F(KvCacheManagerV2PageLockTest, PartialReuseCopiesSharedHostPrefixToPrivateGpuPage) -{ - auto manager = std::make_shared(sparseConfig()); - auto const apiLock = manager->lockExclusive(); - auto& storage = manager->storage(); - auto page = seedPrefix(*manager, kSparseHistoryLevel); - SlotId const hostSlot = page->slotId(); - PoolGroupIndex const hostPool = storage.getPoolGroupIndex(kSparseHistoryLevel, page->lifeCycle); - auto const hostAddress - = std::get(storage.slotAddress(kSparseHistoryLevel, hostPool, hostSlot, PoolIndex{0})); - size_t const bytes = storage.slotSize(kSparseHistoryLevel, hostPool).at(PoolIndex{0}); + EXPECT_EQ(codecPtr->decoded(layerGroup), 1U) << "sparse history must reach the codec on promotion"; page->readyEvent.synchronize(); - constexpr uint8_t kPattern = 0xA7; - std::memset(reinterpret_cast(hostAddress), kPattern, bytes); - - auto full = manager->createKvCache({}, tokens()); - auto closeFull = FuncGuard([&]() { full->close(); }); - ASSERT_TRUE(full->resume(stream(), true)); - auto partial = manager->createKvCache({}, tokens(2)); - auto closePartial = FuncGuard([&]() { partial->close(); }); - ASSERT_EQ(partial->historyLength(), 2); - ASSERT_TRUE(partial->resume(stream())); - auto privatePage = pageAt(*partial); - EXPECT_NE(privatePage, page); - EXPECT_FALSE(privatePage->isCommitted()); - EXPECT_EQ(privatePage->cacheLevel, kHotLevel); - EXPECT_EQ(page->cacheLevel, kSparseHistoryLevel); - EXPECT_EQ(page->slotId(), hostSlot); - EXPECT_EQ(page->status(), PageStatus::LOCKED); - EXPECT_EQ(full->getBasePageIndices(LifeCycleId{0})[0], slotIdToPageIndexValue(hostSlot)); - - auto const gpuPool = storage.getPoolGroupIndex(kHotLevel, privatePage->lifeCycle); - auto const gpuAddress - = std::get(storage.slotAddress(kHotLevel, gpuPool, privatePage->slotId(), PoolIndex{0})); - std::vector copied(bytes); - ASSERT_EQ(cudaMemcpyAsync( - copied.data(), reinterpret_cast(gpuAddress), bytes, cudaMemcpyDeviceToHost, mStream), - cudaSuccess); - ASSERT_EQ(cudaStreamSynchronize(mStream), cudaSuccess); - EXPECT_TRUE(std::all_of(copied.begin(), copied.end(), [](uint8_t byte) { return byte == kPattern; })); -} - -TEST_F(KvCacheManagerV2PageLockTest, DiskSparsePrefixRestoresToHostBeforeLocking) -{ - auto config = sparseConfig(); - config.cacheTiers.emplace_back(DiskCacheTierConfig{4 << 20, "/tmp"}); - auto manager = std::make_shared(std::move(config)); - auto const apiLock = manager->lockExclusive(); - auto page = seedPrefix(*manager, CacheLevel{2}); - EXPECT_EQ(page->queryLockLevel(), kSparseHistoryLevel); - EXPECT_THROW(makeShared(page->hold()), LogicError); - auto cache = manager->createKvCache({}, tokens()); - auto closeCache = FuncGuard([&]() { cache->close(); }); - ASSERT_TRUE(cache->resume(stream(), true)); - EXPECT_EQ(page->cacheLevel, kSparseHistoryLevel); - EXPECT_EQ(page->status(), PageStatus::LOCKED); - auto const gpuStats = manager->getStorageStatistics(kHotLevel).at(PoolGroupIndex{0}); - EXPECT_EQ(gpuStats.free, gpuStats.total); -} - -TEST_F(KvCacheManagerV2PageLockTest, WritableSparsePageRestoresToGpuButFullHistoryStaysOnHost) -{ - for (int const historyLength : {0, 4}) - { - SCOPED_TRACE(historyLength); - auto manager = std::make_shared(sparseConfig()); - auto const apiLock = manager->lockExclusive(); - auto& storage = manager->storage(); - auto cache = manager->createKvCache(); - auto closeCache = FuncGuard([&]() { cache->close(); }); - ASSERT_TRUE(cache->resume(stream())); - ASSERT_TRUE(cache->resize(4, historyLength)); - auto page = pageAt(*cache); - // Tier-aware locking alone must not demote a GPU page. - EXPECT_EQ(page->cacheLevel, kHotLevel); - cache->suspend(); - TypedVec evictOne(storage.numPoolGroups(kHotLevel), 1); - storage.forceEvict(kHotLevel, evictOne); - ASSERT_EQ(page->cacheLevel, kSparseHistoryLevel); - ASSERT_TRUE(cache->resume(std::nullopt, true)); - EXPECT_EQ(page->cacheLevel, historyLength == 0 ? kHotLevel : kSparseHistoryLevel); - EXPECT_EQ(page->status(), PageStatus::LOCKED); - } -} - -TEST_F(KvCacheManagerV2PageLockTest, ResizeOomRestoresOriginalHostLock) -{ - auto config = sparseConfig(); - std::get(config.layers.front()).slidingWindowSize = 4; - auto manager = std::make_shared(std::move(config)); - auto const apiLock = manager->lockExclusive(); - auto page = seedPrefix(*manager, kSparseHistoryLevel); - auto cache = manager->createKvCache({}, tokens()); - auto closeCache = FuncGuard([&]() { cache->close(); }); - ASSERT_TRUE(cache->resume(stream(), true)); - ASSERT_TRUE(cache->resize(8, 4)); - SlotId const hostSlot = page->slotId(); - auto const gpuFree = manager->getStorageStatistics(kHotLevel).at(PoolGroupIndex{0}).free; - - // Sparse SWA history stays pinned. A failed request for three GPU pages - // must preserve the host lock and the previous eligible-history count. - EXPECT_FALSE(cache->resize(20, 8)); - EXPECT_EQ(cache->capacity(), 8); - EXPECT_EQ(cache->historyLength(), 4); - EXPECT_EQ(pageAt(*cache), page); - EXPECT_EQ(page->cacheLevel, kSparseHistoryLevel); - EXPECT_EQ(page->slotId(), hostSlot); - EXPECT_EQ(page->status(), PageStatus::LOCKED); - EXPECT_FALSE(page->scheduledForEviction()); - EXPECT_EQ(manager->getStorageStatistics(kHotLevel).at(PoolGroupIndex{0}).free, gpuFree); -} - -TEST_F(KvCacheManagerV2PageLockTest, MixedSparseAndDensePrefixUsesSeparateLockLevels) -{ - auto config = sparseConfig(); - auto dense = std::get(config.layers.front()); - dense.layerId = 1; - dense.buffers.front().isSparse = false; - config.layers.emplace_back(std::move(dense)); - auto manager = std::make_shared(std::move(config)); - auto const apiLock = manager->lockExclusive(); - auto sparsePage = seedPrefix(*manager, kSparseHistoryLevel, LifeCycleId{0}); - auto densePage = seedPrefix(*manager, kSparseHistoryLevel, LifeCycleId{1}); - auto cache = manager->createKvCache({}, tokens()); - auto closeCache = FuncGuard([&]() { cache->close(); }); - ASSERT_TRUE(cache->resume(stream(), true)); - EXPECT_EQ(pageAt(*cache, 0, LifeCycleId{0}), sparsePage); - EXPECT_EQ(pageAt(*cache, 0, LifeCycleId{1}), densePage); - EXPECT_EQ(sparsePage->cacheLevel, kSparseHistoryLevel); - EXPECT_EQ(densePage->cacheLevel, kHotLevel); - EXPECT_EQ(sparsePage->status(), PageStatus::LOCKED); - EXPECT_EQ(densePage->status(), PageStatus::LOCKED); -} - -TEST_F(KvCacheManagerV2PageLockTest, CommitRebasesOntoSharedHostPrefix) -{ - auto manager = std::make_shared(sparseConfig()); - auto const apiLock = manager->lockExclusive(); - auto page = seedPrefix(*manager, kSparseHistoryLevel); - auto first = manager->createKvCache({}, tokens()); - auto closeFirst = FuncGuard([&]() { first->close(); }); - ASSERT_TRUE(first->resume(stream(), true)); - auto second = manager->createKvCache(); - auto closeSecond = FuncGuard([&]() { second->close(); }); - ASSERT_TRUE(second->resume(stream())); - ASSERT_TRUE(second->resize(4, 4)); - ASSERT_EQ(pageAt(*second)->cacheLevel, kHotLevel); - ASSERT_TRUE(second->enterDecode()); - second->commit(tokens()); - EXPECT_EQ(second->numCommittedBlocks(), 1); - EXPECT_EQ(pageAt(*second), page); - EXPECT_EQ(page->cacheLevel, kSparseHistoryLevel); - EXPECT_EQ(page->status(), PageStatus::LOCKED); - auto const gpuStats = manager->getStorageStatistics(kHotLevel).at(PoolGroupIndex{0}); - EXPECT_EQ(gpuStats.free, gpuStats.total); -} - -TEST_F(KvCacheManagerV2PageLockTest, ScratchSlotReturnsToGpuPool) -{ - auto manager = std::make_shared(sparseConfig()); - auto const apiLock = manager->lockExclusive(); - auto page = seedPrefix(*manager, kSparseHistoryLevel); - auto cache = manager->createKvCache({}, tokens()); - auto closeCache = FuncGuard([&]() { cache->close(); }); - ASSERT_TRUE(cache->resume(stream(), true)); - auto const gpuFree = manager->getStorageStatistics(kHotLevel).at(PoolGroupIndex{0}).free; - auto const hostFree = manager->getStorageStatistics(kSparseHistoryLevel).at(PoolGroupIndex{0}).free; - auto slots = manager->storage().newGpuSlots(TypedVec(LifeCycleId{1}, 1)); - ScratchSlotLock scratch(std::move(slots[LifeCycleId{0}].front()), *cache, LifeCycleId{0}); - EXPECT_EQ(manager->getStorageStatistics(kHotLevel).at(PoolGroupIndex{0}).free, gpuFree - 1); - { - auto scope = cache->recordEventScope(); - scratch.unlock(); - } - EXPECT_EQ(manager->getStorageStatistics(kHotLevel).at(PoolGroupIndex{0}).free, gpuFree); - EXPECT_EQ(manager->getStorageStatistics(kSparseHistoryLevel).at(PoolGroupIndex{0}).free, hostFree); } } // namespace diff --git a/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2ConcurrencyTest.cpp b/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2ConcurrencyTest.cpp index 00fae7fb7cdd..42bdbc4542f5 100644 --- a/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2ConcurrencyTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2ConcurrencyTest.cpp @@ -25,7 +25,6 @@ #include "tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.h" #include "tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.h" #include "tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.h" -#include "tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/funcGuard.h" #include #include @@ -114,42 +113,6 @@ TEST(KvCacheManagerV2ConcurrencyTest, LevelStatsAreConservedAcrossConcurrentDrai EXPECT_EQ(manager->getAndResetIterationDiskPrefetchBlocks(), 0); } -TEST(KvCacheManagerV2ConcurrencyTest, PageStorageSnapshotWaitsForSharedApiReaders) -{ - ASSERT_EQ(cudaSetDevice(0), cudaSuccess); - cudaStream_t stream{}; - ASSERT_EQ(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking), cudaSuccess); - auto destroyStream = FuncGuard([&]() { EXPECT_EQ(cudaStreamDestroy(stream), cudaSuccess); }); - auto config = makeTieredConfig(); - std::get(config.layers.front()).buffers.front().isSparse = true; - auto manager = std::make_shared(std::move(config)); - auto cache = manager->createKvCache(); - auto closeCache = FuncGuard([&]() { cache->close(); }); - ASSERT_TRUE(cache->resume(reinterpret_cast(stream))); - ASSERT_TRUE(cache->resize(4, 4)); - ASSERT_TRUE(cache->enterDecode()); - - std::promise started; - auto startedFuture = started.get_future(); - std::future snapshotFuture; - { - auto const apiLock = manager->lockShared(); - snapshotFuture = std::async(std::launch::async, - [&] - { - started.set_value(); - return cache->getPageStorageSnapshot(LifeCycleId{0}); - }); - EXPECT_EQ(startedFuture.wait_for(std::chrono::seconds{5}), std::future_status::ready); - // Inspecting a locked host page changes non-atomic holder refcounts, so even a - // shared API reader must exclude the snapshot until it releases its lock. - EXPECT_EQ(snapshotFuture.wait_for(std::chrono::milliseconds{100}), std::future_status::timeout); - } - auto const snapshot = snapshotFuture.get(); - EXPECT_EQ(snapshot.eligibleHistoryBlocks(), 1); - EXPECT_EQ(snapshot.cacheLevels(), (std::vector>{kSparseHistoryLevel})); -} - // Creates `count` roots and proposes all of them for erasure, leaving that many pending entries and // an unchanged root map. Returns the number of roots now present. size_t seedPendingRootErases(BlockRadixTree& tree, size_t count) diff --git a/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2ConfigTest.cpp b/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2ConfigTest.cpp index 99756c07befc..78995508390d 100644 --- a/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2ConfigTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2ConfigTest.cpp @@ -34,8 +34,13 @@ using tensorrt_llm::batch_manager::kv_cache_manager_v2::test::makeTieredConfig; AttentionLayerConfig makeAttentionLayer(LayerId layerId, bool isSparse, size_t bufferSize = 4096) { - return AttentionLayerConfig{ + AttentionLayerConfig layer{ .layerId = layerId, .buffers = {BufferConfig{.role = "key", .size = bufferSize, .isSparse = isSparse}}}; + if (isSparse) + { + layer.sparseCache = SparseCacheConfig{/*numSelectedBlocks=*/1, /*maxCachedUnits=*/2, /*unitTokens=*/4}; + } + return layer; } TEST(KvCacheManagerV2ConfigTest, ExistingBufferConfigurationDefaultsToDense) @@ -46,33 +51,26 @@ TEST(KvCacheManagerV2ConfigTest, ExistingBufferConfigurationDefaultsToDense) EXPECT_NO_THROW(makeConfig().validate()); } -TEST(KvCacheManagerV2ConfigTest, SparseBuffersRequireHostMemoryAtLevelOne) -{ - auto config = makeConfig(); - config.layers = {makeAttentionLayer(0, true)}; - EXPECT_THROW(config.validate(), std::invalid_argument); - - config.cacheTiers.emplace_back(DiskCacheTierConfig{4 << 20, "/tmp"}); - EXPECT_THROW(config.validate(), std::invalid_argument); - - config.cacheTiers[1] = GpuCacheTierConfig{4 << 20}; - config.cacheTiers.emplace_back(HostCacheTierConfig{4 << 20}); - EXPECT_THROW(config.validate(), std::invalid_argument); - - config.cacheTiers[1] = HostCacheTierConfig{4 << 20}; - EXPECT_NO_THROW(config.validate()); -} - -TEST(KvCacheManagerV2ConfigTest, MixedBuffersInOneLayerAreRejected) +TEST(KvCacheManagerV2ConfigTest, OneLayerMayMixSparseAndDenseBuffers) { + // Sparsity picks a buffer's pool and nothing above it, so an indexer and the data it + // selects over stay in the one layer they belong to. auto config = makeTieredConfig(); for (bool const firstIsSparse : {false, true}) { auto layer = makeAttentionLayer(0, firstIsSparse); layer.buffers.push_back(BufferConfig{.role = "value", .size = 4096, .isSparse = !firstIsSparse}); config.layers = {layer}; - EXPECT_THROW(config.validate(), std::invalid_argument); - EXPECT_THROW(createStorageConfig(config), std::invalid_argument); + ASSERT_NO_THROW(config.validate()); + auto const storage = createStorageConfig(config); + auto const attributes = storage.bufferAttributes(); + auto const& key = attributes.at(BufferId{0, "key"}); + auto const& value = attributes.at(BufferId{0, "value"}); + // One life cycle, one pool group, but a pool each: equal sizes do not coalesce + // across sparsity, because the two pools differ in medium. + EXPECT_EQ(key.lifeCycleId, value.lifeCycleId); + EXPECT_EQ(storage.slotDescList.size(), PoolGroupIndex{1}); + EXPECT_NE(key.poolIndex, value.poolIndex); } } @@ -86,7 +84,7 @@ TEST(KvCacheManagerV2ConfigTest, SparseSsmBuffersAreRejected) EXPECT_THROW(createStorageConfig(config), std::invalid_argument); } -TEST(KvCacheManagerV2ConfigTest, SparseAndDenseBuffersHaveSeparateLifecyclesAndPools) +TEST(KvCacheManagerV2ConfigTest, SparseAndDenseLayersShareALifeCycleButNotAPool) { auto config = makeTieredConfig(); config.layers = {makeAttentionLayer(0, false), makeAttentionLayer(1, true), makeAttentionLayer(2, false), @@ -100,24 +98,24 @@ TEST(KvCacheManagerV2ConfigTest, SparseAndDenseBuffersHaveSeparateLifecyclesAndP auto const& dense = attributes.at(BufferId{0, "key"}); auto const& sparse = attributes.at(BufferId{1, "key"}); - EXPECT_EQ(registry.size(), LifeCycleId{2}); - EXPECT_EQ(storage.slotDescList.size(), PoolGroupIndex{2}); - EXPECT_NE(dense.lifeCycleId, sparse.lifeCycleId); - EXPECT_NE(grouping[dense.lifeCycleId], grouping[sparse.lifeCycleId]); - EXPECT_FALSE(std::get(registry[dense.lifeCycleId]).isSparse); - EXPECT_TRUE(std::get(registry[sparse.lifeCycleId]).isSparse); - EXPECT_FALSE(registry[dense.lifeCycleId] == registry[sparse.lifeCycleId]); + // Sparsity is not part of a life cycle's identity, so layers alike in window and sinks + // share one however their buffers are backed. + EXPECT_EQ(registry.size(), LifeCycleId{1}); + EXPECT_EQ(storage.slotDescList.size(), PoolGroupIndex{1}); + EXPECT_EQ(dense.lifeCycleId, sparse.lifeCycleId); + EXPECT_EQ(grouping[dense.lifeCycleId], grouping[sparse.lifeCycleId]); + // Separate pools, because only one of them is host-backed. + EXPECT_NE(dense.poolIndex, sparse.poolIndex); auto const& otherDense = attributes.at(BufferId{2, "key"}); auto const& otherSparse = attributes.at(BufferId{3, "key"}); - EXPECT_EQ(dense.lifeCycleId, otherDense.lifeCycleId); - EXPECT_EQ(sparse.lifeCycleId, otherSparse.lifeCycleId); + EXPECT_EQ(dense.poolIndex, otherDense.poolIndex); EXPECT_EQ(sparse.poolIndex, otherSparse.poolIndex); EXPECT_EQ(otherSparse.offset, sparse.offset + sparse.size); EXPECT_EQ(otherDense.offset, dense.offset + dense.size); } -TEST(KvCacheManagerV2ConfigTest, MatchingSlotShapesSharePoolsWithinEachSparsity) +TEST(KvCacheManagerV2ConfigTest, MatchingSlotShapesSharePoolsAcrossWindows) { auto config = makeTieredConfig(); config.layers.clear(); @@ -133,18 +131,26 @@ TEST(KvCacheManagerV2ConfigTest, MatchingSlotShapesSharePoolsWithinEachSparsity) auto const attributes = storage.bufferAttributes(); auto const grouping = storage.lifeCycleGrouping(); auto poolGroup = [&](LayerId layerId) { return grouping[attributes.at(BufferId{layerId, "key"}).lifeCycleId]; }; - EXPECT_EQ(storage.numLifeCycles(), LifeCycleId{4}); - EXPECT_EQ(storage.slotDescList.size(), PoolGroupIndex{2}); + // Two life cycles, one per window; sparsity is not part of their identity, so layer 0 + // (dense) and layer 2 (sparse) share one. Both life cycles present the same pools -- + // a dense 4096 and a sparse 4096 -- so one pool group serves them. + EXPECT_EQ(storage.numLifeCycles(), LifeCycleId{2}); + EXPECT_EQ(storage.slotDescList.size(), PoolGroupIndex{1}); EXPECT_EQ(poolGroup(0), poolGroup(1)); + EXPECT_EQ(poolGroup(0), poolGroup(2)); EXPECT_EQ(poolGroup(2), poolGroup(3)); - EXPECT_NE(poolGroup(0), poolGroup(2)); + // Sparsity still separates the pools within that group. + auto poolOf = [&](LayerId layerId) { return attributes.at(BufferId{layerId, "key"}).poolIndex; }; + EXPECT_NE(poolOf(0), poolOf(2)); } -TEST(KvCacheManagerV2ConfigTest, SparseBuffersCoalesceByExpandedSize) +TEST(KvCacheManagerV2ConfigTest, SparseBuffersCoalesceBySize) { + // A sparse buffer's `size` is its whole page, as for any other buffer, so two of equal + // size coalesce however finely each divides its page into selection units. auto config = makeTieredConfig(); - auto firstLayer = makeAttentionLayer(0, true, 2048); - firstLayer.buffers.front().tokensPerBlockOverride = 2; + auto firstLayer = makeAttentionLayer(0, true, 4096); + firstLayer.sparseCache = SparseCacheConfig{/*numSelectedBlocks=*/1, /*maxCachedUnits=*/2, /*unitTokens=*/2}; config.layers = {firstLayer, makeAttentionLayer(1, true, 4096)}; ASSERT_NO_THROW(config.validate()); @@ -158,7 +164,8 @@ TEST(KvCacheManagerV2ConfigTest, SparseBuffersCoalesceByExpandedSize) EXPECT_EQ(first.poolIndex, second.poolIndex); EXPECT_EQ(first.size, second.size); EXPECT_EQ(second.offset, first.offset + first.size); - EXPECT_EQ(first.expansion, 2); + // tokensPerBlockOverride is the index converter's, and a sparse buffer does not use it. + EXPECT_EQ(first.expansion, 1); EXPECT_EQ(second.expansion, 1); } diff --git a/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2GeneratorTest.cpp b/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2GeneratorTest.cpp new file mode 100644 index 000000000000..612ec32a9e4d --- /dev/null +++ b/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2GeneratorTest.cpp @@ -0,0 +1,208 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/generator.h" + +#include + +#include +#include +#include +#include + +namespace +{ + +using namespace tensorrt_llm::batch_manager::kv_cache_manager_v2; + +//! Sets the flag it was given when it goes, so a test can see whether a frame was unwound. +class Flag +{ +public: + explicit Flag(bool* destroyed) noexcept + : mDestroyed(destroyed) + { + } + + Flag(Flag const&) = delete; + Flag(Flag&&) = delete; + Flag& operator=(Flag const&) = delete; + Flag& operator=(Flag&&) = delete; + + ~Flag() + { + *mDestroyed = true; + } + +private: + bool* mDestroyed; +}; + +Generator evens(int limit) +{ + for (int value = 0; value < limit; value += 2) + { + co_yield value; + } +} + +// The bodies that need state are free functions taking a pointer, not capturing lambdas: a +// coroutine copies its parameters into the frame, but a lambda's closure is only referred to, +// so an immediately-invoked capturing lambda leaves the frame pointing at a destroyed object. +Generator marksEntry(bool* entered) +{ + *entered = true; + co_yield 1; +} + +Generator holdsAFlagAcrossTwoYields(bool* destroyed) +{ + Flag const flag{destroyed}; + static_cast(co_yield 1); + static_cast(co_yield 2); +} + +Generator holdsAFlagFromItsFirstYield(bool* destroyed) +{ + Flag const flag{destroyed}; + co_yield 1; +} + +TEST(KvCacheManagerV2Generator, OneWayYieldsLazilyAndInOrder) +{ + bool entered = false; + auto body = marksEntry(&entered); + EXPECT_FALSE(entered) << "the body ran before the first pull"; + EXPECT_EQ(body.value(), 1); + EXPECT_TRUE(entered); + + std::vector seen; + for (auto value : evens(7)) + { + seen.push_back(value); + } + EXPECT_EQ(seen, (std::vector{0, 2, 4, 6})); +} + +TEST(KvCacheManagerV2Generator, OneWayOverAnEmptySequenceYieldsNothing) +{ + auto sequence = evens(0); + EXPECT_EQ(sequence.begin(), sequence.end()); + EXPECT_TRUE(sequence.done()); +} + +TEST(KvCacheManagerV2Generator, TwoWayCarriesAValueBackIntoTheBody) +{ + // Two rounds, so the second request is answered from what the first was given. + auto body = []() -> Generator + { + int const first = co_yield 10; + int const second = co_yield first + 1; + static_cast(co_yield second * 2); + }(); + + ASSERT_EQ(body.value(), 10); + body.send(5); + ASSERT_EQ(body.value(), 6); + body.send(7); + EXPECT_EQ(body.value(), 14); +} + +TEST(KvCacheManagerV2Generator, TwoWayReportsDoneOnlyAfterTheBodyReturns) +{ + auto body = []() -> Generator { static_cast(co_yield 1); }(); + + ASSERT_FALSE(body.done()); + body.send(0); + EXPECT_TRUE(body.done()); + EXPECT_FALSE(body.value().has_value()) << "a finished body is told apart by having no value"; +} + +TEST(KvCacheManagerV2Generator, ReturnValueArrivesOnlyWhenTheBodyFinishes) +{ + auto body = []() -> Generator + { + int const given = co_yield 1; + co_return std::string(static_cast(given), 'x'); + }(); + + ASSERT_EQ(body.value(), 1); + EXPECT_FALSE(body.result().has_value()); + body.send(3); + ASSERT_TRUE(body.result().has_value()); + EXPECT_EQ(*body.result(), "xxx"); +} + +TEST(KvCacheManagerV2Generator, ExceptionSurfacesOutOfWhicheverCallResumedTheBody) +{ + auto beforeFirstYield = []() -> Generator + { + throw std::runtime_error("priming"); + co_return; + }(); + EXPECT_THROW((void) beforeFirstYield.value(), std::runtime_error); + + auto afterFirstYield = []() -> Generator + { + static_cast(co_yield 1); + throw std::runtime_error("resuming"); + }(); + ASSERT_EQ(afterFirstYield.value(), 1); + EXPECT_THROW(afterFirstYield.send(0), std::runtime_error); +} + +TEST(KvCacheManagerV2Generator, DestroyingASuspendedBodyRunsItsLocals) +{ + // A layout body holds its guards and buffers across the yield, so abandoning it -- which + // is what a throw between asking and answering does -- has to release them. + bool destroyed = false; + { + auto body = holdsAFlagAcrossTwoYields(&destroyed); + ASSERT_EQ(body.value(), 1); + ASSERT_FALSE(destroyed); + } + EXPECT_TRUE(destroyed); +} + +TEST(KvCacheManagerV2Generator, DestroyingAnUnstartedBodyRunsNothing) +{ + bool destroyed = false; + { + auto body = holdsAFlagFromItsFirstYield(&destroyed); + (void) body; + } + EXPECT_FALSE(destroyed) << "the local is not alive until the body has run to its first yield"; +} + +TEST(KvCacheManagerV2Generator, MovingAPrimedBodyCarriesThatItHasStarted) +{ + // Moving an already-suspended body and then resuming it must not re-prime it: that would + // resume it with nothing sent, and the body would read whatever the empty slot held. + auto body = []() -> Generator + { + int const given = co_yield 1; + static_cast(co_yield given * 10); + }(); + ASSERT_EQ(body.value(), 1); + + std::vector> moved; + moved.push_back(std::move(body)); + moved.back().send(4); + EXPECT_EQ(moved.back().value(), 40); +} + +} // namespace diff --git a/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2HostMemBackingTest.cpp b/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2HostMemBackingTest.cpp index 98e15e0cff15..e40c860ff354 100644 --- a/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2HostMemBackingTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2HostMemBackingTest.cpp @@ -40,17 +40,16 @@ class HostMemBackingTest : public ::testing::Test void SetUp() override { ASSERT_EQ(cudaSetDevice(0), cudaSuccess); - // Forces primary context creation; reserve() captures the current - // context and refuses to run without one. + // Forces primary context creation; the backing captures the current + // context and refuses to be built without one. ASSERT_EQ(cudaFree(nullptr), cudaSuccess); - mBacking = std::make_unique(HostMemBackingOptions{kCommitUnit}); - // Deliberately larger than the allocation unit, so the span a case - // commits in differs from the unit the backing allocates in. - mChunkSize = mBacking->commitGranularity() * kUnitsPerChunk; + // Deliberately larger than the allocation unit, so the span a case moves + // the target by differs from the unit the backing allocates in. + mChunkSize = hostMemSizeGranularity() * kUnitsPerChunk; mTotalSize = mChunkSize * kNumChunks; + mBacking = std::make_unique(mTotalSize); } - static constexpr size_t kCommitUnit = size_t{2} << 20; static constexpr size_t kUnitsPerChunk = 4; static constexpr size_t kNumChunks = 4; @@ -121,9 +120,9 @@ bool backingHoldsMemory(MemAddress addr) //! released. TEST_F(HostMemBackingTest, DestroyReleasesTheMemory) { - MemAddress const base = mBacking->reserve(mTotalSize); + MemAddress const base = mBacking->address(); ASSERT_NE(base, 0U); - mBacking->commit(0, mChunkSize); + mBacking->resizeTo(mChunkSize); ASSERT_TRUE(backingHoldsMemory(base)) << "nothing was held to begin with"; mBacking->destroy(); @@ -136,9 +135,9 @@ TEST_F(HostMemBackingTest, DestroyReleasesTheMemory) //! is already permanently unusable. TEST_F(HostMemBackingTest, DestroyUnderPoisonKeepsTheMemory) { - MemAddress const base = mBacking->reserve(mTotalSize); + MemAddress const base = mBacking->address(); ASSERT_NE(base, 0U); - mBacking->commit(0, mChunkSize); + mBacking->resizeTo(mChunkSize); ASSERT_TRUE(backingHoldsMemory(base)) << "nothing was held to begin with"; Poison::set("DestroyUnderPoisonKeepsTheMemory", "simulated invariant violation"); @@ -150,12 +149,12 @@ TEST_F(HostMemBackingTest, DestroyUnderPoisonKeepsTheMemory) TEST_F(HostMemBackingTest, CommittedRangeIsReadableAndWritable) { - MemAddress const base = mBacking->reserve(mTotalSize); + MemAddress const base = mBacking->address(); ASSERT_NE(base, 0U); for (size_t chunk = 0; chunk < kNumChunks; ++chunk) { - mBacking->commit(chunk * mChunkSize, mChunkSize); + mBacking->resizeTo((chunk + 1) * mChunkSize); stamp(base, chunk * mChunkSize, mChunkSize, 0); } for (size_t chunk = 0; chunk < kNumChunks; ++chunk) @@ -168,15 +167,15 @@ TEST_F(HostMemBackingTest, CommittedRangeIsReadableAndWritable) //! hands out pointers derived from it before the fill has finished. TEST_F(HostMemBackingTest, BaseAddressIsStableAcrossCommits) { - MemAddress const base = mBacking->reserve(mTotalSize); + MemAddress const base = mBacking->address(); ASSERT_NE(base, 0U); - mBacking->commit(0, mChunkSize); + mBacking->resizeTo(mChunkSize); stamp(base, 0, mChunkSize, 0); for (size_t chunk = 1; chunk < kNumChunks; ++chunk) { - mBacking->commit(chunk * mChunkSize, mChunkSize); + mBacking->resizeTo((chunk + 1) * mChunkSize); // Data committed earlier survives later commits, and stays where it was. ASSERT_TRUE(verify(base, 0, mChunkSize, 0)) << "after committing chunk " << chunk; } @@ -186,18 +185,19 @@ TEST_F(HostMemBackingTest, BaseAddressIsStableAcrossCommits) //! alone: data written between populate() and commit() survives it. TEST_F(HostMemBackingTest, CommitOfPopulatedRangeKeepsContents) { - MemAddress const base = mBacking->reserve(mTotalSize); + MemAddress const base = mBacking->address(); ASSERT_NE(base, 0U); mBacking->open(0, mTotalSize); mBacking->populate(0, mTotalSize); stamp(base, 0, mTotalSize, 3); - for (size_t chunk = 0; chunk < kNumChunks; ++chunk) + mBacking->setTargetSize(mTotalSize); + while (mBacking->committedSize() != mBacking->targetSize()) { - mBacking->commit(chunk * mChunkSize, mChunkSize, /*populated=*/true); + mBacking->resizeStep(/*populated=*/true); } EXPECT_TRUE(verify(base, 0, mTotalSize, 3)); - for (size_t offset = 0; offset < mTotalSize; offset += kCommitUnit) + for (size_t offset = 0; offset < mTotalSize; offset += mBacking->registerUnit()) { unsigned int flags = 0; ASSERT_EQ(cudaHostGetFlags(&flags, reinterpret_cast(base + offset)), cudaSuccess) @@ -207,20 +207,17 @@ TEST_F(HostMemBackingTest, CommitOfPopulatedRangeKeepsContents) TEST_F(HostMemBackingTest, DecommitFreesAndRecommitWorks) { - MemAddress const base = mBacking->reserve(mTotalSize); + MemAddress const base = mBacking->address(); ASSERT_NE(base, 0U); - for (size_t chunk = 0; chunk < kNumChunks; ++chunk) - { - mBacking->commit(chunk * mChunkSize, mChunkSize); - } + mBacking->resizeTo(mTotalSize); stamp(base, 0, mTotalSize, 0); // Drop the tail, then bring it back. Contents need not survive, but the // range must be usable again and the base must not have moved. size_t const tailOffset = mChunkSize * (kNumChunks - 1); - mBacking->decommit(tailOffset, mChunkSize); - mBacking->commit(tailOffset, mChunkSize); + mBacking->resizeTo(tailOffset); + mBacking->resizeTo(mTotalSize); stamp(base, tailOffset, mChunkSize, 7); EXPECT_TRUE(verify(base, tailOffset, mChunkSize, 7)); @@ -233,20 +230,20 @@ TEST_F(HostMemBackingTest, DecommitFreesAndRecommitWorks) TEST_F(HostMemBackingTest, LargeCommitIsReleasablePiecewise) { size_t const unit = mBacking->commitGranularity(); - MemAddress const base = mBacking->reserve(mTotalSize); + MemAddress const base = mBacking->address(); ASSERT_NE(base, 0U); - // One commit spanning many allocation units. - mBacking->commit(0, mTotalSize); + // One target spanning many allocation units. + mBacking->resizeTo(mTotalSize); stamp(base, 0, mTotalSize, 4); - // Release a single unit from the end, well inside what was one commit call. + // Release a single unit from the end, well inside what one step populated. size_t const tail = mTotalSize - unit; - mBacking->decommit(tail, unit); + mBacking->resizeTo(tail); EXPECT_TRUE(verify(base, 0, tail, 4)) << "releasing one unit disturbed the rest"; // And take it back. - mBacking->commit(tail, unit); + mBacking->resizeTo(mTotalSize); stamp(base, tail, unit, 8); EXPECT_TRUE(verify(base, tail, unit, 8)); EXPECT_TRUE(verify(base, 0, tail, 4)) << "recommit clobbered the retained range"; @@ -254,58 +251,211 @@ TEST_F(HostMemBackingTest, LargeCommitIsReleasablePiecewise) TEST_F(HostMemBackingTest, DestroyIsIdempotent) { - MemAddress const base = mBacking->reserve(mTotalSize); + MemAddress const base = mBacking->address(); ASSERT_NE(base, 0U); - mBacking->commit(0, mChunkSize); + mBacking->resizeTo(mChunkSize); mBacking->destroy(); mBacking->destroy(); SUCCEED(); } -//! The driver rejects a copy whose host range spans two registrations. Nothing -//! reports this at registration time -- it surfaces only when the copy is made -- -//! so it is pinned down here: it is the reason the commit unit is a whole number -//! of slots, and a regression would otherwise show up as a failed eviction. -TEST_F(HostMemBackingTest, CopyAcrossAUnitBoundaryIsRejected) +//! The constraint the whole register-unit scheme exists for: the driver rejects a copy that +//! spans two registrations. Registrations are chunked well above a commit unit, so this +//! needs a register unit large enough to make the chunk small enough to straddle. +TEST_F(HostMemBackingTest, CopyAcrossARegistrationBoundaryIsRejected) { - MemAddress const base = mBacking->reserve(mTotalSize); + // Above the chunk the backing aims for, so a registration is exactly one register unit and + // the boundary is known here without naming the backing's own constant. + constexpr size_t kRegisterUnit = size_t{40} << 20; + size_t const total = 2 * kRegisterUnit; + HostMemBacking backing(total, HostMemBackingOptions{.registerUnit = kRegisterUnit}); + MemAddress const base = backing.address(); ASSERT_NE(base, 0U); - mBacking->commit(0, mChunkSize); - size_t const unit = mBacking->commitGranularity(); - ASSERT_GE(mChunkSize, 2 * unit) << "a chunk must hold at least two units for this to mean anything"; + backing.resizeTo(total); void* deviceBuffer = nullptr; - ASSERT_EQ(cudaMalloc(&deviceBuffer, 2 * unit), cudaSuccess); - - // Wholly inside the first unit: accepted. - EXPECT_EQ(cudaMemcpy(deviceBuffer, reinterpret_cast(base), unit, cudaMemcpyHostToDevice), cudaSuccess); - // Straddling the first boundary: rejected. - EXPECT_NE(cudaMemcpy(deviceBuffer, reinterpret_cast(base + unit / 2), unit, cudaMemcpyHostToDevice), + ASSERT_EQ(cudaMalloc(&deviceBuffer, kRegisterUnit), cudaSuccess); + + // Wholly inside the first registration: accepted. + EXPECT_EQ(cudaMemcpy(deviceBuffer, reinterpret_cast(base), kRegisterUnit, cudaMemcpyHostToDevice), + cudaSuccess); + // Straddling the boundary: rejected. + EXPECT_NE(cudaMemcpy(deviceBuffer, reinterpret_cast(base + kRegisterUnit / 2), kRegisterUnit, + cudaMemcpyHostToDevice), cudaSuccess) << "a copy spanning two registrations should not be accepted"; cudaGetLastError(); // clear the sticky error the rejected copy leaves EXPECT_EQ(cudaFree(deviceBuffer), cudaSuccess); + backing.destroy(); } -//! The point of the whole exercise: the GPU must reach the committed range. -TEST_F(HostMemBackingTest, GpuRoundTrip) +//! A register unit that does not divide the commit unit puts registration bounds +//! somewhere other than the unit edges, so a decommit can land inside one. The +//! registration has to go -- it cannot be shortened -- and its surviving prefix +//! be registered again, or it pins pages madvise is about to discard. +TEST_F(HostMemBackingTest, DecommitSplitsTheRegistrationItLandsInside) { - MemAddress const base = mBacking->reserve(mTotalSize); + constexpr size_t kRegisterUnit = 9216; // a whole number of neither pages nor commit units + size_t const unit = hostMemSizeGranularity(); + ASSERT_NE(unit % kRegisterUnit, 0U) + << "the register unit must not divide the commit unit for this to mean anything"; + + size_t const total = unit * 4; + HostMemBacking backing(total, HostMemBackingOptions{.registerUnit = kRegisterUnit}); + MemAddress const base = backing.address(); ASSERT_NE(base, 0U); - for (size_t chunk = 0; chunk < kNumChunks; ++chunk) + backing.resizeTo(total); + stamp(base, 0, total, 7); + + backing.resizeTo(unit * 2); + + void* deviceBuffer = nullptr; + ASSERT_EQ(cudaMalloc(&deviceBuffer, kRegisterUnit), cudaSuccess); + // The last register unit that still fits below the new end: it sits in the + // registration the decommit split, so reaching it proves the prefix came back. + size_t const survivor = (unit * 2 / kRegisterUnit - 1) * kRegisterUnit; + EXPECT_EQ( + cudaMemcpy(deviceBuffer, reinterpret_cast(base + survivor), kRegisterUnit, cudaMemcpyHostToDevice), + cudaSuccess) + << "the prefix below the decommit should still be registered"; + cudaGetLastError(); + EXPECT_EQ(cudaFree(deviceBuffer), cudaSuccess); + + // The half that was dropped takes commits again, and the GPU must see what the + // host writes there. A registration left over the discarded pages would still + // name the physical pages madvise took, so this reads back stale bytes rather + // than failing -- which is why the check is on the data and not on the status. + backing.resizeTo(unit * 4); + stamp(base, unit * 2, unit * 2, 9); + + size_t const recommitted = (unit * 2 / kRegisterUnit + 1) * kRegisterUnit; + ASSERT_LT(recommitted + kRegisterUnit, unit * 4); + void* checkBuffer = nullptr; + ASSERT_EQ(cudaMalloc(&checkBuffer, kRegisterUnit), cudaSuccess); + ASSERT_EQ(cudaMemcpy(checkBuffer, reinterpret_cast(base + recommitted), kRegisterUnit, + cudaMemcpyHostToDevice), + cudaSuccess); + std::vector readBack(kRegisterUnit, 0); + ASSERT_EQ(cudaMemcpy(readBack.data(), checkBuffer, kRegisterUnit, cudaMemcpyDeviceToHost), cudaSuccess); + for (size_t i = 0; i < kRegisterUnit; ++i) { - mBacking->commit(chunk * mChunkSize, mChunkSize); + ASSERT_EQ(readBack[i], static_cast((recommitted + i + 9) & 0xFF)) + << "byte " << i << " of the recommitted range"; } + EXPECT_EQ(cudaFree(checkBuffer), cudaSuccess); + backing.destroy(); +} + +//! A step that knows the target is still ahead must stop registering on a chunk +//! boundary. Registering whatever it happened to populate would leave a remainder +//! pinned on its own, and the next step would start from there, so a long fill +//! would end up with twice the registrations it needs. +TEST_F(HostMemBackingTest, AStepShortOfTheTargetRegistersWholeChunksOnly) +{ + constexpr size_t kRegisterUnit = 9216; // divides neither a page nor a commit chunk + // Several commit steps' worth, so most steps have a successor. + size_t const total = hostMemSizeGranularity() * 64; + HostMemBacking backing(total, HostMemBackingOptions{.registerUnit = kRegisterUnit}); + ASSERT_NE(backing.address(), 0U); + + backing.setTargetSize(total); + size_t chunk = 0; + while (backing.committedSize() != total) + { + size_t const registered = backing.resizeStep(); + if (registered == 0) + { + continue; + } + if (chunk == 0) + { + chunk = registered; + } + if (backing.committedSize() != total) + { + EXPECT_EQ(registered % chunk, 0U) << "a step left a partial registration behind"; + } + } + EXPECT_GT(chunk, 0U) << "nothing was registered at all"; + EXPECT_EQ(chunk % kRegisterUnit, 0U) << "a registration bound fell inside a register unit"; + // The final step registers the remainder, down to the last whole register unit. + EXPECT_EQ(backing.registeredSize(), total / kRegisterUnit * kRegisterUnit); +} + +//! A fill resumed after a target that left a partial tail starts its first +//! registration where that tail ended, which is the one start that does not sit +//! on a chunk boundary and so cannot be derived. A later shrink has to find it +//! anyway: unregistering at anything other than a registration's own base is +//! rejected by the driver, which poisons KVCM2. +TEST_F(HostMemBackingTest, ShrinkFindsARegistrationThatResumedOffAChunkBoundary) +{ + constexpr size_t kRegisterUnit = 9216; + size_t const unit = hostMemSizeGranularity(); + size_t const total = unit * 64; + // Not a whole number of register units, so the first fill stops short of its + // target and the second resumes from there rather than from a boundary. + size_t const first = unit * 3; + ASSERT_NE(first % kRegisterUnit, 0U); + + HostMemBacking backing(total, HostMemBackingOptions{.registerUnit = kRegisterUnit}); + MemAddress const base = backing.address(); + ASSERT_NE(base, 0U); + + backing.resizeTo(first); + size_t const resumeFrom = backing.registeredSize(); + ASSERT_LT(resumeFrom, first) << "the first target was meant to leave a partial tail"; + backing.resizeTo(total); + + // Down past the resumed registration's base and back up, twice, so a start + // the object has mislaid shows up as a failed unregister rather than as + // memory that merely happens to still work. + backing.resizeTo(first); + backing.resizeTo(total); + backing.resizeTo(unit); + backing.resizeTo(total); + ASSERT_FALSE(Poison::poisoned()) << "a registration was dropped at the wrong address"; + EXPECT_EQ(backing.registeredSize(), total / kRegisterUnit * kRegisterUnit); + + // The register units either side of where the fill resumed must still hold + // what the host writes, which a registration left over discarded pages would + // not: it would name the physical pages madvise took. + stamp(base, 0, total, 5); + void* deviceBuffer = nullptr; + ASSERT_EQ(cudaMalloc(&deviceBuffer, kRegisterUnit), cudaSuccess); + std::vector readBack(kRegisterUnit, 0); + for (size_t index = resumeFrom / kRegisterUnit - 1; index <= resumeFrom / kRegisterUnit + 1; ++index) + { + size_t const offset = index * kRegisterUnit; + ASSERT_EQ(cudaMemcpy(deviceBuffer, reinterpret_cast(base + offset), kRegisterUnit, + cudaMemcpyHostToDevice), + cudaSuccess) + << "register unit at " << offset; + ASSERT_EQ(cudaMemcpy(readBack.data(), deviceBuffer, kRegisterUnit, cudaMemcpyDeviceToHost), cudaSuccess); + for (size_t i = 0; i < kRegisterUnit; ++i) + { + ASSERT_EQ(readBack[i], static_cast((offset + i + 5) & 0xFF)) << "byte " << i << " at " << offset; + } + } + EXPECT_EQ(cudaFree(deviceBuffer), cudaSuccess); + backing.destroy(); + EXPECT_FALSE(Poison::poisoned()); +} + +//! The point of the whole exercise: the GPU must reach the committed range. +TEST_F(HostMemBackingTest, GpuRoundTrip) +{ + MemAddress const base = mBacking->address(); + ASSERT_NE(base, 0U); + mBacking->resizeTo(mTotalSize); stamp(base, 0, mTotalSize, 3); void* deviceBuffer = nullptr; ASSERT_EQ(cudaMalloc(&deviceBuffer, mTotalSize), cudaSuccess); - // One copy per unit. Each unit is a separate registration and the driver - // rejects a copy that straddles two, which is why the pools size their unit - // to a whole number of slots: every copy they make is one slot. + // One copy per unit, each well inside a registration. A caller whose copies are + // larger declares them as the register unit, which keeps the bounds off them. size_t const unit = mBacking->commitGranularity(); for (size_t offset = 0; offset < mTotalSize; offset += unit) { diff --git a/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2HostMemTest.cpp b/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2HostMemTest.cpp index ac959de6b815..e726bd5c4644 100644 --- a/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2HostMemTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2HostMemTest.cpp @@ -104,9 +104,9 @@ void* countMemset(void* ptr, int value, size_t size) TEST(KvCacheManagerV2HostMemTest, SelectsConfiguredPageMode) { - hostMadvisePageMode(MemAddress{1}, HostMem::kAlignment, true, captureMadvise); + hostMadvisePageMode(MemAddress{1}, HostMem::sizeGranularity(), true, captureMadvise); EXPECT_EQ(gCapturedAdvice, MADV_HUGEPAGE); - hostMadvisePageMode(MemAddress{1}, HostMem::kAlignment, false, captureMadvise); + hostMadvisePageMode(MemAddress{1}, HostMem::sizeGranularity(), false, captureMadvise); EXPECT_EQ(gCapturedAdvice, MADV_NOHUGEPAGE); ScopedEnv defaultThp("TLLM_KV_CACHE_MANAGER_V2_THP", std::nullopt); @@ -127,7 +127,7 @@ class PrefaultFallbackTest : public testing::TestWithParam TEST_P(PrefaultFallbackTest, TouchesMemory) { - std::vector data(HostMem::kAlignment, 0xFF); + std::vector data(HostMem::sizeGranularity(), 0xFF); gMadviseErrno = GetParam(); gMemsetCalls = 0; hostPrefaultChunk(reinterpret_cast(data.data()), data.size(), failMadvise, countMemset); @@ -139,7 +139,7 @@ INSTANTIATE_TEST_SUITE_P(UnsupportedPopulateWrite, PrefaultFallbackTest, testing TEST(KvCacheManagerV2HostMemTest, ConvertsPrefaultEnomem) { - std::vector data(HostMem::kAlignment); + std::vector data(HostMem::sizeGranularity()); gMadviseErrno = ENOMEM; EXPECT_THROW(hostPrefaultChunk(reinterpret_cast(data.data()), data.size(), failMadvise, countMemset), HostOOMError); @@ -147,7 +147,7 @@ TEST(KvCacheManagerV2HostMemTest, ConvertsPrefaultEnomem) TEST(KvCacheManagerV2HostMemTest, PropagatesOtherPrefaultErrors) { - std::vector data(HostMem::kAlignment); + std::vector data(HostMem::sizeGranularity()); gMadviseErrno = EIO; try { @@ -168,11 +168,11 @@ TEST_P(HostMemPageModeTest, AllocatesAndResizes) { ScopedEnv thp("TLLM_KV_CACHE_MANAGER_V2_THP", std::string(GetParam())); ASSERT_EQ(cudaSetDevice(0), cudaSuccess); - HostMem memory(4 * HostMem::kAlignment, HostMem::kAlignment); + HostMem memory(4 * HostMem::sizeGranularity(), HostMem::sizeGranularity()); EXPECT_NE(memory.address(), 0); - EXPECT_EQ(memory.size(), HostMem::kAlignment); - memory.resize(2 * HostMem::kAlignment); - EXPECT_EQ(memory.size(), 2 * HostMem::kAlignment); + EXPECT_EQ(memory.size(), HostMem::sizeGranularity()); + memory.resize(2 * HostMem::sizeGranularity()); + EXPECT_EQ(memory.size(), 2 * HostMem::sizeGranularity()); } INSTANTIATE_TEST_SUITE_P(ThpModes, HostMemPageModeTest, testing::Values("0", "1")); @@ -190,7 +190,7 @@ TEST(KvCacheManagerV2HostMemTest, AllocationSupportsGpuRoundTrip) ASSERT_EQ(cudaMalloc(&devicePtr, kSize), cudaSuccess); // A copy at a time, because the backing registers the range unit by unit and // the driver rejects one that straddles two registrations. - size_t const unit = HostMem::kAlignment; + size_t const unit = HostMem::sizeGranularity(); for (size_t offset = 0; offset < kSize; offset += unit) { ASSERT_EQ(cudaMemcpy(static_cast(devicePtr) + offset, @@ -218,7 +218,7 @@ TEST(KvCacheManagerV2HostMemTest, DestructionAfterPoisonDoesNotTerminate) { ASSERT_EQ(cudaSetDevice(0), cudaSuccess); { - HostMem memory(HostMem::kAlignment); + HostMem memory(HostMem::sizeGranularity()); Poison::set("DestructionAfterPoisonDoesNotTerminate", "simulated invariant violation"); ASSERT_TRUE(Poison::poisoned()); } @@ -230,10 +230,11 @@ TEST(KvCacheManagerV2HostMemTest, DestructionAfterPoisonDoesNotTerminate) TEST(KvCacheManagerV2HostMemTest, SingleSizeConstructorRoundsUpToAlignment) { ASSERT_EQ(cudaSetDevice(0), cudaSuccess); - for (size_t const requested : {size_t{8}, HostMem::kAlignment + 12}) + for (size_t const requested : {size_t{8}, HostMem::sizeGranularity() + 12}) { HostMem memory(requested); - size_t const expected = (requested + HostMem::kAlignment - 1) / HostMem::kAlignment * HostMem::kAlignment; + size_t const expected + = (requested + HostMem::sizeGranularity() - 1) / HostMem::sizeGranularity() * HostMem::sizeGranularity(); EXPECT_EQ(memory.size(), expected); std::memset(reinterpret_cast(memory.address()), 0xA5, memory.size()); } @@ -272,55 +273,18 @@ TEST(KvCacheManagerV2HostMemTest, ResizeGrowPreservesContentsAndAddress) EXPECT_TRUE(std::all_of(tail, tail + (kGrown - kInitial), [](unsigned char v) { return v == 0xC3; })); } -//! Sizes are quantized to the unit memory is allocated in, which grows with the -//! tier. Nothing rounds to the fill chunk -- the last fill step is just short. -TEST(KvCacheManagerV2HostMemTest, GranularityGrowsWithTierSize) +//! One unit for every host tier, whatever its quota. madvise and mprotect take any +//! page-aligned range, so nothing about a larger tier asks for a larger unit -- unlike +//! the device tier, where the unit is also the mapping granularity. +TEST(KvCacheManagerV2HostMemTest, HostGranularityDoesNotFollowTierSize) { - // Below the first threshold every tier shares the minimum unit. - EXPECT_EQ(tierAllocationGranularity(CacheTier::HOST_MEM, 256 << 20), HostMem::kAlignment); - // Above it the unit grows, and stops growing once the per-allocation cost - // has been amortized away. - EXPECT_GT(tierAllocationGranularity(CacheTier::HOST_MEM, size_t{4} << 30), HostMem::kAlignment); - EXPECT_EQ(tierAllocationGranularity(CacheTier::HOST_MEM, size_t{1} << 40), size_t{32} << 20); + EXPECT_EQ(tierAllocationGranularity(CacheTier::HOST_MEM, 256 << 20), HostMem::sizeGranularity()); + EXPECT_EQ(tierAllocationGranularity(CacheTier::HOST_MEM, size_t{4} << 30), HostMem::sizeGranularity()); + EXPECT_EQ(tierAllocationGranularity(CacheTier::HOST_MEM, size_t{1} << 40), HostMem::sizeGranularity()); + // The device tier does follow it: there the unit is what a mapping costs. + EXPECT_GT(tierAllocationGranularity(CacheTier::GPU_MEM, size_t{1} << 40), HostMem::sizeGranularity()); } -//! A commit unit larger than kAlignment must work end to end. Tiers of a few -//! GiB and up get one, so a size rounded to kAlignment instead would be -//! rejected by HostMem for every realistic host tier. -TEST(KvCacheManagerV2HostMemTest, HonoursCommitUnitLargerThanAlignment) -{ - ASSERT_EQ(cudaSetDevice(0), cudaSuccess); - constexpr size_t kUnit = size_t{32} << 20; - ASSERT_GT(kUnit, HostMem::kAlignment); - - HostMem memory(4 * kUnit, 2 * kUnit, HostMemBackingOptions{kUnit}); - EXPECT_EQ(memory.size(), 2 * kUnit); - auto* bytes = reinterpret_cast(memory.address()); - std::memset(bytes, 0x3C, 2 * kUnit); - EXPECT_TRUE(std::all_of(bytes, bytes + 2 * kUnit, [](unsigned char v) { return v == 0x3C; })); - - // A size that is only kAlignment-aligned is not a valid size for this unit. - EXPECT_THROW(memory.resize(2 * kUnit + HostMem::kAlignment), std::exception); -} - -//! The pool must quantize its size to the same unit, or constructing any tier -//! whose granularity exceeds kAlignment throws. -TEST(KvCacheManagerV2HostMemTest, SlotPoolAlignsToCommitUnit) -{ - ASSERT_EQ(cudaSetDevice(0), cudaSuccess); - size_t const unit = tierAllocationGranularity(CacheTier::HOST_MEM, size_t{8} << 30); - ASSERT_GT(unit, HostMem::kAlignment); - - // A slot size that divides kAlignment but not the larger unit, so a pool - // rounding to the wrong one produces a size HostMem rejects. - constexpr size_t kSlotSize = 1 << 20; - HostSlotPool pool(kSlotSize, SlotCount{3}, /*vmSize=*/unit * 4, HostMemBackingOptions{unit}); - EXPECT_GE(pool.numSlots(), SlotCount{3}); - EXPECT_NE(std::get(pool.slotAddress(SlotId{0})), MemAddress{0}); -} - -//! The reservation is a hard bound. Serving a larger resize would mean moving -//! the base address, which callers are entitled to assume never happens. TEST(KvCacheManagerV2HostMemTest, ResizeBeyondReservationIsRejected) { ASSERT_EQ(cudaSetDevice(0), cudaSuccess); @@ -330,7 +294,7 @@ TEST(KvCacheManagerV2HostMemTest, ResizeBeyondReservationIsRejected) EXPECT_EQ(memory.vmSize(), kMax); EXPECT_NO_THROW(memory.resize(kMax)); - EXPECT_THROW(memory.resize(kMax + HostMem::kAlignment), std::exception); + EXPECT_THROW(memory.resize(kMax + HostMem::sizeGranularity()), std::exception); // The rejected resize left the object usable and unmoved. EXPECT_EQ(memory.address(), base); @@ -445,24 +409,25 @@ TEST(KvCacheManagerV2HostMemTest, ConcurrentWaitUsableFromManyThreads) TEST(KvCacheManagerV2HostMemTest, DeferredParallelFillKeepsEarlyWrites) { ASSERT_EQ(cudaSetDevice(0), cudaSuccess); - constexpr size_t kChunk = HostMem::kFillChunkSize; - constexpr size_t kSize = 4 * kChunk; + // A fill advances the watermark one commit step at a time. + size_t const kChunk = HostMemBacking::commitChunkSize(); + size_t const kSize = 4 * kChunk; HostMem memory(kSize, kSize, HostMemBackingOptions{}, /*waitForFill=*/false); auto* bytes = reinterpret_cast(memory.address()); auto const pattern - = [](size_t offset) { return static_cast((offset / HostMem::kAlignment) * 5 + 3); }; + = [](size_t offset) { return static_cast((offset / HostMem::sizeGranularity()) * 5 + 3); }; for (size_t chunkEnd = kChunk; chunkEnd <= kSize; chunkEnd += kChunk) { memory.waitUsable(chunkEnd); - for (size_t offset = chunkEnd - kChunk; offset < chunkEnd; offset += HostMem::kAlignment) + for (size_t offset = chunkEnd - kChunk; offset < chunkEnd; offset += HostMem::sizeGranularity()) { bytes[offset] = pattern(offset); } } memory.waitUsable(kSize); - for (size_t offset = 0; offset < kSize; offset += HostMem::kAlignment) + for (size_t offset = 0; offset < kSize; offset += HostMem::sizeGranularity()) { ASSERT_EQ(bytes[offset], pattern(offset)) << "byte at " << offset << " lost during the fill"; } @@ -474,23 +439,23 @@ TEST(KvCacheManagerV2HostMemTest, DeferredParallelFillKeepsEarlyWrites) TEST(KvCacheManagerV2HostMemTest, ParallelInitialFillThenGrow) { ASSERT_EQ(cudaSetDevice(0), cudaSuccess); - constexpr size_t kChunk = HostMem::kFillChunkSize; - constexpr size_t kInitial = 2 * kChunk + kChunk / 2; - constexpr size_t kGrown = 4 * kChunk + kChunk / 4; + size_t const kChunk = HostMemBacking::commitChunkSize(); + size_t const kInitial = 2 * kChunk + kChunk / 2; + size_t const kGrown = 4 * kChunk + kChunk / 4; HostMem memory(kGrown, kInitial); auto* bytes = reinterpret_cast(memory.address()); - for (size_t offset = 0; offset < kInitial; offset += HostMem::kAlignment) + for (size_t offset = 0; offset < kInitial; offset += HostMem::sizeGranularity()) { - bytes[offset] = static_cast((offset / HostMem::kAlignment) * 7 + 1); + bytes[offset] = static_cast((offset / HostMem::sizeGranularity()) * 7 + 1); } memory.resize(kGrown); - for (size_t offset = 0; offset < kGrown; offset += HostMem::kAlignment) + for (size_t offset = 0; offset < kGrown; offset += HostMem::sizeGranularity()) { if (offset < kInitial) { - ASSERT_EQ(bytes[offset], static_cast((offset / HostMem::kAlignment) * 7 + 1)) + ASSERT_EQ(bytes[offset], static_cast((offset / HostMem::sizeGranularity()) * 7 + 1)) << "byte at " << offset << " lost by the grow"; } else diff --git a/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2StagingBufferTest.cu b/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2StagingBufferTest.cu index 9e2597909e99..369b71eba752 100644 --- a/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2StagingBufferTest.cu +++ b/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2StagingBufferTest.cu @@ -226,7 +226,7 @@ TEST(KvCacheManagerV2StagingBufferTest, EphemeralHostIndicesUploadIntoDeviceRing auto device = deviceManager.acquire(sizeof(PageIndexPair), std::min(remainingBytes, kChunkBytes), sizeof(PageIndexPair), alignof(PageIndexPair), reinterpret_cast(stream)); size_t const chunkPairs = std::min(input.size() - offset, device.size() / sizeof(PageIndexPair)); - detail::copyPageIndicesToDevice(static_cast(device.address()), input.data() + offset, chunkPairs, + copyToDeviceSyncSource(static_cast(device.address()), input.data() + offset, chunkPairs, reinterpret_cast(stream)); // The helper captures the pageable host source before returning. diff --git a/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2StatsTest.cpp b/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2StatsTest.cpp index c60c2795d413..eed2adf2cf9f 100644 --- a/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2StatsTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2StatsTest.cpp @@ -336,7 +336,7 @@ TEST_F(ReplayCacheTest, OptionalWindowClaimsSurviveIndependentPageRemoval) // Declining the claim allocates writable pages without replacing // the candidates available to a later reader. auto declined = manager->createKvCache({}, prefix, 3); - ASSERT_TRUE(declined->resume(mStream, std::nullopt, std::vector{})); + ASSERT_TRUE(declined->resume(mStream, std::vector{})); for (size_t i = 0; i < original.size(); ++i) { EXPECT_NE(declined->getBasePageIndices(encoder)[c.ordinals[i]], original[i]); @@ -349,13 +349,13 @@ TEST_F(ReplayCacheTest, OptionalWindowClaimsSurviveIndependentPageRemoval) EXPECT_FALSE(missing->reuseStatus()[encoder.value()].complete); EXPECT_TRUE(missing->reuseStatus()[encoder.value()].coverage.empty()); EXPECT_TRUE(missing->reuseStatus()[peer.value()].complete); - EXPECT_THROW(missing->resume(mStream, std::nullopt, selected), std::invalid_argument); + EXPECT_THROW(missing->resume(mStream, selected), std::invalid_argument); EXPECT_FALSE(missing->isActive()); - ASSERT_TRUE(missing->resume(mStream, std::nullopt, std::vector{peer})); + ASSERT_TRUE(missing->resume(mStream, std::vector{peer})); missing->close(); // The earlier claim owns every candidate even after tree unlink. manager->getAndResetIterationReusedBlocksByLevel(); - ASSERT_TRUE(reader->resume(mStream, std::nullopt, selected)); + ASSERT_TRUE(reader->resume(mStream, selected)); reader->commitPendingStats(); auto const reusedByLevel = manager->getAndResetIterationReusedBlocksByLevel(); ASSERT_NE(reusedByLevel.find(encoder), reusedByLevel.end()); @@ -438,7 +438,7 @@ TEST_F(ReplayCacheTest, DecliningOptionalClaimReleasesCapacityBeforeAllocation) auto reservations = storage.newSlotsForPoolGroup(kHotLevel, group, storage.getStatistics(kHotLevel, group).free); // No host tier, no free slots: the rejected candidate is the only capacity // available to the request's fresh writable Encoder page. - EXPECT_TRUE(reader->resume(mStream, std::nullopt, std::vector{})); + EXPECT_TRUE(reader->resume(mStream, std::vector{})); EXPECT_EQ(reader->numCommittedTokens(), 8); reader->close(); for (auto& slot : reservations) diff --git a/cpp/tests/unit_tests/multi_gpu/kvCacheManagerV2HostMemNumaTest.cpp b/cpp/tests/unit_tests/multi_gpu/kvCacheManagerV2HostMemNumaTest.cpp index 2a586edd8050..ff40e66a1c1d 100644 --- a/cpp/tests/unit_tests/multi_gpu/kvCacheManagerV2HostMemNumaTest.cpp +++ b/cpp/tests/unit_tests/multi_gpu/kvCacheManagerV2HostMemNumaTest.cpp @@ -133,8 +133,8 @@ TEST_P(HostMemNumaTest, MmapPagesLandOnTheGpusNode) constexpr size_t kSize = size_t{8} << 20; HostMemBackingOptions options; options.allowRemoteNumaFallback = allowRemoteFallback; - HostMemBacking backing{options}; - MemAddress const base = backing.reserve(kSize); + HostMemBacking backing{kSize, options}; + MemAddress const base = backing.address(); ASSERT_NE(base, 0U); // The flag picks between a strict binding and a preference. Both place the @@ -144,7 +144,7 @@ TEST_P(HostMemNumaTest, MmapPagesLandOnTheGpusNode) ASSERT_EQ(::get_mempolicy(&mode, nullptr, 0, reinterpret_cast(base), MPOL_F_ADDR), 0); EXPECT_EQ(mode, allowRemoteFallback ? MPOL_PREFERRED : MPOL_BIND); - backing.commit(0, kSize); + backing.resizeTo(kSize); stamp(base, 0, kSize, 0); auto const nodes = pageNodes(base, kSize, size_t{2} << 20); diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/AGENTS.md b/tensorrt_llm/runtime/kv_cache_manager_v2/AGENTS.md index 9cfef24d5ec3..af9c3756b0b1 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/AGENTS.md +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/AGENTS.md @@ -26,7 +26,9 @@ This directory contains only: means adding a rebind here and an entry in `__all__`. - `_introspection.py` — white-box hooks for tests and accuracy harnesses. Every function forwards to the bindings' native `_introspection` submodule; the indirection exists so callers - import a stable Python path and get plain Python containers back. + import a stable Python path and get plain Python containers back. The one exception is + `batch_gpu_coalesced_page_table`, which returns a DLPack view: its callers check the + device-array contract and pointer stability, which a host copy would erase. ## Commands diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py index 3eb903af9973..8834aee6181f 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py @@ -78,6 +78,7 @@ class _BatchDescFieldSpec: BatchDesc.__dataclass_fields__ = _BatchDescFieldSpec.__dataclass_fields__ del _BatchDescFieldSpec, _dataclasses_bd BufferConfig = _cpp.BufferConfig +SparseCacheConfig = _cpp.SparseCacheConfig BufferId = _cpp.BufferId CoalescedBuffer = _cpp.CoalescedBuffer CacheTier = _cpp.CacheTier @@ -136,6 +137,7 @@ class _KVCacheManagerConfigFieldSpec: enable_partial_commit: bool = True pool_rebalance: object = None cold_tier_offload_min_priority: int = DEFAULT_COLD_TIER_OFFLOAD_MIN_PRIORITY + host_sparse_pool_quota: object = None KVCacheManagerConfig.__dataclass_fields__ = _KVCacheManagerConfigFieldSpec.__dataclass_fields__ @@ -146,15 +148,11 @@ class _KVCacheManagerConfigFieldSpec: KVCacheStoredData = _cpp.KVCacheStoredData KVCacheUpdatedData = _cpp.KVCacheUpdatedData KvCacheStatus = _cpp.KvCacheStatus -LogicError = _cpp.LogicError OutOfMemoryError = _cpp.OutOfMemoryError OutOfPagesError = _cpp.OutOfPagesError PageIndexConverter = _cpp.PageIndexConverter PageIndexMode = _cpp.PageIndexMode PageStatus = _cpp.PageStatus -PageStorageSnapshot = _cpp.PageStorageSnapshot -Batch = _cpp.Batch -BatchDeviceArray = _cpp.BatchDeviceArray PlannedDropHandle = _cpp.PlannedDropHandle PoolDesc = _cpp.PoolDesc PoolGroupDesc = _cpp.PoolGroupDesc @@ -163,7 +161,8 @@ class _KVCacheManagerConfigFieldSpec: ReuseGroupStatus = _cpp.ReuseGroupStatus ReuseScope = _cpp.ReuseScope ReusedBlocksByLevel = _cpp.ReusedBlocksByLevel -ScratchDesc = _cpp.ScratchDesc +Batch = _cpp.Batch +IndexTableSpec = _cpp.IndexTableSpec SlotDesc = _cpp.SlotDesc SlotDescVariant = _cpp.SlotDescVariant SsmLayerConfig = _cpp.SsmLayerConfig @@ -184,6 +183,7 @@ class _KVCacheManagerConfigFieldSpec: BeamIndex = int CacheLevel = int CacheTierConfig = Union[GpuCacheTierConfig, HostCacheTierConfig, DiskCacheTierConfig] +BatchRow = int CudaStream = int DataRole = str LayerGroupId = int @@ -222,6 +222,7 @@ def typed_range(*args: int) -> range: "BatchDesc", "BeamIndex", "BufferConfig", + "SparseCacheConfig", "BufferId", "CoalescedBuffer", "CacheLevel", @@ -251,9 +252,9 @@ def typed_range(*args: int) -> range: "KVCacheUpdatedData", "KvCacheStatus", "LayerGroupId", + "BatchRow", "LayerId", "LifeCycleId", - "LogicError", "MemAddress", "NDEBUG", "OutOfPagesError", @@ -262,9 +263,6 @@ def typed_range(*args: int) -> range: "PoolGroupPeakBlockStats", "PageIndexMode", "PageStatus", - "PageStorageSnapshot", - "Batch", - "BatchDeviceArray", "PoolDesc", "PoolGroupDesc", "PoolGroupIndex", @@ -272,7 +270,8 @@ def typed_range(*args: int) -> range: "Priority", "ReuseGroupStatus", "ReuseScope", - "ScratchDesc", + "Batch", + "IndexTableSpec", "KVCacheIterationStatsDelta", "ReusedBlocksByLevel", "KVCacheStatsDelta", diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi index e418a261853e..32ea799b14da 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi @@ -30,6 +30,7 @@ from typing import ( Type, TypeAlias, Union, + overload, ) from tensorrt_llm.bindings.executor import KvCacheRetentionConfig @@ -50,9 +51,6 @@ class CuError(Exception): error_code: Any -class LogicError(Exception): - """An operation violates the cache or batch's state and usage requirements.""" - class OutOfMemoryError(Exception): ... class HostOOMError(OutOfMemoryError): @@ -93,13 +91,12 @@ LifeCycleId = NewType("LifeCycleId", int) LayerGroupId: TypeAlias = LifeCycleId class AttnLifeCycle: - """The attention life cycle, keyed by window, sink tokens, sparsity, and reuse policy.""" + """The attention life cycle, keyed by window, sink tokens, and reuse policy.""" def __init__( self, window_size: int | None, num_sink_blocks: int, - is_sparse: bool = False, reuse_policy: AttentionReusePolicy = ..., ) -> None: ... @staticmethod @@ -107,7 +104,6 @@ class AttnLifeCycle: window_size: int | None, num_sink_tokens: int | None, tokens_per_block: int, - is_sparse: bool = False, reuse_policy: AttentionReusePolicy = ..., ) -> "AttnLifeCycle": ... @property @@ -115,8 +111,6 @@ class AttnLifeCycle: @property def num_sink_blocks(self) -> int: ... @property - def is_sparse(self) -> bool: ... - @property def reuse_policy(self) -> AttentionReusePolicy: ... def get_stale_range(self, history_length: int, tokens_per_block: int) -> HalfOpenRange: ... @@ -153,6 +147,7 @@ MemAddress = NewType("MemAddress", int) Priority = NewType("Priority", int) PoolGroupIndex = NewType("PoolGroupIndex", int) PoolIndex = NewType("PoolIndex", int) +BatchRow = NewType("BatchRow", int) # From _storage_manager.py class StorageStatistics: @@ -224,6 +219,21 @@ class PoolGroupPeakBlockStats: unavailable: int evictable: int +class DeviceArray(Protocol): + """A dense array in GPU memory, borrowed for the duration of a call. + + Structural, not nominal: anything exporting the DLPack protocol over a CUDA buffer + satisfies it, so a ``torch.Tensor`` is accepted without this package depending on + torch. Callers may equally pass a CuPy, JAX or numba device array. + + Shape and dtype are stated per parameter rather than in the type, and are checked at + the boundary. Contents are read or written on the caller's stream; the array must + stay alive and unmodified until that work completes. + """ + + def __dlpack__(self, *, stream: int | None = ...) -> Any: ... + def __dlpack_device__(self) -> tuple[int, int]: ... + # From _config.py DataRole = NewType("DataRole", str) @@ -257,14 +267,130 @@ class DiskCacheTierConfig: def tier(self) -> CacheTier: ... def validate(self) -> None: ... -@dataclass(slots=True) +@dataclass(slots=True, frozen=True) class BufferConfig: + """One buffer of a layer's KV cache. + + ``is_sparse`` marks a buffer whose history is too large for device memory and is read + through a per-request device cache rather than directly (see + :meth:`KVCacheManager.is_sparse`). + + It changes the *medium* its pages are made of, not where they lock. Every buffer locks + its pages at the hot level; a sparse buffer's hot pool is host-backed, device-mapped so + a kernel can read it. Its pages are otherwise ordinary hot pages -- coalesced, + addressed by page index, visible to the transceiver -- and the page-index table stays + uniform. A sparse buffer places no requirement on the cold + tiers, which may be host, disk, or absent. + + Backing is decided per pool, and a pool is uniformly sparse or uniformly dense because + buffers group on ``(size, is_sparse)``, so one layer's buffers may land in both. A group + holds whichever its buffers grouped into, and reports GPU memory either way. A pool + group must hold at least one non-sparse pool, since slot counts are sized from device + bytes; a wholly sparse pool group is refused by the ``KVCacheManager`` constructor. + A wholly sparse layer is allowed when its group also contains dense buffers from + another layer. + + The layer's :attr:`AttentionLayerConfig.sparse_cache` configures its fetch path, + including selection width, cache capacity and unit size. Without it, sparse storage + and locking still work, but the buffer has no GPU cache to read through. + + Device residency is a lease on the life cycle's arena owned by a :class:`Batch`. + Each Batch row has one lease, held by its seated request. Both the unit being written + and the selected history occupy cells in it. The selection is below the watermark; + the single unit in flight is above it and cannot be evicted. + + A lease holds contiguous runs of buffer-sized cells, addressed by arithmetic from + the buffer and entry index. :meth:`Batch.fetch_sparse_pages` fills selected units. + ``sparse_unit_chunks``, the layer's ``sparse_cache.unit_tokens`` and the buffer's + size determine where each unit's bytes are in a densely packed page. + + :meth:`Batch.flush_sparse_units` writes completed units back to their host pages + before :meth:`_KVCache.resize` advances history or the request leaves its Batch. + The flush leaves the cells cached; advancing history makes them evictable. + + Fixed at construction: the storage layout derived from it is built once. + """ + role: DataRole size: int tokens_per_block_override: int | None = None is_sparse: bool = False coalescing_group: int = 0 + sparse_unit_chunks: int = 1 + """Pieces a selection unit is gathered from within *this buffer's* page. + + Per buffer, because it is this buffer's own dimension order: + + - **1** for a sequence-major page (``[token][head][dim]``), where a unit's tokens carry + their heads with them and the whole unit is one contiguous range. Also 1 for any + single-head page -- a compressed or pooled latent -- where the orders coincide. + - **numHeads** for a head-major page (``[head][token][dim]``), where each head + contributes its own run and a unit is that many strided pieces. + + Every sparse model in this repository today is one of the first two cases, so 1 is + almost always right. The rest of the layout follows from this, the layer's + ``sparse_cache.unit_tokens`` and the page being densely packed, so a padded page is not + expressible. + """ + +@dataclass(slots=True) +class SparseCacheConfig: + """One layer's sparse fetch path. + + Every sparse buffer of a layer shares one resolve, one resident set and one recency + order. Within each Batch row's lease, metadata is laid out ``[selection][entry]``. + Layers with the same ``selection_id`` share that metadata while keeping separate + data cells; otherwise each layer has its own selection. + + Its presence is what configures the fetch path. A sparse buffer in a layer without it + still gets sparse storage and locking; it simply has no GPU cache to read through. + + Layers sharing a life cycle must agree on width, capacity and unit size, since they + share one arena; their ``selection_id`` values may differ. Batch construction checks + the rounded cache capacity and selection width against the resolve kernel's available + shared memory on the current device. + """ + + num_selected_blocks: int + """Selection width: units a sparse-attention step selects for this layer.""" + + max_cached_units: int + """Cells one buffer gets in a request's GPU cache, counting both populations: the + history it caches and every unit still in flight. + + Two to four times ``num_selected_blocks`` is the usual starting point: a cache only as + wide as the selection retains nothing between steps, and retention of units not + selected this step is where the hit rate comes from. It must exceed + ``num_selected_blocks``, since equality would leave no cell for the unit being written. + """ + + unit_tokens: int + """Original tokens one selection unit spans. + + A whole page is ``tokens_per_block``, which a caller wanting that states rather than + omits -- there is no unit size that suits a layer which did not name one. + + Must divide ``tokens_per_block``; the quotient is how many units a page holds. + Distinct from ``BufferConfig.tokens_per_block_override``, which reinterprets a buffer's + page for *storage*: this names the sparse cache's granularity and nothing else, so + pages, page indices and index spaces all keep the manager's ``tokens_per_block``. + Setting both on one buffer is rejected. + """ + + selection_id: int | None + """Layers whose selection is one selection, named so they can say so. + + A model running one indexer across several layers hands them all the same units every + step. Naming that lets those layers share one resident set and one resolve, where + otherwise each repeats the other's answer against its own copy of it. + + ``None`` means the layer selects for itself, and is read as the layer's own id -- so a + layer that names nothing never shares, and an explicit id that names such a layer is + rejected. Layers that do share must agree on ``num_selected_blocks``, + ``max_cached_units`` and ``unit_tokens``. + """ + class AttentionReusePolicy(enum.IntEnum): REQUIRED = 0 PRIVATE = 1 @@ -272,11 +398,25 @@ class AttentionReusePolicy(enum.IntEnum): @dataclass(slots=True) class AttentionLayerConfig: + """One attention layer's buffers and its history rules. + + A sparse layer may also declare a sliding window, though the combination is unusual: + sparsity already bounds how much history a step reads. Where both apply, the window + bounds what is *selectable* as well as what is retained -- the indexer producing the + selection applies the same window, so it never names a block the window has made stale. + """ + layer_id: LayerId buffers: list[BufferConfig] sliding_window_size: int | None = None num_sink_tokens: int | None = None reuse_policy: AttentionReusePolicy = AttentionReusePolicy.REQUIRED + + sparse_cache: SparseCacheConfig | None = None + """The sparse fetch path for this layer's sparse buffers, if it has one. + + Setting it on a layer with no sparse buffer to serve is rejected. + """ @property def window_size(self) -> int | None: ... @@ -333,6 +473,16 @@ class KVCacheManagerConfig: enable_partial_commit: bool = True pool_rebalance: PoolRebalanceConfig = ... cold_tier_offload_min_priority: int = DEFAULT_COLD_TIER_OFFLOAD_MIN_PRIORITY + # Host bytes the GPU tier's host-backed pools may take, or None for unbounded. A sparse + # buffer's pages are host memory, so the device quota does not pay for them and nothing + # else bounds them: a slot spans every pool of its group, so their row count follows from + # the device-backed pools the quota does size. This only lowers a slot count, and pool + # groups holding such pools share it by the ratio that divides the device quota. + # + # Kept exactly, as is the feasibility floor ``constraints`` implies; the pool ratio is what + # gives way between them. The two can only conflict if the floors alone exceed this + # budget, which construction rejects. + host_sparse_pool_quota: int | None = None @property def enable_swa_scratch_reuse(self) -> bool: ... @@ -496,83 +646,6 @@ KvCacheStatus: TypeAlias = _Status IndexSeq = array.array[int] | memoryview[int] -class BatchDeviceArray: - """DLPack view of int32 CUDA metadata; the view keeps its allocation alive. - - Treat these arrays as read-only. Convert with ``torch.from_dlpack(view)`` or - another DLPack consumer, after publication. Holding a view does not pin KV pages. - """ - - def __dlpack_device__(self) -> tuple[int, int]: ... - def __dlpack__(self, stream: int | None = None, **kwargs: object) -> object: ... - -class Batch: - """Stable request slots and raw GPU metadata across all layer groups. - - Membership is non-owning and exclusive. Closing/destroying a request removes - it; closing/destroying the batch detaches live requests without closing them. - Use the requests' owning thread. ``publish`` runs outside graph capture after - mutations; ``wait_ready`` orders readers after upload and KV-copy completion. - Call ``record_read`` after submitting reads (or graph replay), before mutating, - suspending, removing, or closing requests. Device addresses remain stable. - """ - - def __init__( - self, manager: KVCacheManager, max_rows: int, max_blocks: int, max_beam_width: int = 1 - ) -> None: - """Allocate fixed device tables. Currently supports beam width 1.""" - @property - def max_rows(self) -> int: ... - @property - def max_blocks(self) -> int: ... - @property - def max_beam_width(self) -> int: ... - @property - def num_layer_groups(self) -> int: ... - @property - def dirty_rows(self) -> list[int]: ... - def add(self, kv_cache: _KVCache, row: int | None = None) -> int: ... - def remove(self, kv_cache: _KVCache) -> None: ... - def close(self) -> None: ... - def publish(self, cuda_stream: CudaStream) -> list[int]: - """Queue dirty rows and return their slots. Failures stay dirty; retry before reading.""" - def wait_ready(self, cuda_stream: CudaStream) -> None: ... - def record_read(self, cuda_stream: CudaStream) -> None: ... - def resize( - self, - capacities: list[int | None], - history_lengths: list[int | None], - cuda_stream: CudaStream, - ) -> list[bool | None]: - """Lists use stable row slots. Return per-request success (None for holes), then publish.""" - def page_table(self, layer_group_id: LayerGroupId) -> BatchDeviceArray: - """Raw slot IDs, shape [max_rows, max_beam_width, max_blocks]; preserves BAD_PAGE_INDEX.""" - def num_blocks(self, layer_group_id: LayerGroupId) -> BatchDeviceArray: - """Eligible history counts, shape [max_rows, max_beam_width]; zero for inactive/dense rows.""" - -class PageStorageSnapshot: - """Copied host metadata for one layer group and beam; indices are raw mixed-tier slot IDs. - - ``BAD_PAGE_INDEX`` is preserved and has no cache level. Eligibility is zero for - prefill, inactive requests and dense groups. Readiness events are retained internally; - they do not pin storage. Use indices only while the request is active and the version - matches. Submit reads on the request's stream, or call ``record_page_storage_read`` - after submission on another stream, before mutating or closing the request. - """ - - @property - def version(self) -> int: ... - @property - def row(self) -> int | None: ... - @property - def base_page_indices(self) -> list[int]: ... - @property - def cache_levels(self) -> list[CacheLevel | None]: ... - @property - def eligible_history_blocks(self) -> int: ... - def wait_ready(self, cuda_stream: CudaStream) -> None: - """Queue copy-completion waits without blocking the CPU or uploading metadata.""" - class ReuseGroupStatus: """Group availability at the initial prefix claim, not subsequent request progress.""" @@ -610,7 +683,17 @@ class _KVCache: ) -> None: ... def set_base_page_index_buf( self, beam_idx: BeamIndex, layer_group_id: LayerGroupId, buf: memoryview | None - ) -> None: ... + ) -> None: + """Write page indices through ``buf`` instead of this cache's own storage. + + Raises while this cache belongs to a :class:`Batch`: membership already binds every + layer group to a row of the batch's table, and a second buffer would take the writes + with it, leaving that row holding whatever it held when the cache joined. Remove it + from the batch first, or give the batch the buffer worth writing to. + + This holds for the batch too -- it ends the membership before handing the indices + back, so nothing binds a seated member's buffer. + """ @property def manager(self) -> "KVCacheManager": ... @property @@ -619,6 +702,14 @@ class _KVCache: def cuda_stream(self, cuda_stream: CudaStream) -> None: ... @property def num_blocks(self) -> int: ... + @property + def batch_row(self) -> BatchRow | None: + """Row this cache occupies in its :class:`Batch`, or ``None`` when it belongs to none. + + Also the answer to who owns its page index buffers: a member's are the batch's + rows, so :meth:`set_base_page_index_buf` refuses and ``suspend()`` leaves them + bound rather than copying them back out. + """ def commit_pending_stats(self) -> KVCacheStatsDelta: ... def discard_pending_stats(self) -> None: ... def record_connector_reuse(self, token_begin: int, token_end: int) -> None: ... @@ -645,6 +736,27 @@ class _KVCache: valid_only: bool = False, ) -> Iterator[int]: ... def resize(self, capacity: int | None, history_length: int | None = None) -> bool: ... + @property + def max_input_length(self) -> int | None: + """Bound on ``capacity - history_length``: how far capacity may outrun the watermark. + + ``None`` unless a sparse buffer's fetch path is configured, since a bound on + in-flight blocks describes nothing without a cache to hold them. :meth:`resize` + rejects a gap wider than this. + + **One token**, and read-only, because it is a consequence of the supported shape + rather than of how wide a cache is. A sparse request contributes exactly one query + token per step: :meth:`Batch.fetch_sparse_pages` resolves one selection + per request, while an attention kernel indexes its selection per query token, and + the two agree only at one token each. Prefill and speculative decoding give a + request several query tokens per step and are out of scope; this bound is the guard + that keeps them out. + + It follows that at most one unit is ever in flight, so essentially all of + :attr:`SparseCacheConfig.max_cached_units` is available to the selection rather than + being reserved as headroom. + """ + @property def capacity(self) -> int: ... @capacity.setter @@ -672,37 +784,39 @@ class _KVCache: def reuse_scope(self) -> ReuseScope: ... def plan_committed_block_drop(self) -> PlannedDropHandle | None: ... def stop_committing(self) -> None: ... - def suspend(self) -> None: ... + def suspend(self) -> None: + """Unlock this cache's pages and detach from its stream. + + Raises while it belongs to a :class:`Batch`: the unlocks write ``BAD_PAGE_INDEX`` + through the bound buffer, which is the batch's row, so doing it while seated would + write a row a publish may still be reading. Remove it from the batch first. + """ + def resume( self, cuda_stream: CudaStream | None = None, - is_decoding: bool | None = None, optional_reuse_groups: Sequence[LayerGroupId] | None = None, - ) -> bool: ... - def enter_decode(self) -> bool: ... - @property - def is_decoding(self) -> bool: ... - @property - def page_storage_version(self) -> int: ... - @property - def page_storage_dirty(self) -> bool: ... - @property - def page_storage_row(self) -> int | None: ... - def bind_page_storage_row(self, row: int | None) -> None: - """Bind a standalone consumer's row; Batch members must use Batch.add/remove.""" - def acknowledge_page_storage(self, version: int) -> bool: - """Clear dirty state after all groups/beams use this same version, if it is still current.""" - def get_page_priorities( - self, layer_group_id: LayerGroupId, beam_index: BeamIndex = ... - ) -> list[Priority]: ... - def get_page_storage_snapshot( - self, layer_group_id: LayerGroupId, beam_id: BeamIndex = DEFAULT_BEAM_INDEX - ) -> PageStorageSnapshot: - """Read the final state under the manager lock, including after a failed operation's rollback.""" - def record_page_storage_read(self, cuda_stream: CudaStream) -> None: - """Join submitted reader work into the active request's stream before any cache mutation.""" + ) -> bool: + """Lock this cache's pages at their required levels. + + Returns ``False`` when utilisation is too high or pages cannot be had. Raises while + it belongs to a :class:`Batch`, for the same reason as :meth:`suspend` -- locking + writes page indices through the batch's row. A member is resumed before it joins. + + Takes no sparse cache arena lease: that follows membership, so a request receiving + its history costs the GPU cache nothing until it is seated. + """ def prefetch(self, target: CacheLevel) -> bool: ... - def get_scratch_desc(self, layer_group_id: LayerGroupId) -> ScratchDesc | None: ... + def get_swa_scratch_range(self, layer_group_id: LayerGroupId) -> HalfOpenRange: + """Block ordinal range this life cycle serves from scratch slots. + + Empty when none of its blocks do. Pulled per step rather than mirrored: it is two + integers, and keeping it pulled is what stops a range and its slots from being + updated at different moments. + """ + + def get_scratch_slots(self, layer_group_id: LayerGroupId) -> list[int]: + """Page indices of this life cycle's scratch slots, in slot order.""" @property def has_scratch_slots(self) -> bool: ... @property @@ -713,7 +827,12 @@ class _KVCache: def text_only(self) -> bool: ... @text_only.setter def text_only(self, text_only: bool) -> None: ... + # Per life cycle, which is where the scratch slots live. The one-argument form + # collapses them, so one life cycle's scratch makes it answer False for every other. + @overload def supports_index_mode(self, mode: PageIndexMode) -> bool: ... + @overload + def supports_index_mode(self, mode: PageIndexMode, layer_group_id: LayerGroupId) -> bool: ... @property def status(self) -> _Status: ... @property @@ -727,6 +846,16 @@ class PoolDesc: base_address: MemAddress slot_bytes: int + storage_medium: CacheTier + """What this pool's pages actually are. + + The same as the owning group's tier, except for a sparse pool of the GPU tier: those + are host-backed, and the group still reports ``GPU_MEM`` because its other pools are + not. Ask per pool rather than per group wherever the medium of an address decides + something -- registering it for a disaggregated transfer, or choosing a copy + implementation -- since a group is not uniform in this. + """ + class BufferId(NamedTuple): layer_id: LayerId role: DataRole @@ -790,23 +919,19 @@ class HalfOpenRange: def __len__(self) -> int: ... def __eq__(self, other: object) -> bool: ... -@dataclass(slots=True, frozen=True) -class ScratchDesc: - range: HalfOpenRange - slot_ids: Sequence[int] - def __bool__(self) -> bool: ... - @dataclass(slots=True, frozen=True) class PageIndexConverter: scale: int expansion: int layer_offset: int + scratch_pages_per_block: int def __call__( self, base_indices: Sequence[int], index_mode: PageIndexMode | None = None, - scratch: ScratchDesc | None = None, + scratch_range: HalfOpenRange | None = None, + scratch_slots: Sequence[int] = (), ) -> list[int]: ... class IKvCacheColdPageCodec: ... @@ -820,6 +945,700 @@ def create_default_kv_cache_cold_page_codec() -> IKvCacheColdPageCodec: consumes an explicitly supplied codec, including an attempt that fails. """ +class IndexTableSpec: + """Buffers a consumer must address from one table, under one base pointer. + + States a kernel's requirement, not a preference: some attention ops take a single block + table covering K and V rather than one each. Order defines the table's buffer dimension, + so a kernel reads buffer ``i`` of this spec at entry ``i``. + + Every buffer must share a life cycle and a page-index scale, since no one table can + serve buffers whose converted indices differ. A spec spanning either is refused when + the ``Batch`` is constructed. + """ + + def __init__(self, buffers: Sequence[BufferId]) -> None: ... + @property + def buffers(self) -> list[BufferId]: ... + +class Batch: + """The ``_KVCache`` objects stepped together, and their device-side mirror. + + Owns the batch row order and every per-row device array the sparse path reads: a page + table per layer group, the eligible-block counts, and the staging buffers behind + them. Every such array is indexed by row, and each must agree on which row a request + occupies, so the order lives here rather than being re-established per call. + + Rows are slots rather than positions in a compacted list: removing a member leaves a + hole for reuse instead of shifting later rows, keeping array shapes and addresses + stable across steps, which is what CUDA graph capture requires. + + Membership is exclusive -- a ``_KVCache`` belongs to at most one ``Batch``. + + :meth:`publish` restages every row. Mutating a member directly therefore costs a + republish rather than producing a stale table; refreshing only the rows whose locks + changed is a planned refinement, and a cost rather than a correctness question. + """ + + def __init__( + self, + manager: KVCacheManager, + stream: CudaStream, + index_tables: Sequence[IndexTableSpec], + max_batch_size: int | None = None, + max_blocks_per_seq: int | None = None, + mode: PageIndexMode | None = None, + max_total_input_tokens: int | None = None, + ) -> None: + """Bound the device arrays, whose shape must not change between steps. + + ``index_tables`` names what a consumer needs addressable together, and is + required: ``[]`` is how a caller that converts no indices says so, which is a + choice about what the consumer reads rather than an absence the batch could fill + in. + + ``max_total_input_tokens`` is the tokens one step may ingest across every row. It + sizes the write-back's work list: a row advancing ``t`` tokens crosses at most + ``t / unit_tokens + 1`` units, so the batch crosses at most + ``max_total_input_tokens / unit_tokens + max_batch_size`` of them. Left unset it + assumes a decode step, where each row advances one token and so crosses at most + one unit boundary. + + ``stream`` is the stream every member runs on. A batch owns it rather than taking + one per call: a member's own work and the batch's transfers have to be ordered + against each other, and two streams would leave that to the caller. Joining binds + the member to it, and a member cannot be moved off it. + + Neither can be grown later, so neither is inferred: left unset each derives from + the manager's ``constraints``, which fails when none were declared. A pool-sizing + hint is a default, never a substitute for a batch bound. + + The batch stages the host page indices itself; :meth:`coalesced_page_table` exposes + them for a consumer that needs the same numbers. One buffer with one writer is what + keeps a member's binding alive for the whole decode, and which side allocates it + never mattered to that. + + ``mode`` applies to every life cycle. Left unset, each takes the smaller table: + ``SHARED`` where :meth:`KVCacheManager.supports_index_mode` reports it always + available, ``PER_LAYER`` otherwise. + + Set, it **overrides** that answer rather than being checked against it, so + ``SHARED`` may be named where the manager returns ``None``. The manager answers for + every request a life cycle could carry, so a sliding window leaves it undecided; + naming ``SHARED`` is the narrower claim that no member of *this* batch will hold + scratch slots there. A decode-only batch is the case it exists for: scratch covers + the blocks a prefill chunk adds, and a one-token request adds none. That is what + lets prefill and decode run as two batches over one manager with SWA scratch + enabled, the decode one still taking ``SHARED`` for the sliding life cycles -- where + the saving is worth having, since that batch holds the most rows. + + The claim is validated rather than trusted. Every :meth:`publish` checks each member + against its own :meth:`_KVCache.supports_index_mode` and raises if one holds scratch + after all -- mid-step, with the batch already built around the smaller table. Seat + only members that keep the claim. + + It is fixed here because a converted table's shape is baked into a captured graph, + so it cannot follow a member's scratch state, which varies per step. + """ + + def add_kv_cache(self, kv_cache: _KVCache, row: BatchRow = BatchRow(-1)) -> BatchRow: + """Seat a cache in ``row``, or in the lowest free row when it is negative. + + Returns the row. Raises if the row is taken or out of range, if the batch is full, + if the cache already belongs to a batch, if it has not been resumed, or if its beam + width exceeds one -- only beam 0 is mirrored, so a wider request would have every + beam read beam 0's pages. + + Takes the sparse cache arena lease, because membership is what being stepped means: + a resumed request that is not seated cannot fetch, and a disaggregated one receives + its whole history into host pages without touching the cache. The arena has one + lease per Batch row, so a free row always has a free lease and this cannot fail + for want of one. + + So seat a request only once it is ready to be stepped: one still receiving its + disaggregated history, or sitting a step out, costs the arena nothing while it stays + out, which is why the arena is sized for the rows rather than for every request the + scheduler admits. Nothing checks it -- an early seat shows up as arena pressure. + """ + + def remove_kv_cache(self, kv_cache: _KVCache) -> _KVCache: + """Free the cache's row, leaving a hole for reuse, and return the cache. + + The batch held a reference while the cache was seated; this hands it back, so a + caller that kept only the batch still gets the cache. Raises when it is not a + member. + + Returns the arena lease, which **discards whatever the GPU cache held**: a lease + comes back with every entry free, so a unit whose only copy was in a cell can no + longer be addressed. Call :meth:`flush_sparse_units` first for anything worth + keeping. A batch is meant to persist across steps -- churning membership per step + cold-starts every request's cache. + """ + + def remove_row(self, row: BatchRow) -> _KVCache | None: + """The same, by the row :meth:`add_kv_cache` returned. + + ``None`` when nothing is seated there, so a caller sweeping rows need not ask twice + about each one. + """ + + def move(self, dst_row: BatchRow, src_row: BatchRow) -> None: + """Move the member at ``src_row`` to the empty ``dst_row``, leaving ``src_row`` free. + + Membership does not lapse: the sparse cache lease, its contents and the stream + binding all survive, which is what separates this from :meth:`remove_row` followed + by :meth:`add_kv_cache` -- that pair hands the lease back with every entry free. It + also writes each life cycle's index-table row once instead of twice. + + Raises when either row is out of range, ``src_row`` holds nothing, or ``dst_row`` is + occupied. Waits for the previous staging upload. Device tables keep their addresses + and reach the new row on the next ``publish``, so a step may resize, move and publish + once rather than publishing per operation. + """ + + @property + def num_index_tables(self) -> int: ... + def index_table(self, table_id: int) -> list[IndexTableSpec]: + """Every request this table satisfies. + + A request's buffer order is its own mapping into the buffer dimension, so the first + buffer of each sits at entry zero. Requests whose converted indices coincide are + merged onto one table, which is how a quantised layer's data and block scales come + to share one. + """ + + def index_table_id(self, spec: IndexTableSpec) -> int: + """Which table serves this spec, which must be one the ``Batch`` was built with. + + The query to prefer: a caller that asked for a table by naming its buffers gets it + back without going through a buffer, and equivalent specs resolve to the same id -- + a layer's data and its block scales name one table between them. The spec's buffers + occupy entries ``[0, len(spec.buffers))`` in its own order, so an entry needs no + separate lookup. + + Specs may not overlap, so a buffer belongs to exactly one. Read a subset of a table + by slicing it rather than by declaring the subset as a spec of its own. + """ + + def pool_spec_for(self, spec: IndexTableSpec) -> tuple[MemAddress, int]: + """``(base, page_stride)`` for this spec's converted indices. + + A page index ``p`` of this spec addresses ``base + p * page_stride``. + + Keyed by the spec because the base is. Specs whose converted entries coincide are + merged onto one table, and reading that one set of entries through different bases + is what keeps a quantised layer's data and its block scales apart -- so a table has + as many bases as it serves specs, and only the spec picks one. The spec's buffers + share a page stride, which is what lets one base address them all. + """ + + def converted_index_table(self, table_id: int) -> DeviceArray: + """``int32[max_batch, entries, max_blocks_per_seq]`` converted indices for one table. + + A row carries every entry before the next row starts, which is the shape an + attention kernel takes a key and value page list in. One entry's plane is therefore + strided rather than contiguous: ``table[:, entry, :]``. + + A DLPack array over memory the ``Batch`` owns; it keeps the ``Batch`` alive for as + long as the array exists. A buffer reads the entry its position in the spec gives + it: a spec's buffers occupy entries ``[0, len(spec.buffers))`` in the spec's own + order, so a caller holding the spec needs no further lookup. + + The address is stable for the ``Batch``'s lifetime, so a captured graph may bake it. + """ + + def index_mode(self, layer_group_id: LayerGroupId) -> PageIndexMode: + """Page index mode of one life cycle, as chosen at construction. + + Left unset at construction, each life cycle takes the smaller table: ``SHARED`` + where the manager reports it always available, ``PER_LAYER`` otherwise. Fixed here + because a converted table's shape is baked into a captured graph, so it cannot + follow a member's scratch state. + """ + + def row_of(self, kv_cache: _KVCache) -> BatchRow | None: + """Row of a member, or ``None`` when it is not seated.""" + + def __getitem__(self, row: BatchRow) -> _KVCache | None: + """The member seated in ``row``, or ``None`` when that row is free. + + The inverse of :meth:`row_of`. Raises ``IndexError`` when ``row`` is beyond + ``max_batch_size``, so ``None`` always means an empty row rather than a bad index. + """ + + @property + def num_kv_caches(self) -> int: + """Members seated, which is the rows occupied. + + Also the leases held: seating takes one from every life cycle's arena and leaving + returns it, so the free leases are the arena's capacity less this. + """ + + def peek_hit_miss_counter(self) -> tuple[int, int]: + """Cumulative ``(hits, misses)`` in selection units, since build or the last reset. + + A hit is a selected unit the GPU cache already held; a miss is one the fetch had to + copy from its host page. Padded selection entries and in-flight units are neither, + so they appear in neither figure -- which is why the pair comes from the resolve + rather than from a count the caller could keep itself. + + Reads on a temporary stream and synchronises only that, so the pair may be a few + launches stale. That does not move a ratio; synchronise the device first if an exact + pair matters. + """ + + def reset_hit_miss_counter(self) -> None: + """Zero the cumulative pair, enqueued on the Batch's stream.""" + + @property + def sparse_arena_bytes(self) -> int: + """Device bytes every life cycle's sparse arena holds, ``0`` without sparse offload. + + The arenas sit outside the cache tiers' quota, so a caller budgeting device memory + accounts for them separately and otherwise has to predict the layout to do so. This + reports what was actually allocated, which makes that prediction checkable rather + than something only a profiler disagrees with. + """ + + @property + def is_contiguous(self) -> bool: + """Whether the occupied rows are exactly ``[0, num_kv_caches)``. + + A batch does not require it: removing a member leaves its row free and the next + join takes the lowest free one, so a hole is legal and transient. An engine that + maps rows onto its own request slots usually does require it, and asking here beats + keeping a mirror of the membership to check against. + """ + + @property + def cuda_stream(self) -> CudaStream: + """The stream this batch and every member run on. + + Assigning moves the batch and every member onto the new stream, which first waits + on what the old one has queued, so no explicit synchronization is needed. A member + cannot be moved on its own: its own setter refuses while it is in a batch. + """ + + @cuda_stream.setter + def cuda_stream(self, cuda_stream: CudaStream) -> None: ... + @property + def max_batch_size(self) -> int: ... + def get_sparse_cache_base_address(self, layer_id: LayerId, data_role: DataRole) -> MemAddress: + """Base of the sparse cache arena backing this buffer. Raises if it has no fetch path. + + The arena is this Batch's: it is sized by :attr:`max_batch_size`, one lease per + row, so the address is meaningful only for caches seated here. + + The indices from :meth:`fetch_sparse_pages` and :meth:`prepare_sparse_write_indices` + are relative to this base. The arena holds the GPU-resident units that model kernels + read and write. + + Distinct from :meth:`KVCacheManager.get_mem_pool_base_address`, which names this + buffer's *pages* for every buffer alike -- what the transceiver and the KV transfer + path address. + That meaning does not change when a buffer is sparse; the arena is a second + allocation, so it gets a second accessor rather than overloading the first. + + Role-dependent: the arena mirrors the life cycle's hot pool group, so buffers of + different sizes -- data and its block scales, say -- sit in regions of their own, and + an index is relative to the region of the buffer it was fetched for. Buffers sharing + a pool share a region, and the index already carries which of them it belongs to. + """ + @property + def max_blocks_per_seq(self) -> int: ... + def resize( + self, + capacities: Sequence[int | None], + history_lengths: Sequence[int | None], + ) -> list[int]: + """Resize every member, then publish once. + + One result per row, in row order, ``0`` for a free row: a member rolls back on its + own, so a batch can partially fail. A table published after a partial failure + describes the post-rollback state. + + Rejects, before touching any member, a history length that falls or that outruns + the capacity its row already had -- capacity granted by this call makes room for + tokens written *after* it, so history can only cover tokens whose space was already + there. A request arriving with a whole prompt as history is the one shape that + excludes, so seed it before it joins. + + A row whose growth fails still takes its history advance, which needs no slots and + is what lets the units below it be written back; its result stays ``0``, since the + row is not at the state that was asked for. + """ + + def flush_sparse_units(self, new_history_lengths: Sequence[int]) -> None: + """Write back the units each row's watermark is about to pass. + + One entry per row, giving the history you are about to set; a row may repeat its + current history, which asks for nothing. **Call this before advancing history, not + after**: a block must be in its host page before it can be selected, and + :meth:`_KVCache.resize` refuses a history this has not covered. + + Bounded by the capacity each row started from, not by how far history travels, so + a resumed request moving history across a whole prompt writes back nothing -- its + units reached their host pages directly rather than through a cell. A move wider + than the in-flight headroom is refused. + + Reads the index table as of the last :meth:`publish`, which covers the blocks in + range because they were allocated by an earlier step. + """ + + def prepare_sparse_write_indices( + self, + buffers: Sequence[BufferId], + out: DeviceArray, # int32[batch, len(buffers)] + stream: CudaStream | None = None, + ) -> None: + """Give each request's in-flight unit a cell, and report where to write it. + + The unit between the two watermarks is the one the model is about to produce. It + has no host copy yet -- a unit reaches its page only when :meth:`_KVCache.resize` + carries the watermark past it -- so nothing is fetched here. The call assigns a + cell and leaves its bytes to the model; reading one before the model writes it + yields whatever the cell last held. + + Separate from :meth:`fetch_sparse_pages` because the orders differ. A fetch needs + the selection, which the indexer produces from the previous layer's output; a write + needs only the watermarks. So this can run first, which is the order the model + requires, and it needs no top-k. + + Idempotent while the watermarks and the reservation hold. The resolve probes + residency before taking a victim, so a later call -- including the one inside + :meth:`fetch_sparse_pages` -- finds the entry and reuses it rather than assigning a + second. The entry is pinned for as long as its unit is in flight, so no fetch can + evict it, and a fetch never copies over it: a unit at or above the watermark has no + written-back host page, so there is nothing it could be refilled from. + + The Batch's published rows own the arena leases; each seated cache uses its row's + lease until it leaves the Batch. + + Args: + buffers: exactly one layer's sparse buffers, in the cache's order -- the same + rule and the same order as :meth:`fetch_sparse_pages`. + out: int32 ``[batch.max_batch_size, len(buffers)]``. ``out[b][s]`` is the index + of the in-flight unit's cell, and ``BAD_PAGE_INDEX`` where the row holds no + reservation or has nothing in flight. ``s`` is the buffer's output slice + from :meth:`sparse_out_slots`, the same assignment a fetch uses: buffers + that share a slice are written once, and unused slices are left untouched. + One entry per request, not per query token: :meth:`_KVCache.resize` + bounds ``capacity - history_length`` to a single token, so at most one unit + is ever in flight. + + ``BAD_PAGE_INDEX`` is an ordinary outcome, not an error. A unit is in flight + only once it is whole: with ``SparseCacheConfig.unit_tokens`` of 4, that is one + step in four, which is also the only step on which a ratio-4 compressor has + produced anything to write. With one token per unit it is every step. + stream: the CUDA stream to enqueue on. + + Indices are in the same expanded index space :meth:`fetch_sparse_pages` produces + and relative to the same base -- :meth:`get_sparse_cache_base_address` for that + buffer, **not** :meth:`get_mem_pool_base_address`. The latter names the buffer's + host-backed *pages*; the cells live in a separate GPU arena. + + Call it after publication and before the model's KV writes, ordering the two + through the stream. Regenerate the output whenever a watermark or a reservation + changes, and keep it unmodified until its consumers are done. Complete the writes + before :meth:`flush_sparse_units` writes a finished unit back to its host page, + then advance history through :meth:`_KVCache.resize`. + + Device-driven and stream-ordered, with static shapes, so a CUDA graph can capture + it. + + + Order it after batch publication and before consumers, on the same stream or + through events across streams. The call reads no device values on the host, and + with fixed batch and output shapes it is CUDA-graph-capturable. + + The entry is shared by all of the layer's sparse buffers, but each buffer's index + is relative to its own :meth:`get_sparse_cache_base_address`, including when data + and scales sit in different pools. The indices carry request and buffer placement; + they are neither request-local cache-entry numbers nor byte addresses. + """ + def fetch_sparse_pages( + self, + buffers: Sequence[BufferId], + selected: DeviceArray, # int32[batch, topk] + out: DeviceArray, # int32[batch, len(buffers), topk] + stream: CudaStream | None = None, + ) -> None: + """Resolve a batch's selections against each buffer's per-request history cache. + + Device-driven and stream-ordered throughout: the selection is read from device + memory, so the host neither inspects it nor decides anything, and the call + enqueues onto ``stream`` and returns. Order it against the selection kernel by + sharing a stream, or across streams with an event. Every shape is static across + steps, so the call is capturable in a CUDA graph. + + ``stream`` defaults to the batch's own, which is the only correct choice while + layers are fetched in order, since a layer's selection comes from the one before + it. Name a stream only to run ahead of that -- a model that fetches several layers + early can do so on its own -- and then order it against the batch's stream + yourself: the resolve's scratch and counters are batch state the step clears on + that stream. + + ``buffers`` must be exactly one layer's sparse buffers, in the order + :meth:`get_mem_pool_base_address` numbers them. One layer, not one layer group: + a layer's selection is normally computed from the previous layer's output, so a + step's selections are not all known at its start, and giving each layer its own + call lets the first layer's copies start as early as possible. + + The whole layer, not a subset, because the resolve marks entries resident on + behalf of every buffer at once: filling only some would leave the rest stale while + the metadata claimed otherwise. + + **Only for a layer whose selection is its own.** A model running one indexer + across several layers hands them all the same units, and those layers name one + ``selection_id`` and share a resident set. Their selection is known for all of + them at once, so they resolve it once -- with + :meth:`resolve_sparse_selection` -- and then fetch a layer at a time with + :meth:`fetch_resolved_pages`. Calling this per layer instead would have the + second resolve read the first's admissions as hits, issue no copies, and leave + the second layer's cells unfilled; it raises rather than doing that. + + ``selected`` holds selection-unit ordinals, ``topk`` uniform across requests and ``-1`` + past a request's valid count. The page table and the per-request count of + eligible blocks come from ``batch``, whose row order ``selected`` and ``out`` + both inherit. + + Eligible means below the history watermark: a block at or above it holds input + tokens whose host page has not been written back yet, so selecting it would read + stale bytes. Not a memory-safety problem -- the page table is uniform, so the + ordinal resolves to a real page of the same request -- but a wrong answer. The call + cannot detect an ineligible ordinal, so naming only eligible ones is the selector's + responsibility. + + Each selected unit is looked up in the owning request's cache for that layer. + A hit is reused in place; a miss is copied from the locked host page into a slot + taken from the free list, evicting least-recently-used when it is empty. + The Batch row's lease guarantees room for a full selection and the in-flight + unit. Eviction is local to this selection's resident set; it cannot evict another + selection's entries. Layers sharing a ``selection_id`` use the split resolve/fetch + calls described above and keep their cells in matching slots. + + Copies leave the host page authoritative rather than moving it, so evicting a + fetched cell moves no data. The copies are a snapshot, which is sound only because a + block below the watermark is already written back and immutable. A cell holding an + in-flight block is the one exception, and is never an eviction candidate. + + The in-flight unit is reserved internally but is not appended to ``out``. Use + :meth:`prepare_sparse_write_indices` before the model's KV writes to obtain its + per-buffer physical write indices independently of the selection. + + ``out`` is **compact**: entry ``[b][s][k]`` is the cache index holding + ``selected[b][k]``, and ``BAD_PAGE_INDEX`` where that slot carries no selection. + Padding stays where the selection put it rather than being packed out, so slot + ``k`` of ``out`` and slot ``k`` of ``selected`` describe the same unit. + + ``s`` is the buffer's *output slice*, from :meth:`sparse_out_slots` -- not its + position in ``buffers``. Buffers whose arena indices coincide resolve to one + slice, so a slice may serve several buffers and the slices no buffer maps to are + left untouched. Read a buffer's row at its slice; ``out``'s second extent stays + ``len(buffers)`` so the shape is independent of how the buffers happen to merge. + + Compact, in the selection's own order, rather than a table indexed by unit + ordinal, because that is the shape sparse attention kernels take: a token-sparse + kernel is handed the index list in place of its page table and reads one unit per + slot, and a block-sparse kernel takes the list beside a page table it still uses. + None indexes a table by ordinal, so producing one would only oblige the caller to + compact it again. + + The buffer dimension is a dimension rather than a separate array per buffer + because attention kernels commonly want one tensor covering K and V. It cannot be + dropped: a layer's buffers share the resolve, but not their cells. A lease gives + each buffer its own contiguous run of cells, so one selection resolves to a + different index per buffer, and each needs its own row. + + Indices are in the expanded index space :class:`PageIndexConverter` produces, and + are relative to :meth:`get_sparse_cache_base_address` for the buffer they were + fetched for -- **not** to :meth:`get_mem_pool_base_address`, and not to a cache + level. The pool names the buffer's pages, which for a sparse life cycle are host + memory; the cells a kernel reads live in a separate GPU arena, with a region per + pool. Pairing these indices with the pool base reads the wrong allocation. + + The selection is per layer, not per KV head: every head of a layer shares one + resident set. Algorithms that let different KV heads select different units are + therefore not expressible against this cache. + + **Decode only.** ``selected`` carries one row per batch row, so a request + contributes a single selection per call. Sparse attention kernels index their + selection by *query token*, which coincides with the request only while a request + contributes one token -- that is, during generation. + + The restriction is not just the row count. The resolve makes the selected units + resident in the request's cache, which holds ``max_cached_units``; residency is a + property of the request, not of a query token. A prefill chunk of N query tokens + each choosing their own units has a union of up to N times the selection width, so + the reservation could no longer cover a call, and a unit evicted for a later query + token would invalidate an index already returned for an earlier one. Widening this + to one row per token would make that reachable rather than fix it. + + Serving prefill from a host-resident history is an open design question; see + SPARSE_ATTENTION_DESIGN.md. + """ + + def resolve_sparse_selection( + self, + buffers: Sequence[BufferId], + selected: DeviceArray, # int32[batch, topk] + stream: CudaStream | None = None, + ) -> None: + """Resolve one selection, admitting what it names and listing what to copy. + + The first half of :meth:`fetch_sparse_pages`, split out for the layers that share + a ``selection_id``: they resolve once between them and then fetch a layer at a + time with :meth:`fetch_resolved_pages`. + + ``buffers`` names any one member layer in full. The residency and the miss list + belong to the selection rather than to the layer that happened to name it, so + which member is passed makes no difference -- but it must still be a whole + layer's sparse buffers, as :meth:`fetch_sparse_pages` requires. + + The miss list this leaves behind is read by every :meth:`fetch_resolved_pages` + that follows and is overwritten by the next resolve of the same selection, so + fetch each member layer before resolving the selection again. Both run on + ``stream``, so that ordering is the caller's only across streams. + """ + + def fetch_resolved_pages( + self, + buffers: Sequence[BufferId], + selected: DeviceArray, # int32[batch, topk] + out: DeviceArray, # int32[batch, len(buffers), topk] + stream: CudaStream | None = None, + ) -> None: + """Copy one layer's share of a resolved selection, and write its indices. + + The second half of :meth:`fetch_sparse_pages`. ``selected`` must be the array the + matching :meth:`resolve_sparse_selection` read, since the indices written are + those units'; ``out`` is exactly as :meth:`fetch_sparse_pages` describes it, for + this layer's buffers alone. + + Every layer of the selection copies the same units, each into its own cells: the + entry index is shared, the bytes behind it are not. + """ + + def sparse_out_slots(self, buffers: Sequence[BufferId]) -> list[int]: + """Which slice of ``out`` each buffer's indices land in. + + ``buffers`` is one layer's sparse buffers, as :meth:`fetch_sparse_pages` takes + them, and the result is parallel to it: entry ``i`` is the second-dimension index + to read ``buffers[i]`` at in the ``out`` of :meth:`fetch_sparse_pages`, and in the + ``out`` of :meth:`prepare_sparse_write_indices`. Both use the same assignment. + + Buffers that resolve a selection to the same arena index share a slice, so the + copy runs once and the kernel reads one row for both. Mirrored pools do this -- + equal width gives equal indices, as for a K and a V of the same shape, or a key + and a value scale -- and the slots then repeat, e.g. ``[0, 1, 0, 1]``. Buffers of + unequal width keep slices of their own, so a scale shared by K and V together + yields ``[0, 1, 2]``. + + Slices are numbered by first use, so the distinct slices are ``0`` upwards and + ``max(slots) + 1`` of the ``len(buffers)`` rows carry data. The rest are left + untouched rather than packed out, keeping ``out``'s shape independent of the + merge. + + The grouping follows the index table spec the batch was built with, so two + buffers merge only where the caller asked for one table over both. It is fixed + for the batch's lifetime: ask once and keep the answer. + """ + + def flushed_history(self, row: BatchRow) -> int: + """The history a row has been written back to. + + A row that has never needed a flush reads its own history, so a caller may always + compare against it. + """ + + def publish(self) -> None: + """Copy the staged rows to the device, enqueuing on the batch's stream. + + The staged page indices are slot ids as stored -- unscaled, with + ``BAD_PAGE_INDEX`` preserved rather than folded to a safe slot, since only the + sentinel separates an absent block from slot zero. Publishing also converts them + into the index tables declared at construction, which is what a sparse fetch and + write-back read: there the buffer's own offset within its page is already in the + value, so an entry counts whole buffers from its table's base. + + The page table is copied, not read in place: the snapshot is the whole point, since + the host has to be free to write the next step's rows long before the kernels + reading this step's have retired. + + Two stages. The host-to-device fill runs on an upload stream of the batch's own, + so it overlaps the previous step -- the host enqueues this while that step's + forward is still executing. The batch's stream then carries one linear + device-to-device copy into the tables the fetch reads, and the conversion. Everything + expensive is in the first stage, and only the second waits for anything. + + Raises if a member holds scratch slots in a life cycle this batch gave ``SHARED``: + its layers no longer share one set of converted indices, so the single table the + batch was built around cannot serve them. Where ``mode`` was derived that means the + phases were mixed after construction fixed it; where ``SHARED`` was named, this is + the check that holds the caller to the claim. + + Raises if the previous publish has not been waited on. This call stages the + watermarks and views itself, so it would be writing into the same hazard it is + meant to close. It refuses rather than waiting because a wait here covers only its + own write: a member mutated before the call has already raced by then, and waiting + would hide that. :meth:`resize` publishes, so republishing after one needs + :meth:`wait_published` in between. + + **The second stage must be ordered after the kernels that read the previous + publish.** Those live in a captured graph, so the tables cannot be double-buffered + out of the hazard -- the capture baked their addresses, which is why the staging + sits upstream of them instead. Issuing this on the stream the graph replays on is + what provides that ordering today. A torn publish is not a stale answer: the cache views + carry pointers, so it can pair one request's residency with another's lease. + """ + + def wait_published(self) -> None: + """Block until the last :meth:`publish` has finished reading the staging rows. + + The copy is asynchronous and reads pinned memory, so it is still in flight when + :meth:`publish` returns. Every later host write to a member's page indices -- a + resize, a commit, a beam fork, an SWA recycle -- must be separated from it by this + call, or that write races the copy and publishes a torn table. + + :meth:`resize` already does it, which is why driving a batched step through that + rather than through the members is the sanctioned path. For the operations with no + ``Batch`` equivalent the simple rule is to call this once at the top of a step: it + is idempotent and free once the copy has landed. + """ + + def coalesced_page_table(self, layer_group_id: LayerGroupId) -> DeviceArray: + """``int32[max_batch, max_blocks_per_seq]`` coalesced page indices, on the host. + + A read-only DLPack array over the pinned staging the members write into, which a + consumer needing these numbers reads instead of keeping a buffer of its own. It keeps + the ``Batch`` alive for as long as the array exists. + + Read-only because the entries belong to the members, written through their page + locks: a write here would race a publish already reading the same memory, and would + be lost anyway, since ``resize`` rewrites only the entries it changes. + + An entry names a coalesced slot, which may hold several buffers; a converted index + names one buffer inside one. See :meth:`converted_index_table`. + """ + + def history_lengths(self) -> DeviceArray: + """``int32[max_batch]`` per-row history length, in tokens. + + Published by the same transfer as the page tables, because it bounds which entries + of a row are valid: pairing one step's table with another step's lengths would read + past the end of a row that may since have been reused. + + The counts a consumer wants are all derived from this and :meth:`capacities`: + ``num_blocks`` is ``history_length // tokens_per_block``, while the sparse resolve + works in units and divides by ``unit_tokens`` instead. Publishing the source rather + than the three derived counts keeps one truth. + + A free row reads zero -- a request with nothing eligible. + """ + + def capacities(self) -> DeviceArray: + """``int32[max_batch]`` per-row capacity, in tokens. See :meth:`history_lengths`.""" + class KVCacheManager: def __init__( self, @@ -830,12 +1649,77 @@ class KVCacheManager: def __del__(self) -> None: ... def shutdown(self) -> None: ... def clear_reusable_blocks(self) -> None: ... - def is_sparse(self, layer_id: LayerId, data_role: DataRole) -> bool: - """Whether the named buffer uses sparse attention. Rejects unknown buffers.""" def get_mem_pool_base_address( self, layer_id: LayerId, data_role: DataRole, index_mode: PageIndexMode | None = None ) -> MemAddress: ... + def get_sparse_cache_max_units(self, layer_id: LayerId) -> int: + """Cells a request's cache holds per buffer of this layer, in selection units. + + Counts both populations -- the history it caches and everything in flight -- so the + lease follows from this alone. Only one unit is ever in flight, so the surplus above + the selection width is retention rather than headroom. Zero when the layer has no + sparse fetch path. + """ + + def get_sparse_cache_units_per_page(self, layer_id: LayerId) -> int: + """Selection units one page of this layer's buffers holds; one when a unit is a page. + + Distinct from :attr:`PageIndexConverter.expansion`, which is how much a buffer's page + *index space* is stretched by ``tokens_per_block_override``. A sparse buffer does not + stretch its index space -- that stays 1 -- while this counts the sparse cache's own + granularity, set by the layer's ``sparse_cache.unit_tokens``. + """ + + def get_num_bytes(self, level: int, medium: CacheTier) -> int: + """Bytes this level holds in pools whose pages are ``medium``. + + A level's quota pays only for its own tier, so bytes of a second medium are drawn + from no quota and appear in no other total a caller can read -- :meth:`get_quota` + excludes them. Report them separately when accounting for what a deployment pins. + + Pair this with :meth:`sparse_backing_medium` to count what sparse buffers hold. A + sparse pool's row count follows from the device-backed pools of its group, since a + slot spans every pool, so that figure moves only when the hot level is resized. + """ + + def sparse_backing_medium(self) -> CacheTier | None: + """The medium a sparse buffer's pages are, or ``None`` when none is sparse. + + The hot level's second medium: host memory the GPU reads through a device mapping, + which is what sparse loading pulls into the per-request GPU cache. + """ + + def get_sparse_cache_buffer_count(self, layer_id: LayerId) -> int: + """Sparse buffers of this layer's life cycle, one per ``(layer, role)``. + + Each has its own cells, so a lease is this many runs of + :meth:`get_sparse_cache_max_units` cells. + """ + + def is_sparse(self, layer_id: LayerId, data_role: DataRole) -> bool: + """Whether this buffer was declared ``BufferConfig.is_sparse``. + + Every buffer locks its pages at the hot level, sparse or not, so + :meth:`_KVCache.get_base_page_indices` reports one level's numbering throughout. + What a sparse buffer changes is the *medium*: its pages are host memory, + device-mapped, in a pool of their own beside the layer's device-backed ones. + + That level is not what a caller addresses to read a unit resolved by + :meth:`fetch_sparse_pages`. Those live in the request's lease on the life cycle's + cache arena, which is a separate allocation reached through + :meth:`get_sparse_cache_base_address`. + """ + def get_page_stride(self, layer_id: LayerId, data_role: DataRole) -> int: ... + def sparse_index_table_specs(self) -> list[IndexTableSpec]: + """The index table specs every sparse life cycle needs, one per (layer, pool) run. + + A :class:`Batch` still has to be given these: it refuses a sparse buffer it was told + nothing about, because a fetch with no converted indices would read the wrong pages + rather than fail. Their shape is not a choice -- a copy walks one run per launch + from one base pointer, so the run is the request -- so this saves re-deriving the + pool layout to restate it. + """ def get_page_index_upper_bound(self, layer_id: LayerId, data_role: DataRole) -> int: ... def get_page_index_scale(self, layer_id: LayerId, data_role: DataRole) -> int: ... def get_page_index_converter( @@ -924,7 +1808,26 @@ class KVCacheManager: def enable_partial_match(self) -> bool: ... @property def enable_partial_commit(self) -> bool: ... - def supports_index_mode(self, mode: PageIndexMode) -> bool | None: ... + def supports_index_mode( + self, mode: PageIndexMode, layer_group_id: LayerGroupId | None = None + ) -> bool | None: + """``None`` means it depends on the request. + + Given a layer group the answer is definitive for every life cycle that cannot hold + scratch slots -- anything without a sliding window -- which the collapsing form + cannot express. Configuration only, so it never varies with what a request is doing. + """ + + def max_scratch_slots(self, max_blocks_per_seq: int) -> int: + """The most scratch slots one request can need, over every life cycle and buffer. + + Given the widest page table a row will get; zero where nothing can hold scratch. + + Sizes scratch staging. A caller that stages scratch slots itself asks this rather + than deriving it: which blocks can be stale turns on the window, the sinks and the + rewind margin, and a second copy of that reasoning is a second thing to keep right. + """ + @property def num_layers(self) -> int: ... @property diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_introspection.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_introspection.py index 4374facb4052..9dc70bf7797e 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_introspection.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_introspection.py @@ -246,14 +246,307 @@ def storage_utilization(manager: Any, cache_level: int = 0) -> list[float]: return list(_cpp().storage_utilization(manager, cache_level)) -def grains_for_slots(num_slots: int, slot_size_list: list[int], granularity: int) -> int: - """Return the grain count required for a pool group slot count.""" - return int(_cpp().grains_for_slots(num_slots, slot_size_list, granularity)) +def sparse_cache_resident_ordinals(kv_cache: Any, lc_id: int) -> list[int]: + """Resident unit ordinal per sparse cache entry, laid out [buffer][entry]. + + A free entry holds the bad-ordinal sentinel, which is -1, the same value as + BAD_PAGE_INDEX. Empty when the life cycle reserves no cache. + """ + return list(_cpp().sparse_cache_resident_ordinals(kv_cache, lc_id)) + + +def sparse_cache_base_cell(kv_cache: Any, lc_id: int, pool: int = 0) -> int | None: + """First cell of this request's lease in one pool's arena region, or None if it holds none. + + Identifies the lease the way the page list used to identify the reservation, now that a + lease is one contiguous run of cells rather than scattered slots. Per pool, because the + arena mirrors the hot pool group: buffers of different sizes sit in regions of their own. + """ + cell = _cpp().sparse_cache_base_cell(kv_cache, lc_id, pool) + return None if cell is None else int(cell) + + +def sparse_flush_counts(kv_cache: Any, lc_id: int) -> list[int]: + """Units the last write-back moved for this cache's row, per layer. + + -1 where no write-back has run for that layer, which distinguishes "never flushed" + from "flushed and found nothing resident". + + The lists belong to the Batch rather than the request -- one scan covers every row -- + so this reads the cache's own row out of it, and an unbatched cache has none. + """ + return list(_cpp().sparse_flush_counts(kv_cache, lc_id)) + + +def batch_page_table(batch: Any, layer_group_id: int) -> list[int]: + """A Batch's published device page table, row-major [max_batch][max_blocks_per_seq]. + + Reads the device copy, so it shows what the last publish() put there rather than what + the staging rows currently hold. Synchronises. + """ + return list(_cpp().batch_page_table(batch, layer_group_id)) + + +def batch_index_table_slot(batch: Any, layer_id: int, role: str) -> tuple[int, int]: + """``(table, entry)`` serving one buffer. + + A consumer names a table by the spec it asked for -- :meth:`Batch.index_table_id` -- + and slices it at the buffer's position in that spec, which it already knows. This is + the same answer reached from the buffer, and exists so a test can compare the two. + """ + return tuple(_cpp().batch_index_table_slot(batch, layer_id, role)) + + +def batch_index_table_base(batch: Any, layer_id: int, role: str) -> int: + """What a converted index for one buffer is added to. + + Buffers sharing a pool share it, their own offsets being already in the entries. + """ + return _cpp().batch_index_table_base(batch, layer_id, role) + + +def batch_index_table(batch: Any, table_id: int) -> list[int]: + """A Batch's converted index table, row-major [max_batch][entries][blocks_per_seq]. + + Reads the device copy, so it shows what the last publish() converted rather than what + the raw staging holds. Synchronises. + """ + return list(_cpp().batch_index_table(batch, table_id)) + + +def batch_gpu_coalesced_page_table(batch: Any, layer_group_id: int) -> Any: + """One layer group's coalesced page table on the device, as a DLPack array. + + ``int32[max_batch, max_blocks_per_seq]``. A view rather than a copy, so unlike the + other hooks here it neither synchronises nor returns a plain list: callers check the + device-array contract and that the pointer is stable across resizes. The indices a + consumer actually reads come from ``Batch.coalesced_page_table()``, on the host. + """ + return _cpp().batch_gpu_coalesced_page_table(batch, layer_group_id) + + +def batch_sparse_entries_per_row(batch: Any, layer_group_id: int) -> list[int]: + """Metadata entries per layer of the lease each row holds. + + Zero where the row holds none, so this is non-zero exactly where a member holds a + reservation for that layer group. Synchronises. + """ + return list(_cpp().batch_sparse_entries_per_row(batch, layer_group_id)) -def grains_to_slots(pg_grains: int, slot_size_list: list[int], granularity: int) -> tuple[int, int]: +def batch_watermarks(batch: Any) -> tuple[list[int], list[int]]: + """A Batch's published per-row (history_length, capacity), in tokens. + + Reads the device copy, so it shows what the last publish() put there. Synchronises. + """ + history, capacity = _cpp().batch_watermarks(batch) + return list(history), list(capacity) + + +class SparseResolveResult(NamedTuple): + """State and outputs of one sparse resolve launch.""" + + resident_ordinal: list[int] + lru_order: list[int] + out_slot: list[int] + #: The compacted work list, one run per layer, of which the first ``miss_total[layer]`` + #: are valid. Entry order is whatever the CTAs raced to, so compare as a set. + miss_row: list[int] + miss_ordinal: list[int] + miss_slot: list[int] + miss_total: list[int] + miss_count: list[int] + in_flight_slot: list[int] + + +def sparse_unit_copy( + host_pool: bytes | list[int], + host_page_stride: int, + page_table: list[int], + device_arena: bytes | list[int], + base_cell: int, + num_blocks: int, + cells_per_buffer: int, + buffer_bytes: int, + num_chunks: int, + chunk_bytes: int, + chunk_stride: int, + unit_stride: int, + expansion: int, + miss_ordinal: list[int], + miss_slot: list[int], + miss_count: list[int], + selection_width: int, + buffer_layer: list[int], + write_back: bool = False, + grid_ctas: int = 1, +) -> list[int]: + """Move units between a caller-supplied cold pool and cache entries. + + Exercises the copy's address arithmetic against a known byte pattern. Returns the + destination pool: the cache pool when fetching, the cold pool when writing back. + + ``page_table`` holds converted indices, [num_buffers][num_blocks]: one entry per + buffer, each counting whole buffers rather than coalesced slots, so a buffer's offset + within its page is already in the value and ``host_page_stride`` is one buffer. + + ``buffer_layer`` maps each buffer to the work-list row it reads, so buffers of one + layer share a list; its length is the number of buffers copied. + + ``grid_ctas`` is how many CTAs walk the compacted list. Only throughput depends on it, + so a result must not. Synchronises. + """ + return list( + _cpp().sparse_unit_copy( + list(host_pool), + host_page_stride, + page_table, + list(device_arena), + base_cell, + num_blocks, + cells_per_buffer, + buffer_bytes, + num_chunks, + chunk_bytes, + chunk_stride, + unit_stride, + expansion, + miss_ordinal, + miss_slot, + miss_count, + selection_width, + buffer_layer, + write_back, + grid_ctas, + ) + ) + + +class SparseFlushList(NamedTuple): + """Entries holding units a watermark move has just passed.""" + + #: The compacted list, one run per layer, of which the first ``total[layer]`` are + #: valid. Entry order is whatever the CTAs raced to, so compare as a set. + row: list[int] + ordinal: list[int] + slot: list[int] + total: list[int] + count: list[int] + + +def sparse_flush_list( + resident_ordinal: list[int], + num_layers: int, + entries_per_layer: int, + from_units: int, + to_units: int, + max_flush_units: int, +) -> SparseFlushList: + """Collect the entries whose units fall in ``[from_units, to_units)``. + + These are the units the watermark has just made selectable, so they must reach their + host pages before any later fetch can read them. Synchronises. + """ + row, ordinal, slot, total, count = _cpp().sparse_flush_list( + resident_ordinal, + num_layers, + entries_per_layer, + from_units, + to_units, + max_flush_units, + ) + return SparseFlushList(list(row), list(ordinal), list(slot), list(total), list(count)) + + +def sparse_resolve( + resident_ordinal: list[int], + lru_order: list[int], + num_layers: int, + entries_per_layer: int, + selected: list[int], + selection_width: int, + history_units: int | list[int], + capacity_units: int | list[int] = 0, + max_in_flight_units: int = 0, + inactive_rows: list[int] | None = None, +) -> SparseResolveResult: + """Run the sparse resolve kernel over caller-supplied cache state. + + Issues one launch per layer, as a step does: a layer's selection is not known until + the previous layer has run. Cache state is therefore [row][layer][entry], following + the view each row owns, while the selection and every output are [layer][row][width] + so that each launch sees a contiguous slice of its own. + + A negative selection entry is padding. ``inactive_rows`` get a default view, standing + in for a row holding no reservation. Ordinals in ``[history_units, capacity_units)`` are in + flight: they are made resident before the selection takes any victim, are never + chosen as victims themselves, and never enter the miss list because their host pages + hold nothing to fetch. Synchronises. + """ + # One row unless the caller passes per-row lists; the kernel is batched either way. + history = [history_units] if isinstance(history_units, int) else list(history_units) + if isinstance(capacity_units, int): + capacity = [max(capacity_units, h) for h in history] + else: + capacity = list(capacity_units) + + ( + resident, + lru, + out, + miss_row, + miss_ord, + miss_slot, + miss_total, + miss_count, + in_flight, + ) = _cpp().sparse_resolve( + resident_ordinal, + lru_order, + num_layers, + entries_per_layer, + selected, + selection_width, + history, + capacity, + max_in_flight_units, + list(inactive_rows or []), + ) + return SparseResolveResult( + list(resident), + list(lru), + list(out), + list(miss_row), + list(miss_ord), + list(miss_slot), + list(miss_total), + list(miss_count), + list(in_flight), + ) + + +def grains_for_slots( + num_slots: int, + slot_size_list: list[int], + granularity: int, + pools: list[int] | None = None, +) -> int: + """Return the grain count required for a pool group slot count. + + ``pools`` names which of ``slot_size_list``'s pools to count, defaulting to all of + them -- which is what a caller with a single budget means. + """ + return int(_cpp().grains_for_slots(num_slots, slot_size_list, granularity, pools)) + + +def grains_to_slots( + pg_grains: int, + slot_size_list: list[int], + granularity: int, + pools: list[int] | None = None, +) -> tuple[int, int]: """Return (slot count, consumed grains) for a pool group grain budget.""" - slots, used = _cpp().grains_to_slots(pg_grains, slot_size_list, granularity) + slots, used = _cpp().grains_to_slots(pg_grains, slot_size_list, granularity, pools) return int(slots), int(used) @@ -263,10 +556,28 @@ def ratio_to_slot_count_list( ratio_list: list[float], granularity: int, min_slots: list[int], + second_quota: int | None = None, + second_granularity: int | None = None, + second_pools: list[list[int]] | None = None, ) -> list[int]: - """Return slot counts by pool group for a quota and ratio list.""" + """Return slot counts by pool group for a quota and ratio list. + + ``slot_size_lists`` holds every pool's size. A second medium's budget is given by + ``second_quota`` with its own ``second_granularity`` and ``second_pools``; the quota + pays for every pool that budget does not name, so the two cannot overlap or leave a + pool unpaid. + """ return list( - _cpp().ratio_to_slot_count_list(quota, slot_size_lists, ratio_list, granularity, min_slots) + _cpp().ratio_to_slot_count_list( + quota, + slot_size_lists, + ratio_list, + granularity, + min_slots, + second_quota, + second_granularity, + second_pools, + ) ) @@ -337,3 +648,8 @@ def compute_slots_for_batch( ) -> list[int]: """Return the minimum per-pool-group slot counts to support ``batch``.""" return list(_cpp().compute_slots_for_batch(manager, batch, tokens_per_block, swa_scratch_reuse)) + + +def tier_allocation_granularity(tier: Any, quota: int) -> int: + """Return the unit a ``tier`` of ``quota`` bytes allocates its pools in.""" + return _cpp().tier_allocation_granularity(tier, quota) diff --git a/tests/unittest/kv_cache_manager_v2_tests/fake_engine.py b/tests/unittest/kv_cache_manager_v2_tests/fake_engine.py index 08c6e1aa04eb..6bd39a3513b6 100644 --- a/tests/unittest/kv_cache_manager_v2_tests/fake_engine.py +++ b/tests/unittest/kv_cache_manager_v2_tests/fake_engine.py @@ -156,14 +156,14 @@ def _check_pages( is_ssm = isinstance(layer_cfg, SsmLayerConfig) tokens_per_block = 1 if is_ssm else self.tokens_per_block_map[layer_id][buf_id] token_bytes = buf.size if is_ssm else exact_div(buf.size, tokens_per_block) + lc_id = manager.get_layer_group_id(layer_id) index_mode = ( PageIndexMode.SHARED - if kv_cache.supports_index_mode(PageIndexMode.SHARED) + if kv_cache.supports_index_mode(PageIndexMode.SHARED, lc_id) else PageIndexMode.PER_LAYER ) pool = manager.get_mem_pool_base_address(layer_id, role, index_mode) stride = manager.get_page_stride(layer_id, role) - lc_id = manager.get_layer_group_id(layer_id) if is_ssm: # SSM: only one page (the SSM slot), only check the last history token if not history: @@ -174,8 +174,12 @@ def _check_pages( else: base_pages = kv_cache.get_base_page_indices(lc_id, beam) page_converter = manager.get_page_index_converter(layer_id, role) - scratch_desc = kv_cache.get_scratch_desc(lc_id) - pages = page_converter(base_pages, index_mode, scratch_desc) + pages = page_converter( + base_pages, + index_mode, + kv_cache.get_swa_scratch_range(lc_id) or None, + kv_cache.get_scratch_slots(lc_id), + ) capacity = kv_cache.capacity history_len = len(history) if is_ssm: @@ -232,14 +236,14 @@ def _write_new_tokens( is_ssm = isinstance(layer_cfg, SsmLayerConfig) tokens_per_block = 1 if is_ssm else self.tokens_per_block_map[layer_id][buf_id] token_bytes = buf.size if is_ssm else exact_div(buf.size, tokens_per_block) + lc_id = manager.get_layer_group_id(layer_id) index_mode = ( PageIndexMode.SHARED - if kv_cache.supports_index_mode(PageIndexMode.SHARED) + if kv_cache.supports_index_mode(PageIndexMode.SHARED, lc_id) else PageIndexMode.PER_LAYER ) pool = manager.get_mem_pool_base_address(layer_id, role, index_mode) stride = manager.get_page_stride(layer_id, role) - lc_id = manager.get_layer_group_id(layer_id) if is_ssm: # SSM: write only the last input token at position 0 of the SSM page ssm_idx = kv_cache.get_ssm_block_base_index(lc_id, beam) @@ -252,8 +256,12 @@ def _write_new_tokens( : div_up(history_len + len(input), tokens_per_block) ] page_converter = manager.get_page_index_converter(layer_id, role) - scratch_desc = kv_cache.get_scratch_desc(lc_id) - pages = page_converter(base_pages, index_mode, scratch_desc) + pages = page_converter( + base_pages, + index_mode, + kv_cache.get_swa_scratch_range(lc_id) or None, + kv_cache.get_scratch_slots(lc_id), + ) capacity = kv_cache.capacity input_range = HalfOpenRange(history_len, history_len + len(input)) if not is_ssm: diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py index 923f4767eda2..14c540318766 100755 --- a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py +++ b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py @@ -13,6 +13,7 @@ # limitations under the License. import array +import contextlib import functools import gc import hashlib @@ -45,12 +46,14 @@ BufferConfig, BufferId, CacheLevel, + CacheTier, CorruptedError, CudaStream, DataRole, DiskCacheTierConfig, GpuCacheTierConfig, HostCacheTierConfig, + IndexTableSpec, KVCacheDesc, KVCacheEventManager, KVCacheManager, @@ -58,7 +61,6 @@ KVCacheUpdatedData, LayerGroupId, LayerId, - LogicError, MemAddress, OutOfPagesError, PageIndexMode, @@ -66,6 +68,7 @@ PoolRebalanceConfig, ReuseScope, SlidingWindowSize, + SparseCacheConfig, SsmLayerConfig, SwaScratchReuseConfig, TokenId, @@ -104,12 +107,14 @@ BufferConfig, BufferId, CacheLevel, + CacheTier, CorruptedError, CudaStream, DataRole, DiskCacheTierConfig, GpuCacheTierConfig, HostCacheTierConfig, + IndexTableSpec, KVCacheDesc, KVCacheEventManager, KVCacheManager, @@ -117,7 +122,6 @@ KVCacheUpdatedData, LayerGroupId, LayerId, - LogicError, MemAddress, OutOfPagesError, PageIndexMode, @@ -125,6 +129,7 @@ PoolRebalanceConfig, ReuseScope, SlidingWindowSize, + SparseCacheConfig, SsmLayerConfig, SwaScratchReuseConfig, TokenId, @@ -171,6 +176,26 @@ from kernels import HostGate, enable_kernel_delay +# A Batch that is only inspected, never published, still has to name a stream. The null +# stream is the honest answer for those: nothing is ever enqueued on it. +NULL_STREAM = cast(CudaStream, 0) + + +class _ScratchDesc(NamedTuple): + """A life cycle's scratch range and the slot ids serving it.""" + + range: HalfOpenRange + slot_ids: list[int] + + +def _scratch_desc(kv_cache, lg_id): + """The scratch range and the slots serving it, or None when none of the blocks do.""" + scratch_range = kv_cache.get_swa_scratch_range(lg_id) + if not scratch_range: + return None + return _ScratchDesc(scratch_range, kv_cache.get_scratch_slots(lg_id)) + + def get_cached_cuda_event_type(): try: from bindings.internal.batch_manager.kv_cache_manager_v2 import _introspection @@ -219,146 +244,5488 @@ def on_gc_event(phase, info): gc.set_debug(0) -def assert_no_ref_cycle(func): - """Decorator to wrap test methods with GC debugging context.""" +def assert_no_ref_cycle(func): + """Decorator to wrap test methods with GC debugging context.""" + + @functools.wraps(func) + def wrapper(self, *args, **kwargs): + with ref_cycle_check_context(): + result = func(self, *args, **kwargs) + return result + + return wrapper + + +class TestCacheLevelStorage(unittest.TestCase): + def test_grains_to_slots_rejects_zero_divisors(self) -> None: + invalid_inputs = [ + (1, [0], 16 << 20), + (1, [16 << 20], 0), + ] + + for grains, slot_sizes, granularity in invalid_inputs: + with self.subTest(slot_sizes=slot_sizes, granularity=granularity): + with self.assertRaisesRegex((ValueError, RuntimeError), "must be positive"): + _introspection.grains_to_slots(grains, slot_sizes, granularity) + + def test_grains_to_slots_refines_proportional_lower_bound(self) -> None: + granularity = 16 << 20 + slot_size_list = [16_252_928, 4_063_232] + min_slots = 157 + + grains = _introspection.grains_for_slots(min_slots, slot_size_list, granularity) + slots, used = _introspection.grains_to_slots(grains, slot_size_list, granularity) + + self.assertGreaterEqual(slots, min_slots) + self.assertLessEqual(used, grains) + + # One mebibyte per slot against a one-mebibyte grain, so a slot is a grain and the + # arithmetic a reader has to follow is the sharing rather than the rounding. + GRAIN = 1 << 20 + TWO_MEDIA = [1 << 20, 1 << 20] + ONE_MEDIUM = [1 << 20] + HOST = [1] + + def _split(self, quota, host_quota, min_slots, ratio, groups=2, device_only=()): + """Slot counts for `groups` groups of one device pool and, usually, one host pool. + + A group named in `device_only` holds a device pool alone, which is how a group that + costs the host budget nothing is really shaped -- the quota pays for every pool the + host budget does not name, so there is no way to leave one unpaid. + """ + return _introspection.ratio_to_slot_count_list( + quota, + [self.ONE_MEDIUM if g in device_only else self.TWO_MEDIA for g in range(groups)], + ratio, + self.GRAIN, + min_slots, + second_quota=host_quota, + second_granularity=self.GRAIN, + second_pools=[[] if g in device_only else self.HOST for g in range(groups)], + ) + + def test_a_group_with_no_pool_on_the_second_budget_is_never_bound_by_it(self) -> None: + """An empty selector is how "this group costs that medium nothing" is said.""" + slots = self._split( + quota=200 << 20, + host_quota=10 << 20, + min_slots=[1, 1], + ratio=[0.5, 0.5], + device_only=(0,), + ) + self.assertGreater(slots[0], 10, "group 0 pays no host budget, so it is not held to it") + self.assertLessEqual(slots[1], 10) + + def test_a_second_budget_nobody_exceeds_changes_nothing(self) -> None: + generous = self._split( + quota=60 << 20, host_quota=1000 << 20, min_slots=[1, 1], ratio=[0.5, 0.5] + ) + # Without a second budget the quota pays for both pools, so the comparison is + # against a one-pool group of the same device cost per slot. + unbounded = _introspection.ratio_to_slot_count_list( + 60 << 20, [self.ONE_MEDIUM] * 2, [0.5, 0.5], self.GRAIN, [1, 1] + ) + self.assertEqual(generous, unbounded) + + def test_a_host_bound_group_releases_its_device_budget(self) -> None: + """The reason this is one pass and not a quota followed by a clamp. + + Group 0 has no host pool and group 1 does. Starve the host budget and group 1 is + held down -- the device grains it then cannot use must reach group 0 rather than + going unspent, which is precisely what clamping afterwards failed to do. + """ + held = self._split( + quota=200 << 20, + host_quota=4 << 20, + min_slots=[1, 1], + ratio=[0.5, 0.5], + device_only=(0,), + ) + free = self._split( + quota=200 << 20, + host_quota=1000 << 20, + min_slots=[1, 1], + ratio=[0.5, 0.5], + device_only=(0,), + ) + self.assertLessEqual(held[1], 4, "the host budget bound group 1") + self.assertGreater( + held[0], free[0], "and group 0 grew into the device budget group 1 gave up" + ) + + def test_both_budgets_are_kept_and_the_ratio_gives_way(self) -> None: + """Both bounds hold exactly; only the ratio is negotiable.""" + for ratio in ([0.5, 0.5], [0.9, 0.1], [0.01, 0.99]): + with self.subTest(ratio=ratio): + slots = self._split( + quota=120 << 20, host_quota=60 << 20, min_slots=[5, 5], ratio=ratio + ) + self.assertLessEqual(sum(slots), 60, "the host budget was never exceeded") + self.assertLessEqual(sum(slots), 120, "nor the device quota") + self.assertGreaterEqual(slots[0], 5) + self.assertGreaterEqual(slots[1], 5) + + def test_a_second_budget_too_small_for_the_constraints_is_rejected(self) -> None: + """Both budgets are hard, so their only true conflict is a configuration error. + + It names neither the ratio nor the device quota, which is what makes it a + construction-time property: a rebalance changes those and cannot start failing. + """ + with self.assertRaisesRegex((ValueError, RuntimeError), "cannot hold"): + self._split(quota=200 << 20, host_quota=40 << 20, min_slots=[50, 50], ratio=[0.5, 0.5]) + + def test_an_empty_pool_selector_is_refused_rather_than_hung(self) -> None: + """grains_to_slots searches for a slot count; with no pool there is none to find. + + It used to run past SlotCount's range looking for one, which is a hang rather than + an answer -- the kind of failure that is only ever seen as a stuck process. + """ + with self.assertRaisesRegex((ValueError, RuntimeError), "at least one pool"): + _introspection.grains_to_slots(16, [1 << 20], 1 << 20, pools=[]) + + def test_slot_counts_are_maximal_for_the_grains_they_cost(self) -> None: + """Grains -> slots -> grains has to be a fixed point, or sizing disagrees with itself. + + grains_to_slots can return fewer slots than the grains it consumed would buy, when + those grains divide badly across a group's pools. A count left that way is short of + what it already paid for, and is not the count a re-derivation from the same bytes + returns -- which is what the bump at the end of the allocator exists to prevent. + A second budget must not cost the primary one that property. + """ + # Sizes that do not divide the grain evenly are where the under-count shows up. + uneven = [[3 << 20, 5 << 20], [7 << 20, 1 << 20]] + # The quota pays for whatever `second` does not name, so each case says only which + # pools the second budget takes. + cases = { + "one budget": dict(second=None, second_quota=None), + "two budgets, ample second": dict(second=[[1], [1]], second_quota=256 << 20), + "one group outside the second budget": dict(second=[[1], []], second_quota=256 << 20), + } + for name, case in cases.items(): + with self.subTest(case=name): + slots = _introspection.ratio_to_slot_count_list( + 512 << 20, + uneven, + [0.5, 0.5], + self.GRAIN, + [1, 1], + second_quota=case["second_quota"], + second_granularity=self.GRAIN if case["second"] else None, + second_pools=case["second"], + ) + for group, count in enumerate(slots): + taken = set(case["second"][group]) if case["second"] else set() + primary = [p for p in range(len(uneven[group])) if p not in taken] + cost = _introspection.grains_for_slots( + count, uneven[group], self.GRAIN, pools=primary + ) + recovered, _ = _introspection.grains_to_slots( + cost, uneven[group], self.GRAIN, pools=primary + ) + self.assertEqual( + recovered, + count, + f"{name}: group {group} holds {count} slots but its own grains buy " + f"{recovered}, so re-deriving from the same bytes would disagree", + ) + if case["second"]: + spent = sum( + _introspection.grains_for_slots( + slots[g], uneven[g], self.GRAIN, pools=case["second"][g] + ) + for g in range(len(slots)) + ) + self.assertLessEqual( + spent * self.GRAIN, + case["second_quota"], + f"{name}: the bump spent past the second budget", + ) + + def test_the_division_settles_promptly_on_mismatched_grains(self) -> None: + """A dead loop is not the only way the greedy finish can go wrong. + + A dense pool far smaller than the device grain beside an unbudgeted host pool is the + shape where one device-grain step spans thousands of host-grain steps, and a step that + costs the quota nothing cannot be refused on budget grounds. Such a loop still + terminates, so only a clock catches it -- an assertion on the slot counts never will. + """ + sizes = [[256, 1024], [4096]] + start = time.monotonic() + slots = _introspection.ratio_to_slot_count_list( + 1 << 30, + sizes, + [0.5, 0.5], + 32 << 20, + [1, 1], + second_quota=None, + second_granularity=4096, + second_pools=[[1], []], + ) + elapsed = time.monotonic() - start + self.assertLess( + elapsed, + 2.0, + f"the division took {elapsed:.1f}s, so it is stepping rather than jumping", + ) + self.assertGreater(slots[0], 0) + self.assertGreater(slots[1], 0) + + def _all_bytes(self, slots, sizes, second_pools=None, second_granularity=None): + """What a slot count costs, each pool in the grain of the medium backing it. + + The measure CacheLevelStorage::allBytes() reports, rebuilt from the two selectors + the allocator is given. + """ + taken = set(second_pools or ()) + primary = [p for p in range(len(sizes)) if p not in taken] + total = ( + _introspection.grains_for_slots(slots, sizes, self.GRAIN, pools=primary) * self.GRAIN + ) + if taken: + total += ( + _introspection.grains_for_slots( + slots, sizes, second_granularity, pools=sorted(taken) + ) + * second_granularity + ) + return total + + def test_feeding_back_the_layouts_own_ratio_changes_nothing(self) -> None: + """The allocator must be a fixed point of its own output, or rebalancing never stops. + + _needAdjustment compares the ratio a rebalance would target against the ratio the + current layout reports, and fires when they differ. Both are all-bytes measures, so + if the slot counts this returns imply a ratio that it would answer differently, the + manager re-adjusts to the same layout forever on an unchanged profile. + """ + cases = { + "single budget, uneven sizes": dict( + sizes=[[3 << 20, 5 << 20], [7 << 20, 1 << 20], [2 << 20]], + ratio=[0.2, 0.5, 0.3], + min_slots=[1, 1, 1], + quota=512 << 20, + second_quota=None, + second_pools=None, + ), + # The case the device-only ratio got wrong: a host-backed pool is most of the + # group's bytes and none of its quota, so the two measures diverge sharply. + "host-backed pool dominates the bytes": dict( + sizes=[[1 << 20, 64 << 20], [1 << 20]], + ratio=[0.5, 0.5], + min_slots=[1, 1], + quota=200 << 20, + second_quota=4096 << 20, + second_pools=[[1], []], + ), + "second budget binds": dict( + sizes=[[1 << 20, 8 << 20], [1 << 20, 1 << 20]], + ratio=[0.7, 0.3], + min_slots=[2, 2], + quota=400 << 20, + second_quota=64 << 20, + second_pools=[[1], [1]], + ), + } + for name, case in cases.items(): + with self.subTest(case=name): + second_granularity = self.GRAIN if case["second_pools"] else None + + def solve(ratio, case=case, second_granularity=second_granularity): + return _introspection.ratio_to_slot_count_list( + case["quota"], + case["sizes"], + ratio, + self.GRAIN, + case["min_slots"], + second_quota=case["second_quota"], + second_granularity=second_granularity, + second_pools=case["second_pools"], + ) + + first = solve(case["ratio"]) + achieved = [ + self._all_bytes( + first[g], + case["sizes"][g], + case["second_pools"][g] if case["second_pools"] else None, + second_granularity, + ) + for g in range(len(first)) + ] + total = sum(achieved) + again = solve([b / total for b in achieved]) + self.assertEqual( + again, + first, + f"{name}: {first} reports ratio {[b / total for b in achieved]}, which " + f"the allocator answers with {again} -- a rebalance that never settles", + ) + + def test_ratio_to_slot_count_list_preserves_min_slots(self) -> None: + granularity = 16 << 20 + slot_size_lists = [ + [16_252_928, 4_063_232], + [491_520, 126_720, 15_872], + [31_457_280, 7_864_320], + ] + min_slots = [157, 3907, 157] + total_min_grains = sum( + _introspection.grains_for_slots(slots, sizes, granularity) + for slots, sizes in zip(min_slots, slot_size_lists) + ) + + slot_counts = _introspection.ratio_to_slot_count_list( + total_min_grains * granularity, + slot_size_lists, + [0.2, 0.5, 0.3], + granularity, + min_slots, + ) + + for slot_count, min_slot in zip(slot_counts, min_slots): + self.assertGreaterEqual(slot_count, min_slot) + + +def create_config( + tokens_per_block: int, + gpu_quota: int, + host_quota: int, + disk_quota: int, + num_layers: int, + window_size: SlidingWindowSize, + sink_tokens: int, + kv_buf_size: int = 8192, + block_quant_buf_size: int | None = None, +) -> KVCacheManagerConfig: + layer_buffers = [ + BufferConfig(role=Role.KEY, size=kv_buf_size), + BufferConfig(role=Role.VALUE, size=kv_buf_size), + ] + if block_quant_buf_size is not None: + layer_buffers.extend( + [ + BufferConfig(role=Role.KEY_BLOCK_QUANT, size=block_quant_buf_size), + BufferConfig(role=Role.VALUE_BLOCK_QUANT, size=block_quant_buf_size), + ] + ) + disk_path_candidates = ["/workspace/", "/tmp/nvidia-mps/", "/tmp"] + disk_path = next(p for p in disk_path_candidates if os.path.exists(p)) + assert gpu_quota > 0 + cache_tiers = [ + GpuCacheTierConfig(quota=gpu_quota), + HostCacheTierConfig(quota=host_quota), + DiskCacheTierConfig(quota=disk_quota, path=disk_path), + ] + cache_tiers = [t for t in cache_tiers if t.quota > 0] + return KVCacheManagerConfig( + tokens_per_block=tokens_per_block, + cache_tiers=[t for t in cache_tiers if t.quota > 0], + layers=[ + AttentionLayerConfig( + layer_id=layer_id, + buffers=deepcopy(layer_buffers), + sliding_window_size=window_size if layer_id % 2 == 0 else None, + num_sink_tokens=sink_tokens if layer_id % 2 == 0 else None, + ) + for layer_id in typed_range(LayerId(num_layers)) + ], + ) + + +class TestSparseBufferConfig(unittest.TestCase): + def test_flag_defaults_and_copy(self): + self.assertFalse(BufferConfig("key", 4096).is_sparse) + buffer = BufferConfig("key", 4096, 2, is_sparse=True) + copied = deepcopy(buffer) + self.assertTrue(copied.is_sparse) + self.assertEqual(copied.tokens_per_block_override, 2) + buffer.is_sparse = False + self.assertTrue(copied.is_sparse) + + def test_the_quantised_shape_gives_each_pool_its_own_arena_region(self): + """Data and block scales differ in size, so they land in pools of their own. + + The arena mirrors the hot pool group, so an index the fetch emits is relative to the + region of the buffer it was emitted for -- which is why the base address is per + buffer rather than per life cycle. + """ + init_cuda_once() + data = BufferConfig( + "key", + 4096, + is_sparse=True, + ) + scales = BufferConfig( + "key_scale", + 128, + is_sparse=True, + ) + config = KVCacheManagerConfig( + tokens_per_block=128, + cache_tiers=[GpuCacheTierConfig(64 << 20), HostCacheTierConfig(64 << 20)], + constraints=[BatchDesc(kv_caches=[KVCacheDesc(capacity=1024, history_length=0)] * 2)], + # A pool group sizes its slots from its device-backed pools. + layers=[ + AttentionLayerConfig( + 0, + [data, scales, BufferConfig("index", 256)], + sparse_cache=SparseCacheConfig( + num_selected_blocks=2, max_cached_units=4, unit_tokens=128 + ), + ) + ], + ) + manager = KVCacheManager(config) + try: + with TemporaryCudaStream([]) as stream: + batch = Batch( + manager, + cast(CudaStream, stream.handle), + max_batch_size=1, + max_blocks_per_seq=1, + index_tables=manager.sparse_index_table_specs(), + ) + bases = { + role: int(batch.get_sparse_cache_base_address(LayerId(0), role)) + for role in ("key", "key_scale") + } + self.assertEqual(len(set(bases.values())), 2, bases) + finally: + manager.shutdown() + + def test_sparse_sizing_defaults_and_roundtrip(self): + plain = BufferConfig("key", 4096) + self.assertFalse(plain.is_sparse) + self.assertEqual(plain.sparse_unit_chunks, 1) + + buffer = BufferConfig("key", 4096, is_sparse=True, sparse_unit_chunks=4) + copied = deepcopy(buffer) + buffer.sparse_unit_chunks = 1 + self.assertEqual(copied.sparse_unit_chunks, 4) + self.assertTrue(copied.is_sparse) + + cache = SparseCacheConfig(num_selected_blocks=16, max_cached_units=64, unit_tokens=4) + copied_cache = deepcopy(cache) + cache.max_cached_units = 0 + self.assertEqual(copied_cache.max_cached_units, 64) + self.assertEqual(copied_cache.num_selected_blocks, 16) + self.assertEqual(copied_cache.unit_tokens, 4) + + def test_rejects_unit_chunks_on_a_dense_buffer(self): + layer = AttentionLayerConfig(0, [BufferConfig("key", 4096, sparse_unit_chunks=2)]) + with self.assertRaisesRegex(RuntimeError, "sparseUnitChunks set on a non-sparse buffer"): + self._make_config(layer) + + def test_rejects_a_cache_on_a_layer_with_no_sparse_buffer(self): + """sparse_cache configures a fetch path, so there must be something to fetch for.""" + layer = AttentionLayerConfig( + 0, + [BufferConfig("key", 4096)], + sparse_cache=SparseCacheConfig( + num_selected_blocks=2, max_cached_units=4, unit_tokens=128 + ), + ) + with self.assertRaisesRegex(RuntimeError, "no sparse buffer to serve"): + self._make_config(layer) + + def test_rejects_cache_narrower_than_selection(self): + layer = AttentionLayerConfig( + 0, + [ + BufferConfig("key", 4096, is_sparse=True), + # A pool group sizes its slots from its device-backed pools, so a sparse + # layer carries one -- which is also the shape a model uses. + BufferConfig("index", 256), + ], + sparse_cache=SparseCacheConfig( + num_selected_blocks=16, max_cached_units=8, unit_tokens=128 + ), + ) + with self.assertRaisesRegex(RuntimeError, "maxCachedUnits"): + self._make_config(layer) + + def test_rejects_a_page_its_units_do_not_divide(self): + layer = AttentionLayerConfig( + 0, + [ + BufferConfig( + "key", + 4096, + is_sparse=True, + # 3 chunks do not divide 4096, so a unit has no whole-byte size. + sparse_unit_chunks=3, + ), + # A pool group sizes its slots from its device-backed pools, so a sparse + # layer carries one -- which is also the shape a model uses. + BufferConfig("index", 256), + ], + sparse_cache=SparseCacheConfig( + num_selected_blocks=16, max_cached_units=64, unit_tokens=4 + ), + ) + with self.assertRaisesRegex(RuntimeError, "divide evenly"): + self._make_config(layer) + + def test_rejects_both_unit_knobs_on_one_buffer(self): + layer = AttentionLayerConfig( + 0, + [ + BufferConfig( + "key", + 4096, + tokens_per_block_override=4, + is_sparse=True, + ), + # A pool group sizes its slots from its device-backed pools, so a sparse + # layer carries one -- which is also the shape a model uses. + BufferConfig("index", 256), + ], + sparse_cache=SparseCacheConfig( + num_selected_blocks=16, max_cached_units=64, unit_tokens=4 + ), + ) + with self.assertRaisesRegex(RuntimeError, "tokensPerBlockOverride is not supported"): + self._make_config(layer) + + @staticmethod + def _make_config(layer: "AttentionLayerConfig") -> "KVCacheManagerConfig": + return KVCacheManagerConfig( + tokens_per_block=4, + cache_tiers=[GpuCacheTierConfig(4 << 20), HostCacheTierConfig(4 << 20)], + layers=[layer], + ) + + def test_rejects_a_layer_whose_buffers_are_all_sparse(self): + """A slot is one row across a group's pools, priced in the device bytes it costs. + + Sparse pools are host memory and cost none, so a group holding only those has no + slot count to derive. Raised when storage is built rather than by config + validation: a pool group forms across layers, so no single layer can be judged. + """ + layer = AttentionLayerConfig( + 0, + [BufferConfig("key", 4096, is_sparse=True)], + sparse_cache=SparseCacheConfig( + num_selected_blocks=2, max_cached_units=4, unit_tokens=4 + ), + ) + with self.assertRaisesRegex((ValueError, RuntimeError), "wholly sparse"): + KVCacheManager(self._make_config(layer)) + + def test_a_layer_may_mix_sparse_and_dense_buffers(self): + """Sparsity picks a buffer's pool and nothing above it. + + An indexer and the data it selects over are one attention computation, so forcing + them into separate layers to satisfy storage would be the tail wagging the dog. + """ + layer = AttentionLayerConfig( + 0, [BufferConfig("key", 4096, is_sparse=True), BufferConfig("value", 4096)] + ) + config = KVCacheManagerConfig( + tokens_per_block=4, + cache_tiers=[GpuCacheTierConfig(4 << 20), HostCacheTierConfig(4 << 20)], + # The cache arena is reserved up front, so a sparse deployment must say + # how many requests run together. + constraints=[BatchDesc(kv_caches=[KVCacheDesc(capacity=4096, history_length=0)] * 4)], + layers=[layer], + ) + self.assertEqual([b.is_sparse for b in config.layers[0].buffers], [True, False]) + # Built, not just configured: the pool grouping that has to tolerate the mix runs + # during construction, so a config-only check would pass even if storage refused it. + manager = KVCacheManager(config) + try: + self.assertTrue(manager.is_sparse(0, "key")) + self.assertFalse(manager.is_sparse(0, "value")) + finally: + manager.shutdown() + + +class TestSparseTierFlexibility(unittest.TestCase): + """What a sparse deployment requires of its cold tiers: nothing. + + History lives on a host-backed pool group at the *hot* level, so level 1 is an ordinary + codec target. It may be host, disk, or absent. An earlier revision required level 1 to + be HOST_MEM, because that was where sparse history resided. + """ + + def setUp(self) -> None: + init_cuda_once() + + @staticmethod + def _manager(tiers) -> KVCacheManager: + return KVCacheManager( + KVCacheManagerConfig( + tokens_per_block=128, + cache_tiers=tiers, + constraints=[ + BatchDesc(kv_caches=[KVCacheDesc(capacity=1024, history_length=0)] * 2) + ], + layers=[ + AttentionLayerConfig( + 0, + [ + BufferConfig( + "key", + 4096, + is_sparse=True, + ), + # A pool group is sized from its device-backed pools, so a sparse layer + # carries one. This is also the shape a model uses: an indexer and the + # data it selects over are one attention computation. + BufferConfig("index", 256), + ], + sparse_cache=SparseCacheConfig( + num_selected_blocks=2, max_cached_units=4, unit_tokens=128 + ), + ) + ], + ) + ) + + def test_any_cold_tier_arrangement_is_accepted(self): + for name, tiers in ( + ("gpu only", [GpuCacheTierConfig(64 << 20)]), + ("gpu + disk", [GpuCacheTierConfig(64 << 20), DiskCacheTierConfig(64 << 20, "/tmp")]), + ("gpu + host", [GpuCacheTierConfig(64 << 20), HostCacheTierConfig(64 << 20)]), + ): + with self.subTest(tiers=name): + manager = self._manager(tiers) + try: + self.assertGreater(manager.get_sparse_cache_max_units(0), 0) + finally: + manager.shutdown() + + @staticmethod + def _manager_with_host_quota(tiers, host_quota) -> KVCacheManager: + config = KVCacheManagerConfig( + tokens_per_block=128, + cache_tiers=tiers, + constraints=[BatchDesc(kv_caches=[KVCacheDesc(capacity=1024, history_length=0)] * 2)], + layers=[ + AttentionLayerConfig( + 0, + [ + BufferConfig("key", 4096, is_sparse=True), + BufferConfig("index", 256), + ], + sparse_cache=SparseCacheConfig( + num_selected_blocks=2, max_cached_units=4, unit_tokens=128 + ), + ) + ], + host_sparse_pool_quota=host_quota, + ) + return KVCacheManager(config) + + @staticmethod + def _sparse_pinned_bytes(manager: KVCacheManager) -> int: + """Bytes the hot level holds in the medium backing its sparse buffers.""" + return manager.get_num_bytes(0, manager.sparse_backing_medium()) + + def test_the_host_quota_bounds_what_the_sparse_pools_pin(self): + """The host-backed pools have no budget of their own; this is the one that bounds them. + + Their row count follows from the device-backed pools the GPU quota sizes, so without + this a deployment can only reach them by shrinking the device cache it does want. + """ + tiers = [GpuCacheTierConfig(64 << 20)] + unbounded = self._manager_with_host_quota(tiers, None) + try: + pinned = self._sparse_pinned_bytes(unbounded) + finally: + unbounded.shutdown() + self.assertGreater(pinned, 0) + + capped = self._manager_with_host_quota(tiers, pinned // 2) + try: + self.assertLessEqual(self._sparse_pinned_bytes(capped), pinned // 2) + self.assertLess(self._sparse_pinned_bytes(capped), pinned) + finally: + capped.shutdown() + + def test_a_host_quota_above_what_is_pinned_changes_nothing(self): + tiers = [GpuCacheTierConfig(64 << 20)] + unbounded = self._manager_with_host_quota(tiers, None) + try: + pinned = self._sparse_pinned_bytes(unbounded) + finally: + unbounded.shutdown() + + generous = self._manager_with_host_quota(tiers, pinned * 4) + try: + self.assertEqual(self._sparse_pinned_bytes(generous), pinned) + finally: + generous.shutdown() + + def test_a_zero_host_quota_is_refused(self): + """Zero asks for host-backed pools with no room for a slot, which is a mistake. + + Whether a pool is host-backed is declared per buffer, so this is not the way to + turn them off. + """ + with self.assertRaisesRegex(Exception, "must be positive"): + self._manager_with_host_quota([GpuCacheTierConfig(64 << 20)], 0) + + def test_a_sparse_pool_names_the_host_memory_it_actually_is(self): + """A GPU-tier group is not uniform, so the medium has to be read per pool. + + The sparse pool's pages are host memory while the dense pool beside it is device + memory, and the group reports one tier for both. A consumer that has to name the + medium of an address -- registering it for a disaggregated transfer, say -- gets + it wrong for exactly one pool if it asks the group. + """ + manager = self._manager([GpuCacheTierConfig(64 << 20)]) + try: + media = [ + pool.storage_medium for group in manager.pool_group_descs for pool in group.pools + ] + self.assertEqual( + media.count(CacheTier.HOST_MEM), 1, "the one sparse buffer is the host-backed pool" + ) + self.assertEqual( + media.count(CacheTier.GPU_MEM), + len(media) - 1, + "every other pool of a GPU tier is device memory", + ) + finally: + manager.shutdown() + + +class TestSparseCacheReservation(unittest.TestCase): + """GPU pages a request reserves to back its sparse history cache.""" + + def setUp(self) -> None: + init_cuda_once() + + @staticmethod + def _manager( + max_cached_units: int, *, num_selected_blocks: int = 2, unit_tokens: int = 128 + ) -> KVCacheManager: + # sparse_cache is absent when the fetch path is unconfigured: the buffer still gets + # sparse storage and locking, it simply has no GPU cache to read through. + # The cache must exceed the selection, since the surplus is the in-flight headroom. + sparse_cache = ( + SparseCacheConfig( + num_selected_blocks=num_selected_blocks, + max_cached_units=max_cached_units, + unit_tokens=unit_tokens, + ) + if max_cached_units > 0 + else None + ) + return KVCacheManager( + KVCacheManagerConfig( + tokens_per_block=128, + cache_tiers=[GpuCacheTierConfig(64 << 20), HostCacheTierConfig(64 << 20)], + # The cache arena is reserved up front, so a sparse deployment must say + # how many requests run together. + constraints=[ + BatchDesc(kv_caches=[KVCacheDesc(capacity=4096, history_length=0)] * 4) + ], + layers=[ + AttentionLayerConfig( + 0, + [ + BufferConfig( + "key", + 4096, + is_sparse=True, + # One contiguous run per unit; a cached buffer must say how + # to find a unit inside itself. + ), + # A pool group is sized from its device-backed pools, so a sparse layer + # carries one. This is also the shape a model uses: an indexer and the + # data it selects over are one attention computation. + BufferConfig("index", 256), + ], + sparse_cache=sparse_cache, + ) + ], + ) + ) + + @staticmethod + def _free(manager: KVCacheManager) -> int: + return sum(stats.free for stats in manager.get_storage_statistics()) + + def test_batch_rejects_resolve_shared_memory_overflow_before_fetch(self) -> None: + for max_cached_units, selection_width in ((1 << 20, 2), (1 << 16, (1 << 16) - 1)): + with self.subTest(max_cached_units=max_cached_units, selection_width=selection_width): + # Token-sized units keep the arena small while either residency metadata + # or the selection's probe table exceeds a device's per-block limit. + manager = self._manager( + max_cached_units, + num_selected_blocks=selection_width, + unit_tokens=1, + ) + try: + with self.assertRaisesRegex(RuntimeError, "above this device's limit"): + Batch( + manager, + NULL_STREAM, + max_batch_size=1, + max_blocks_per_seq=4, + index_tables=manager.sparse_index_table_specs(), + ) + finally: + manager.shutdown() + + def test_batch_accepts_resolve_shared_memory_that_fits(self) -> None: + # Three token-sized units round up to a full 128-unit cell. + manager = self._manager(max_cached_units=3, unit_tokens=1) + batch = None + try: + batch = Batch( + manager, + NULL_STREAM, + max_batch_size=1, + max_blocks_per_seq=4, + index_tables=manager.sparse_index_table_specs(), + ) + self.assertEqual(batch.max_batch_size, 1) + finally: + del batch + manager.shutdown() + + def _free_delta_across_resume(self, max_cached_units: int) -> int: + """GPU slots a single resume consumes, for the given cache width.""" + manager = self._manager(max_cached_units) + try: + before = self._free(manager) + cache = manager.create_kv_cache() + with TemporaryCudaStream([]) as stream: + self.assertTrue(cache.resume(cast(CudaStream, stream.handle))) + used = before - self._free(manager) + cache.close() + self.assertEqual(self._free(manager), before) + return used + finally: + manager.shutdown() + + def test_the_input_bound_is_one_token_whatever_the_cache_width(self): + """A sparse request contributes one query token per step, so the gap is one token. + + The bound follows from the supported shape, not from how much spare room the + reservation has: widening the cache buys reuse, not a longer input. + """ + for width in (4, 6): + manager = self._manager(max_cached_units=width) + try: + cache = manager.create_kv_cache() + try: + self.assertEqual(cache.max_input_length, 1) + finally: + cache.close() + finally: + manager.shutdown() + + dense = self._manager(max_cached_units=0) + try: + cache = dense.create_kv_cache() + try: + # None: without a cache to hold in-flight blocks the bound describes nothing. + self.assertIsNone(cache.max_input_length) + finally: + cache.close() + finally: + dense.shutdown() + + def test_shape_constants_describe_one_buffer(self): + manager = self._manager(max_cached_units=4) + try: + self.assertEqual(manager.get_sparse_cache_max_units(0), 4) + self.assertEqual(manager.get_sparse_cache_units_per_page(0), 1) + self.assertEqual(manager.get_sparse_cache_buffer_count(0), 1) + finally: + manager.shutdown() + + def test_the_cache_does_not_draw_on_the_storage_pool(self): + """Cells come from the life cycle's own arena, so a resume costs the pool nothing. + + This is what lets the hot pool group become host memory without the cache following + it there -- the two are no longer the same allocation. + """ + baseline = self._free_delta_across_resume(0) + for width in (4, 6): + self.assertEqual( + self._free_delta_across_resume(width), + baseline, + "a wider cache must not consume more storage slots", + ) + + def test_unconfigured_fetch_path_reserves_nothing(self): + manager = self._manager(max_cached_units=0) + try: + self.assertEqual(manager.get_sparse_cache_max_units(0), 0) + finally: + manager.shutdown() + + def test_resize_rejects_more_than_one_token_in_flight(self): + """The guard that keeps prefill and speculative decoding out. + + Either would give a request several query tokens in one step, which the fetch + cannot serve: it resolves one selection per request, and the units a whole step's + worth of query tokens wanted could not all be made resident at once. + """ + manager = self._manager(max_cached_units=4) + try: + cache = manager.create_kv_cache() + try: + with TemporaryCudaStream([]) as stream: + self.assertTrue(cache.resume(cast(CudaStream, stream.handle))) + # One token ahead of the watermark: an ordinary decode step. + self.assertTrue(cache.resize(1, 0)) + # Two is already a step that writes two tokens. + with self.assertRaisesRegex(RuntimeError, "can hold in flight"): + cache.resize(3, 1) + finally: + cache.close() + finally: + manager.shutdown() + + def test_a_bulk_history_jump_writes_nothing_back(self): + """How a resumed request arrives: its whole prompt becomes history in one step. + + Nothing is written back, because nothing was ever in flight -- those units reached + their host pages directly rather than through a cell. The write-back must therefore + be bounded by the capacity the resize started from, not by how far history moved, + which is a span no reservation could cover. + """ + manager = self._manager(max_cached_units=4) + cache = None + try: + cache = manager.create_kv_cache() + with TemporaryCudaStream([]) as stream: + self.assertTrue(cache.resume(cast(CudaStream, stream.handle))) + # Far more blocks than the cache could hold in flight, and no Batch at all. + # Both are the point: nothing was ever in flight, so no write-back is owed, + # and resize() may take the move without one. That is what lets a resumed + # request be seeded before it joins a Batch -- which it must be, since a + # Batch will not accept a history beyond the capacity its row already had. + self.assertTrue(cache.resize(4096, 4096)) + self.assertEqual(cache.history_length, 4096) + finally: + if cache is not None: + cache.close() + manager.shutdown() + + def test_cache_state_starts_free(self): + manager = self._manager(max_cached_units=4) + try: + lc_id = manager.get_layer_group_id(0) + buffers = manager.get_sparse_cache_buffer_count(0) + # max_cached_units counts every cell, so 4 units at expansion 1 is 4 slots and + # 4 entries per buffer. + entries = manager.get_sparse_cache_max_units(0) + self.assertEqual(buffers, 1) + self.assertEqual(entries, 4) + + cache = manager.create_kv_cache() + with TemporaryCudaStream([]) as stream: + handle = cast(CudaStream, stream.handle) + batch = Batch( + manager, + handle, + max_batch_size=1, + max_blocks_per_seq=4, + index_tables=manager.sparse_index_table_specs(), + ) + self.assertTrue(cache.resume(handle)) + # The lease comes with membership, so the cache state exists from the add. + batch.add_kv_cache(cache) + ordinals = _introspection.sparse_cache_resident_ordinals(cache, lc_id) + self.assertEqual(len(ordinals), buffers * entries) + self.assertTrue(all(o == BAD_PAGE_INDEX for o in ordinals), ordinals) + batch.remove_kv_cache(cache) + cache.close() + finally: + manager.shutdown() + + def test_cache_state_is_dropped_when_the_member_leaves(self): + manager = self._manager(max_cached_units=4) + try: + lc_id = manager.get_layer_group_id(0) + cache = manager.create_kv_cache() + with TemporaryCudaStream([]) as stream: + handle = cast(CudaStream, stream.handle) + batch = Batch( + manager, + handle, + max_batch_size=1, + max_blocks_per_seq=4, + index_tables=manager.sparse_index_table_specs(), + ) + self.assertTrue(cache.resume(handle)) + batch.add_kv_cache(cache) + self.assertNotEqual(_introspection.sparse_cache_resident_ordinals(cache, lc_id), []) + batch.remove_kv_cache(cache) + # The metadata describes cells the request no longer leases, so it goes too. + self.assertEqual(_introspection.sparse_cache_resident_ordinals(cache, lc_id), []) + batch.add_kv_cache(cache) + self.assertNotEqual(_introspection.sparse_cache_resident_ordinals(cache, lc_id), []) + batch.remove_kv_cache(cache) + cache.close() + finally: + manager.shutdown() + + def test_move_keeps_the_lease_that_leaving_would_drop(self): + """The whole reason move() exists rather than remove_row() + add_kv_cache(). + + That pair hands the arena lease back, and a lease comes back with every entry free, + so it discards residency the member is still entitled to -- which is exactly what + the test above pins. A move is not a departure, so the cells, and the base they + start at, are the same ones afterwards. + """ + manager = self._manager(max_cached_units=4) + try: + lc_id = manager.get_layer_group_id(0) + cache = manager.create_kv_cache() + with TemporaryCudaStream([]) as stream: + handle = cast(CudaStream, stream.handle) + batch = Batch( + manager, + handle, + max_batch_size=3, + max_blocks_per_seq=4, + index_tables=manager.sparse_index_table_specs(), + ) + self.assertTrue(cache.resume(handle)) + self.assertEqual(batch.add_kv_cache(cache, 2), 2) + before = _introspection.sparse_cache_resident_ordinals(cache, lc_id) + base = _introspection.sparse_cache_base_cell(cache, lc_id) + self.assertNotEqual(before, []) + + batch.move(0, 2) + + self.assertEqual(batch.row_of(cache), 0) + self.assertEqual( + _introspection.sparse_cache_resident_ordinals(cache, lc_id), + before, + "a move is not a departure, so the cache state stays", + ) + self.assertEqual( + _introspection.sparse_cache_base_cell(cache, lc_id), + base, + "and it is the same lease, not a fresh one taken from the arena", + ) + self.assertEqual(batch.num_kv_caches, 1, "one lease held throughout") + batch.remove_kv_cache(cache) + cache.close() + finally: + manager.shutdown() + + @staticmethod + def _hot_slots(manager: KVCacheManager) -> int: + return sum(s.total for s in manager.get_storage_statistics(CacheLevel(0))) + + @classmethod + def _sized_manager(cls, gpu_mb: int, capacity: int, sparse_bytes: int = 4096) -> KVCacheManager: + return KVCacheManager( + KVCacheManagerConfig( + tokens_per_block=128, + cache_tiers=[GpuCacheTierConfig(gpu_mb << 20), HostCacheTierConfig(64 << 20)], + constraints=[ + BatchDesc(kv_caches=[KVCacheDesc(capacity=capacity, history_length=0)] * 4) + ], + layers=[ + AttentionLayerConfig( + 0, + [ + BufferConfig( + "key", + sparse_bytes, + is_sparse=True, + ), + # A pool group is sized from its device-backed pools, so a sparse layer + # carries one. This is also the shape a model uses: an indexer and the + # data it selects over are one attention computation. + BufferConfig("index", 256), + ], + sparse_cache=SparseCacheConfig( + num_selected_blocks=2, max_cached_units=4, unit_tokens=128 + ), + ) + ], + ) + ) + + def test_a_sparse_pool_does_not_consume_the_gpu_quota(self): + """A level's quota measures one medium, and these pages are not made of it. + + A slot spans every pool of its group, so the row count comes from the pools the + quota does measure. Widening the sparse pool therefore buys more host bytes per + row without costing a device byte, which is the point of putting history on the + host at all. + """ + narrow = self._sized_manager(gpu_mb=64, capacity=8192, sparse_bytes=4096) + try: + wide = self._sized_manager(gpu_mb=64, capacity=8192, sparse_bytes=65536) + try: + self.assertEqual( + self._hot_slots(narrow), + self._hot_slots(wide), + "a sparse pool's bytes must not be charged to the GPU quota", + ) + finally: + wide.shutdown() + finally: + narrow.shutdown() + + def test_the_slot_count_still_meets_the_declared_floor(self): + """The constraints' floor is a minimum the quota is raised to meet.""" + manager = self._sized_manager(gpu_mb=1, capacity=8192) + try: + self.assertGreaterEqual(self._hot_slots(manager), 8192 // 128) + finally: + manager.shutdown() + + def test_the_arena_and_the_pages_have_different_bases(self): + """Two allocations, two accessors. + + get_mem_pool_base_address keeps naming the buffer's pages whether or not it is + sparse, which is what the transceiver and the KV transfer path address. The arena + is where the fetch's indices point, and it is a different allocation entirely. + """ + manager = self._manager(max_cached_units=4) + try: + pages = manager.get_mem_pool_base_address(LayerId(0), "key") + with TemporaryCudaStream([]) as stream: + batch = Batch( + manager, + cast(CudaStream, stream.handle), + max_batch_size=1, + max_blocks_per_seq=1, + index_tables=manager.sparse_index_table_specs(), + ) + arena = batch.get_sparse_cache_base_address(LayerId(0), "key") + self.assertNotEqual(int(pages), int(arena)) + self.assertNotEqual(int(arena), 0) + finally: + manager.shutdown() + + def test_the_arena_accessor_rejects_a_buffer_without_a_fetch_path(self): + manager = self._manager(max_cached_units=0) + try: + with TemporaryCudaStream([]) as stream: + batch = Batch( + manager, + cast(CudaStream, stream.handle), + max_batch_size=1, + max_blocks_per_seq=1, + index_tables=manager.sparse_index_table_specs(), + ) + with self.assertRaises(RuntimeError): + batch.get_sparse_cache_base_address(LayerId(0), "key") + finally: + manager.shutdown() + + def test_leases_are_distinct_and_within_the_arena(self): + """Two live requests must not be handed the same cells.""" + manager = self._manager(max_cached_units=4) + caches = [] + try: + lc_id = manager.get_layer_group_id(0) + cells_per_lease = manager.get_sparse_cache_buffer_count(0) * ( + manager.get_sparse_cache_max_units(0) // manager.get_sparse_cache_units_per_page(0) + ) + with TemporaryCudaStream([]) as stream: + handle = cast(CudaStream, stream.handle) + batch = Batch( + manager, + handle, + max_batch_size=2, + max_blocks_per_seq=4, + index_tables=manager.sparse_index_table_specs(), + ) + for _ in range(2): + cache = manager.create_kv_cache() + caches.append(cache) + self.assertTrue(cache.resume(handle)) + batch.add_kv_cache(cache) + bases = [_introspection.sparse_cache_base_cell(c, lc_id) for c in caches] + self.assertTrue(all(b >= 0 for b in bases), bases) + self.assertEqual(len(set(bases)), len(bases), "leases must not overlap") + for base in bases: + self.assertEqual( + base % cells_per_lease, 0, "a lease starts on a block boundary" + ) + for cache in caches: + batch.remove_kv_cache(cache) + finally: + for cache in caches: + cache.close() + manager.shutdown() + + def test_a_lease_is_returned_when_the_member_leaves(self): + manager = self._manager(max_cached_units=4) + try: + lc_id = manager.get_layer_group_id(0) + cache = manager.create_kv_cache() + try: + with TemporaryCudaStream([]) as stream: + handle = cast(CudaStream, stream.handle) + batch = Batch( + manager, + handle, + max_batch_size=1, + max_blocks_per_seq=4, + index_tables=manager.sparse_index_table_specs(), + ) + self.assertTrue(cache.resume(handle)) + # Resuming alone takes none: a request that is not being stepped cannot + # fetch, and a resumed one may still be receiving its history. + self.assertIsNone(_introspection.sparse_cache_base_cell(cache, lc_id)) + batch.add_kv_cache(cache) + self.assertIsNotNone(_introspection.sparse_cache_base_cell(cache, lc_id)) + batch.remove_kv_cache(cache) + self.assertIsNone(_introspection.sparse_cache_base_cell(cache, lc_id)) + # Returned rather than leaked: the same cells are available again. + batch.add_kv_cache(cache) + self.assertIsNotNone(_introspection.sparse_cache_base_cell(cache, lc_id)) + batch.remove_kv_cache(cache) + finally: + cache.close() + finally: + manager.shutdown() + + def test_the_watermark_can_cross_a_unit_boundary(self): + """resize() must take the long path and write back, not take the shortcut. + + Nothing is resident until the fetch path is wired, so the write-back list comes + back empty; what this covers is that the shortcut declines, the addressing is + assembled without error, and the launches are accepted. + """ + manager = self._manager(max_cached_units=4) + cache = None + batch = None + try: + cache = manager.create_kv_cache() + with TemporaryCudaStream([]) as stream: + handle = cast(CudaStream, stream.handle) + # Driven through a Batch because the write-back reads the converted indices + # a publish produces, and Batch::resize publishes at the end of every call. + batch = Batch( + manager, + handle, + max_batch_size=1, + max_blocks_per_seq=4, + index_tables=manager.sparse_index_table_specs(), + ) + self.assertTrue(cache.resume(handle)) + lc_id = manager.get_layer_group_id(0) + # tokens_per_block is 128 and the unit is a whole block. Reach the end of + # unit 0 with the last token still in flight, so that the next step's + # watermark move genuinely carries a unit from in-flight to history -- + # which is the only way a write-back has anything to consider. Seeded + # before joining, since a Batch only accepts a history its row already had + # the capacity for; no unit is crossed getting here, so nothing is written + # back yet and no published table is needed. + self.assertTrue(cache.resize(128, 127)) + batch.add_kv_cache(cache) + # The crossing below writes back block 0, so its converted index has to + # exist by then. + batch.publish() + self.assertEqual( + _introspection.sparse_flush_counts(cache, lc_id), + [-1], + "no write-back has run for this row yet", + ) + # Before the move, not after: the unit has to reach its host page while it + # is still below capacity and above history. + batch.flush_sparse_units([128]) + self.assertEqual(list(batch.resize([129], [128])), [1]) + self.assertEqual(cache.history_length, 128) + self.assertEqual( + _introspection.sparse_flush_counts(cache, lc_id), + [0], + "a write-back ran and found nothing resident to move", + ) + finally: + if cache is not None: + if batch is not None and batch.row_of(cache) is not None: + batch.remove_kv_cache(cache) + cache.close() + manager.shutdown() + + def test_one_flush_covers_every_row_that_moved(self): + """The write-back is the step's, not a request's: one call serves the batch. + + Rows move independently, and a row standing still must not be given a stale range + from whichever row did move. + """ + manager = self._manager(max_cached_units=4) + caches = [] + batch = None + try: + with TemporaryCudaStream([]) as stream: + handle = cast(CudaStream, stream.handle) + batch = Batch( + manager, + handle, + max_batch_size=3, + max_blocks_per_seq=4, + index_tables=manager.sparse_index_table_specs(), + ) + lc_id = manager.get_layer_group_id(0) + for _ in range(2): + cache = manager.create_kv_cache() + caches.append(cache) + self.assertTrue(cache.resume(handle)) + self.assertTrue(cache.resize(128, 127)) + batch.add_kv_cache(cache) + batch.publish() + + # Row 0 crosses unit 0; row 1 stays put; row 2 is free. + batch.flush_sparse_units([128, 127, 0]) + self.assertEqual( + _introspection.sparse_flush_counts(caches[0], lc_id), + [0], + "the row that moved was scanned", + ) + self.assertEqual( + _introspection.sparse_flush_counts(caches[1], lc_id), + [0], + "a row that did not move reports none rather than keeping a stale count", + ) + + # Only the row the flush covered may advance. + self.assertEqual(list(batch.resize([129, 128, None], [128, 127, None])), [1, 1, 0]) + with self.assertRaisesRegex(Exception, "written back"): + caches[1].resize(129, 128) + finally: + for cache in caches: + if batch is not None and batch.row_of(cache) is not None: + batch.remove_kv_cache(cache) + cache.close() + manager.shutdown() + + def test_a_move_inside_one_unit_still_takes_the_shortcut(self): + manager = self._manager(max_cached_units=4) + cache = None + batch = None + try: + with TemporaryCudaStream([]) as stream: + handle = cast(CudaStream, stream.handle) + batch = Batch( + manager, + handle, + max_batch_size=1, + max_blocks_per_seq=4, + index_tables=manager.sparse_index_table_specs(), + ) + cache = manager.create_kv_cache() + self.assertTrue(cache.resume(handle)) + lc_id = manager.get_layer_group_id(0) + self.assertTrue(cache.resize(64, 63)) + batch.add_kv_cache(cache) + batch.publish() + + batch.flush_sparse_units([64]) + self.assertEqual(list(batch.resize([64], [64])), [1]) + self.assertEqual(cache.history_length, 64) + # Still -1: no unit boundary was crossed, so there was nothing to write back + # and the shortcut was free to take the move. + self.assertEqual(_introspection.sparse_flush_counts(cache, lc_id), [-1]) + finally: + if cache is not None: + if batch is not None and batch.row_of(cache) is not None: + batch.remove_kv_cache(cache) + cache.close() + manager.shutdown() + + def test_unconfigured_life_cycle_has_no_cache_state(self): + manager = self._manager(max_cached_units=0) + try: + lc_id = manager.get_layer_group_id(0) + cache = manager.create_kv_cache() + with TemporaryCudaStream([]) as stream: + self.assertTrue(cache.resume(cast(CudaStream, stream.handle))) + self.assertEqual(_introspection.sparse_cache_resident_ordinals(cache, lc_id), []) + cache.close() + finally: + manager.shutdown() + + +class TestSparseResolve(unittest.TestCase): + """The sparse resolve kernel: probe index, victim choice and recency rebuild.""" + + def setUp(self) -> None: + init_cuda_once() + + @staticmethod + def _free(count: int) -> list[int]: + return [BAD_PAGE_INDEX] * count + + def _resolve( + self, + resident: list[int], + selected: list[int], + *, + num_layers: int = 1, + lru: list[int] | None = None, + history_units: int = 1 << 20, + capacity_units: int | None = None, + max_in_flight_units: int = 0, + num_rows: int = 1, + inactive_rows: list[int] | None = None, + ): + entries = len(resident) // (num_layers * num_rows) + width = len(selected) // (num_layers * num_rows) + if lru is None: + lru = list(range(entries)) * num_layers * num_rows + history = [history_units] * num_rows + capacity = [history_units if capacity_units is None else capacity_units] * num_rows + return _introspection.sparse_resolve( + resident, + lru, + num_layers, + entries, + selected, + width, + history, + capacity, + max_in_flight_units, + inactive_rows, + ) + + def test_an_empty_cache_misses_everything(self): + result = self._resolve(self._free(8), [5, 3, 7]) + self.assertEqual(result.miss_count, [3]) + # Every selection got a distinct entry, and the entries now hold those units. + self.assertEqual(len(set(result.out_slot)), 3) + for ordinal, slot in zip([5, 3, 7], result.out_slot): + self.assertEqual(result.resident_ordinal[slot], ordinal) + + def test_resident_units_hit(self): + resident = [10, 11, 12] + self._free(5) + result = self._resolve(resident, [11, 10]) + self.assertEqual(result.miss_count, [0]) + self.assertEqual(result.out_slot, [1, 0]) + # A hit leaves the resident set untouched. + self.assertEqual(result.resident_ordinal, resident) + + def test_padding_resolves_to_no_slot(self): + result = self._resolve(self._free(8), [4, -1, -1]) + self.assertEqual(result.miss_count, [1]) + self.assertEqual(result.out_slot[1:], [-1, -1]) + + def test_hits_outrank_units_fetched_this_step(self): + # Entry 0 holds unit 10 and is selected; unit 99 is new. + resident = [10] + self._free(7) + result = self._resolve(resident, [10, 99]) + fetched = result.out_slot[1] + self.assertEqual(result.lru_order[0], 0, "the hit should be most recently used") + self.assertEqual(result.lru_order[1], fetched, "the fetch ranks below the hit") + + def test_victims_come_from_the_lru_tail(self): + # A full cache, with entry 2 placed last in recency so it is the victim. + resident = [10, 11, 12, 13] + lru = [0, 1, 3, 2] + result = self._resolve(resident, [99], lru=lru) + self.assertEqual(result.miss_count, [1]) + self.assertEqual(result.out_slot, [2]) + self.assertEqual(result.resident_ordinal[2], 99) + # The units that were not selected keep their data. + self.assertEqual([result.resident_ordinal[i] for i in (0, 1, 3)], [10, 11, 13]) + + def test_in_flight_entries_are_never_evicted(self): + # Unit 12 is at or above the watermark, so it holds the only copy of its data. + resident = [10, 11, 12, 13] + lru = [0, 1, 3, 2] # entry 2 (unit 12) is least recently used + result = self._resolve(resident, [99], lru=lru, history_units=12) + self.assertEqual(result.resident_ordinal[2], 12, "an in-flight unit must survive") + self.assertEqual(result.resident_ordinal[result.out_slot[0]], 99) + + def test_in_flight_units_become_resident_without_a_fetch(self): + # Units 10 and 11 are in flight: the model is about to write them, so they need + # entries but their host pages hold nothing to copy. + result = self._resolve( + self._free(4), + [-1], + history_units=10, + capacity_units=12, + max_in_flight_units=2, + ) + self.assertEqual(result.miss_count, [0], "an in-flight unit is never fetched") + slots = result.in_flight_slot + self.assertEqual(len(set(slots)), 2, "each in-flight unit got its own entry") + self.assertEqual(result.resident_ordinal[slots[0]], 10) + self.assertEqual(result.resident_ordinal[slots[1]], 11) + + def test_an_in_flight_unit_keeps_its_entry_across_steps(self): + # Unit 10 is already resident in entry 2 and still in flight. + resident = [1, 2, 10, 3] + result = self._resolve( + resident, + [-1], + history_units=10, + capacity_units=11, + max_in_flight_units=1, + ) + self.assertEqual(result.in_flight_slot, [2], "a resident in-flight unit hits") + self.assertEqual(result.miss_count, [0]) + + def test_in_flight_units_are_resolved_before_the_selection(self): + # Two entries, one in-flight unit and one selected history unit: both fit, and the + # in-flight unit must not be the one that loses. + result = self._resolve( + self._free(2), + [5], + history_units=10, + capacity_units=11, + max_in_flight_units=1, + ) + in_flight = result.in_flight_slot[0] + selected = result.out_slot[0] + self.assertNotEqual(in_flight, selected) + self.assertEqual(result.resident_ordinal[in_flight], 10) + self.assertEqual(result.resident_ordinal[selected], 5) + self.assertEqual(result.miss_count, [1], "only the history unit is fetched") + + def test_rows_resolve_independently(self): + """Each row reads its own cache, so one row's misses must not disturb another's.""" + # Row 0 holds unit 7; row 1 is empty. Both select 7. + resident = [7, BAD_PAGE_INDEX] + [BAD_PAGE_INDEX, BAD_PAGE_INDEX] + result = self._resolve(resident, [7, 7], num_rows=2) + self.assertEqual(result.miss_count, [0, 1], "row 0 hits, row 1 fetches") + self.assertEqual(result.out_slot[0], 0, "row 0 reuses its own entry") + # Row 1's unit landed in row 1's own cache, leaving row 0 untouched. + self.assertEqual(result.resident_ordinal[:2], [7, BAD_PAGE_INDEX]) + self.assertEqual(result.resident_ordinal[2 + result.out_slot[1]], 7) + + def test_misses_compact_across_rows(self): + """Every row's misses land in one dense list, each entry naming its own row.""" + # Three rows of two entries each; rows 0 and 2 miss, row 1 hits what it selects. + resident = ( + [BAD_PAGE_INDEX, BAD_PAGE_INDEX] + + [4, BAD_PAGE_INDEX] + + [BAD_PAGE_INDEX, BAD_PAGE_INDEX] + ) + result = self._resolve(resident, [3, 4, 5], num_rows=3, history_units=8) + self.assertEqual(result.miss_count, [1, 0, 1], "row 1 hits its selection") + self.assertEqual( + result.miss_total, [2], "the list holds exactly the misses, not a slot per row" + ) + + # Position says nothing about which row an entry came from, and the CTAs race, so + # the list is only defined as a set. + entries = { + (result.miss_row[i], result.miss_ordinal[i]) for i in range(result.miss_total[0]) + } + self.assertEqual(entries, {(0, 3), (2, 5)}) + for i in range(result.miss_total[0]): + row = result.miss_row[i] + slot = result.miss_slot[i] + self.assertEqual( + result.resident_ordinal[row * 2 + slot], + result.miss_ordinal[i], + "the entry names the cell the unit was given", + ) + + def test_a_row_holding_no_reservation_is_skipped(self): + """A free row, or one that has suspended, stages a default view the kernel skips.""" + resident = [BAD_PAGE_INDEX] * 4 + result = self._resolve(resident, [5, 5], num_rows=2, inactive_rows=[1]) + # Row 0 resolved; row 1 was left exactly as the caller initialised it. + self.assertEqual(result.miss_count[0], 1) + self.assertEqual(result.miss_count[1], 0, "an inactive row does no work") + self.assertEqual(result.resident_ordinal[2:], [BAD_PAGE_INDEX] * 2) + + def test_layers_resolve_independently(self): + # Layer 0 hits, layer 1 misses. + resident = [7, BAD_PAGE_INDEX] + self._free(2) + selected = [7, -1, 42, -1] + result = self._resolve(resident, selected, num_layers=2) + self.assertEqual(result.miss_count, [0, 1]) + self.assertEqual(result.out_slot[0], 0) + self.assertEqual(result.resident_ordinal[2 + result.out_slot[2]], 42) + + def test_large_caches_preserve_victims_and_recency_order(self): + """Hits, fetched victims, and untouched entries retain their relative order.""" + for entries in (31, 32, 33, 255, 256, 257, 2048, 4096, 4224, 8192): + with self.subTest(entries=entries): + rng = random.Random(entries) + history = 100000 + in_flight = 4 + resident = list(range(entries - in_flight)) + list( + range(history, history + in_flight) + ) + resident[: entries // 8] = self._free(entries // 8) + rng.shuffle(resident) + lru = list(range(entries)) + rng.shuffle(lru) + count = min(entries // 4, 512) + hits = rng.sample( + [ordinal for ordinal in resident if 0 <= ordinal < history], count + ) + misses = list(range(50000, 50000 + count)) + selected = hits + misses + [-1] * 7 + rng.shuffle(selected) + hit_slots = { + slot + for slot, ordinal in enumerate(resident) + if ordinal in hits or history <= ordinal < history + in_flight + } + eligible = [slot for slot in lru if slot not in hit_slots] + victims = set(eligible[-count:]) + expected_lru = ( + [slot for slot in lru if slot in hit_slots] + + [slot for slot in lru if slot in victims] + + [slot for slot in lru if slot not in hit_slots | victims] + ) + try: + result = self._resolve( + resident, + selected, + lru=lru, + history_units=history, + capacity_units=history + in_flight, + max_in_flight_units=in_flight, + ) + except ValueError as error: + if "above this device's limit" in str(error): + raise unittest.SkipTest(str(error)) from error + raise + self.assertEqual(result.miss_count, [count]) + self.assertEqual(result.lru_order, expected_lru) + fetched_slots = set() + for ordinal, slot in zip(selected, result.out_slot): + if ordinal < 0: + self.assertEqual(slot, -1) + else: + self.assertEqual(result.resident_ordinal[slot], ordinal) + if ordinal in misses: + fetched_slots.add(slot) + self.assertEqual(fetched_slots, victims) + for slot, ordinal in enumerate(resident): + if slot not in victims: + self.assertEqual(result.resident_ordinal[slot], ordinal) + self.assertEqual( + result.in_flight_slot, + [resident.index(history + i) for i in range(in_flight)], + ) + + def test_colliding_selections_preserve_slots_and_recency(self) -> None: + """Sparse ordinals and padding preserve the same eviction and stable-order contract.""" + for entries in (63, 257, 2080): + with self.subTest(entries=entries): + history = 1 << 28 + resident = [i * 2048 for i in range(entries - 2)] + [history, history + 1] + lru = list(range(entries)) + random.Random(entries).shuffle(lru) + count = min(entries // 4, 32) + hits = resident[:count] + misses = [(entries + i) * 2048 for i in range(count)] + selected = [-1, *misses, -1, *reversed(hits), -1] + hit_slots = set(range(count)) | {entries - 2, entries - 1} + eligible = [slot for slot in lru if slot not in hit_slots] + victims = set(eligible[-count:]) + expected_lru = ( + [slot for slot in lru if slot in hit_slots] + + [slot for slot in lru if slot in victims] + + [slot for slot in lru if slot not in hit_slots | victims] + ) + result = self._resolve( + resident, + selected, + lru=lru, + history_units=history, + capacity_units=history + 2, + max_in_flight_units=4, + ) + self.assertEqual(result.lru_order, expected_lru) + self.assertEqual(result.miss_count, [count]) + self.assertEqual(result.miss_total, [count]) + self.assertEqual(result.in_flight_slot, [entries - 2, entries - 1, -1, -1]) + for ordinal, slot in zip(selected, result.out_slot): + if ordinal < 0: + self.assertEqual(slot, -1) + else: + self.assertEqual(result.resident_ordinal[slot], ordinal) + if ordinal in hits: + self.assertEqual(slot, resident.index(ordinal)) + for slot in set(range(entries)) - victims: + self.assertEqual(result.resident_ordinal[slot], resident[slot]) + + +class TestSparseUnitCopy(unittest.TestCase): + """Moving a unit between its cold page and its cache cell, in both directions.""" + + def setUp(self) -> None: + init_cuda_once() + + @staticmethod + def _pattern(size: int) -> list[int]: + return [(i * 7 + 13) % 251 for i in range(size)] + + def test_a_contiguous_unit_lands_at_its_entry(self): + # One layer, token-major: a unit is one contiguous run of bytes. + unit, units_per_page, pages = 32, 2, 3 + page_stride = unit * units_per_page + host = self._pattern(page_stride * pages) + # Block 2 lives in host page 0, and its unit 1 is selected. + page_table = [9, 9, 0] + ordinal = 2 * units_per_page + 1 + result = _introspection.sparse_unit_copy( + host_pool=host, + host_page_stride=page_stride, + page_table=page_table, + device_arena=[0] * (page_stride * 2), + # A non-zero lease base: entry 0 belongs to the second cell of the arena, which + # is what every request past the first gets. + base_cell=1, + num_blocks=len(page_table), + cells_per_buffer=1, + buffer_bytes=page_stride, + num_chunks=1, + chunk_bytes=unit, + chunk_stride=0, + unit_stride=unit, + expansion=units_per_page, + miss_ordinal=[ordinal], + miss_slot=[0], + miss_count=[1], + selection_width=1, + buffer_layer=list(range(1)), + ) + # Entry 0 of a lease based at cell 1 is the arena's cell 1, unit 0. + expected = host[unit : unit * 2] + self.assertEqual(result[page_stride : page_stride + unit], expected) + # Nothing else was written. + self.assertEqual(result[:page_stride], [0] * page_stride) + + def test_chunked_layout_follows_the_strides(self): + # Two heads: a unit is two runs, chunk_stride apart on both sides. + chunk, chunks, units = 16, 2, 2 + chunk_stride = chunk * units + page_stride = chunk_stride * chunks + host = self._pattern(page_stride) + result = _introspection.sparse_unit_copy( + host_pool=host, + host_page_stride=page_stride, + page_table=[0], + device_arena=[0] * (page_stride), + base_cell=0, + num_blocks=1, + cells_per_buffer=1, + buffer_bytes=page_stride, + num_chunks=chunks, + chunk_bytes=chunk, + chunk_stride=chunk_stride, + unit_stride=chunk, + expansion=units, + miss_ordinal=[1], + miss_slot=[1], + miss_count=[1], + selection_width=1, + buffer_layer=list(range(1)), + ) + # Unit 1 of each chunk moves to unit 1 of the same chunk, verbatim. + for c in range(chunks): + base = c * chunk_stride + chunk + self.assertEqual(result[base : base + chunk], host[base : base + chunk]) + # Unit 0 of each chunk was not selected. + for c in range(chunks): + base = c * chunk_stride + self.assertEqual(result[base : base + chunk], [0] * chunk) + + def test_layers_land_in_their_own_sub_slot(self): + # A cold page concatenates two layers' buffers; the cache cell is derived. + unit, layers = 16, 2 + host = self._pattern(unit * layers) + result = _introspection.sparse_unit_copy( + host_pool=host, + # One buffer, not the coalesced page: a converted index counts whole buffers. + host_page_stride=unit, + # Deliberately crossed, which is now the table's doing rather than a separate + # offset array: slot 0 of a pool coalescing two buffers converts to index 0 or + # 1, and layer 0 is given the second. + page_table=[1, 0], + device_arena=[0] * (unit * layers), + base_cell=0, + num_blocks=1, + cells_per_buffer=1, + buffer_bytes=unit, + num_chunks=1, + chunk_bytes=unit, + chunk_stride=0, + unit_stride=unit, + expansion=1, + miss_ordinal=[0, 0], + miss_slot=[0, 0], + miss_count=[1, 1], + selection_width=1, + buffer_layer=list(range(layers)), + ) + # The cold page comes from the entry and the hot cell is arithmetic on the sub-slot, + # so layer 0 reads the second cold buffer into cell 0, and layer 1 the reverse. + self.assertEqual(result[:unit], host[unit:]) + self.assertEqual(result[unit:], host[:unit]) + + def test_a_layer_without_misses_copies_nothing(self): + unit = 16 + host = self._pattern(unit) + result = _introspection.sparse_unit_copy( + host_pool=host, + host_page_stride=unit, + page_table=[0], + device_arena=[0] * (unit), + base_cell=0, + num_blocks=1, + cells_per_buffer=1, + buffer_bytes=unit, + num_chunks=1, + chunk_bytes=unit, + chunk_stride=0, + unit_stride=unit, + expansion=1, + miss_ordinal=[0], + miss_slot=[0], + miss_count=[0], + selection_width=1, + buffer_layer=list(range(1)), + ) + self.assertEqual(result, [0] * unit) + + def test_an_unaligned_unit_still_copies(self): + # 20 bytes is not a multiple of 16, so the wide path must not be taken. + unit = 20 + host = self._pattern(unit * 2) + result = _introspection.sparse_unit_copy( + host_pool=host, + host_page_stride=unit * 2, + page_table=[0], + device_arena=[0] * (unit * 2), + base_cell=0, + num_blocks=1, + cells_per_buffer=1, + buffer_bytes=unit * 2, + num_chunks=1, + chunk_bytes=unit, + chunk_stride=0, + unit_stride=unit, + expansion=2, + miss_ordinal=[1], + miss_slot=[1], + miss_count=[1], + selection_width=1, + buffer_layer=list(range(1)), + ) + self.assertEqual(result[unit:], host[unit:]) + + def test_the_work_splits_across_the_grid(self): + """Giving the copy more CTAs only divides the work list, never changes the result.""" + unit, units_per_page, misses = 16, 8, 5 + page_stride = unit * units_per_page + host = self._pattern(page_stride) + # More warps than work on the last pass: the surplus must find nothing to do. + for grid_ctas in (1, 2, 4, 8): + with self.subTest(grid_ctas=grid_ctas): + result = _introspection.sparse_unit_copy( + host_pool=host, + host_page_stride=page_stride, + page_table=[0], + device_arena=[0] * page_stride, + base_cell=0, + num_blocks=1, + cells_per_buffer=1, + buffer_bytes=page_stride, + num_chunks=1, + chunk_bytes=unit, + chunk_stride=0, + unit_stride=unit, + expansion=units_per_page, + miss_ordinal=list(range(misses)), + miss_slot=list(range(misses)), + miss_count=[misses], + selection_width=misses, + buffer_layer=list(range(1)), + grid_ctas=grid_ctas, + ) + self.assertEqual(result[: unit * misses], host[: unit * misses]) + self.assertEqual(result[unit * misses :], [0] * (page_stride - unit * misses)) + + +class TestSparseFetch(unittest.TestCase): + """KVCacheManager.fetch_sparse_pages: one layer's selections, resolved for a batch.""" + + TOKENS_PER_BLOCK = 128 + SELECTED = 2 + CACHED_UNITS = 4 + + def setUp(self) -> None: + init_cuda_once() + + # Bytes per sparse buffer. A pool holds buffers of one size, so giving two roles a + # different size is how a test puts them in a pool of their own. + SPARSE_BYTES = 4096 + + @classmethod + def _manager( + cls, + roles: tuple[str, ...] = ("key", "value"), + sizes: dict[str, int] | None = None, + ) -> KVCacheManager: + return KVCacheManager( + KVCacheManagerConfig( + tokens_per_block=cls.TOKENS_PER_BLOCK, + cache_tiers=[GpuCacheTierConfig(64 << 20), HostCacheTierConfig(64 << 20)], + # The cache arena is reserved up front, so a sparse deployment must say + # how many requests run together. + constraints=[ + BatchDesc(kv_caches=[KVCacheDesc(capacity=4096, history_length=0)] * 4) + ], + layers=[ + AttentionLayerConfig( + 0, + [ + BufferConfig( + role, + (sizes or {}).get(role, cls.SPARSE_BYTES), + is_sparse=True, + ) + for role in roles + ] + + [ + # A pool group sizes its slots from its device-backed pools, so a + # sparse layer carries one -- the shape a model uses anyway. + BufferConfig("index", 256) + ], + sparse_cache=SparseCacheConfig( + num_selected_blocks=cls.SELECTED, + max_cached_units=cls.CACHED_UNITS, + unit_tokens=cls.TOKENS_PER_BLOCK, + ), + ) + ], + ) + ) + + @staticmethod + def _buffers(roles: tuple[str, ...]) -> list[BufferId]: + return [BufferId(0, role) for role in roles] + + def _torch(self): + try: + import torch + except ImportError: # pragma: no cover - fast mode without torch + raise unittest.SkipTest("torch is needed to pass device arrays") + return torch + + @contextlib.contextmanager + def _fixture( + self, + roles: tuple[str, ...] = ("key", "value"), + rows: int = 2, + blocks: int = 4, + sizes: dict[str, int] | None = None, + ): + """A resumed member in row 0 of a two-row batch, with history to select from.""" + torch = self._torch() + manager = self._manager(roles, sizes) + cache = None + batch = None + try: + with TemporaryCudaStream([]) as stream: + handle = cast(CudaStream, stream.handle) + batch = Batch( + manager, + handle, + max_batch_size=rows, + max_blocks_per_seq=blocks, + index_tables=manager.sparse_index_table_specs(), + ) + cache = manager.create_kv_cache() + self.assertTrue(cache.resume(handle)) + # Two whole blocks of history, so ordinals 0 and 1 are eligible. + tokens = 2 * self.TOKENS_PER_BLOCK + # Seeded before joining: a Batch only accepts a history its row already had + # the capacity for, and this one is jumping straight past both. + self.assertTrue(cache.resize(tokens, tokens)) + batch.add_kv_cache(cache) + # resize() publishes, so the batch is already current here. + self.assertEqual(batch.resize([tokens, None], [tokens, None]), [True, False]) + yield torch, manager, batch, cache, handle + finally: + if cache is not None: + if batch is not None and batch.row_of(cache) is not None: + batch.remove_kv_cache(cache) + cache.close() + manager.shutdown() + + def _out(self, torch, rows: int, num_buffers: int, width: int | None = None): + """Out is compact: one index per selection slot, not per unit ordinal.""" + return torch.full( + (rows, num_buffers, self.SELECTED if width is None else width), + 12345, + dtype=torch.int32, + device="cuda", + ) + + def test_every_selection_slot_is_written(self): + """A stale index must never survive, so padded slots take the sentinel too. + + Of the slots a buffer maps to. Buffers whose indices coincide share one, and the + entries of ``out`` no buffer maps to are not written at all -- see + ``sparse_out_slots``. These two buffers do not coincide, so every entry is live. + """ + roles = ("key", "value") + with self._fixture(roles) as (torch, manager, batch, _cache, handle): + selected = torch.tensor([[0, -1], [-1, -1]], dtype=torch.int32, device="cuda") + out = self._out(torch, 2, len(roles)) + batch.fetch_sparse_pages(self._buffers(roles), selected, out, handle) + torch.cuda.synchronize() + + # Slot 0 carried a selection and resolved for both buffers of the member's row. + self.assertTrue(bool((out[0, :, 0] != BAD_PAGE_INDEX).all())) + # Slot 1 was padded, so it reads as absent rather than as the fill we started from. + self.assertTrue(bool((out[0, :, 1] == BAD_PAGE_INDEX).all())) + + def test_mirrored_pools_resolve_a_selection_to_one_index(self): + """Mirrored pools give a selection one index, so one slice of ``out`` serves both. + + K and V share a pool; their block scales share another. The two pools hold the same + number of buffers, so a lease spans the same cells in each and a selection resolves + to the same index in both. ``sparse_out_slots`` reports that as a shared slot, and + the fetch then resolves it once. + + The equality itself is no longer observable here: the merge is what removes the + second copy, so there is nothing left to compare it against. It is enforced instead + where the slots are built, which refuses to share a slot unless the two buffers' + arena geometry -- cellsPerLease, cellsPerBuffer and slotInPool -- agrees. + + Checked on the *second* member, whose lease is not zero. An index is + ``leaseIndex * cellsPerLease + slotInPool * cellsPerBuffer + ...``, and + ``cellsPerLease`` is the only term that differs between pools of unequal width, so + at lease zero a wrongly shared slot would look identical to a right one. + """ + roles = ("key", "value", "key_scale", "value_scale") + sizes = {"key_scale": 512, "value_scale": 512} + with self._fixture(roles, sizes=sizes) as (torch, manager, batch, _cache, handle): + second = manager.create_kv_cache() + try: + self.assertTrue(second.resume(handle)) + tokens = 2 * self.TOKENS_PER_BLOCK + self.assertTrue(second.resize(tokens, tokens)) + batch.add_kv_cache(second) + self.assertEqual(batch.row_of(second), 1, "the second member takes row 1") + self.assertEqual(batch.resize([tokens, tokens], [tokens, tokens]), [True, True]) + + buffers = self._buffers(roles) + self.assertEqual(batch.sparse_out_slots(buffers), [0, 1, 0, 1]) + + selected = torch.tensor([[0, 1], [0, 1]], dtype=torch.int32, device="cuda") + out = self._out(torch, 2, len(roles)) + batch.fetch_sparse_pages(buffers, selected, out, handle) + torch.cuda.synchronize() + + for row in (0, 1): + key, value = out[row, 0, :], out[row, 1, :] + self.assertTrue( + bool((key != BAD_PAGE_INDEX).all()), f"row {row} should resolve" + ) + # The two buffers of a pool sit at consecutive slots, so they differ by + # a constant. Stated as a difference because the stride is the arena's + # business, not this test's. + deltas = (value - key).tolist() + self.assertEqual(len(set(deltas)), 1, f"row {row}: one stride, {deltas}") + self.assertNotEqual(deltas[0], 0, f"row {row}: key and value share a cell") + + # Different leases, or the geometry term the sharing turns on is multiplied + # away and this proves nothing. + self.assertNotEqual(out[0, 0, :].tolist(), out[1, 0, :].tolist()) + finally: + if batch.row_of(second) is not None: + batch.remove_kv_cache(second) + second.close() + + def test_mirrored_pools_share_one_output_slice(self): + """Buffers whose indices coincide are resolved once, into one slice of ``out``. + + The scatter runs per pool, and a pool whose slices an earlier one already wrote is + not launched: the work and the writes are what the sharing saves. ``out`` keeps one + entry per buffer, so the entries no buffer maps to are left as they were, which is + what this checks -- a merged layout that still wrote them would save nothing. + """ + roles = ("key", "value", "key_scale", "value_scale") + sizes = {"key_scale": 512, "value_scale": 512} + with self._fixture(roles, sizes=sizes) as (torch, manager, batch, _cache, handle): + buffers = self._buffers(roles) + # Mirrored pools: key with its scale, value with its scale. + self.assertEqual(batch.sparse_out_slots(buffers), [0, 1, 0, 1]) + + selected = torch.tensor([[0, 1], [-1, -1]], dtype=torch.int32, device="cuda") + out = self._out(torch, 2, len(roles)) + batch.fetch_sparse_pages(buffers, selected, out, handle) + torch.cuda.synchronize() + + # The slots buffers map to carry real indices. + self.assertTrue(bool((out[0, 0, :] != BAD_PAGE_INDEX).all())) + self.assertTrue(bool((out[0, 1, :] != BAD_PAGE_INDEX).all())) + # The rest were never written, so they still hold the fill the test started + # from -- not the sentinel, which would mean the scatter had run over them. + self.assertTrue(bool((out[:, 2:, :] == 12345).all()), "a shared slice was rewritten") + + def test_a_scale_shared_by_two_buffers_keeps_its_own_slice(self): + """Coinciding is the test, not pooling: an unmirrored pool gets a slice of its own. + + One scale serving both key and value is a pool of one buffer, so its entries do not + line up with the key-and-value pool's and nothing may be shared with it. + """ + roles = ("key", "value", "kv_scale") + sizes = {"kv_scale": 512} + with self._fixture(roles, sizes=sizes) as (torch, manager, batch, _cache, handle): + buffers = self._buffers(roles) + self.assertEqual(batch.sparse_out_slots(buffers), [0, 1, 2]) + + selected = torch.tensor([[0, 1], [-1, -1]], dtype=torch.int32, device="cuda") + out = self._out(torch, 2, len(roles)) + batch.fetch_sparse_pages(buffers, selected, out, handle) + torch.cuda.synchronize() + # Nothing is shared, so every entry is a slot some buffer maps to. + self.assertTrue(bool((out[0, :, :] != BAD_PAGE_INDEX).all())) + + def test_the_hit_rate_counts_selected_units_and_nothing_else(self): + """Hits and misses are selection entries naming real history, and only those. + + A padded entry is neither, which is why the pair cannot be derived on the host from + the selection's width -- the padding is the model's and lives on the device. + """ + roles = ("key",) + with self._fixture(roles) as (torch, manager, batch, _cache, handle): + self.assertEqual(batch.peek_hit_miss_counter(), (0, 0)) + buffers = self._buffers(roles) + out = self._out(torch, 2, len(roles)) + + # Row 0 selects two real units; row 1 holds no reservation and resolves nothing. + both = torch.tensor([[0, 1], [-1, -1]], dtype=torch.int32, device="cuda") + batch.fetch_sparse_pages(buffers, both, out, handle) + torch.cuda.synchronize() + self.assertEqual(batch.peek_hit_miss_counter(), (0, 2), "a cold cache misses both") + + # The same units are resident now, so the second fetch copies nothing. + batch.fetch_sparse_pages(buffers, both, out, handle) + torch.cuda.synchronize() + self.assertEqual(batch.peek_hit_miss_counter(), (2, 2)) + + # One real selection and one padded: the padding moves neither figure. + padded = torch.tensor([[0, -1], [-1, -1]], dtype=torch.int32, device="cuda") + batch.fetch_sparse_pages(buffers, padded, out, handle) + torch.cuda.synchronize() + self.assertEqual(batch.peek_hit_miss_counter(), (3, 2)) + + batch.reset_hit_miss_counter() + torch.cuda.synchronize() + self.assertEqual(batch.peek_hit_miss_counter(), (0, 0)) + + def test_a_write_back_leaves_the_hit_rate_alone(self): + """The flush shares the copy kernel; its units moved are not cache misses.""" + roles = ("key",) + with self._fixture(roles) as (torch, manager, batch, _cache, handle): + buffers = self._buffers(roles) + out = self._out(torch, 2, len(roles)) + both = torch.tensor([[0, 1], [-1, -1]], dtype=torch.int32, device="cuda") + batch.fetch_sparse_pages(buffers, both, out, handle) + torch.cuda.synchronize() + before = batch.peek_hit_miss_counter() + + batch.flush_sparse_units([2 * self.TOKENS_PER_BLOCK, 0]) + torch.cuda.synchronize() + self.assertEqual(batch.peek_hit_miss_counter(), before) + + def test_the_output_follows_the_selection_order(self): + """Slot k of out and slot k of the selection describe the same unit. + + Padding is left where the selection put it rather than packed out, which is what + lets a caller pair the two arrays without a second index. + """ + roles = ("key",) + with self._fixture(roles) as (torch, manager, batch, _cache, handle): + forward = torch.tensor([[0, 1], [-1, -1]], dtype=torch.int32, device="cuda") + first = self._out(torch, 2, len(roles)) + batch.fetch_sparse_pages(self._buffers(roles), forward, first, handle) + + # The same two units, named in the other order: the indices must follow. + reverse = torch.tensor([[1, 0], [-1, -1]], dtype=torch.int32, device="cuda") + second = self._out(torch, 2, len(roles)) + batch.fetch_sparse_pages(self._buffers(roles), reverse, second, handle) + torch.cuda.synchronize() + + self.assertEqual(int(first[0, 0, 0]), int(second[0, 0, 1])) + self.assertEqual(int(first[0, 0, 1]), int(second[0, 0, 0])) + + def test_a_row_holding_no_reservation_reads_as_absent(self): + roles = ("key", "value") + with self._fixture(roles) as (torch, manager, batch, _cache, handle): + selected = torch.tensor([[0, 1], [0, 1]], dtype=torch.int32, device="cuda") + out = self._out(torch, 2, len(roles)) + batch.fetch_sparse_pages(self._buffers(roles), selected, out, handle) + torch.cuda.synchronize() + self.assertTrue( + bool((out[1] == BAD_PAGE_INDEX).all()), + "a free row selects nothing, however the caller filled its selection", + ) + + def test_each_buffer_of_a_layer_gets_its_own_index(self): + """A layer's buffers share the resolve but not their cells. + + The cache packs one buffer's entries densely across the pages it reserved rather + than leaving each at its nominal cell, so one selection lands somewhere different + in every buffer. That is why out carries a buffer dimension at all. + """ + roles = ("key", "value") + with self._fixture(roles) as (torch, manager, batch, _cache, handle): + selected = torch.tensor([[0, 1], [-1, -1]], dtype=torch.int32, device="cuda") + out = self._out(torch, 2, len(roles)) + batch.fetch_sparse_pages(self._buffers(roles), selected, out, handle) + torch.cuda.synchronize() + + for slot in (0, 1): + key, value = int(out[0, 0, slot]), int(out[0, 1, slot]) + self.assertNotEqual(key, BAD_PAGE_INDEX) + self.assertNotEqual(key, value, f"selection slot {slot} shares a cell") + + def test_a_hit_resolves_to_the_same_index_as_the_fetch_that_filled_it(self): + """A second call must reuse the entry rather than fetch it again.""" + roles = ("key",) + with self._fixture(roles) as (torch, manager, batch, _cache, handle): + selected = torch.tensor([[0, -1], [-1, -1]], dtype=torch.int32, device="cuda") + first = self._out(torch, 2, len(roles)) + batch.fetch_sparse_pages(self._buffers(roles), selected, first, handle) + second = self._out(torch, 2, len(roles)) + batch.fetch_sparse_pages(self._buffers(roles), selected, second, handle) + torch.cuda.synchronize() + self.assertEqual(int(first[0, 0, 0]), int(second[0, 0, 0])) + + def test_the_resolve_makes_the_selected_units_resident(self): + """The indices alone could be anything; residency is what the fetch actually changed.""" + roles = ("key",) + with self._fixture(roles) as (torch, manager, batch, cache, handle): + lg_id = manager.get_layer_group_id(0) + before = _introspection.sparse_cache_resident_ordinals(cache, lg_id) + self.assertNotIn(0, before, "nothing is cached before the first fetch") + self.assertNotIn(1, before) + + selected = torch.tensor([[0, 1], [-1, -1]], dtype=torch.int32, device="cuda") + out = self._out(torch, 2, len(roles)) + batch.fetch_sparse_pages(self._buffers(roles), selected, out, handle) + torch.cuda.synchronize() + + after = _introspection.sparse_cache_resident_ordinals(cache, lg_id) + self.assertIn(0, after, "the selected units are resident afterwards") + self.assertIn(1, after) + + def test_buffers_must_be_one_whole_layer(self): + roles = ("key", "value") + with self._fixture(roles) as (torch, manager, batch, _cache, handle): + selected = torch.tensor([[0, -1], [-1, -1]], dtype=torch.int32, device="cuda") + # A subset would leave the buffers it omits stale while the metadata, which the + # resolve updates for the whole layer, claimed they were resident. + with self.assertRaises(RuntimeError): + batch.fetch_sparse_pages( + [BufferId(0, "key")], selected, self._out(torch, 2, 1), handle + ) + + def test_selected_must_carry_one_row_per_batch_row(self): + """Per query token is what a kernel's selection looks like; this call is not that. + + A decode step gives a request one query token, so the two agree there. A prefill + chunk would give it many, and the resolve has no way to hold the union resident. + """ + roles = ("key",) + with self._fixture(roles, rows=2) as (torch, manager, batch, _cache, handle): + # Four query tokens for a two-row batch: what a prefill caller would pass. + selected = torch.full((4, self.SELECTED), -1, dtype=torch.int32, device="cuda") + out = self._out(torch, 4, len(roles)) + with self.assertRaises(ValueError): + batch.fetch_sparse_pages(self._buffers(roles), selected, out, handle) + + def test_out_must_carry_one_row_per_buffer(self): + roles = ("key", "value") + with self._fixture(roles) as (torch, manager, batch, _cache, handle): + selected = torch.tensor([[0, -1], [-1, -1]], dtype=torch.int32, device="cuda") + with self.assertRaises(ValueError): + batch.fetch_sparse_pages( + self._buffers(roles), selected, self._out(torch, 2, 1), handle + ) + + +class TestIndexTables(unittest.TestCase): + """Which converted index tables a Batch builds, and which specs share one. + + A spec states what a kernel must address together. Two specs share a table when + their converted entries coincide, which needs the same life cycle, the same page index + scale, and the same offsets -- the condition a quantised layer's data and block scales + satisfy because their pool group mirrors their slots. + """ + + def setUp(self) -> None: + init_cuda_once() + + @staticmethod + def _manager(*, roles: tuple[str, ...], layers: int = 2) -> KVCacheManager: + return KVCacheManager( + KVCacheManagerConfig( + tokens_per_block=32, + cache_tiers=[GpuCacheTierConfig(quota=8 << 20)], + layers=[ + AttentionLayerConfig( + layer_id=LayerId(layer), + buffers=[BufferConfig(role=role, size=1024) for role in roles], + ) + for layer in range(layers) + ], + constraints=[BatchDesc(kv_caches=[KVCacheDesc(capacity=32, history_length=0)])], + ) + ) + + @staticmethod + def _batch(manager, specs, max_batch_size=1): + return Batch( + manager, + NULL_STREAM, + max_batch_size=max_batch_size, + max_blocks_per_seq=4, + index_tables=specs, + ) + + def test_a_rows_entries_sit_together(self): + """A row holds its entries side by side, rather than one plane per entry. + + The shape an attention kernel takes a key and value page list in: XQA's + kvCachePageList is [batchSize][beamWidth][2][maxNbPagesPerSeq]. A single-row + batch cannot tell the two layouts apart, so this one has several. + """ + try: + import torch + except ImportError: # pragma: no cover - fast mode without torch + raise unittest.SkipTest("torch is needed to read the view") + manager = self._manager(roles=("key", "value")) + try: + spec = IndexTableSpec([BufferId(0, "key"), BufferId(0, "value")]) + batch = self._batch(manager, [spec], max_batch_size=3) + view = torch.from_dlpack(batch.converted_index_table(0)) + self.assertEqual(tuple(view.shape), (3, 2, 4)) + # One row spans every entry; one entry spans one row of blocks. + self.assertEqual(view.stride(), (2 * 4, 4, 1)) + finally: + manager.shutdown() + + def test_the_view_is_shaped_for_the_entries_the_table_serves(self): + """One entry per buffer of the spec, adjacent within each row.""" + try: + import torch + except ImportError: # pragma: no cover - fast mode without torch + raise unittest.SkipTest("torch is needed to read the view") + manager = self._manager(roles=("key", "value")) + try: + spec = IndexTableSpec([BufferId(0, "key"), BufferId(0, "value")]) + batch = self._batch(manager, [spec]) + view = torch.from_dlpack(batch.converted_index_table(0)) + # Row-major in the entries, so a row carries both buffers side by side: two + # buffers in the spec gives [max_batch, 2, max_blocks_per_seq]. + self.assertEqual(tuple(view.shape), (1, 2, 4)) + self.assertEqual(view.dtype, torch.int32) + self.assertTrue(view.is_cuda) + # The entry a buffer reads is its position in the spec that named it. + self.assertEqual(_introspection.batch_index_table_slot(batch, 0, "value")[1], 1) + with self.assertRaises(Exception): + batch.converted_index_table(batch.num_index_tables) + finally: + manager.shutdown() + + def test_a_spec_is_reachable_from_every_buffer_it_names(self): + """Entry order follows the spec's buffer order, which is what a kernel indexes.""" + manager = self._manager(roles=("key", "value")) + try: + spec = IndexTableSpec([BufferId(0, "key"), BufferId(0, "value")]) + batch = self._batch(manager, [spec]) + self.assertEqual(batch.num_index_tables, 1) + self.assertEqual(_introspection.batch_index_table_slot(batch, 0, "key"), (0, 0)) + self.assertEqual(_introspection.batch_index_table_slot(batch, 0, "value"), (0, 1)) + # One base for the whole spec: the offsets are already in the entries. + self.assertEqual( + _introspection.batch_index_table_base(batch, 0, "key"), + _introspection.batch_index_table_base(batch, 0, "value"), + ) + finally: + manager.shutdown() + + def test_layers_of_one_life_cycle_share_a_table_under_shared(self): + """SHARED is exactly this: entries that every layer of the group can read.""" + manager = self._manager(roles=("key",), layers=2) + try: + lg = manager.get_layer_group_id(0) + specs = [IndexTableSpec([BufferId(layer, "key")]) for layer in range(2)] + batch = self._batch(manager, specs) + self.assertEqual(batch.index_mode(lg), PageIndexMode.SHARED) + self.assertEqual(batch.num_index_tables, 1) + self.assertEqual( + _introspection.batch_index_table_slot(batch, 0, "key")[0], + _introspection.batch_index_table_slot(batch, 1, "key")[0], + ) + # Two specs merged onto it, one per layer. + self.assertEqual(len(batch.index_table(0)), 2) + finally: + manager.shutdown() + + def test_per_layer_keeps_the_layers_apart(self): + """However alike the entries look, a PER_LAYER table serves one layer.""" + manager = self._manager(roles=("key",), layers=2) + try: + specs = [IndexTableSpec([BufferId(layer, "key")]) for layer in range(2)] + batch = Batch( + manager, + NULL_STREAM, + max_batch_size=1, + max_blocks_per_seq=4, + mode=PageIndexMode.PER_LAYER, + index_tables=specs, + ) + self.assertEqual(batch.num_index_tables, 2) + self.assertNotEqual( + _introspection.batch_index_table_slot(batch, 0, "key")[0], + _introspection.batch_index_table_slot(batch, 1, "key")[0], + ) + finally: + manager.shutdown() + + @staticmethod + def _quantised_manager(*, paired_scales: bool) -> KVCacheManager: + """One layer whose data and block-scale buffers land in separate pools. + + Pools group buffers of equal size, and a pool's scale is how many buffers it + coalesces. So two same-size scale buffers give scale 2 like the data pair, while a + lone one gives scale 1. + """ + scale_roles = ["k_scale", "v_scale"] if paired_scales else ["k_scale"] + return KVCacheManager( + KVCacheManagerConfig( + tokens_per_block=32, + cache_tiers=[GpuCacheTierConfig(quota=8 << 20)], + layers=[ + AttentionLayerConfig( + layer_id=LayerId(0), + buffers=[BufferConfig(role=role, size=1024) for role in ("key", "value")] + + [BufferConfig(role=role, size=256) for role in scale_roles], + ) + ], + constraints=[BatchDesc(kv_caches=[KVCacheDesc(capacity=32, history_length=0)])], + ) + ) + + def test_data_and_block_scales_share_a_table_when_their_scales_match(self): + """The quantised case the merge rule exists for. + + Both pools coalesce two buffers and the pool group mirrors their slots, so the + converted entries coincide and one table serves both -- through different bases. + """ + manager = self._quantised_manager(paired_scales=True) + try: + self.assertEqual( + manager.get_page_index_scale(0, "key"), manager.get_page_index_scale(0, "k_scale") + ) + data = IndexTableSpec([BufferId(0, "key"), BufferId(0, "value")]) + scales = IndexTableSpec([BufferId(0, "k_scale"), BufferId(0, "v_scale")]) + batch = self._batch(manager, [data, scales]) + + self.assertEqual(batch.num_index_tables, 1) + self.assertEqual(len(batch.index_table(0)), 2) + # Same entries, so the same table at the same slots... + self.assertEqual( + _introspection.batch_index_table_slot(batch, 0, "key"), + _introspection.batch_index_table_slot(batch, 0, "k_scale"), + ) + self.assertEqual( + _introspection.batch_index_table_slot(batch, 0, "value"), + _introspection.batch_index_table_slot(batch, 0, "v_scale"), + ) + # ...reached through different bases, which is what keeps them distinct. + self.assertNotEqual( + _introspection.batch_index_table_base(batch, 0, "key"), + _introspection.batch_index_table_base(batch, 0, "k_scale"), + ) + finally: + manager.shutdown() + + def test_a_spec_names_its_own_table_without_going_through_a_buffer(self): + """The spec a caller asked for is the handle it gets its table back by. + + Mirrored specs resolve to one id, which is the same statement + ``index_table_slot`` makes per buffer -- asserted here from the spec side, since + that is the query a caller holding specs actually has. + """ + manager = self._quantised_manager(paired_scales=True) + try: + data = IndexTableSpec([BufferId(0, "key"), BufferId(0, "value")]) + scales = IndexTableSpec([BufferId(0, "k_scale"), BufferId(0, "v_scale")]) + batch = self._batch(manager, [data, scales]) + + table = batch.index_table_id(data) + self.assertEqual(table, batch.index_table_id(scales), "mirrored specs are one table") + # The spec's buffers sit at entries [0, n) in its own order, so the per-buffer + # query adds nothing a holder of the spec did not already know. + for entry, role in enumerate(("key", "value")): + self.assertEqual( + _introspection.batch_index_table_slot(batch, 0, role), (table, entry) + ) + for entry, role in enumerate(("k_scale", "v_scale")): + self.assertEqual( + _introspection.batch_index_table_slot(batch, 0, role), (table, entry) + ) + finally: + manager.shutdown() + + def test_a_buffer_and_its_spec_reach_the_same_table_memory(self): + """The two queries differ only in whether the caller still has to slice. + + A buffer's table is the one its spec names, so the pointer is the same and the + entry is the whole of the difference -- which is what lets a caller hold specs and + never ask about a buffer. + """ + try: + import torch + except ImportError: # pragma: no cover - fast mode without torch + raise unittest.SkipTest("torch is needed to read the view") + # Unpaired, so the scale keeps a table of its own: with one table every buffer + # would reach the right memory by having nowhere else to go. + manager = self._quantised_manager(paired_scales=False) + try: + data = IndexTableSpec([BufferId(0, "key"), BufferId(0, "value")]) + scale = IndexTableSpec([BufferId(0, "k_scale")]) + batch = self._batch(manager, [data, scale]) + self.assertEqual(batch.num_index_tables, 2, "the two specs must not have merged") + + for spec, roles in ((data, ("key", "value")), (scale, ("k_scale",))): + owned = torch.from_dlpack(batch.converted_index_table(batch.index_table_id(spec))) + other = next(s for s in (data, scale) if s is not spec) + disowned = torch.from_dlpack( + batch.converted_index_table(batch.index_table_id(other)) + ) + self.assertNotEqual(owned.data_ptr(), disowned.data_ptr()) + for entry, role in enumerate(roles): + table, slot = _introspection.batch_index_table_slot(batch, 0, role) + from_buffer = torch.from_dlpack(batch.converted_index_table(table)) + self.assertEqual(from_buffer.data_ptr(), owned.data_ptr(), role) + self.assertEqual(slot, entry, role) + # The entry addresses a plane of that same memory rather than another + # allocation, so a subset is a slice and never a table of its own. + self.assertEqual( + from_buffer[:, slot, :].data_ptr(), + owned.data_ptr() + slot * owned.stride(1) * owned.element_size(), + role, + ) + finally: + manager.shutdown() + + def test_merged_specs_keep_pool_specs_of_their_own(self): + """One table, two bases -- which is the whole mechanism of the merge. + + The data and its block scales share one set of converted entries and are told + apart only by what they add those entries to, so the pool spec is a property of + the spec and could not be answered for the table. + """ + manager = self._quantised_manager(paired_scales=True) + try: + data = IndexTableSpec([BufferId(0, "key"), BufferId(0, "value")]) + scales = IndexTableSpec([BufferId(0, "k_scale"), BufferId(0, "v_scale")]) + batch = self._batch(manager, [data, scales]) + self.assertEqual(batch.num_index_tables, 1, "the premise is that they merged") + + data_base, data_stride = batch.pool_spec_for(data) + scale_base, scale_stride = batch.pool_spec_for(scales) + self.assertNotEqual(data_base, scale_base, "one base would read one pool twice") + # Separate pools of unequal buffer size, so the pages they count differ too. + self.assertNotEqual(data_stride, scale_stride) + for spec, base in ((data, data_base), (scales, scale_base)): + for buffer in spec.buffers: + self.assertEqual( + _introspection.batch_index_table_base(batch, buffer.layer_id, buffer.role), + base, + "every buffer of a spec is reached through the spec's own base", + ) + finally: + manager.shutdown() + + def test_a_pool_spec_needs_a_spec_the_batch_was_built_with(self): + manager = self._quantised_manager(paired_scales=True) + try: + data = IndexTableSpec([BufferId(0, "key"), BufferId(0, "value")]) + scales = IndexTableSpec([BufferId(0, "k_scale"), BufferId(0, "v_scale")]) + batch = self._batch(manager, [data, scales]) + with self.assertRaises(IndexError): + batch.pool_spec_for(IndexTableSpec([BufferId(0, "key")])) + finally: + manager.shutdown() + + def test_a_spec_the_batch_was_not_built_with_has_no_table(self): + manager = self._quantised_manager(paired_scales=True) + try: + data = IndexTableSpec([BufferId(0, "key"), BufferId(0, "value")]) + scales = IndexTableSpec([BufferId(0, "k_scale"), BufferId(0, "v_scale")]) + batch = self._batch(manager, [data, scales]) + # A subset is not a spec: it names a slice of a table, not a table. + with self.assertRaises(IndexError): + batch.index_table_id(IndexTableSpec([BufferId(0, "key")])) + finally: + manager.shutdown() + + def test_specs_may_not_overlap(self): + """A buffer belongs to one spec; a subset is read by slicing, not by declaring it. + + Both orders are refused, so this is a property of the pair rather than of + whichever spec happened to be built first. + """ + manager = self._quantised_manager(paired_scales=True) + try: + both = IndexTableSpec([BufferId(0, "key"), BufferId(0, "value")]) + lone = IndexTableSpec([BufferId(0, "key")]) + for order in ([both, lone], [lone, both]): + with self.assertRaises(RuntimeError) as caught: + self._batch(manager, order) + self.assertIn("may not overlap", str(caught.exception)) + finally: + manager.shutdown() + + def test_a_spec_spanning_scales_is_refused(self): + """A spec whose buffers disagree on scale cannot be served by one table. + + Scale sets the entry stride, so buffers with different scales have different + converted indices for the same slot. + """ + manager = self._quantised_manager(paired_scales=False) + try: + data_scale = manager.get_page_index_scale(0, "key") + lone_scale = manager.get_page_index_scale(0, "k_scale") + self.assertNotEqual(data_scale, lone_scale, "config does not produce differing scales") + with self.assertRaises(Exception) as caught: + self._batch(manager, [IndexTableSpec([BufferId(0, "key"), BufferId(0, "k_scale")])]) + self.assertIn("scales", str(caught.exception)) + finally: + manager.shutdown() + + def test_a_spec_spanning_life_cycles_is_refused(self): + """Different life cycles number their slots separately, so the entries cannot agree.""" + manager = KVCacheManager( + KVCacheManagerConfig( + tokens_per_block=32, + cache_tiers=[GpuCacheTierConfig(quota=8 << 20)], + layers=[ + AttentionLayerConfig( + layer_id=LayerId(0), + buffers=[BufferConfig(role="key", size=1024)], + sliding_window_size=32, + num_sink_tokens=0, + ), + AttentionLayerConfig( + layer_id=LayerId(1), + buffers=[BufferConfig(role="key", size=1024)], + sliding_window_size=None, + ), + ], + constraints=[BatchDesc(kv_caches=[KVCacheDesc(capacity=32, history_length=0)])], + ) + ) + try: + self.assertNotEqual(manager.get_layer_group_id(0), manager.get_layer_group_id(1)) + with self.assertRaises(Exception) as caught: + self._batch(manager, [IndexTableSpec([BufferId(0, "key"), BufferId(1, "key")])]) + self.assertIn("life cycles", str(caught.exception)) + finally: + manager.shutdown() + + def _assert_converted_matches_reference( + self, manager, batch, cache, roles, mode, scratch=None, layer=0 + ): + """Rebuild the concat table from PageIndexConverter and compare entry for entry. + + The reference gives each buffer its own table and puts the per-role offset in that + table's base. A concat spec shares one base across its buffers, so it must carry + the offset in the entry instead: + + expected[i] = reference[i] + (role_base[i] - spec_base) / stride + + Under SHARED that difference is the K-to-V gap; under PER_LAYER both bases are the + pool group's, so it is zero and the reference's own layer_offset already separates + the roles. Either way the reference supplies the arithmetic, so this tests the + conversion rather than restating it. + """ + lg = manager.get_layer_group_id(0) + raw = [int(x) for x in cache.get_base_page_indices(lg)] + self.assertTrue(any(x != BAD_PAGE_INDEX for x in raw), "no blocks to convert") + + converted = _introspection.batch_index_table(batch, 0) + rows, blocks = batch.max_batch_size, batch.max_blocks_per_seq + self.assertEqual(len(converted), len(roles) * rows * blocks) + + for entry, role in enumerate(roles): + stride = manager.get_page_stride(layer, role) + shift = ( + manager.get_mem_pool_base_address(layer, role, mode) + - _introspection.batch_index_table_base(batch, layer, role) + ) // stride + reference = list( + manager.get_page_index_converter(layer, role)( + raw[:blocks], + mode, + scratch.range if scratch else None, + scratch.slot_ids if scratch else [], + ) + ) + want = [x if x == BAD_PAGE_INDEX else x + shift for x in reference] + # Row-major in the entries: a row carries every entry before the next row + # starts, which is what puts a spec's key and value planes side by side. + base = (0 * len(roles) + entry) * blocks + got = converted[base : base + blocks] + self.assertEqual(got, want, f"entry {entry} ({role})") + + @contextlib.contextmanager + def _published(self, manager, mode): + """A one-member batch whose {key, value} spec has been converted once.""" + spec = IndexTableSpec([BufferId(0, "key"), BufferId(0, "value")]) + cache = None + batch = None + try: + with TemporaryCudaStream([]) as stream: + handle = cast(CudaStream, stream.handle) + batch = Batch( + manager, + handle, + max_batch_size=1, + max_blocks_per_seq=4, + mode=mode, + index_tables=[spec], + ) + cache = manager.create_kv_cache() + self.assertTrue(cache.resume(handle)) + tokens = 2 * 32 + # Seeded before joining: the Batch only accepts a history its row already + # had the capacity for. + self.assertTrue(cache.resize(tokens, tokens)) + batch.add_kv_cache(cache) + batch.resize([tokens], [tokens]) + yield batch, cache + finally: + if cache is not None: + if batch.row_of(cache) is not None: + batch.remove_kv_cache(cache) + cache.close() + + def test_shared_entries_match_the_reference_converter(self): + manager = self._manager(roles=("key", "value"), layers=1) + try: + with self._published(manager, PageIndexMode.SHARED) as (batch, cache): + self._assert_converted_matches_reference( + manager, batch, cache, ("key", "value"), PageIndexMode.SHARED + ) + finally: + manager.shutdown() + + def test_per_layer_entries_match_the_reference_converter(self): + """A separate case because the base differs between the modes. + + The pool group's base here, carrying no per-layer offset, against the spec's + first buffer under SHARED, which carries one. Same formula, different entries. + """ + manager = self._manager(roles=("key", "value"), layers=1) + try: + with self._published(manager, PageIndexMode.PER_LAYER) as (batch, cache): + self._assert_converted_matches_reference( + manager, batch, cache, ("key", "value"), PageIndexMode.PER_LAYER + ) + finally: + manager.shutdown() + + def test_scratch_staging_covers_the_layer_taking_the_most_sub_pages(self): + """A life cycle whose layers put unequal buffer counts in one pool. + + The slots a request takes come from the life cycle's widest sub-pages-per-slot + ratio, so the staging has to be sized by that same ratio. Sizing it from whichever + buffer an index table happens to name first under-counts whenever that buffer's + layer is not the widest, and the request then holds more slots than its row has. + """ + window, tpb, prompt_len = 128, 32, 1024 + # Layer 0 puts four buffers in the pool, layers 1-3 put two. Same size, so one + # pool; same window, so one life cycle. + layers = [ + AttentionLayerConfig( + layer_id=LayerId(0), + buffers=[ + BufferConfig(role=DataRole(role), size=8192) + for role in ("key", "value", "key2", "value2") + ], + sliding_window_size=window, + num_sink_tokens=0, + ) + ] + [ + AttentionLayerConfig( + layer_id=LayerId(i), + buffers=[ + BufferConfig(role=DataRole("key"), size=8192), + BufferConfig(role=DataRole("value"), size=8192), + ], + sliding_window_size=window, + num_sink_tokens=0, + ) + for i in range(1, 4) + ] + cfg = KVCacheManagerConfig( + tokens_per_block=tpb, + cache_tiers=[GpuCacheTierConfig(quota=64 << 20)], + layers=layers, + swa_scratch_reuse=SwaScratchReuseConfig(max_rewind_len=0), + ) + manager = KVCacheManager(cfg) + cache = None + batch = None + try: + lg = manager.get_layer_group_id(0) + max_blocks = div_up(prompt_len, tpb) + + # The premise: the widest layer takes strictly more sub-pages per slot than a + # two-buffer one, so the two sizings genuinely differ. + widest = manager.get_page_index_converter(0, "key") + narrow = manager.get_page_index_converter(1, "key") + self.assertEqual(widest.scale, narrow.scale, "expected one pool for the group") + self.assertGreater( + widest.scratch_pages_per_block, + narrow.scratch_pages_per_block, + "config is uniform, so it cannot distinguish the two sizings", + ) + + # A table naming only a narrow layer's buffers must still stage the widest + # layer's slots, because the request takes them for the whole life cycle. + spec = IndexTableSpec([BufferId(1, "key"), BufferId(1, "value")]) + with TemporaryCudaStream([]) as stream: + handle = cast(CudaStream, stream.handle) + batch = Batch( + manager, + handle, + index_tables=[spec], + max_batch_size=1, + max_blocks_per_seq=max_blocks, + mode=PageIndexMode.PER_LAYER, + ) + cache = manager.create_kv_cache(None, [TokenId(i) for i in range(prompt_len)]) + self.assertTrue(cache.resume(handle)) + batch.add_kv_cache(cache) + # Taking the slots is what overruns a row sized from the narrow buffer. + self.assertTrue(cache.resize(prompt_len)) + self.assertTrue(cache.has_scratch_slots, "config produced no scratch") + + slots = cache.get_scratch_slots(lg) + self.assertLessEqual( + len(slots), + manager.max_scratch_slots(max_blocks), + "the request holds more scratch slots than the manager sizes for", + ) + batch.wait_published() + batch.publish() + batch.wait_published() + finally: + if batch is not None and cache is not None and batch.row_of(cache) is not None: + batch.remove_kv_cache(cache) + if cache is not None: + cache.close() + manager.shutdown() + + def test_scratch_blocks_match_the_reference_converter(self): + """The substituting path, against PageIndexConverter given the same scratch range. + + A scratch block's slot comes from the spec rather than the page table, and the + layer's offset wraps inside the shared slot instead of being added to it. Both are + easy to transcribe wrongly, and neither shows up on a config without scratch, so + this drives a real scratch range and defers the arithmetic to the reference. + """ + num_layers, window, tpb, prompt_len = 4, 128, 32, 1024 + cfg = KVCacheManagerConfig( + tokens_per_block=tpb, + cache_tiers=[GpuCacheTierConfig(quota=64 << 20)], + layers=[ + AttentionLayerConfig( + layer_id=LayerId(i), + buffers=[ + BufferConfig(role=DataRole("key"), size=8192), + BufferConfig(role=DataRole("value"), size=8192), + ], + sliding_window_size=window, + num_sink_tokens=0, + ) + for i in range(num_layers) + ], + swa_scratch_reuse=SwaScratchReuseConfig(max_rewind_len=0), + ) + manager = KVCacheManager(cfg) + cache = None + batch = None + try: + lg = manager.get_layer_group_id(0) + # Scratch makes the layers stop sharing one table, so PER_LAYER is required. + self.assertIsNone(manager.supports_index_mode(PageIndexMode.SHARED, lg)) + + # The last layer, because its sub-page offsets are the largest: the wrap only + # bites when total % scale + offset reaches scale, which layer zero's offsets + # of {0, 1} never do. + layer = num_layers - 1 + converter = manager.get_page_index_converter(layer, "value") + scale, spb = converter.scale, converter.scratch_pages_per_block + self.assertTrue( + any((block * spb) % scale + converter.layer_offset >= scale for block in range(32)), + "config never wraps the layer offset, so the scratch path is only half tested", + ) + + spec = IndexTableSpec([BufferId(layer, "key"), BufferId(layer, "value")]) + with TemporaryCudaStream([]) as stream: + handle = cast(CudaStream, stream.handle) + batch = Batch( + manager, + handle, + max_batch_size=1, + max_blocks_per_seq=div_up(prompt_len, tpb), + mode=PageIndexMode.PER_LAYER, + index_tables=[spec], + ) + cache = manager.create_kv_cache(None, [TokenId(i) for i in range(prompt_len)]) + self.assertTrue(cache.resume(handle)) + batch.add_kv_cache(cache) + self.assertTrue(cache.resize(prompt_len)) + self.assertTrue(cache.has_scratch_slots, "config produced no scratch") + + scratch = _scratch_desc(cache, lg) + self.assertIsNotNone(scratch) + self.assertGreater(scratch.range.end - scratch.range.beg, 0) + + batch.wait_published() + batch.publish() + self._assert_converted_matches_reference( + manager, + batch, + cache, + ("key", "value"), + PageIndexMode.PER_LAYER, + scratch=scratch, + layer=layer, + ) + finally: + if cache is not None: + if batch.row_of(cache) is not None: + batch.remove_kv_cache(cache) + cache.close() + manager.shutdown() + + def test_a_named_shared_serves_a_member_that_keeps_the_claim(self): + """The case the escape hatch exists for: manager undecided, member clear. + + The manager answers for every request this sliding life cycle could carry, so it + will not promise SHARED. A member whose capacity never reaches the window holds + nothing stale, so it takes no scratch and says so itself -- and the batch that + named SHARED publishes it against one table for all four layers. + """ + num_layers, window, tpb = 4, 128, 32 + manager = KVCacheManager( + KVCacheManagerConfig( + tokens_per_block=tpb, + cache_tiers=[GpuCacheTierConfig(quota=64 << 20)], + layers=[ + AttentionLayerConfig( + layer_id=LayerId(i), + buffers=[BufferConfig(role=DataRole("key"), size=8192)], + sliding_window_size=window, + num_sink_tokens=0, + ) + for i in range(num_layers) + ], + swa_scratch_reuse=SwaScratchReuseConfig(max_rewind_len=0), + ) + ) + cache = None + batch = None + try: + lg = manager.get_layer_group_id(0) + # The manager will not promise it... + self.assertIsNone(manager.supports_index_mode(PageIndexMode.SHARED, lg)) + with TemporaryCudaStream([]) as stream: + handle = cast(CudaStream, stream.handle) + batch = Batch( + manager, + handle, + max_batch_size=1, + max_blocks_per_seq=4, + mode=PageIndexMode.SHARED, + index_tables=[IndexTableSpec([BufferId(0, "key")])], + ) + # One table for the life cycle, not one per layer: the saving being claimed. + self.assertEqual(batch.num_index_tables, 1) + + tokens = tpb # One block, far short of the window, so nothing goes stale. + cache = manager.create_kv_cache(None, [TokenId(i) for i in range(tokens)]) + self.assertTrue(cache.resume(handle)) + self.assertTrue(cache.resize(tokens)) + # ...but the member itself does, which is what makes the claim true. + self.assertFalse(cache.has_scratch_slots) + self.assertTrue(cache.supports_index_mode(PageIndexMode.SHARED, lg)) + + batch.add_kv_cache(cache) + batch.publish() + batch.wait_published() + # Published, and every layer still reads the one table. + self.assertEqual(batch.index_mode(lg), PageIndexMode.SHARED) + self.assertEqual(batch.num_index_tables, 1) + finally: + if cache is not None: + if batch is not None and batch.row_of(cache) is not None: + batch.remove_kv_cache(cache) + cache.close() + manager.shutdown() + + def test_a_named_shared_is_held_to_its_claim_at_publish(self): + """Naming SHARED asserts no member holds scratch; publish is what enforces it. + + Construction lets the claim stand because configuration cannot settle it -- the + manager answers for every request the life cycle could carry. So the batch admits + a member that does hold scratch, and the publish that would hand its layers one + table for indices that differ between them refuses instead. + """ + num_layers, window, tpb, prompt_len = 4, 128, 32, 1024 + manager = KVCacheManager( + KVCacheManagerConfig( + tokens_per_block=tpb, + cache_tiers=[GpuCacheTierConfig(quota=64 << 20)], + layers=[ + AttentionLayerConfig( + layer_id=LayerId(i), + buffers=[BufferConfig(role=DataRole("key"), size=8192)], + sliding_window_size=window, + num_sink_tokens=0, + ) + for i in range(num_layers) + ], + swa_scratch_reuse=SwaScratchReuseConfig(max_rewind_len=0), + ) + ) + cache = None + batch = None + try: + lg = manager.get_layer_group_id(0) + self.assertIsNone(manager.supports_index_mode(PageIndexMode.SHARED, lg)) + with TemporaryCudaStream([]) as stream: + handle = cast(CudaStream, stream.handle) + # Named, so construction allows what the manager would not derive. + batch = Batch( + manager, + handle, + max_batch_size=1, + max_blocks_per_seq=div_up(prompt_len, tpb), + mode=PageIndexMode.SHARED, + index_tables=[IndexTableSpec([BufferId(0, "key")])], + ) + self.assertEqual(batch.index_mode(lg), PageIndexMode.SHARED) + + # A prefill-sized grow is what takes scratch slots, so this member breaks + # the claim the batch was built on. + cache = manager.create_kv_cache(None, [TokenId(i) for i in range(prompt_len)]) + self.assertTrue(cache.resume(handle)) + self.assertTrue(cache.resize(prompt_len)) + self.assertTrue(cache.has_scratch_slots, "config produced no scratch") + + # Seating is outside the assertion: the claim is that the batch takes the + # member and publish() is what refuses it, which a block covering both + # would not tell apart. + batch.add_kv_cache(cache) + with self.assertRaises(RuntimeError) as caught: + batch.publish() + self.assertIn("SHARED life cycle", str(caught.exception)) + finally: + if cache is not None: + if batch is not None and batch.row_of(cache) is not None: + batch.remove_kv_cache(cache) + cache.close() + manager.shutdown() + + def test_an_unrequested_buffer_has_no_slot(self): + manager = self._manager(roles=("key", "value")) + try: + batch = self._batch(manager, [IndexTableSpec([BufferId(0, "key")])]) + with self.assertRaises(Exception) as caught: + _introspection.batch_index_table_slot(batch, 0, "value") + self.assertIn("no index table spec", str(caught.exception)) + finally: + manager.shutdown() + + def test_no_specs_means_no_tables(self): + manager = self._manager(roles=("key",)) + try: + self.assertEqual(self._batch(manager, []).num_index_tables, 0) + finally: + manager.shutdown() + + +class TestIndexMode(unittest.TestCase): + """Which page index mode a life cycle gets, and why it is a per-life-cycle question. + + Only scratch-slot substitution can make one layer's converted indices differ from + another's in the same life cycle. A life cycle with no sliding window yields an empty + scratch range whatever the watermarks, so no request can hold scratch slots there and + SHARED is always available -- even with scratch reuse enabled for the manager. + """ + + #: Layer 0 slides and so may hold scratch; layer 1 does not and so never can. + WINDOWED_LAYER = 0 + PLAIN_LAYER = 1 + + def setUp(self) -> None: + init_cuda_once() + + @staticmethod + def _manager(*, scratch_reuse: bool) -> KVCacheManager: + return KVCacheManager( + KVCacheManagerConfig( + tokens_per_block=32, + cache_tiers=[GpuCacheTierConfig(quota=4 << 20)], + layers=[ + AttentionLayerConfig( + layer_id=LayerId(0), + buffers=[BufferConfig(role=Role.KEY, size=1024)], + sliding_window_size=32, + num_sink_tokens=0, + ), + AttentionLayerConfig( + layer_id=LayerId(1), + buffers=[BufferConfig(role=Role.KEY, size=1024)], + sliding_window_size=None, + ), + ], + constraints=[BatchDesc(kv_caches=[KVCacheDesc(capacity=32, history_length=0)])], + swa_scratch_reuse=(SwaScratchReuseConfig() if scratch_reuse else None), + ) + ) + + def test_a_life_cycle_without_a_window_keeps_shared_under_scratch_reuse(self): + """The collapsing query cannot say this: it answers for the manager, not a group.""" + manager = self._manager(scratch_reuse=True) + try: + windowed = manager.get_layer_group_id(self.WINDOWED_LAYER) + plain = manager.get_layer_group_id(self.PLAIN_LAYER) + self.assertNotEqual(windowed, plain) + + # Per life cycle: only the sliding one is in doubt. + self.assertIsNone(manager.supports_index_mode(PageIndexMode.SHARED, windowed)) + self.assertTrue(manager.supports_index_mode(PageIndexMode.SHARED, plain)) + # PER_LAYER never depends on anything. + self.assertTrue(manager.supports_index_mode(PageIndexMode.PER_LAYER, windowed)) + finally: + manager.shutdown() + + def test_without_scratch_reuse_every_life_cycle_is_shared(self): + manager = self._manager(scratch_reuse=False) + try: + for layer in (self.WINDOWED_LAYER, self.PLAIN_LAYER): + lg = manager.get_layer_group_id(layer) + self.assertTrue(manager.supports_index_mode(PageIndexMode.SHARED, lg)) + finally: + manager.shutdown() + + def test_a_batch_takes_the_smaller_table_wherever_it_can(self): + """Per life cycle, not per batch: the sliding group alone pays for PER_LAYER.""" + manager = self._manager(scratch_reuse=True) + try: + batch = Batch( + manager, NULL_STREAM, index_tables=[], max_batch_size=1, max_blocks_per_seq=4 + ) + self.assertEqual( + batch.index_mode(manager.get_layer_group_id(self.WINDOWED_LAYER)), + PageIndexMode.PER_LAYER, + ) + self.assertEqual( + batch.index_mode(manager.get_layer_group_id(self.PLAIN_LAYER)), + PageIndexMode.SHARED, + ) + finally: + manager.shutdown() + + def test_an_explicit_mode_overrides_every_life_cycle(self): + manager = self._manager(scratch_reuse=False) + try: + batch = Batch( + manager, + NULL_STREAM, + index_tables=[], + max_batch_size=1, + max_blocks_per_seq=4, + mode=PageIndexMode.PER_LAYER, + ) + for layer in (self.WINDOWED_LAYER, self.PLAIN_LAYER): + lg = manager.get_layer_group_id(layer) + # SHARED was available and still not taken. + self.assertTrue(manager.supports_index_mode(PageIndexMode.SHARED, lg)) + self.assertEqual(batch.index_mode(lg), PageIndexMode.PER_LAYER) + finally: + manager.shutdown() + + def test_naming_shared_takes_it_where_the_manager_will_not_promise_it(self): + """A named SHARED is a claim about this batch's members, not about requests. + + The manager answers for every request a life cycle could carry, so a sliding one + stays undecided. A batch that will only ever hold decode members knows better: + scratch covers the blocks a prefill chunk adds, and a one-token request adds none. + """ + manager = self._manager(scratch_reuse=True) + try: + windowed = manager.get_layer_group_id(self.WINDOWED_LAYER) + self.assertIsNone(manager.supports_index_mode(PageIndexMode.SHARED, windowed)) + + spec = IndexTableSpec([BufferId(self.WINDOWED_LAYER, Role.KEY)]) + batch = Batch( + manager, + NULL_STREAM, + max_batch_size=1, + max_blocks_per_seq=4, + mode=PageIndexMode.SHARED, + index_tables=[spec], + ) + self.assertEqual(batch.index_mode(windowed), PageIndexMode.SHARED) + # One table for the life cycle, which is the whole point of taking SHARED. + self.assertEqual(batch.num_index_tables, 1) + finally: + manager.shutdown() + + def test_a_derived_shared_is_still_refused_where_scratch_is_possible(self): + """Only a named mode may outrun the manager: nothing would check a derived one.""" + manager = self._manager(scratch_reuse=True) + try: + batch = Batch( + manager, + NULL_STREAM, + max_batch_size=1, + max_blocks_per_seq=4, + index_tables=[IndexTableSpec([BufferId(self.WINDOWED_LAYER, Role.KEY)])], + ) + # Derived, so the sliding life cycle took the table it can prove correct. + self.assertEqual( + batch.index_mode(manager.get_layer_group_id(self.WINDOWED_LAYER)), + PageIndexMode.PER_LAYER, + ) + finally: + manager.shutdown() + + +class TestBatchPageTable(unittest.TestCase): + """The Batch's own host page-index staging: who writes it, and what publish does with it. + + A member writes its indices into the Batch's row and nowhere else, which is what keeps + one buffer with one writer and a binding that survives the whole decode. A consumer + needing the same numbers reads coalesced_page_table() rather than keeping a buffer of its + own. + """ + + ROWS = 2 + BLOCKS = 4 + GROUPS = 1 + + def setUp(self) -> None: + init_cuda_once() + + def _torch(self): + try: + import torch + except ImportError: # pragma: no cover - fast mode without torch + raise unittest.SkipTest("torch is needed to read the staging as an array") + return torch + + @contextlib.contextmanager + def _fixture(self): + """A manager and a Batch staging its own page indices.""" + torch = self._torch() + manager = TestSparseFetch._manager(("key",)) + try: + # The stream belongs to the fixture because the Batch takes it at construction: + # a member joins the stream it was built for, so one cannot outlive the other. + with TemporaryCudaStream([]) as stream: + handle = cast(CudaStream, stream.handle) + batch = Batch( + manager, + handle, + max_batch_size=self.ROWS, + max_blocks_per_seq=self.BLOCKS, + index_tables=manager.sparse_index_table_specs(), + ) + yield torch, manager, batch, handle + finally: + manager.shutdown() + + def _row(self, torch, batch, manager, row: int, group: int = 0): + lg = manager.get_layer_group_id(group) + host = torch.from_dlpack(batch.coalesced_page_table(lg)) + return host[row].tolist() + + @contextlib.contextmanager + def _member(self, manager, batch, row: int, handle): + """A resumed cache seated in `row`, with two whole blocks of history.""" + cache = manager.create_kv_cache() + try: + self.assertTrue(cache.resume(handle)) + tokens = 2 * TestSparseFetch.TOKENS_PER_BLOCK + # Seeded before joining: the Batch only accepts a history its row already had + # the capacity for. + self.assertTrue(cache.resize(tokens, tokens)) + batch.add_kv_cache(cache, row) + sizes = [None] * self.ROWS + sizes[row] = tokens + batch.resize(list(sizes), list(sizes)) + yield cache + finally: + if batch.row_of(cache) is not None: + batch.remove_kv_cache(cache) + cache.close() + + def test_a_member_writes_its_indices_into_its_own_row(self): + """The row it was seated in, and no other.""" + with self._fixture() as (torch, manager, batch, handle): + lg = manager.get_layer_group_id(0) + with self._member(manager, batch, 1, handle) as cache: + expected = [int(x) for x in cache.get_base_page_indices(lg)][: self.BLOCKS] + self.assertEqual(self._row(torch, batch, manager, 1), expected) + # A page index, not the fill: the member really has pages here. + self.assertNotEqual(expected[0], BAD_PAGE_INDEX) + # The row nobody occupies still reads as absent. + self.assertEqual( + self._row(torch, batch, manager, 0), [BAD_PAGE_INDEX] * self.BLOCKS + ) + + def test_the_host_view_is_what_publish_uploads(self): + """The device table is a snapshot of the host rows, row for row. + + It is copied rather than read in place because publish() is a snapshot: the host + has to be free to write the next step's rows long before the kernels reading this + step's have retired. + """ + with self._fixture() as (torch, manager, batch, handle): + lg = manager.get_layer_group_id(0) + with self._member(manager, batch, 1, handle): + # _member's resize() already published; republishing needs the wait. + batch.wait_published() + batch.publish() + torch.cuda.synchronize() + table = _introspection.batch_page_table(batch, lg) + self.assertEqual(len(table), self.ROWS * self.BLOCKS) + expected = self._row(torch, batch, manager, 1) + self.assertEqual(table[self.BLOCKS : 2 * self.BLOCKS], expected) + self.assertNotEqual(expected[0], BAD_PAGE_INDEX) + self.assertEqual(table[: self.BLOCKS], [BAD_PAGE_INDEX] * self.BLOCKS) + + def test_publishing_over_a_publish_in_flight_is_refused(self): + """A wait here would be too late, so it is the caller's to make. + + Everything that writes the staged rows has to be separated from the copy reading + them: the members' own mutations, and the watermarks and views publish() stages + itself. Waiting inside publish() would cover only the second, and would do it while + hiding that a member mutated before the call had already raced. + """ + with self._fixture() as (torch, manager, batch, handle): + with self._member(manager, batch, 0, handle): + # _member's resize() published and nobody has waited on it. + with self.assertRaises(Exception) as caught: + batch.publish() + self.assertIn("still reading the staged rows", str(caught.exception)) + + batch.wait_published() + batch.publish() + + def test_a_member_refuses_a_second_page_index_buffer(self): + """Its indices live in the Batch's row, so a second buffer is refused. + + That buffer would take the writes with it and strand the row at whatever it held + when the cache joined -- silently, which is the one outcome worth a hard failure. + """ + with self._fixture() as (torch, manager, batch, handle): + lg = manager.get_layer_group_id(0) + with self._member(manager, batch, 0, handle) as cache: + self.assertEqual(cache.batch_row, 0) + spare = array.array("i", [BAD_PAGE_INDEX] * self.BLOCKS) + with self.assertRaises(Exception) as caught: + cache.set_base_page_index_buf(DEFAULT_BEAM_INDEX, lg, memoryview(spare)) + self.assertIn("Batch member", str(caught.exception)) + + def test_a_seated_member_cannot_suspend(self): + """Unlocking pages writes through the binding, which is the Batch's row. + + Doing that while seated would write a row a publish may still be reading, so the + member leaves first and the row is cleared by the removal. + """ + with self._fixture() as (torch, manager, batch, handle): + with self._member(manager, batch, 1, handle) as cache: + with self.assertRaisesRegex(Exception, "remove it from the Batch first"): + cache.suspend() + batch.remove_kv_cache(cache) + self.assertEqual(cache.batch_row, None) + self.assertEqual( + self._row(torch, batch, manager, 1), [BAD_PAGE_INDEX] * self.BLOCKS + ) + cache.suspend() + + +class TestSelectionId(unittest.TestCase): + """What a selection_id may name, which decides when two layers share a selection.""" + + TOKENS_PER_BLOCK = 128 + + @classmethod + def _config(cls, ids: list[int | None], **overrides) -> KVCacheManagerConfig: + """One sparse layer per entry of `ids`, each naming that selection or none.""" + + def cache(layer: int, selection_id: int | None) -> SparseCacheConfig: + return SparseCacheConfig( + num_selected_blocks=overrides.get("num_selected_blocks", {}).get(layer, 2), + max_cached_units=overrides.get("max_cached_units", {}).get(layer, 4), + unit_tokens=overrides.get("unit_tokens", {}).get(layer, cls.TOKENS_PER_BLOCK), + selection_id=selection_id, + ) + + return KVCacheManagerConfig( + tokens_per_block=cls.TOKENS_PER_BLOCK, + cache_tiers=[GpuCacheTierConfig(64 << 20)], + constraints=[BatchDesc(kv_caches=[KVCacheDesc(capacity=4096, history_length=0)] * 2)], + layers=[ + AttentionLayerConfig( + LayerId(layer), + [ + BufferConfig("key", 4096, is_sparse=True), + BufferConfig("index", 256), + ], + sparse_cache=cache(layer, selection_id), + ) + for layer, selection_id in enumerate(ids) + ], + ) + + def test_an_unset_id_defaults_to_the_layers_own(self): + self.assertIsNone( + SparseCacheConfig( + num_selected_blocks=2, max_cached_units=4, unit_tokens=self.TOKENS_PER_BLOCK + ).selection_id + ) + named = SparseCacheConfig( + num_selected_blocks=2, + max_cached_units=4, + unit_tokens=self.TOKENS_PER_BLOCK, + selection_id=7, + ) + self.assertEqual(named.selection_id, 7) + self.assertEqual(deepcopy(named).selection_id, 7) + + def test_layers_may_share_one_selection(self): + self._config([3, 3, None]) + + def test_silence_is_never_shared(self): + """Layer 1 names nothing, so it reads as selection 1; layer 0 may not take it.""" + with self.assertRaises(AssertionError) as caught: + self._config([1, None]) + self.assertIn("names a layer that set none of its own", str(caught.exception)) + + def test_sharers_must_agree_on_what_they_select(self): + with self.assertRaises(AssertionError) as caught: + self._config([5, 5], max_cached_units={1: 8}) + self.assertIn("must agree on", str(caught.exception)) + + def test_a_layer_may_name_its_own_id(self): + """Naming yourself is not sharing, so it collides with nothing.""" + self._config([0, None]) + + +class TestSparseSelectionSharing(unittest.TestCase): + """What two layers given one selection resolve to. + + Which decides whether a shared selection may share one resident set. + + A layer's resolve maintains that layer's own residency and recency. When several + layers of a life cycle are handed the same selection -- which is what a model with + one indexer across several layers does -- the question is whether those per-layer + answers differ. They do not, and this pins that: same starting state, same input, + same evolution. It is the premise a shared selection group rests on, so it is + checked here rather than assumed. + """ + + TOKENS_PER_BLOCK = 128 + SELECTED = 2 + CACHED_UNITS = 4 + SPARSE_BYTES = 4096 + LAYERS = 2 + + def setUp(self) -> None: + init_cuda_once() + + @classmethod + def _manager(cls, selection_id: int | None = None) -> KVCacheManager: + """Two sparse layers alike in every way, so one life cycle holds both. + + With `selection_id` they also name one selection, which is the whole + difference between resolving twice and resolving once. + """ + sparse_cache = SparseCacheConfig( + num_selected_blocks=cls.SELECTED, + max_cached_units=cls.CACHED_UNITS, + unit_tokens=cls.TOKENS_PER_BLOCK, + selection_id=selection_id, + ) + return KVCacheManager( + KVCacheManagerConfig( + tokens_per_block=cls.TOKENS_PER_BLOCK, + cache_tiers=[GpuCacheTierConfig(64 << 20), HostCacheTierConfig(64 << 20)], + constraints=[ + BatchDesc(kv_caches=[KVCacheDesc(capacity=4096, history_length=0)] * 4) + ], + layers=[ + AttentionLayerConfig( + LayerId(layer), + [ + BufferConfig("key", cls.SPARSE_BYTES, is_sparse=True), + BufferConfig("index", 256), + ], + sparse_cache=sparse_cache, + ) + for layer in range(cls.LAYERS) + ], + ) + ) + + def test_one_selection_resolves_alike_in_every_layer(self): + try: + import torch + except ImportError: # pragma: no cover - fast mode without torch + raise unittest.SkipTest("torch is needed to pass device arrays") + + manager = self._manager() + cache = None + batch = None + try: + with TemporaryCudaStream([]) as stream: + handle = cast(CudaStream, stream.handle) + lc_id = manager.get_layer_group_id(LayerId(0)) + self.assertEqual(lc_id, manager.get_layer_group_id(LayerId(1))) + + batch = Batch( + manager, + handle, + max_batch_size=2, + max_blocks_per_seq=4, + index_tables=manager.sparse_index_table_specs(), + ) + cache = manager.create_kv_cache() + self.assertTrue(cache.resume(handle)) + tokens = 2 * self.TOKENS_PER_BLOCK + self.assertTrue(cache.resize(tokens, tokens)) + batch.add_kv_cache(cache) + self.assertEqual(batch.resize([tokens, None], [tokens, None]), [True, False]) + + selected = torch.tensor([[0, 1], [-1, -1]], dtype=torch.int32, device="cuda") + for layer in range(self.LAYERS): + out = torch.full((2, 1, self.SELECTED), 12345, dtype=torch.int32, device="cuda") + batch.fetch_sparse_pages([BufferId(layer, "key")], selected, out, handle) + torch.cuda.synchronize() + + ordinals = _introspection.sparse_cache_resident_ordinals(cache, lc_id) + entries = len(ordinals) // self.LAYERS + self.assertEqual(len(ordinals), entries * self.LAYERS) + first = ordinals[:entries] + # Something was admitted, or the comparison below is vacuous. + self.assertNotEqual(first, [BAD_PAGE_INDEX] * entries, "no unit became resident") + for layer in range(1, self.LAYERS): + self.assertEqual( + ordinals[layer * entries : (layer + 1) * entries], + first, + f"layer {layer} resolved the same selection differently", + ) + finally: + if cache is not None: + if batch is not None and batch.row_of(cache) is not None: + batch.remove_kv_cache(cache) + cache.close() + del batch + manager.shutdown() + + def test_layers_naming_one_selection_keep_one_resident_set(self): + """One resolve, one metadata row, and each layer's own mirror of its entries.""" + try: + import torch + except ImportError: # pragma: no cover - fast mode without torch + raise unittest.SkipTest("torch is needed to pass device arrays") + + manager = self._manager(selection_id=0) + cache = None + batch = None + try: + with TemporaryCudaStream([]) as stream: + handle = cast(CudaStream, stream.handle) + lc_id = manager.get_layer_group_id(LayerId(0)) + + batch = Batch( + manager, + handle, + max_batch_size=2, + max_blocks_per_seq=4, + index_tables=manager.sparse_index_table_specs(), + ) + cache = manager.create_kv_cache() + self.assertTrue(cache.resume(handle)) + tokens = 2 * self.TOKENS_PER_BLOCK + self.assertTrue(cache.resize(tokens, tokens)) + batch.add_kv_cache(cache) + self.assertEqual(batch.resize([tokens, None], [tokens, None]), [True, False]) + + layers = [[BufferId(layer, "key")] for layer in range(self.LAYERS)] + selected = torch.tensor([[0, 1], [-1, -1]], dtype=torch.int32, device="cuda") + + # The fused call cannot serve a shared selection, and says so rather + # than resolving it a second time. + with self.assertRaises(Exception) as caught: + out = torch.empty((2, 1, self.SELECTED), dtype=torch.int32, device="cuda") + batch.fetch_sparse_pages(layers[0], selected, out, handle) + self.assertIn("shares its selection_id", str(caught.exception)) + + # Resolved once for both layers, then copied a layer at a time. + batch.resolve_sparse_selection(layers[0], selected, handle) + outs = [] + for buffers in layers: + out = torch.full((2, 1, self.SELECTED), 12345, dtype=torch.int32, device="cuda") + batch.fetch_resolved_pages(buffers, selected, out, handle) + outs.append(out) + torch.cuda.synchronize() + + # One metadata row for the two layers, where selecting apart gave two. + entries = _introspection.batch_sparse_entries_per_row(batch, lc_id)[0] + ordinals = _introspection.sparse_cache_resident_ordinals(cache, lc_id) + self.assertEqual(len(ordinals), entries) + self.assertNotEqual(ordinals, [BAD_PAGE_INDEX] * entries, "no unit became resident") + + # The entry each unit got is the selection's, so the two layers differ + # only by where their cells sit: one buffer's worth of them apart. + cells_per_buffer = entries # unit_tokens == tokens_per_block + first = outs[0].tolist() + for layer in range(1, self.LAYERS): + shifted = [ + [ + [ + index - layer * cells_per_buffer + if index != BAD_PAGE_INDEX + else index + for index in slots + ] + for slots in row + ] + for row in outs[layer].tolist() + ] + self.assertEqual(shifted, first, f"layer {layer} landed on a different entry") + finally: + if cache is not None: + if batch is not None and batch.row_of(cache) is not None: + batch.remove_kv_cache(cache) + cache.close() + del batch + manager.shutdown() + + +class TestSparseDataPath(unittest.TestCase): + """The bytes themselves, end to end, against real hot pages. + + Every other sparse test asserts residency, index shape or counts. Those pass whatever + base address the copy uses, so they cannot see a page side pointed at the wrong + allocation -- which is exactly what moving history to the hot tier changed. These two + read and write the actual memory on both sides of the copy. + """ + + TPB = 128 + BUF = 4096 + SELECTED = 2 + CELLS = 4 + + def setUp(self) -> None: + init_cuda_once() + + @classmethod + def _manager(cls) -> KVCacheManager: + return KVCacheManager( + KVCacheManagerConfig( + tokens_per_block=cls.TPB, + cache_tiers=[GpuCacheTierConfig(64 << 20), HostCacheTierConfig(64 << 20)], + constraints=[ + BatchDesc(kv_caches=[KVCacheDesc(capacity=4096, history_length=0)] * 2) + ], + layers=[ + AttentionLayerConfig( + 0, + [ + BufferConfig( + "key", + cls.BUF, + is_sparse=True, + ), + # A pool group is sized from its device-backed pools, so a sparse layer + # carries one. This is also the shape a model uses: an indexer and the + # data it selects over are one attention computation. + BufferConfig("index", 256), + ], + sparse_cache=SparseCacheConfig( + num_selected_blocks=cls.SELECTED, + max_cached_units=cls.CELLS, + unit_tokens=cls.TPB, + ), + ) + ], + ) + ) + + def _deps(self): + try: + import ctypes + + import cuda.bindings.runtime as cudart + import torch + except ImportError: # pragma: no cover - fast mode without torch + raise unittest.SkipTest("torch and the CUDA bindings are needed to address raw memory") + return ctypes, cudart, torch + + @staticmethod + def _pattern(size: int, seed: int) -> bytes: + return bytes((i * seed + 13) % 251 for i in range(size)) + + def test_a_fetched_cell_holds_the_bytes_of_its_hot_page(self): + ctypes, cudart, torch = self._deps() + manager = self._manager() + cache = None + try: + lg = manager.get_layer_group_id(0) + pool = int(manager.get_mem_pool_base_address(LayerId(0), "key")) + stride = manager.get_page_stride(LayerId(0), "key") + cache = manager.create_kv_cache() + with TemporaryCudaStream([]) as stream: + handle = cast(CudaStream, stream.handle) + batch = Batch( + manager, + handle, + max_batch_size=1, + max_blocks_per_seq=8, + index_tables=manager.sparse_index_table_specs(), + ) + arena = int(batch.get_sparse_cache_base_address(LayerId(0), "key")) + self.assertTrue(cache.resume(handle)) + batch.add_kv_cache(cache) + self.assertTrue(cache.resize(4 * self.TPB, 4 * self.TPB)) + + # Paint one block's hot page. It is host memory, so the host writes it + # directly; the kernel reaches the same bytes through the device mapping. + ordinal = 1 + pattern = self._pattern(self.BUF, 7) + pages = [int(x) for x in cache.get_base_page_indices(lg)] + ctypes.memmove(pool + pages[ordinal] * stride, pattern, self.BUF) + + batch.publish() + selected = torch.tensor([[ordinal, -1]], dtype=torch.int32, device="cuda") + out = torch.full((1, 1, self.SELECTED), 12345, dtype=torch.int32, device="cuda") + batch.fetch_sparse_pages([BufferId(0, "key")], selected, out, handle) + torch.cuda.synchronize() + + index = int(out[0][0][0]) + self.assertNotEqual(index, BAD_PAGE_INDEX) + got = bytearray(self.BUF) + sink = (ctypes.c_ubyte * self.BUF).from_buffer(got) + (err,) = cudart.cudaMemcpy( + ctypes.addressof(sink), + arena + index * self.BUF, + self.BUF, + cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost, + ) + self.assertEqual(int(err), 0) + self.assertEqual(bytes(got), pattern, "the cell does not hold its page's bytes") + # Removal returns the lease and discards the GPU cache, so the + # write-back record has to match the history being given up. + batch.flush_sparse_units([4 * self.TPB]) + batch.remove_kv_cache(cache) + finally: + if cache is not None: + cache.close() + manager.shutdown() + + def test_a_flushed_cell_lands_in_its_hot_page(self): + ctypes, cudart, torch = self._deps() + manager = self._manager() + cache = None + try: + lg = manager.get_layer_group_id(0) + pool = int(manager.get_mem_pool_base_address(LayerId(0), "key")) + stride = manager.get_page_stride(LayerId(0), "key") + cache = manager.create_kv_cache() + with TemporaryCudaStream([]) as stream: + handle = cast(CudaStream, stream.handle) + batch = Batch( + manager, + handle, + max_batch_size=1, + max_blocks_per_seq=8, + index_tables=manager.sparse_index_table_specs(), + ) + arena = int(batch.get_sparse_cache_base_address(LayerId(0), "key")) + self.assertTrue(cache.resume(handle)) + batch.add_kv_cache(cache) + # Capacity one token past a watermark that has not reached unit 0, so unit 0 + # is in flight and the resolve pins a cell for it. + self.assertTrue(cache.resize(self.TPB, self.TPB - 1)) + batch.publish() + selected = torch.tensor([[-1, -1]], dtype=torch.int32, device="cuda") + out = torch.full((1, 1, self.SELECTED), 12345, dtype=torch.int32, device="cuda") + batch.fetch_sparse_pages([BufferId(0, "key")], selected, out, handle) + torch.cuda.synchronize() + + resident = _introspection.sparse_cache_resident_ordinals(cache, lg) + self.assertIn(0, resident, "the in-flight unit must hold a cell") + cell = _introspection.sparse_cache_base_cell(cache, lg) + resident.index(0) + + pattern = self._pattern(self.BUF, 11) + source = (ctypes.c_ubyte * self.BUF).from_buffer(bytearray(pattern)) + (err,) = cudart.cudaMemcpy( + arena + cell * self.BUF, + ctypes.addressof(source), + self.BUF, + cudart.cudaMemcpyKind.cudaMemcpyHostToDevice, + ) + self.assertEqual(int(err), 0) + + batch.wait_published() + # The watermark crosses unit 0, so its cell is written back. The flush runs + # before the move, since a block must reach its host page before it becomes + # selectable. + batch.flush_sparse_units([self.TPB]) + self.assertTrue(cache.resize(self.TPB + 1, self.TPB)) + torch.cuda.synchronize() + + pages = [int(x) for x in cache.get_base_page_indices(lg)] + landed = ctypes.string_at(pool + pages[0] * stride, self.BUF) + self.assertEqual(landed, pattern, "the write-back did not reach the hot page") + batch.remove_kv_cache(cache) + finally: + if cache is not None: + cache.close() + manager.shutdown() + + +class TestSparseWriteIndices(unittest.TestCase): + """prepare_sparse_write_indices: the cell the model writes its new unit into. + + The unit between the two watermarks has no host page yet, so nothing can be fetched for + it; what the model needs is only where to put it. These cover that the reported cell is + the one the write-back later reads, that asking twice does not move it, and that a fetch + running afterwards leaves it alone. + """ + + TPB = 128 + BUF = 4096 + SELECTED = 2 + CELLS = 4 + + def setUp(self) -> None: + init_cuda_once() + + @classmethod + def _manager(cls) -> KVCacheManager: + return KVCacheManager( + KVCacheManagerConfig( + tokens_per_block=cls.TPB, + cache_tiers=[GpuCacheTierConfig(64 << 20), HostCacheTierConfig(64 << 20)], + constraints=[ + BatchDesc(kv_caches=[KVCacheDesc(capacity=4096, history_length=0)] * 2) + ], + layers=[ + AttentionLayerConfig( + 0, + [ + BufferConfig( + "key", + cls.BUF, + is_sparse=True, + ), + # A pool group is sized from its device-backed pools, so a sparse layer + # carries one. This is also the shape a model uses: an indexer and the + # data it selects over are one attention computation. + BufferConfig("index", 256), + ], + sparse_cache=SparseCacheConfig( + num_selected_blocks=cls.SELECTED, + max_cached_units=cls.CELLS, + unit_tokens=cls.TPB, + ), + ) + ], + ) + ) + + def _deps(self): + try: + import ctypes + + import cuda.bindings.runtime as cudart + import torch + except ImportError: # pragma: no cover - fast mode without torch + raise unittest.SkipTest("torch and the CUDA bindings are needed to address raw memory") + return ctypes, cudart, torch + + @staticmethod + def _pattern(size: int, seed: int) -> bytes: + return bytes((i * seed + 13) % 251 for i in range(size)) + + def test_the_prepared_cell_is_the_one_written_back(self): + """The whole contract in one path: write where it says, and the page gets it.""" + ctypes, cudart, torch = self._deps() + manager = self._manager() + cache = None + try: + lg = manager.get_layer_group_id(0) + pool = int(manager.get_mem_pool_base_address(LayerId(0), "key")) + stride = manager.get_page_stride(LayerId(0), "key") + cache = manager.create_kv_cache() + with TemporaryCudaStream([]) as stream: + handle = cast(CudaStream, stream.handle) + batch = Batch( + manager, + handle, + max_batch_size=1, + max_blocks_per_seq=8, + index_tables=manager.sparse_index_table_specs(), + ) + arena = int(batch.get_sparse_cache_base_address(LayerId(0), "key")) + self.assertTrue(cache.resume(handle)) + batch.add_kv_cache(cache) + # Capacity one token past a history that has not reached unit 0, so unit 0 + # is the unit in flight. + self.assertTrue(cache.resize(self.TPB, self.TPB - 1)) + batch.publish() + + out = torch.full((1, 1), 12345, dtype=torch.int32, device="cuda") + batch.prepare_sparse_write_indices([BufferId(0, "key")], out, handle) + torch.cuda.synchronize() + index = int(out[0][0]) + self.assertNotEqual(index, BAD_PAGE_INDEX, "a unit is in flight") + + pattern = self._pattern(self.BUF, 41) + source = (ctypes.c_ubyte * self.BUF).from_buffer(bytearray(pattern)) + (err,) = cudart.cudaMemcpy( + arena + index * self.BUF, + ctypes.addressof(source), + self.BUF, + cudart.cudaMemcpyKind.cudaMemcpyHostToDevice, + ) + self.assertEqual(int(err), 0) + + batch.wait_published() + # The watermark crosses unit 0, so its cell is written back. The flush runs + # before the move, since a block must reach its host page before it becomes + # selectable. + batch.flush_sparse_units([self.TPB]) + self.assertTrue(cache.resize(self.TPB + 1, self.TPB)) + torch.cuda.synchronize() + + pages = [int(x) for x in cache.get_base_page_indices(lg)] + landed = ctypes.string_at(pool + pages[0] * stride, self.BUF) + self.assertEqual(landed, pattern, "the prepared cell was not the one written back") + batch.remove_kv_cache(cache) + finally: + if cache is not None: + cache.close() + manager.shutdown() + + def test_preparing_twice_reports_the_same_cell(self): + """Idempotent while the watermarks hold: the second call must not move the unit.""" + _, _, torch = self._deps() + manager = self._manager() + cache = None + try: + cache = manager.create_kv_cache() + with TemporaryCudaStream([]) as stream: + handle = cast(CudaStream, stream.handle) + batch = Batch( + manager, + handle, + max_batch_size=1, + max_blocks_per_seq=8, + index_tables=manager.sparse_index_table_specs(), + ) + self.assertTrue(cache.resume(handle)) + batch.add_kv_cache(cache) + self.assertTrue(cache.resize(self.TPB, self.TPB - 1)) + batch.publish() + + seen = [] + for _ in range(3): + out = torch.full((1, 1), 12345, dtype=torch.int32, device="cuda") + batch.prepare_sparse_write_indices([BufferId(0, "key")], out, handle) + torch.cuda.synchronize() + seen.append(int(out[0][0])) + self.assertNotEqual(seen[0], BAD_PAGE_INDEX) + self.assertEqual(len(set(seen)), 1, f"the cell moved between calls: {seen}") + # Removal returns the lease and discards the GPU cache, so the + # write-back record has to match the history being given up. + batch.flush_sparse_units([self.TPB - 1]) + batch.remove_kv_cache(cache) + finally: + if cache is not None: + cache.close() + manager.shutdown() + + def test_a_row_with_nothing_in_flight_reports_bad_page_index(self): + """Both watermarks on a unit boundary: there is no unit to write.""" + _, _, torch = self._deps() + manager = self._manager() + cache = None + try: + cache = manager.create_kv_cache() + with TemporaryCudaStream([]) as stream: + handle = cast(CudaStream, stream.handle) + batch = Batch( + manager, + handle, + max_batch_size=2, + max_blocks_per_seq=8, + index_tables=manager.sparse_index_table_specs(), + ) + self.assertTrue(cache.resume(handle)) + batch.add_kv_cache(cache) + self.assertTrue(cache.resize(self.TPB, self.TPB)) + batch.publish() + + out = torch.full((2, 1), 12345, dtype=torch.int32, device="cuda") + batch.prepare_sparse_write_indices([BufferId(0, "key")], out, handle) + torch.cuda.synchronize() + self.assertEqual( + int(out[0][0]), BAD_PAGE_INDEX, "capacity holds no unit above history" + ) + # Written in full: the unoccupied row must be overwritten, not left at 12345. + self.assertEqual( + int(out[1][0]), BAD_PAGE_INDEX, "a free row carries no reservation" + ) + # Removal returns the lease and discards the GPU cache, so the + # write-back record has to match the history being given up. + batch.flush_sparse_units([self.TPB, 0]) + batch.remove_kv_cache(cache) + finally: + if cache is not None: + cache.close() + manager.shutdown() + + def test_every_buffer_of_a_layer_gets_its_own_cell(self): + """Two buffers, and a headroom wider than what is actually in flight. + + The resolve's in-flight array is a row per request of max_cached_units minus the + selection width, while at most one unit is ever in flight. Reading it with the wrong + row stride writes one buffer's slots over the next buffer's output, which this sees + as the second buffer losing its index. + """ + _, _, torch = self._deps() + manager = KVCacheManager( + KVCacheManagerConfig( + tokens_per_block=self.TPB, + cache_tiers=[GpuCacheTierConfig(64 << 20), HostCacheTierConfig(64 << 20)], + constraints=[ + BatchDesc(kv_caches=[KVCacheDesc(capacity=4096, history_length=0)] * 2) + ], + layers=[ + AttentionLayerConfig( + 0, + [ + BufferConfig( + role, + self.BUF, + is_sparse=True, + # Headroom of CELLS - SELECTED = 2, wider than the one unit + # that can be in flight. + ) + for role in ("key", "value") + ] + + [ + # A pool group sizes its slots from its device-backed pools. + BufferConfig("index", 256) + ], + sparse_cache=SparseCacheConfig( + num_selected_blocks=self.SELECTED, + max_cached_units=self.CELLS, + unit_tokens=self.TPB, + ), + ) + ], + ) + ) + cache = None + try: + cache = manager.create_kv_cache() + with TemporaryCudaStream([]) as stream: + handle = cast(CudaStream, stream.handle) + batch = Batch( + manager, + handle, + max_batch_size=1, + max_blocks_per_seq=8, + index_tables=manager.sparse_index_table_specs(), + ) + self.assertTrue(cache.resume(handle)) + batch.add_kv_cache(cache) + self.assertTrue(cache.resize(self.TPB, self.TPB - 1)) + batch.publish() + + out = torch.full((1, 2), 12345, dtype=torch.int32, device="cuda") + batch.prepare_sparse_write_indices( + [BufferId(0, "key"), BufferId(0, "value")], out, handle + ) + torch.cuda.synchronize() + got = [int(out[0][0]), int(out[0][1])] + self.assertNotIn( + BAD_PAGE_INDEX, got, f"both buffers hold the in-flight unit: {got}" + ) + self.assertEqual(len(set(got)), 2, f"buffers must not share a cell: {got}") + # Removal returns the lease and discards the GPU cache, so the + # write-back record has to match the history being given up. + batch.flush_sparse_units([self.TPB - 1]) + batch.remove_kv_cache(cache) + finally: + if cache is not None: + cache.close() + manager.shutdown() + + def test_each_layer_reports_a_cell_of_its_own(self): + """One unit ordinal, one index per layer -- never one answer for the step. + + A layer resolves against its own resident set and its own recency order, so the same + ordinal lands on a different entry in each; and each layer's buffer owns its own run + of cells, so even equal entries would not collide. Both are why the call takes one + layer rather than a group. + """ + _, _, torch = self._deps() + layers = 3 + manager = KVCacheManager( + KVCacheManagerConfig( + tokens_per_block=self.TPB, + cache_tiers=[GpuCacheTierConfig(64 << 20), HostCacheTierConfig(64 << 20)], + constraints=[ + BatchDesc(kv_caches=[KVCacheDesc(capacity=4096, history_length=0)] * 2) + ], + layers=[ + AttentionLayerConfig( + layer, + [ + BufferConfig( + "key", + self.BUF, + is_sparse=True, + ), + # A pool group is sized from its device-backed pools, so a sparse layer + # carries one. This is also the shape a model uses: an indexer and the + # data it selects over are one attention computation. + BufferConfig("index", 256), + ], + sparse_cache=SparseCacheConfig( + num_selected_blocks=self.SELECTED, + max_cached_units=self.CELLS, + unit_tokens=self.TPB, + ), + ) + for layer in range(layers) + ], + ) + ) + cache = None + try: + cache = manager.create_kv_cache() + with TemporaryCudaStream([]) as stream: + handle = cast(CudaStream, stream.handle) + batch = Batch( + manager, + handle, + max_batch_size=1, + max_blocks_per_seq=8, + index_tables=manager.sparse_index_table_specs(), + ) + self.assertTrue(cache.resume(handle)) + batch.add_kv_cache(cache) + self.assertTrue(cache.resize(self.TPB, self.TPB - 1)) + batch.publish() + + indices = [] + for layer in range(layers): + out = torch.full((1, 1), 12345, dtype=torch.int32, device="cuda") + batch.prepare_sparse_write_indices( + [BufferId(LayerId(layer), "key")], out, handle + ) + torch.cuda.synchronize() + indices.append(int(out[0][0])) + + self.assertNotIn(BAD_PAGE_INDEX, indices, "every layer has the same unit in flight") + self.assertEqual( + len(set(indices)), layers, f"layers must not share a cell: {indices}" + ) + # Each layer's buffer owns an adjacent run of cells_per_buffer cells, so the + # runs the three indices fall in are consecutive. + runs = sorted(index // self.CELLS for index in indices) + self.assertEqual( + runs, + list(range(runs[0], runs[0] + layers)), + f"cell runs are not per buffer: {indices}", + ) + # Removal returns the lease and discards the GPU cache, so the + # write-back record has to match the history being given up. + batch.flush_sparse_units([self.TPB - 1]) + batch.remove_kv_cache(cache) + finally: + if cache is not None: + cache.close() + manager.shutdown() + + def test_a_fetch_neither_moves_nor_overwrites_the_prepared_cell(self): + """A fetch runs its own resolve; the prepared cell must come through unchanged. + + It covers the copy list: an in-flight unit's host page holds nothing yet, so a fetch + that treated its entry as an ordinary miss would copy stale bytes over what the model + wrote. Removing that guard fails this. + + It does not cover the separate `pinned` term in the victim test. Every in-flight + ordinal is probed in the resolve's required prefix, so its entry is always classified + hit or fetched and never `kClassUntouched` -- which the victim test already excludes. + Dropping `pinned` leaves every test here passing, so it appears to be redundant with + that classification rather than load-bearing. + """ + ctypes, cudart, torch = self._deps() + manager = self._manager() + cache = None + try: + lg = manager.get_layer_group_id(0) + pool = int(manager.get_mem_pool_base_address(LayerId(0), "key")) + stride = manager.get_page_stride(LayerId(0), "key") + cache = manager.create_kv_cache() + with TemporaryCudaStream([]) as stream: + handle = cast(CudaStream, stream.handle) + batch = Batch( + manager, + handle, + max_batch_size=1, + max_blocks_per_seq=8, + index_tables=manager.sparse_index_table_specs(), + ) + arena = int(batch.get_sparse_cache_base_address(LayerId(0), "key")) + self.assertTrue(cache.resume(handle)) + batch.add_kv_cache(cache) + # resize() bounds capacity - history to one token, so a unit is in flight + # only as history reaches a unit boundary: two units of history behind the + # watermark, and unit 2 in flight. + self.assertTrue(cache.resize(3 * self.TPB, 3 * self.TPB - 1)) + batch.publish() + + out = torch.full((1, 1), 12345, dtype=torch.int32, device="cuda") + batch.prepare_sparse_write_indices([BufferId(0, "key")], out, handle) + torch.cuda.synchronize() + index = int(out[0][0]) + self.assertNotEqual(index, BAD_PAGE_INDEX) + + written = self._pattern(self.BUF, 53) + source = (ctypes.c_ubyte * self.BUF).from_buffer(bytearray(written)) + (err,) = cudart.cudaMemcpy( + arena + index * self.BUF, + ctypes.addressof(source), + self.BUF, + cudart.cudaMemcpyKind.cudaMemcpyHostToDevice, + ) + self.assertEqual(int(err), 0) + + # Paint the history pages differently, so a copy over the in-flight cell + # would be visible rather than coincidentally matching. + pages = [int(x) for x in cache.get_base_page_indices(lg)] + for ordinal in range(2): + ctypes.memmove( + pool + pages[ordinal] * stride, + self._pattern(self.BUF, 60 + ordinal), + self.BUF, + ) + + # A full-width selection of the two eligible units, so the fetch fills every + # cell the reservation has beyond the one in flight. + selected = torch.tensor([[0, 1]], dtype=torch.int32, device="cuda") + fetched = torch.full((1, 1, self.SELECTED), 12345, dtype=torch.int32, device="cuda") + batch.fetch_sparse_pages([BufferId(0, "key")], selected, fetched, handle) + torch.cuda.synchronize() + + after = torch.full((1, 1), 12345, dtype=torch.int32, device="cuda") + batch.prepare_sparse_write_indices([BufferId(0, "key")], after, handle) + torch.cuda.synchronize() + self.assertEqual(int(after[0][0]), index, "the fetch moved the in-flight cell") + + got = bytearray(self.BUF) + sink = (ctypes.c_ubyte * self.BUF).from_buffer(got) + (err,) = cudart.cudaMemcpy( + ctypes.addressof(sink), + arena + index * self.BUF, + self.BUF, + cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost, + ) + self.assertEqual(int(err), 0) + self.assertEqual(bytes(got), written, "the fetch overwrote the in-flight cell") + # Removal returns the lease and discards the GPU cache, so the + # write-back record has to match the history being given up. + batch.flush_sparse_units([3 * self.TPB - 1]) + batch.remove_kv_cache(cache) + finally: + if cache is not None: + cache.close() + manager.shutdown() + + +class TestSparseQuantisedDataPath(unittest.TestCase): + """The quantised shape end to end: one layer, two buffers, two pools. + + Every other sparse data-path test runs a life cycle whose buffers are alike, so one pool + serves them all and a single set of strides is right by accident. Here the data and its + block scales differ in size, so they sit in pools -- and arena regions -- of their own, + and the fetch is one launch per pool. A copy that used the wrong pool's cell size, slot + stride or region base would read or write the wrong bytes, which is what this reads back. + """ + + TPB = 128 + DATA = 4096 + SCALE = 128 + SELECTED = 2 + CELLS = 4 + + def setUp(self) -> None: + init_cuda_once() + + @classmethod + def _manager(cls) -> KVCacheManager: + def buffer(role: str, size: int) -> BufferConfig: + return BufferConfig( + role, + size, + is_sparse=True, + ) + + return KVCacheManager( + KVCacheManagerConfig( + tokens_per_block=cls.TPB, + cache_tiers=[GpuCacheTierConfig(64 << 20), HostCacheTierConfig(64 << 20)], + constraints=[ + BatchDesc(kv_caches=[KVCacheDesc(capacity=4096, history_length=0)] * 2) + ], + layers=[ + AttentionLayerConfig( + 0, + [ + buffer("key", cls.DATA), + buffer("key_scale", cls.SCALE), + # A pool group sizes its slots from its device-backed pools. + BufferConfig("index", 256), + ], + sparse_cache=SparseCacheConfig( + num_selected_blocks=cls.SELECTED, + max_cached_units=cls.CELLS, + unit_tokens=cls.TPB, + ), + ) + ], + ) + ) + + def _deps(self): + try: + import ctypes + + import cuda.bindings.runtime as cudart + import torch + except ImportError: # pragma: no cover - fast mode without torch + raise unittest.SkipTest("torch and the CUDA bindings are needed to address raw memory") + return ctypes, cudart, torch + + @staticmethod + def _pattern(size: int, seed: int) -> bytes: + return bytes((i * seed + 13) % 251 for i in range(size)) + + def test_each_buffer_is_fetched_into_its_own_pool_region(self): + ctypes, cudart, torch = self._deps() + manager = self._manager() + cache = None + roles = (("key", self.DATA, 7), ("key_scale", self.SCALE, 23)) + try: + lg = manager.get_layer_group_id(0) + cache = manager.create_kv_cache() + with TemporaryCudaStream([]) as stream: + handle = cast(CudaStream, stream.handle) + batch = Batch( + manager, + handle, + max_batch_size=1, + max_blocks_per_seq=8, + index_tables=manager.sparse_index_table_specs(), + ) + self.assertTrue(cache.resume(handle)) + batch.add_kv_cache(cache) + self.assertTrue(cache.resize(4 * self.TPB, 4 * self.TPB)) + + # Paint one block's hot page in each pool with a pattern of its own, so a + # copy that crossed pools reads bytes that cannot pass for the right ones. + ordinal = 1 + pages = [int(x) for x in cache.get_base_page_indices(lg)] + patterns = {} + for role, size, seed in roles: + patterns[role] = self._pattern(size, seed) + pool = int(manager.get_mem_pool_base_address(LayerId(0), role)) + stride = manager.get_page_stride(LayerId(0), role) + ctypes.memmove(pool + pages[ordinal] * stride, patterns[role], size) + + batch.publish() + selected = torch.tensor([[ordinal, -1]], dtype=torch.int32, device="cuda") + out = torch.full((1, 2, self.SELECTED), 12345, dtype=torch.int32, device="cuda") + batch.fetch_sparse_pages( + [BufferId(0, "key"), BufferId(0, "key_scale")], + selected, + out, + handle, + ) + torch.cuda.synchronize() + + slots = batch.sparse_out_slots([BufferId(0, "key"), BufferId(0, "key_scale")]) + for position, (role, size, _) in enumerate(roles): + # At the buffer's slice, not its position: mirrored pools resolve to + # one index and share a slice, leaving the other holding the fill. + index = int(out[0][slots[position]][0]) + self.assertNotEqual(index, 12345, f"{role}'s slice was never written") + self.assertNotEqual(index, BAD_PAGE_INDEX, role) + arena = int(batch.get_sparse_cache_base_address(LayerId(0), role)) + got = bytearray(size) + sink = (ctypes.c_ubyte * size).from_buffer(got) + (err,) = cudart.cudaMemcpy( + ctypes.addressof(sink), + arena + index * size, + size, + cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost, + ) + self.assertEqual(int(err), 0, role) + self.assertEqual( + bytes(got), patterns[role], f"{role}'s cell does not hold its page" + ) + # Removal returns the lease and discards the GPU cache, so the + # write-back record has to match the history being given up. + batch.flush_sparse_units([4 * self.TPB]) + batch.remove_kv_cache(cache) + finally: + if cache is not None: + cache.close() + manager.shutdown() + + def test_each_buffer_is_written_back_to_its_own_hot_page(self): + ctypes, cudart, torch = self._deps() + manager = self._manager() + cache = None + roles = (("key", self.DATA, 31), ("key_scale", self.SCALE, 37)) + try: + lg = manager.get_layer_group_id(0) + cache = manager.create_kv_cache() + with TemporaryCudaStream([]) as stream: + handle = cast(CudaStream, stream.handle) + batch = Batch( + manager, + handle, + max_batch_size=1, + max_blocks_per_seq=8, + index_tables=manager.sparse_index_table_specs(), + ) + self.assertTrue(cache.resume(handle)) + batch.add_kv_cache(cache) + # Capacity one token past a watermark that has not reached unit 0, so unit 0 + # is in flight and the resolve pins a cell for it in both pools. + self.assertTrue(cache.resize(self.TPB, self.TPB - 1)) + batch.publish() + selected = torch.tensor([[-1, -1]], dtype=torch.int32, device="cuda") + out = torch.full((1, 2, self.SELECTED), 12345, dtype=torch.int32, device="cuda") + batch.fetch_sparse_pages( + [BufferId(0, "key"), BufferId(0, "key_scale")], + selected, + out, + handle, + ) + torch.cuda.synchronize() + + resident = _introspection.sparse_cache_resident_ordinals(cache, lg) + self.assertIn(0, resident, "the in-flight unit must hold a cell") + entry = resident.index(0) + + patterns = {} + for pool_index, (role, size, seed) in enumerate(roles): + patterns[role] = self._pattern(size, seed) + arena = int(batch.get_sparse_cache_base_address(LayerId(0), role)) + cell = _introspection.sparse_cache_base_cell(cache, lg, pool_index) + entry + source = (ctypes.c_ubyte * size).from_buffer(bytearray(patterns[role])) + (err,) = cudart.cudaMemcpy( + arena + cell * size, + ctypes.addressof(source), + size, + cudart.cudaMemcpyKind.cudaMemcpyHostToDevice, + ) + self.assertEqual(int(err), 0, role) + + batch.wait_published() + # The watermark crosses unit 0, so both pools' cells are written back. The + # flush runs before the move, since a block must reach its host page before + # it becomes selectable. + batch.flush_sparse_units([self.TPB]) + self.assertTrue(cache.resize(self.TPB + 1, self.TPB)) + torch.cuda.synchronize() + + pages = [int(x) for x in cache.get_base_page_indices(lg)] + for role, size, _ in roles: + pool = int(manager.get_mem_pool_base_address(LayerId(0), role)) + stride = manager.get_page_stride(LayerId(0), role) + landed = ctypes.string_at(pool + pages[0] * stride, size) + self.assertEqual(landed, patterns[role], f"{role}'s write-back went astray") + batch.remove_kv_cache(cache) + finally: + if cache is not None: + cache.close() + manager.shutdown() + + +class TestSparseFlush(unittest.TestCase): + """Writing a unit back to its host page once the watermark passes it.""" + + def setUp(self) -> None: + init_cuda_once() + + def test_only_units_the_watermark_crossed_are_listed(self): + # Entries hold units 3, 7, 11 and one free entry. + resident = [3, 7, 11, BAD_PAGE_INDEX] + result = _introspection.sparse_flush_list( + resident, + num_layers=1, + entries_per_layer=4, + from_units=5, + to_units=11, + max_flush_units=4, + ) + self.assertEqual(result.count, [1], "only unit 7 is in [5, 11)") + self.assertEqual( + result.total, [1], "the list is compacted, so its total says how far it is filled" + ) + self.assertEqual(result.ordinal[0], 7) + self.assertEqual(result.slot[0], 1) + self.assertEqual( + result.row[0], 0, "the entry carries its row, since position no longer does" + ) + + def test_a_watermark_that_moved_nothing_lists_nothing(self): + resident = [3, 7, 11, BAD_PAGE_INDEX] + result = _introspection.sparse_flush_list( + resident, num_layers=1, entries_per_layer=4, from_units=8, to_units=8, max_flush_units=4 + ) + self.assertEqual(result.count, [0]) + + def test_layers_are_scanned_independently(self): + # Layer 0 holds unit 4, layer 1 holds unit 9; only layer 0's is in range. + resident = [4, BAD_PAGE_INDEX, 9, BAD_PAGE_INDEX] + result = _introspection.sparse_flush_list( + resident, num_layers=2, entries_per_layer=2, from_units=0, to_units=5, max_flush_units=2 + ) + self.assertEqual(result.count, [1, 0]) + # One list per layer, each with its own counter, so layer 1's stays empty. + self.assertEqual(result.total, [1, 0]) + self.assertEqual(result.ordinal[0], 4) + + def test_write_back_inverts_the_fetch(self): + """Fetching a unit and writing it back must reproduce the cold page.""" + unit, units = 16, 2 + page_stride = unit * units + cold = [(i * 5 + 3) % 251 for i in range(page_stride)] + common = dict( + host_page_stride=page_stride, + page_table=[0], + base_cell=0, + num_blocks=1, + cells_per_buffer=1, + buffer_bytes=page_stride, + num_chunks=1, + chunk_bytes=unit, + chunk_stride=0, + unit_stride=unit, + expansion=units, + miss_ordinal=[1], + miss_slot=[1], + miss_count=[1], + selection_width=1, + buffer_layer=list(range(1)), + ) + hot = _introspection.sparse_unit_copy( + host_pool=cold, device_arena=[0] * page_stride, write_back=False, **common + ) + self.assertEqual(hot[unit:], cold[unit:], "the fetch landed the unit") + + # Write that same entry back over a blanked cold page; the unit must return. + restored = _introspection.sparse_unit_copy( + host_pool=[0] * page_stride, device_arena=hot, write_back=True, **common + ) + self.assertEqual(restored[unit:], cold[unit:], "the write-back restored the unit") + self.assertEqual(restored[:unit], [0] * unit, "only the listed unit moved") + + def test_write_back_moves_into_the_cold_page(self): + # The device pool is zeroed by the hook, so a write-back must blank the cold unit + # it targets and leave the rest of the cold page alone. + unit, units = 16, 2 + page_stride = unit * units + cold = [(i * 5 + 3) % 251 for i in range(page_stride)] + result = _introspection.sparse_unit_copy( + host_pool=cold, + host_page_stride=page_stride, + page_table=[0], + device_arena=[0] * (page_stride), + base_cell=0, + num_blocks=1, + cells_per_buffer=1, + buffer_bytes=page_stride, + num_chunks=1, + chunk_bytes=unit, + chunk_stride=0, + unit_stride=unit, + expansion=units, + miss_ordinal=[1], + miss_slot=[1], + miss_count=[1], + selection_width=1, + buffer_layer=list(range(1)), + write_back=True, + ) + self.assertEqual(result[unit:], [0] * unit, "unit 1 was overwritten from the entry") + self.assertEqual(result[:unit], cold[:unit], "unit 0 was untouched") + + +class TestBatch(unittest.TestCase): + """Row order and the device-side page table the sparse fetch reads.""" + + BAD = BAD_PAGE_INDEX + + def setUp(self) -> None: + init_cuda_once() + + @staticmethod + def _manager() -> KVCacheManager: + return KVCacheManager( + KVCacheManagerConfig( + tokens_per_block=128, + cache_tiers=[GpuCacheTierConfig(64 << 20), HostCacheTierConfig(64 << 20)], + layers=[AttentionLayerConfig(0, [BufferConfig("key", 4096)])], + ) + ) + + def test_device_arrays_are_consumable_as_dlpack(self): + """The page table and watermarks are what a caller actually receives. + + Imported lazily: the fast-mode path deliberately avoids loading the whole package. + """ + try: + import torch + except ImportError: # pragma: no cover - fast mode without torch + raise unittest.SkipTest("torch is needed to consume a DLPack array") + + manager = self._manager() + try: + batch = Batch( + manager, NULL_STREAM, index_tables=[], max_batch_size=2, max_blocks_per_seq=4 + ) + lg_id = manager.get_layer_group_id(0) + table = torch.from_dlpack(_introspection.batch_gpu_coalesced_page_table(batch, lg_id)) + history = torch.from_dlpack(batch.history_lengths()) + capacity = torch.from_dlpack(batch.capacities()) + + self.assertEqual(tuple(table.shape), (2, 4)) + self.assertEqual(tuple(history.shape), (2,)) + self.assertEqual(tuple(capacity.shape), (2,)) + for arr in (table, history, capacity): + self.assertEqual(arr.dtype, torch.int32) + self.assertEqual(arr.device.type, "cuda") + # Unoccupied rows: absent blocks, and nothing eligible. + self.assertTrue(bool((table == BAD_PAGE_INDEX).all())) + self.assertTrue(bool((history == 0).all())) + finally: + manager.shutdown() + + def test_sparse_leases_are_staged_per_row(self): + """What is mirrored is the handle, not the memory: the reservation stays per request. + + A view appears for a member holding one and stays absent for a free row, so the + fetch can tell the two apart. + """ + manager = KVCacheManager( + KVCacheManagerConfig( + tokens_per_block=128, + cache_tiers=[GpuCacheTierConfig(64 << 20), HostCacheTierConfig(64 << 20)], + # The cache arena is reserved up front, so a sparse deployment must say + # how many requests run together. + constraints=[ + BatchDesc(kv_caches=[KVCacheDesc(capacity=4096, history_length=0)] * 4) + ], + layers=[ + AttentionLayerConfig( + 0, + [ + BufferConfig( + "key", + 4096, + is_sparse=True, + ), + # A pool group is sized from its device-backed pools, so a sparse layer + # carries one. This is also the shape a model uses: an indexer and the + # data it selects over are one attention computation. + BufferConfig("index", 256), + ], + sparse_cache=SparseCacheConfig( + num_selected_blocks=2, max_cached_units=4, unit_tokens=128 + ), + ) + ], + ) + ) + cache = None + batch = None + try: + lg_id = manager.get_layer_group_id(0) + with TemporaryCudaStream([]) as stream: + handle = cast(CudaStream, stream.handle) + batch = Batch( + manager, + handle, + max_batch_size=2, + max_blocks_per_seq=4, + index_tables=manager.sparse_index_table_specs(), + ) + # Before anyone resumes, no reservation exists, so no view does either. + batch.publish() + self.assertEqual(_introspection.batch_sparse_entries_per_row(batch, lg_id), [0, 0]) + + cache = manager.create_kv_cache() + self.assertTrue(cache.resume(handle)) + batch.add_kv_cache(cache) + batch.publish() + entries = _introspection.batch_sparse_entries_per_row(batch, lg_id) + self.assertEqual(entries[0], manager.get_sparse_cache_max_units(0)) + self.assertEqual(entries[1], 0, "the free row stays absent") + + # Leaving releases the reservation, so the view must go with it rather than + # leaving the fetch a handle into freed memory. A member is removed before + # it suspends -- suspending in place would write its row behind the Batch's + # back. + batch.wait_published() + batch.remove_kv_cache(cache) + cache.suspend() + batch.publish() + self.assertEqual(_introspection.batch_sparse_entries_per_row(batch, lg_id), [0, 0]) + finally: + if cache is not None: + if batch is not None and batch.row_of(cache) is not None: + batch.remove_kv_cache(cache) + cache.close() + manager.shutdown() + + def test_membership_is_countable_and_a_hole_is_visible(self): + """An engine mapping rows onto request slots asks the Batch, not a mirror of it.""" + manager = self._manager() + try: + with TemporaryCudaStream([]) as stream: + handle = cast(CudaStream, stream.handle) + batch = Batch( + manager, handle, index_tables=[], max_batch_size=3, max_blocks_per_seq=4 + ) + self.assertEqual(batch.num_kv_caches, 0) + self.assertTrue(batch.is_contiguous, "an empty batch has no hole") + + caches = [manager.create_kv_cache() for _ in range(3)] + for cache in caches: + self.assertTrue(cache.resume(handle)) + batch.add_kv_cache(cache) + self.assertEqual(batch.num_kv_caches, 3) + self.assertTrue(batch.is_contiguous) + + # A hole is legal, and is exactly what an engine wants to be told about. + batch.remove_kv_cache(caches[1]) + self.assertEqual(batch.num_kv_caches, 2) + self.assertFalse(batch.is_contiguous) + + # Refilling it closes the hole; dropping the tail never opens one. + batch.add_kv_cache(caches[1]) + self.assertTrue(batch.is_contiguous) + batch.remove_kv_cache(caches[2]) + self.assertEqual(batch.num_kv_caches, 2) + self.assertTrue(batch.is_contiguous) + + for cache in caches: + if batch.row_of(cache) is not None: + batch.remove_kv_cache(cache) + cache.close() + finally: + manager.shutdown() + + def test_rows_are_slots_so_a_hole_is_reused(self): + manager = self._manager() + try: + with TemporaryCudaStream([]) as stream: + handle = cast(CudaStream, stream.handle) + batch = Batch( + manager, handle, index_tables=[], max_batch_size=3, max_blocks_per_seq=4 + ) + first = manager.create_kv_cache() + second = manager.create_kv_cache() + third = manager.create_kv_cache() + for cache in (first, second, third): + self.assertTrue(cache.resume(handle)) + self.assertEqual(batch.add_kv_cache(first), 0) + self.assertEqual(batch.add_kv_cache(second), 1) + self.assertEqual(batch.add_kv_cache(third), 2) + + # Removing the middle member must not shift the ones after it: addresses and + # row indices have to survive a graph capture. + batch.remove_kv_cache(second) + self.assertEqual(batch.row_of(third), 2) + self.assertEqual(batch.row_of(second), None) + + fourth = manager.create_kv_cache() + self.assertTrue(fourth.resume(handle)) + self.assertEqual(batch.add_kv_cache(fourth), 1, "the hole is reused") + + for cache in (first, third, fourth): + batch.remove_kv_cache(cache) + for cache in (first, second, third, fourth): + cache.close() + finally: + manager.shutdown() + + def test_a_taken_row_is_refused_rather_than_overwritten(self): + """Seating over a member would rebind its row, which nothing downstream would notice.""" + manager = self._manager() + try: + with TemporaryCudaStream([]) as stream: + handle = cast(CudaStream, stream.handle) + batch = Batch( + manager, handle, index_tables=[], max_batch_size=2, max_blocks_per_seq=4 + ) + first = manager.create_kv_cache() + second = manager.create_kv_cache() + for cache in (first, second): + self.assertTrue(cache.resume(handle)) + self.assertEqual(batch.add_kv_cache(first, 0), 0) + + with self.assertRaisesRegex(Exception, "already occupied"): + batch.add_kv_cache(second, 0) + with self.assertRaisesRegex(Exception, "beyond max_batch_size"): + batch.add_kv_cache(second, 2) + # Neither attempt seated it, and neither disturbed the member already there. + self.assertEqual(batch.row_of(second), None) + self.assertEqual(batch.row_of(first), 0) + self.assertEqual(batch.num_kv_caches, 1) + + # The row it was refused is still free to whoever asks for it properly. + self.assertEqual(batch.add_kv_cache(second), 1) + batch.remove_kv_cache(first) + batch.remove_kv_cache(second) + first.close() + second.close() + finally: + manager.shutdown() + + def test_indexing_a_row_gives_its_member_or_none(self) -> None: + """The inverse of row_of(). None is an empty row, never a bad index.""" + manager = self._manager() + caches = [] + batch = None + try: + with TemporaryCudaStream([]) as stream: + handle = cast(CudaStream, stream.handle) + batch = Batch( + manager, handle, index_tables=[], max_batch_size=3, max_blocks_per_seq=4 + ) + cache = manager.create_kv_cache() + caches.append(cache) + self.assertTrue(cache.resume(handle)) + batch.add_kv_cache(cache, 1) + + seated = batch[1] + self.assertIsNotNone(seated) + # Compared through row_of rather than identity: whether nanobind hands back + # the same Python object is its business, not this API's contract. + self.assertEqual(batch.row_of(seated), 1) + self.assertIsNone(batch[0], "a free row reads as empty") + self.assertIsNone(batch[2]) + + # Out of range is a different answer from empty, and must not be confusable. + for bad in (3, -1, 99): + with self.assertRaises(IndexError): + batch[bad] + finally: + for cache in caches: + if batch is not None and batch.row_of(cache) is not None: + batch.remove_kv_cache(cache) + cache.close() + manager.shutdown() + + def test_move_carries_bindings_and_watermarks_to_the_empty_row(self) -> None: + try: + import torch + except ImportError: # pragma: no cover - fast mode without torch + raise unittest.SkipTest("torch is needed to pass device arrays") + + manager = self._manager() + caches = [] + batch = None + try: + with TemporaryCudaStream([]) as stream: + handle = cast(CudaStream, stream.handle) + batch = Batch( + manager, handle, index_tables=[], max_batch_size=3, max_blocks_per_seq=4 + ) + lg = manager.get_layer_group_id(0) + table = torch.from_dlpack(_introspection.batch_gpu_coalesced_page_table(batch, lg)) + address = table.data_ptr() + for history in (17, 31): + cache = manager.create_kv_cache() + caches.append(cache) + self.assertTrue(cache.resume(handle)) + self.assertTrue(cache.resize(128, history)) + batch.add_kv_cache(cache) + batch.flush_sparse_units([17, 31, 0]) + batch.publish() + + # Row 1 to the free row 2. Row 0 is a bystander and must not move. + batch.move(2, 1) + self.assertEqual([batch.row_of(cache) for cache in caches], [0, 2]) + self.assertEqual(batch.flushed_history(2), 31, "the watermark travels") + self.assertEqual(batch.flushed_history(0), 17, "the bystander is untouched") + + batch.publish() + torch.cuda.synchronize() + self.assertEqual(table.data_ptr(), address, "device tables keep their addresses") + self.assertEqual(table[0].tolist(), list(caches[0].get_base_page_indices(lg))) + self.assertEqual(table[2].tolist(), list(caches[1].get_base_page_indices(lg))) + self.assertTrue( + bool((table[1] == BAD_PAGE_INDEX).all()), + "the vacated row publishes as free, not as a stale copy of the member", + ) + self.assertEqual(torch.from_dlpack(batch.history_lengths()).tolist(), [17, 0, 31]) + self.assertEqual(torch.from_dlpack(batch.capacities()).tolist(), [128, 0, 128]) + finally: + for cache in caches: + if batch is not None and batch.row_of(cache) is not None: + batch.remove_kv_cache(cache) + cache.close() + manager.shutdown() + + def test_move_rejects_a_bad_pair_without_changing_rows(self) -> None: + manager = self._manager() + caches = [] + batch = None + try: + with TemporaryCudaStream([]) as stream: + handle = cast(CudaStream, stream.handle) + batch = Batch( + manager, handle, index_tables=[], max_batch_size=3, max_blocks_per_seq=4 + ) + for _ in range(2): + cache = manager.create_kv_cache() + caches.append(cache) + self.assertTrue(cache.resume(handle)) + batch.add_kv_cache(caches[0], 0) + batch.add_kv_cache(caches[1], 1) + + # Matched on the message rather than the type, so each case is pinned to the + # reason it is refused: a type alone would not distinguish them. + # + # (dst, src). dst == src is both faults at once, and is reported as the + # occupied destination -- a row cannot be occupied and empty together. + rejected = { + (0, 1): "destination row is already occupied", + (1, 1): "destination row is already occupied", + (0, 0): "destination row is already occupied", + (1, 2): "source row holds no member", + (3, 0): "beyond max_batch_size", + (2, 3): "beyond max_batch_size", + (-1, 0): "beyond max_batch_size", + (2, -1): "beyond max_batch_size", + } + for (dst, src), message in rejected.items(): + with self.assertRaisesRegex(Exception, message): + batch.move(dst, src) + self.assertEqual([batch.row_of(cache) for cache in caches], [0, 1]) + + # Still usable afterwards: a refusal must not have half-moved anything. + batch.move(2, 1) + self.assertEqual([batch.row_of(cache) for cache in caches], [0, 2]) + finally: + for cache in caches: + if batch is not None and batch.row_of(cache) is not None: + batch.remove_kv_cache(cache) + cache.close() + manager.shutdown() + + def test_the_batch_keeps_a_seated_member_alive(self): + """A caller may keep only the Batch: seating is what holds the cache up. - @functools.wraps(func) - def wrapper(self, *args, **kwargs): - with ref_cycle_check_context(): - result = func(self, *args, **kwargs) - return result + Dropping the caller's reference and then getting the cache back from the Batch is + the whole contract -- if the row held a raw pointer the object would already be + gone, and what came back would name freed memory. + """ + manager = self._manager() + try: + with TemporaryCudaStream([]) as stream: + handle = cast(CudaStream, stream.handle) + batch = Batch( + manager, handle, index_tables=[], max_batch_size=2, max_blocks_per_seq=4 + ) + cache = manager.create_kv_cache() + self.assertTrue(cache.resume(handle)) + row = batch.add_kv_cache(cache) + + # Drop the caller's only reference. Were the row raw, this would destroy the + # cache and poison on the way out, since a member must not outlive its row. + del cache + self.assertEqual(batch.num_kv_caches, 1) + + # Leaving hands it back, still usable -- which a dangling row could not do. + returned = batch.remove_row(row) + self.assertIsNotNone(returned) + self.assertEqual(batch.num_kv_caches, 0) + self.assertEqual(batch.row_of(returned), None) + returned.close() + + # A row nobody occupies has nothing to hand back. + self.assertIsNone(batch.remove_row(0)) + finally: + manager.shutdown() - return wrapper + def test_membership_is_exclusive(self): + manager = self._manager() + try: + with TemporaryCudaStream([]) as stream: + handle = cast(CudaStream, stream.handle) + one = Batch( + manager, handle, index_tables=[], max_batch_size=2, max_blocks_per_seq=4 + ) + two = Batch( + manager, handle, index_tables=[], max_batch_size=2, max_blocks_per_seq=4 + ) + cache = manager.create_kv_cache() + self.assertTrue(cache.resume(handle)) + one.add_kv_cache(cache) + with self.assertRaisesRegex(Exception, "at most one Batch"): + two.add_kv_cache(cache) + one.remove_kv_cache(cache) + # Free to join another batch once released. + two.add_kv_cache(cache) + two.remove_kv_cache(cache) + cache.close() + finally: + manager.shutdown() + def test_beam_search_is_refused_rather_than_half_mirrored(self): + """Only beam 0 is staged, so a wider request would have every beam read beam 0.""" + # Partial commit forbids beam_width changes, so this manager disables it. + manager = KVCacheManager( + KVCacheManagerConfig( + tokens_per_block=128, + cache_tiers=[GpuCacheTierConfig(64 << 20), HostCacheTierConfig(64 << 20)], + layers=[AttentionLayerConfig(0, [BufferConfig("key", 4096)])], + enable_partial_commit=False, + ) + ) + cache = None + try: + batch = Batch( + manager, NULL_STREAM, index_tables=[], max_batch_size=1, max_blocks_per_seq=4 + ) + cache = manager.create_kv_cache() + cache.beam_width = 2 + with self.assertRaisesRegex(Exception, "beam"): + batch.add_kv_cache(cache) + self.assertEqual(batch.row_of(cache), None, "a refused member takes no row") + finally: + if cache is not None: + cache.close() + manager.shutdown() -class TestCacheLevelStorage(unittest.TestCase): - def test_grains_to_slots_rejects_zero_divisors(self) -> None: - invalid_inputs = [ - (1, [0], 16 << 20), - (1, [16 << 20], 0), - ] + def test_a_full_batch_refuses_rather_than_growing(self): + manager = self._manager() + try: + with TemporaryCudaStream([]) as stream: + handle = cast(CudaStream, stream.handle) + batch = Batch( + manager, handle, index_tables=[], max_batch_size=1, max_blocks_per_seq=4 + ) + first = manager.create_kv_cache() + second = manager.create_kv_cache() + self.assertTrue(first.resume(handle)) + self.assertTrue(second.resume(handle)) + batch.add_kv_cache(first) + # Growing would move the device arrays, which a captured graph cannot tolerate. + with self.assertRaisesRegex(Exception, "full"): + batch.add_kv_cache(second) + batch.remove_kv_cache(first) + first.close() + second.close() + finally: + manager.shutdown() - for grains, slot_sizes, granularity in invalid_inputs: - with self.subTest(slot_sizes=slot_sizes, granularity=granularity): - with self.assertRaisesRegex((ValueError, RuntimeError), "must be positive"): - _introspection.grains_to_slots(grains, slot_sizes, granularity) + def test_an_unoccupied_row_reads_as_absent(self): + manager = self._manager() + try: + lg_id = manager.get_layer_group_id(0) + with TemporaryCudaStream([]) as stream: + batch = Batch( + manager, + cast(CudaStream, stream.handle), + index_tables=[], + max_batch_size=2, + max_blocks_per_seq=3, + ) + batch.publish() + table = _introspection.batch_page_table(batch, lg_id) + # Every cell is the sentinel, not slot zero: only the sentinel distinguishes + # an absent block from a real page. + self.assertEqual(table, [self.BAD] * 6) + finally: + manager.shutdown() - def test_grains_to_slots_refines_proportional_lower_bound(self) -> None: - granularity = 16 << 20 - slot_size_list = [16_252_928, 4_063_232] - min_slots = 157 + def test_publish_mirrors_a_members_page_indices(self): + manager = self._manager() + cache = None + batch = None + try: + lg_id = manager.get_layer_group_id(0) + cache = manager.create_kv_cache() + with TemporaryCudaStream([]) as stream: + handle = cast(CudaStream, stream.handle) + batch = Batch( + manager, handle, index_tables=[], max_batch_size=2, max_blocks_per_seq=4 + ) + self.assertTrue(cache.resume(handle)) + batch.add_kv_cache(cache) + cache.resize(256, 0) # 128-token blocks, so two of them + batch.publish() + table = _introspection.batch_page_table(batch, lg_id) + + # The installed buffer is max_blocks_per_seq long, so the member reports + # that whole span: its blocks first, then the sentinel tail. + row = list(cache.get_base_page_indices(lg_id)) + self.assertEqual(len(row), 4) + self.assertEqual(table[:4], row, "row 0 mirrors the member") + self.assertTrue(all(i != self.BAD for i in row[:2]), row) + self.assertEqual(row[2:], [self.BAD] * 2, "past the member's blocks") + # Row 1 is unoccupied and stays absent. + self.assertEqual(table[4:], [self.BAD] * 4) + finally: + if cache is not None: + if batch is not None and batch.row_of(cache) is not None: + batch.remove_kv_cache(cache) + cache.close() + manager.shutdown() - grains = _introspection.grains_for_slots(min_slots, slot_size_list, granularity) - slots, used = _introspection.grains_to_slots(grains, slot_size_list, granularity) + def test_batch_resize_advances_every_member_and_publishes(self): + manager = self._manager() + caches = [] + batch = None + try: + lg_id = manager.get_layer_group_id(0) + with TemporaryCudaStream([]) as stream: + handle = cast(CudaStream, stream.handle) + batch = Batch( + manager, handle, index_tables=[], max_batch_size=3, max_blocks_per_seq=8 + ) + for _ in range(2): + cache = manager.create_kv_cache() + caches.append(cache) + self.assertTrue(cache.resume(handle)) + batch.add_kv_cache(cache) + + results = batch.resize([256, 384, None], [0, 0, None]) + # Row 2 is free, so it reports False rather than raising. + self.assertEqual(results, [True, True, False]) + self.assertEqual(caches[0].capacity, 256) + self.assertEqual(caches[1].capacity, 384) + + # Published as part of the resize, so no separate publish is needed. + table = _introspection.batch_page_table(batch, lg_id) + self.assertEqual(sum(1 for i in table[0:8] if i != self.BAD), 2) + self.assertEqual(sum(1 for i in table[8:16] if i != self.BAD), 3) + self.assertEqual(table[16:24], [self.BAD] * 8, "the free row stays absent") + finally: + for cache in caches: + if batch is not None and batch.row_of(cache) is not None: + batch.remove_kv_cache(cache) + cache.close() + manager.shutdown() - self.assertGreaterEqual(slots, min_slots) - self.assertLessEqual(used, grains) + def test_watermarks_publish_with_the_page_table(self): + """One transfer carries both, so a kernel can never pair mismatched steps.""" + manager = self._manager() + caches = [] + batch = None + try: + lg_id = manager.get_layer_group_id(0) + with TemporaryCudaStream([]) as stream: + handle = cast(CudaStream, stream.handle) + batch = Batch( + manager, handle, index_tables=[], max_batch_size=3, max_blocks_per_seq=8 + ) + # Seeded before joining: the Batch only accepts a history its row already + # had the capacity for, and these start at nothing. + for capacity, history in ((256, 128), (384, 256)): + cache = manager.create_kv_cache() + caches.append(cache) + self.assertTrue(cache.resume(handle)) + self.assertTrue(cache.resize(capacity, history)) + batch.add_kv_cache(cache) + + batch.resize([256, 384, None], [128, 256, None]) + history, capacity = _introspection.batch_watermarks(batch) + self.assertEqual(history, [128, 256, 0]) + self.assertEqual(capacity, [256, 384, 0], "a free row reads as nothing eligible") + + # num_blocks is derived, not published: history_length // tokens_per_block. + table = _introspection.batch_page_table(batch, lg_id) + for row, hist in enumerate(history): + num_blocks = hist // 128 + valid = table[row * 8 : row * 8 + num_blocks] + self.assertTrue( + all(i != self.BAD for i in valid), + f"row {row} must have {num_blocks} valid entries, got {valid}", + ) + finally: + for cache in caches: + if batch is not None and batch.row_of(cache) is not None: + batch.remove_kv_cache(cache) + cache.close() + manager.shutdown() - def test_ratio_to_slot_count_list_preserves_min_slots(self) -> None: - granularity = 16 << 20 - slot_size_lists = [ - [16_252_928, 4_063_232], - [491_520, 126_720, 15_872], - [31_457_280, 7_864_320], - ] - min_slots = [157, 3907, 157] - total_min_grains = sum( - _introspection.grains_for_slots(slots, sizes, granularity) - for slots, sizes in zip(min_slots, slot_size_lists) - ) + def test_a_removed_member_leaves_no_watermark_behind(self): + manager = self._manager() + cache = None + batch = None + try: + with TemporaryCudaStream([]) as stream: + handle = cast(CudaStream, stream.handle) + batch = Batch( + manager, handle, index_tables=[], max_batch_size=1, max_blocks_per_seq=4 + ) + cache = manager.create_kv_cache() + self.assertTrue(cache.resume(handle)) + # Seeded before joining: the Batch only accepts a history its row already + # had the capacity for. + self.assertTrue(cache.resize(256, 128)) + batch.add_kv_cache(cache) + batch.resize([256], [128]) + self.assertEqual(_introspection.batch_watermarks(batch)[0], [128]) + + batch.remove_kv_cache(cache) + batch.publish() + # Otherwise the row would still claim blocks the departed request owned. + history, capacity = _introspection.batch_watermarks(batch) + self.assertEqual((history, capacity), ([0], [0])) + finally: + if cache is not None: + cache.close() + manager.shutdown() - slot_counts = _introspection.ratio_to_slot_count_list( - total_min_grains * granularity, - slot_size_lists, - [0.2, 0.5, 0.3], - granularity, - min_slots, - ) + def test_a_joining_member_publishes_its_own_watermarks(self): + """A row holds whatever its last occupant left until the new one seeds it.""" + manager = self._manager() + cache = None + batch = None + try: + with TemporaryCudaStream([]) as stream: + handle = cast(CudaStream, stream.handle) + batch = Batch( + manager, handle, index_tables=[], max_batch_size=1, max_blocks_per_seq=8 + ) + cache = manager.create_kv_cache() + self.assertTrue(cache.resume(handle)) + # Advanced while unbatched, which is what makes the row's own state stale. + self.assertTrue(cache.resize(256, 128)) - for slot_count, min_slot in zip(slot_counts, min_slots): - self.assertGreaterEqual(slot_count, min_slot) + batch.add_kv_cache(cache) + batch.publish() + self.assertEqual( + _introspection.batch_watermarks(batch), + ([128], [256]), + "the row must describe the member that now occupies it", + ) + finally: + if cache is not None: + if batch is not None and batch.row_of(cache) is not None: + batch.remove_kv_cache(cache) + cache.close() + manager.shutdown() + def test_batch_resize_needs_one_entry_per_row(self): + manager = self._manager() + try: + with TemporaryCudaStream([]) as stream: + handle = cast(CudaStream, stream.handle) + batch = Batch( + manager, handle, index_tables=[], max_batch_size=2, max_blocks_per_seq=4 + ) + # Indexed by row, so a short list would silently misalign members. + with self.assertRaisesRegex(Exception, "per row"): + batch.resize([128], [0]) + finally: + manager.shutdown() -def create_config( - tokens_per_block: int, - gpu_quota: int, - host_quota: int, - disk_quota: int, - num_layers: int, - window_size: SlidingWindowSize, - sink_tokens: int, - kv_buf_size: int = 8192, - block_quant_buf_size: int | None = None, -) -> KVCacheManagerConfig: - layer_buffers = [ - BufferConfig(role=Role.KEY, size=kv_buf_size), - BufferConfig(role=Role.VALUE, size=kv_buf_size), - ] - if block_quant_buf_size is not None: - layer_buffers.extend( - [ - BufferConfig(role=Role.KEY_BLOCK_QUANT, size=block_quant_buf_size), - BufferConfig(role=Role.VALUE_BLOCK_QUANT, size=block_quant_buf_size), - ] - ) - disk_path_candidates = ["/workspace/", "/tmp/nvidia-mps/", "/tmp"] - disk_path = next(p for p in disk_path_candidates if os.path.exists(p)) - assert gpu_quota > 0 - cache_tiers = [ - GpuCacheTierConfig(quota=gpu_quota), - HostCacheTierConfig(quota=host_quota), - DiskCacheTierConfig(quota=disk_quota, path=disk_path), - ] - cache_tiers = [t for t in cache_tiers if t.quota > 0] - return KVCacheManagerConfig( - tokens_per_block=tokens_per_block, - cache_tiers=[t for t in cache_tiers if t.quota > 0], - layers=[ - AttentionLayerConfig( - layer_id=layer_id, - buffers=deepcopy(layer_buffers), - sliding_window_size=window_size if layer_id % 2 == 0 else None, - num_sink_tokens=sink_tokens if layer_id % 2 == 0 else None, + def test_history_may_not_outrun_the_capacity_the_row_already_had(self): + """Capacity granted now makes room for tokens written later, not already written.""" + manager = self._manager() + cache = None + batch = None + try: + with TemporaryCudaStream([]) as stream: + handle = cast(CudaStream, stream.handle) + batch = Batch( + manager, handle, index_tables=[], max_batch_size=1, max_blocks_per_seq=8 + ) + cache = manager.create_kv_cache() + self.assertTrue(cache.resume(handle)) + # Seeded before joining: the Batch only accepts a history its row already + # had the capacity for. + self.assertTrue(cache.resize(256, 128)) + batch.add_kv_cache(cache) + batch.resize([256], [128]) + + # Capacity 256 now, so token 256 has no space yet and cannot be history, + # however much capacity this same call would add. + with self.assertRaisesRegex(Exception, "capacity the row already had"): + batch.resize([512], [384]) + # Refused before anything moved. + self.assertEqual(_introspection.batch_watermarks(batch), ([128], [256])) + + with self.assertRaisesRegex(Exception, "cannot be decreased"): + batch.resize([256], [64]) + self.assertEqual(_introspection.batch_watermarks(batch), ([128], [256])) + finally: + if cache is not None: + if batch is not None and batch.row_of(cache) is not None: + batch.remove_kv_cache(cache) + cache.close() + manager.shutdown() + + def test_a_row_whose_growth_fails_still_advances_its_history(self): + """The advance needs no slots, and is what frees the cells the growth wanted.""" + tpb, page = 128, 4096 + # The pool bottoms out at 512 blocks whatever the quota, so two rows of 256 fill + # it and a third-party grow to 512 has nowhere to come from. Measured, not derived + # from the quota: a small quota does not shrink the pool below that floor. + blocks, pool = 256, 512 + manager = KVCacheManager( + KVCacheManagerConfig( + tokens_per_block=tpb, + cache_tiers=[GpuCacheTierConfig(pool * page)], + layers=[AttentionLayerConfig(0, [BufferConfig("key", page)])], ) - for layer_id in typed_range(LayerId(num_layers)) - ], - ) + ) + caches = [] + batch = None + try: + with TemporaryCudaStream([]) as stream: + handle = cast(CudaStream, stream.handle) + batch = Batch( + manager, handle, index_tables=[], max_batch_size=2, max_blocks_per_seq=pool + ) + for _ in range(2): + cache = manager.create_kv_cache() + caches.append(cache) + self.assertTrue(cache.resume(handle)) + batch.add_kv_cache(cache) + + each = blocks * tpb + self.assertEqual(list(batch.resize([each, each], [0, 0])), [1, 1]) + + # The two rows now hold the whole pool, so row 1's grow has nowhere to come + # from and rolls back. Row 0 only advances history, which never allocates. + results = list(batch.resize([each, pool * tpb], [each, each])) + self.assertEqual(results[0], 1) + self.assertEqual( + results[1], 0, "the grow must actually fail or this proves nothing" + ) + history, capacity = _introspection.batch_watermarks(batch) + self.assertEqual(history, [each, each], "both advanced; neither needed a slot") + self.assertEqual(capacity, [each, each], "row 1's growth rolled itself back") + finally: + for cache in caches: + if batch is not None and batch.row_of(cache) is not None: + batch.remove_kv_cache(cache) + cache.close() + manager.shutdown() -class TestSparseBufferConfig(unittest.TestCase): - def test_flag_defaults_and_copy(self): - self.assertFalse(BufferConfig("key", 4096).is_sparse) - buffer = BufferConfig("key", 4096, 2, is_sparse=True) - copied = deepcopy(buffer) - self.assertTrue(copied.is_sparse) - self.assertEqual(copied.tokens_per_block_override, 2) - buffer.is_sparse = False - self.assertTrue(copied.is_sparse) + def test_wait_published_is_a_noop_when_nothing_is_in_flight(self): + manager = self._manager() + try: + batch = Batch( + manager, NULL_STREAM, index_tables=[], max_batch_size=1, max_blocks_per_seq=4 + ) + # Safe to call at the top of a step whether or not a publish happened. + batch.wait_published() + batch.wait_published() + finally: + manager.shutdown() - def test_requires_host_at_level_one(self): - layer = AttentionLayerConfig(0, [BufferConfig("key", 4096, is_sparse=True)]) - gpu = GpuCacheTierConfig(4 << 20) - for tiers in ([gpu], [gpu, gpu], [gpu, DiskCacheTierConfig(4 << 20, "/tmp")]): - with self.subTest(tiers=tiers): - with self.assertRaisesRegex(ValueError, "level 1.*HOST_MEM"): - KVCacheManagerConfig(tokens_per_block=4, cache_tiers=tiers, layers=[layer]) - config = KVCacheManagerConfig( - tokens_per_block=4, - cache_tiers=[gpu, HostCacheTierConfig(4 << 20)], - layers=[layer], - ) - self.assertTrue(config.layers[0].buffers[0].is_sparse) + def test_a_republish_after_waiting_shows_the_new_indices(self): + """A republish after waiting must carry the change. - def test_rejects_mixed_layer_buffers(self): - layer = AttentionLayerConfig( - 0, [BufferConfig("key", 4096, is_sparse=True), BufferConfig("value", 4096)] - ) - with self.assertRaisesRegex(ValueError, "cannot share an attention layer lifecycle"): - KVCacheManagerConfig( - tokens_per_block=4, - cache_tiers=[GpuCacheTierConfig(4 << 20), HostCacheTierConfig(4 << 20)], - layers=[layer], - ) + publish() copies asynchronously from pinned staging, so the host has to wait + before touching a member again. + """ + manager = self._manager() + cache = None + batch = None + try: + lg_id = manager.get_layer_group_id(0) + cache = manager.create_kv_cache() + with TemporaryCudaStream([]) as stream: + handle = cast(CudaStream, stream.handle) + batch = Batch( + manager, handle, index_tables=[], max_batch_size=1, max_blocks_per_seq=8 + ) + self.assertTrue(cache.resume(handle)) + batch.add_kv_cache(cache) + cache.resize(256, 0) + batch.publish() + first = _introspection.batch_page_table(batch, lg_id) + self.assertEqual(sum(1 for i in first if i != self.BAD), 2) + + # The contract: wait before the resize, which writes the staging row the + # copy above is still reading. + batch.wait_published() + cache.resize(512, 0) + batch.publish() + second = _introspection.batch_page_table(batch, lg_id) + self.assertEqual(sum(1 for i in second if i != self.BAD), 4, second) + finally: + if cache is not None: + if batch is not None and batch.row_of(cache) is not None: + batch.remove_kv_cache(cache) + cache.close() + manager.shutdown() + + def test_a_removed_members_row_is_cleared(self): + manager = self._manager() + try: + lg_id = manager.get_layer_group_id(0) + cache = manager.create_kv_cache() + with TemporaryCudaStream([]) as stream: + handle = cast(CudaStream, stream.handle) + batch = Batch( + manager, handle, index_tables=[], max_batch_size=1, max_blocks_per_seq=4 + ) + self.assertTrue(cache.resume(handle)) + batch.add_kv_cache(cache) + cache.resize(256, 0) + batch.remove_kv_cache(cache) + batch.publish() + # Otherwise the row would still name pages the departed request owns. + self.assertEqual(_introspection.batch_page_table(batch, lg_id), [self.BAD] * 4) + finally: + cache.close() + manager.shutdown() class TestKVCacheManagerV2(unittest.TestCase): @@ -704,7 +6071,8 @@ def test_offload_reload_preserves_priority_and_partial_reuse(self) -> None: pressure.close() self.assertTrue(cached.resume(self.stream.handle)) self.assertEqual(cached.get_page_priorities(0), [80]) - self.assertEqual(cached.get_page_storage_snapshot(0).cache_levels, [GPU_LEVEL]) + # The resume brought the page back: active on the hot level, gone from host. + self.assertEqual(_introspection.active_page_stats(cached)[0], [1, 0]) self.engine.execute([Step(cached, [], tokens)], self.stream.handle) self.stream.synchronize() for config in (None, self._config(10)): @@ -885,212 +6253,6 @@ def test_side_stream_waits_for_a_busy_recycled_slot(self) -> None: loader.close() execution_stream.synchronize() - def test_batch_publishes_sparse_gpu_metadata(self) -> None: - import torch - - self.manager = KVCacheManager( - KVCacheManagerConfig( - tokens_per_block=4, - cache_tiers=[GpuCacheTierConfig(4 << 20), HostCacheTierConfig(4 << 20)], - layers=[ - AttentionLayerConfig(0, [BufferConfig("key", 4096, is_sparse=True)]), - AttentionLayerConfig(1, [BufferConfig("key", 4096)]), - ], - ) - ) - batch = Batch(self.manager, max_rows=3, max_blocks=4) - cache = self.manager.create_kv_cache() - sparse_group = self.manager.get_layer_group_id(0) - dense_group = self.manager.get_layer_group_id(1) - stream = torch.cuda.Stream() - with torch.cuda.stream(stream): - try: - self.assertTrue(cache.resume(stream.cuda_stream)) - self.assertTrue(cache.resize(12, 6)) - self.assertEqual(batch.add(cache, row=2), 2) - with self.assertRaises(LogicError): - torch.from_dlpack(batch.page_table(sparse_group)) - self.assertEqual(batch.publish(stream.cuda_stream), [0, 1, 2]) - table = torch.from_dlpack(batch.page_table(sparse_group)) - counts = torch.from_dlpack(batch.num_blocks(sparse_group)) - dense_counts = torch.from_dlpack(batch.num_blocks(dense_group)) - self.assertEqual(tuple(table.shape), (3, 1, 4)) - self.assertEqual(table.dtype, torch.int32) - self.assertEqual(counts.cpu().tolist(), [[0], [0], [0]]) - address = table.data_ptr() - batch.record_read(stream.cuda_stream) - - old_version = cache.page_storage_version - self.assertTrue(cache.enter_decode()) - self.assertFalse(cache.acknowledge_page_storage(old_version)) - self.assertEqual(batch.dirty_rows, [2]) - self.assertEqual(batch.publish(stream.cuda_stream), [2]) - batch.wait_ready(stream.cuda_stream) - expected = [ - [[BAD_PAGE_INDEX] * 4], - [[BAD_PAGE_INDEX] * 4], - [list(cache.get_base_page_indices(sparse_group)) + [BAD_PAGE_INDEX]], - ] - self.assertEqual(table.cpu().tolist(), expected) - self.assertEqual(counts.cpu().tolist(), [[0], [0], [1]]) - self.assertEqual(dense_counts.cpu().tolist(), [[0], [0], [0]]) - self.assertFalse(cache.page_storage_dirty) - self.assertEqual(batch.publish(stream.cuda_stream), []) - batch.record_read(stream.cuda_stream) - - cache.suspend() - self.assertEqual(batch.publish(stream.cuda_stream), [2]) - self.assertEqual(counts.cpu().tolist(), [[0], [0], [0]]) - self.assertEqual(table.cpu().tolist(), [[[BAD_PAGE_INDEX] * 4]] * 3) - self.assertTrue(cache.resume()) - self.assertEqual(batch.publish(stream.cuda_stream), [2]) - self.assertEqual(counts.cpu().tolist(), [[0], [0], [1]]) - self.assertEqual(table.data_ptr(), address) - batch.record_read(stream.cuda_stream) - cache.close() - self.assertIsNone(cache.page_storage_row) - self.assertEqual(batch.publish(stream.cuda_stream), [2]) - self.assertEqual(table.cpu().tolist(), [[[BAD_PAGE_INDEX] * 4]] * 3) - self.assertEqual(counts.cpu().tolist(), [[0], [0], [0]]) - batch.close() - self.assertEqual(table.data_ptr(), address) - self.assertEqual(table.cpu().tolist(), [[[BAD_PAGE_INDEX] * 4]] * 3) - with self.assertRaises(LogicError): - torch.from_dlpack(batch.page_table(sparse_group)) - finally: - stream.synchronize() - batch.close() - cache.close() - - @parameterized.expand(["decode", "suspend", "close"]) - def test_shared_prefill_defers_sparse_offload(self, release: str) -> None: - import torch - - self.manager = KVCacheManager( - KVCacheManagerConfig( - tokens_per_block=4, - cache_tiers=[GpuCacheTierConfig(4 << 20), HostCacheTierConfig(4 << 20)], - layers=[AttentionLayerConfig(0, [BufferConfig("key", 4096, is_sparse=True)])], - ) - ) - decoder = self.manager.create_kv_cache() - prefill = None - batch = Batch(self.manager, max_rows=1, max_blocks=3) - group = self.manager.get_layer_group_id(0) - stream = torch.cuda.Stream() - with torch.cuda.stream(stream): - try: - self.assertTrue(decoder.resume(stream.cuda_stream)) - self.assertTrue(decoder.resize(12, 8)) - decoder.commit([0, 1, 2, 3]) - prefill = self.manager.create_kv_cache(ReuseScope(), [0, 1, 2, 3]) - self.assertTrue(prefill.resume(stream.cuda_stream)) - self.assertEqual( - decoder.get_base_page_indices(group)[0], prefill.get_base_page_indices(group)[0] - ) - self.assertTrue(decoder.enter_decode()) - self.assertTrue(decoder.is_decoding) - deferred = decoder.get_page_storage_snapshot(group) - self.assertEqual(deferred.cache_levels, [0, 1, 0]) - self.assertEqual(deferred.eligible_history_blocks, 0) - batch.add(decoder) - self.assertEqual(batch.publish(stream.cuda_stream), [0]) - table = torch.from_dlpack(batch.page_table(group)) - counts = torch.from_dlpack(batch.num_blocks(group)) - address = table.data_ptr() - self.assertEqual(table.cpu().tolist(), [[deferred.base_page_indices]]) - self.assertEqual(counts.cpu().tolist(), [[0]]) - batch.record_read(stream.cuda_stream) - self.assertEqual(batch.publish(stream.cuda_stream), []) - - if release == "decode": - self.assertTrue(prefill.enter_decode()) - elif release == "suspend": - prefill.suspend() - else: - prefill.close() - self.assertEqual(batch.publish(stream.cuda_stream), [0]) - batch.wait_ready(stream.cuda_stream) - self.assertEqual(decoder.history_length, 8) - completed = decoder.get_page_storage_snapshot(group) - self.assertEqual(completed.cache_levels, [1, 1, 0]) - self.assertEqual(completed.eligible_history_blocks, 2) - self.assertGreater(completed.version, deferred.version) - self.assertEqual(table.data_ptr(), address) - self.assertEqual(table.cpu().tolist(), [[completed.base_page_indices]]) - self.assertEqual(counts.cpu().tolist(), [[2]]) - batch.record_read(stream.cuda_stream) - self.assertEqual(batch.publish(stream.cuda_stream), []) - finally: - stream.synchronize() - batch.close() - decoder.close() - if prefill is not None: - prefill.close() - - def test_sparse_page_storage_metadata(self) -> None: - self.manager = KVCacheManager( - KVCacheManagerConfig( - tokens_per_block=4, - cache_tiers=[GpuCacheTierConfig(4 << 20), HostCacheTierConfig(4 << 20)], - layers=[ - AttentionLayerConfig(0, [BufferConfig("key", 4096, is_sparse=True)]), - AttentionLayerConfig(1, [BufferConfig("key", 4096)]), - ], - ) - ) - self.assertTrue(self.manager.is_sparse(0, "key")) - self.assertFalse(self.manager.is_sparse(1, "key")) - sparse_group = self.manager.get_layer_group_id(0) - dense_group = self.manager.get_layer_group_id(1) - cache = self.manager.create_kv_cache() - with TemporaryCudaStream([]) as stream_holder: - stream = cast(CudaStream, stream_holder.handle) - try: - self.assertTrue(cache.resume(stream)) - self.assertTrue(cache.resize(12, 6)) - cache.bind_page_storage_row(3) - prefill = cache.get_page_storage_snapshot(sparse_group) - self.assertEqual(prefill.row, 3) - self.assertEqual(prefill.eligible_history_blocks, 0) - self.assertEqual(prefill.cache_levels, [0, 0, 0]) - self.assertTrue(cache.acknowledge_page_storage(prefill.version)) - self.assertFalse(cache.page_storage_dirty) - self.assertTrue(cache.enter_decode()) - self.assertTrue(cache.page_storage_dirty) - self.assertFalse(cache.acknowledge_page_storage(prefill.version)) - decode = cache.get_page_storage_snapshot(sparse_group) - self.assertEqual(decode.version, cache.page_storage_version) - self.assertEqual(decode.eligible_history_blocks, 1) - self.assertEqual(decode.cache_levels, [1, 0, 0]) - self.assertEqual( - decode.base_page_indices, list(cache.get_base_page_indices(sparse_group)) - ) - self.assertEqual( - cache.get_page_storage_snapshot(dense_group).eligible_history_blocks, 0 - ) - copied_indices = decode.base_page_indices - copied_indices[0] = BAD_PAGE_INDEX - self.assertNotEqual(decode.base_page_indices[0], BAD_PAGE_INDEX) - decode.wait_ready(stream) - cache.record_page_storage_read(stream) - self.assertTrue(cache.acknowledge_page_storage(decode.version)) - cache.commit([0, 1, 2, 3]) - self.assertTrue(cache.page_storage_dirty) - cache.suspend() - suspended = cache.get_page_storage_snapshot(sparse_group) - self.assertEqual(suspended.base_page_indices, [BAD_PAGE_INDEX] * 3) - self.assertEqual(suspended.cache_levels, [None] * 3) - self.assertEqual(suspended.eligible_history_blocks, 0) - self.assertEqual(cache.page_storage_row, 3) - self.assertTrue(cache.resume()) - cache.bind_page_storage_row(None) - self.assertIsNone(cache.page_storage_row) - finally: - cache.close() - stream_holder.take_finish_event().synchronize() - self.assertEqual(cache.get_page_storage_snapshot(sparse_group).base_page_indices, []) - class Request(NamedTuple): id: int kv_cache: _KVCache @@ -4330,7 +9492,13 @@ def test_invalid_max_util_for_resume(self, _name: str, max_util_for_resume: floa with self.assertRaisesRegex((ValueError, RuntimeError), "max_util_for_resume must be in"): KVCacheManager(cfg) - def test_ratio_slot_count_rounding_matches_python(self): + def test_a_slot_one_byte_under_a_grain_still_costs_a_whole_grain(self): + """Rounding is per pool and per grain, so a slightly smaller slot buys nothing. + + Both groups cost one grain a slot -- `grain - 1` rounds up -- so five grains buy + five slots however they are split. Which group takes the odd one is decided by the + ratio, and these two tie exactly, so this does not pin that. + """ grain = 2 << 20 cfg = KVCacheManagerConfig( tokens_per_block=self.TOKENS_PER_BLOCK, @@ -4361,8 +9529,10 @@ def stat_slot_sizes(stat) -> list[int]: slots_by_size = { tuple(stat_slot_sizes(stat)): stat.total for stat in manager.get_storage_statistics() } - self.assertEqual(slots_by_size[(grain - 1,)], 2) - self.assertEqual(slots_by_size[(grain,)], 3) + small = slots_by_size[(grain - 1,)] + full = slots_by_size[(grain,)] + self.assertEqual(small + full, 5, "five grains buy five slots, not one more") + self.assertEqual({small, full}, {2, 3}, "and they are split as evenly as five allows") manager.shutdown() @parameterized.expand([(0,), (64,), (50,), (256,)]) @@ -4807,7 +9977,7 @@ def test_request_scratch_toggle_for_two_round_inference(self): self.assertTrue(kv.resume(stream)) self.assertTrue(kv.resize(len(prompt))) self.assertTrue(kv.enable_swa_scratch_reuse) - self.assertIsNotNone(kv.get_scratch_desc(lg_id)) + self.assertIsNotNone(_scratch_desc(kv, lg_id)) self.assertTrue(kv.has_scratch_slots) with self.assertRaisesRegex(ValueError, "scratch blocks are needed"): kv.enable_swa_scratch_reuse = False @@ -4815,14 +9985,14 @@ def test_request_scratch_toggle_for_two_round_inference(self): self.engine.execute([Step(kv, prompt, history)], stream) kv.commit(prompt) history.extend(prompt) - self.assertIsNone(kv.get_scratch_desc(lg_id)) + self.assertIsNone(_scratch_desc(kv, lg_id)) self.assertFalse(kv.has_scratch_slots) kv.enable_swa_scratch_reuse = False self.assertFalse(kv.enable_swa_scratch_reuse) kv.capacity = len(history) + 1 self.assertFalse(kv.has_scratch_slots) - self.assertIsNone(kv.get_scratch_desc(lg_id)) + self.assertIsNone(_scratch_desc(kv, lg_id)) self.engine.execute([Step(kv, [decode_token], history)], stream) kv.commit([decode_token]) @@ -4832,13 +10002,13 @@ def test_request_scratch_toggle_for_two_round_inference(self): kv.enable_swa_scratch_reuse = True self.assertTrue(kv.enable_swa_scratch_reuse) self.assertTrue(kv.resize(len(history) + len(second_prompt), len(history))) - self.assertIsNotNone(kv.get_scratch_desc(lg_id)) + self.assertIsNotNone(_scratch_desc(kv, lg_id)) self.assertTrue(kv.has_scratch_slots) self.engine.execute([Step(kv, second_prompt, history)], stream) kv.commit(second_prompt) history.extend(second_prompt) - self.assertIsNone(kv.get_scratch_desc(lg_id)) + self.assertIsNone(_scratch_desc(kv, lg_id)) self.assertFalse(kv.has_scratch_slots) self.engine.execute([Step(kv, [], history)], stream) kv.stop_committing() @@ -4886,7 +10056,7 @@ def test_scratch_slot_count(self): self.assertEqual(len(layer_groups), 1) lg_id = LayerGroupId(0) - scratch_desc = kv.get_scratch_desc(lg_id) + scratch_desc = _scratch_desc(kv, lg_id) assert scratch_desc is not None self.assertIsNotNone(scratch_desc) num_blocks = div_up(prompt_len, tokens_per_block) # 32 @@ -4904,7 +10074,7 @@ def test_scratch_slot_count(self): # Verify much less than 32 total slots self.assertLess(expected_total, num_blocks) - # Scratch slot count in ScratchDesc matches expected + # Scratch slot count matches expected self.assertEqual(len(scratch_desc.slot_ids), expected_scratch_slots) # Check base page indices: scratch blocks have BAD_PAGE_INDEX, normal have valid @@ -4965,7 +10135,7 @@ def test_scratch_slot_count(self): @parameterized.expand([(0, 7), (64, 5)]) def test_scratch_shared_slot_ids(self, rewind_len: int, expected_scratch_blocks: int): - """Verify that scratch blocks share coalesced slot IDs via ScratchDesc.""" + """Verify that scratch blocks share coalesced slot IDs.""" # 8 layers, window=32, tokens_per_block=32, prompt=256 # num_sub_pages = 8 (all layers in one group) # rewind_len=0: blocks 0-6 are scratch, block 7 is in-window. @@ -4990,7 +10160,7 @@ def test_scratch_shared_slot_ids(self, rewind_len: int, expected_scratch_blocks: self.assertTrue(success) lg_id = LayerGroupId(0) - scratch_desc = kv.get_scratch_desc(lg_id) + scratch_desc = _scratch_desc(kv, lg_id) assert scratch_desc is not None self.assertIsNotNone(scratch_desc) @@ -5017,7 +10187,9 @@ def test_scratch_shared_slot_ids(self, rewind_len: int, expected_scratch_blocks: # Verify PageIndexConverter.convert_all produces correct per-layer indices layer_id = LayerId(0) converter = self.manager.get_page_index_converter(layer_id, DataRole("key")) - page_indices = converter(indices, PageIndexMode.PER_LAYER, scratch_desc) + page_indices = converter( + indices, PageIndexMode.PER_LAYER, scratch_desc.range, scratch_desc.slot_ids + ) # All scratch blocks should produce valid (non-BAD) page indices for i in range(scratch_desc.range.beg, scratch_desc.range.end): self.assertNotEqual( @@ -5104,7 +10276,7 @@ def test_scratch_chunk_size_variation( lg_id_0 = LayerGroupId(0) lg_id_1 = LayerGroupId(1) - scratch_desc_0 = kv.get_scratch_desc(lg_id_0) + scratch_desc_0 = _scratch_desc(kv, lg_id_0) self.assertIsNotNone(scratch_desc_0) self.assertEqual(scratch_desc_0.range.beg, chunk1_scratch_range[0]) self.assertEqual(scratch_desc_0.range.end, chunk1_scratch_range[1]) @@ -5114,7 +10286,7 @@ def test_scratch_chunk_size_variation( ) # Layer 1 should have 0 scratch blocks - scratch_desc_1 = kv.get_scratch_desc(lg_id_1) + scratch_desc_1 = _scratch_desc(kv, lg_id_1) self.assertIsNone(scratch_desc_1) self.engine.execute([Step(kv, prompt1, [])], stream) @@ -5130,7 +10302,7 @@ def test_scratch_chunk_size_variation( success = kv.resize(352, 256) self.assertTrue(success) - scratch_desc_0 = kv.get_scratch_desc(lg_id_0) + scratch_desc_0 = _scratch_desc(kv, lg_id_0) if chunk2_scratch_range is None: self.assertIsNone(scratch_desc_0) else: @@ -5139,7 +10311,7 @@ def test_scratch_chunk_size_variation( self.assertEqual(scratch_desc_0.range.end, chunk2_scratch_range[1]) # Layer 1 should still have 0 scratch blocks - scratch_desc_1 = kv.get_scratch_desc(lg_id_1) + scratch_desc_1 = _scratch_desc(kv, lg_id_1) self.assertIsNone(scratch_desc_1) self.engine.execute([Step(kv, prompt2, prompt1)], stream) From e360c15147c04d1da56d9367b39e961fc7ce86d6 Mon Sep 17 00:00:00 2001 From: Yao Yao Date: Tue, 27 Oct 2026 19:37:39 +0000 Subject: [PATCH 02/12] [None][chore] KVCM2: keep ScratchDesc bound while consumers still read it KvCache reports a life cycle's scratch range and its slots separately, through getSwaScratchRange() and getScratchSlots(). Consumers in this tree still expect the single object the manager used to return, so the binding layer pairs the two accessors back into a ScratchDesc with `range` and `slot_ids`. This lives only in nanobind: the core keeps the two accessors as its API. It is a work-around that lets the KVCM2 core change stand alone, and should be removed once the full sparse attention support is merged and its consumers read the two accessors directly. Signed-off-by: Yao Yao --- .../batch_manager/kvCacheManagerV2.cpp | 39 +++++++++++++++++++ 1 file changed, 39 insertions(+) diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp index 27f355afafbd..005dcdb6d00d 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp @@ -69,6 +69,18 @@ namespace tensorrt_llm::nanobind::batch_manager namespace { +//! A scratch range paired with the slots covering it. +//! +//! KvCache reports the two separately, through getSwaScratchRange() and getScratchSlots(). +//! This pairs them for Python consumers that still expect one object, and exists only so +//! this change stands alone; remove it, and the get_scratch_desc() binding, once the full +//! sparse attention support lands and those consumers read the two accessors directly. +struct ScratchDescCompat +{ + kv::HalfOpenRange range; + std::vector slotIds; +}; + // Exposed via introspection sub-module for tests. class TestPaddingColdPageCodec final : public kv::IKvCacheColdPageCodec { @@ -1465,6 +1477,17 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) "__contains__", [](BlockRange const& self, int item) { return self.contains(kv::BlockOrdinal{item}); }, nb::arg("item")); + // A scratch range and its slots as one object, which is how KvCache used to report + // them. Kept for consumers that still read `.range` and `.slot_ids`; the manager + // itself reports the two separately through get_swa_scratch_range() and + // get_scratch_slots(). Remove this class and get_scratch_desc() once the full sparse + // attention support lands and those consumers move to the two accessors -- they exist + // only so this change stands alone. + nb::class_(m, "ScratchDesc") + .def_ro("range", &ScratchDescCompat::range) + .def_ro("slot_ids", &ScratchDescCompat::slotIds) + .def("__bool__", [](ScratchDescCompat const& self) { return static_cast(self.range); }); + nb::enum_(m, "PageIndexMode") .value("SHARED", kv::PageIndexMode::SHARED) .value("PER_LAYER", kv::PageIndexMode::PER_LAYER); @@ -2043,6 +2066,22 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) }, nb::arg("layer_group_id"), nb::call_guard(), "Page indices of this life cycle's scratch slots, in slot order.") + .def( + "get_scratch_desc", + [](kv::KvCache const& self, int layerGroupId) -> std::optional + { + auto const lgId = kv::LayerGroupId{layerGroupId}; + auto const range = self.getSwaScratchRange(lgId); + if (!range) + { + return std::nullopt; + } + auto const slots = self.getScratchSlots(lgId); + return ScratchDescCompat{range, std::vector(slots.begin(), slots.end())}; + }, + nb::arg("layer_group_id"), nb::call_guard(), + "get_swa_scratch_range() and get_scratch_slots() as one object, or None when this life " + "cycle serves nothing from scratch. Remove with the ScratchDesc class.") .def_prop_ro("has_scratch_slots", &kv::KvCache::hasScratchSlots) .def_prop_rw("enable_swa_scratch_reuse", &kv::KvCache::isSwaScratchReuseEnabled, [](kv::KvCache& self, bool enable) From 10eb17fb93ba6dbfba572f4cfab0665041128a3e Mon Sep 17 00:00:00 2001 From: Yao Yao Date: Tue, 27 Oct 2026 19:42:34 +0000 Subject: [PATCH 03/12] [None][chore] KVCM2: admit decode through the binding while callers still ask KvCache no longer marks a decode transition. What enter_decode() did beyond setting that flag was offload sparse history, and the sparse fetch path owns that now; a cache with no sparse life cycle had nothing to offload, so an active request was always admitted. The binding keeps answering that question so the callers in this tree work unchanged. Like the ScratchDesc pairing, this lives only in nanobind and is a work-around that lets the KVCM2 core change stand alone. Remove it once the full sparse attention support is merged and its callers stop asking. Signed-off-by: Yao Yao --- .../nanobind/batch_manager/kvCacheManagerV2.cpp | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp index 005dcdb6d00d..61f562778fc9 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp @@ -2102,6 +2102,23 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) nb::arg("mode"), nb::arg("layer_group_id"), "Per life cycle, which is where the scratch slots live.") .def_prop_ro("status", [](kv::KvCache const& kvc) { return kvc.status(); }) .def_prop_ro("is_active", &kv::KvCache::isActive) + // Admission into decode, which the manager no longer marks: beyond setting a flag, + // this offloaded sparse history, and the sparse fetch path owns that now. A cache + // with no sparse life cycle had nothing to offload, so always admitting is what it + // already did. Remove once the full sparse attention support is merged and its + // callers drop this; it exists only so this change stands alone. + .def( + "enter_decode", + [](kv::KvCache const& self) + { + if (!self.isActive()) + { + throw kv::LogicError("Decode admission requires an active request"); + } + return true; + }, + nb::call_guard(), + "Admit this request into decode. Always true for an active request.") .def_prop_ro("num_blocks", [](kv::KvCache const& self) { return self.numBlocks().value(); }) .def_prop_ro( "batch_row", From c9d7652d85688865b09c0389fa9b70963c64b5f3 Mon Sep 17 00:00:00 2001 From: Yao Yao Date: Tue, 27 Oct 2026 19:57:40 +0000 Subject: [PATCH 04/12] [None][chore] KVCM2: leave sparse metadata publication to the fetch path The manager built a Batch to publish sparse page-storage metadata through. That Batch is now the sparse fetch path's, built where the fetch happens, and its constructor takes what the fetch needs rather than what the manager had. Every use of `sparse_metadata_batch` is already guarded on it being set, so not building it leaves those paths inert and the rest of the manager unchanged. Like the two binding shims, this is a work-around that lets the KVCM2 core change stand alone. Restore the construction once the full sparse attention support is merged and brings the fetch path with it. Signed-off-by: Yao Yao --- .../_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index a963738f9ff8..b7fb2642d89f 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py @@ -2583,10 +2583,12 @@ def create_cold_page_codec(cache_config: object) -> Optional[object]: } ) ) - if self._sparse_layer_group_ids: - self.sparse_metadata_batch = Batch( - self.impl, index_mapper_capacity, self.max_blocks_per_seq, self.max_beam_width - ) + # Metadata publication belongs to the sparse fetch path, which builds the Batch + # it publishes through; this manager does not. Every use of + # `sparse_metadata_batch` is guarded on it being set, so leaving it unset keeps + # those paths inert. Restore this with the full sparse attention support, which + # brings that fetch path with it -- it is dropped here only so the KVCM2 core + # change stands alone. self._log_kv_cache_pool_lifecycle_mapping() self._reserve_guard_page() From 5f990c8eb7a836ae4cbb088d0bb22248a9d7a41d Mon Sep 17 00:00:00 2001 From: Yao Yao Date: Wed, 28 Oct 2026 00:59:00 +0000 Subject: [PATCH 05/12] [None][chore] KVCM2: accept is_decoding on resume while callers still pass it Resuming no longer distinguishes a decoding request. What that selected was sparse history offload, which the sparse fetch path owns now, so the parameter has nothing left to choose and the binding accepts and ignores it. Like the other binding shims, this is a work-around that lets the KVCM2 core change stand alone. Drop the parameter once the full sparse attention support is merged and its callers stop passing it. Signed-off-by: Yao Yao --- .../nanobind/batch_manager/kvCacheManagerV2.cpp | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp index 61f562778fc9..89fed2ed61d8 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp @@ -1949,8 +1949,15 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) nb::class_(m, "_KVCache") .def( "resume", - [](kv::KvCache& self, nb::object stream, std::optional> groups) - { + // `is_decoding` is accepted and ignored: resuming no longer distinguishes a + // decoding request, since what that selected -- sparse history offload -- is the + // sparse fetch path's now. Drop the parameter once the full sparse attention + // support is merged and its callers stop passing it; it is here only so this + // change stands alone. + [](kv::KvCache& self, nb::object stream, std::optional isDecoding, + std::optional> groups) + { + (void) isDecoding; std::optional optStream; if (!stream.is_none()) optStream = reinterpret_cast(nb::cast(stream)); @@ -1966,7 +1973,8 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) nb::gil_scoped_release rel; return self.resume(optStream, std::move(selected)); }, - nb::arg("cuda_stream") = nb::none(), nb::arg("optional_reuse_groups") = nb::none()) + nb::arg("cuda_stream") = nb::none(), nb::arg("is_decoding") = nb::none(), + nb::arg("optional_reuse_groups") = nb::none()) .def( "get_page_priorities", [](kv::KvCache const& self, int layerGroupId, int beamIndex) From a6f70282a2143bae87bb2c889baafd810d65f570 Mon Sep 17 00:00:00 2001 From: Yao Yao Date: Wed, 28 Oct 2026 07:37:00 +0000 Subject: [PATCH 06/12] [None][fix] KVCM2: link tllm_kv_cache against common_environment Co-Authored-By: Claude Opus 5.5 Signed-off-by: Yao Yao --- cpp/tensorrt_llm/batch_manager/CMakeLists.txt | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/cpp/tensorrt_llm/batch_manager/CMakeLists.txt b/cpp/tensorrt_llm/batch_manager/CMakeLists.txt index 0f033e83ec3b..c14e7c9147f0 100644 --- a/cpp/tensorrt_llm/batch_manager/CMakeLists.txt +++ b/cpp/tensorrt_llm/batch_manager/CMakeLists.txt @@ -86,8 +86,10 @@ add_library(tllm::kv_cache ALIAS tllm_kv_cache) target_include_directories( tllm_kv_cache PUBLIC ${KV_CACHE_MANAGER_V2_INCLUDE_DIR} ${SHA256_DIR}) target_link_libraries( - tllm_kv_cache PUBLIC tllm::kv_cache_host_memory tllm::common_cuda - tllm::executor_types CUDA::cuda_driver) + tllm_kv_cache + PUBLIC tllm::kv_cache_host_memory tllm::common_cuda tllm::executor_types + CUDA::cuda_driver + PRIVATE tllm::common_environment) add_library( tllm_kv_cache_legacy STATIC From ed6ab6f6b802fdfa9d713b29a8cb44b498002d43 Mon Sep 17 00:00:00 2001 From: Yao Yao Date: Wed, 28 Oct 2026 11:04:14 +0000 Subject: [PATCH 07/12] [None][fix] KVCM2: declare the resume stub's is_decoding slot The binding keeps `is_decoding` as resume()'s second positional parameter while callers still pass it, but the stub declared `optional_reuse_groups` there and omitted it. A positional call written against the stub sent its reuse-group list to `is_decoding`, and a keyword call passing `is_decoding` did not match the stub at all. Goes with the binding shim it describes. Signed-off-by: Yao Yao --- tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi index 32ea799b14da..4a187007f2ff 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi @@ -795,6 +795,7 @@ class _KVCache: def resume( self, cuda_stream: CudaStream | None = None, + is_decoding: bool | None = None, optional_reuse_groups: Sequence[LayerGroupId] | None = None, ) -> bool: """Lock this cache's pages at their required levels. @@ -805,6 +806,10 @@ class _KVCache: Takes no sparse cache arena lease: that follows membership, so a request receiving its history costs the GPU cache nothing until it is seated. + + ``is_decoding`` is accepted and ignored. Resuming no longer distinguishes a decoding + request; the parameter is kept so callers that still pass it keep working, and goes + with the full sparse attention support. """ def prefetch(self, target: CacheLevel) -> bool: ... def get_swa_scratch_range(self, layer_group_id: LayerGroupId) -> HalfOpenRange: From eea8b3d069459f1ea36e63a39f3645fe758c31d8 Mon Sep 17 00:00:00 2001 From: Yao Yao Date: Wed, 28 Oct 2026 11:21:35 +0000 Subject: [PATCH 08/12] [None][fix] KVCM2: re-export ScratchDesc and use HostMem::sizeGranularity in the NVFP4 cold-page test Co-Authored-By: Claude Opus 5.5 Signed-off-by: Yao Yao --- cpp/tests/unit_tests/kernels/nvfp4ColdPageKernelsTest.cpp | 2 +- tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py | 2 ++ tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi | 5 +++++ 3 files changed, 8 insertions(+), 1 deletion(-) diff --git a/cpp/tests/unit_tests/kernels/nvfp4ColdPageKernelsTest.cpp b/cpp/tests/unit_tests/kernels/nvfp4ColdPageKernelsTest.cpp index 37a5a66f8d1f..fd2db655e4ec 100644 --- a/cpp/tests/unit_tests/kernels/nvfp4ColdPageKernelsTest.cpp +++ b/cpp/tests/unit_tests/kernels/nvfp4ColdPageKernelsTest.cpp @@ -302,7 +302,7 @@ class MappedHostRegion { public: explicit MappedHostRegion(std::size_t payloadBytes) - : mMemory(roundUp(kGuardBytes + payloadBytes + kGuardBytes, HostMem::kAlignment)) + : mMemory(roundUp(kGuardBytes + payloadBytes + kGuardBytes, HostMem::sizeGranularity())) , mPayloadBytes(payloadBytes) { TLLM_CHECK_WITH_INFO(kGuardBytes + payloadBytes + kGuardBytes <= mMemory.size(), diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py index 8834aee6181f..b5433ecdf335 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py @@ -163,6 +163,7 @@ class _KVCacheManagerConfigFieldSpec: ReusedBlocksByLevel = _cpp.ReusedBlocksByLevel Batch = _cpp.Batch IndexTableSpec = _cpp.IndexTableSpec +ScratchDesc = _cpp.ScratchDesc SlotDesc = _cpp.SlotDesc SlotDescVariant = _cpp.SlotDescVariant SsmLayerConfig = _cpp.SsmLayerConfig @@ -270,6 +271,7 @@ def typed_range(*args: int) -> range: "Priority", "ReuseGroupStatus", "ReuseScope", + "ScratchDesc", "Batch", "IndexTableSpec", "KVCacheIterationStatsDelta", diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi index 4a187007f2ff..b02ad96eb9fc 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi @@ -924,6 +924,11 @@ class HalfOpenRange: def __len__(self) -> int: ... def __eq__(self, other: object) -> bool: ... +class ScratchDesc: + range: HalfOpenRange + slot_ids: Sequence[int] + def __bool__(self) -> bool: ... + @dataclass(slots=True, frozen=True) class PageIndexConverter: scale: int From 37de9108f0a0b9f9eba1dd1ab9fd69bb84dba940 Mon Sep 17 00:00:00 2001 From: Yao Yao Date: Wed, 28 Oct 2026 16:48:55 +0000 Subject: [PATCH 09/12] [None][fix] KVCM2: let PageIndexConverter take the ScratchDesc that get_scratch_desc returns Co-Authored-By: Claude Opus 5.5 Signed-off-by: Yao Yao --- .../batch_manager/kvCacheManagerV2.cpp | 23 +++++++++++++++---- .../runtime/kv_cache_manager_v2/__init__.pyi | 2 +- 2 files changed, 20 insertions(+), 5 deletions(-) diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp index 89fed2ed61d8..f7d9ab23f7e0 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp @@ -3723,13 +3723,28 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) .def_ro("scratch_pages_per_block", &kv::PageIndexConverter::scratchPagesPerBlock) .def( "__call__", + // `scratch_range` also takes the ScratchDesc that get_scratch_desc() returns, which + // carries its slots with it; see ScratchDescCompat. [](kv::PageIndexConverter const& self, std::vector const& baseIndices, - std::optional indexMode, - std::optional> scratchRange, + std::optional indexMode, nb::object scratch, std::vector const& scratchSlots) -> std::vector { - return self(baseIndices, indexMode, scratchRange.value_or(kv::HalfOpenRange{}), - kv::Span{scratchSlots.data(), static_cast(scratchSlots.size())}); + kv::HalfOpenRange range{}; + std::vector const* slots = &scratchSlots; + if (nb::isinstance(scratch)) + { + TLLM_CHECK_WITH_INFO( + scratchSlots.empty(), "scratch_slots must be empty when scratch_range is a ScratchDesc"); + auto const& desc = nb::cast(scratch); + range = desc.range; + slots = &desc.slotIds; + } + else if (!scratch.is_none()) + { + range = nb::cast>(scratch); + } + return self( + baseIndices, indexMode, range, kv::Span{slots->data(), static_cast(slots->size())}); }, nb::arg("base_indices"), nb::arg("index_mode") = nb::none(), nb::arg("scratch_range") = nb::none(), nb::arg("scratch_slots") = std::vector{}); diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi index b02ad96eb9fc..3d88b4a4fa3f 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi @@ -940,7 +940,7 @@ class PageIndexConverter: self, base_indices: Sequence[int], index_mode: PageIndexMode | None = None, - scratch_range: HalfOpenRange | None = None, + scratch_range: HalfOpenRange | ScratchDesc | None = None, scratch_slots: Sequence[int] = (), ) -> list[int]: ... From 0b81bd8c8d9c4e753502e158382ef9f9c2d01e15 Mon Sep 17 00:00:00 2001 From: Yao Yao Date: Wed, 28 Oct 2026 17:22:15 +0000 Subject: [PATCH 10/12] [None][test] KVCM2: shrink the GPU tier to its floors before the DSv4 cold-page pressure request Co-Authored-By: Claude Opus 5.5 Signed-off-by: Yao Yao --- .../deepseek_v4/test_deepseek_v4_cache_manager.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py index e03bb4e02f5b..43820e811229 100644 --- a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py +++ b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py @@ -62,6 +62,7 @@ BAD_PAGE_INDEX, BatchDesc, BufferConfig, + CacheTier, GpuCacheTierConfig, HostCacheTierConfig, KVCacheDesc, @@ -990,6 +991,15 @@ def test_nvfp4_cold_page_codec_migrates_real_csa_hca_through_host( first_cache = cache_manager.kv_cache_map[first.py_request_id] cache_manager.suspend_request(first) + # The pressure request displaces the suspended one only when the GPU tier holds + # no slack beyond its slot floors, which depends on how the allocator packed the + # pools. Shrink the tier to the least the floors allow; resize refuses anything + # smaller. + gpu_quota = cache_manager.impl.get_quota(0) + grain = _introspection.tier_allocation_granularity(CacheTier.GPU_MEM, gpu_quota) + while cache_manager.impl.resize(0, gpu_quota - grain): + gpu_quota -= grain + pressure = self._create_request(request_id=1, prompt_len=pressure_len) requests.append(pressure) assert cache_manager.prepare_context(pressure) From bfc7cb1680834eb68bdff93017c710010bb2181e Mon Sep 17 00:00:00 2001 From: Yao Yao Date: Wed, 28 Oct 2026 17:22:28 +0000 Subject: [PATCH 11/12] [None][test] KVCM2: skip the large sparse-resolve test on GPUs with under 160 KiB of shared memory per block Co-Authored-By: Claude Opus 5.5 Signed-off-by: Yao Yao --- .../test_kv_cache_manager_v2.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py index 14c540318766..829aab663454 100755 --- a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py +++ b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py @@ -1842,6 +1842,20 @@ def test_layers_resolve_independently(self): def test_large_caches_preserve_victims_and_recency_order(self): """Hits, fetched victims, and untouched entries retain their relative order.""" + import cuda.bindings.runtime as cudart + + # The 8192-entry case needs 147552 bytes of dynamic shared memory per resolve block, + # on top of what the kernel declares statically. + err, device = cudart.cudaGetDevice() + self.assertEqual(int(err), int(cudart.cudaError_t.cudaSuccess)) + err, shared_optin = cudart.cudaDeviceGetAttribute( + cudart.cudaDeviceAttr.cudaDevAttrMaxSharedMemoryPerBlockOptin, device + ) + self.assertEqual(int(err), int(cudart.cudaError_t.cudaSuccess)) + if shared_optin < 160 << 10: + self.skipTest( + f"needs 160 KiB of shared memory per block; this device allows {shared_optin}" + ) for entries in (31, 32, 33, 255, 256, 257, 2048, 4096, 4224, 8192): with self.subTest(entries=entries): rng = random.Random(entries) From 8df9cace60cf9dfbedf75b3d8a8600fedc829a73 Mon Sep 17 00:00:00 2001 From: Yao Yao Date: Wed, 28 Oct 2026 17:58:09 +0000 Subject: [PATCH 12/12] [None][test] KVCM2: give the lender's sparse refusal test whole sparse-attention life cycles with a dense indexer Co-Authored-By: Claude Opus 5.5 Signed-off-by: Yao Yao --- .../kv_cache/sharing/test_staging_lender.py | 51 ++++++++++++------- 1 file changed, 32 insertions(+), 19 deletions(-) diff --git a/tests/unittest/_torch/executor/kv_cache/sharing/test_staging_lender.py b/tests/unittest/_torch/executor/kv_cache/sharing/test_staging_lender.py index 46df1d1206fe..79717b94f868 100644 --- a/tests/unittest/_torch/executor/kv_cache/sharing/test_staging_lender.py +++ b/tests/unittest/_torch/executor/kv_cache/sharing/test_staging_lender.py @@ -175,36 +175,49 @@ def test_recurrent_state_is_refused(kit, hybrid_manager): assert len(kit.retained()) == before, "a refused attach allocated staging" -def sparse_layer(monkeypatch, index): - """Managers built from now on mark the buffers of their layer ``index`` sparse, which no model - here does: the runtime may then lock that layer group's read-only pages in host memory.""" +def sparse_layers(monkeypatch, indices): + """Managers built from now on give their layers ``indices`` sparse attention, which no model + here does: their K/V buffers become sparse, and each also carries the dense indexer buffer + that selects from them and the sparse cache it selects into. ``indices`` must cover whole + life cycles, as a model with sparse attention does; the runtime may then lock those layer + groups' read-only pages in host memory.""" from dataclasses import replace from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 - from tensorrt_llm.runtime.kv_cache_manager_v2 import AttentionLayerConfig, BufferConfig + from tensorrt_llm.runtime.kv_cache_manager_v2 import ( + AttentionLayerConfig, + BufferConfig, + SparseCacheConfig, + ) build = KVCacheManagerV2._build_cache_config - def with_a_sparse_layer(self, config): + def with_sparse_layers(self, config): layers = list(config.layers) - layer = layers[index] - layers[index] = AttentionLayerConfig( - layer_id=layer.layer_id, - buffers=[ + for index in indices: + layer = layers[index] + buffers = [ BufferConfig(b.role, b.size, b.tokens_per_block_override, is_sparse=True) for b in layer.buffers - ], - sliding_window_size=layer.sliding_window_size, - num_sink_tokens=layer.num_sink_tokens, - ) + ] + buffers.append(BufferConfig("indexer", layer.buffers[0].size)) + layers[index] = AttentionLayerConfig( + layer_id=layer.layer_id, + buffers=buffers, + sliding_window_size=layer.sliding_window_size, + num_sink_tokens=layer.num_sink_tokens, + sparse_cache=SparseCacheConfig( + num_selected_blocks=1, max_cached_units=2, unit_tokens=self.tokens_per_block + ), + ) return build(self, replace(config, layers=layers)) - monkeypatch.setattr(KVCacheManagerV2, "_build_cache_config", with_a_sparse_layer) + monkeypatch.setattr(KVCacheManagerV2, "_build_cache_config", with_sparse_layers) -# case: (manager keywords, the layer whose buffers are sparse); "later" puts them in a layer group -# after a dense one -SPARSE_CASES = {"first": ({}, 0), "later": ({"windows": [WINDOW, 256]}, 1)} +# case: (manager keywords, the layers given sparse attention); "first" makes the only life cycle +# sparse, "later" the one after a dense life cycle +SPARSE_CASES = {"first": ({}, (0, 1)), "later": ({"windows": [WINDOW, 256]}, (1,))} @pytest.mark.parametrize("case", list(SPARSE_CASES)) @@ -213,8 +226,8 @@ def test_a_layer_group_of_sparse_buffers_is_refused(kit, host_tier_manager, monk names a host slot that staging would address as a device slot: the attach refuses such a manager, naming the group and keeping nothing, also where it follows a dense one. The runtime builds sparse buffers only with host memory as the cache level below the GPU, as here.""" - keywords, index = SPARSE_CASES[case] - sparse_layer(monkeypatch, index) + keywords, indices = SPARSE_CASES[case] + sparse_layers(monkeypatch, indices) before = len(kit.retained()) with host_tier_manager(**keywords) as mgr: layers = {int(layer.layer_id): layer for layer in mgr.impl.init_config.layers}