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 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..f7d9ab23f7e0 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp @@ -69,12 +69,16 @@ namespace tensorrt_llm::nanobind::batch_manager namespace { -//! A DLPack view keeps the batch's device allocation alive without importing torch. -struct BatchDeviceArray +//! 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 { - std::shared_ptr batch; - kv::LayerGroupId group; - bool counts; + kv::HalfOpenRange range; + std::vector slotIds; }; // Exposed via introspection sub-module for tests. @@ -315,6 +319,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 +394,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); @@ -1392,32 +1477,35 @@ 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); - 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 +1651,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 +1708,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 +1861,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 +1887,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 +1899,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 +1911,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 +1938,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,9 +1949,15 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) nb::class_(m, "_KVCache") .def( "resume", + // `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)); @@ -1949,31 +1971,15 @@ 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()) .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 +1999,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 +2060,36 @@ 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.getScratchDesc(kv::LayerGroupId{layerGroupId}); }, - nb::arg("layer_group_id"), nb::call_guard()) + { 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) + { + 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( + "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) @@ -2063,10 +2103,41 @@ 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) + // 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", + [](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 +2204,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 +2223,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 +2472,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 +2712,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 +2776,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 +2807,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 +3365,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 +3450,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 +3638,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 +3669,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( @@ -2970,19 +3723,31 @@ 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, nb::object scratchObj) -> std::vector + std::optional indexMode, nb::object scratch, + std::vector const& scratchSlots) -> std::vector { - kv::ScratchDesc const* scratch = nullptr; - std::optional scratchHolder; - if (!scratchObj.is_none()) + 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()) { - scratchHolder = nb::cast(scratchObj); - scratch = &*scratchHolder; + range = nb::cast>(scratch); } - return self(baseIndices, indexMode, 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") = 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/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/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/_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() 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..b5433ecdf335 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,6 +161,8 @@ class _KVCacheManagerConfigFieldSpec: ReuseGroupStatus = _cpp.ReuseGroupStatus ReuseScope = _cpp.ReuseScope ReusedBlocksByLevel = _cpp.ReusedBlocksByLevel +Batch = _cpp.Batch +IndexTableSpec = _cpp.IndexTableSpec ScratchDesc = _cpp.ScratchDesc SlotDesc = _cpp.SlotDesc SlotDescVariant = _cpp.SlotDescVariant @@ -184,6 +184,7 @@ class _KVCacheManagerConfigFieldSpec: BeamIndex = int CacheLevel = int CacheTierConfig = Union[GpuCacheTierConfig, HostCacheTierConfig, DiskCacheTierConfig] +BatchRow = int CudaStream = int DataRole = str LayerGroupId = int @@ -222,6 +223,7 @@ def typed_range(*args: int) -> range: "BatchDesc", "BeamIndex", "BufferConfig", + "SparseCacheConfig", "BufferId", "CoalescedBuffer", "CacheLevel", @@ -251,9 +253,9 @@ def typed_range(*args: int) -> range: "KVCacheUpdatedData", "KvCacheStatus", "LayerGroupId", + "BatchRow", "LayerId", "LifeCycleId", - "LogicError", "MemAddress", "NDEBUG", "OutOfPagesError", @@ -262,9 +264,6 @@ def typed_range(*args: int) -> range: "PoolGroupPeakBlockStats", "PageIndexMode", "PageStatus", - "PageStorageSnapshot", - "Batch", - "BatchDeviceArray", "PoolDesc", "PoolGroupDesc", "PoolGroupIndex", @@ -273,6 +272,8 @@ def typed_range(*args: int) -> range: "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..3d88b4a4fa3f 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,44 @@ 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. + + ``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_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 +832,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 +851,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,7 +924,6 @@ 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] @@ -801,12 +934,14 @@ 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 | ScratchDesc | None = None, + scratch_slots: Sequence[int] = (), ) -> list[int]: ... class IKvCacheColdPageCodec: ... @@ -820,6 +955,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 +1659,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 +1818,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/_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) 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} 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..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 @@ -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,5502 @@ 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.""" + 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) + 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 +6085,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 +6267,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 +9506,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 +9543,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 +9991,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 +9999,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 +10016,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 +10070,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 +10088,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 +10149,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 +10174,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 +10201,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 +10290,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 +10300,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 +10316,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 +10325,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)