From 4bfc3904fbafba1ea0237c02fb7c803c029fd8cb Mon Sep 17 00:00:00 2001 From: Guiju Zhang <7135567+cascade812@users.noreply.github.com> Date: Fri, 30 Oct 2026 03:13:21 +0000 Subject: [PATCH 01/11] [TRTLLM-16282][feat] Integrate sparse KV offload for DeepSeek-V4 and DSA on KVCM2 Sparse KV cache offload for DeepSeek-V4 and DSA (GLM-5.x) layers, on the KVCM2 sparse attention core (#21129). A decode-only worker keeps the KV history that sparse attention selects from in host memory, and each decode step fetches the selected units into a per-request GPU cache. - SparseKvCacheOffloadConfig, enabled with enable_kv_cache_offload on the DSA and DeepSeek-V4 sparse attention configs. The serve schemas and the llm_args golden manifest follow. With offload on, the indexer's paged MQA logits default to DeepGEMM. - KVCacheManagerV2 owns the attention Batch, with one arena lease per Batch row. Offload sizing (sparse_offload_memory.py) takes the arena and offload buffers from the device budget before max_tokens sizes the pools, and measures host memory from MemAvailable. - One fetch path (sparse/offload.py, offload_kernels.py) serves the DeepSeek-V4 and DSA attention backends. The MLA RoPE write reads the token-major block table, PrimTS sparse MLA stays off offloaded ratio-4 layers, and the offload NVFP4 scratch carries the zeroed KV-tile tail. - An offloading worker is decode-only: it skips KV cache estimation, warms up generation shapes only, and captures no prefill CUDA graphs. Block reuse, and offload without a cache transceiver, are rejected. - Sparse KV transfers into an offloading generation worker. - KVCM2: the sparse resolve keeps one 32-bit eviction word per cached entry and picks its CTA width at dispatch; a negative Batch row is rejected; the ScratchDesc / get_scratch_desc(), is_decoding and enter_decode compatibility shims are removed; SPARSE_ATTENTION_DESIGN.md describes the design. - For resident-versus-offload parity runs: FORCE_DETERMINISTIC=1 and the FORCE_TOPK_KERNEL_DETERMINISTIC / FORCE_MHC_KERNEL_DETERMINISTIC switches make TopK, mHC and the trtllm-gen FMHA reduction repeatable, the disaggregated sender keeps unclaimed REQUEST_DATA for the transfer timeout when that exceeds 120 s, and token_parity.md records the runs. Signed-off-by: Guiju Zhang <7135567+cascade812@users.noreply.github.com> --- .../SPARSE_ATTENTION_DESIGN.md | 372 ++-- .../kv_cache_manager_v2/batch.cpp | 26 +- .../batch_manager/kv_cache_manager_v2/batch.h | 17 + .../kv_cache_manager_v2/introspection.cpp | 101 +- .../kv_cache_manager_v2/introspection.h | 16 +- .../kv_cache_manager_v2/kvCacheManager.cpp | 2 +- .../kv_cache_manager_v2/sparseCache.cpp | 73 +- .../kv_cache_manager_v2/sparseCache.h | 122 +- .../kv_cache_manager_v2/sparseCacheKernels.cu | 664 +++--- .../kv_cache_manager_v2/sparseCacheKernels.h | 13 +- .../trtllmGenKernels/fmha/fmhaKernels.h | 12 +- .../batch_manager/kvCacheManagerV2.cpp | 105 +- cpp/tests/unit_tests/kernels/CMakeLists.txt | 1 + .../kernels/fmhaDeterministicTest.cpp | 246 +++ scripts/generate_trtllm_serve_schemas.py | 23 + .../backends/fmha/prims_ts_sparse_mla.py | 7 + .../backends/sparse/csa2/cache_manager.py | 23 +- .../backends/sparse/deepseek_v4/backend.py | 185 +- .../sparse/deepseek_v4/cache_manager.py | 720 +++---- .../backends/sparse/deepseek_v4/compressor.py | 59 +- .../backends/sparse/deepseek_v4/kernels.py | 615 ++---- .../backends/sparse/deepseek_v4/metadata.py | 83 +- .../backends/sparse/deepseek_v4/offload.py | 88 - .../backends/sparse/deepseek_v4/params.py | 2 - .../attention/backends/sparse/dsa/backend.py | 188 +- .../backends/sparse/dsa/cache_manager.py | 369 +++- .../attention/backends/sparse/dsa/metadata.py | 104 + .../attention/backends/sparse/dsa/module.py | 8 + .../attention/backends/sparse/dsa/params.py | 2 + .../attention/backends/sparse/offload.py | 1101 ++++++++++ .../backends/sparse/offload_kernels.py | 358 ++++ .../_torch/attention/backends/trtllm.py | 89 +- .../_torch/disaggregation/native/transfer.py | 209 +- .../disaggregation/resource/kv_extractor.py | 23 +- .../_torch/disaggregation/resource/page.py | 14 +- .../_torch/disaggregation/resource/utils.py | 27 +- tensorrt_llm/_torch/model_config.py | 26 +- tensorrt_llm/_torch/modules/mhc/mhc_cuda.py | 34 +- tensorrt_llm/_torch/modules/top_k.py | 12 + tensorrt_llm/_torch/pyexecutor/_util.py | 14 +- .../_torch/pyexecutor/config_utils.py | 13 + .../engine/runners/decoder/graphs.py | 7 + .../engine/runners/decoder/warmup.py | 38 +- .../kv_cache/kv_cache_manager_v2.py | 973 +++++++-- .../pyexecutor/kv_cache/sharing/_layout.py | 6 +- .../kv_cache/sparse_offload_memory.py | 839 ++++++++ tensorrt_llm/_torch/pyexecutor/py_executor.py | 38 +- .../_torch/pyexecutor/py_executor_creator.py | 17 +- tensorrt_llm/llmapi/__init__.py | 5 +- tensorrt_llm/llmapi/llm_args.py | 92 +- .../runtime/kv_cache_manager_v2/__init__.py | 2 - .../runtime/kv_cache_manager_v2/__init__.pyi | 37 +- .../kv_cache_manager_v2/_introspection.py | 19 +- .../schemas/trtllm-serve-config.schema.json | 112 +- .../trtllm-serve-disagg-config.schema.json | 112 +- .../usage/llm_args_golden_manifest.json | 10 + .../test_lists/test-db/l0_b200.yml | 1 + .../test_lists/test-db/l0_gb300.yml | 1 + .../deepseek_v4/test_compressor_kernel.py | 177 ++ .../deepseek_v4/test_compressor_module.py | 104 +- .../test_deepseek_v4_cache_manager.py | 1848 +++++++++++++++-- .../deepseek_v4/test_deepseek_v4_kernels.py | 575 +++-- .../deepseek_v4/test_deepseek_v4_offload.py | 1270 +++++------ .../test_prims_ts_sparse_mla_fmha.py | 35 + .../attention/sparse/dsa/test_dsa_indexer.py | 5 + .../dsa/test_dsa_offload_cache_manager.py | 789 +++++++ .../sparse/dsa/test_dsa_offload_config.py | 134 ++ .../sparse/test_sparse_offload_kernels.py | 262 +++ .../engine/runners/decoder/test_warmup.py | 95 +- .../kv_cache/sharing/test_staging_lender.py | 15 +- .../kv_cache/test_kv_cache_v2_scheduler.py | 289 --- .../kv_cache/test_kvcm2_integration.py | 1545 ++++++++++++-- .../kv_cache/test_kvv2_host_tier_sizing.py | 16 + .../kv_cache/test_mixed_precision_kv_cache.py | 2 +- .../kv_cache/test_sparse_offload_memory.py | 767 +++++++ .../_torch/executor/test_py_executor.py | 71 + ...y_executor_creator_mla_cache_reuse_sync.py | 24 + .../modeling/test_modeling_deepseekv4.py | 31 + tests/unittest/_torch/modules/test_mhc.py | 53 +- tests/unittest/_torch/modules/test_top_k.py | 41 + .../disaggregated/test_chunked_transfer.py | 12 +- .../test_deepseek_v4_kv_transfer.py | 228 ++ .../disaggregated/test_sparse_kv_transfer.py | 341 +++ .../test_transfer_ownership_regressions.py | 9 +- .../disaggregated/test_unsubmitted_session.py | 27 + .../test_kv_cache_manager_v2.py | 118 +- tests/unittest/llmapi/test_llm_args.py | 136 +- .../test_trtllm_serve_config_schemas.py | 42 + token_parity.md | 295 +++ 89 files changed, 14090 insertions(+), 3741 deletions(-) create mode 100644 cpp/tests/unit_tests/kernels/fmhaDeterministicTest.cpp delete mode 100644 tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/offload.py create mode 100644 tensorrt_llm/_torch/attention/backends/sparse/offload.py create mode 100644 tensorrt_llm/_torch/attention/backends/sparse/offload_kernels.py create mode 100644 tensorrt_llm/_torch/pyexecutor/kv_cache/sparse_offload_memory.py create mode 100644 tests/unittest/_torch/attention/sparse/dsa/test_dsa_offload_cache_manager.py create mode 100644 tests/unittest/_torch/attention/sparse/dsa/test_dsa_offload_config.py create mode 100644 tests/unittest/_torch/attention/sparse/test_sparse_offload_kernels.py create mode 100644 tests/unittest/_torch/executor/kv_cache/test_sparse_offload_memory.py create mode 100644 tests/unittest/disaggregated/test_sparse_kv_transfer.py create mode 100644 token_parity.md 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 index be085a74c87b..8803aecf43d5 100644 --- 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 @@ -1,3 +1,8 @@ + + # Sparse attention in KVCacheManagerV2 Target: the **decode instance of a disaggregated deployment**. @@ -88,7 +93,7 @@ entry and the selection unit is already a page. 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. + replaced by the same eviction. 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 @@ -125,15 +130,15 @@ already be in its host page. 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. +The flush is a call of its own rather than a step inside `resize()`, for two reasons. + +It carries an ordering requirement that only a caller can satisfy: the write-back reads +converted page indices, and only a publish produces those. A flush reached through +`resize()` would depend on when the last publish happened without saying so; a named call +makes it a stated precondition. -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 +It is also what lets the write-back be **batch-wide**. Per request it would be 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. @@ -142,13 +147,29 @@ row more than once, and a row whose growth failed rolls back to a history below 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. +**The neighbouring hooks are not flush sites, and each would be wrong as one.** + +- `resize()` is where the watermark moves, and it is the only such place: `setHistoryLength` + is `resize(nullopt, hist)` and `setCapacity()` is `resize(cap, nullopt)`, and one call may + move capacity and history together. A check placed anywhere else would see only half the + moves. +- `_unlockStaleBlocks` exists for SWA blocks outside the window, whose content is *dead*; + sparse history is live and will be read again. The trigger and the range arithmetic match, + the disposition does not — routing sparse layers through it would drop the data rather + than write it back. It also gates on `windowSize.has_value()`, so for a full-attention + sparse layer it fires on nothing. +- `_shortcutSetHistoryLength` returns "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`, returns true, and `resize()` returns before reaching + its body. Work placed in that body does not run for a sparse layer. +- `commit()` needs no flush site: 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()`. + ### How much a step may cross The write-back's work list is sized by `maxTotalInputTokens`, the tokens one step may ingest @@ -213,7 +234,7 @@ watermark — and the second is written by definition. So the two populations ar distinguished by the ordinal alone: ``` -unflushed <=> resident_ordinal >= history_length / unit_tokens +unflushed <=> entry_ordinal >= history_length / unit_tokens ``` Nothing records this. History is monotonic (`resize()` rejects a decrease), so an ordinal @@ -231,8 +252,9 @@ 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. +- **The unflushed set is at most one unit**, bounded by `max_input_length` below rather + than observed at runtime. 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 @@ -317,7 +339,7 @@ filled together; the cache took whole coalesced slots only while it drew them fr 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 +The consequence 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 @@ -325,8 +347,9 @@ 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. +Per `(request, layer)`, managed entirely by kernels: one 32-bit word per entry, holding +the unit resident there and when that unit was last used. A free entry is a value of the +same word, so there is no separate 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. A layer @@ -339,17 +362,30 @@ is a real exclusion rather than a hypothetical: Rocket's selection indices are 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. +**Scope is per request.** Two requests sharing a prefix do not share a cached copy. That +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; not supported. + +**Replacement is by age, not "evict whatever was not selected this step".** Retention +across steps is the point: holding hot entries beyond the current top-k gives a large +miss-rate reduction (HiSparse reports the measurement), and that is only available when +the cache is larger than the selection width. + +An entry's word carries a 7-bit step timestamp and one history bit saying whether its last +use found it already resident. Eviction takes the highest `2 * age - history`, so age +decides and, within one age, a unit fetched that step leaves before one that was hit — which +is what makes a repeatedly selected unit outrank a first-time selection. The order is +therefore approximate where an exact recency permutation would be exact: two entries of the +same age and the same history bit are interchangeable, and which of them leaves is whichever +thread reaches the oversubscribed bucket first. + +The clock wraps, so a timestamp left untouched for longer than its range would read as a +recent one. Every resolve streams every entry anyway, so one resolve in every +`kRebaseStep` pulls the timestamps it passes up to the oldest age the field still +distinguishes. That is idempotent, which is what lets the write-index path apply it too +without the two having to agree on which of them does. **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 @@ -361,7 +397,8 @@ ordinal to slot. The two scale with different things: | 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. +stop paying for. The dense reverse form is also what the eviction pass wants: it reads every +entry once either way. ### Probe index @@ -380,61 +417,74 @@ by the cache's capacity `N`: 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(...); } + table.insert(ord, mask); // one CAS per probe step } __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 + uint32_t word = entry[s]; + entryCache[s] = word; + int32_t ord = SparseEntry::ordinal(word); + bool matched = ord >= 0 && table.publish(ord, s, mask); // claims its query + flag[s] = matched ? kFlagMatched : (ord >= historyUnits ? kFlagHeld : kFlagEvictable); + if (flag[s] == kFlagEvictable) atomicAdd(&hist[staleness(word)], 1); } ``` 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. +empty entry. A position is a single 32-bit word holding both the ordinal and the entry that +claimed it: a `__syncthreads()` separates build from match, so no thread reads a position +mid-insert and the entry index needs no atomicity of its own. Packing them halves the +shared-memory traffic of the publish pass, which walks the whole resident set and is where +most of the kernel's time goes. The split is 14 bits of entry index against 18 of ordinal, +and both bounds are checked rather than truncated into. + +Multiplicative hashing rather than identity, because ordinals are small and dense, and dense +runs make an unsuccessful probe walk the whole run. + +The table stays at `next_pow2(2 * required)` rather than being right-sized to a load factor +of one half. Right-sizing measured *slower*: the publish pass is dominated by *unsuccessful* +probes, whose expected cost is `0.5 * (1 + 1/(1-a)^2)` and so rises much faster in the load +factor `a` than a successful probe's does. + +The resident sweep is also where each entry is classified and its staleness binned, so the +pass the match needs is the pass the eviction ranking needs; see "Capacity and eviction". + +**The CTA width is chosen per life cycle.** Every pass but the probe is a stride loop over +the resident set, so doubling the CTA halves what a thread walks — but only pays while +shared memory already holds a multiprocessor to one or two blocks of the kernel. Past that +the wider CTA buys nothing per block and costs residency. The kernel is therefore +instantiated at 256 and 512 threads, and the launcher picks between them by asking the +device for the narrow instantiation's occupancy at this shape's shared-memory size. The +three extents are the life cycle's, so the answer is computed once and held. + +**The index is keyed by the selection, not by the resident set.** The alternative — keying +by ordinal and valuing by slot, sized `next_pow2(2N)` — returns a hit's slot straight from +the probe and so saves one pass. It does not make the miss list available any sooner, as +might be supposed: that build has to walk every slot and cross a `__syncthreads()` before +its first probe, so the misses are known no earlier either way. + +What keying by the selection costs, and what it buys: + +- **Shared memory is O(k), not O(N).** At the DSA geometry the table is 2048 buckets rather + than 8192. This is the point of the choice: `N` 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 +- **The index build cannot precede the dependency wait.** A resident-keyed build reads only + the layer's own residency table and 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. + produces. Only shared-memory initialisation 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. +On DSA at the geometry above, the resolve kernel takes 0.833 ms per decode step summed +across 21 sparse layers, against 1.468 ms for the resident-keyed form. 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". +the table entirely and runs a separate kernel that scans the entry words for its single +in-flight ordinal, reusing it where a previous step already pinned it and taking the stalest +candidate otherwise. See "Writing the in-flight unit". ## Extraction interface @@ -509,14 +559,14 @@ 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. +The descriptor path crosses **no plugin boundary**. KVCM2 resolves selections, updates the +eviction state, 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. +copies start while the resolve is still finishing. 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. @@ -548,7 +598,7 @@ NVFP4 while leaving the active GPU cache alone, and `triattention`, which evicts 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: +it does. Which pages, precisely, because two true sentences sound like one here: 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. @@ -571,13 +621,13 @@ logic on identical input. 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, +layers, each against its own resident set — 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 +set, 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 @@ -600,7 +650,7 @@ 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 +write-back: `sparseResolveKernel` is the only writer of the lease's entry words, and `sparseFlushListKernel` only reads them — so collapsing the resolve is sufficient, and a per-layer flush against shared metadata stays correct. @@ -624,9 +674,10 @@ 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; +4. allocate a slot per miss, taking the stalest entries first; 5. extract and copy each miss from the device-mapped host page; -6. promote hits above new misses in recency order; +6. stamp every entry the step used with the current step, flagged by whether it was already + holding its unit; 7. write the output index list: one index per selection slot, in the selection's own order, `kBadPageIndex` where a slot is padded. @@ -642,9 +693,9 @@ One index per buffer, though. A layer's buffers share the resolve but not their 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. +Steps 4 and 6 are one pass: an entry's new word is written where its fate is decided, so +nothing is left to do after the work list is complete. The copies — the only step that +touches PCIe, and so the long pole — are released as soon as that pass ends. No host involvement, so the call enqueues and returns. Every shape is static across steps. @@ -746,7 +797,7 @@ there is nothing to write back and no converted table is needed yet. 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. +the step clock delivers — layer N's entries are at age zero 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 @@ -909,16 +960,15 @@ 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. +Direct mutation of a seated member is permitted rather than refused. Refusing it would need +a batched twin for every mutator — `commit` changes locks and runs every decode step — and +the failure it would prevent is a restaged table, which costs a step rather than giving 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. +**Batched resize** resizes every member and publishes once at the end. The contract is row +order and publish-on-completion, which holds however the resize itself is implemented. The +result is per-request rather than a single flag, since `_KVCache.resize()` rolls back per +cache and 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. @@ -943,7 +993,7 @@ 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. +where the hit rate comes from. That range is the one HiSparse measures as useful. 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 @@ -966,7 +1016,8 @@ Declared or planned in `tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi`: 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. +- `Batch` — membership, row order, device arrays, `resize()` and `publish()`, which + transfers the whole table; dirty-row tracking is planned, not declared. - `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. @@ -979,47 +1030,27 @@ 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()`. +## Lock levels + +Slot release is level-aware: `Page::~Page()` releases with the page's own `cacheLevel`, and +the `kHotLevel` in `ScratchSlotLock::unlock()` holds because 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. No per-target lock level exists: there is no `queryLockLevel`, no +`BatchedLockTarget::cacheLevel`, no `StaleBackup::cacheLevel` and no `residencyLevel`. -Also: no attention kernel addresses a host page. A sparse layer's kernel reads its units -from the arena, which is GPU memory bounded by the lease count and the cells a lease holds, -and a non-sparse layer reads the device pool through its block table as it always has. The -host pool is reached only by this manager's own copies, which compute their offsets in -`size_t`. +No attention kernel addresses a host page. A sparse layer's kernel reads its units from the +arena, which is GPU memory bounded by the lease count and the cells a lease holds, and a +non-sparse layer reads the device pool through its block table as it always has. The host +pool is reached only by this manager's own copies, which compute their offsets in `size_t`. ## 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 +stride — a selection's rows live at `entry + slot * entriesPerSelection` — so one entry count serves every layer, and `cellsPerBuffer` follows from `max_cached_units` and the unit size together. @@ -1031,8 +1062,8 @@ That is an implementation limit, not a property of the design. Lifting it: 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. +- The entry words 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. @@ -1185,11 +1216,10 @@ block that has not yet waited may still be running beside a grid that could read 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 end of its eviction loop, which is also where the entry words are written, so it calls +`cudaTriggerProgrammaticLaunchCompletion()` there and only its miss count follows. The +flush-list scan is the same shape — the list is the last thing it writes — so in both cases +the release brings the copy's launch forward rather than overlapping a tail. 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 @@ -1350,22 +1380,22 @@ 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. +unallocated. No host tier's quota holds them either: `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. +It is the only coverage that constrains *which* allocation the copy reads. Offsetting the +page base by one slot leaves 57 other sparse tests passing, because residency, index shape +and counts are all insensitive to that address — so nothing else in the suite would catch a +pool addressed at the wrong base. -**No cell is ever read from, or written to, an encoded page.** 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: +**No cell is ever read from, or written to, an encoded page.** Sparse pages sit at the hot +level, so they are eligible for demotion, and demotion is what invokes the codec. A page +cannot be encoded while the cache still holds a cell copied from it: - An `ACTIVE` request locks the pages it requires (`AGENTS.md`, *Per-request cache*), and `LOCKED` forbids eviction outright (*Page status*). @@ -1386,49 +1416,45 @@ a block the window has made stale is never named. Nothing reads an unlocked page can read one that has been encoded or recycled. A cell fetched before the window moved does outlive its page's lock there. It is harmless: -nothing names it again, so the LRU evicts it within a few steps. The write-back is +nothing names it again, so it ages out within a few steps. The write-back is unaffected in either case — the units it crosses lie in `[history, capacity)`, above the watermark `_unlockStaleBlocks` works from, so their pages are locked for as long as the request is active. -That is worth stating because it is the *selector*, not this package, that makes it true. +It is the *selector*, not this package, that makes that 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 does not arise today in any case: a sliding window already bounds a step's reads to a narrow fixed span, which is the span a selection would otherwise have to discover. -## Settled by construction +## Standing limits and invariants -Recorded because each looks like an open question and is not: +Properties the rest of the design is entitled to assume. +- **Decode only.** One manager cannot serve prefill and decode together: a sparse kernel + indexes its selection per query token, and a multi-token step's selections cannot all be + resident in one per-request cache at once. The decode instance of a disaggregated + deployment is the supported target — see *Supported scope*. +- **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 therefore cannot be + rewound into, and no further rule is needed to guarantee it. +- **A selection unit never exceeds `tokens_per_block`.** Sub-page units are supported + through `expansion`; a unit spanning several pages is rejected, which is what keeps slot + contiguity and k-at-a-time eviction out of the design. +- **`N` is a uniform upper bound**, not per-request. Ragged device arrays would cost the + static shapes graph capture needs. - **`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. +- **Pinning covers the in-flight unit only.** That unit is private because unflushed blocks + are not in the radix tree, and bounded to one by `max_input_length`. Ordering among the + rest comes from staleness 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), 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 e797c5c20d85..b3063b09bd36 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 @@ -106,8 +106,9 @@ Batch::Batch(KvCacheManager& manager, CUstream stream, std::vector(BufName::kPageTables), kBadPageIndex.value()); - // Watermarks of a free row are zero: a request with nothing eligible. - for (auto const name : {BufName::kHistory, BufName::kCapacity}) + // Watermarks of a free row are zero: a request with nothing eligible. The eviction + // clock starts there too, which publish() advances off before any resolve reads it. + for (auto const name : {BufName::kHistory, BufName::kCapacity, BufName::kSparseClock}) { std::ranges::fill(_buf(name), 0); } @@ -1096,6 +1097,7 @@ void Batch::_layoutBufs(size_t scratchElements) noexcept {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::kSparseClock, sizeof(int32_t), alignof(int32_t)}, {BufName::kLeases, sizeof(SparseLeaseIndex) * rows, alignof(SparseLeaseIndex)}, {BufName::kScratch, sizeof(int32_t) * scratchElements, alignof(int32_t)}, }}; @@ -1115,6 +1117,11 @@ SparseLeaseIndex const* Batch::_leaseMirror() const noexcept return reinterpret_cast(sparseLeases()); } +int32_t const* Batch::_sparseStepAgeMirror() const noexcept +{ + return reinterpret_cast(_mirrorAt(BufName::kSparseClock)); +} + size_t Batch::_mirrorBytes() const noexcept { // From the histories through the leases: the page tables ahead of them are read from the @@ -1220,8 +1227,7 @@ BatchRow Batch::addKvCache(KvCache& kvCache, std::optional requestedRo else { row = *requestedRow; - TLLM_CHECK_WITH_INFO( - row != kNoBatchRow && row < mMaxBatchSize, "Batch::addKvCache: row is beyond max_batch_size"); + TLLM_CHECK_WITH_INFO(row >= 0 && row < mMaxBatchSize, "Batch::addKvCache: row is beyond max_batch_size"); if (mRows[row] != nullptr) { throw LogicError("Batch::addKvCache: row is already occupied"); @@ -1471,6 +1477,16 @@ void Batch::publish() } mFlushCountersDirty = false; + // The sparse eviction clock, which the resolves of this step then read. It wraps into + // SparseEntry's timestamp field, and a rebase every kRebaseStep steps is what keeps a + // timestamp from falling far enough behind the wrap to read as a recent one. + // + // Staged rather than passed by value, with the watermarks and for the same reason: a + // captured resolve replays against whatever the device holds, not against what the host + // held when the graph was recorded. + mSparseStepAge = (mSparseStepAge + 1) & SparseEntry::kTimestampMask; + _buf(BufName::kSparseClock).front() = static_cast(mSparseStepAge); + // 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 @@ -1677,6 +1693,7 @@ void Batch::prepareSparseWriteIndices( resolve.inFlightSlot = scratch.inFlightSlot; resolve.maxInFlightUnits = scratch.maxInFlightUnits; resolve.missCount = scratch.missCount; + resolve.stepAge = _sparseStepAgeMirror(); // No selection on this path, so there is nothing to resolve into a slot. resolve.selected = nullptr; resolve.outSlot = nullptr; @@ -1865,6 +1882,7 @@ void Batch::resolveSparseSelection(std::vector const& buffers, int32_t resolve.maxInFlightUnits = scratch.maxInFlightUnits; resolve.work = _sparseFetchWork(scratch, span.selectionSlot); resolve.missCount = scratch.missCount; + resolve.stepAge = _sparseStepAgeMirror(); // Every row adds its own hits and misses as it resolves them. resolve.hitMiss = scratch.hitMiss; launchSparseResolve(resolve, stream); 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 d00ee29f5cac..7d5ef069da70 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 @@ -719,6 +719,7 @@ class Batch kPageTables = 0, kHistory, kCapacity, + kSparseClock, kLeases, kScratch, kCount @@ -799,6 +800,9 @@ class Batch //! The same array where a kernel reads it, which is what the launches are handed. [[nodiscard]] SparseLeaseIndex const* _leaseMirror() const noexcept; + //! The staged eviction clock where a resolve reads it, for the same reason. + [[nodiscard]] int32_t const* _sparseStepAgeMirror() 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. @@ -880,6 +884,19 @@ class Batch //! because the step's clear must not touch it. CudaUniqPtr mHitMiss; + //! The eviction clock every resolve of this Batch reads. Advanced in publish(), which + //! is the one place a step begins: a step resolves once per layer, so advancing per + //! resolve would age an entry by the layer count, and advancing from the watermarks + //! would quantise the clock by a life cycle's unit size. + //! + //! Shared by the life cycles. A single clock is what lets one Batch-wide advance serve + //! them, and nothing compares an entry of one against an entry of another. + //! + //! The host counter only; publish() stages it into BufName::kSparseClock, which is what + //! a resolve reads. See SparseResolveArgs::stepAge for why a kernel cannot take it by + //! value. + uint32_t mSparseStepAge = 0; + //! Device storage behind every life cycle's per-BufferId copy descriptors. CudaUniqPtr mCopyDesc; 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 b4eade03c7e4..0bb69533c733 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 @@ -263,6 +263,34 @@ CudaUniqPtr<> deviceCopy(std::vector const& values) return owned; } +CudaUniqPtr<> deviceCopyWords(std::vector const& values) +{ + CUstream const stream = nullptr; + auto owned = allocateDeviceMemory(sizeof(uint32_t) * values.size()); + if (values.empty()) + { + return owned; + } + cuCheck(cuMemcpyHtoDAsync( + reinterpret_cast(owned.get()), values.data(), sizeof(uint32_t) * values.size(), stream)); + cuCheck(cuStreamSynchronize(stream)); + return owned; +} + +std::vector readBackWords(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(uint32_t) * count, stream)); + cuCheck(cuStreamSynchronize(stream)); + return out; +} + std::vector readBack(CudaUniqPtr<> const& buffer, size_t count) { CUstream const stream = nullptr; @@ -280,12 +308,15 @@ std::vector readBack(CudaUniqPtr<> const& buffer, size_t count) } // namespace KvCacheIntrospection::SparseResolveProbe KvCacheIntrospection::sparseResolve( - std::vector const& residentOrdinal, std::vector const& lruOrder, int numLayers, + std::vector const& residentOrdinal, std::vector const& entryAge, 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; + // Far enough into the clock that every age the probe accepts has a timestamp of its own + // below it, so the state it builds needs no wrap. + auto const stepAge = SparseEntry::kMaxAge; auto const numRows = static_cast(historyUnits.size()); auto const perRowEntries = static_cast(numLayers) * entriesPerLayer; auto const entryCount = perRowEntries * static_cast(numRows); @@ -293,9 +324,14 @@ KvCacheIntrospection::SparseResolveProbe KvCacheIntrospection::sparseResolve( 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 + && residentOrdinal.size() == entryCount && (entryAge.empty() || entryAge.size() == entryCount) && selected.size() == selectionCount, "sparseResolve: input sizes must match the declared shape"); + for (auto const age : entryAge) + { + TLLM_CHECK_WITH_INFO(age >= 0 && static_cast(age) <= SparseEntry::kMaxAge, + "sparseResolve: an entry age must be within the range the eviction clock keeps distinguishable"); + } auto selection = deviceCopy(selected); auto outSlot = deviceCopy(std::vector(selectionCount, 0)); @@ -310,19 +346,19 @@ KvCacheIntrospection::SparseResolveProbe KvCacheIntrospection::sparseResolve( 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) + // One metadata block per lease: the two caller arrays packed into the entry words the + // arena actually holds. The probe takes and returns them separately, since an ordinal + // and an age are what a test states. + std::vector metadataHost(entryCount); + for (size_t i = 0; i < entryCount; ++i) { - 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 const ordinal = residentOrdinal[i]; + auto const age = entryAge.empty() ? 0U : static_cast(entryAge[i]); + metadataHost[i] = ordinal < 0 + ? SparseEntry::kFreeWord + : SparseEntry::pack(SparseEntry::usedByte(stepAge - age, false), static_cast(ordinal)); } - auto metadata = deviceCopy(metadataHost); + auto metadata = deviceCopyWords(metadataHost); // Row r takes lease r, except an inactive row, which holds none for the kernel to skip. std::vector leaseHost(static_cast(numRows)); @@ -333,11 +369,13 @@ KvCacheIntrospection::SparseResolveProbe KvCacheIntrospection::sparseResolve( : row; } auto leases = deviceCopy(leaseHost); + // Device-resident, as the Batch publishes it: the resolve reads the clock rather than + // taking it by value. + auto stepAgeBuffer = deviceCopy(std::vector(1, static_cast(stepAge))); SparseArenaView arena; arena.metadataBase = metadata.get(); - arena.leaseStride = sizeof(int32_t) * leaseInts; - arena.lruOffset = static_cast(perRowEntries); + arena.leaseStride = sizeof(uint32_t) * 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. @@ -380,21 +418,24 @@ KvCacheIntrospection::SparseResolveProbe KvCacheIntrospection::sparseResolve( args.work.capacity = numRows * selectionWidth; args.work.numLists = 1; args.missCount = reinterpret_cast(missCount.get()) + countOffset; + args.stepAge = reinterpret_cast(stepAgeBuffer.get()); launchSparseResolve(args, stream); } cuCheck(cuStreamSynchronize(stream)); SparseResolveProbe result; - auto const metadataOut = readBack(metadata, leaseInts * static_cast(numRows)); + auto const metadataOut = readBackWords(metadata, entryCount); result.residentOrdinal.resize(entryCount); - result.lruOrder.resize(entryCount); - for (int row = 0; row < numRows; ++row) + result.entryAge.resize(entryCount); + result.entryHit.resize(entryCount); + for (size_t i = 0; i < entryCount; ++i) { - 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); + auto const word = metadataOut[i]; + auto const ordinal = SparseEntry::ordinal(word); + result.residentOrdinal[i] = ordinal; + result.entryAge[i] = ordinal < 0 ? -1 : static_cast(SparseEntry::age(word, stepAge)); + result.entryHit[i] = ordinal < 0 ? 0 : static_cast(SparseEntry::wasHit(word)); } result.outSlot = readBack(outSlot, selectionCount); result.missRow = readBack(missRow, selectionCount); @@ -475,7 +516,6 @@ std::vector KvCacheIntrospection::sparseUnitCopy(std::vector c // 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; @@ -552,11 +592,13 @@ KvCacheIntrospection::SparseFlushList KvCacheIntrospection::sparseFlushList(std: 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); + // one lease laid out as the arena lays one. The scan reads only the ordinals, so the + // age each word carries is left at zero. + std::vector metadataHost(entryCount); + std::transform(residentOrdinal.begin(), residentOrdinal.end(), metadataHost.begin(), + [](int32_t ordinal) + { return ordinal < 0 ? SparseEntry::kFreeWord : SparseEntry::pack(0, static_cast(ordinal)); }); + auto metadata = deviceCopyWords(metadataHost); auto leases = deviceCopy(std::vector(1, 0)); // A layer per metadata row, in order: the probe drives unshared selections. @@ -566,8 +608,7 @@ KvCacheIntrospection::SparseFlushList KvCacheIntrospection::sparseFlushList(std: SparseArenaView arena; arena.metadataBase = metadata.get(); - arena.leaseStride = sizeof(int32_t) * 2 * entryCount; - arena.lruOffset = static_cast(entryCount); + arena.leaseStride = sizeof(uint32_t) * entryCount; arena.numLayers = numLayers; arena.numSelections = numLayers; arena.entriesPerSelection = entriesPerLayer; 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 0d5be223805d..b4fa08e93d8f 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 @@ -110,13 +110,18 @@ class KvCacheIntrospection 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. + // probe index and victim choice 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 + //! [layer][entry] steps since each entry's last use, and whether that use found it + //! already resident. Together they are the eviction order: an entry is given away + //! before an entry of the same age that was hit, and before any younger one. -1 and + //! false where the entry is free. + std::vector entryAge; + std::vector entryHit; + 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; @@ -127,8 +132,11 @@ class KvCacheIntrospection std::vector inFlightSlot; //!< [layer][maxInFlightUnits] }; + //! @param entryAge [layer][entry] steps since each entry's last use, as the state the + //! resolve starts from. Empty leaves every entry at zero, which makes the eviction order + //! among them arbitrary. static SparseResolveProbe sparseResolve(std::vector const& residentOrdinal, - std::vector const& lruOrder, int numLayers, int entriesPerLayer, std::vector const& selected, + std::vector const& entryAge, int numLayers, int entriesPerLayer, std::vector const& selected, int selectionWidth, std::vector const& historyUnits, std::vector const& capacityUnits, int maxInFlightUnits, std::vector const& inactiveRows); 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 77a7e427b91f..1ae2defe95d2 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 @@ -181,7 +181,7 @@ void KvCacheManager::_computeSparseCacheShape() // 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 + // `entry + 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]; 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 index b92ea80706e8..e0dfd7820236 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/sparseCache.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/sparseCache.cpp @@ -20,7 +20,7 @@ #include "kv_cache_manager_v2/exceptions.h" #include "kv_cache_manager_v2/utils/math.h" -#include +#include #include namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 @@ -40,26 +40,12 @@ size_t alignedBytes(size_t bytes) noexcept namespace { -//! Per-lease metadata, in carve order. Arrays shared by every lease of a life cycle live +//! Bytes one lease's entry words take. 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) +size_t leaseBytes(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; + return alignedBytes(sizeof(uint32_t) * entries); } } // namespace @@ -107,9 +93,8 @@ SparseCellArena::SparseCellArena(SparseCacheInit const& shape, int maxLeases) 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(); + mMetadataBytes = leaseBytes(mShape); // One region per pool, each aligned so that a buffer's cell indices stay relative to its own // pool's base. @@ -121,12 +106,10 @@ SparseCellArena::SparseCellArena(SparseCacheInit const& shape, int maxLeases) 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; + size_t const sharedBytes = layerBytes + slotBytes + selectionBytes; mNumBytes = cellBytes + sharedBytes + static_cast(maxLeases) * mMetadataBytes; @@ -142,7 +125,6 @@ SparseCellArena::SparseCellArena(SparseCacheInit const& shape, int maxLeases) 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(), @@ -152,25 +134,15 @@ SparseCellArena::SparseCellArena(SparseCacheInit const& shape, int maxLeases) 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. + // lease back clean. 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. + // stream only, and the host vectors it reads die with this scope. cuCheck(cuStreamSynchronize(stream)); } @@ -184,24 +156,9 @@ void SparseCellArena::initialiseLease(SparseLeaseIndex slot, CUstream stream) "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)); + cuCheck(cuMemsetD32Async(reinterpret_cast(mMetadataBase + static_cast(index) * mMetadataBytes), + SparseEntry::kFreeWord, entries, stream)); } SparseArenaView SparseCellArena::view() const noexcept @@ -209,8 +166,6 @@ 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; @@ -232,11 +187,13 @@ std::vector SparseCellArena::peekResidentOrdinals(SparseLeaseIndex slot // knows which stream last touched the slot. cuCheck(cuCtxSynchronize()); auto const arena = view(); - std::vector out(arena.numEntries()); - cuCheck(cuMemcpyDtoHAsync(out.data(), + std::vector words(arena.numEntries()); + cuCheck(cuMemcpyDtoHAsync(words.data(), reinterpret_cast(arena.metadataBase + static_cast(slot.value()) * arena.leaseStride), - sizeof(int32_t) * out.size(), stream)); + sizeof(uint32_t) * words.size(), stream)); cuCheck(cuStreamSynchronize(stream)); + std::vector out(words.size()); + std::transform(words.begin(), words.end(), out.begin(), [](uint32_t word) { return SparseEntry::ordinal(word); }); return out; } 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 index ddd95d21e45b..ed723ec7481f 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/sparseCache.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/sparseCache.h @@ -42,10 +42,109 @@ namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 // // 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. +// resident set serves 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. // --------------------------------------------------------------------------- + +//! One cache entry's whole persistent state, packed into a single word: +//! +//! bits 31..25 timestamp, the step clock value at the entry's last use +//! bit 24 history, set when that use found the entry already resident +//! bits 23..0 ordinal, the unit held; kFreeOrdinal means the entry is free +//! +//! The top byte is addressed together as the "age byte", `(timestamp << 1) | history`. +//! +//! The history flag says the entry was already resident at its last use, so eviction by +//! highest `staleness()` -- which is `2 * age - history` -- orders by age first and, within +//! one age, gives away what was fetched before what was hit. +//! +//! The step clock wraps, so a timestamp more than `kMaxAge + kRebaseStep` behind it would +//! read as a recent one. `rebase()` is what keeps it from getting that far: it snaps a +//! timestamp forward to `kMaxAge` and is idempotent, so the kernels that stream a row may +//! each apply it without having to agree on which of them does. +struct SparseEntry +{ + static constexpr uint32_t kOrdinalBits = 24; + static constexpr uint32_t kOrdinalMask = (1U << kOrdinalBits) - 1U; + //! The ordinal field of an entry holding no unit. + static constexpr uint32_t kFreeOrdinal = kOrdinalMask; + //! An entry holding no unit, which is the state initialiseLease() leaves a lease in. + static constexpr uint32_t kFreeWord = kFreeOrdinal; + + //! Values the timestamp field holds, and the mask that wraps a step count into it. + static constexpr uint32_t kTimestampValues = 128; + static constexpr uint32_t kTimestampMask = kTimestampValues - 1; + + //! Steps between rebases, and the age a rebase leaves an older entry at. Their sum is + //! the oldest age still distinguishable from a recent one, so it must not exceed the + //! timestamp's own range. + static constexpr uint32_t kRebaseStep = 32; + static constexpr uint32_t kMaxAge = kTimestampValues - 1 - kRebaseStep; + + //! Whether the resolves of `stepAge` rebase the timestamps they pass. Derived from the + //! clock rather than published beside it, so the two cannot disagree. + [[nodiscard]] static constexpr bool rebaseDue(uint32_t stepAge) noexcept + { + return stepAge % kRebaseStep == 0; + } + + //! `ordinal`'s high bits are dropped, so a word may be passed where one is wanted. + [[nodiscard]] static constexpr uint32_t pack(uint32_t ageByte, uint32_t ordinal) noexcept + { + return (ageByte << kOrdinalBits) | (ordinal & kOrdinalMask); + } + + [[nodiscard]] static constexpr uint32_t ageByte(uint32_t word) noexcept + { + return word >> kOrdinalBits; + } + + //! The unit held, or -1 -- kBadBlockOrdinal's value -- when the entry is free. + [[nodiscard]] static constexpr int32_t ordinal(uint32_t word) noexcept + { + return (word & kOrdinalMask) == kFreeOrdinal ? -1 : static_cast(word & kOrdinalMask); + } + + //! Steps since the entry's last use. Exact while that is at most kTimestampMask, which + //! is what the rebase maintains. + [[nodiscard]] static constexpr uint32_t age(uint32_t word, uint32_t stepAge) noexcept + { + return (stepAge - (ageByte(word) >> 1)) & kTimestampMask; + } + + [[nodiscard]] static constexpr bool wasHit(uint32_t word) noexcept + { + return (ageByte(word) & 1U) != 0U; + } + + //! Eviction key, highest first. Stays below the free entry's own key for every age the + //! rebase leaves representable, so a free entry is always given away first. + [[nodiscard]] static constexpr uint32_t staleness(uint32_t word, uint32_t stepAge) noexcept + { + return static_cast((stepAge << 1) - ageByte(word)); + } + + //! The age byte of an entry used at `stepAge`. `hit` is set when the entry already held + //! the unit and clear when the unit was fetched into it. + [[nodiscard]] static constexpr uint32_t usedByte(uint32_t stepAge, bool hit) noexcept + { + return ((stepAge & kTimestampMask) << 1) | (hit ? 1U : 0U); + } + + //! The word with its age pulled up to kMaxAge if it had fallen further behind. Applying + //! this twice at one `stepAge` changes nothing the second time. + [[nodiscard]] static constexpr uint32_t rebase(uint32_t word, uint32_t stepAge) noexcept + { + return age(word, stepAge) <= kMaxAge + ? word + : pack((((stepAge - kMaxAge) & kTimestampMask) << 1) | (ageByte(word) & 1U), word); + } +}; + +static_assert(SparseEntry::kMaxAge + SparseEntry::kRebaseStep < SparseEntry::kTimestampValues, + "A timestamp must stay distinguishable for a whole rebase period past the oldest age the rebase leaves"); + //! 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 @@ -60,15 +159,13 @@ inline constexpr SparseLeaseIndex kNoSparseLease{}; //! 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. + //! Lease i's metadata begins at `metadataBase + i * leaseStride`: one SparseEntry word + //! per entry, holding the unit resident there and when it was last used. //! - //! 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". + //! The words are dense rather than keyed: a fetch builds a hash over their ordinals 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. @@ -184,9 +281,9 @@ class SparseCellArena 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. + //! Clear `slot`'s metadata on `stream`, which is what taking it means here: every entry + //! becomes free. 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 @@ -218,7 +315,6 @@ class SparseCellArena 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 index 45700f319f91..5b2cfcb817a6 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/sparseCacheKernels.cu +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/sparseCacheKernels.cu @@ -27,8 +27,12 @@ #include #include +#include +#include +#include #include #include +#include namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 { @@ -47,32 +51,38 @@ namespace } \ } while (0) +//! Threads a copy, out-index or flush-list CTA takes. The resolve is templated on its own +//! width instead; see resolveThreads(). constexpr int kThreads = 256; +//! The two widths the resolve is instantiated at. Nothing wider was ever faster at any +//! shape measured, and the narrow one must still cover the eviction histogram. +constexpr int kResolveThreadsNarrow = 256; +constexpr int kResolveThreadsWide = 512; 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`. +//! A lease's metadata block: one SparseEntry word per entry. //! //! 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) +__device__ inline uint32_t* leaseMetadata(SparseArenaView const& arena, SparseLeaseIndex lease) { - return reinterpret_cast(arena.metadataBase + static_cast(lease.value()) * arena.leaseStride); + 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; +//! What the resolve knows about an entry while it is choosing victims. +constexpr uint8_t kFlagEvictable = 0; //!< Not selected this step and holding no in-flight unit. +constexpr uint8_t kFlagMatched = 1; //!< Selected this step, so its unit stays and its age resets. +constexpr uint8_t kFlagHeld = 2; //!< Holding in-flight data, which only this entry has. + +//! Buckets the eviction histogram takes, which is the whole range of SparseEntry::staleness. +constexpr int32_t kAgeBuckets = 256; +//! The bucket a free entry falls in. No resident entry reaches it, so the free ones are +//! always given away before any unit is evicted. +constexpr int32_t kStalest = kAgeBuckets - 1; using BlockScan = cub::BlockScan; @@ -81,19 +91,17 @@ __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. +//! Byte offsets of the shared-memory regions one resolve CTA carves. Computed identically +//! on the host, which sizes the launch, and in the kernel, which carves it. Every region is +//! 16-byte aligned: flag 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 table = 0; + size_t entry = 0; + size_t flag = 0; size_t missIdx = 0; size_t missOrd = 0; + size_t hist = 0; size_t total = 0; }; @@ -103,22 +111,18 @@ __host__ __device__ inline SmemLayout smemLayout(int32_t tableSize, int32_t entr { 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; + layout.table = offset; + offset = alignUp16(offset + sizeof(uint32_t) * static_cast(tableSize)); + layout.entry = offset; + offset = alignUp16(offset + sizeof(uint32_t) * static_cast(entries)); + layout.flag = 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.hist = offset; + offset = alignUp16(offset + sizeof(int32_t) * static_cast(kAgeBuckets)); layout.total = offset; return layout; } @@ -161,54 +165,100 @@ __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) +//! Bits of a probe-index position given to the entry holding the keyed ordinal, leaving the +//! rest for the ordinal itself. An ordinal of every one bit would be indistinguishable from +//! an empty position, and a slot of every one bit is what "keyed, but not resident" is, so +//! neither field reaches its own maximum. +constexpr uint32_t kTableSlotBits = 14; +constexpr uint32_t kTableSlotMask = (1U << kTableSlotBits) - 1U; +constexpr uint32_t kTableEmpty = 0xFFFFFFFFU; +constexpr int32_t kMaxTableOrdinal = static_cast((1U << (32 - kTableSlotBits)) - 2U); +constexpr int32_t kMaxTableSlot = static_cast(kTableSlotMask) - 1; + +//! The resolve's open-addressed probe index, one word per position. +//! +//! Packed rather than two arrays because every pass over it is a shared-memory round trip +//! per probe step, and the publish pass walks the whole resident set: halving the bytes a +//! position takes halves the traffic and leaves twice the index in one cache line. +//! +//! Insert races; publish and lookup do not, because they run between syncs that separate +//! them from the inserts. A position is written once by insert and at most once by publish, +//! each by the single thread owning the ordinal. +struct QueryTable { - 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) + uint32_t* word; + + __device__ explicit QueryTable(unsigned char* base) + : word(reinterpret_cast(base)) + { + } + + __device__ void clear(int32_t position) const { - int32_t const idx = begin + tile * kWarpSize + lane; - uint32_t const votes = __ballot_sync(0xffffffffU, idx < n && flags[idx] != 0); - if (idx < n) + word[position] = kTableEmpty; + } + + __device__ uint32_t find(int32_t ordinal, uint32_t mask, uint32_t& position) const + { + position = slotOf(ordinal, mask); + uint32_t current = word[position]; + while (current != kTableEmpty && (current >> kTableSlotBits) != static_cast(ordinal)) { - flags[idx] = carry + __popc(votes & precedingLanes); + position = (position + 1U) & mask; + current = word[position]; } - carry += __popc(votes); + return current; } - if (lane == 0) + + //! Key `ordinal`, with no entry named yet. + __device__ void insert(int32_t ordinal, uint32_t mask) const { - warpTotals[warp] = carry; + uint32_t const want = (static_cast(ordinal) << kTableSlotBits) | kTableSlotMask; + uint32_t position = slotOf(ordinal, mask); + uint32_t previous = atomicCAS(&word[position], kTableEmpty, want); + while (previous != kTableEmpty && (previous >> kTableSlotBits) != static_cast(ordinal)) + { + position = (position + 1U) & mask; + previous = atomicCAS(&word[position], kTableEmpty, want); + } } - __syncthreads(); - int32_t offset = 0; - int32_t total = 0; -#pragma unroll - for (int32_t w = 0; w < kWarps; ++w) + //! Name `slot` as the entry holding `ordinal`, if the ordinal was keyed. Returns whether + //! it was, which is what says the entry is selected this step. + __device__ bool publish(int32_t ordinal, int32_t slot, uint32_t mask) const { - int32_t const count = warpTotals[w]; - offset += w < warp ? count : 0; - total += count; + uint32_t position = 0; + if (find(ordinal, mask, position) == kTableEmpty) + { + return false; + } + word[position] = (static_cast(ordinal) << kTableSlotBits) | static_cast(slot); + return true; } - for (int32_t tile = 0; tile < tiles; ++tile) + + //! The entry holding `ordinal`, or kNoSparseSlot when nothing does. + __device__ int32_t lookup(int32_t ordinal, uint32_t mask) const { - int32_t const idx = begin + tile * kWarpSize + lane; - if (idx < n) + uint32_t position = 0; + uint32_t const current = find(ordinal, mask, position); + if (current == kTableEmpty) { - flags[idx] += offset; + return kNoSparseSlot; } + uint32_t const slot = current & kTableSlotMask; + return slot == kTableSlotMask ? kNoSparseSlot : static_cast(slot); } - __syncthreads(); - return total; +}; + +//! Bits of the prepare-write reduction key given to the entry index, the rest to the +//! staleness plus one -- so that zero is a key no candidate produces and can stand for a +//! block that found none. +constexpr uint32_t kVictimSlotBits = 23; +constexpr uint32_t kVictimSlotMask = (1U << kVictimSlotBits) - 1U; + +__device__ inline uint32_t victimKey(uint32_t staleness, int32_t slot) noexcept +{ + return ((staleness + 1U) << kVictimSlotBits) | static_cast(slot); } //! Reserve one in-flight unit without building a selection probe index. @@ -225,6 +275,9 @@ __global__ __launch_bounds__(kThreads) void sparsePrepareWriteKernel(SparseResol 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); + // The same bound the resolve's probe index imposes, since this ordinal is resolved + // through it on every later step. + KVCM2_SPARSE_REQUIRE(historyUnits <= kMaxTableOrdinal); if (threadIdx.x == 0) { args.missCount[row] = 0; @@ -236,57 +289,52 @@ __global__ __launch_bounds__(kThreads) void sparsePrepareWriteKernel(SparseResol } 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; + auto const stepAge = static_cast(*args.stepAge); + bool const rebase = SparseEntry::rebaseDue(stepAge); + auto* entry = leaseMetadata(args.arena, lease) + static_cast(args.selectionSlot) * entries; + 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) + // The entry already holding the in-flight unit, if one does, and the stalest of those + // that could be given to it otherwise -- free entries first, then history by age. An + // entry above the watermark is another unit's only copy and is never a candidate. + // + // Both reductions take a key that is zero where a thread saw no candidate, which is + // below every key one did produce. + uint32_t match = 0; + uint32_t stalest = 0; + for (int32_t s = static_cast(threadIdx.x); s < entries; s += kThreads) { - int32_t const slot = order[p]; - int32_t const ordinal = resident[slot]; - previousOrder[p] = slot; + uint32_t word = entry[s]; + if (rebase) + { + word = SparseEntry::rebase(word, stepAge); + entry[s] = word; + } + int32_t const ordinal = SparseEntry::ordinal(word); if (ordinal == historyUnits) { - match = p; + match = static_cast(s) + 1U; } - if (ordinal == kFreeOrdinal || ordinal < historyUnits) + else if (ordinal < historyUnits) { - victim = p; + auto const staleness + = ordinal < 0 ? static_cast(kStalest) : SparseEntry::staleness(word, stepAge); + stalest = max(stalest, victimKey(staleness, s)); } } - int32_t const matched = Reduce(temp).Reduce(match, cuda::maximum<>{}); - if (threadIdx.x == 0) - { - matchPosition = matched; - } + uint32_t const victim = Reduce(temp).Reduce(stalest, cuda::maximum<>{}); __syncthreads(); - int32_t const lastVictim = Reduce(temp).Reduce(victim, cuda::maximum<>{}); + uint32_t const matched = Reduce(temp).Reduce(match, 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; + // The reservation covers a full selection plus every in-flight unit, so a candidate + // exists whenever the unit is not already held. + KVCM2_SPARSE_REQUIRE(matched != 0 || victim != 0); + auto const slot = static_cast(matched != 0 ? matched - 1U : (victim & kVictimSlotMask)); + entry[slot] + = SparseEntry::pack(SparseEntry::usedByte(stepAge, matched != 0), static_cast(historyUnits)); + args.inFlightSlot[row] = slot; } } @@ -294,8 +342,16 @@ __global__ __launch_bounds__(kThreads) void sparsePrepareWriteKernel(SparseResol //! //! 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) +//! +//! `kResolveThreads` is chosen per shape by resolveThreads(): every pass but the probe is a +//! stride loop over the resident set, so a wider CTA shortens the kernel wherever a +//! multiprocessor has the room to hold more warps of it. +template +__global__ __launch_bounds__(kResolveThreads) void sparseResolveKernel(SparseResolveArgs args, int32_t tableSize) { + // The threshold search below gives one thread to each histogram bucket. + static_assert(kResolveThreads >= kAgeBuckets, "A resolve CTA must cover the eviction histogram"); + using Scan = cub::BlockScan; extern __shared__ unsigned char smem[]; int32_t const selection = args.selectionSlot; int32_t const row = static_cast(blockIdx.x); @@ -309,6 +365,9 @@ __global__ __launch_bounds__(kThreads) void sparseResolveKernel(SparseResolveArg int32_t const entries = args.arena.entriesPerSelection; int32_t const width = args.selectionWidth; uint32_t const mask = static_cast(tableSize) - 1U; + int32_t const tid = static_cast(threadIdx.x); + auto const stepAge = static_cast(*args.stepAge); + bool const rebase = SparseEntry::rebaseDue(stepAge); int32_t const historyUnits = args.historyTokens[row] / args.unitTokens; int32_t const inFlight = args.capacityTokens[row] / args.unitTokens - historyUnits; @@ -319,25 +378,28 @@ __global__ __launch_bounds__(kThreads) void sparseResolveKernel(SparseResolveArg 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); + QueryTable const table(smem + layout.table); + auto* entryCache = reinterpret_cast(smem + layout.entry); + auto* flag = reinterpret_cast(smem + layout.flag); auto* missIdx = reinterpret_cast(smem + layout.missIdx); auto* missOrd = reinterpret_cast(smem + layout.missOrd); + auto* hist = reinterpret_cast(smem + layout.hist); - __shared__ typename BlockScan::TempStorage temp; + __shared__ typename Scan::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 + __shared__ int32_t copyBase; + // The staleness at which the misses run out, and how many of the entries holding + // exactly that staleness are wanted. A threshold past the histogram evicts nothing. + __shared__ int32_t thresholdStaleness; + __shared__ int32_t neededAtThreshold; + __shared__ int32_t takenAtThreshold; + __shared__ int32_t victimCounter; + + // The lease's own words 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; + uint32_t* entry = leaseMetadata(args.arena, lease) + 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; @@ -346,16 +408,23 @@ __global__ __launch_bounds__(kThreads) void sparseResolveKernel(SparseResolveArg int32_t* outSlot = args.outSlot + selectionBase; int32_t* inFlightSlot = args.inFlightSlot + inFlightBase; - if (threadIdx.x == 0) + if (tid == 0) { numMiss = 0; numCopy = 0; numResolved = 0; + thresholdStaleness = kAgeBuckets; + neededAtThreshold = 0; + takenAtThreshold = 0; + victimCounter = 0; + } + for (int32_t i = tid; i < tableSize; i += kResolveThreads) + { + table.clear(i); } - for (int32_t i = static_cast(threadIdx.x); i < tableSize; i += kThreads) + for (int32_t b = tid; b < kAgeBuckets; b += kResolveThreads) { - keys[i] = kEmptyKey; - vals[i] = kNoSparseSlot; + hist[b] = 0; } __syncthreads(); @@ -365,57 +434,57 @@ __global__ __launch_bounds__(kThreads) void sparseResolveKernel(SparseResolveArg // 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) + for (int32_t i = tid; i < required; i += kResolveThreads) { 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); - } + // A position of the index holds the ordinal beside the entry naming it, so an + // ordinal past this would be truncated into a different one. + KVCM2_SPARSE_REQUIRE(ord <= kMaxTableOrdinal); + table.insert(ord, mask); + } + + // 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 = tid + inFlight; i < args.maxInFlightUnits; i += kResolveThreads) + { + inFlightSlot[i] = kNoSparseSlot; } __syncthreads(); - for (int32_t s = static_cast(threadIdx.x); s < entries; s += kThreads) + // Classify the resident set against the index, and build the eviction histogram over + // what the step may give away. The words land in shared memory on the way past: two + // passes below read them again, and each read would otherwise cross to global memory. + for (int32_t s = tid; s < entries; s += kResolveThreads) { - int32_t const ord = residentOrdinal[s]; - cls[s] = kClassUntouched | (ord >= historyUnits ? kPinned : 0); - previousLru[s] = lruOrder[s]; - if (ord == kFreeOrdinal) + uint32_t word = entry[s]; + if (rebase) { - continue; + word = SparseEntry::rebase(word, stepAge); + entry[s] = word; } - uint32_t pos = slotOf(ord, mask); - int32_t key = keys[pos]; - while (key != kEmptyKey && key != ord) + entryCache[s] = word; + int32_t const ordinal = SparseEntry::ordinal(word); + bool const matched = ordinal >= 0 && table.publish(ordinal, s, mask); + bool const held = ordinal >= historyUnits; + uint8_t const cls = matched ? kFlagMatched : (held ? kFlagHeld : kFlagEvictable); + flag[s] = cls; + if (cls == kFlagEvictable) { - pos = (pos + 1U) & mask; - key = keys[pos]; + auto const staleness + = ordinal < 0 ? static_cast(kStalest) : SparseEntry::staleness(word, stepAge); + atomicAdd(&hist[staleness], 1); } - 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) + for (int32_t i = tid; i < required; i += kResolveThreads) { int32_t const ord = i < inFlight ? historyUnits + i : selected[i - inFlight]; if (ord < 0) @@ -424,18 +493,10 @@ __global__ __launch_bounds__(kThreads) void sparseResolveKernel(SparseResolveArg continue; } resolvedMine += i < inFlight ? 0 : 1; - uint32_t pos = slotOf(ord, mask); - int32_t key = keys[pos]; - while (key != kEmptyKey && key != ord) + int32_t const slot = table.lookup(ord, mask); + if (slot != kNoSparseSlot) { - 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 { @@ -454,17 +515,16 @@ __global__ __launch_bounds__(kThreads) void sparseResolveKernel(SparseResolveArg // 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) + for (int32_t j = tid; j < numMiss; j += kResolveThreads) { mine += missOrd[j] < historyUnits ? 1 : 0; } int32_t scanned = 0; int32_t claimed = 0; - BlockScan(temp).ExclusiveSum(mine, scanned, claimed); - if (threadIdx.x == 0) + Scan(temp).ExclusiveSum(mine, scanned, claimed); + if (tid == 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 @@ -479,130 +539,72 @@ __global__ __launch_bounds__(kThreads) void sparseResolveKernel(SparseResolveArg } __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) + // Where the misses run out, walking the histogram down from its stalest bucket. One + // thread per bucket, so the scan that totals them is the whole search. { - int32_t const slot = previousLru[p]; - scratch[p] = cls[slot] == kClassUntouched ? 1 : 0; + int32_t const bucket = tid < kAgeBuckets ? hist[kStalest - tid] : 0; + int32_t inclusive = 0; + Scan(temp).InclusiveSum(bucket, inclusive); + int32_t const before = inclusive - bucket; + if (tid < kAgeBuckets && numMiss > 0 && before < numMiss && inclusive >= numMiss) + { + thresholdStaleness = kStalest - tid; + neededAtThreshold = numMiss - before; + } } __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(); + KVCM2_SPARSE_REQUIRE(numMiss == 0 || thresholdStaleness < kAgeBuckets); - for (int32_t p = static_cast(threadIdx.x); p < entries; p += kThreads) + // Give the misses the stalest entries. Everything staler than the threshold goes; the + // bucket at it holds more entries than are left to fill, so those race for the places. + for (int32_t s = tid; s < entries; s += kResolveThreads) { - int32_t const slot = previousLru[p]; - bool const evictable = cls[slot] == kClassUntouched; - if (!evictable) + uint32_t const word = entryCache[s]; + uint8_t const cls = flag[s]; + if (cls == kFlagMatched) { + entry[s] = SparseEntry::pack(SparseEntry::usedByte(stepAge, true), word); 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) + if (cls != kFlagEvictable) { - uint32_t const votes = __ballot_sync(0xffffffffU, c == group); - if (c == group) - { - order[p] = carry[group] + __popc(votes & precedingLanes); - } - carry[group] += __popc(votes); + continue; } - } - if (lane == 0) - { -#pragma unroll - for (uint8_t c = 0; c < 3; ++c) + auto const staleness = static_cast( + SparseEntry::ordinal(word) < 0 ? static_cast(kStalest) : SparseEntry::staleness(word, stepAge)); + bool victim = staleness > thresholdStaleness; + if (!victim && staleness == thresholdStaleness) { - classTotals[c][warp] = carry[c]; + victim = atomicAdd(&takenAtThreshold, 1) < neededAtThreshold; } - } - __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) + if (!victim) { - int32_t const total = classTotals[c][w]; - before += w < warp ? total : 0; - count += total; + continue; } - 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 j = atomicAdd(&victimCounter, 1); + int32_t const ordinal = missOrd[j]; + int32_t const position = missIdx[j]; + entry[s] = SparseEntry::pack(SparseEntry::usedByte(stepAge, false), static_cast(ordinal)); + (position < inFlight ? inFlightSlot[position] : outSlot[position - inFlight]) = s; + // 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) { - 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; + auto const c = copyBase + atomicAdd(&numCopy, 1); + args.work.row[c] = row; + args.work.ordinal[c] = ordinal; + args.work.slot[c] = s; } } __syncthreads(); - for (int32_t p = static_cast(threadIdx.x); p < entries; p += kThreads) - { - lruOrder[p] = scratch[p]; - } - if (threadIdx.x == 0) + // The work list and the entry indices are final, so a copy may start. + gridDependencyRelease(); + + if (tid == 0) { args.missCount[row] = numCopy; } @@ -801,11 +803,12 @@ __global__ __launch_bounds__(kThreads) void sparseFlushListKernel(SparseFlushLis // 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 + uint32_t const* entry = leaseMetadata(args.arena, lease) + static_cast(args.arena.layerSelection[layer]) * entries; + auto const ordinalAt = [&](int32_t slot) { return SparseEntry::ordinal(entry[slot]); }; auto const crossed = [&](int32_t slot) { - int32_t const ord = residentOrdinal[slot]; + int32_t const ord = ordinalAt(slot); return ord >= from && ord < to; }; @@ -851,7 +854,7 @@ __global__ __launch_bounds__(kThreads) void sparseFlushListKernel(SparseFlushLis } auto const j = listBase + atomicAdd(&count, 1); args.work.row[j] = row; - args.work.ordinal[j] = residentOrdinal[s]; + args.work.ordinal[j] = ordinalAt(s); args.work.slot[j] = s; } @@ -992,9 +995,12 @@ 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. +//! The device is asked once and each instantiation 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. +//! +//! The narrower allowance of the two is what a caller is told, since the shape is sized +//! before the width that will run it is chosen. int32_t sparseResolveSharedLimit() { static int32_t const limit = [] @@ -1003,18 +1009,74 @@ int32_t sparseResolveSharedLimit() 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); + int allowed = deviceMax; + for (void const* function : {reinterpret_cast(&sparseResolveKernel), + reinterpret_cast(&sparseResolveKernel)}) + { + // The opt-in names dynamic bytes, and what an instantiation declares statically + // is spent from the same per-block budget, so only the rest of it can be asked + // for. The two differ: cub's scan storage follows the thread count. + cudaFuncAttributes attr{}; + TLLM_CUDA_CHECK(cudaFuncGetAttributes(&attr, function)); + int const value = deviceMax - static_cast(attr.sharedSizeBytes); + TLLM_CUDA_CHECK(cudaFuncSetAttribute(function, cudaFuncAttributeMaxDynamicSharedMemorySize, value)); + allowed = std::min(allowed, value); + } + return static_cast(allowed); }(); return limit; } +//! Threads a resolve CTA of this shape takes. +//! +//! Every pass but the probe is a stride loop over the resident set, so doubling the CTA +//! halves the entries a thread walks. A CTA's work is one row and does not grow with its +//! width, so that only pays where the wider block raises the warps a multiprocessor holds; +//! where it does not, the wider CTA costs residency and the kernel loses more to waves than +//! the shorter loops win. The device's own occupancy answer for each instantiation -- at +//! the dynamic shared memory this shape asks for, which is what drives the count -- is +//! therefore the whole decision, and a tie goes to the narrow one, whose smaller blocks +//! leave a smaller tail on a grid that does not fill the device. +//! +//! The three extents are the life cycle's, so the answer never changes within one; it is +//! held rather than asked per launch. +int resolveThreads(int32_t entriesPerSelection, int32_t selectionWidth, int32_t maxInFlightUnits, size_t shared) +{ + using Shape = std::tuple; + static std::mutex mutex; + static std::map decided; + + Shape const shape{entriesPerSelection, selectionWidth, maxInFlightUnits}; + std::lock_guard const guard(mutex); + auto const found = decided.find(shape); + if (found != decided.end()) + { + return found->second; + } + // The query adds the instantiation's own static shared memory to what is passed here, + // so the figure it works from is the block's whole footprint. + auto residentBlocks = [shared](void const* kernel, int threads) + { + int blocksPerSm = 0; + TLLM_CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor(&blocksPerSm, kernel, threads, shared)); + return blocksPerSm; + }; + int const narrowWarps = residentBlocks(reinterpret_cast(&sparseResolveKernel), + kResolveThreadsNarrow) + * (kResolveThreadsNarrow / static_cast(kWarpSize)); + int const wideWarps + = residentBlocks(reinterpret_cast(&sparseResolveKernel), kResolveThreadsWide) + * (kResolveThreadsWide / static_cast(kWarpSize)); + + // A CTA's work is one row and does not grow with its width, so a wider one is worth + // taking only where it raises the warps a multiprocessor holds. Where both widths reach + // the same occupancy the narrow one wins on grid quantisation, since a row count that + // does not fill the device leaves a smaller tail. + int const threads = wideWarps > narrowWarps ? kResolveThreadsWide : kResolveThreadsNarrow; + decided.emplace(shape, threads); + return threads; +} + } // namespace size_t checkedSparseResolveSharedBytes(int32_t entriesPerSelection, int32_t selectionWidth, int32_t maxInFlightUnits) @@ -1045,6 +1107,8 @@ void launchSparseResolve(SparseResolveArgs const& args, CUstream stream) && (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"); + TLLM_CHECK_WITH_INFO(args.stepAge != nullptr, + "launchSparseResolve: the eviction clock is device-resident, so a captured graph reads the step's own value"); if (args.selectionWidth > 0) { checkWorkList(args.work, "launchSparseResolve"); @@ -1052,14 +1116,18 @@ void launchSparseResolve(SparseResolveArgs const& args, CUstream stream) 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) + // A position of the probe index holds an entry index beside the ordinal keyed there, so + // a cache wider than the field would resolve two entries onto one position. + TLLM_CHECK_WITH_INFO(args.arena.entriesPerSelection <= kMaxTableSlot, + "Sparse resolve: the cache holds " + std::to_string(args.arena.entriesPerSelection) + + " entries per selection, above the " + std::to_string(kMaxTableSlot) + + " its probe index can name; reduce max_cached_units"); + + dim3 const grid(static_cast(args.numRows)); + if (args.selectionWidth == 0 && args.maxInFlightUnits == 1) { - 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); + tensorrt_llm::common::launchWithPdlWhenEnabled("sparsePrepareWrite", sparsePrepareWriteKernel, grid, + dim3(kThreads), 0, reinterpret_cast(stream), args); return; } @@ -1067,9 +1135,15 @@ void launchSparseResolve(SparseResolveArgs const& args, CUstream stream) 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); + if (resolveThreads(args.arena.entriesPerSelection, args.selectionWidth, args.maxInFlightUnits, shared) + == kResolveThreadsWide) + { + tensorrt_llm::common::launchWithPdlWhenEnabled("sparseResolve", sparseResolveKernel, grid, + dim3(kResolveThreadsWide), shared, reinterpret_cast(stream), args, tableSize); + return; + } + tensorrt_llm::common::launchWithPdlWhenEnabled("sparseResolve", sparseResolveKernel, grid, + dim3(kResolveThreadsNarrow), 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 index cde52f078d65..3fd33b7fd9bd 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/sparseCacheKernels.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/sparseCacheKernels.h @@ -164,6 +164,17 @@ struct SparseResolveArgs //! compacted list -- but it is what makes a row's hit rate observable. int32_t* missCount; + //! [1] the step this resolve belongs to, as a count that wraps into SparseEntry's + //! timestamp field. Advanced once per step for the whole Batch, not once per resolve: a + //! step runs one resolve per layer, and an entry's age is in steps. + //! + //! Device-resident for the reason the watermarks are: it moves every step, so a captured + //! graph must read the current value rather than one baked in at capture. + //! + //! Whether the step also rebases follows from the value -- see SparseEntry::rebaseDue -- + //! so there is nothing beside it that could disagree with it. + int32_t const* stepAge; + //! 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 @@ -178,7 +189,7 @@ struct SparseResolveArgs //! 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. +//! the staged entry words and their classification grow linearly with entriesPerSelection. [[nodiscard]] size_t sparseResolveSharedBytes( int32_t entriesPerSelection, int32_t selectionWidth, int32_t maxInFlightUnits); diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h index deb8884ad99d..4b303a0be53f 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h @@ -626,7 +626,7 @@ class TllmGenFmhaKernel TLLM_CHECK_WITH_INFO(findIter != functions.end(), "Trtllm-gen kernels not found: " + info); auto const& kernelMeta = mKernelMeta[findIter->second.mMetaInfoIndex]; - const CUfunction func = findIter->second.mDeviceFunction; + CUfunction const func = findIter->second.mDeviceFunction; // mGroupsHeadsQ, mGroupsTokensHeadsQ, and mNumInsts{Q,Kv} are part of the hashID, so the // matched cubin already agrees with the AutoTuner decision. Re-assign from kernelMeta @@ -1215,6 +1215,12 @@ class TllmGenFmhaKernel // Multi-CTA KV mode (bool -> enum conversion) options.mMultiCtasKvMode = params.mMultiCtasKvMode ? tensorrt_llm::kernels::MultiCtasKvMode::GmemReduction : tensorrt_llm::kernels::MultiCtasKvMode::Disabled; + if (tensorrt_llm::common::getEnvForceDeterministicAttention()) + { + // Keep one KV reduction partition even when the selector considers CGA reduction. + options.mMultiCtasKvMode = tensorrt_llm::kernels::MultiCtasKvMode::Disabled; + options.mMaxNumCtasPerSeqKv = 1; + } options.mStoresSoftmaxStats = true; // Attention features @@ -1577,7 +1583,7 @@ class TllmFmhaKernelFactory public: using KernelType = TllmGenFmhaKernel; - KernelType* getKernels(const typename KernelType::KernelMeta* pKernelList, unsigned int nbKernels, Data_type dtypeQ, + KernelType* getKernels(typename KernelType::KernelMeta const* pKernelList, unsigned int nbKernels, Data_type dtypeQ, Data_type dtypeK, Data_type dtypeV, Data_type dtypeOut, unsigned int sm, int numEltsPerSageAttnBlkQ = 0, int numEltsPerSageAttnBlkK = 0, int numEltsPerSageAttnBlkP = 0, int numEltsPerSageAttnBlkV = 0, bool fusesDsv4InvRopeFp8Quant = false) @@ -1650,7 +1656,7 @@ class TllmFmhaKernelFactory | (static_cast(fusesDsv4InvRopeFp8Quant) << 44); } - std::unordered_map> mKernels; + std::unordered_map const> mKernels; }; inline TllmGenFmhaKernel* getTllmFmhaKernels(Data_type dtypeQ, Data_type dtypeK, Data_type dtypeV, Data_type dtypeOut, diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp index f7d9ab23f7e0..04816843d636 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp @@ -69,18 +69,6 @@ namespace tensorrt_llm::nanobind::batch_manager namespace { -//! A scratch range paired with the slots covering it. -//! -//! KvCache reports the two separately, through getSwaScratchRange() and getScratchSlots(). -//! This pairs them for Python consumers that still expect one object, and exists only so -//! this change stands alone; remove it, and the get_scratch_desc() binding, once the full -//! sparse attention support lands and those consumers read the two accessors directly. -struct ScratchDescCompat -{ - kv::HalfOpenRange range; - std::vector slotIds; -}; - // Exposed via introspection sub-module for tests. class TestPaddingColdPageCodec final : public kv::IKvCacheColdPageCodec { @@ -1477,17 +1465,6 @@ 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); @@ -1949,15 +1926,8 @@ 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; + [](kv::KvCache& self, nb::object stream, std::optional> groups) + { std::optional optStream; if (!stream.is_none()) optStream = reinterpret_cast(nb::cast(stream)); @@ -1973,8 +1943,7 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) nb::gil_scoped_release rel; return self.resume(optStream, std::move(selected)); }, - nb::arg("cuda_stream") = nb::none(), nb::arg("is_decoding") = nb::none(), - nb::arg("optional_reuse_groups") = nb::none()) + nb::arg("cuda_stream") = nb::none(), nb::arg("optional_reuse_groups") = nb::none()) .def( "get_page_priorities", [](kv::KvCache const& self, int layerGroupId, int beamIndex) @@ -2074,22 +2043,6 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) }, nb::arg("layer_group_id"), nb::call_guard(), "Page indices of this life cycle's scratch slots, in slot order.") - .def( - "get_scratch_desc", - [](kv::KvCache const& self, int layerGroupId) -> std::optional - { - auto const lgId = kv::LayerGroupId{layerGroupId}; - auto const range = self.getSwaScratchRange(lgId); - if (!range) - { - return std::nullopt; - } - auto const slots = self.getScratchSlots(lgId); - return ScratchDescCompat{range, std::vector(slots.begin(), slots.end())}; - }, - nb::arg("layer_group_id"), nb::call_guard(), - "get_swa_scratch_range() and get_scratch_slots() as one object, or None when this life " - "cycle serves nothing from scratch. Remove with the ScratchDesc class.") .def_prop_ro("has_scratch_slots", &kv::KvCache::hasScratchSlots) .def_prop_rw("enable_swa_scratch_reuse", &kv::KvCache::isSwaScratchReuseEnabled, [](kv::KvCache& self, bool enable) @@ -2110,23 +2063,6 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) nb::arg("mode"), nb::arg("layer_group_id"), "Per life cycle, which is where the scratch slots live.") .def_prop_ro("status", [](kv::KvCache const& kvc) { return kvc.status(); }) .def_prop_ro("is_active", &kv::KvCache::isActive) - // Admission into decode, which the manager no longer marks: beyond setting a flag, - // this offloaded sparse history, and the sparse fetch path owns that now. A cache - // with no sparse life cycle had nothing to offload, so always admitting is what it - // already did. Remove once the full sparse attention support is merged and its - // callers drop this; it exists only so this change stands alone. - .def( - "enter_decode", - [](kv::KvCache const& self) - { - if (!self.isActive()) - { - throw kv::LogicError("Decode admission requires an active request"); - } - return true; - }, - nb::call_guard(), - "Admit this request into decode. Always true for an active request.") .def_prop_ro("num_blocks", [](kv::KvCache const& self) { return self.numBlocks().value(); }) .def_prop_ro( "batch_row", @@ -2532,20 +2468,20 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) nb::arg("batch")); mIntrospection.def( "sparse_resolve", - [](std::vector const& residentOrdinal, std::vector const& lruOrder, int numLayers, + [](std::vector const& residentOrdinal, std::vector const& entryAge, 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, + auto r = kv::KvCacheIntrospection::sparseResolve(residentOrdinal, entryAge, 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)); + return std::make_tuple(std::move(r.residentOrdinal), std::move(r.entryAge), std::move(r.entryHit), + 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("resident_ordinal"), nb::arg("entry_age"), 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"), @@ -3723,28 +3659,13 @@ 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 scratch, + std::optional indexMode, + std::optional> scratchRange, std::vector const& scratchSlots) -> std::vector { - kv::HalfOpenRange range{}; - std::vector const* slots = &scratchSlots; - if (nb::isinstance(scratch)) - { - TLLM_CHECK_WITH_INFO( - scratchSlots.empty(), "scratch_slots must be empty when scratch_range is a ScratchDesc"); - auto const& desc = nb::cast(scratch); - range = desc.range; - slots = &desc.slotIds; - } - else if (!scratch.is_none()) - { - range = nb::cast>(scratch); - } - return self( - baseIndices, indexMode, range, kv::Span{slots->data(), static_cast(slots->size())}); + return self(baseIndices, indexMode, scratchRange.value_or(kv::HalfOpenRange{}), + kv::Span{scratchSlots.data(), static_cast(scratchSlots.size())}); }, nb::arg("base_indices"), nb::arg("index_mode") = nb::none(), nb::arg("scratch_range") = nb::none(), nb::arg("scratch_slots") = std::vector{}); diff --git a/cpp/tests/unit_tests/kernels/CMakeLists.txt b/cpp/tests/unit_tests/kernels/CMakeLists.txt index f75cb1de4299..5086a317b1e2 100644 --- a/cpp/tests/unit_tests/kernels/CMakeLists.txt +++ b/cpp/tests/unit_tests/kernels/CMakeLists.txt @@ -228,6 +228,7 @@ add_tllm_gtest( LABELS "unit;kernels;attention-mask;gpu") add_gtest(fmhaDsv4EpilogueFusionTest fmhaDsv4EpilogueFusionTest.cpp) +add_gtest(fmhaDeterministicTest fmhaDeterministicTest.cpp NO_GTEST_MAIN) add_tllm_gtest( nvfp4ColdPageKernelsTest diff --git a/cpp/tests/unit_tests/kernels/fmhaDeterministicTest.cpp b/cpp/tests/unit_tests/kernels/fmhaDeterministicTest.cpp new file mode 100644 index 000000000000..3275b612a086 --- /dev/null +++ b/cpp/tests/unit_tests/kernels/fmhaDeterministicTest.cpp @@ -0,0 +1,246 @@ +/* + * 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 + +#include +#include +#include +#include +#include +#include +#include + +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaRunner.h" + +namespace +{ + +using tensorrt_llm::kernels::DATA_TYPE_BF16; +using tensorrt_llm::kernels::DATA_TYPE_E4M3; +using tensorrt_llm::kernels::FmhaKernelType; +using tensorrt_llm::kernels::QkvLayout; +using tensorrt_llm::kernels::SparseType; +using tensorrt_llm::kernels::TileScheduler; +using tensorrt_llm::kernels::TllmGenFmhaRunner; +using tensorrt_llm::kernels::TllmGenFmhaRunnerParams; +using tensorrt_llm::kernels::TrtllmGenAttentionMaskType; + +class FmhaDeterministicTest : public ::testing::Test +{ +protected: + static constexpr int32_t kHeadDim = 512; + static constexpr int32_t kNumHeads = 32; + static constexpr int32_t kSparseTopK = 640; + static constexpr int32_t kCaptureKvLength = 9216; + + void SetUp() override + { + if (!tensorrt_llm::common::isSM100Family(tensorrt_llm::common::getSMVersion())) + { + GTEST_SKIP() << "TRTLLM-Gen sparse MLA requires Blackwell or newer."; + } + TLLM_CUDA_CHECK(cudaStreamCreate(&mStream)); + } + + void TearDown() override + { + if (mStream != nullptr) + { + cudaStreamSynchronize(mStream); + } + if (mGraphExec != nullptr) + { + cudaGraphExecDestroy(mGraphExec); + } + if (mGraph != nullptr) + { + cudaGraphDestroy(mGraph); + } + for (void* ptr : mAllocations) + { + cudaFree(ptr); + } + if (mStream != nullptr) + { + cudaStreamDestroy(mStream); + } + } + + template + T* allocate(size_t count) + { + T* ptr = nullptr; + TLLM_CUDA_CHECK(cudaMalloc(&ptr, count * sizeof(T))); + mAllocations.push_back(ptr); + return ptr; + } + + template + void copyToDevice(T* dst, std::vector const& src) + { + TLLM_CUDA_CHECK(cudaMemcpy(dst, src.data(), src.size() * sizeof(T), cudaMemcpyHostToDevice)); + } + + std::vector readOutput(uint16_t const* output) + { + std::vector result(kNumHeads * kHeadDim); + TLLM_CUDA_CHECK( + cudaMemcpyAsync(result.data(), output, result.size() * sizeof(uint16_t), cudaMemcpyDeviceToHost, mStream)); + TLLM_CUDA_CHECK(cudaStreamSynchronize(mStream)); + return result; + } + + TllmGenFmhaRunnerParams makeParams(int32_t numHeads = kNumHeads) + { + TllmGenFmhaRunnerParams params{}; + params.mQkvLayout = QkvLayout::PagedKv; + params.mMaskType = TrtllmGenAttentionMaskType::Dense; + params.mKernelType = FmhaKernelType::Generation; + params.mTileScheduler = TileScheduler::Persistent; + params.mMultiCtasKvMode = false; + params.mHeadDimQk = kHeadDim; + params.mHeadDimV = kHeadDim; + params.mNumHeadsQ = numHeads; + params.mNumHeadsKv = 1; + params.mNumHeadsQPerKv = numHeads; + params.mBatchSize = 1; + params.mMaxSeqLenCacheKv = kCaptureKvLength; + params.mMaxSeqLenQ = 1; + params.mMaxSeqLenKv = kCaptureKvLength; + // Negative reaches are unbounded, which is what a full-attention decode wants. + params.mLeftSlidingWindow = -1; + params.mRightSlidingWindow = -1; + params.mChunkedAttentionSize = std::numeric_limits::max(); + params.mSumOfSeqLensQ = 1; + params.mSumOfSeqLensKv = kCaptureKvLength; + params.mMaxNumPagesPerSeqKv = kSparseTopK; + params.mNumTokensPerPage = 1; + params.mNumPagesInMemPool = kSparseTopK; + params.mMultiProcessorCount = tensorrt_llm::common::getMultiProcessorCount(); + params.mScaleQ = 1.0F; + params.mSparseAttention = SparseType::DynamicTokenSparse; + params.mSparseTopK = kSparseTopK; + params.stream = mStream; + return params; + } + + cudaStream_t mStream{}; + cudaGraph_t mGraph{}; + cudaGraphExec_t mGraphExec{}; + std::vector mAllocations; +}; + +TEST_F(FmhaDeterministicTest, LongSparseKvKeepsOneReductionPartition) +{ + TllmGenFmhaRunner runner(DATA_TYPE_E4M3, DATA_TYPE_E4M3, DATA_TYPE_E4M3, DATA_TYPE_BF16, 0, 0, 0, 0, false); + for (int32_t const numHeads : {32, 64}) + { + for (int32_t const kvLength : {129, 640, kCaptureKvLength}) + { + auto params = makeParams(numHeads); + params.mMaxSeqLenKv = kvLength; + auto const [supported, info] = runner.isSupportedWithInfo(params); + ASSERT_TRUE(supported) << info; + EXPECT_NE(info.find("multiCtasKvMode=0"), std::string::npos) + << "heads=" << numHeads << ", kvLength=" << kvLength << ": " << info; + } + } +} + +TEST_F(FmhaDeterministicTest, CudaGraphMatchesEagerForChangingSparseLengths) +{ + auto params = makeParams(); + auto* q = allocate(kNumHeads * kHeadDim); + auto* kv = allocate(kSparseTopK * kHeadDim); + auto* output = allocate(kNumHeads * kHeadDim); + auto* sequenceLength = allocate(1); + auto* topkLength = allocate(1); + auto* indices = allocate(kSparseTopK); + auto* scale = allocate(1); + auto* softmaxScale = allocate(1); + std::vector qValues(kNumHeads * kHeadDim); + std::vector kvValues(kSparseTopK * kHeadDim); + for (size_t i = 0; i < qValues.size(); ++i) + { + qValues[i] = static_cast(24 + (i * 7 % 24) + ((i / 3 % 2) << 7)); + } + for (size_t i = 0; i < kvValues.size(); ++i) + { + kvValues[i] = static_cast(24 + (i * 11 % 24) + ((i / 5 % 2) << 7)); + } + std::vector indexValues(kSparseTopK); + for (int32_t i = 0; i < kSparseTopK; ++i) + { + indexValues[i] = i; + } + copyToDevice(q, qValues); + copyToDevice(kv, kvValues); + copyToDevice(indices, indexValues); + copyToDevice(scale, std::vector{1.0F}); + copyToDevice(softmaxScale, std::vector{static_cast(std::log2(std::exp(1.0)) / std::sqrt(kHeadDim))}); + copyToDevice(sequenceLength, std::vector{kCaptureKvLength}); + copyToDevice(topkLength, std::vector{kSparseTopK}); + params.qPtr = q; + params.kvPtr = kv; + params.slidingWindowKvPoolBasePtr = kv; + params.oPtr = output; + params.seqLensKvPtr = sequenceLength; + params.kvPageIdxPtr = indices; + params.ptrSparseMlaTopKLens = topkLength; + params.outputScalePtr = scale; + params.scaleSoftmaxLog2Ptr = softmaxScale; + params.kvSfScalePtr = scale; + params.oSfScalePtr = scale; + TllmGenFmhaRunner runner(DATA_TYPE_E4M3, DATA_TYPE_E4M3, DATA_TYPE_E4M3, DATA_TYPE_BF16, 0, 0, 0, 0, false); + runner.run(params); + TLLM_CUDA_CHECK(cudaStreamSynchronize(mStream)); + TLLM_CUDA_CHECK(cudaStreamBeginCapture(mStream, cudaStreamCaptureModeGlobal)); + runner.run(params); + TLLM_CUDA_CHECK(cudaStreamEndCapture(mStream, &mGraph)); + TLLM_CUDA_CHECK(cudaGraphInstantiate(&mGraphExec, mGraph, nullptr, nullptr, 0)); + + for (int32_t const kvLength : {127, 128, 129, 160, 513, 2049, 8193}) + { + SCOPED_TRACE(kvLength); + int32_t const selected = std::min(kvLength, 128) + std::min(kvLength / 4, 512); + copyToDevice(sequenceLength, std::vector{kvLength}); + copyToDevice(topkLength, std::vector{selected}); + params.mMaxSeqLenKv = kvLength; + params.mSumOfSeqLensKv = kvLength; + runner.run(params); + auto const eager = readOutput(output); + ASSERT_TRUE( + std::all_of(eager.begin(), eager.end(), [](uint16_t value) { return (value & 0x7f80U) != 0x7f80U; })); + TLLM_CUDA_CHECK(cudaGraphLaunch(mGraphExec, mStream)); + auto const graph = readOutput(output); + EXPECT_EQ(eager, graph); + } +} + +} // namespace + +int main(int argc, char** argv) +{ + if (setenv("FORCE_ATTENTION_KERNEL_DETERMINISTIC", "1", 1) != 0) + { + return EXIT_FAILURE; + } + ::testing::InitGoogleTest(&argc, argv); + return RUN_ALL_TESTS(); +} diff --git a/scripts/generate_trtllm_serve_schemas.py b/scripts/generate_trtllm_serve_schemas.py index a4e3522d8dd1..a99ff13fb978 100644 --- a/scripts/generate_trtllm_serve_schemas.py +++ b/scripts/generate_trtllm_serve_schemas.py @@ -104,6 +104,7 @@ def _schema_for(annotation: Any, parent: dict | None = None) -> dict: from pydantic import TypeAdapter schema = TypeAdapter(annotation).json_schema(schema_generator=_schema_generator()) + _apply_sparse_indexer_defaults(schema) if parent is not None: definitions = parent.setdefault("$defs", {}) for name, definition in schema.pop("$defs", {}).items(): @@ -113,6 +114,28 @@ def _schema_for(annotation: Any, parent: dict | None = None) -> dict: return schema +def _apply_sparse_indexer_defaults(schema: dict) -> None: + """Describe the offload-dependent default resolved by the config validator.""" + for definition in [schema, *schema.get("$defs", {}).values()]: + if definition.get("title") not in { + "DeepSeekSparseAttentionConfig", + "DeepSeekV4SparseAttentionConfig", + }: + continue + field = "use_cute_dsl_paged_mqa_logits" + resident_default = definition["properties"][field].pop("default") + definition.setdefault("allOf", []).append( + { + "if": { + "properties": {"enable_kv_cache_offload": {"const": True}}, + "required": ["enable_kv_cache_offload"], + }, + "then": {"properties": {field: {"default": False}}}, + "else": {"properties": {field: {"default": resident_default}}}, + } + ) + + def _nullable(schema: dict) -> dict: return {"anyOf": [schema, {"type": "null"}]} diff --git a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_sparse_mla.py b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_sparse_mla.py index 2cb531e21c3d..c2a36d8ce512 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_sparse_mla.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_sparse_mla.py @@ -253,6 +253,13 @@ def _unsupported_reason( metadata.kv_cache_manager, "swa_pool_ptr" ): return "a DeepSeek-V4 KV cache manager is required." + if ( + int(attn.compress_ratio) == 4 + and getattr(metadata, "sparse_offload_state", None) is not None + ): + # The kernel gathers compressed rows from the manager's compress pool, but + # offloaded rows sit in the sparse cache arena that the aux pool pointer names. + return "sparse KV cache offload keeps the compressed rows in the sparse cache arena." if metadata.beam_width != 1: return "beam search is not supported." if getattr(metadata, "token_positions_cuda", None) is None: diff --git a/tensorrt_llm/_torch/attention/backends/sparse/csa2/cache_manager.py b/tensorrt_llm/_torch/attention/backends/sparse/csa2/cache_manager.py index 7d87e291667f..49f49e56c173 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/csa2/cache_manager.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/csa2/cache_manager.py @@ -51,6 +51,14 @@ ) +def _scratch_entry(cache, pool_id): + """The scratch range and the slots serving it, or None when none of the blocks do.""" + scratch_range = cache.get_swa_scratch_range(pool_id) + if not scratch_range: + return None + return scratch_range, cache.get_scratch_slots(pool_id) + + def _validate_context_swa_layer_limit(limit: int | None, num_layers: int) -> None: if limit is not None and (type(limit) is not int or not 0 < limit < num_layers): raise ValueError("CSA2 context SWA layer limit must be an integer inside the model") @@ -792,7 +800,8 @@ def get_cache_indices(self, request_id: int, layer_idx: int, role: CSA2CacheRole return converter( cache.get_base_page_indices(pool_id).tolist(), role.index_mode, - cache.get_scratch_desc(pool_id), + cache.get_swa_scratch_range(pool_id) or None, + cache.get_scratch_slots(pool_id), ) def compute_batch_page_tables(self, request_ids, num_contexts) -> None: @@ -856,17 +865,17 @@ def _scratch_descriptors(self, request_ids): if not holders: return None descs = [ - (index, row, desc) + (index, row, entry) for row, cache in holders for index, pool in enumerate(self._scratch_pools.tolist()) - if (desc := cache.get_scratch_desc(pool)) + if (entry := _scratch_entry(cache, pool)) ] ranges = np.zeros((len(self._scratch_pools), len(request_ids), 2), dtype=np.int32) - width = max((len(desc.slot_ids) for _, _, desc in descs), default=1) + width = max((len(slot_ids) for _, _, (_, slot_ids) in descs), default=1) slots = np.zeros((len(self._scratch_pools), len(request_ids), width), dtype=np.int32) - for index, row, desc in descs: - ranges[index, row] = (desc.range.beg, desc.range.end) - slots[index, row, : len(desc.slot_ids)] = desc.slot_ids + for index, row, (scratch_range, slot_ids) in descs: + ranges[index, row] = (scratch_range.beg, scratch_range.end) + slots[index, row, : len(slot_ids)] = slot_ids return ranges, slots def batch_page_table(self, spec, width: int) -> torch.Tensor: diff --git a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/backend.py b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/backend.py index 1337b7c747f9..7cf4c96d1f68 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/backend.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/backend.py @@ -221,7 +221,12 @@ def _prepare_sparse_forward_args( sparse_args.sparse_attn_kv_lens = metadata.sparse_mla_topk_lens[self.compress_ratio][ start_idx:end_idx ] - if self.compress_ratio > 1: + if self.compress_ratio == 4 and metadata.sparse_offload_state is not None: + descriptor = metadata.sparse_offload_state.layers[self.layer_idx] + sparse_args.aux_kv_cache_pool_ptr = descriptor.arena_base_addresses[ + descriptor.data_buffer_index + ] + elif self.compress_ratio > 1: sparse_args.aux_kv_cache_pool_ptr = metadata.sparse_mla_base_ptrs[self.compress_ratio] else: sparse_args.aux_kv_cache_pool_ptr = None @@ -305,7 +310,8 @@ def sparse_attn_predict( local_layer_idx, DeepseekV4AttentionType.SWA.value ] - active_request_count = None + sparse_state = None + sparse_index_kwargs = {} if self.compress_ratio > 1: compressed_buffer_ptr = metadata.compressed_buffer_ptrs[layer_idx] compress_pool_base_ptr = metadata.sparse_mla_base_ptrs[self.compress_ratio] @@ -318,28 +324,44 @@ def sparse_attn_predict( topk_indices = sparse_backend_args.topk_indices assert topk_indices is not None, "topk_indices is required when compress_ratio=4" compressed_local_indices = topk_indices - # Metadata setup and prepare() validate the dtype and batch layout. - # Fresh prefill already uses the resident write table. - if ( - self.sparse_attention_config.enable_kv_cache_offload - and attention_input_type == AttentionInputType.generation_only - ): - state = metadata.sparse_offload_state - assert state is not None, "Sparse offload metadata must be initialized" - # module.forward_sparse_attn has joined both producer - # streams. Keep fetch before the conversion kernel's - # FMHA scheduler prologue on this consuming stream. - block_table_compressed = kv_cache_manager.fetch_sparse_read_table( - state, - layer_idx, - topk_indices, - req_id, - metadata.compressed_kv_lens_cuda[self.compress_ratio], - block_table_compressed, + if self.sparse_attention_config.enable_kv_cache_offload: + sparse_state = metadata.sparse_offload_state + if ( + sparse_state is None + or attention_input_type != AttentionInputType.generation_only + ): + raise NotImplementedError( + "Sparse arena attention requires prepared decode metadata" + ) + descriptor = sparse_state.layers[layer_idx] + if sparse_state.write_layer != layer_idx: + raise RuntimeError( + "Prepare this layer's compressor writes before sparse attention" + ) + compressed_lengths = metadata.compressed_kv_lens_cuda[self.compress_ratio][ + : metadata.num_seqs + ] + fetched = kv_cache_manager.fetch_sparse_indices( + sparse_state, layer_idx, topk_indices, req_id, compressed_lengths + ) + data_index = descriptor.data_buffer_index + # Read at the slot: buffers whose arena indices coincide share one + # slice, and the slices no buffer maps to are never written. + data_slot = sparse_state.out_slots[layer_idx][data_index] + sparse_index_kwargs = dict( + compressed_fetched_indices=fetched[:, data_slot, :], + compressed_write_indices=sparse_state.write_indices[ + len(descriptor.buffer_ids) + ][:, data_slot], + compressed_history_lengths=sparse_state.history_lengths, + compressed_lengths=compressed_lengths, + compressed_batch_rows=sparse_state.batch_rows, + active_request_count=sparse_state.active_request_count, + ) + compress_pool_base_ptr = descriptor.arena_base_addresses[data_index] + forward_args.sparse_runtime_params.aux_kv_cache_pool_ptr = ( + compress_pool_base_ptr ) - # Metadata counts include CUDA graph dummy rows. The active - # count is refreshed in place before replay to mask them. - active_request_count = state.active_request_count else: compressed_local_indices = metadata.compressed_local_indices_cuda[start_idx:end_idx] else: @@ -404,8 +426,8 @@ def sparse_attn_predict( num_compressed_indices=metadata.max_compressed_indices[self.compress_ratio], **sched_kwargs, split_extra=self.use_fp8_ds_mla, - active_request_count=active_request_count, - mask_invalid_pages=active_request_count is not None, + mask_invalid_pages=sparse_state is not None, + **sparse_index_kwargs, ) if self.use_fp8_ds_mla: @@ -415,69 +437,90 @@ def sparse_attn_predict( if self._uses_nvfp4_compress and self.compress_ratio > 1: num_compressed_indices = metadata.max_compressed_indices[self.compress_ratio] compressed_indices = global_indices[:, -num_compressed_indices:].contiguous() - data_pool, scale_pool = kv_cache_manager.get_compress_pool_buffers(self.compress_ratio) - num_pool_tokens = data_pool.numel() // ( - (self.head_dim + kv_cache_manager.nvfp4_residual_dim) // 2 - ) - # The previous layer has already enqueued its attention on this stream. - # Release its Python reference before allocating the next scratch so - # the caching allocator can reuse the storage in stream order. - metadata.nvfp4_compress_fp8_scratch = None - if attention_input_type == AttentionInputType.generation_only: - num_scratch_rows = compressed_indices.numel() - scratch = _alloc_nvfp4_compress_fp8_scratch( - num_scratch_rows, self.head_dim, global_indices.device - ) + if sparse_state is not None: + if descriptor.scale_buffer_index is None or sparse_state.nvfp4_scratch is None: + raise RuntimeError( + "Sparse NVFP4 attention requires block scales and gather scratch" + ) + scratch = sparse_state.nvfp4_scratch[: compressed_indices.shape[0]] + data_pool, scale_pool = kv_cache_manager.get_sparse_compress_buffers(layer_idx) + # Paired arena geometry gives data and scales the same row indices. + # The existing gather replaces those indices with FP8 scratch indices. torch.ops.trtllm.nvfp4_mla_kv_cache_gather_direct( data_pool, scale_pool, compressed_indices, - scratch[:num_scratch_rows].view(*compressed_indices.shape, self.head_dim), + scratch, self._nvfp4_compress_scale_quant_orig, kv_cache_manager.nvfp4_residual_dim, - num_pool_tokens, + data_pool.shape[0], ) else: - request_count = ( - metadata.num_contexts - if attention_input_type == AttentionInputType.context_only - else metadata.num_seqs - ) - compressed_kv_lengths = metadata.compressed_kv_lens_cuda[self.compress_ratio][ - :request_count - ].contiguous() - total_raw_kv_tokens = int(metadata.host_total_kv_lens[0].item()) - if attention_input_type == AttentionInputType.mixed: - total_raw_kv_tokens += int(metadata.host_total_kv_lens[1].item()) - # Sum(floor(kv_len / ratio)) <= floor(sum(kv_len) / ratio). - # The latter is a host-known upper bound for the logical-token - # compaction workspace and avoids a device synchronization. - max_compressed_kv_tokens = total_raw_kv_tokens // self.compress_ratio - # Dynamic sparse FMHA may prefetch a full fixed-width Top-K - # tile even when only a short prefix is valid. Match the DSA - # context gather contract by keeping at least one index row's - # width addressable in the scratch pool. - scratch_capacity = max( - compressed_indices.shape[1], - min(max_compressed_kv_tokens, compressed_indices.numel()), + data_pool, scale_pool = kv_cache_manager.get_compress_pool_buffers( + self.compress_ratio ) - scratch = _alloc_nvfp4_compress_fp8_scratch( - scratch_capacity, self.head_dim, global_indices.device + num_pool_tokens = data_pool.numel() // ( + (self.head_dim + kv_cache_manager.nvfp4_residual_dim) // 2 ) - if max_compressed_kv_tokens > 0: - torch.ops.trtllm.nvfp4_mla_context_kv_cache_gather_direct( + # The previous layer has already enqueued its attention on this stream. + # Release its Python reference before allocating the next scratch so + # the caching allocator can reuse the storage in stream order. + metadata.nvfp4_compress_fp8_scratch = None + if attention_input_type == AttentionInputType.generation_only: + num_scratch_rows = compressed_indices.numel() + scratch = _alloc_nvfp4_compress_fp8_scratch( + num_scratch_rows, self.head_dim, global_indices.device + ) + torch.ops.trtllm.nvfp4_mla_kv_cache_gather_direct( data_pool, scale_pool, - compressed_local_indices.contiguous(), - req_id.contiguous(), - compressed_kv_lengths, compressed_indices, - scratch[:scratch_capacity].unsqueeze(1), + scratch[:num_scratch_rows].view(*compressed_indices.shape, self.head_dim), self._nvfp4_compress_scale_quant_orig, kv_cache_manager.nvfp4_residual_dim, - max_compressed_kv_tokens, num_pool_tokens, ) + else: + request_count = ( + metadata.num_contexts + if attention_input_type == AttentionInputType.context_only + else metadata.num_seqs + ) + compressed_kv_lengths = metadata.compressed_kv_lens_cuda[self.compress_ratio][ + :request_count + ].contiguous() + total_raw_kv_tokens = int(metadata.host_total_kv_lens[0].item()) + if attention_input_type == AttentionInputType.mixed: + total_raw_kv_tokens += int(metadata.host_total_kv_lens[1].item()) + # Sum(floor(kv_len / ratio)) <= floor(sum(kv_len) / ratio). + # The latter is a host-known upper bound for the logical-token + # compaction workspace and avoids a device synchronization. + max_compressed_kv_tokens = total_raw_kv_tokens // self.compress_ratio + # Dynamic sparse FMHA may prefetch a full fixed-width Top-K + # tile even when only a short prefix is valid. Match the DSA + # context gather contract by keeping at least one index row's + # width addressable in the scratch pool. + scratch_capacity = max( + compressed_indices.shape[1], + min(max_compressed_kv_tokens, compressed_indices.numel()), + ) + scratch = _alloc_nvfp4_compress_fp8_scratch( + scratch_capacity, self.head_dim, global_indices.device + ) + if max_compressed_kv_tokens > 0: + torch.ops.trtllm.nvfp4_mla_context_kv_cache_gather_direct( + data_pool, + scale_pool, + compressed_local_indices.contiguous(), + req_id.contiguous(), + compressed_kv_lengths, + compressed_indices, + scratch[:scratch_capacity].unsqueeze(1), + self._nvfp4_compress_scale_quant_orig, + kv_cache_manager.nvfp4_residual_dim, + max_compressed_kv_tokens, + num_pool_tokens, + ) metadata.nvfp4_compress_fp8_scratch = scratch global_indices[:, -num_compressed_indices:] = compressed_indices forward_args.sparse_runtime_params.aux_kv_cache_pool_ptr = scratch.data_ptr() diff --git a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/cache_manager.py b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/cache_manager.py index 1d97a1f2f1cd..14a3644dd661 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/cache_manager.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/cache_manager.py @@ -15,8 +15,8 @@ from collections import defaultdict from dataclasses import replace -from math import gcd -from typing import TYPE_CHECKING, Dict, Iterable, List, Optional, Tuple +from math import ceil, gcd +from typing import Dict, Iterable, List, Optional, Tuple import torch @@ -24,13 +24,13 @@ from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import ( _GUARD_PAGE_REQUEST_ID, _RESERVED_REQUEST_IDS, - CUDA_GRAPH_DUMMY_REQUEST_ID, GPU_LEVEL, KVCacheManagerV2, _fill_kv_pages, + _get_generation_kv_capacity, ) +from tensorrt_llm._torch.pyexecutor.kv_cache.sparse_offload_memory import sparse_cache_config from tensorrt_llm._torch.pyexecutor.resource_manager import scan_kv_buffer_for_invalid_values -from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests from tensorrt_llm._utils import ( TensorWrapper, convert_to_torch_tensor, @@ -44,30 +44,24 @@ from tensorrt_llm.llmapi.llm_args import DeepSeekV4SparseAttentionConfig, KvCacheConfig from tensorrt_llm.logger import logger from tensorrt_llm.mapping import Mapping +from tensorrt_llm.math_utils import ceil_div from tensorrt_llm.runtime import ModelConfig from tensorrt_llm.runtime.kv_cache_manager_v2 import ( BAD_PAGE_INDEX, AttentionLayerConfig, BatchDesc, BufferConfig, - BufferId, - CudaStream, DataRole, + HalfOpenRange, KVCacheDesc, LayerId, PageIndexMode, - ScratchDesc, + SparseCacheConfig, ) from tensorrt_llm.runtime.kv_cache_manager_v2 import KVCacheManagerConfig as KVCacheManagerConfigPy +from ..offload import SparseOffloadCacheMixin, SparseOffloadLayerDescriptor, SparseOffloadState from .compressor import KVCacheDtype, get_nvfp4_compress_residual_dim -from .kernels import ( - check_sparse_read_table, - merge_sparse_read_table, - prepare_sparse_write_table, - select_sparse_history_pages, -) -from .offload import SparseOffloadLayerDescriptor, SparseOffloadState from .params import ( DEEPSEEK_V4_NON_SLIDING_ATTENTION, DEEPSEEK_V4_SLIDING_ATTENTION, @@ -77,9 +71,6 @@ is_overlap_compressor, ) -if TYPE_CHECKING: - from tensorrt_llm._torch.attention.backends.interface import AttentionMetadata - COMPRESS_BLOCK_SCALE_ROLE = DataRole("deepseek_v4_compress_block_scale") KV_CACHE_COPY_ALIGNMENT = 16 NVFP4_VECTOR_SIZE = 16 @@ -177,10 +168,24 @@ def _estimate_non_sliding_attn_size_per_token( use_fp8_ds_mla: bool = False, has_nvfp4_compress: bool = False, nvfp4_residual_dim: int | None = None, -) -> int: + offload_sparse_history: bool = False, +) -> float: + """Return the GPU bytes per token of compressed-layer history. + + With ``offload_sparse_history``, size a worker that offloads ratio-4 + history: leave out the offloaded rows and count the rest at their exact + bytes per token instead of rounding each down to whole bytes. + """ total_bytes = 0 for compress_ratio in compress_ratios: for attn_type in DEEPSEEK_V4_NON_SLIDING_ATTENTION: + if ( + offload_sparse_history + and compress_ratio == DEEPSEEK_V4_SPARSE_RATIO + and attn_type == DeepseekV4AttentionType.COMPRESS + ): + # Offloaded history lives in host memory. + continue if compress_ratio_has_attention(compress_ratio, attn_type): total_bytes += _get_attn_bytes_per_token( head_dim, @@ -192,6 +197,7 @@ def _estimate_non_sliding_attn_size_per_token( use_fp8_ds_mla=use_fp8_ds_mla, has_nvfp4_compress=has_nvfp4_compress, nvfp4_residual_dim=nvfp4_residual_dim, + exact=offload_sparse_history, ) return total_bytes @@ -208,7 +214,15 @@ def _estimate_swa_cache_size( scratch: bool, indexer_k_dtype: str = "fp8", use_fp8_ds_mla: bool = False, + generation_capacity_headroom: int | None = None, ) -> Tuple[int, int]: + """Return the sliding-window bytes per token and per request. + + By default a window costs one page per ``tokens_per_block`` tokens of + window. With ``generation_capacity_headroom``, the tokens a generation step + reserves above history, it costs the most pages its live interval touches + at any page offset, which is what the cache manager holds. + """ tokens_per_block = int(tokens_per_block) size_per_token = 0 size_per_request = 0 @@ -226,9 +240,15 @@ def _estimate_swa_cache_size( window_size = state_factor * compress_ratio if window_size <= 0: continue - window_tokens = ( - (int(window_size) + tokens_per_block - 1) // tokens_per_block - ) * tokens_per_block + if generation_capacity_headroom is None: + window_pages = ceil_div(int(window_size), tokens_per_block) + else: + # The live interval spans window + headroom - 1 tokens. + window_pages = ( + ceil_div(int(window_size) + generation_capacity_headroom - 2, tokens_per_block) + + 1 + ) + window_tokens = window_pages * tokens_per_block token_bytes = _get_attn_bytes_per_token( head_dim, index_head_dim, @@ -262,7 +282,13 @@ def _get_attn_bytes_per_token( use_fp8_ds_mla: bool = False, has_nvfp4_compress: bool = False, nvfp4_residual_dim: int | None = None, -) -> int: + exact: bool = False, +) -> float: + """Return the bytes one token adds to a buffer. + + A compressed row covers ``compress_ratio`` tokens, so its bytes are divided + among them: rounded down to whole bytes, or exactly when ``exact`` is set. + """ token_bytes = get_token_bytes( head_dim, index_head_dim, @@ -275,10 +301,20 @@ def _get_attn_bytes_per_token( nvfp4_residual_dim=nvfp4_residual_dim, ) if attn_type in [DeepseekV4AttentionType.COMPRESS, DeepseekV4AttentionType.INDEXER_COMPRESS]: + if exact: + return token_bytes / compress_ratio token_bytes //= compress_ratio return token_bytes +def _scratch_entry(kv_cache, pool_id) -> tuple[HalfOpenRange, list[int]] | None: + """The scratch range and the slots serving it, or None when none of the blocks do.""" + scratch_range = kv_cache.get_swa_scratch_range(pool_id) + if not scratch_range: + return None + return scratch_range, kv_cache.get_scratch_slots(pool_id) + + def _get_index_mode(attn_type: DeepseekV4AttentionType) -> PageIndexMode: if attn_type in DEEPSEEK_V4_SLIDING_ATTENTION: return PageIndexMode.PER_LAYER @@ -286,7 +322,7 @@ def _get_index_mode(attn_type: DeepseekV4AttentionType) -> PageIndexMode: return PageIndexMode.SHARED -class DeepseekV4CacheManager(KVCacheManagerV2): +class DeepseekV4CacheManager(SparseOffloadCacheMixin, KVCacheManagerV2): _supports_reuse_match_backoff = False # This tensor is for compatibility with AttentionOp, it only contains swa attention. @@ -339,13 +375,17 @@ def __init__( "for DeepseekV4CacheManager" ) - self._enable_kv_cache_offload = sparse_attn_config.enable_kv_cache_offload - if self._enable_kv_cache_offload: - # The attention adapter is connected; enable only after validating - # the real runtime's transfer, addressing, and scratch contracts. - raise NotImplementedError( - "DeepSeek-V4 KV cache offload requires the KVCM v2 sparse runtime " - "and validated fetch addressing, scratch lifetime, and transfer ordering." + self._init_sparse_offload( + enabled=sparse_attn_config.enable_kv_cache_offload, + topk=sparse_attn_config.index_topk, + ) + sparse_kv_offload_config = sparse_attn_config.get_kv_cache_offload_config() + if sparse_kv_offload_config is not None: + self._check_sparse_kv_offload_config( + sparse_kv_offload_config, + kv_cache_config, + mapping, + spec_config=kwargs.get("spec_config"), ) # DeepSeek-V4 specific attributes initialization @@ -422,6 +462,7 @@ def __init__( mapping=mapping, dtype=cache_dtype, max_num_tokens=max_num_tokens, + sparse_kv_offload_config=sparse_kv_offload_config, **kwargs, ) self.is_vswa = True # DeepSeek-V4 must has VSWA @@ -430,6 +471,11 @@ def __init__( # to be in the same pool and have the same scale. self._assert_layer_pool_scale() self._assert_nvfp4_compress_pool_layout() + self._sparse_arena_buffers: dict[int, tuple[torch.Tensor, torch.Tensor | None]] = {} + if self._enable_kv_cache_offload: + descriptors = self.get_sparse_offload_descriptors() + self._initialize_sparse_arena_buffers(descriptors) + self._initialize_sparse_arena_padding(descriptors) # For DeepSeek-V4 Attention, the base pointer for SWA pool # Use first PP layer instead of hardcoded 0 for pipeline parallelism. @@ -487,6 +533,10 @@ def _assert_nvfp4_compress_pool_layout(self) -> None: return for layer_idx in self.pp_layers: compress_ratio = self._compress_ratios[layer_idx] + if self._enable_kv_cache_offload and compress_ratio == DEEPSEEK_V4_SPARSE_RATIO: + # Arena geometry is validated when building the sparse descriptors; + # the host page converters do not address the arena. + continue if not compress_ratio_has_attention(compress_ratio, DeepseekV4AttentionType.COMPRESS): continue layer_id = self._layer_attn_to_layer_id[layer_idx, DeepseekV4AttentionType.COMPRESS] @@ -544,6 +594,10 @@ def get_buffers(self, layer_idx: int, attn_type: DeepseekV4AttentionType) -> tor Returns: The buffer tensor (shape: [num_blocks, tokens_per_block, attn_dim]) For blockwise FP8 layers, shape is [num_blocks, tokens_per_block, attn_dim + scale_size] + Sparse COMPRESS storage is device-mapped host memory and keeps the ordinary + page shape, [num_blocks, compressed_block_size, dim_per_token]: a sparse + buffer's size is its whole page, which the cache divides into selection units. + GPU arena views are returned by get_sparse_compress_buffers(). """ layer_id = self._layer_attn_to_layer_id[(layer_idx, attn_type)] data_role = attn_type.role @@ -604,7 +658,11 @@ def get_buffers(self, layer_idx: int, attn_type: DeepseekV4AttentionType) -> tor return convert_to_torch_tensor(TensorWrapper(addr, dtype, shape)) def get_compress_scale_buffers(self, layer_idx: int) -> torch.Tensor: - """Return E4M3 per-16 scales for an NVFP4 COMPRESS layer.""" + """Return E4M3 per-16 scales from the NVFP4 COMPRESS storage pool. + + Sparse storage is device-mapped host memory, indexed one compressed row + at a time. GPU arena scales are returned by get_sparse_compress_buffers(). + """ if not self._use_nvfp4_compress: raise RuntimeError("COMPRESS block scales exist only for NVFP4 KV cache.") layer_id = self._layer_attn_to_layer_id[layer_idx, DeepseekV4AttentionType.COMPRESS] @@ -614,9 +672,10 @@ def get_compress_scale_buffers(self, layer_idx: int) -> torch.Tensor: page_index_upper_bound = self.impl.get_page_index_upper_bound( layer_id, COMPRESS_BLOCK_SCALE_ROLE ) + block_size = self.compressed_block_sizes[layer_idx] shape = ( page_index_upper_bound, - self.compressed_block_sizes[layer_idx], + block_size, (self.head_dim + self.nvfp4_residual_dim) // NVFP4_VECTOR_SIZE, ) return convert_to_torch_tensor(TensorWrapper(addr, DataType.FP8, shape)) @@ -942,9 +1001,16 @@ def get_cache_indices( return converter( base_indices, page_index_mode, - kv_cache.get_scratch_desc(pool_id), + kv_cache.get_swa_scratch_range(pool_id) or None, + kv_cache.get_scratch_slots(pool_id), ) + def _get_window_page_headroom(self) -> int | None: + """Headroom for counting the pages a window holds with offload, else None.""" + if not self._enable_kv_cache_offload: + return None + return self._generation_kv_capacity_headroom + def _get_extra_quota_padding(self) -> int: """Ensure each attention type has minimal space when max_tokens is small.""" return len(DeepseekV4AttentionType) * (2 << 20) @@ -961,6 +1027,7 @@ def _get_quota_from_max_tokens(self, max_tokens: int) -> int: use_fp8_ds_mla=self.use_fp8_ds_mla, has_nvfp4_compress=self._use_nvfp4_compress, nvfp4_residual_dim=self.nvfp4_residual_dim, + offload_sparse_history=self._enable_kv_cache_offload, ) ( context_swa_size_per_token, @@ -976,6 +1043,7 @@ def _get_quota_from_max_tokens(self, max_tokens: int) -> int: scratch=self.enable_swa_scratch_reuse, indexer_k_dtype=self._indexer_k_dtype, use_fp8_ds_mla=self.use_fp8_ds_mla, + generation_capacity_headroom=self._get_window_page_headroom(), ) ( generation_swa_size_per_token, @@ -991,6 +1059,7 @@ def _get_quota_from_max_tokens(self, max_tokens: int) -> int: scratch=False, indexer_k_dtype=self._indexer_k_dtype, use_fp8_ds_mla=self.use_fp8_ds_mla, + generation_capacity_headroom=self._get_window_page_headroom(), ) max_context_tokens = ( self._max_num_tokens if self._max_num_tokens is not None else max_tokens @@ -1018,6 +1087,7 @@ def _get_max_tokens_from_quota(self, quota: int) -> float: use_fp8_ds_mla=self.use_fp8_ds_mla, has_nvfp4_compress=self._use_nvfp4_compress, nvfp4_residual_dim=self.nvfp4_residual_dim, + offload_sparse_history=self._enable_kv_cache_offload, ) context_swa_size_per_token, _ = _estimate_swa_cache_size( self.head_dim, @@ -1030,6 +1100,7 @@ def _get_max_tokens_from_quota(self, quota: int) -> float: scratch=self.enable_swa_scratch_reuse, indexer_k_dtype=self._indexer_k_dtype, use_fp8_ds_mla=self.use_fp8_ds_mla, + generation_capacity_headroom=self._get_window_page_headroom(), ) ( generation_swa_size_per_token, @@ -1045,6 +1116,7 @@ def _get_max_tokens_from_quota(self, quota: int) -> float: scratch=False, indexer_k_dtype=self._indexer_k_dtype, use_fp8_ds_mla=self.use_fp8_ds_mla, + generation_capacity_headroom=self._get_window_page_headroom(), ) padding = self._get_extra_quota_padding() size_per_batch = self.max_batch_size * generation_swa_size_per_request + padding @@ -1091,13 +1163,50 @@ def _build_cache_config(self, config: KVCacheManagerConfigPy) -> KVCacheManagerC if padding_pages: nvfp4_padding_pages[grouped_layers[-1]] = padding_pages + def _make_buffer_config( + layer_idx: int, + role: DataRole, + page_bytes: int, + *, + is_sparse: bool, + ) -> BufferConfig: + if not is_sparse: + return BufferConfig(role=role, size=page_bytes) + if self.use_fp8_ds_mla: + raise NotImplementedError( + "Sparse fp8_ds_mla compressed KV requires a selection layout with " + "separate data and footer-scale regions" + ) + # `size` is the whole page, as for any other buffer; the layer's sparse_cache + # says how many selection units it divides into. num_kv_heads is 1 here, so a + # unit is one contiguous range and the default sparse_unit_chunks of 1 is right. + return BufferConfig(role=role, size=page_bytes, is_sparse=True) + + def _sparse_cache_config() -> SparseCacheConfig: + topk = self._sparse_offload_topk + if topk is None or topk <= 0: + raise ValueError("Sparse compressed KV requires a positive index_topk") + # A selection names compressed rows of four original tokens each. + return sparse_cache_config( + num_selected_units=topk, + unit_tokens=DEEPSEEK_V4_SPARSE_RATIO, + tokens_per_block=self.tokens_per_block, + gpu_cache_factor=self._sparse_kv_offload.gpu_cache_factor, + ) + def _add_layer( layer_idx: int, attention_types: List[DeepseekV4AttentionType], sliding_window_size: int | None, *, - is_sparse: bool = False, + sparse_roles: frozenset[DeepseekV4AttentionType] = frozenset(), ) -> None: + # Sparsity is per buffer, so one layer may hold the compressed KV the cache + # fetches through beside the indexer that selects over it -- they are one + # attention computation and belong together. + def _is_sparse(attn_type: DeepseekV4AttentionType) -> bool: + return attn_type in sparse_roles + layer_id = LayerId(len(layers)) for attn_type in attention_types: layer_attn_to_layer_id[layer_idx, attn_type] = layer_id @@ -1108,12 +1217,11 @@ def _add_layer( buffers = [] for attn_type in attention_types: size = self._get_attn_bytes_per_block(attn_type, layer_idx) - if is_sparse: - buffer = BufferConfig(role=attn_type.role, size=size, is_sparse=True) - else: - # The current runtime does not yet accept the sparse keyword. - buffer = BufferConfig(role=attn_type.role, size=size) - buffers.append(buffer) + buffers.append( + _make_buffer_config( + layer_idx, attn_type.role, size, is_sparse=_is_sparse(attn_type) + ) + ) if self._use_nvfp4_compress and DeepseekV4AttentionType.COMPRESS in attention_types: scale_page_size = ( self.compressed_block_sizes[layer_idx] @@ -1121,9 +1229,11 @@ def _add_layer( // NVFP4_VECTOR_SIZE ) buffers.append( - BufferConfig( - role=COMPRESS_BLOCK_SCALE_ROLE, - size=scale_page_size, + _make_buffer_config( + layer_idx, + COMPRESS_BLOCK_SCALE_ROLE, + scale_page_size, + is_sparse=_is_sparse(DeepseekV4AttentionType.COMPRESS), ) ) # Reuse snapshots copy complete coalesced slots in 16-byte @@ -1139,13 +1249,17 @@ def _add_layer( ) buffers.extend( [ - BufferConfig( - role=data_padding_role, - size=data_page_size, + _make_buffer_config( + layer_idx, + data_padding_role, + data_page_size, + is_sparse=_is_sparse(DeepseekV4AttentionType.COMPRESS), ), - BufferConfig( - role=scale_padding_role, - size=scale_page_size, + _make_buffer_config( + layer_idx, + scale_padding_role, + scale_page_size, + is_sparse=_is_sparse(DeepseekV4AttentionType.COMPRESS), ), ] ) @@ -1155,6 +1269,7 @@ def _add_layer( buffers=buffers, sliding_window_size=sliding_window_size, num_sink_tokens=None, + sparse_cache=_sparse_cache_config() if sparse_roles else None, ) layers.append(layer_config) @@ -1173,18 +1288,21 @@ def _add_layer( [DeepseekV4AttentionType.SWA], self._get_window_size(compress_ratio, DeepseekV4AttentionType.SWA), ) - if self._enable_kv_cache_offload: - _add_layer(layer, [DeepseekV4AttentionType.COMPRESS], None, is_sparse=True) - _add_layer(layer, [DeepseekV4AttentionType.INDEXER_COMPRESS], None) - else: - _add_layer( - layer, - [ - DeepseekV4AttentionType.COMPRESS, - DeepseekV4AttentionType.INDEXER_COMPRESS, - ], - None, - ) + # One layer either way. Only the compressed KV is fetched through the + # sparse cache; the indexer that selects over it stays device-resident. + _add_layer( + layer, + [ + DeepseekV4AttentionType.COMPRESS, + DeepseekV4AttentionType.INDEXER_COMPRESS, + ], + None, + sparse_roles=( + frozenset({DeepseekV4AttentionType.COMPRESS}) + if self._enable_kv_cache_offload + else frozenset() + ), + ) _add_layer( layer, [ @@ -1389,15 +1507,18 @@ def get_cache_bytes_per_token(self) -> int: """Get the average cache bytes per token for DeepSeek-V4.""" has_fp8_kv_cache = self.dtype == DataType.FP8 compress_ratios = [self._compress_ratios[layer] for layer in self.pp_layers] - return _estimate_non_sliding_attn_size_per_token( - self.head_dim, - self.index_head_dim, - compress_ratios, - has_fp8_kv_cache, - indexer_k_dtype=self._indexer_k_dtype, - use_fp8_ds_mla=self.use_fp8_ds_mla, - has_nvfp4_compress=self._use_nvfp4_compress, - nvfp4_residual_dim=self.nvfp4_residual_dim, + return ceil( + _estimate_non_sliding_attn_size_per_token( + self.head_dim, + self.index_head_dim, + compress_ratios, + has_fp8_kv_cache, + indexer_k_dtype=self._indexer_k_dtype, + use_fp8_ds_mla=self.use_fp8_ds_mla, + has_nvfp4_compress=self._use_nvfp4_compress, + nvfp4_residual_dim=self.nvfp4_residual_dim, + offload_sparse_history=self._enable_kv_cache_offload, + ) ) def get_max_resource_count(self) -> int: @@ -1444,6 +1565,7 @@ def _get_cache_bytes_for_tokens(self, total_tokens: int, *, context: bool) -> in use_fp8_ds_mla=self.use_fp8_ds_mla, has_nvfp4_compress=self._use_nvfp4_compress, nvfp4_residual_dim=self.nvfp4_residual_dim, + offload_sparse_history=self._enable_kv_cache_offload, ) swa_size_per_token, swa_size_per_request = _estimate_swa_cache_size( self.head_dim, @@ -1456,12 +1578,25 @@ def _get_cache_bytes_for_tokens(self, total_tokens: int, *, context: bool) -> in scratch=self.enable_swa_scratch_reuse, indexer_k_dtype=self._indexer_k_dtype, use_fp8_ds_mla=self.use_fp8_ds_mla, + generation_capacity_headroom=self._get_window_page_headroom(), ) return int( total_tokens * (non_sliding_attn_size_per_token + swa_size_per_token) + swa_size_per_request ) + def get_request_device_bytes(self, num_tokens: int) -> int: + """Return the device pool bytes one decoding request of ``num_tokens`` tokens holds.""" + return self._get_cache_bytes_for_tokens(num_tokens, context=False) + + def _sparse_offload_nvfp4_head_dim(self) -> int | None: + return self.head_dim if self._use_nvfp4_compress else None + + def _has_sparse_offload_layers(self) -> bool: + return self._enable_kv_cache_offload and any( + self._compress_ratios[layer] == DEEPSEEK_V4_SPARSE_RATIO for layer in self.pp_layers + ) + def get_needed_resource_to_completion(self, request: llm_request.LlmRequest) -> int: if self._is_generation_request(request): return self._get_generation_bytes(request) @@ -1508,7 +1643,7 @@ def _copy_idx_to_device(self, copy_idx: torch.Tensor) -> torch.Tensor: def _copy_scratch_metadata_to_device( self, - scratch_descs_by_pool: list[list[ScratchDesc | None]], + scratch_descs_by_pool: list[list[tuple[HalfOpenRange, list[int]] | None]], num_contexts: int, host_begs_staging: torch.Tensor, host_ends_staging: torch.Tensor, @@ -1525,17 +1660,17 @@ def _copy_scratch_metadata_to_device( host_slots.zero_() for pool_idx, scratch_descs in enumerate(scratch_descs_by_pool): - for context_idx, desc in enumerate(scratch_descs): - if desc is None: + for context_idx, entry in enumerate(scratch_descs): + if entry is None: continue - slot_ids = desc.slot_ids + scratch_range, slot_ids = entry if len(slot_ids) > self._max_scratch_slots: raise RuntimeError( f"Scratch slot count {len(slot_ids)} exceeds staging capacity " f"{self._max_scratch_slots}" ) - host_begs[pool_idx, context_idx] = int(desc.range.beg) - host_ends[pool_idx, context_idx] = int(desc.range.end) + host_begs[pool_idx, context_idx] = int(scratch_range.beg) + host_ends[pool_idx, context_idx] = int(scratch_range.end) for slot_idx, slot_id in enumerate(slot_ids): host_slots[pool_idx, context_idx, slot_idx] = int(slot_id) @@ -1564,7 +1699,7 @@ def compute_sliding_block_tables( if self.enable_swa_scratch_reuse and num_contexts > 0: scratch_descs_by_pool = [ [ - self.kv_cache_map[req].get_scratch_desc(pool_id) + _scratch_entry(self.kv_cache_map[req], pool_id) for req in request_ids[:num_contexts] ] for pool_id in range(self.num_pools) @@ -1653,82 +1788,10 @@ def copy_batch_sliding_block_tables( non_blocking=True, ) - def prepare_resources(self, scheduled_batch: ScheduledRequests) -> None: - if self._enable_kv_cache_offload and any( - not (request.is_first_context_chunk and request.is_last_context_chunk) - for request in scheduled_batch.context_requests - ): - # Attention metadata's prompt_lens contains the current chunk's - # length, so only the scheduled request can identify a first chunk - # that does not complete its prompt. Reject before model writes. - raise NotImplementedError("Sparse offload does not support chunked prefill") - super().prepare_resources(scheduled_batch) - - def _validate_sparse_history_policy(self) -> None: - """Check fixed history policies during attention metadata initialization.""" - if ( - self.enable_block_reuse - or self.kv_compression_manages_history - or self._has_cp_helix - or self.is_draft - ): - raise NotImplementedError( - "Sparse offload requires block reuse, KV compression, CP Helix, and draft caches disabled" - ) - - def _order_sparse_history_update(self, requests: Iterable[llm_request.LlmRequest]) -> None: - """Order demotion after all model producers and attention readers. - - The non-overlap executor orders its calling stream after model forward. - Join it onto the cache stream before resize may demote/recycle pages. - prepare_sparse_offload() supplies the reverse dependency for the next - forward. No host/device synchronization or per-layer history mutation - is needed. Other executor schedules remain behind the inference guard. - """ - if not self._enable_kv_cache_offload: - return - active = False - for request in requests: - cache = self.kv_cache_map.get(request.py_request_id) - if cache is None or not cache.is_active: - continue - if cache.cuda_stream != self._stream.cuda_stream: - raise RuntimeError("Sparse history updates require the manager's cache stream") - active = True - if active: - if torch.cuda.is_current_stream_capturing(): - raise RuntimeError("Sparse history updates must run outside CUDA graph capture") - self._stream.wait_stream(torch.cuda.current_stream(self._stream.device)) - - def update_context_resources(self, scheduled_batch: ScheduledRequests) -> None: - self._order_sparse_history_update(scheduled_batch.context_requests) - super().update_context_resources(scheduled_batch) - - def update_resources( - self, - scheduled_batch: ScheduledRequests, - attn_metadata: "AttentionMetadata | None" = None, - kv_cache_dtype_byte_size: float | None = None, - ) -> None: - self._order_sparse_history_update(scheduled_batch.generation_requests) - # The base computes max_beam_num_tokens - 1, excluding the sampled - # output token whose KV has not been produced yet. Never increment here. - super().update_resources(scheduled_batch, attn_metadata, kv_cache_dtype_byte_size) - def get_sparse_offload_descriptors(self) -> dict[int, SparseOffloadLayerDescriptor]: - """Validate the ratio-4 shared write-table contract using runtime IDs. - - Called once when constructing each metadata instance. Different layer - offsets are allowed in SHARED mode, where the buffer pointer supplies - the offset. Different groups/scales require per-layer write tables and - are not supported by the current ratio-shared compressor interface. - """ + """Describe each local ratio-4 layer's complete sparse fetch buffer set.""" if not self._enable_kv_cache_offload: raise ValueError("Sparse offload metadata requires an offload-enabled cache manager") - if self._use_nvfp4_compress or self.use_fp8_ds_mla: - raise NotImplementedError( - "Sparse offload metadata supports BF16/per-tensor FP8 KV only" - ) local_layers = [ layer for layer in self.pp_layers @@ -1736,230 +1799,93 @@ def get_sparse_offload_descriptors(self) -> dict[int, SparseOffloadLayerDescript ] if not local_layers: return {} - if not callable(getattr(self.impl, "is_sparse", None)) or not callable( - getattr(self.impl, "copy_base_page_indices_to_device", None) - ): - raise NotImplementedError("KVCM v2 sparse page-table APIs are not implemented") - - descriptors = {} - shared_layout = None - for layer in local_layers: - layer_id = self._layer_attn_to_layer_id[layer, DeepseekV4AttentionType.COMPRESS] - role = DeepseekV4AttentionType.COMPRESS.role - if not self.impl.is_sparse(layer_id, role): - raise ValueError(f"Layer {layer}'s COMPRESS buffer must be sparse") - converter = self.impl.get_page_index_converter(layer_id, role) - if converter.expansion != 1 or not 0 < converter.scale <= 0x7FFFFFFF: - raise NotImplementedError( - "Sparse offload requires expansion 1 and a positive int32 scale" - ) - group_id = self.impl.get_layer_group_id(layer_id) - layout = ( - group_id, - int(converter.scale), - # SHARED includes this buffer's offset. PER_LAYER exposes - # the common pool origin, even when the layer offsets differ. - self.impl.get_mem_pool_base_address(layer_id, role, PageIndexMode.PER_LAYER), + if self.use_fp8_ds_mla: + raise NotImplementedError( + "Sparse fp8_ds_mla compressed KV requires a selection layout with " + "separate data and footer-scale regions" ) - if shared_layout is not None and layout != shared_layout: - raise NotImplementedError( - "Ratio-4 sparse layers must share a layer group, page scale, and GPU pool" - ) - shared_layout = layout - page_bound = int(self.impl.get_page_index_upper_bound(layer_id, role)) - if not 0 < page_bound <= 0x7FFFFFFF: - raise ValueError("Sparse GPU pool must have a positive int32 physical-page bound") - descriptors[layer] = SparseOffloadLayerDescriptor( - buffer_id=BufferId(layer_id, role), - group_id=group_id, - page_scale=int(converter.scale), - page_index_upper_bound=page_bound, + layout = self._sparse_arena_layout() + return { + layer: self._build_sparse_offload_descriptor( + layout, + self._layer_attn_to_layer_id[layer, DeepseekV4AttentionType.COMPRESS], + DeepseekV4AttentionType.COMPRESS.role, + # A selection names compressed rows of four original tokens each. + unit_tokens=DEEPSEEK_V4_SPARSE_RATIO, + scale_role=COMPRESS_BLOCK_SCALE_ROLE if self._use_nvfp4_compress else None, + name=f"Layer {layer}", ) + for layer in local_layers + } - sparse_groups = {descriptor.group_id for descriptor in descriptors.values()} - for (layer, attn_type), layer_id in self._layer_attn_to_layer_id.items(): - if layer in descriptors and attn_type == DeepseekV4AttentionType.COMPRESS: - continue - if self.impl.get_layer_group_id(layer_id) in sparse_groups: - raise ValueError( - "Sparse COMPRESS must have a separate lifecycle from ordinary cache buffers" - ) - return descriptors - - @nvtx_range_debug("dsv4_prepare_sparse_offload") - def prepare_sparse_offload( - self, - state: SparseOffloadState, - request_ids: list[int], - write_page_table: torch.Tensor, - *, - beam_width: int, + def _initialize_sparse_arena_buffers( + self, layers: dict[int, SparseOffloadLayerDescriptor] ) -> None: - """Upload the current mixed-tier mapping and history frontier in request order. - - The caller must keep request locks stable through the forward. Uploads - and write-table conversion use the current stream, ordered after the - manager's stream. This host preparation is outside graph capture; - subsequent replay consumes the persistent device buffers in place. - Only trailing CUDA-graph dummy requests are padding. Real requests - must be unique, with one beam, and retain their supplied order. + """Cache tensor views over KVCM-owned arena regions before graph capture. + + Each view is [pool_rows, 1, row_elements], so an arena index is also a + one-row page index. Pool rows cover all leases and buffers, including + padding roles. Layers sharing a region reuse the same tensor storage. + KVCM must keep the allocation alive while these views or graphs use it. """ - state.prepared = False - if not self._enable_kv_cache_offload or beam_width != 1: - raise ValueError("Sparse offload preparation requires offload enabled and beam width 1") - batch = state.history_blocks.numel() - if len(request_ids) > batch: - raise ValueError("Request batch exceeds sparse offload metadata capacity") - if ( - write_page_table.shape != (batch, self.max_blocks_per_seq) - or write_page_table.dtype != torch.int32 - or write_page_table.device != state.history_blocks.device - ): - raise ValueError( - "Sparse write table must match the metadata's CUDA int32 [B, M] layout" - ) - active_count = len(request_ids) - if CUDA_GRAPH_DUMMY_REQUEST_ID in request_ids: - active_count = request_ids.index(CUDA_GRAPH_DUMMY_REQUEST_ID) - if any(req != CUDA_GRAPH_DUMMY_REQUEST_ID for req in request_ids[active_count:]): - raise ValueError("CUDA-graph dummy requests must form a trailing suffix") - live_ids = request_ids[:active_count] - if len(set(live_ids)) != active_count: - raise ValueError("Sparse offload requires unique live request IDs") - caches = [self.kv_cache_map[request_id] for request_id in live_ids] - histories = [] - for cache in caches: - if cache.beam_width != 1: - raise ValueError("Sparse offload supports one beam per request") - history = cache.history_length - if not 0 <= history <= cache.capacity: - raise ValueError("Invalid KVCM history frontier") - if cache.num_blocks > self.max_blocks_per_seq: - raise ValueError("Request allocation exceeds sparse offload page-table capacity") - histories.append(history // self.tokens_per_block) - - with torch.cuda.device(state.history_blocks.device): + if not layers: + return + with torch.cuda.device(self._stream.device): if torch.cuda.is_current_stream_capturing(): - raise RuntimeError( - "Sparse offload page tables must be prepared outside graph capture" - ) - stream = torch.cuda.current_stream() - stream.wait_stream(self._stream) - # A queued H2D copy borrows the pinned host staging. Wait only for - # that upload before reusing its storage, never for the whole GPU. - if state.history_upload_pending: - state.history_upload_done.synchronize() - state.history_blocks_host.zero_() - state.history_blocks_host[:active_count].copy_( - torch.tensor(histories, dtype=torch.int32, device="cpu") - ) - state.history_blocks.copy_(state.history_blocks_host, non_blocking=True) - state.history_upload_done.record(stream) - state.history_upload_pending = True - state.active_request_count.fill_(active_count) - state.selected_history_pages.fill_(BAD_PAGE_INDEX) - state.fetched_page_table.fill_(BAD_PAGE_INDEX) - state.compress_read_table.fill_(BAD_PAGE_INDEX) - - for group_id, raw_table in state.base_page_tables.items(): - raw_table.fill_(BAD_PAGE_INDEX) - if active_count: - self.impl.copy_base_page_indices_to_device( - kv_caches=caches, - layer_group_id=group_id, - out=raw_table[:active_count], - stream=CudaStream(stream.cuda_stream), + raise RuntimeError("Sparse arena views must be initialized before graph capture") + pool_views: dict[int, torch.Tensor] = {} + for layer, descriptor in layers.items(): + specs = [ + ( + descriptor.data_buffer_index, + DataType.UINT8 if self._use_nvfp4_compress else self.dtype, + (self.head_dim + self.nvfp4_residual_dim) // 2 + if self._use_nvfp4_compress + else self.head_dim, + ) + ] + if descriptor.scale_buffer_index is not None: + specs.append( + ( + descriptor.scale_buffer_index, + DataType.FP8, + (self.head_dim + self.nvfp4_residual_dim) // NVFP4_VECTOR_SIZE, + ) ) + views = [] + for index, dtype, row_elements in specs: + base = descriptor.arena_base_addresses[index] + if base not in pool_views: + pool_views[base] = convert_to_torch_tensor( + TensorWrapper( + base, dtype, (descriptor.arena_row_counts[index], 1, row_elements) + ) + ) + views.append(pool_views[base]) + self._sparse_arena_buffers[layer] = (views[0], views[1] if len(views) > 1 else None) - # Descriptor validation guarantees one compatible ratio-4 table. - descriptor = next(iter(state.layers.values())) - prepare_sparse_write_table( - state.base_page_tables[descriptor.group_id], - state.history_blocks, - state.active_request_count, - write_page_table, - page_scale=descriptor.page_scale, - ) - state.prepared = True + def get_sparse_compress_buffers( + self, layer_idx: int + ) -> tuple[torch.Tensor, torch.Tensor | None]: + """Return cached data/scales arena views indexed by prepared or fetched rows.""" + return self._sparse_arena_buffers[layer_idx] - @nvtx_range_debug("dsv4_fetch_sparse_read_table") - def fetch_sparse_read_table( - self, - state: SparseOffloadState, - layer_idx: int, - topk_indices: torch.Tensor, - request_rows: torch.Tensor, - compressed_lengths: torch.Tensor, - write_page_table: torch.Tensor, + def prepare_sparse_compressor_write_indices( + self, state: SparseOffloadState, layer_idx: int ) -> torch.Tensor: - """Stage one layer's selected history and return its physical read table. - - The caller has joined indexer/compressor producers on the consuming - stream. Selection, KVCM fetch, merge, index conversion, and attention - must remain ordered on that stream, including during graph replay. - The shared workspace and KVCM scratch cannot be reused by another - layer/forward until attention has consumed them. No host tensor reads - or active-count branches are allowed here. + """Reserve this layer's pinned destinations on the compressor's stream. + + Return int32[B, F] in Batch and descriptor buffer order. The caller + must join Batch publication before this call, and join the compressor + stream before fetch/attention or a history update. No KV is fetched here. + Padding roles were initialized during manager construction. The fused + compressor uses the data column as a one-row page table for data and + NVFP4 scales, whose matching geometry is checked during descriptor construction + and, once per layer outside graph capture, by ``prepare_sparse_write_indices`` + against the indices native returns. Each buffer still uses its own arena base + and byte stride. """ - if not self._enable_kv_cache_offload or not state.prepared or state.is_prefill: - raise RuntimeError("Sparse fetch requires prepared decode metadata") - descriptor = state.layers[layer_idx] - fetched_page_scale = descriptor.fetched_page_scale - if fetched_page_scale is None: - raise NotImplementedError("KVCM sparse fetch output page units are not bound") - if not 0 < fetched_page_scale <= 0x7FFFFFFF: - raise ValueError("Sparse fetch page scale must be a positive int32") - if not callable(getattr(self.impl, "fetch_sparse_pages", None)): - raise NotImplementedError("KVCM v2 fetch_sparse_pages is not implemented") - - with torch.cuda.device(state.history_blocks.device): - raw_table = state.base_page_tables[descriptor.group_id] - select_sparse_history_pages( - topk_indices, - request_rows, - state.active_request_count, - compressed_lengths, - raw_table, - state.history_blocks, - self.tokens_per_block // DEEPSEEK_V4_SPARSE_RATIO, - state.selected_history_pages, - ) - # Fetch even when all selections are -1. KVCM must overwrite the - # entire output, so a replay cannot reuse the previous layer's map. - self.impl.fetch_sparse_pages( - buffers=[descriptor.buffer_id], - page_table=raw_table, - selected=state.selected_history_pages, - num_blocks=state.history_blocks, - out=state.fetched_page_table, - stream=CudaStream(torch.cuda.current_stream().cuda_stream), - ) - merge_sparse_read_table( - state.fetched_page_table, - write_page_table, - state.history_blocks, - state.active_request_count, - state.compress_read_table, - fetched_page_scale=fetched_page_scale, - ) - check_sparse_read_table( - topk_indices, - request_rows, - state.active_request_count, - compressed_lengths, - state.compress_read_table, - self.tokens_per_block // DEEPSEEK_V4_SPARSE_RATIO, - descriptor.page_index_upper_bound, - state.read_table_valid, - ) - if state.debug_assert: - # KVCM must return every required page. Keep the assertion opt-in - # because a failed device assertion poisons the CUDA context. - torch._assert_async( - state.read_table_valid, - "Sparse offload: invalid top-k request/token or missing/out-of-bounds KV page", - ) - return state.compress_read_table + return self.prepare_sparse_write_indices(state, layer_idx) @nvtx_range_debug("dsv4_copy_batch_compress_block_tables") def copy_batch_compress_block_tables( @@ -1973,7 +1899,7 @@ def copy_batch_compress_block_tables( ) -> None: """Build the COMPRESS block table for one compression ratio and copy it to the destination.""" if self._enable_kv_cache_offload and compress_ratio == DEEPSEEK_V4_SPARSE_RATIO: - raise RuntimeError("Use prepare_sparse_offload() for the mixed-tier ratio-4 table") + raise RuntimeError("Sparse COMPRESS requires per-buffer arena indices") assert beam_width == 1, "DSV4 only supports beam width 1 now" copy_idx = self.index_mapper.get_copy_index(request_ids, num_contexts, beam_width) staging = self._host_compress_block_tables_staging[compress_ratio] @@ -2036,6 +1962,7 @@ def get_cache_size_per_token(model_config: ModelConfig, mapping: Mapping, **kwar kv_cache_config is not None and getattr(kv_cache_config, "dtype", "auto") == "fp8_ds_mla" ) + offload = model_config.sparse_attention_config.enable_kv_cache_offload non_sliding_attn_size_per_token = _estimate_non_sliding_attn_size_per_token( head_dim, index_head_dim, @@ -2044,7 +1971,13 @@ def get_cache_size_per_token(model_config: ModelConfig, mapping: Mapping, **kwar indexer_k_dtype=indexer_k_dtype, use_fp8_ds_mla=use_fp8_ds_mla, has_nvfp4_compress=has_nvfp4_compress, + offload_sparse_history=offload, ) + generation_capacity_headroom = None + if offload: + _, generation_capacity_headroom = _get_generation_kv_capacity( + kwargs.get("spec_config"), is_draft=kwargs.get("is_draft", False) + ) swa_size_per_token, swa_size_per_request = _estimate_swa_cache_size( head_dim, index_head_dim, @@ -2056,10 +1989,11 @@ def get_cache_size_per_token(model_config: ModelConfig, mapping: Mapping, **kwar scratch=False, indexer_k_dtype=indexer_k_dtype, use_fp8_ds_mla=use_fp8_ds_mla, + generation_capacity_headroom=generation_capacity_headroom, ) max_batch_size = int(kwargs.get("max_batch_size") or 0) return ( - non_sliding_attn_size_per_token + swa_size_per_token, + ceil(non_sliding_attn_size_per_token + swa_size_per_token), swa_size_per_request * max_batch_size, ) @@ -2125,12 +2059,13 @@ def _guard_reservation_detail(self) -> str: f">1 attn buffer)" ) - @torch.inference_mode() - def check_invalid_values_in_kv_cache(self, fill_with_zero: bool = False) -> bool: - some_checks_unavailable = False - has_invalid_values = torch.tensor( - [False], dtype=torch.bool, device=torch.cuda.current_device() - ) + def _iter_cache_cleanup_buffers(self) -> Iterable[tuple[torch.Tensor, int | None]]: + """Yield device storage and arena views with their optional storage guard-page index. + + Host-backed sparse history is left out. Nothing reads a row of it before + a transfer or a write-back fills it, and walking it would cross PCIe for + every page of the host budget. + """ buffers_handled = set() # Handle each attention buffer from start to end to traverse the whole @@ -2147,19 +2082,38 @@ def check_invalid_values_in_kv_cache(self, fill_with_zero: bool = False) -> bool for buffer_key, buffer in buffers: if buffer_key in buffers_handled: continue - if scan_kv_buffer_for_invalid_values(buffer, has_invalid_values): - some_checks_unavailable = True - if fill_with_zero: - buffer.zero_() - # Warmup scrubs the cache with fill_with_zero before inference; - # restore this buffer's guard-page sentinel so the startup - # warning does not claim a guard that was just wiped (mirrors - # the base check_invalid_values_in_kv_cache). The guard key is - # this buffer's (layer_id, role) identity == buffer_key. - guard_page = self._guard_page_by_layer.get(buffer_key) - if guard_page is not None and self._guard_page_value is not None: - _fill_kv_pages(buffer, [guard_page], self._guard_page_value) buffers_handled.add(buffer_key) + if self._enable_kv_cache_offload and self.impl.is_sparse(*buffer_key): + continue + yield buffer, self._guard_page_by_layer.get(buffer_key) + + # Arena views span complete GPU pools and can be shared by several layers. + # Their physical row indices are independent of the host storage guard pages. + arenas_handled = set() + for arenas in self._sparse_arena_buffers.values(): + for arena in arenas: + if arena is None or arena.data_ptr() in arenas_handled: + continue + arenas_handled.add(arena.data_ptr()) + yield arena, None + + @torch.inference_mode() + def check_invalid_values_in_kv_cache(self, fill_with_zero: bool = False) -> bool: + if self._enable_kv_cache_offload: + # Finish cache-stream writeback, which reads the arenas, before + # scanning or clearing them. + torch.cuda.current_stream(self._stream.device).wait_stream(self._stream) + some_checks_unavailable = False + has_invalid_values = torch.tensor( + [False], dtype=torch.bool, device=torch.cuda.current_device() + ) + for buffer, guard_page in self._iter_cache_cleanup_buffers(): + if scan_kv_buffer_for_invalid_values(buffer, has_invalid_values): + some_checks_unavailable = True + if fill_with_zero: + buffer.zero_() + if guard_page is not None and self._guard_page_value is not None: + _fill_kv_pages(buffer, [guard_page], self._guard_page_value) torch.cuda.synchronize() if some_checks_unavailable: diff --git a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/compressor.py b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/compressor.py index 64d90b780996..849d4e2373b4 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/compressor.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/compressor.py @@ -14,6 +14,7 @@ from tensorrt_llm._torch.modules.linear import Linear from tensorrt_llm._torch.modules.rms_norm import RMSNorm +from ..offload_kernels import gather_sparse_batch_rows from .params import DeepseekV4AttentionType if TYPE_CHECKING: @@ -74,7 +75,7 @@ def resolve_kv_cache_dtype(kv_cache_dtype: Union[str, KVCacheDtype]) -> KVCacheD class Compressor(nn.Module): - """KV compressor using Triton kernels with paged memory management. + """KV compressor with paged state and paged or sparse-arena output. Args: mla_params: MLA parameters containing hidden_size and head dimensions @@ -190,13 +191,36 @@ def forward( kv_type = DeepseekV4AttentionType.COMPRESSOR_KV score_type = DeepseekV4AttentionType.COMPRESSOR_SCORE - # Get cache buffers - kv_cache = metadata.kv_cache_manager.get_buffers(self.layer_idx, compress_type) - kv_cache_scale = ( - metadata.kv_cache_manager.get_compress_scale_buffers(self.layer_idx) - if not self.is_indexer and self.kv_cache_dtype == KVCacheDtype.NVFP4_BLOCKWISE - else None - ) + # Reserve on the producer stream. The module joins this stream before + # fetch/attention, and scheduling joins it before advancing history. + sparse_state = metadata.sparse_offload_state + sparse_write_indices = None + if not self.is_indexer and self.compress_ratio == 4 and sparse_state is not None: + if num_contexts or metadata.num_gen_tokens_per_seq != 1 or self.footer_scale_cache: + raise NotImplementedError( + "Sparse compressor writes require single-token decode without footer scales" + ) + descriptor = sparse_state.layers[self.layer_idx] + writes = metadata.kv_cache_manager.prepare_sparse_compressor_write_indices( + sparse_state, self.layer_idx + ) + # Native write rows are by Batch row; the scatter addresses requests + # by input position. + # At the slot: coinciding buffers are resolved once into one slice. + data_slot = sparse_state.out_slots[self.layer_idx][descriptor.data_buffer_index] + sparse_write_indices = gather_sparse_batch_rows( + writes[:, data_slot], sparse_state.batch_rows[:bsz] + ) + kv_cache, kv_cache_scale = metadata.kv_cache_manager.get_sparse_compress_buffers( + self.layer_idx + ) + else: + kv_cache = metadata.kv_cache_manager.get_buffers(self.layer_idx, compress_type) + kv_cache_scale = ( + metadata.kv_cache_manager.get_compress_scale_buffers(self.layer_idx) + if not self.is_indexer and self.kv_cache_dtype == KVCacheDtype.NVFP4_BLOCKWISE + else None + ) paged_kv_state = metadata.kv_cache_manager.get_buffers(self.layer_idx, kv_type) paged_score_state = metadata.kv_cache_manager.get_buffers(self.layer_idx, score_type) @@ -205,7 +229,11 @@ def forward( if self.is_indexer: block_table = metadata.indexer_k_cache_block_offsets else: - block_table = metadata.compress_block_tables[self.compress_ratio] + block_table = ( + sparse_write_indices.view(bsz, 1) + if sparse_write_indices is not None + else metadata.compress_block_tables[self.compress_ratio] + ) block_table_kv_state = metadata.sliding_block_tables[local_layer_idx, kv_type.value] block_table_score_state = metadata.sliding_block_tables[local_layer_idx, score_type.value] @@ -213,7 +241,11 @@ def forward( # state_tokens_per_block: for compressor kv/score state caches (used in compress kernels) # compress_tokens_per_block: for compressed KV cache (used in scatter) state_tokens_per_block = metadata.kv_cache_manager.tokens_per_block - compress_tokens_per_block = metadata.kv_cache_manager.compressed_block_sizes[self.layer_idx] + compress_tokens_per_block = ( + 1 + if sparse_write_indices is not None + else metadata.kv_cache_manager.compressed_block_sizes[self.layer_idx] + ) # Get compression metadata cu_new_comp_kv = metadata.cu_new_comp_kv_cuda[self.compress_ratio] @@ -282,6 +314,11 @@ def forward( # Scatter to cache with appropriate quantization (all modes fused) start_pos = metadata.past_kv_lens_cuda[self.compress_ratio][:bsz] + if sparse_write_indices is not None: + # One-row page addressing uses zero storage positions; RoPE IDs stay + # unchanged. The kernel checks num_outputs before dereferencing a page. + start_pos = sparse_state.scatter_start_positions[:bsz] + num_comp_tokens = torch.where(sparse_write_indices >= 0, num_comp_tokens, 0) total_tokens = kv_comp.shape[0] # Allocate optional returned postprocess buffers for indexer paths. @@ -326,7 +363,7 @@ def forward( ) return kv_comp, None - # Fused postprocess + scatter: RMSNorm + RoPE + Hadamard + paged cache write + # Fused RMSNorm + RoPE + Hadamard + cache write, including sparse arenas. torch.ops.trtllm.compressor_postprocess_scatter( kv_comp, kv_out, diff --git a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/kernels.py b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/kernels.py index 0e64ab823449..a4fc581cb8f1 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/kernels.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/kernels.py @@ -9,12 +9,14 @@ import triton import triton.language as tl +from ..offload_kernels import check_sparse_offload_tensor + ######################################################## # Index gather kernel ######################################################## -@triton.jit +@triton.jit(do_not_specialize=["sparse_num_lengths"]) def _deepseek_v4_local_to_global_kernel( req_id_ptr, block_table_swa_ptr, @@ -58,6 +60,21 @@ def _deepseek_v4_local_to_global_kernel( active_request_count_ptr, MASK_DECODE_PADDING: tl.constexpr, MASK_INVALID_PAGES: tl.constexpr, + fetched_indices_ptr, + write_indices_ptr, + history_lengths_ptr, + compressed_lengths_ptr, + batch_rows_ptr, + fetch_stride0, + fetch_stride1, + write_stride, + history_stride, + length_stride, + batch_rows_stride, + MERGE_SPARSE_INDICES: tl.constexpr, + SPARSE_BATCH_SIZE: tl.constexpr, + sparse_num_lengths, + UNIT_TOKENS: tl.constexpr, ): """ Triton kernel for converting local indices to global KV cache pool indices. @@ -99,6 +116,13 @@ def _deepseek_v4_local_to_global_kernel( if MASK_DECODE_PADDING: active = tl.load(active_request_count_ptr) live = (token_id < active) & (req >= 0) & (req < active) + # Sparse offload state belongs to the request's native Batch row, which need + # not equal its input position; block tables and lengths follow the position. + row = req + if MERGE_SPARSE_INDICES: + live = live & (req < SPARSE_BATCH_SIZE) & (req < sparse_num_lengths) + row = tl.load(batch_rows_ptr + req * batch_rows_stride, mask=live, other=-1) + live = live & (row >= 0) & (row < SPARSE_BATCH_SIZE) # Load all SWA local indices for this token swa_ids = tl.arange(0, num_swa_indices) @@ -136,28 +160,42 @@ def _deepseek_v4_local_to_global_kernel( ) compressed_local_idx = tl.load(compressed_ptr) - # Compute global indices for all compressed positions - compressed_valid_mask = compressed_local_idx >= 0 - compressed_block_ordinal = compressed_local_idx // tokens_per_block_compressed - compressed_token_in_block = compressed_local_idx % tokens_per_block_compressed - compressed_valid_block = compressed_block_ordinal < max_blocks_compressed - compressed_full_mask = live & compressed_valid_mask & compressed_valid_block - - compressed_bt_ptr = ( - block_table_compressed_ptr - + req * bt_compressed_stride0 - + compressed_block_ordinal * bt_compressed_stride1 - ) - compressed_page_index = tl.load(compressed_bt_ptr, mask=compressed_full_mask, other=0) - if MASK_INVALID_PAGES: - compressed_full_mask = compressed_full_mask & (compressed_page_index >= 0) - - compressed_global_index = ( - compressed_buffer_offset_in_tokens - + compressed_page_index * tokens_per_block_compressed - + compressed_token_in_block - ) - compressed_global_index = tl.where(compressed_full_mask, compressed_global_index, -1) + if MERGE_SPARSE_INDICES: + history = tl.load(history_lengths_ptr + row * history_stride, mask=live, other=0) + history_units = history // UNIT_TOKENS + length = tl.load(compressed_lengths_ptr + req * length_stride, mask=live, other=0) + valid = live & (compressed_local_idx >= 0) & (compressed_local_idx < length) + fetched = tl.load( + fetched_indices_ptr + row * fetch_stride0 + compressed_ids * fetch_stride1, + mask=valid & (compressed_local_idx < history_units), + other=-1, + ) + write = tl.load(write_indices_ptr + row * write_stride, mask=live, other=-1) + compressed_global_index = tl.where( + valid & (compressed_local_idx == history_units), write, fetched + ) + else: + compressed_valid_mask = compressed_local_idx >= 0 + compressed_block_ordinal = compressed_local_idx // tokens_per_block_compressed + compressed_token_in_block = compressed_local_idx % tokens_per_block_compressed + compressed_valid_block = compressed_block_ordinal < max_blocks_compressed + compressed_full_mask = live & compressed_valid_mask & compressed_valid_block + + compressed_bt_ptr = ( + block_table_compressed_ptr + + req * bt_compressed_stride0 + + compressed_block_ordinal * bt_compressed_stride1 + ) + compressed_page_index = tl.load(compressed_bt_ptr, mask=compressed_full_mask, other=0) + if MASK_INVALID_PAGES: + compressed_full_mask = compressed_full_mask & (compressed_page_index >= 0) + + compressed_global_index = ( + compressed_buffer_offset_in_tokens + + compressed_page_index * tokens_per_block_compressed + + compressed_token_in_block + ) + compressed_global_index = tl.where(compressed_full_mask, compressed_global_index, -1) if SPLIT_EXTRA: compressed_out_ptr = ( @@ -202,6 +240,11 @@ def deepseek_v4_local_to_global_indices( split_extra: bool = False, active_request_count: torch.Tensor | None = None, mask_invalid_pages: bool = False, + compressed_fetched_indices: torch.Tensor | None = None, + compressed_write_indices: torch.Tensor | None = None, + compressed_history_lengths: torch.Tensor | None = None, + compressed_lengths: torch.Tensor | None = None, + compressed_batch_rows: torch.Tensor | None = None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor | None]: """ Convert local token indices to global KV cache pool indices. @@ -241,6 +284,24 @@ def deepseek_v4_local_to_global_indices( mask_invalid_pages: Mask negative physical page mappings in both pools when consuming an offload read table. Compiled out by default to preserve the original non-offload indexing behavior. + compressed_fetched_indices: CUDA int32 [B, K] physical history indices, + in Batch row order and original top-k column order. Supply all five + sparse tensors and active_request_count together for single-token + decode. The kernel merges history and in-flight destinations into + query order while preparing SWA indices; no compressed page table + or additional address conversion is used. + compressed_write_indices: CUDA int32 [B] prepared in-flight destinations, + by Batch row. May be a strided data-column view shared by paired + NVFP4 pools. + compressed_history_lengths: CUDA int32 [B] published history token counts, + by Batch row. + compressed_lengths: CUDA int32 [R] complete compressed-row counts by + request, in req_id's numbering, with Q <= R <= B. A selected unit + equal to history // compress_ratio and below this count uses the + write index. Invalid indices remain -1. + compressed_batch_rows: CUDA int32 [B] native Batch row of each request, + in req_id's numbering, or -1 for no live request. A request whose + row is -1 gets no indices in either pool. Returns: A combined index tensor, or separate SWA and compressed tensors when @@ -257,7 +318,7 @@ def deepseek_v4_local_to_global_indices( num_tokens = req_id.shape[0] num_swa_indices = swa_local_indices.shape[1] if active_request_count is not None: - _check_sparse_offload_tensor(active_request_count, "active_request_count", 1, req_id.device) + check_sparse_offload_tensor(active_request_count, "active_request_count", 1, req_id.device) if active_request_count.shape != (1,): raise ValueError("active_request_count must be [1]") @@ -272,7 +333,58 @@ def deepseek_v4_local_to_global_indices( # Compute SWA buffer offset relative to swa_pool_base_ptr in tokens swa_buffer_offset_in_tokens = (swa_buffer_ptr - swa_pool_base_ptr) // token_stride - if has_compressed: + sparse_tensors = ( + (compressed_fetched_indices, "compressed_fetched_indices", 2), + (compressed_write_indices, "compressed_write_indices", 1), + (compressed_history_lengths, "compressed_history_lengths", 1), + (compressed_lengths, "compressed_lengths", 1), + (compressed_batch_rows, "compressed_batch_rows", 1), + ) + merge_sparse_indices = any(tensor is not None for tensor, _, _ in sparse_tensors) + sparse_batch_size = sparse_num_lengths = 0 + fetch_stride0 = fetch_stride1 = write_stride = history_stride = length_stride = 0 + batch_rows_stride = 0 + if merge_sparse_indices: + if ( + not has_compressed + or active_request_count is None + or any(tensor is None for tensor, _, _ in sparse_tensors) + ): + raise ValueError( + "Sparse indices require a compressed layer, all mappings and an active count" + ) + for tensor, name, ndim in sparse_tensors: + check_sparse_offload_tensor(tensor, name, ndim, req_id.device) + sparse_batch_size, width = compressed_fetched_indices.shape + sparse_num_lengths = compressed_lengths.numel() + if ( + width != num_compressed_indices + or compressed_write_indices.shape != (sparse_batch_size,) + or compressed_history_lengths.shape != (sparse_batch_size,) + or compressed_batch_rows.shape != (sparse_batch_size,) + or not num_tokens <= sparse_num_lengths <= sparse_batch_size + ): + raise ValueError( + "Sparse mappings must cover the native Batch and configured top-k width" + ) + if compressed_local_indices is None: + raise ValueError("Sparse index preparation requires the original top-k selection") + check_sparse_offload_tensor( + compressed_local_indices, "compressed_local_indices", 2, req_id.device + ) + if compressed_local_indices.shape != (num_tokens, num_compressed_indices): + raise ValueError("Sparse top-k indices must be [Q, K]") + fetch_stride0, fetch_stride1 = compressed_fetched_indices.stride() + write_stride = compressed_write_indices.stride(0) + history_stride = compressed_history_lengths.stride(0) + length_stride = compressed_lengths.stride(0) + batch_rows_stride = compressed_batch_rows.stride(0) + tokens_per_block_compressed = 1 + compressed_buffer_offset_in_tokens = 0 + max_blocks_compressed = 0 + block_table_compressed_c = block_table_swa + compressed_local_indices_c = compressed_local_indices + elif has_compressed: assert block_table_compressed is not None, ( "block_table_compressed required when compress_ratio > 1" ) @@ -406,447 +518,24 @@ def deepseek_v4_local_to_global_indices( active_request_count_ptr=active_request_count, MASK_DECODE_PADDING=active_request_count is not None, MASK_INVALID_PAGES=mask_invalid_pages, + fetched_indices_ptr=compressed_fetched_indices, + write_indices_ptr=compressed_write_indices, + history_lengths_ptr=compressed_history_lengths, + compressed_lengths_ptr=compressed_lengths, + batch_rows_ptr=compressed_batch_rows, + fetch_stride0=fetch_stride0, + fetch_stride1=fetch_stride1, + write_stride=write_stride, + history_stride=history_stride, + length_stride=length_stride, + batch_rows_stride=batch_rows_stride, + MERGE_SPARSE_INDICES=merge_sparse_indices, + SPARSE_BATCH_SIZE=sparse_batch_size, + sparse_num_lengths=sparse_num_lengths, + UNIT_TOKENS=compress_ratio, launch_pdl=launch_with_pdl, ) if split_extra: return out, out_extra return out - - -def _check_sparse_offload_tensor( - tensor: torch.Tensor, name: str, ndim: int, device: torch.device -) -> None: - if tensor.dtype != torch.int32 or tensor.ndim != ndim: - raise ValueError(f"{name} must be a {ndim}D int32 tensor") - if not tensor.is_cuda or tensor.device != device: - raise ValueError(f"{name} must be on CUDA device {device}") - - -@triton.jit -def _select_sparse_history_pages_kernel( - topk_ptr, - request_rows_ptr, - active_request_count_ptr, - compressed_lengths_ptr, - raw_page_table_ptr, - history_blocks_ptr, - out_ptr, - topk_stride0, - topk_stride1, - request_rows_stride, - compressed_lengths_stride, - raw_stride0, - raw_stride1, - history_blocks_stride, - out_stride0, - out_stride1, - NUM_REQUESTS: tl.constexpr, - MAX_BLOCKS: tl.constexpr, - TOPK: tl.constexpr, - TOKENS_PER_PAGE: tl.constexpr, - BLOCK_K: tl.constexpr, -): - query = tl.program_id(0) - active_count = tl.load(active_request_count_ptr) - request = tl.load(request_rows_ptr + query * request_rows_stride) - active = (query < active_count) & (request >= 0) & (request < active_count) - active = active & (request < NUM_REQUESTS) - length = tl.load(compressed_lengths_ptr + request * compressed_lengths_stride, active, other=0) - history = tl.load(history_blocks_ptr + request * history_blocks_stride, active, other=0) - - k = tl.arange(0, BLOCK_K) - token = tl.load(topk_ptr + query * topk_stride0 + k * topk_stride1, k < TOPK, other=-1) - ordinal = token // TOKENS_PER_PAGE - valid = active & (token >= 0) & (token < length) - valid = valid & (ordinal < history) & (ordinal < MAX_BLOCKS) - page = tl.load( - raw_page_table_ptr + request * raw_stride0 + ordinal * raw_stride1, valid, other=-1 - ) - # Sort ordinals, never host slot IDs. Slot 0 is valid in either tier. - sorted_ordinals = tl.sort(tl.where(valid & (page >= 0), ordinal, 0x7FFFFFFF), descending=False) - previous = tl.gather(sorted_ordinals, tl.maximum(k - 1, 0), axis=0) - unique = (sorted_ordinals != 0x7FFFFFFF) & ((k == 0) | (sorted_ordinals != previous)) - position = tl.cumsum(unique.to(tl.int32), axis=0) - 1 - tl.store(out_ptr + request * out_stride0 + position * out_stride1, sorted_ordinals, mask=unique) - - -def select_sparse_history_pages( - topk_indices: torch.Tensor, - request_rows: torch.Tensor, - active_request_count: torch.Tensor, - compressed_lengths: torch.Tensor, - raw_page_table: torch.Tensor, - history_blocks: torch.Tensor, - compressed_tokens_per_page: int, - out: torch.Tensor, -) -> None: - """Select unique host-history block ordinals for single-token decode. - - All tensors are CUDA int32 on one device. ``topk_indices`` is [Q, K] - and contains logical compressed-token indices; ``request_rows`` is [Q] - and maps query rows to request rows in ``raw_page_table`` ([B, M]). - ``active_request_count`` is [1]: only the first N query/request rows are - live. Their request IDs must be a permutation of [0, N), with one query - per request. Multi-query requests require a separate union operation. - - ``compressed_lengths`` and ``history_blocks`` are [B], respectively in - compressed tokens and original KVCM blocks. Only ordinals below the - history frontier with a nonnegative raw page entry are eligible. Neither - resident tail pages nor padding/invalid tokens are fetched. - - The caller owns ``out`` ([B, S], S >= min(K, M)) and must provide storage - disjoint from all inputs. Each row is sorted, compacted, and padded with - -1. Every output entry is rewritten, including inactive or empty rows. - Inputs and their token order are unchanged. No allocation, device-to-host - read, or synchronization is performed; all work uses the current stream. - """ - device = raw_page_table.device - for tensor, name, ndim in ( - (topk_indices, "topk_indices", 2), - (request_rows, "request_rows", 1), - (active_request_count, "active_request_count", 1), - (compressed_lengths, "compressed_lengths", 1), - (raw_page_table, "raw_page_table", 2), - (history_blocks, "history_blocks", 1), - (out, "out", 2), - ): - _check_sparse_offload_tensor(tensor, name, ndim, device) - num_requests, max_blocks = raw_page_table.shape - num_queries, topk = topk_indices.shape - if compressed_tokens_per_page <= 0: - raise ValueError("compressed_tokens_per_page must be positive") - if request_rows.shape != (num_queries,) or active_request_count.shape != (1,): - raise ValueError("request_rows must be [Q] and active_request_count must be [1]") - if compressed_lengths.shape != (num_requests,) or history_blocks.shape != (num_requests,): - raise ValueError("compressed_lengths and history_blocks must be [B]") - if out.shape[0] != num_requests or out.shape[1] < min(topk, max_blocks): - raise ValueError("out must be [B, S] with S >= min(K, M); selection must not be truncated") - - with torch.cuda.device(device): - # Clear separately: query rows can be permuted and graph padding need - # not contain a usable request ID. The following scatter has one writer - # per live request under the single-token decode contract. - out.fill_(-1) - if num_queries == 0 or topk == 0 or num_requests == 0 or max_blocks == 0: - return - _select_sparse_history_pages_kernel[(num_queries,)]( - topk_indices, - request_rows, - active_request_count, - compressed_lengths, - raw_page_table, - history_blocks, - out, - *topk_indices.stride(), - request_rows.stride(0), - compressed_lengths.stride(0), - *raw_page_table.stride(), - history_blocks.stride(0), - *out.stride(), - NUM_REQUESTS=num_requests, - MAX_BLOCKS=max_blocks, - TOPK=topk, - TOKENS_PER_PAGE=compressed_tokens_per_page, - BLOCK_K=triton.next_power_of_2(topk), - ) - - -@triton.jit -def _merge_sparse_read_table_kernel( - fetched_ptr, - write_ptr, - history_blocks_ptr, - active_request_count_ptr, - out_ptr, - fetched_stride0, - fetched_stride1, - write_stride0, - write_stride1, - history_blocks_stride, - out_stride0, - out_stride1, - MAX_BLOCKS: tl.constexpr, - FETCHED_PAGE_SCALE: tl.constexpr, - BLOCK: tl.constexpr, -): - request = tl.program_id(0) - ordinal = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK) - active = request < tl.load(active_request_count_ptr) - history = tl.load(history_blocks_ptr + request * history_blocks_stride) - mask = active & (ordinal < MAX_BLOCKS) - fetched = tl.load( - fetched_ptr + request * fetched_stride0 + ordinal * fetched_stride1, - mask & (ordinal < history), - other=-1, - ) - resident = tl.load( - write_ptr + request * write_stride0 + ordinal * write_stride1, - mask & (ordinal >= history), - other=-1, - ) - # The adapter must supply the agreed fetch units. Resident write pages - # already use the converter's physical SHARED-page convention. - normalized = fetched.to(tl.int64) * FETCHED_PAGE_SCALE - normalized = tl.where((fetched >= 0) & (normalized <= 0x7FFFFFFF), normalized, -1) - page = tl.where(ordinal < history, normalized, tl.where(resident >= 0, resident, -1)) - tl.store( - out_ptr + request * out_stride0 + ordinal * out_stride1, - tl.where(active, page, -1), - mask=ordinal < MAX_BLOCKS, - ) - - -def merge_sparse_read_table( - fetched_page_table: torch.Tensor, - write_page_table: torch.Tensor, - history_blocks: torch.Tensor, - active_request_count: torch.Tensor, - out: torch.Tensor, - *, - fetched_page_scale: int, -) -> None: - """Merge fetched history and the resident tail into a separate read table. - - All tensors are CUDA int32 on one device. Both input tables and ``out`` - are [B, M]; ``history_blocks`` is [B], and ``active_request_count`` is [1]. - ``write_page_table`` already contains physical GPU pages in the compressed - buffer's SHARED convention. ``fetched_page_scale`` is mandatory: pass 1 for - fetched physical pages, or the converter's scale for raw GPU slots with - the same origin. Other fetch origins must be adapted before this call. - - History uses only fetched entries; the entire tail uses only the write - table. Negative/overflowed pages and inactive rows become -1. No missing - history entry falls back to a write address. The caller must validate - required-page coverage before enabling attention; masking is not recovery. - - ``out`` must have storage disjoint from the inputs. All entries are - overwritten on the current stream without allocation or synchronization. - The writable table and the original token selections remain unchanged. - """ - device = write_page_table.device - for tensor, name, ndim in ( - (fetched_page_table, "fetched_page_table", 2), - (write_page_table, "write_page_table", 2), - (history_blocks, "history_blocks", 1), - (active_request_count, "active_request_count", 1), - (out, "out", 2), - ): - _check_sparse_offload_tensor(tensor, name, ndim, device) - num_requests, max_blocks = write_page_table.shape - if fetched_page_table.shape != write_page_table.shape or out.shape != write_page_table.shape: - raise ValueError("fetched_page_table, write_page_table, and out must all be [B, M]") - if history_blocks.shape != (num_requests,) or active_request_count.shape != (1,): - raise ValueError("history_blocks must be [B] and active_request_count must be [1]") - if not 0 < fetched_page_scale <= 0x7FFFFFFF: - raise ValueError("fetched_page_scale must be a positive int32 scale") - if num_requests == 0 or max_blocks == 0: - return - with torch.cuda.device(device): - _merge_sparse_read_table_kernel[(num_requests, triton.cdiv(max_blocks, 256))]( - fetched_page_table, - write_page_table, - history_blocks, - active_request_count, - out, - *fetched_page_table.stride(), - *write_page_table.stride(), - history_blocks.stride(0), - *out.stride(), - MAX_BLOCKS=max_blocks, - FETCHED_PAGE_SCALE=fetched_page_scale, - BLOCK=256, - ) - - -@triton.jit -def _check_sparse_read_table_kernel( - topk_ptr, - requests_ptr, - active_ptr, - lengths_ptr, - read_ptr, - valid_ptr, - topk_stride0, - topk_stride1, - requests_stride, - lengths_stride, - read_stride0, - read_stride1, - NUM_REQUESTS: tl.constexpr, - NUM_QUERIES: tl.constexpr, - TOPK: tl.constexpr, - MAX_BLOCKS: tl.constexpr, - TOKENS_PER_PAGE: tl.constexpr, - PAGE_BOUND: tl.constexpr, - BLOCK_K: tl.constexpr, -): - query = tl.program_id(0) - active = tl.load(active_ptr) - live = query < active - req = tl.load(requests_ptr + query * requests_stride) - request_ok = (req >= 0) & (req < active) & (req < NUM_REQUESTS) - length = tl.load(lengths_ptr + req * lengths_stride, mask=live & request_ok, other=0) - k = tl.arange(0, BLOCK_K) - token = tl.load(topk_ptr + query * topk_stride0 + k * topk_stride1, mask=k < TOPK, other=-1) - required = live & (token >= 0) - ordinal = token // TOKENS_PER_PAGE - token_ok = (token < length) & (ordinal < MAX_BLOCKS) - page = tl.load( - read_ptr + req * read_stride0 + ordinal * read_stride1, - mask=required & request_ok & token_ok, - other=-1, - ) - page_ok = (page >= 0) & (page < PAGE_BOUND) - valid = ( - (active >= 0) - & (active <= NUM_REQUESTS) - & (active <= NUM_QUERIES) - & (~live | request_ok) - & tl.min((~required | (token_ok & page_ok)).to(tl.int32), axis=0) - ) - tl.atomic_and(valid_ptr, valid.to(tl.int32), sem="relaxed") - - -def check_sparse_read_table( - topk_indices: torch.Tensor, - request_rows: torch.Tensor, - active_request_count: torch.Tensor, - compressed_lengths: torch.Tensor, - read_page_table: torch.Tensor, - compressed_tokens_per_page: int, - page_index_upper_bound: int, - out: torch.Tensor, -) -> None: - """Check every live decode selection has an addressable physical KV page. - - Inputs are CUDA int32: top-k [Q, K], request rows [Q], active count [1], - lengths [B], and read table [B, M]. The positive, exclusive page bound is - relative to the layer's SHARED pointer and must include KVCM scratch. - ``out`` is a disjoint persistent int32 [1], rewritten to 1 on success or - 0 on invalid request/token, missing page, or out-of-bounds page. Padded - queries and negative token sentinels are ignored. The caller may pass it - to torch._assert_async before launching attention; no host read is needed. - """ - device = read_page_table.device - for tensor, name, ndim in ( - (topk_indices, "topk_indices", 2), - (request_rows, "request_rows", 1), - (active_request_count, "active_request_count", 1), - (compressed_lengths, "compressed_lengths", 1), - (read_page_table, "read_page_table", 2), - (out, "out", 1), - ): - _check_sparse_offload_tensor(tensor, name, ndim, device) - batch, max_blocks = read_page_table.shape - queries, topk = topk_indices.shape - if queries <= 0 or topk <= 0: - raise ValueError( - "Sparse read validation requires nonempty fixed-capacity query/top-k dimensions" - ) - if request_rows.shape != (queries,) or compressed_lengths.shape != (batch,): - raise ValueError("request_rows must be [Q] and compressed_lengths must be [B]") - if active_request_count.shape != (1,) or out.shape != (1,): - raise ValueError("active_request_count and out must be [1]") - if compressed_tokens_per_page <= 0 or not 0 < page_index_upper_bound <= 0x7FFFFFFF: - raise ValueError("Tokens per page and physical page bound must be positive") - with torch.cuda.device(device): - out.fill_(1) - _check_sparse_read_table_kernel[(queries,)]( - topk_indices, - request_rows, - active_request_count, - compressed_lengths, - read_page_table, - out, - *topk_indices.stride(), - request_rows.stride(0), - compressed_lengths.stride(0), - *read_page_table.stride(), - NUM_REQUESTS=batch, - NUM_QUERIES=queries, - TOPK=topk, - MAX_BLOCKS=max_blocks, - TOKENS_PER_PAGE=compressed_tokens_per_page, - PAGE_BOUND=page_index_upper_bound, - BLOCK_K=triton.next_power_of_2(topk), - ) - - -@triton.jit -def _prepare_sparse_write_table_kernel( - raw_ptr, - history_ptr, - active_count_ptr, - out_ptr, - raw_stride0, - raw_stride1, - history_stride, - out_stride0, - out_stride1, - MAX_BLOCKS: tl.constexpr, - PAGE_SCALE: tl.constexpr, - BLOCK: tl.constexpr, -): - request = tl.program_id(0) - ordinal = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK) - active = request < tl.load(active_count_ptr) - history = tl.load(history_ptr + request * history_stride) - resident = active & (ordinal >= history) & (ordinal < MAX_BLOCKS) - raw = tl.load(raw_ptr + request * raw_stride0 + ordinal * raw_stride1, resident, other=-1) - page = raw.to(tl.int64) * PAGE_SCALE - valid = resident & (raw >= 0) & (page <= 0x7FFFFFFF) - tl.store( - out_ptr + request * out_stride0 + ordinal * out_stride1, - tl.where(valid, page, -1), - ordinal < MAX_BLOCKS, - ) - - -def prepare_sparse_write_table( - raw_page_table: torch.Tensor, - history_blocks: torch.Tensor, - active_request_count: torch.Tensor, - out: torch.Tensor, - *, - page_scale: int, -) -> None: - """Convert only resident raw slots to SHARED pages, preserving invalids. - - All tensors are CUDA int32 on one device: tables are [B, M], history - counts [B], and the active count [1]. ``out`` must be disjoint from inputs. - Host history and padded requests become -1, including host slot zero. - The caller supplies the resident converter scale (expansion must be 1); - the layer offset is already part of the compressed buffer pointer. - No tensor allocation or host synchronization is performed. - """ - device = raw_page_table.device - for tensor, name, ndim in ( - (raw_page_table, "raw_page_table", 2), - (history_blocks, "history_blocks", 1), - (active_request_count, "active_request_count", 1), - (out, "out", 2), - ): - _check_sparse_offload_tensor(tensor, name, ndim, device) - batch, max_blocks = raw_page_table.shape - if out.shape != raw_page_table.shape: - raise ValueError("raw_page_table and out must have the same [B, M] shape") - if history_blocks.shape != (batch,) or active_request_count.shape != (1,): - raise ValueError("history_blocks must be [B] and active_request_count must be [1]") - if not 0 < page_scale <= 0x7FFFFFFF: - raise ValueError("page_scale must be a positive int32 scale") - if batch == 0 or max_blocks == 0: - return - with torch.cuda.device(device): - _prepare_sparse_write_table_kernel[(batch, triton.cdiv(max_blocks, 256))]( - raw_page_table, - history_blocks, - active_request_count, - out, - *raw_page_table.stride(), - history_blocks.stride(0), - *out.stride(), - MAX_BLOCKS=max_blocks, - PAGE_SCALE=page_scale, - BLOCK=256, - ) diff --git a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/metadata.py b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/metadata.py index 4ed9a04e968b..414962a4b0de 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/metadata.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/metadata.py @@ -14,8 +14,12 @@ from tensorrt_llm._utils import prefer_pinned from ..dsa.metadata import DSAtrtllmAttentionMetadata +from ..offload import ( + SparseOffloadState, + allocate_sparse_offload_state, + get_sparse_offload_metadata_layers, +) from .indexer import DeepseekV4Indexer -from .offload import SparseOffloadState from .params import ( DEEPSEEK_V4_SLIDING_ATTENTION, DEEPSEEK_V4_SPARSE_RATIO, @@ -369,61 +373,25 @@ def _init_sparse_offload_state(self) -> None: if not self.sparse_metadata_params.enable_kv_cache_offload: return manager = self.kv_cache_manager - # Per-batch beam/draft values start at defaults. Also check the manager's - # configured limits before allocating any offload buffers. - if ( - self.beam_width != 1 - or manager.max_beam_width != 1 - or self.max_draft_tokens - or manager.max_draft_len - or manager.max_total_draft_tokens - or self.draft_kv_cache_manager is not None - ): - raise NotImplementedError( - "Sparse offload requires single-beam, non-speculative execution" - ) - # Compression policies are bound after manager construction, before - # metadata initialization. Validate them once here, not on each batch. - manager._validate_sparse_history_policy() - # Descriptor setup also validates the fixed KV dtype and pool layout. - layers = manager.get_sparse_offload_descriptors() + # Descriptors include every sparse role in native fetch order. + layers = get_sparse_offload_metadata_layers(self) if not layers: # This PP stage has no ratio-4 main KV to stage. return - if self.sparse_mla_topk <= 0: - raise ValueError("Sparse offload metadata requires positive top-k") - batch = self.max_num_sequences - max_blocks = self.kv_cache_manager.max_blocks_per_seq - - def allocate(name: str, shape: tuple[int, ...], value: int) -> torch.Tensor: - tensor = self.get_empty( + # Native fetch always processes every Batch row, including for smaller + # graph profiles. Its selection width is fixed by BufferConfig. + self.sparse_offload_state = allocate_sparse_offload_state( + lambda name, shape, dtype: self.get_empty( self.cuda_graph_buffers, shape, - dtype=torch.int32, + dtype=dtype, cache_name=f"dsv4_sparse_offload_{name}", capture_graph=self.is_cuda_graph, - ) - tensor.fill_(value) - return tensor - - history_blocks = allocate("history_blocks", (batch,), 0) - self.sparse_offload_state = SparseOffloadState( - layers=layers, - base_page_tables={ - group: allocate(f"base_page_table_{group}", (batch, max_blocks), -1) - for group in sorted({descriptor.group_id for descriptor in layers.values()}) - }, - history_blocks=history_blocks, - history_blocks_host=torch.zeros( - batch, dtype=torch.int32, device="cpu", pin_memory=prefer_pinned() - ), - active_request_count=allocate("active_request_count", (1,), 0), - selected_history_pages=allocate( - "selected_history_pages", (batch, min(self.sparse_mla_topk, max_blocks)), -1 ), - fetched_page_table=allocate("fetched_page_table", (batch, max_blocks), -1), - compress_read_table=allocate("compress_read_table", (batch, max_blocks), -1), - read_table_valid=allocate("read_table_valid", (1,), 1), + layers, + max_batch_size=manager.max_batch_size, + topk=self.sparse_mla_topk, + nvfp4_head_dim=manager.head_dim if manager._use_nvfp4_compress else None, ) def mla_prepare_fused_q_gen_cu_seqlens(self) -> Optional[torch.Tensor]: @@ -500,13 +468,13 @@ def prepare_for_block_tables(self): compress_ratio == DEEPSEEK_V4_SPARSE_RATIO and self.sparse_metadata_params.enable_kv_cache_offload ): - if self.sparse_offload_state is None: - compress_block_table.fill_(-1) - else: + # Host page indices cannot provide GPU write destinations. + # Sparse writes require a separate per-layer arena mapping. + compress_block_table.fill_(-1) + if self.sparse_offload_state is not None: self.kv_cache_manager.prepare_sparse_offload( self.sparse_offload_state, self.request_ids, - compress_block_table, beam_width=self.beam_width, ) continue @@ -632,7 +600,14 @@ def _build_cache_buffer_data_pointers( swa_buffer_ptrs = {layer_idx: manager.swa_pool_ptr for layer_idx in manager.pp_layers} compressed_buffer_ptrs = { - layer_idx: manager.get_buffers(layer_idx, DeepseekV4AttentionType.COMPRESS).data_ptr() + layer_idx: ( + manager._ensure_sparse_batch().get_sparse_cache_base_address( + manager._layer_attn_to_layer_id[layer_idx, DeepseekV4AttentionType.COMPRESS], + DeepseekV4AttentionType.COMPRESS.role, + ) + if manager._enable_kv_cache_offload and compress_ratios_by_layer[layer_idx] == 4 + else manager.get_buffers(layer_idx, DeepseekV4AttentionType.COMPRESS).data_ptr() + ) for layer_idx in manager.pp_layers if is_compress_layer(compress_ratios_by_layer[layer_idx]) } @@ -758,6 +733,8 @@ def prepare(self) -> None: assert self.kv_cache_manager is not None assert self.request_ids is not None self.validate_sparse_offload_batch() + if self.sparse_offload_state is not None: + self.kv_cache_manager.prepare_sparse_batch(self.request_ids) self.kv_cache_manager.compute_sliding_block_tables( self.request_ids, diff --git a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/offload.py b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/offload.py deleted file mode 100644 index ee2157dbd6cc..000000000000 --- a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/offload.py +++ /dev/null @@ -1,88 +0,0 @@ -# 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. - -"""Attention-owned descriptors and persistent sparse offload metadata.""" - -from __future__ import annotations - -import os -from dataclasses import dataclass, field -from typing import TYPE_CHECKING - -import torch - -if TYPE_CHECKING: - from tensorrt_llm.runtime.kv_cache_manager_v2 import BufferId, LayerGroupId - - -@dataclass(frozen=True, slots=True) -class SparseOffloadLayerDescriptor: - """One local model layer's main compressed buffer in KVCM's namespace. - - ``page_scale`` converts resident raw GPU slots to SHARED physical pages. - ``fetched_page_scale`` must be bound separately once KVCM specifies fetch - output units: 1 for SHARED physical pages, or the appropriate converter - scale for raw slots in the same pool. None keeps decode disabled rather - than guessing the unfinished API's units. Buffer pointers supply the - layer offset in either case. - """ - - buffer_id: BufferId - group_id: LayerGroupId - page_scale: int - # Exclusive physical-page bound relative to this layer's SHARED pointer, - # as returned by KVCM. Its runtime implementation must include scratch. - page_index_upper_bound: int - fetched_page_scale: int | None = None - - -@dataclass(kw_only=True) -class SparseOffloadState: - """Fixed-capacity storage owned by one serial attention metadata instance. - - All device tensors are int32. Tables are [B, M], selections are [B, S], - history counts are [B], and the active count is [1]. History staging is - pinned CPU memory. Layer workspaces are reused only after the previous - layer's attention has consumed them on the same stream. Concurrent - forwards require separate metadata and KVCM scratch ownership. - - ``prepared`` means preparation was enqueued successfully; stream ordering - establishes device readiness. ``is_prefill`` is set by host batch validation - before execution, never inferred from a device count inside a layer. - - ``read_table_valid`` is a private int32 [1] diagnostic flag, refreshed by - GPU validation after each decode fetch. Only the device assertion requires - ``TLLM_DSV4_OFFLOAD_DEBUG_ASSERT=1``, read at initialization before graph - capture. A failed assertion makes the CUDA context unusable. By default, - the flag does not trigger an error or request recovery; production trusts - KVCM to provide valid mappings for every required page. - """ - - layers: dict[int, SparseOffloadLayerDescriptor] - base_page_tables: dict[LayerGroupId, torch.Tensor] - history_blocks: torch.Tensor - history_blocks_host: torch.Tensor - active_request_count: torch.Tensor - selected_history_pages: torch.Tensor - fetched_page_table: torch.Tensor - compress_read_table: torch.Tensor - read_table_valid: torch.Tensor - debug_assert: bool = field( - default_factory=lambda: os.environ.get("TLLM_DSV4_OFFLOAD_DEBUG_ASSERT", "0") == "1" - ) - history_upload_done: torch.cuda.Event = field(default_factory=torch.cuda.Event) - history_upload_pending: bool = False - prepared: bool = False - is_prefill: bool = False diff --git a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/params.py b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/params.py index baa60f4f5154..7a1b7e25138e 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/params.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/params.py @@ -90,7 +90,6 @@ class DeepSeekV4Params(DSAParams): algorithm: Literal["deepseek_v4"] = field(init=False, default="deepseek_v4") compress_ratios: List[int] = field(default_factory=list) window_size: int = 128 - enable_kv_cache_offload: bool = False @dataclass(frozen=True) @@ -99,4 +98,3 @@ class DeepSeekV4MetadataParams(DSAMetadataParams): compress_ratios: List[int] = field(default_factory=list) window_size: int = 128 - enable_kv_cache_offload: bool = False diff --git a/tensorrt_llm/_torch/attention/backends/sparse/dsa/backend.py b/tensorrt_llm/_torch/attention/backends/sparse/dsa/backend.py index 1c14f3421a51..62f21e1ed0e4 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/dsa/backend.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/dsa/backend.py @@ -18,9 +18,11 @@ MLAParams, PositionalEmbeddingParams, ) -from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention +from tensorrt_llm._torch.attention.backends.trtllm import MlaLatentWriteTarget, TrtllmAttention +from tensorrt_llm._torch.modules.multi_stream_utils import do_multi_stream from tensorrt_llm.models.modeling_utils import QuantConfig +from ..offload_kernels import gather_sparse_batch_rows, merge_sparse_arena_indices from .indexer import ( Indexer, transform_local_topk_and_prepare_pool_view, @@ -182,6 +184,14 @@ def sparse_attn_predict( : topk_indices.shape[1], ].copy_(topk_indices) + if self.sparse_params.enable_kv_cache_offload: + if not is_generation: + raise NotImplementedError( + "DSA KV cache offload runs decode batches only; context requests reach " + "an offloading worker through disaggregated KV transfer" + ) + return self._sparse_offload_arena_indices(topk_indices, metadata, forward_args), None + local_layer_idx = self.get_local_layer_idx(metadata) kv_cache_dtype = getattr(metadata.kv_cache_manager, "dtype", None) if kv_cache_dtype == tensorrt_llm.bindings.DataType.NVFP4 and not is_generation: @@ -266,6 +276,182 @@ def sparse_attn_predict( return topk_indices_global, None + def _sparse_offload_arena_indices( + self, + topk_indices: torch.Tensor, + metadata: DSAtrtllmAttentionMetadata, + forward_args: AttentionForwardArgs, + ) -> torch.Tensor: + """Fetch this layer's selected history and return its latent arena rows. + + The selection holds token ordinals in query order, one query per request. + History tokens come from the fetch, the token written this step from its + write row, and attention reads both through the layer's arena base. An + NVFP4 latent is first dequantized to FP8 scratch rows, and attention reads + those instead, as it does from the GPU-resident NVFP4 cache. + """ + state = metadata.sparse_offload_state + if state is None or state.write_layer != self.layer_idx: + raise RuntimeError("Write this layer's latent row before sparse attention") + if metadata.num_contexts: + raise NotImplementedError("DSA KV cache offload runs decode batches only") + num_tokens = topk_indices.shape[0] + topk = metadata.sparse_mla_topk + topk_indices = topk_indices[:, :topk] + request_indices = metadata.req_idx_per_token[:num_tokens] + kv_lengths = metadata.kv_lens_cuda[: metadata.num_seqs] + descriptor = state.layers[self.layer_idx] + if getattr(metadata, "sparse_offload_fetch_pending", None) == self.layer_idx: + metadata.sparse_offload_fetch_event.wait() + fetched = state.fetched_indices[len(descriptor.buffer_ids)] + metadata.sparse_offload_fetch_pending = None + else: + fetched = metadata.kv_cache_manager.fetch_sparse_indices( + state, self.layer_idx, topk_indices, request_indices, kv_lengths + ) + data_index = descriptor.data_buffer_index + # Both paths resolve buffers whose arena indices coincide once, into one slice, so + # the data buffer is read at its slot rather than at its own position. + data_slot = state.out_slots[self.layer_idx][data_index] + indices = merge_sparse_arena_indices( + topk_indices, + request_indices, + state.batch_rows, + state.active_request_count, + kv_lengths, + state.history_lengths, + fetched[:, data_slot, :], + state.write_indices[len(descriptor.buffer_ids)][:, data_slot], + metadata.sparse_offload_attn_indices[:num_tokens], + unit_tokens=descriptor.unit_tokens, + ) + if state.nvfp4_scratch is None: + forward_args.sparse_runtime_params.aux_kv_cache_pool_ptr = ( + descriptor.arena_base_addresses[data_index] + ) + return indices + # The data and scale arenas share their rows, so one index reads both. The + # gather replaces each arena row with its FP8 scratch row, -1 staying -1. + data_rows, scale_rows = metadata.kv_cache_manager.get_sparse_latent_nvfp4_arenas( + self.layer_idx + ) + scratch = state.nvfp4_scratch[:num_tokens] + torch.ops.trtllm.nvfp4_mla_kv_cache_gather_direct( + data_rows, + scale_rows, + indices, + scratch, + self._nvfp4_mla_kv_scale_quant_orig, + metadata.kv_cache_manager.mla_kv_cache_residual_dim, + data_rows.shape[0], + ) + forward_args.sparse_runtime_params.aux_kv_cache_pool_ptr = scratch.data_ptr() + return indices + + def prepare_sparse_write(self, metadata: DSAtrtllmAttentionMetadata) -> None: + """Overlap this layer's write mapping with its independent projections.""" + if not self.sparse_params.enable_kv_cache_offload or not do_multi_stream(): + return + stream = metadata.sparse_offload_stream + if stream is None: + return + pending = metadata.sparse_offload_write_pending + if pending == self.layer_idx: + return + if pending is not None: + raise RuntimeError("The preceding sparse write mapping has not been consumed") + metadata.sparse_offload_start_event.record() + with torch.cuda.stream(stream): + metadata.sparse_offload_start_event.wait() + self._prepare_sparse_write_mapping(metadata) + metadata.sparse_offload_write_event.record() + metadata.sparse_offload_write_pending = self.layer_idx + + def prefetch_next_sparse_history(self, metadata: DSAtrtllmAttentionMetadata) -> None: + """Fetch the next shared-indexer layer during the current layer's MLP. + + Call after this layer's attention consumes the shared offload scratch. + A shared layer already has its top-k; its reservation protects the new + token's row while the history fetch runs independently of its writer. + The writer joins reservation, and attention joins the fetch. + """ + if not self.sparse_params.enable_kv_cache_offload or not do_multi_stream(): + return + layer_idx = metadata.sparse_offload_next_shared_layer.get(self.layer_idx) + if layer_idx is None or metadata.sparse_offload_stream is None: + return + if ( + metadata.sparse_offload_write_pending is not None + or metadata.sparse_offload_fetch_pending is not None + ): + raise RuntimeError("The preceding sparse offload work has not been consumed") + metadata.sparse_offload_start_event.record() + with torch.cuda.stream(metadata.sparse_offload_stream): + metadata.sparse_offload_start_event.wait() + self._prepare_sparse_write_mapping(metadata, layer_idx) + metadata.sparse_offload_write_event.record() + metadata.kv_cache_manager.fetch_sparse_indices( + metadata.sparse_offload_state, + layer_idx, + metadata.shared_topk_indices[: metadata.num_tokens, : metadata.sparse_mla_topk], + metadata.req_idx_per_token[: metadata.num_tokens], + metadata.kv_lens_cuda[: metadata.num_seqs], + ) + metadata.sparse_offload_fetch_event.record() + metadata.sparse_offload_write_pending = layer_idx + metadata.sparse_offload_fetch_pending = layer_idx + + def _prepare_sparse_write_mapping( + self, metadata: DSAtrtllmAttentionMetadata, layer_idx: Optional[int] = None + ) -> None: + """Publish the reserved arena rows and the writer's one-token page table.""" + if layer_idx is None: + layer_idx = self.layer_idx + state = metadata.sparse_offload_state + manager = metadata.kv_cache_manager + write_indices = manager.prepare_sparse_write_indices(state, layer_idx) + descriptor = state.layers[layer_idx] + # At the slot: the write resolves coinciding buffers once, as the fetch does. + data_slot = state.out_slots[layer_idx][descriptor.data_buffer_index] + write_rows = gather_sparse_batch_rows( + write_indices[:, data_slot], + state.batch_rows, + out=metadata.sparse_latent_block_offsets[0, :, 0, 0], + ) + metadata.sparse_latent_write_slots.copy_((write_rows >= 0).to(torch.int32) - 1) + + def mla_rope_generation( + self, + fused_q: Optional[torch.Tensor], + q_pe: Optional[torch.Tensor], + latent_cache: torch.Tensor, + metadata: DSAtrtllmAttentionMetadata, + *args, + **kwargs, + ) -> None: + """Store the new latent row in the layer's arena when the latent is offloaded.""" + state = getattr(metadata, "sparse_offload_state", None) + if not self.sparse_params.enable_kv_cache_offload or state is None: + return super().mla_rope_generation( + fused_q, q_pe, latent_cache, metadata, *args, **kwargs + ) + manager = metadata.kv_cache_manager + if getattr(metadata, "sparse_offload_write_pending", None) == self.layer_idx: + metadata.sparse_offload_write_event.wait() + metadata.sparse_offload_write_pending = None + else: + self._prepare_sparse_write_mapping(metadata) + num_tokens = latent_cache.shape[0] + write_slots = metadata.sparse_latent_write_slots[:num_tokens] + pool_pointers, pool_mapping = manager.get_sparse_latent_write_pool(self.layer_idx) + kwargs["latent_write_target"] = MlaLatentWriteTarget( + block_offsets=metadata.sparse_latent_block_offsets, + pool_pointers=pool_pointers, + pool_mapping=pool_mapping, + write_slots=write_slots, + ) + return super().mla_rope_generation(fused_q, q_pe, latent_cache, metadata, *args, **kwargs) + def _remap_topk_to_global( self, topk_indices: torch.Tensor, diff --git a/tensorrt_llm/_torch/attention/backends/sparse/dsa/cache_manager.py b/tensorrt_llm/_torch/attention/backends/sparse/dsa/cache_manager.py index dc653921e1eb..23f012eb8cc2 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/dsa/cache_manager.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/dsa/cache_manager.py @@ -6,6 +6,8 @@ import math import os +from collections import Counter +from dataclasses import replace from typing import TYPE_CHECKING, List, Optional, Tuple, Union import torch @@ -13,17 +15,37 @@ import tensorrt_llm import tensorrt_llm.bindings from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2, Role +from tensorrt_llm._torch.pyexecutor.kv_cache.sparse_offload_memory import sparse_cache_config from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager -from tensorrt_llm._utils import TensorWrapper, convert_to_torch_tensor, get_size_in_bytes +from tensorrt_llm._utils import ( + TensorWrapper, + convert_to_torch_tensor, + get_size_in_bytes, + get_sm_version, +) from tensorrt_llm.bindings import DataType from tensorrt_llm.bindings.executor import KvCacheConfig from tensorrt_llm.bindings.internal.batch_manager import CacheType as CacheTypeCpp from tensorrt_llm.logger import logger from tensorrt_llm.mapping import Mapping -from tensorrt_llm.runtime.kv_cache_manager_v2 import BufferConfig, DataRole, PageIndexMode - +from tensorrt_llm.runtime.kv_cache_manager_v2 import ( + AttentionLayerConfig, + BufferConfig, + DataRole, + PageIndexMode, +) +from tensorrt_llm.runtime.kv_cache_manager_v2 import KVCacheManagerConfig as KVCacheManagerConfigPy + +from ..offload import ( + SPARSE_OFFLOAD_STATE_COPIES, + SparseOffloadCacheMixin, + SparseOffloadLayerDescriptor, +) from .params import DSAParams +# A DSA selection names single tokens of the MLA latent. +DSA_SPARSE_UNIT_TOKENS = 1 + ModelConfig = tensorrt_llm.bindings.ModelConfig if TYPE_CHECKING: @@ -95,6 +117,44 @@ def _get_indexer_k_cache_size_per_token( return num_attention_layers * bytes_per_layer +def _get_full_indexer_cache_size_per_token( + model_config: ModelConfig, + mapping: Mapping, + num_layers: Optional[int] = None, +) -> int: + """Estimate the indexer K-cache cost of the local full-indexer layers. + + Shared-indexer layers own no indexer row. With ``num_layers`` every one of + those layers is costed as a full-indexer layer. + """ + sparse_attention_config = model_config.sparse_attention_config + config = model_config.pretrained_config + sparse_params = sparse_attention_config.to_sparse_params(pretrained_config=config) + if not isinstance(sparse_params, DSAParams): + raise ValueError("DSA cache requires DSA sparse parameters") + if num_layers is not None: + num_indexer_layers = max(num_layers, 1) + else: + local_layer_ids = mapping.pp_layers(model_config.get_num_attention_layers()) + num_indexer_layers = sum( + 1 + for layer_id in local_layer_ids + if getattr( + sparse_attention_config.to_sparse_params( + pretrained_config=config, layer_idx=layer_id + ), + "is_full_indexer_layer", + True, + ) + ) + # Indexer K cache: physically allocated as raw UINT8, so it is NOT scaled by + # the KV cache dtype (unlike the latent). Under FP4 the indexer stores two + # E2M1 codes per byte; the scale bytes are unchanged. + return num_indexer_layers * _get_indexer_k_cache_bytes_per_token( + sparse_params.index_head_dim, 128, sparse_params.indexer_k_dtype == "fp4" + ) + + def derive_indexer_k_cache_layer_mask( sparse_attention_config: "SparseAttentionConfig", pretrained_config, @@ -153,6 +213,10 @@ def __init__( ) if not isinstance(sparse_params, DSAParams): raise ValueError("DSA cache requires DSA sparse parameters") + if sparse_attention_config.get_kv_cache_offload_config() is not None: + raise ValueError( + "DSA KV cache offload requires kv_cache_config.use_kv_cache_manager_v2=True." + ) self.quant_block_size = 128 self.index_head_dim = sparse_params.index_head_dim # V1 pool geometry is fixed by the legacy C++ manager. Residual @@ -279,14 +343,6 @@ def get_cache_size_per_token( ) if not isinstance(sparse_params, DSAParams): raise ValueError("DSA cache requires DSA sparse parameters") - index_head_dim = sparse_params.index_head_dim - quant_block_size = 128 - # Under FP4 the indexer stores two E2M1 codes per byte, so the - # per-token data footprint halves (132 B -> 68 B at index_head_dim=128); - # the scale bytes are unchanged (4 per token, one int32 holding four - # UE8M0 exponents at quant_block_size=32 after packing). - use_fp4 = sparse_params.indexer_k_dtype == "fp4" - indexer_data_dim = index_head_dim // 2 if use_fp4 else index_head_dim # Get latent KV-cache bytes. NVFP4 stores two values per data byte # plus one E4M3 block scale for every 16 values. @@ -313,32 +369,9 @@ def get_cache_size_per_token( ) mem_per_token = num_attention_layers * head_dim * bytes_per_element - if num_layers is not None: - num_indexer_layers = max(num_layers, 1) - else: - local_layer_ids = mapping.pp_layers(model_config.get_num_attention_layers()) - num_indexer_layers = sum( - 1 - for layer_id in local_layer_ids - if getattr( - sparse_attention_config.to_sparse_params( - pretrained_config=config, layer_idx=layer_id - ), - "is_full_indexer_layer", - True, - ) - ) - - # Indexer K cache: physically allocated as raw UINT8 in - # WindowBlockManager::allocatePools (poolDtype = kUINT8), so we assume - # 1 byte/element here -- it is NOT scaled by the KV cache dtype (unlike - # the latent above). The data-portion byte count already reflects fp8 vs - # fp4 via indexer_data_dim. - indexer_bytes_per_token = num_indexer_layers * ( - indexer_data_dim + index_head_dim // quant_block_size * 4 + return mem_per_token + _get_full_indexer_cache_size_per_token( + model_config, mapping, num_layers ) - mem_per_token += indexer_bytes_per_token - return mem_per_token def get_cache_bytes_per_token(self): """Compute actual cache bytes per token from instance configuration.""" @@ -387,8 +420,15 @@ def get_cache_bytes_per_token(self): return cache_size_bytes_per_token -class DSACacheManagerV2(KVCacheManagerV2): - """KVCacheManagerV2-backed cache manager with a DSA indexer K-cache.""" +class DSACacheManagerV2(SparseOffloadCacheMixin, KVCacheManagerV2): + """KVCacheManagerV2-backed cache manager with a DSA indexer K-cache. + + With KV cache offload, every layer's MLA latent (``KEY``, with its + ``KEY_BLOCK_SCALE`` under NVFP4) keeps its history in host memory and is + read through a per-request GPU cache of single-token units; the indexer + K-cache (``INDEX_KEY``) stays on the GPU. Full-indexer and shared-indexer + layers differ only in where their selection comes from. + """ def __init__( self, @@ -453,6 +493,12 @@ def __init__( self.use_fp8_ds_mla, head_dim = _resolve_fp8_ds_mla_head_dim( kv_cache_config, tokens_per_block, head_dim, dtype ) + sparse_kv_offload_config = sparse_attention_config.get_kv_cache_offload_config() + self._init_sparse_offload( + enabled=sparse_kv_offload_config is not None, topk=sparse_params.index_topk + ) + if sparse_kv_offload_config is not None: + self._check_dsa_offload_support(dtype) from tensorrt_llm._torch.speculative import get_num_spec_layers @@ -462,6 +508,7 @@ def __init__( self._indexer_k_cache_layer_mask = derive_indexer_k_cache_layer_mask( sparse_attention_config, pretrained_config, total_num_layers ) + self._selection_leader_cache: dict[int, int | None] | None = None super().__init__( kv_cache_config, @@ -479,10 +526,25 @@ def __init__( max_num_tokens=max_num_tokens, model_config=model_config, max_beam_width=max_beam_width, + sparse_kv_offload_config=sparse_kv_offload_config, **kwargs, ) if self.num_local_layers == 0: raise ValueError("DSA requires at least one local attention layer") + self._sparse_offload_descriptors: dict[int, SparseOffloadLayerDescriptor] = {} + self._sparse_latent_write_pools: dict[int, Tuple[torch.Tensor, torch.Tensor]] = {} + self._sparse_latent_nvfp4_arenas: dict[int, Tuple[torch.Tensor, torch.Tensor]] = {} + if self._enable_kv_cache_offload: + self._sparse_offload_descriptors = self.get_sparse_offload_descriptors() + self._initialize_sparse_arena_padding(self._sparse_offload_descriptors) + self._sparse_latent_write_pools = { + layer: self._latent_arena_pool_tensors(descriptor) + for layer, descriptor in self._sparse_offload_descriptors.items() + } + if self.dtype == DataType.NVFP4: + self._sparse_latent_nvfp4_arenas = self._latent_nvfp4_arena_views( + self._sparse_offload_descriptors + ) primary_converters = [ self.impl.get_page_index_converter(local_layer_idx, Role.KEY) for local_layer_idx in range(self.num_local_layers) @@ -605,6 +667,11 @@ def get_unique_primary_pool(self) -> torch.Tensor: """Return the uniform MLA K pool in the V1-compatible layout.""" if self._unique_primary_pool is not None: return self._unique_primary_pool + if self._enable_kv_cache_offload: + raise RuntimeError( + "With KV cache offload the MLA latent's pages are host memory; attention " + "reads it through the sparse cache arena" + ) if self.num_local_layers == 0: raise ValueError("DSA requires at least one local attention layer") if self.kv_factor != 1: @@ -671,6 +738,11 @@ def get_layer_bytes_per_token(self, local_layer_idx: int, data_role: DataRole) - elif data_role == Role.ALL: cache_bytes += self.mla_kv_cache_residual_dim // 2 cache_bytes += self.mla_kv_cache_residual_dim // 16 + if data_role == Role.ALL and self._sparse_kv_offload is not None: + # Role.ALL is what the device estimators count; the latent's pages + # are host memory. + for role in self._sparse_latent_roles(): + cache_bytes -= self.get_layer_bytes_per_token(local_layer_idx, role) if ( data_role == Role.ALL and self._indexer_k_cache_layer_mask[self.pp_layers[local_layer_idx]] @@ -678,6 +750,219 @@ def get_layer_bytes_per_token(self, local_layer_idx: int, data_role: DataRole) - cache_bytes += self.get_layer_bytes_per_token(local_layer_idx, Role.INDEX_KEY) return cache_bytes + def _check_dsa_offload_support(self, dtype: DataType) -> None: + """Refuse offload on the paths that have no arena read or write yet.""" + if self.use_fp8_ds_mla: + raise NotImplementedError( + "DSA KV cache offload does not support kv_cache_config.dtype='fp8_ds_mla'." + ) + if dtype not in (DataType.FP8, DataType.BF16, DataType.NVFP4): + raise NotImplementedError( + f"DSA KV cache offload supports an FP8, BF16 or NVFP4 MLA latent, got {dtype}." + ) + if get_sm_version() < 100: + raise NotImplementedError( + "DSA KV cache offload reads the latent through the trtllm-gen sparse MLA " + f"kernels, which need SM100 or newer; this GPU is SM{get_sm_version()}." + ) + if self._sparse_offload_topk is None or self._sparse_offload_topk <= 0: + raise ValueError("DSA KV cache offload requires a positive index_topk") + + def _sparse_latent_roles(self) -> Tuple[DataRole, ...]: + """The latent buffers offload keeps in host memory: data, and NVFP4 block scales.""" + if self.dtype == DataType.NVFP4: + return (Role.KEY, Role.KEY_BLOCK_SCALE) + return (Role.KEY,) + + def _selection_leaders(self) -> dict[int, int | None]: + """Per model layer, the model layer whose indexer selects for it. + + A group is a run of local layers starting at one that owns an indexer; + the layers after it reuse its top-k, so the whole run is handed one + selection and shares one resident set in the cache. ``None`` where the + run holds a single layer, which is every layer of a model with no + indexer sharing -- a layer that selects for itself never shares. + + Keyed and valued by model layer, because the local numbering is compacted + after the cache config is built and this answer has to be the same on both + sides of that. Computed once for the same reason. + """ + if self._selection_leader_cache is None: + leader_of: dict[int, int | None] = {} + leader = self.pp_layers[0] + for local, layer in enumerate(self.pp_layers): + # The first layer leads whether or not the mask says it owns an + # indexer: a run has to start somewhere, and + # derive_indexer_k_cache_layer_mask already forces layer 0 to be full. + if self._indexer_k_cache_layer_mask[layer] or local == 0: + leader = layer + leader_of[layer] = leader + members = Counter(leader_of.values()) + self._selection_leader_cache = { + layer: leader if members[leader] > 1 else None + for layer, leader in leader_of.items() + } + return self._selection_leader_cache + + def _build_cache_config(self, config: KVCacheManagerConfigPy) -> KVCacheManagerConfigPy: + """With offload, give every layer's MLA latent host-backed history and a GPU cache.""" + if self._sparse_kv_offload is None: + return config + # Named by the leading local layer, so the layers of one indexer resolve once + # between them. An unshared layer names nothing and keeps a row to itself. + leaders = self._selection_leaders() + + def selection_id(local_layer_idx: int) -> int | None: + leader = leaders[self.pp_layers[local_layer_idx]] + # Negated so the id cannot land in the runtime's layer-id space, where an + # unset id reads as the layer's own and a leading model layer's number may + # well be some other layer's local one. The id only has to be equal within + # a group and distinct between them; it is not a layer reference. + return None if leader is None else -(leader + 1) + + def cache_for(local_layer_idx: int): + return sparse_cache_config( + num_selected_units=self._sparse_offload_topk, + unit_tokens=DSA_SPARSE_UNIT_TOKENS, + tokens_per_block=self.tokens_per_block, + gpu_cache_factor=self._sparse_kv_offload.gpu_cache_factor, + selection_id=selection_id(local_layer_idx), + ) + + sparse_roles = self._sparse_latent_roles() + layers = [] + for layer in config.layers: + if not isinstance(layer, AttentionLayerConfig): + layers.append(layer) + continue + # A sparse buffer's size is still a whole page; the layer's sparse_cache + # divides it into one-token units of the single-head latent. + buffers = [ + BufferConfig(role=buffer.role, size=buffer.size, is_sparse=True) + if buffer.role in sparse_roles + else buffer + for buffer in layer.buffers + ] + layers.append( + AttentionLayerConfig( + layer_id=layer.layer_id, + buffers=buffers, + sliding_window_size=layer.sliding_window_size, + num_sink_tokens=layer.num_sink_tokens, + sparse_cache=cache_for(int(layer.layer_id)), + ) + ) + return replace(config, layers=layers) + + def get_sparse_offload_descriptors(self) -> dict[int, SparseOffloadLayerDescriptor]: + """Describe each local layer's MLA latent arena, by global layer index.""" + if not self._enable_kv_cache_offload: + raise ValueError("Sparse offload metadata requires an offload-enabled cache manager") + layout = self._sparse_arena_layout() + leaders = self._selection_leaders() + return { + layer: self._build_sparse_offload_descriptor( + layout, + self.layer_offsets[layer], + Role.KEY, + unit_tokens=DSA_SPARSE_UNIT_TOKENS, + scale_role=Role.KEY_BLOCK_SCALE if self.dtype == DataType.NVFP4 else None, + selection_owner=leaders[layer], + name=f"Layer {layer}", + ) + for layer in self.pp_layers + } + + def _latent_arena_pool_tensors( + self, descriptor: SparseOffloadLayerDescriptor + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Return a one-pool view of a layer's latent arena for the MLA cache writer. + + The writer addresses its cache through host pool pointers [pool, 2] and a + layer-to-pool mapping [layer, 2]. Pool 0 starts at the arena base, so a + one-token page whose block offset is the write row lands in that row. + NVFP4 pointers are [pool, 2, 2], data then block scales; the scales' + arena pairs its rows with the data's, so the same row addresses both. + """ + data_base = descriptor.arena_base_addresses[descriptor.data_buffer_index] + if descriptor.scale_buffer_index is None: + pointers = [[data_base, 0]] + else: + scale_base = descriptor.arena_base_addresses[descriptor.scale_buffer_index] + pointers = [[[data_base, scale_base], [0, 0]]] + pool_pointers = torch.tensor(pointers, dtype=torch.int64, device="cpu") + pool_mapping = torch.zeros((1, 2), dtype=torch.int32, device="cpu") + return pool_pointers, pool_mapping + + def get_sparse_latent_write_pool(self, layer_idx: int) -> Tuple[torch.Tensor, torch.Tensor]: + """Return the host pool pointers and mapping that address a layer's latent arena.""" + return self._sparse_latent_write_pools[layer_idx] + + def _latent_nvfp4_arena_views( + self, descriptors: dict[int, SparseOffloadLayerDescriptor] + ) -> dict[int, Tuple[torch.Tensor, torch.Tensor]]: + """View each layer's NVFP4 latent arenas as the rows the NVFP4 gather reads. + + Data views are packed bytes [rows, 1, data bytes] and scale views E4M3 + [rows, 1, scale bytes], where a row is one token's unit and an arena row + index addresses it. All layers' buffers of one pool share its region, + and so its view. + """ + views: dict[int, torch.Tensor] = {} + + def view( + descriptor: SparseOffloadLayerDescriptor, + index: int, + dtype: torch.dtype, + row_bytes: int, + ) -> torch.Tensor: + base = descriptor.arena_base_addresses[index] + if base not in views: + rows = descriptor.arena_row_counts[index] + views[base] = convert_to_torch_tensor( + TensorWrapper(base, dtype, (rows, 1, row_bytes)) + ) + return views[base] + + arenas = {} + for layer, descriptor in descriptors.items(): + local_layer_idx = self.layer_offsets[layer] + arenas[layer] = ( + view( + descriptor, + descriptor.data_buffer_index, + torch.uint8, + self.get_layer_bytes_per_token(local_layer_idx, Role.KEY), + ), + view( + descriptor, + descriptor.scale_buffer_index, + torch.float8_e4m3fn, + self.get_layer_bytes_per_token(local_layer_idx, Role.KEY_BLOCK_SCALE), + ), + ) + return arenas + + def get_sparse_latent_nvfp4_arenas(self, layer_idx: int) -> Tuple[torch.Tensor, torch.Tensor]: + """Return a layer's NVFP4 latent data and block-scale arena views, indexed by arena row.""" + return self._sparse_latent_nvfp4_arenas[layer_idx] + + def _sparse_offload_nvfp4_head_dim(self) -> Optional[int]: + return self.head_dim if self.dtype == DataType.NVFP4 else None + + def get_sparse_offload_scratch_bytes(self, config: KVCacheManagerConfigPy) -> int: + """Count the offload state and the latent writer's one-token page table. + + Per Batch row, the page table holds a K and a V offset and the write slots + one entry, since an offloading step has one query token per request. + Both metadata copies hold them. + """ + scratch = super().get_sparse_offload_scratch_bytes(config) + if scratch: + writer_entries = 3 * self.max_batch_size + scratch += SPARSE_OFFLOAD_STATE_COPIES * writer_entries * torch.int32.itemsize + return scratch + def get_cache_bytes_per_token(self) -> int: return sum( self.get_layer_bytes_per_token(local_layer_idx, Role.ALL) @@ -691,6 +976,14 @@ def get_cache_size_per_token( num_layers: Optional[int] = None, **kwargs, ): + sparse_attention_config = model_config.sparse_attention_config + if ( + sparse_attention_config is not None + and sparse_attention_config.get_kv_cache_offload_config() is not None + ): + # The latent's history lives in host memory; only the indexer K-cache + # takes device bytes per token. + return _get_full_indexer_cache_size_per_token(model_config, mapping, num_layers) cache_bytes = DSACacheManager.get_cache_size_per_token( model_config, mapping, num_layers=num_layers, **kwargs ) diff --git a/tensorrt_llm/_torch/attention/backends/sparse/dsa/metadata.py b/tensorrt_llm/_torch/attention/backends/sparse/dsa/metadata.py index e0773630dcb3..01c9912833f5 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/dsa/metadata.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/dsa/metadata.py @@ -20,6 +20,11 @@ from tensorrt_llm.deep_gemm import get_paged_mqa_logits_metadata from tensorrt_llm.logger import logger +from ..offload import ( + SparseOffloadState, + allocate_sparse_offload_state, + get_sparse_offload_metadata_layers, +) from ..params import use_self_sampling_gvr from .cache_manager import is_dsa_cache_manager from .indexer import ( @@ -376,9 +381,100 @@ def __post_init__(self): self.create_buffers_for_mla_rope_append(capture_graph=capture_graph) self.create_buffers_for_indexer(capture_graph=capture_graph) self._create_nvfp4_mla_generation_buffers(capture_graph=capture_graph) + self._init_dsa_sparse_offload_state() + + def _offloads_dsa_latent(self) -> bool: + """Whether attention reads a DSA cache manager's MLA latent through the offload arena.""" + return self.sparse_metadata_params.enable_kv_cache_offload and is_dsa_cache_manager( + self.kv_cache_manager + ) + + def _init_dsa_sparse_offload_state(self) -> None: + """Validate the fixed offload configuration and allocate its persistent state. + + Besides the shared offload state, the MLA latent writer gets a one-token + page table [1, B, 2, 1] and per-token write slots [B]: the writer stores + each request's new latent row at its arena write row through them. An + NVFP4 latent is dequantized for attention into the state's FP8 scratch. + """ + self.sparse_offload_state: Optional[SparseOffloadState] = None + self.sparse_latent_block_offsets: Optional[torch.Tensor] = None + self.sparse_latent_write_slots: Optional[torch.Tensor] = None + self.sparse_offload_attn_indices: Optional[torch.Tensor] = None + self.sparse_offload_stream: Optional[torch.cuda.Stream] = None + self.sparse_offload_write_pending: Optional[int] = None + self.sparse_offload_fetch_pending: Optional[int] = None + self.sparse_offload_next_shared_layer: dict[int, int] = {} + if not self._offloads_dsa_latent(): + return + self.sparse_offload_stream = torch.cuda.Stream() + self.sparse_offload_start_event = torch.cuda.Event() + self.sparse_offload_write_event = torch.cuda.Event() + self.sparse_offload_fetch_event = torch.cuda.Event() + manager = self.kv_cache_manager + layers = get_sparse_offload_metadata_layers(self) + local_layers = manager.pp_layers + indexer_mask = manager.indexer_k_cache_local_layer_mask + self.sparse_offload_next_shared_layer = { + previous: following + for previous, following, has_indexer in zip( + local_layers, local_layers[1:], indexer_mask[1:] + ) + if not has_indexer and previous in layers and following in layers + } + batch = manager.max_batch_size + + def allocate(name: str, shape: tuple[int, ...], dtype: torch.dtype) -> torch.Tensor: + return self.get_empty( + self.cuda_graph_buffers, + shape, + dtype=dtype, + cache_name=f"dsa_sparse_offload_{name}", + capture_graph=self.is_cuda_graph, + ) + + # Native fetch always processes every Batch row, including for smaller + # graph profiles. Its selection width is fixed by the layer's sparse_cache. + self.sparse_offload_state = allocate_sparse_offload_state( + allocate, + layers, + max_batch_size=batch, + topk=self.sparse_mla_topk, + nvfp4_head_dim=manager._sparse_offload_nvfp4_head_dim(), + ) + self.sparse_latent_block_offsets = allocate( + "latent_block_offsets", (1, batch, 2, 1), torch.int32 + ) + self.sparse_latent_block_offsets.fill_(-1) + self.sparse_latent_write_slots = allocate("latent_write_slots", (batch,), torch.int32) + self.sparse_latent_write_slots.fill_(-1) + self.sparse_offload_attn_indices = allocate( + "attn_indices", (batch, self.sparse_mla_topk), torch.int32 + ) + + def _validate_dsa_sparse_offload_batch(self) -> None: + """Reject batches offload cannot run, on the host before any model writes. + + An offloading worker only decodes, one query token per request. + """ + state = self.sparse_offload_state + state.prepared = False + state.is_prefill = False + if self.num_contexts: + raise NotImplementedError( + "DSA KV cache offload runs decode batches only; context requests reach " + "an offloading worker through disaggregated KV transfer" + ) + if self.seq_lens.device.type != "cpu" or self.seq_lens_kv.device.type != "cpu": + raise ValueError("Sparse offload batch validation requires host sequence lengths") + if not torch.equal(self.seq_lens, self.seq_lens_kv) or not torch.all(self.seq_lens == 1): + raise NotImplementedError("Sparse offload decode requires one query per request") def prepare(self): self._validate_ragged_verification() + if self.sparse_offload_state is not None: + self._validate_dsa_sparse_offload_batch() + self.kv_cache_manager.prepare_sparse_batch(self.request_ids) super().prepare() self._invalidate_pool_view_cache() @@ -419,6 +515,11 @@ def prepare(self): # Prepare metadata for indexer Indexer.prepare(metadata=self) + if self.sparse_offload_state is not None: + self.kv_cache_manager.prepare_sparse_offload( + self.sparse_offload_state, self.request_ids, beam_width=self.beam_width + ) + def _prepare_device_window_kv_lens(self, kv_lens: torch.Tensor) -> torch.Tensor: """Bound host-side KV extents until device-selected windows are installed.""" if self.device_windows_mode and self.is_ragged_verify: @@ -925,6 +1026,9 @@ def _create_nvfp4_mla_generation_buffers(self, capture_graph=False): self.nvfp4_mla_compact_indices = None if getattr(self.kv_cache_manager, "dtype", None) != tensorrt_llm.bindings.DataType.NVFP4: return + if self._offloads_dsa_latent(): + # Offloaded rows are dequantized into the offload state's scratch instead. + return # MTP verifies the current token and max_draft_tokens draft tokens for # every generation request. Each query row has its own TopK selection, diff --git a/tensorrt_llm/_torch/attention/backends/sparse/dsa/module.py b/tensorrt_llm/_torch/attention/backends/sparse/dsa/module.py index f1ca16154b73..5514bbf2c41f 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/dsa/module.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/dsa/module.py @@ -177,6 +177,7 @@ def forward_sparse_attn( standard output at index zero. """ output = attn_output[0] + self.mqa.prepare_sparse_write(attn_metadata) proj_outputs = forward_dsa_proj(self, position_ids, hidden_states, attn_metadata) q, compressed_kv, k_pe, latent_cache = proj_outputs[:4] indexer_intermediates = proj_outputs[4:] @@ -204,6 +205,11 @@ def forward_sparse_attn_custom_op( del latent_cache_gen from . import custom_ops # noqa: F401 + # Metadata-dependent cache work stays outside the projection custom op. + if not torch.compiler.is_compiling(): + attn_metadata, _ = custom_ops._extract_extra_attrs(self.layer_idx_str) + self.mqa.prepare_sparse_write(attn_metadata) + if isinstance(hidden_states, Fp4QuantizedTensor): proj_outputs = torch.ops.trtllm.mla_dsa_proj( hidden_states.unquantized_hidden_states, @@ -325,6 +331,7 @@ def _forward_dsa_attn( num_tokens comes from batch-specific metadata and must not be baked into CUDA graph capture. """ + self.mqa.prepare_sparse_write(attn_metadata) num_contexts = attn_metadata.num_contexts num_generations = attn_metadata.num_generations num_ctx_tokens = attn_metadata.num_ctx_tokens @@ -411,6 +418,7 @@ def _forward_dsa_attn( # capture requires the join within the same layer's forward. if self.mqa.indexer is not None: self.mqa.indexer.maybe_join_prev_topk_copy() + self.mqa.prefetch_next_sparse_history(attn_metadata) def should_use_short_mha( diff --git a/tensorrt_llm/_torch/attention/backends/sparse/dsa/params.py b/tensorrt_llm/_torch/attention/backends/sparse/dsa/params.py index 377fe91e4def..bf43da160bd4 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/dsa/params.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/dsa/params.py @@ -44,6 +44,7 @@ class DSAMetadataParams(SparseMetadataParams): use_self_sampling_topk: bool = True use_gvr_emission: bool = False gvr_deterministic_topk: bool = False + enable_kv_cache_offload: bool = False @dataclass(frozen=True) @@ -75,6 +76,7 @@ class DSAParams(SparseParams): # Shared layers reuse the preceding full layer's top-k. is_full_indexer_layer: bool = True mtp_index_share: bool = False + enable_kv_cache_offload: bool = False @property def indices_block_size(self) -> int: diff --git a/tensorrt_llm/_torch/attention/backends/sparse/offload.py b/tensorrt_llm/_torch/attention/backends/sparse/offload.py new file mode 100644 index 000000000000..ed15576c24d3 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/sparse/offload.py @@ -0,0 +1,1101 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Sparse KV cache offload shared by sparse attention models. + +With offload, the history a sparse attention layer selects from lives in host +memory, and each decode step fetches the selected units into per-request rows +of a GPU arena. This module holds what every offloading model shares: + +- the per-layer descriptors and the attention metadata's persistent state; +- ``SparseOffloadCacheMixin``, the KV cache manager side of a decode step: + publishing the native Batch, preparing write rows, fetching selections and + advancing history after the forward. + +A model's cache manager mixes ``SparseOffloadCacheMixin`` in before +``KVCacheManagerV2``, marks its offloaded buffers sparse, and describes each +offloaded layer with ``_build_sparse_offload_descriptor``. +""" + +from __future__ import annotations + +from collections import defaultdict +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Callable, Collection, Iterable + +import torch + +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import CUDA_GRAPH_DUMMY_REQUEST_ID +from tensorrt_llm._utils import ( + TensorWrapper, + convert_to_torch_tensor, + nvtx_range_debug, + prefer_pinned, +) +from tensorrt_llm.runtime.kv_cache_manager_v2 import ( + BAD_PAGE_INDEX, + AttentionLayerConfig, + Batch, + BatchRow, + BufferId, + CudaStream, + DataRole, + LayerGroupId, + LayerId, + _KVCache, +) + +from .offload_kernels import select_sparse_history_units + +if TYPE_CHECKING: + from tensorrt_llm._torch.pyexecutor import llm_request + from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests + from tensorrt_llm.runtime.kv_cache_manager_v2 import KVCacheManagerConfig + +# Attention metadata holding a SparseOffloadState: the eager one, and the CUDA +# graph metadata, which share theirs through the graph buffer pool. +SPARSE_OFFLOAD_STATE_COPIES = 2 + +# The sparse MLA FMHA reads the NVFP4 scratch as its aux KV pool and can read up to +# one KV tile (128 rows) past the last row its indices name (https://nvbugs/6796296), +# so the scratch carries a zeroed tail of that many rows. +NVFP4_SCRATCH_TAIL_ROWS = 128 + + +@dataclass(frozen=True, slots=True) +class SparseOffloadPaddingRegion: + """Unused buffer bytes across arena leases, initialized during manager construction.""" + + base_address: int + bytes_per_lease: int + lease_stride: int + + +@dataclass(frozen=True, slots=True) +class SparseOffloadLayerDescriptor: + """All sparse buffers of one model layer, in KVCM's required fetch order. + + ``unit_tokens`` original tokens make one selection unit. Fetch returns a + separate physical index for each buffer, relative to that buffer's + ``get_sparse_cache_base_address()``. These indices already include the + request and buffer placement; no page scaling or layer offset is needed. + Descriptor construction validates matching data/scale pool counts and + buffer positions, including padding. Their index columns are equal, so a + writer and a gather consume the data column for both pools while retaining + their respective base addresses and byte strides. + ``arena_row_counts`` covers each complete pool region across all leases + and buffers, including padding; it is independent of host-page capacity. + + ``selection_owner`` is the model layer that resolves this layer's + selection, where a model runs one indexer across several layers. ``None`` + means the layer selects for itself, which is the ordinary case. A layer + naming another is relying on being handed that layer's selection verbatim; + nothing checks that, and a different selection would silently resolve + against the owner's residency. + """ + + buffer_ids: tuple[BufferId, ...] + arena_base_addresses: tuple[int, ...] + arena_row_counts: tuple[int, ...] + data_buffer_index: int + scale_buffer_index: int | None + group_id: LayerGroupId + unit_tokens: int + padding_regions: tuple[SparseOffloadPaddingRegion, ...] = () + selection_owner: int | None = None + + +@dataclass(kw_only=True) +class SparseOffloadState: + """Fixed-capacity storage owned by one serial attention metadata instance. + + Index and count tensors are int32. Selections are [B, K], fetch outputs are + [B, F, K], write destinations are [B, F], the active count is [1] and the + Batch row map is [B]. B is the native Batch capacity, including when an + attention graph uses fewer query rows. Storage is shared by layers with the + same buffer count F. Reuse must follow the previous layer's attention on the + consuming stream. + + NVFP4 additionally uses persistent FP8 scratch [B, K, head_dim] for the + dequantized rows consumed by attention. Arena allocations remain owned by + KVCM; the cache manager retains non-owning tensor views sized by arena + geometry. ``scatter_start_positions`` is initialized to zero once and + supplies the one-row-page view a writer can address the arena through. + + ``history_lengths`` is a zero-copy view of Batch's published token counts, + bound outside capture. Batch owns the page tables and watermarks; metadata + does not maintain another copy. A request keeps its Batch row while it stays + scheduled, so rows need not follow the input order or be contiguous. + ``batch_rows`` maps each request's input position to its row, and positions + past the live requests to -1; it is refreshed outside capture each step. + Selections, fetch output, write destinations and watermarks are indexed by + Batch row; per-request inputs and attention outputs by input position. + + ``prepared`` means preparation was enqueued successfully; stream ordering + establishes device readiness. ``is_prefill`` is set by host batch validation + before execution, never inferred from a device count inside a layer. + ``write_layer`` tracks which layer enqueued the shared write mapping during + eager execution or capture. Graph replay uses the captured GPU operations; + this host marker does not determine device readiness. + + An in-flight selection is masked out of the history fetch. Its physical + index must come from the separately prepared write mapping when assembling + attention inputs. + """ + + layers: dict[int, SparseOffloadLayerDescriptor] + active_request_count: torch.Tensor + batch_rows: torch.Tensor + selected_history_units: torch.Tensor + fetched_indices: dict[int, torch.Tensor] + write_indices: dict[int, torch.Tensor] + scatter_start_positions: torch.Tensor + #: Which slice of ``fetched_indices`` and ``write_indices`` each of a layer's buffers + #: reads. Buffers whose arena indices coincide share one, and the slices no buffer + #: maps to are never written, so a buffer must be read at its slot rather than at its + #: own position. Filled on first use, the earliest the native Batch exists to ask. + out_slots: dict[int, tuple[int, ...]] = field(default_factory=dict) + nvfp4_scratch: torch.Tensor | None = None + batch: Batch | None = None + history_lengths: torch.Tensor | None = None + prepared: bool = False + is_prefill: bool = False + write_layer: int | None = None + + @staticmethod + def device_bytes( + *, + max_batch_size: int, + topk: int, + buffer_counts: Collection[int], + nvfp4_head_dim: int | None, + ) -> int: + """Return the device bytes one instance allocates. + + Mirrors the allocations of ``allocate_sparse_offload_state``. + + Args: + max_batch_size: Rows B, the native Batch capacity. + topk: Units one row selects, K. + buffer_counts: The distinct sparse buffer counts F of the layers. + nvfp4_head_dim: Head dim of the dequantized rows with NVFP4, else None. + """ + # The active count, the selections, the scatter start positions and the + # Batch row map, then the fetch outputs and write destinations of each + # buffer count. + index_entries = 1 + max_batch_size * topk + 2 * max_batch_size + index_entries += sum(max_batch_size * count * (topk + 1) for count in buffer_counts) + nvfp4_bytes = 0 + if nvfp4_head_dim is not None: + nvfp4_rows = max_batch_size * topk + NVFP4_SCRATCH_TAIL_ROWS + nvfp4_bytes = nvfp4_rows * nvfp4_head_dim * torch.float8_e4m3fn.itemsize + return index_entries * torch.int32.itemsize + nvfp4_bytes + + +def allocate_sparse_offload_state( + allocate: Callable[[str, tuple[int, ...], torch.dtype], torch.Tensor], + layers: dict[int, SparseOffloadLayerDescriptor], + *, + max_batch_size: int, + topk: int, + nvfp4_head_dim: int | None = None, +) -> SparseOffloadState: + """Allocate an attention metadata's persistent offload state. + + Args: + allocate: Returns an uninitialized tensor of a name, shape and dtype; + the metadata's graph-buffer allocator. + layers: The manager's descriptor of each offloaded local layer. + max_batch_size: Rows B, the native Batch capacity. + topk: Units one row selects. + nvfp4_head_dim: Head dim of the NVFP4 gather's dequantized rows, else None. + """ + + def index_tensor(name: str, shape: tuple[int, ...], value: int) -> torch.Tensor: + tensor = allocate(name, shape, torch.int32) + tensor.fill_(value) + return tensor + + def nvfp4_scratch(head_dim: int) -> torch.Tensor: + rows = max_batch_size * topk + flat = allocate( + "nvfp4_scratch", (rows + NVFP4_SCRATCH_TAIL_ROWS, head_dim), torch.float8_e4m3fn + ) + flat[rows:].zero_() + return flat[:rows].view(max_batch_size, topk, head_dim) + + buffer_counts = sorted({len(descriptor.buffer_ids) for descriptor in layers.values()}) + return SparseOffloadState( + layers=layers, + active_request_count=index_tensor("active_request_count", (1,), 0), + batch_rows=index_tensor("batch_rows", (max_batch_size,), -1), + selected_history_units=index_tensor( + "selected_history_units", (max_batch_size, topk), BAD_PAGE_INDEX + ), + fetched_indices={ + count: index_tensor( + f"fetched_indices_{count}", (max_batch_size, count, topk), BAD_PAGE_INDEX + ) + for count in buffer_counts + }, + write_indices={ + count: index_tensor(f"write_indices_{count}", (max_batch_size, count), BAD_PAGE_INDEX) + for count in buffer_counts + }, + scatter_start_positions=index_tensor("scatter_start_positions", (max_batch_size,), 0), + nvfp4_scratch=nvfp4_scratch(nvfp4_head_dim) if nvfp4_head_dim is not None else None, + ) + + +def get_sparse_offload_metadata_layers(metadata) -> dict[int, SparseOffloadLayerDescriptor]: + """Validate an attention metadata's fixed offload configuration and return its layers. + + Offload runs single-beam and non-speculative, so per-batch beam and draft + values start at defaults; the manager's configured limits are checked too, + before any offload buffer is allocated. History policies are bound after + manager construction, before metadata initialization, so they are checked + once here rather than on each batch. An empty result means this pipeline + stage has no offloaded layer. + """ + manager = metadata.kv_cache_manager + if ( + metadata.beam_width != 1 + or manager.max_beam_width != 1 + or metadata.max_draft_tokens + or manager.max_draft_len + or manager.max_total_draft_tokens + or metadata.draft_kv_cache_manager is not None + ): + raise NotImplementedError("Sparse offload requires single-beam, non-speculative execution") + manager._validate_sparse_history_policy() + layers = manager.get_sparse_offload_descriptors() + if layers and ( + metadata.sparse_mla_topk <= 0 or metadata.sparse_mla_topk != manager._sparse_offload_topk + ): + raise ValueError("Sparse offload metadata must match the configured positive top-k") + return layers + + +@dataclass(frozen=True, slots=True) +class _SparseArenaLayout: + """Where every sparse buffer sits in its pool's arena region.""" + + # Coalesced-buffer position of each buffer, which is its pool index in its group. + pool_indices: dict[BufferId, int] + layer_configs: dict[LayerId, AttentionLayerConfig] + # (buffers in the pool region, this buffer's position among them). + positions: dict[BufferId, tuple[int, int]] + max_units_by_group: dict[LayerGroupId, int] + max_leases: int + + +@dataclass(slots=True) +class _SparseBatchMember: + """A request seated in the native Batch.""" + + cache: _KVCache + # The row add_kv_cache returned. The runtime never moves a seated member, so the + # row holds while the request stays scheduled, and with it the GPU cache lease. + row: BatchRow + # Page-table columns the row may have written since it joined. A cache + # writes page indices only below its block count, so the host snapshot + # copies no further than this. + blocks: int + + +class SparseOffloadCacheMixin: + """KV cache manager side of sparse KV cache offload, model-agnostic. + + Mix in before ``KVCacheManagerV2``. Call ``_init_sparse_offload`` before the + base constructor; it reads the selection width while building the config. + + A request's sparse history advances right after each forward, since one + query token per request writes one unit. Between ``prepare_resources`` and + the post-forward update, the scheduled requests are members of a native + Batch, which owns their host page-index rows; every path that mutates a + request's cache releases the Batch first. + + Batch rows are slots. A request keeps its row, and with it the GPU cache + lease, for as long as it stays scheduled; a departing request frees its row + and an arriving one takes the lowest free row. Attention finds each request's + row through ``SparseOffloadState.batch_rows`` rather than by input position. + """ + + _enable_kv_cache_offload: bool = False + _sparse_offload_topk: int | None = None + _sparse_batch: Batch | None = None + _sparse_batch_prepared: bool = False + + def _init_sparse_offload(self, *, enabled: bool, topk: int | None) -> None: + self._enable_kv_cache_offload = enabled + self._sparse_offload_topk = topk + self._sparse_batch = None + self._sparse_batch_members: dict[int, _SparseBatchMember] = {} + self._sparse_batch_prepared = False + + def _sparse_offload_nvfp4_head_dim(self) -> int | None: + """Head dim of the dequantized rows an NVFP4 gather keeps per selection. + + None when no layer dequantizes into offload scratch. + """ + return None + + def get_sparse_offload_scratch_bytes(self, config: KVCacheManagerConfig) -> int: + """Return the device bytes of the attention metadata's sparse offload state. + + Counts a copy for the CUDA graph metadata even when graphs are off. + """ + buffer_counts = { + sum(buffer.is_sparse for buffer in layer.buffers) + for layer in config.layers + if getattr(layer, "sparse_cache", None) is not None + } + if not buffer_counts: + return 0 + return SPARSE_OFFLOAD_STATE_COPIES * SparseOffloadState.device_bytes( + max_batch_size=self.max_batch_size, + topk=self._sparse_offload_topk, + buffer_counts=buffer_counts, + nvfp4_head_dim=self._sparse_offload_nvfp4_head_dim(), + ) + + def _has_sparse_offload_layers(self) -> bool: + """Whether a local layer fetches through a sparse cache.""" + return self._enable_kv_cache_offload and any( + getattr(layer, "sparse_cache", None) is not None + and any(buffer.is_sparse for buffer in layer.buffers) + for layer in self.impl.init_config.layers + ) + + # ------------------------------------------------------------------ + # Batch lifecycle + # ------------------------------------------------------------------ + + def prepare_resources(self, scheduled_batch: ScheduledRequests) -> None: + self._sparse_batch_prepared = False + if self._sparse_batch is not None: + self._sparse_batch.wait_published() + if self._enable_kv_cache_offload and any( + not request.is_disagg_generation_init_state + for request in scheduled_batch.context_requests + ): + # An offloading worker admits one query token per request per step, so it serves + # the decode phase and runs no context batch. Disaggregated gen-init requests are + # the exception only because they run no forward: they are resized and their KV + # transferred before they join the Batch. Rejected here rather than in the + # compressor, which sees this as a bare context count once the model is already + # writing. + raise NotImplementedError( + "Sparse offload does not support prefill; it serves generation only" + ) + super().prepare_resources(scheduled_batch) + if not self._enable_kv_cache_offload: + return + requests = scheduled_batch.context_requests + scheduled_batch.generation_requests + # Disaggregated gen-init requests are prepared before their KV transfer + # starts and run no forward this step, so they never join the Batch. A + # batch holding only those publishes nothing. + forward_request_ids = [ + request.py_request_id + for request in requests + if not request.is_disagg_generation_init_state + ] + if forward_request_ids or not requests: + self.prepare_sparse_batch(forward_request_ids) + + def _ensure_sparse_batch(self) -> Batch: + """The Batch the sparse caches are seated in, created on first use. + + It owns the cell arenas, so its rows are what they are sized for and its + base addresses are what a fetch's indices are relative to. Both are fixed + once it exists, which is why a single Batch serves the manager's lifetime. + """ + if self._sparse_batch is None: + self._sparse_batch = Batch( + self.impl, + CudaStream(self._stream.cuda_stream), + # The fetch and the write-back read converted indices, so every sparse + # buffer needs a table. Their shape follows from the pool layout. + index_tables=self.impl.sparse_index_table_specs(), + max_batch_size=self.max_batch_size, + max_blocks_per_seq=self.max_blocks_per_seq, + ) + return self._sparse_batch + + def _sparse_arena_base(self, layer_id: int, role: DataRole) -> int: + """Where this buffer's arena region starts, for a consumer of fetched units. + + The arenas belong to the Batch that seats the caches, so the address is + asked of it. + """ + return self._ensure_sparse_batch().get_sparse_cache_base_address(layer_id, role) + + def prepare_sparse_batch(self, request_ids: list[int]) -> None: + """Publish scheduler or warmup requests, or join an already published batch. + + Attention metadata calls this outside capture with its final request order. + Graph padding may append dummy requests; the live requests must be the + scheduled ones, in any order. Warmup requests can reach this without + prepare_resources(). + """ + if self._sparse_batch_prepared: + self.get_sparse_batch(request_ids) + torch.cuda.current_stream(self._stream.device).wait_stream(self._stream) + else: + self._prepare_sparse_batch(request_ids) + + def _prepare_sparse_batch(self, request_ids: list[int]) -> None: + """Publish the scheduled caches before a forward. + + Continuing requests keep their Batch row, membership and sparse GPU residency + across steps, whatever their input order; departing requests free their rows + and arriving ones take the lowest free rows. Members write their page indices + into Batch staging; the runtime's host tables, which dense consumers read, + receive a snapshot after scheduling. + """ + if not self._has_sparse_offload_layers(): + return + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("Sparse Batch publication must run outside CUDA graph capture") + if len(request_ids) > self.max_batch_size: + raise ValueError("Scheduled requests exceed sparse Batch capacity") + caches = {} + padding = False + for request_id in request_ids: + if request_id == CUDA_GRAPH_DUMMY_REQUEST_ID: + padding = True + continue + if padding: + raise ValueError("CUDA-graph dummy requests must form a trailing suffix") + if request_id in caches: + raise ValueError("Sparse Batch requires unique request IDs") + cache = self.kv_cache_map.get(request_id) + if cache is None or not cache.is_active: + raise ValueError("Sparse Batch requires active scheduled request caches") + if cache.beam_width != 1: + raise ValueError("Sparse Batch supports one beam per request") + if cache.cuda_stream != self._stream.cuda_stream: + raise RuntimeError("Sparse Batch requires the manager's cache stream") + if not 0 <= cache.history_length <= cache.capacity: + raise ValueError("Invalid KVCM history frontier") + if cache.num_blocks > self.max_blocks_per_seq: + raise ValueError("Request allocation exceeds sparse Batch page-table capacity") + caches[request_id] = cache + + batch = self._ensure_sparse_batch() + try: + batch.wait_published() + self._release_sparse_batch( + request_id for request_id in self._sparse_batch_members if request_id not in caches + ) + for request_id, cache in caches.items(): + member = self._sparse_batch_members.get(request_id) + if member is None: + self._sparse_batch_members[request_id] = _SparseBatchMember( + cache=cache, row=batch.add_kv_cache(cache), blocks=cache.num_blocks + ) + else: + member.blocks = max(member.blocks, cache.num_blocks) + self._refresh_sparse_page_index_bufs() + stream = torch.cuda.current_stream(self._stream.device) + # Reusing the device mirror follows the previous forward. Publish on + # the cache stream so metadata preparation can also join it there. + self._stream.wait_stream(stream) + batch.publish() + stream.wait_stream(self._stream) + self._sparse_batch_prepared = True + finally: + if not self._sparse_batch_prepared: + self._release_sparse_batch() + + def get_sparse_batch(self, request_ids: list[int]) -> Batch: + """Return the published Batch after checking the attention metadata's requests. + + The live requests must be the Batch's members, each once, in any order: each + keeps its own row, which ``prepare_sparse_offload`` maps from the input + position. CUDA graph padding is a trailing suffix of dummy requests, which map + to no row. Fetch inputs must use Batch.max_batch_size rows, even when a smaller + attention profile is executing. + """ + if self._sparse_batch is None or not self._sparse_batch_prepared: + raise RuntimeError("Sparse Batch must be prepared before model execution") + members = self._sparse_batch_members + live_ids = request_ids[: len(members)] + if ( + len(request_ids) > self._sparse_batch.max_batch_size + or len(live_ids) != len(members) + or set(live_ids) != members.keys() + or any(req != CUDA_GRAPH_DUMMY_REQUEST_ID for req in request_ids[len(members) :]) + ): + raise ValueError("Attention requests do not match the published sparse Batch") + return self._sparse_batch + + def _refresh_sparse_page_index_bufs(self) -> None: + """Snapshot Batch-owned page indices into the runtime's host tables for dense consumers. + + Copies each member's row up to the widest block count any member has had since + joining. Past that, both tables hold BAD_PAGE_INDEX for every member: a cache + writes only below its block count, and its binding filled the rest. + """ + if not self._sparse_batch_members: + return + assert self._sparse_batch is not None + members = sorted(self._sparse_batch_members.items(), key=lambda item: item[1].row) + width = max(member.blocks for _, member in members) + if width == 0: + return + slots = torch.tensor( + [ + self.index_mapper.get_index(request_id) * self.max_beam_width + for request_id, _ in members + ], + dtype=torch.long, + device="cpu", + ) + first, last = members[0][1].row, members[-1][1].row + # Members without free rows between them are one slice; otherwise gather them. + rows = ( + None + if last - first + 1 == len(members) + else torch.tensor([member.row for _, member in members], dtype=torch.long) + ) + for pool_id in range(self.num_pools): + table = torch.from_dlpack(self._sparse_batch.coalesced_page_table(pool_id))[:, :width] + self.host_kv_cache_block_offsets[pool_id, :, 0, :width].index_copy_( + 0, slots, table[first : last + 1] if rows is None else table.index_select(0, rows) + ) + + def _release_sparse_batch(self, request_ids: Iterable[int] | None = None) -> None: + """Detach departing requests while retaining other members' GPU cache leases.""" + if self._sparse_batch is None: + return + self._sparse_batch.wait_published() + request_ids = list(self._sparse_batch_members if request_ids is None else request_ids) + for request_id in request_ids: + member = self._sparse_batch_members.get(request_id) + if member is None: + continue + self._sparse_batch_prepared = False + self._sparse_batch.remove_kv_cache(member.cache) + del self._sparse_batch_members[request_id] + # Removal restores the cache's internal indices. Rebind the runtime's + # host table so subsequent resizes also reach the dense consumers. + self._set_page_index_bufs(request_id, member.cache) + + def suspend_request(self, req: llm_request.LlmRequest) -> None: + self._release_sparse_batch([req.py_request_id]) + self._order_sparse_history_update([req]) + super().suspend_request(req) + + def free_resources(self, request: llm_request.LlmRequest, pin_on_release: bool = False) -> None: + self._release_sparse_batch([request.py_request_id]) + self._order_sparse_history_update([request]) + super().free_resources(request, pin_on_release) + + def release_index_slot(self, request_id: int) -> None: + self._release_sparse_batch([request_id]) + super().release_index_slot(request_id) + + def shutdown(self) -> None: + self._release_sparse_batch() + self._sparse_batch = None + super().shutdown() + + def _validate_sparse_history_policy(self) -> None: + """Check fixed history policies during attention metadata initialization.""" + if ( + self.enable_block_reuse + or self.kv_compression_manages_history + or self._has_cp_helix + or self.is_draft + ): + raise NotImplementedError( + "Sparse offload requires block reuse, KV compression, CP Helix, and draft caches disabled" + ) + + def _order_sparse_history_update(self, requests: Iterable[llm_request.LlmRequest]) -> None: + """Order writeback after all model producers and attention readers. + + The executor orders its calling stream after model forward. + Join it onto the cache stream before flushing completed units or resizing + request caches. prepare_sparse_batch() supplies the reverse dependency + for the next forward. + """ + if not self._enable_kv_cache_offload: + return + active = False + for request in requests: + cache = self.kv_cache_map.get(request.py_request_id) + if cache is None or not cache.is_active: + continue + if cache.cuda_stream != self._stream.cuda_stream: + raise RuntimeError("Sparse history updates require the manager's cache stream") + active = True + if active: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("Sparse history updates must run outside CUDA graph capture") + self._stream.wait_stream(torch.cuda.current_stream(self._stream.device)) + + def update_context_resources(self, scheduled_batch: ScheduledRequests) -> None: + self._release_sparse_batch( + request.py_request_id for request in scheduled_batch.context_requests + ) + self._order_sparse_history_update(scheduled_batch.context_requests) + super().update_context_resources(scheduled_batch) + + def try_allocate_generation(self, req: llm_request.LlmRequest) -> bool: + """Reserve a decode request's next input token unless it is already reserved. + + With offload, history advances right after each forward, so at scheduling + time a cache holding capacity beyond its history has already reserved this + step. Attention-DP padding dummies are created that way. Growing such a + cache again would exceed the sparse cache's one-token bound between + capacity and history. No growth is recorded, so a skipped step reverts + nothing. + """ + if self._enable_kv_cache_offload: + if self._sparse_batch is not None: + self._sparse_batch.wait_published() + cache = self.kv_cache_map.get(req.py_request_id) + if cache is not None and cache.is_active and cache.capacity > cache.history_length: + return True + return super().try_allocate_generation(req) + + def revert_allocate_generation(self, req: llm_request.LlmRequest) -> None: + if self._sparse_batch is not None: + self._sparse_batch.wait_published() + super().revert_allocate_generation(req) + + def update_forward_resources( + self, scheduled_batch: ScheduledRequests, *, wait_for_input_copy: Callable[[], None] + ) -> None: + """Advance sparse decode history from executed inputs before the next schedule. + + Offload decode validates one input token per request, with speculation + disabled. That input's KV is produced by this forward; the sampled + output has no KV yet. Queue writeback after the forward on the cache + stream, independently of when the executor consumes sampling results. + """ + if not self._enable_kv_cache_offload or not scheduled_batch.generation_requests: + return + wait_for_input_copy() + self._order_sparse_history_update(scheduled_batch.generation_requests) + # Before any history moves: the write-back reads + # the converted indices the Batch published, so it has to run while the members are + # still seated. A row that is not advancing repeats its own history, which asks for + # nothing. + batch = self._sparse_batch + if batch is not None: + batch.wait_published() + histories = [0] * batch.max_batch_size + for member in self._sparse_batch_members.values(): + histories[member.row] = member.cache.history_length + for request in scheduled_batch.generation_requests: + member = self._sparse_batch_members.get(request.py_request_id) + if member is not None and member.cache.is_active: + histories[member.row] = member.cache.history_length + 1 + batch.flush_sparse_units(histories) + + # The flush watermark belongs to the attached Batch row. Keep membership + # until resize has checked it and advanced the request's history. + for request in scheduled_batch.generation_requests: + request_id = request.py_request_id + cache = self.kv_cache_map.get(request_id) + if cache is None or not cache.is_active: + continue + history_length = cache.history_length + 1 + if history_length > cache.capacity: + raise ValueError(f"Sparse decode request {request_id} has no reserved input token") + if not cache.resize(None, history_length): + raise ValueError( + f"Failed to advance sparse KV history for request {request_id} " + f"to {history_length} tokens after model forward" + ) + self._allocated_draft_lens.pop(request_id, None) + # This step's allocation is settled here rather than in update_resources. + inflight = self._inflight_draft_lens.get(request_id) + if inflight: + inflight.pop(0) + if not inflight: + del self._inflight_draft_lens[request_id] + self._sparse_batch_prepared = False + + def update_resources( + self, + scheduled_batch: ScheduledRequests, + attn_metadata=None, + kv_cache_dtype_byte_size: float | None = None, + ) -> None: + if self._enable_kv_cache_offload: + # Sparse decode has no speculative slots to rewind. History follows + # forward completion; delayed sampling updates must neither rewind it + # nor detach the Batch prepared for a later forward. + return + super().update_resources(scheduled_batch, attn_metadata, kv_cache_dtype_byte_size) + + # ------------------------------------------------------------------ + # Arena descriptors + # ------------------------------------------------------------------ + + def _sparse_arena_layout(self) -> _SparseArenaLayout: + """Re-derive KVCM's arena walk: configuration layer order, then declaration + order within each pool. Host coalesced-buffer order is not used here.""" + pool_indices = { + buffer_id: pool_index + for group in self.impl.pool_group_descs + for variant in group.slot_desc.variants + for pool_index, coalesced in enumerate(variant.coalesced_buffers) + for buffer_id in coalesced.buffer_ids + } + layer_configs = { + layer.layer_id: layer + for layer in self.impl.init_config.layers + if isinstance(layer, AttentionLayerConfig) + } + sparse_pool_buffers: dict[tuple[LayerGroupId, int], list[BufferId]] = defaultdict(list) + max_units_by_group: dict[LayerGroupId, int] = defaultdict(int) + for layer_config in layer_configs.values(): + cache = layer_config.sparse_cache + if cache is None: + continue + group_id = self.impl.get_layer_group_id(layer_config.layer_id) + for buffer in layer_config.buffers: + if buffer.is_sparse: + buffer_id = BufferId(layer_config.layer_id, buffer.role) + sparse_pool_buffers[group_id, pool_indices[buffer_id]].append(buffer_id) + max_units_by_group[group_id] = max( + max_units_by_group[group_id], cache.max_cached_units + ) + positions = { + buffer_id: (len(buffers), position) + for buffers in sparse_pool_buffers.values() + for position, buffer_id in enumerate(buffers) + } + # The Batch owns the arenas and sizes them for its rows, so the lease count is + # its capacity. A constraint describes a batch the manager must be able to hold, + # not the arena's extent. + max_leases = self.max_batch_size + return _SparseArenaLayout( + pool_indices=pool_indices, + layer_configs=layer_configs, + positions=positions, + max_units_by_group=dict(max_units_by_group), + max_leases=max_leases, + ) + + def _build_sparse_offload_descriptor( + self, + layout: _SparseArenaLayout, + layer_id: LayerId, + data_role: DataRole, + *, + unit_tokens: int, + scale_role: DataRole | None = None, + selection_owner: int | None = None, + name: str, + ) -> SparseOffloadLayerDescriptor: + """Describe one layer's complete sparse fetch buffer set. + + KVCM orders buffers by pool index, retaining declaration order within + each pool. Pool addresses and page-index scales describe host storage; + they do not transform the physical arena indices returned by fetch. + + Args: + layout: The manager's arena layout. + layer_id: The KVCM layer holding the data buffer. + data_role: The buffer attention reads, by arena row. + unit_tokens: Original tokens one selection unit covers. + scale_role: A block-scale buffer that must share the data's arena rows. + selection_owner: The model layer that resolves this layer's selection, + or None where the layer selects for itself. Must agree with the + layer's configured ``selection_id``, since that is what decides + whether the manager gives the layers one resident set. + name: The model layer, for messages. + """ + if not self.impl.is_sparse(layer_id, data_role): + raise ValueError(f"{name}'s {data_role} buffer must be sparse") + units_per_page = self.impl.get_sparse_cache_units_per_page(layer_id) + if units_per_page <= 0 or self.tokens_per_block != units_per_page * unit_tokens: + raise ValueError( + f"{name}: a page of {self.tokens_per_block} tokens must hold whole " + f"selection units of {unit_tokens} tokens" + ) + # A layer may hold a dense buffer beside the one it fetches through, so take + # only the sparse buffers: the dense ones are not fetched through the cache. + layer_config = layout.layer_configs[layer_id] + cache = layer_config.sparse_cache + if ( + cache is None + or cache.max_cached_units <= 0 + or cache.num_selected_blocks != self._sparse_offload_topk + ): + raise ValueError( + "A sparse layer must configure the runtime's selection width and cache capacity" + ) + if (cache.selection_id is None) != (selection_owner is None): + raise ValueError( + f"{name}: a layer shares a selection exactly when it names a selection_id; " + f"got selection_id={cache.selection_id} and owner={selection_owner}" + ) + buffers = [ + BufferId(layer_id, buffer.role) for buffer in layer_config.buffers if buffer.is_sparse + ] + data_id = BufferId(layer_id, data_role) + if data_id not in buffers: + raise ValueError(f"{name}'s sparse buffers must include {data_role}") + buffers.sort(key=layout.pool_indices.__getitem__) + scale_id = BufferId(layer_id, scale_role) if scale_role is not None else None + if scale_id is not None: + if scale_id not in buffers: + raise ValueError(f"{name} requires a sparse block-scale buffer") + data_layout = layout.positions[data_id] + scale_layout = layout.positions[scale_id] + if data_layout != scale_layout: + raise ValueError( + f"{name}: the sparse data and scale arenas must have matching buffer " + "counts and positions, including padding, to share physical indices; " + f"got {data_layout} and {scale_layout}." + ) + # A sparse buffer's size is a whole page; its arena region holds units. + unit_bytes = {buffer.role: buffer.size // units_per_page for buffer in layer_config.buffers} + group_id = self.impl.get_layer_group_id(layer_id) + max_units = layout.max_units_by_group[group_id] + reserved_units = (max_units + units_per_page - 1) // units_per_page * units_per_page + arena_bases = tuple( + self._sparse_arena_base(buffer.layer_id, buffer.role) for buffer in buffers + ) + padding_regions = [] + for buffer, base in zip(buffers, arena_bases): + if buffer == data_id or (scale_id is not None and buffer == scale_id): + continue + pool_count, position = layout.positions[buffer] + buffer_bytes = reserved_units * unit_bytes[buffer.role] + padding_regions.append( + SparseOffloadPaddingRegion( + base_address=base + position * buffer_bytes, + bytes_per_lease=buffer_bytes, + lease_stride=pool_count * buffer_bytes, + ) + ) + return SparseOffloadLayerDescriptor( + buffer_ids=tuple(buffers), + arena_base_addresses=arena_bases, + arena_row_counts=tuple( + layout.max_leases * layout.positions[buffer][0] * reserved_units + for buffer in buffers + ), + data_buffer_index=buffers.index(data_id), + scale_buffer_index=buffers.index(scale_id) if scale_id is not None else None, + group_id=group_id, + unit_tokens=unit_tokens, + padding_regions=tuple(padding_regions), + selection_owner=selection_owner, + ) + + def _initialize_sparse_arena_padding( + self, layers: dict[int, SparseOffloadLayerDescriptor] + ) -> None: + """Initialize unused roles once, before requests can acquire arena leases. + + The arena's per-pool lease strides follow its configuration, independently + of host page capacities. Views cover only padding buffers; real KV and + scale buffers are initialized by their producers. KVCM owns these allocations. + """ + regions = [region for layer in layers.values() for region in layer.padding_regions] + if not regions: + return + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("Sparse arena padding must be initialized before graph capture") + max_leases = self.max_batch_size + with torch.cuda.device(self._stream.device), torch.cuda.stream(self._stream): + for region in regions: + convert_to_torch_tensor( + TensorWrapper( + region.base_address, + torch.uint8, + (max_leases, region.bytes_per_lease), + strides=(region.lease_stride, 1), + ) + ).zero_() + # Construction completes before any producer or captured graph can use the arena. + self._stream.synchronize() + + # ------------------------------------------------------------------ + # Per-step offload operations + # ------------------------------------------------------------------ + + @nvtx_range_debug("prepare_sparse_offload") + def prepare_sparse_offload( + self, + state: SparseOffloadState, + request_ids: list[int], + *, + beam_width: int, + ) -> None: + """Bind the published Batch and refresh the per-step mapping outside capture. + + Scheduling publishes the page tables and watermarks by Batch row. This maps + each request in ``request_ids``, the attention metadata's input order, to its + row, and sets the active count. The native Batch and its device views persist + across steps and graph profiles; only trailing dummy requests may extend the + scheduled requests, and they map to no row. This prepares fetch metadata, not + write destinations. + """ + state.prepared = False + state.write_layer = None + if not self._enable_kv_cache_offload or beam_width != 1: + raise ValueError("Sparse offload preparation requires offload enabled and beam width 1") + with torch.cuda.device(state.selected_history_units.device): + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("Sparse offload metadata must be prepared outside graph capture") + batch = self.get_sparse_batch(request_ids) + if batch.max_batch_size != state.selected_history_units.shape[0]: + raise ValueError("Sparse fetch buffers must use the native Batch capacity") + if state.batch is not None and state.batch is not batch: + raise RuntimeError( + "Sparse Batch allocation must remain stable across graph replays" + ) + torch.cuda.current_stream().wait_stream(self._stream) + if state.batch is None: + state.history_lengths = torch.from_dlpack(batch.history_lengths()) + state.batch = batch + live = len(self._sparse_batch_members) + rows = [self._sparse_batch_members[request_id].row for request_id in request_ids[:live]] + rows += [-1] * (batch.max_batch_size - live) + # Pinned memory returns to the caching host allocator only after the copy runs. + state.batch_rows.copy_( + torch.tensor(rows, dtype=torch.int32, pin_memory=prefer_pinned()), + non_blocking=True, + ) + state.active_request_count.fill_(live) + state.selected_history_units.fill_(BAD_PAGE_INDEX) + for out in state.fetched_indices.values(): + out.fill_(BAD_PAGE_INDEX) + for out in state.write_indices.values(): + out.fill_(BAD_PAGE_INDEX) + state.prepared = True + + def prepare_sparse_write_indices( + self, state: SparseOffloadState, layer_idx: int + ) -> torch.Tensor: + """Reserve this layer's in-flight arena rows on the current stream. + + Return int32[B, F] in Batch row order, with columns assigned by + ``Batch.sparse_out_slots`` and cached in ``state.out_slots[layer_idx]``. + Buffers with matching arena indices share a column; unused columns are + untouched. A used column gets -1 when its row has no unit in flight. + A writer addressing requests by input position + reads request i's entries at Batch row ``state.batch_rows[i]``, for example + through ``gather_sparse_batch_rows``. The caller must join Batch publication + before this call. Reservation must finish before either the writer or + the history fetch; they may run concurrently because fetch excludes the + protected in-flight row. Both must finish before attention or a history + update. No KV is fetched here. + + Data and scale buffers use the same column when their arena indices coincide; + the output-slot mapping follows the arena geometry and requires no device read. + """ + if ( + not self._enable_kv_cache_offload + or not state.prepared + or state.is_prefill + or state.batch is None + ): + raise RuntimeError("Sparse writes require prepared decode metadata") + descriptor = state.layers[layer_idx] + out = state.write_indices[len(descriptor.buffer_ids)] + state.write_layer = None + with torch.cuda.device(out.device): + if layer_idx not in state.out_slots: + state.out_slots[layer_idx] = tuple( + state.batch.sparse_out_slots(list(descriptor.buffer_ids)) + ) + state.batch.prepare_sparse_write_indices( + buffers=descriptor.buffer_ids, + out=out, + stream=CudaStream(torch.cuda.current_stream().cuda_stream), + ) + state.write_layer = layer_idx + return out + + @nvtx_range_debug("fetch_sparse_indices") + def fetch_sparse_indices( + self, + state: SparseOffloadState, + layer_idx: int, + topk_indices: torch.Tensor, + request_indices: torch.Tensor, + unit_lengths: torch.Tensor, + ) -> torch.Tensor: + """Fetch one layer's history into compact per-buffer arena indices. + + Inputs are CUDA int32: top-k unit ordinals [Q, K], each query's request + by input position [Q], and selectable unit counts [R] by input position. + Live queries name distinct requests, one query per request; the prepared + ``state.batch_rows`` maps each to its Batch row. Q and R may be smaller + than Batch.max_batch_size. + + Return int32[B, F, K] in Batch row order and descriptor.buffer_ids order. + Each index is relative to that buffer's sparse arena base. Top-k slots + keep their order; padding and in-flight selections produce -1. Attention + must supply the in-flight write index separately, without treating this + output as an ordinal-indexed page table or rescaling its indices. + + The caller joins selection producers and write reservation before this + call, and joins the fetch before consuming its output on another stream. + Output storage is reused by subsequent layers after attention consumes it. + """ + if ( + not self._enable_kv_cache_offload + or not state.prepared + or state.is_prefill + or state.batch is None + or state.history_lengths is None + ): + raise RuntimeError("Sparse fetch requires prepared decode metadata") + descriptor = state.layers[layer_idx] + with torch.cuda.device(state.selected_history_units.device): + select_sparse_history_units( + topk_indices, + request_indices, + state.batch_rows, + state.active_request_count, + unit_lengths, + state.history_lengths, + state.selected_history_units, + unit_tokens=descriptor.unit_tokens, + ) + out = state.fetched_indices[len(descriptor.buffer_ids)] + if layer_idx not in state.out_slots: + state.out_slots[layer_idx] = tuple( + state.batch.sparse_out_slots(list(descriptor.buffer_ids)) + ) + # Fetch even for an empty selection so every output slot is refreshed + # on graph replay. Native Batch supplies page tables and watermarks. + stream = CudaStream(torch.cuda.current_stream().cuda_stream) + owner = descriptor.selection_owner + if owner is None: + state.batch.fetch_sparse_pages( + buffers=descriptor.buffer_ids, + selected=state.selected_history_units, + out=out, + stream=stream, + ) + else: + # The layers an indexer selects for share one resident set, so the + # selection is resolved once -- by the layer that owns the indexer -- + # and each layer then copies those units into its own cells. Resolving + # per layer instead would have the later resolves read the first's + # admissions as hits and leave their cells unfilled. + if owner == layer_idx: + state.batch.resolve_sparse_selection( + buffers=descriptor.buffer_ids, + selected=state.selected_history_units, + stream=stream, + ) + state.batch.fetch_resolved_pages( + buffers=descriptor.buffer_ids, + selected=state.selected_history_units, + out=out, + stream=stream, + ) + return out diff --git a/tensorrt_llm/_torch/attention/backends/sparse/offload_kernels.py b/tensorrt_llm/_torch/attention/backends/sparse/offload_kernels.py new file mode 100644 index 000000000000..8e1d9a599473 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/sparse/offload_kernels.py @@ -0,0 +1,358 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kernels shared by sparse attention models that offload KV history. + +A decode step's selection names units of a request's history; the KV cache +manager fetches them into per-request arena rows. These kernels mask a +selection to the history the Batch published, and assemble the arena rows the +attention reads from the fetch output and the in-flight unit's write row. + +Model inputs follow the requests' input positions, while the native Batch keeps +each request in a fixed row, its slot, for as long as the request is scheduled. +The kernels read input-ordered tensors by position and Batch-owned tensors by +row, through a per-step map from each request's position to its Batch row. +""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + + +def check_sparse_offload_tensor( + tensor: torch.Tensor, name: str, ndim: int, device: torch.device +) -> None: + """Require a CUDA int32 tensor of ``ndim`` dimensions on ``device``.""" + if tensor.dtype != torch.int32 or tensor.ndim != ndim: + raise ValueError(f"{name} must be a {ndim}D int32 tensor") + if not tensor.is_cuda or tensor.device != device: + raise ValueError(f"{name} must be on CUDA device {device}") + + +@triton.jit(do_not_specialize=["num_lengths"]) +def _select_sparse_history_units_kernel( + topk_ptr, + request_indices_ptr, + batch_rows_ptr, + active_request_count_ptr, + unit_lengths_ptr, + history_lengths_ptr, + out_ptr, + topk_stride0, + topk_stride1, + request_indices_stride, + batch_rows_stride, + unit_lengths_stride, + history_lengths_stride, + out_stride0, + out_stride1, + NUM_ROWS: tl.constexpr, + num_lengths, + TOPK: tl.constexpr, + UNIT_TOKENS: tl.constexpr, + BLOCK_K: tl.constexpr, +): + query = tl.program_id(0) + active_count = tl.load(active_request_count_ptr) + request = tl.load(request_indices_ptr + query * request_indices_stride) + active = (query < active_count) & (request >= 0) & (request < active_count) + active = active & (request < NUM_ROWS) & (request < num_lengths) + row = tl.load(batch_rows_ptr + request * batch_rows_stride, active, other=-1) + active = active & (row >= 0) & (row < NUM_ROWS) + length = tl.load(unit_lengths_ptr + request * unit_lengths_stride, active, other=0) + history = tl.load(history_lengths_ptr + row * history_lengths_stride, active, other=0) + k = tl.arange(0, BLOCK_K) + unit = tl.load(topk_ptr + query * topk_stride0 + k * topk_stride1, k < TOPK, other=-1) + valid = (unit >= 0) & (unit < length) & (unit < history // UNIT_TOKENS) + # Each live query owns a distinct Batch row. Keep top-k positions, including + # holes left by in-flight units, so fetch output has the same slot meaning. + tl.store( + out_ptr + row * out_stride0 + k * out_stride1, + tl.where(valid, unit, -1), + mask=active & (k < TOPK), + ) + + +def select_sparse_history_units( + topk_indices: torch.Tensor, + request_indices: torch.Tensor, + batch_rows: torch.Tensor, + active_request_count: torch.Tensor, + unit_lengths: torch.Tensor, + history_lengths: torch.Tensor, + out: torch.Tensor, + *, + unit_tokens: int, +) -> None: + """Mask non-history selections and scatter them into native Batch rows. + + All tensors are CUDA int32 on one device. ``topk_indices`` is [Q, K] in + selection units, with distinct nonnegative selections per live row. + ``request_indices`` [Q] gives each query's request by input position, one + query per request. ``batch_rows`` [B] gives each request's native Batch row, + or -1 for a position without a live request; live requests have distinct + rows, in any order and with free rows between them. ``active_request_count`` + [1] excludes graph padding. ``unit_lengths`` [R] counts each request's + selectable units by input position, and may cover fewer requests than the + native Batch capacity B. ``history_lengths`` [B] is the Batch-owned, + published history watermark in original tokens. + + ``out`` is caller-owned [B, K], disjoint from the inputs. Only ordinals below + both the unit length and history_lengths // unit_tokens are fetched. + Top-k order and padding positions are preserved; in-flight units and invalid + selections become -1. Batch rows without a live query are overwritten with + -1 on every call, including when a graph's Q is smaller than B. No tensor + allocation, host read, or synchronization is performed. + """ + device = out.device + for tensor, name, ndim in ( + (topk_indices, "topk_indices", 2), + (request_indices, "request_indices", 1), + (batch_rows, "batch_rows", 1), + (active_request_count, "active_request_count", 1), + (unit_lengths, "unit_lengths", 1), + (history_lengths, "history_lengths", 1), + (out, "out", 2), + ): + check_sparse_offload_tensor(tensor, name, ndim, device) + num_rows, width = out.shape + num_queries, topk = topk_indices.shape + if unit_tokens <= 0: + raise ValueError("unit_tokens must be positive") + if topk != width: + raise ValueError("top-k width must match the configured sparse selection width") + if request_indices.shape != (num_queries,) or active_request_count.shape != (1,): + raise ValueError("request_indices must be [Q] and active_request_count must be [1]") + if history_lengths.shape != (num_rows,) or batch_rows.shape != (num_rows,): + raise ValueError("history_lengths and batch_rows must cover the native Batch capacity") + if not num_queries <= unit_lengths.numel() <= num_rows: + raise ValueError("Query and unit-length rows must fit the native Batch capacity") + + with torch.cuda.device(device): + # Clear separately: padded queries can carry stale or invalid request IDs, + # and a smaller graph may leave native Batch rows without any query. + out.fill_(-1) + if num_queries == 0 or topk == 0: + return + _select_sparse_history_units_kernel[(num_queries,)]( + topk_indices, + request_indices, + batch_rows, + active_request_count, + unit_lengths, + history_lengths, + out, + *topk_indices.stride(), + request_indices.stride(0), + batch_rows.stride(0), + unit_lengths.stride(0), + history_lengths.stride(0), + *out.stride(), + NUM_ROWS=num_rows, + num_lengths=unit_lengths.numel(), + TOPK=topk, + UNIT_TOKENS=unit_tokens, + BLOCK_K=triton.next_power_of_2(topk), + ) + + +@triton.jit(do_not_specialize=["num_lengths"]) +def _merge_sparse_arena_indices_kernel( + topk_ptr, + request_indices_ptr, + batch_rows_ptr, + active_request_count_ptr, + unit_lengths_ptr, + history_lengths_ptr, + fetched_ptr, + write_ptr, + out_ptr, + topk_stride0, + topk_stride1, + request_indices_stride, + batch_rows_stride, + unit_lengths_stride, + history_lengths_stride, + fetched_stride0, + fetched_stride1, + write_stride, + out_stride0, + out_stride1, + NUM_ROWS: tl.constexpr, + num_lengths, + TOPK: tl.constexpr, + UNIT_TOKENS: tl.constexpr, + BLOCK_K: tl.constexpr, +): + query = tl.program_id(0) + active_count = tl.load(active_request_count_ptr) + request = tl.load(request_indices_ptr + query * request_indices_stride) + live = (query < active_count) & (request >= 0) & (request < active_count) + live = live & (request < NUM_ROWS) & (request < num_lengths) + row = tl.load(batch_rows_ptr + request * batch_rows_stride, live, other=-1) + live = live & (row >= 0) & (row < NUM_ROWS) + length = tl.load(unit_lengths_ptr + request * unit_lengths_stride, live, other=0) + history_units = ( + tl.load(history_lengths_ptr + row * history_lengths_stride, live, other=0) // UNIT_TOKENS + ) + k = tl.arange(0, BLOCK_K) + in_width = k < TOPK + unit = tl.load(topk_ptr + query * topk_stride0 + k * topk_stride1, in_width, other=-1) + valid = live & (unit >= 0) & (unit < length) + fetched = tl.load( + fetched_ptr + row * fetched_stride0 + k * fetched_stride1, + valid & in_width & (unit < history_units), + other=-1, + ) + write = tl.load(write_ptr + row * write_stride, live, other=-1) + tl.store( + out_ptr + query * out_stride0 + k * out_stride1, + tl.where(valid & (unit == history_units), write, fetched), + mask=in_width, + ) + + +def merge_sparse_arena_indices( + topk_indices: torch.Tensor, + request_indices: torch.Tensor, + batch_rows: torch.Tensor, + active_request_count: torch.Tensor, + unit_lengths: torch.Tensor, + history_lengths: torch.Tensor, + fetched_indices: torch.Tensor, + write_indices: torch.Tensor, + out: torch.Tensor, + *, + unit_tokens: int, +) -> torch.Tensor: + """Return each query's selection as rows of one sparse buffer's arena. + + Inputs follow :func:`select_sparse_history_units`. ``fetched_indices`` is + the fetch output of one buffer, [B, K] in Batch row order and top-k column + order; ``write_indices`` [B] is that buffer's in-flight row by Batch row, + written this step. Both may be strided views of the per-buffer outputs. + + ``out`` is [Q, K] in query order. A selected history unit takes its fetched + row, the in-flight unit takes the write row, and padding queries, invalid + selections and units past the unit length take -1. No tensor allocation, + host read, or synchronization is performed. + """ + device = out.device + for tensor, name, ndim in ( + (topk_indices, "topk_indices", 2), + (request_indices, "request_indices", 1), + (batch_rows, "batch_rows", 1), + (active_request_count, "active_request_count", 1), + (unit_lengths, "unit_lengths", 1), + (history_lengths, "history_lengths", 1), + (fetched_indices, "fetched_indices", 2), + (write_indices, "write_indices", 1), + (out, "out", 2), + ): + check_sparse_offload_tensor(tensor, name, ndim, device) + num_rows, width = fetched_indices.shape + num_queries, topk = topk_indices.shape + if unit_tokens <= 0: + raise ValueError("unit_tokens must be positive") + if topk != width or out.shape != (num_queries, topk): + raise ValueError("Selections, fetch output and out must share the top-k width") + if request_indices.shape != (num_queries,) or active_request_count.shape != (1,): + raise ValueError("request_indices must be [Q] and active_request_count must be [1]") + if ( + history_lengths.shape != (num_rows,) + or write_indices.shape != (num_rows,) + or batch_rows.shape != (num_rows,) + ): + raise ValueError( + "history_lengths, write_indices and batch_rows must cover the native Batch" + ) + if not num_queries <= unit_lengths.numel() <= num_rows: + raise ValueError("Query and unit-length rows must fit the native Batch capacity") + if num_queries == 0 or topk == 0: + return out + with torch.cuda.device(device): + _merge_sparse_arena_indices_kernel[(num_queries,)]( + topk_indices, + request_indices, + batch_rows, + active_request_count, + unit_lengths, + history_lengths, + fetched_indices, + write_indices, + out, + *topk_indices.stride(), + request_indices.stride(0), + batch_rows.stride(0), + unit_lengths.stride(0), + history_lengths.stride(0), + *fetched_indices.stride(), + write_indices.stride(0), + *out.stride(), + NUM_ROWS=num_rows, + num_lengths=unit_lengths.numel(), + TOPK=topk, + UNIT_TOKENS=unit_tokens, + BLOCK_K=triton.next_power_of_2(topk), + ) + return out + + +@triton.jit(do_not_specialize=["num_rows", "num_out"]) +def _gather_sparse_batch_rows_kernel( + values_ptr, + batch_rows_ptr, + out_ptr, + values_stride, + batch_rows_stride, + out_stride, + num_rows, + num_out, + BLOCK: tl.constexpr, +): + position = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + in_range = position < num_out + row = tl.load(batch_rows_ptr + position * batch_rows_stride, in_range, other=-1) + valid = in_range & (row >= 0) & (row < num_rows) + value = tl.load(values_ptr + row * values_stride, valid, other=-1) + tl.store(out_ptr + position * out_stride, value, mask=in_range) + + +def gather_sparse_batch_rows( + values: torch.Tensor, + batch_rows: torch.Tensor, + out: torch.Tensor | None = None, +) -> torch.Tensor: + """Read a per-Batch-row value for each request, in input order. + + ``values`` [B] holds one int32 per native Batch row, such as one buffer's + write rows; ``batch_rows`` [N] gives each request's Batch row, or -1 for a + position without a live request, which reads -1. Both may be strided. ``out`` + [N] receives the result in place, also possibly strided, and is allocated + when None. One kernel launch; no host read or synchronization. + """ + device = values.device + if out is None: + out = torch.empty(batch_rows.shape, dtype=torch.int32, device=device) + for tensor, name in ((values, "values"), (batch_rows, "batch_rows"), (out, "out")): + check_sparse_offload_tensor(tensor, name, 1, device) + if out.shape != batch_rows.shape: + raise ValueError("out must have one entry per request") + num_out = batch_rows.numel() + if num_out == 0: + return out + block = 256 + with torch.cuda.device(device): + _gather_sparse_batch_rows_kernel[(triton.cdiv(num_out, block),)]( + values, + batch_rows, + out, + values.stride(0), + batch_rows.stride(0), + out.stride(0), + values.numel(), + num_out, + BLOCK=block, + ) + return out diff --git a/tensorrt_llm/_torch/attention/backends/trtllm.py b/tensorrt_llm/_torch/attention/backends/trtllm.py index 1a9121037a1d..5b1a470880b2 100644 --- a/tensorrt_llm/_torch/attention/backends/trtllm.py +++ b/tensorrt_llm/_torch/attention/backends/trtllm.py @@ -125,6 +125,26 @@ def generate_spec_decoding_packed_mask(max_num_requests: int, return mask +@dataclass(frozen=True, kw_only=True) +class MlaLatentWriteTarget: + """Where ``mla_rope_generation`` stores each token's new latent row. + + By default the row goes to the token's position in the paged cache. A + target instead stores it in one row of a flat row array, such as a sparse + cache arena: pool 0 of ``pool_pointers`` starts at the array, the layer's + ``pool_mapping`` row selects pool 0, and each request owns a one-token page + whose block offset in ``block_offsets`` [1, requests, 2, 1] is its row. + An NVFP4 cache's pointers are [1, 2, 2]: the block scales go to the same row + of a second array. ``write_slots`` [tokens] is 0 for a token to store and + -1 for a token to skip, such as CUDA graph padding. + """ + + block_offsets: torch.Tensor + pool_pointers: torch.Tensor + pool_mapping: torch.Tensor + write_slots: torch.Tensor + + @dataclass(kw_only=True) class TrtllmAttentionMetadata(AttentionMetadata): workspace: Optional[torch.Tensor] = None @@ -2849,6 +2869,7 @@ def mla_rope_generation( quant_scale_qkv: Optional[torch.Tensor] = None, fuse_fp4_q_quant: bool = False, q_rope_applied: bool = False, + latent_write_target: Optional[MlaLatentWriteTarget] = None, ) -> None: """ fused_q (torch.Tensor): The tensor to store the fused q, with shape (num_tokens, num_heads, kv_lora_rank + qk_rope_head_dim) on GPU. @@ -2871,6 +2892,7 @@ def mla_rope_generation( quant_scale_qkv (torch.Tensor): Non-None means q_nope in quant_q_buffer is already FP8, so the kernel drops the q_nope quantize rows from its grid. fuse_fp4_q_quant (bool): Quantize Q in the fused FP4 RoPE/cache update. q_rope_applied (bool): quant_q_buffer already holds the complete rotated FP8 Q (the DSv4 q_b GEMM fusion wrote it), so the Q half must not rewrite the rope segment. + latent_write_target (MlaLatentWriteTarget): Store the new latent rows there instead of in the paged cache. """ assert self.is_mla_enable and self.mla_params is not None @@ -2881,6 +2903,10 @@ def mla_rope_generation( self._ensure_rope_table_size(metadata.max_seq_len) if self.uses_fp4_mla_attention: + if latent_write_target is not None: + raise NotImplementedError( + "FP4 MLA generation cannot store its latent outside the paged cache" + ) self._fp4_mla_rope_generation( fused_q, q_pe, @@ -2893,7 +2919,16 @@ def mla_rope_generation( helix_tensor_params = [ metadata.helix_position_offsets, metadata.helix_is_inactive_rank ] - if metadata._helix_spec_tokens_valid: + if latent_write_target is not None: + if metadata._helix_spec_tokens_valid or kv_norm_weight is not None: + raise NotImplementedError( + "A latent write target needs the per-token write slots, which Helix " + "speculative verify groups and the fused KV norm do not leave free" + ) + # Each request's one-token page is the target row; the per-token slot + # 0 addresses it, and -1 skips the token. + helix_tensor_params.append(latent_write_target.write_slots) + elif metadata._helix_spec_tokens_valid: # Speculative verify groups: per-token KV write slots (-1 = this # rank does not own the token's position). The append kernel then # gates and addresses per token instead of per sequence. @@ -2905,12 +2940,27 @@ def mla_rope_generation( # request contributes the same number of tokens. with self.presented_token_major_for(metadata): self._mla_rope_generation_impl( - metadata, fused_q, q_pe, latent_cache, cu_q_seqlens, - cu_kv_seqlens, fmha_scheduler_counter, mla_bmm1_scale, - mla_bmm2_scale, quant_q_buffer, out_scale, helix_tensor_params, - kv_norm_weight, kv_norm_eps, precomputed_cu_seqlens, - precomputed_fmha_scheduler, kv_only, kv_done_elsewhere, - quant_scale_qkv, q_rope_applied) + metadata, + fused_q, + q_pe, + latent_cache, + cu_q_seqlens, + cu_kv_seqlens, + fmha_scheduler_counter, + mla_bmm1_scale, + mla_bmm2_scale, + quant_q_buffer, + out_scale, + helix_tensor_params, + kv_norm_weight, + kv_norm_eps, + precomputed_cu_seqlens, + precomputed_fmha_scheduler, + kv_only, + kv_done_elsewhere, + quant_scale_qkv, + q_rope_applied, + latent_write_target=latent_write_target) @contextlib.contextmanager def presented_token_major_for( @@ -2940,7 +2990,22 @@ def _mla_rope_generation_impl( kv_done_elsewhere: bool, quant_scale_qkv: Optional[torch.Tensor], q_rope_applied: bool, + latent_write_target: Optional[MlaLatentWriteTarget] = None, ) -> None: + if latent_write_target is None: + # Read under the caller's presentation: a ragged verify batch + # presents a token-major block table, one row per token. + block_offsets = metadata.kv_cache_block_offsets + pool_pointers = metadata.kv_cache_manager.kv_cache_pool_pointers + pool_mapping = metadata.kv_cache_manager.kv_cache_pool_mapping + layer_idx = self.get_local_layer_idx(metadata) + tokens_per_block = metadata.kv_cache_manager.tokens_per_block + else: + block_offsets = latent_write_target.block_offsets + pool_pointers = latent_write_target.pool_pointers + pool_mapping = latent_write_target.pool_mapping + layer_idx = 0 + tokens_per_block = 1 torch.ops.trtllm.mla_rope_generation( fused_q, q_pe, @@ -2956,9 +3021,9 @@ def _mla_rope_generation_impl( metadata.kv_lens_runtime, # host_past_key_value_lengths metadata.prompt_lens_cpu_runtime, # host_context_lengths, metadata.num_contexts, - metadata.kv_cache_block_offsets, - metadata.kv_cache_manager.kv_cache_pool_pointers, - metadata.kv_cache_manager.kv_cache_pool_mapping, + block_offsets, + pool_pointers, + pool_mapping, None, # kv_scale_orig_quant None, # kv_scale_quant_orig getattr(self, "_nvfp4_mla_kv_scale_orig_quant", None), @@ -2966,12 +3031,12 @@ def _mla_rope_generation_impl( metadata.block_ids_per_seq, helix_tensor_params, self.predicted_tokens_per_seq, - self.get_local_layer_idx(metadata), + layer_idx, self.num_heads, self.num_kv_heads, self.head_dim, getattr(metadata.kv_cache_manager, "mla_kv_cache_residual_dim", 0), - metadata.kv_cache_manager.tokens_per_block, + tokens_per_block, metadata.max_seq_len, # attention_window_size metadata.beam_width, self.quant_mode, diff --git a/tensorrt_llm/_torch/disaggregation/native/transfer.py b/tensorrt_llm/_torch/disaggregation/native/transfer.py index e28ca9d91dd9..eb3b50b45d79 100644 --- a/tensorrt_llm/_torch/disaggregation/native/transfer.py +++ b/tensorrt_llm/_torch/disaggregation/native/transfer.py @@ -22,6 +22,7 @@ import threading import time import weakref +from collections import defaultdict from contextlib import contextmanager, nullcontext from dataclasses import dataclass from enum import Enum @@ -45,6 +46,7 @@ RegMemoryDescs, TransferOp, TransferRequest, + TransferStatus, ) from tensorrt_llm._torch.disaggregation.base.transfer import ( ReceiverBase, @@ -89,7 +91,10 @@ from tensorrt_llm._torch.disaggregation.nixl.agent import NixlTransferAgent from tensorrt_llm._torch.disaggregation.resource.kv_extractor import KVRegionExtractorV1 from tensorrt_llm._torch.disaggregation.resource.page import CacheKind, KVCachePageTable, MapperKind -from tensorrt_llm._torch.disaggregation.resource.utils import get_unique_pool_memory_descs +from tensorrt_llm._torch.disaggregation.resource.utils import ( + get_physical_pool, + get_unique_pool_memory_descs, +) from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest from tensorrt_llm._utils import CUASSERT, nvtx_range @@ -173,6 +178,49 @@ class WriteMetaType(Enum): AUX = "AUX" +@dataclass(frozen=True) +class _TransferGroup: + src_memory_type: str + dst_memory_type: str + start: int + stop: int + + +@dataclass(frozen=True) +class _TransferRequestGroup: + requests: tuple[TransferRequest, ...] + + +class _TransferGroupStatus: + """Retain every submitted write until the whole group has completed. + + An incomplete submission has no proof of quiescence, even if the handles + returned before a submission error have completed. + """ + + def __init__(self) -> None: + self.statuses: list[TransferStatus] = [] + self.submission_complete = False + + def is_completed(self) -> bool: + return self.submission_complete and all( + status.is_completed() is True for status in self.statuses + ) + + def wait(self) -> bool: + completed = self.submission_complete + for status in self.statuses: + # Wait for every handle, including after another handle fails. + completed = status.wait() and completed + return completed + + def last_status_str(self) -> str: + return "; ".join( + getattr(status, "last_status_str", lambda: "")() + for status in self.statuses + ) + + @dataclass class WriteMeta: task: Union[SendTaskBase, "KVRecvTask"] @@ -191,6 +239,7 @@ class WriteMeta: is_last_slice: bool = False meta_type: WriteMetaType = WriteMetaType.KV bounce_dst_base: Optional[int] = None + transfer_groups: tuple[_TransferGroup, ...] = () class MessageType: @@ -540,7 +589,7 @@ class _SessionQuiescence: class Sender(SenderBase): - # Time-to-live for orphaned RecvReqInfo entries (seconds). + # Minimum time-to-live for orphaned RecvReqInfo entries (seconds). # In gen-first ADP broadcast, non-assigned DP ranks accumulate # RecvReqInfo that never gets consumed. Entries older than this # are evicted during periodic sweeps. @@ -553,6 +602,7 @@ def __init__( bounce=None, enforce_physical_ownership: bool = False, retirement_watchdog: Optional[RetirementWatchdog] = None, + request_info_timeout_s: Optional[float] = None, ) -> None: self._registrar = peer_registrar self._device_id = peer_registrar.self_rank_info.device_id @@ -560,6 +610,9 @@ def __init__( self._bounce = bounce self._enforce_physical_ownership = enforce_physical_ownership self._retirement_watchdog = retirement_watchdog + # Generation-first metadata can arrive while prefill is still queued. + # Keep it for at least the configured transfer timeout. + self._request_info_ttl_s = max(self._STALE_REQ_INFO_TTL_S, request_info_timeout_s or 0.0) self._peer_requests: dict = {} self._peer_requests_timestamps: dict[int, float] = {} # unique_rid -> insert time self._peer_requests_lock = threading.Lock() @@ -658,7 +711,7 @@ def sweep_stale_req_infos(self): stale_rids = [ rid for rid, ts in self._peer_requests_timestamps.items() - if now - ts > self._STALE_REQ_INFO_TTL_S + if now - ts > self._request_info_ttl_s ] if not stale_rids: return @@ -809,13 +862,32 @@ def _submit_transfer( self, task: SendTaskBase, peer_rank: int, - request: TransferRequest, + request: TransferRequest | _TransferRequestGroup, ) -> tuple[bool, Optional[str]]: if not self._enforce_physical_ownership: - status = self._agent.submit_transfer_requests(request) - completed = status.wait() - detail = None if completed else getattr(status, "last_status_str", lambda: None)() - return completed, detail + requests = ( + request.requests if isinstance(request, _TransferRequestGroup) else (request,) + ) + # Finish each group before submitting the next so a later submission + # failure cannot abandon an earlier group's in-flight handle. + for member in requests: + status = self._agent.submit_transfer_requests(member) + if not status.wait(): + return False, getattr(status, "last_status_str", lambda: None)() + return True, None + + group_status = ( + _TransferGroupStatus() if isinstance(request, _TransferRequestGroup) else None + ) + + def submit() -> TransferStatus | _TransferGroupStatus: + if group_status is None: + return self._agent.submit_transfer_requests(request) + for member in request.requests: + group_status.statuses.append(self._agent.submit_transfer_requests(member)) + group_status.submission_complete = True + return group_status + with self._ownership_poison_lock: if self._ownership_poisoned is not None: error = self._ownership_poisoned @@ -827,8 +899,13 @@ def _submit_transfer( try: # SUBMITTING already retains the request. Neither backend admission # nor completion may hold the lock needed by deadline containment. - status = self._agent.submit_transfer_requests(request) - task.record_backend_submission(peer_rank, status) + if group_status is not None: + # Keep earlier handles if a later submission raises. Until all + # submissions finish, the group cannot attest quiescence. + task.record_backend_submission(peer_rank, group_status) + status = submit() + if group_status is None: + task.record_backend_submission(peer_rank, status) except Exception as error: with self._ownership_poison_lock: self._ownership_poisoned = error @@ -990,7 +1067,9 @@ def _poll_in_doubt_transfers(self, thread_idx: int) -> None: @staticmethod @nvtx_range("_make_agent_request") - def _make_agent_request(write_meta: WriteMeta, device_id: int) -> "TransferRequest": + def _make_agent_request( + write_meta: WriteMeta, device_id: int + ) -> TransferRequest | _TransferRequestGroup: if not (write_meta.src_ptrs.size == write_meta.dst_ptrs.size == write_meta.sizes.size): raise ValueError( f"Pointer/size mismatch for unique_rid={write_meta.unique_rid}: " @@ -998,34 +1077,43 @@ def _make_agent_request(write_meta: WriteMeta, device_id: int) -> "TransferReque f"{write_meta.dst_ptrs.size=}, " f"{write_meta.sizes.size=}" ) - n = write_meta.src_ptrs.size if write_meta.meta_type == WriteMetaType.AUX: - src_dev, dst_dev, mem_type = 0, 0, MemoryType.DRAM + groups = (_TransferGroup("DRAM", "DRAM", 0, write_meta.src_ptrs.size),) else: if write_meta.dst_device_id is None: raise RuntimeError( f"_make_agent_request: dst_device_id is None for KV transfer " f"unique_rid={write_meta.unique_rid}" ) - src_dev, dst_dev, mem_type = device_id, write_meta.dst_device_id, MemoryType.VRAM + groups = write_meta.transfer_groups or ( + _TransferGroup("VRAM", "VRAM", 0, write_meta.src_ptrs.size), + ) - if n == 0: - src_memory_descs = MemoryDescs(mem_type, []) - dst_memory_descs = MemoryDescs(mem_type, []) - else: + memory_types = {"VRAM": MemoryType.VRAM, "DRAM": MemoryType.DRAM} + requests = [] + for group in groups: + span = slice(group.start, group.stop) + src_dev = device_id if group.src_memory_type == "VRAM" else 0 + dst_dev = write_meta.dst_device_id if group.dst_memory_type == "VRAM" else 0 src_memory_descs = MemoryDescs.from_arrays_uniform_device( - mem_type, write_meta.src_ptrs, write_meta.sizes, src_dev + memory_types[group.src_memory_type], + write_meta.src_ptrs[span], + write_meta.sizes[span], + src_dev, ) dst_memory_descs = MemoryDescs.from_arrays_uniform_device( - mem_type, write_meta.dst_ptrs, write_meta.sizes, dst_dev + memory_types[group.dst_memory_type], + write_meta.dst_ptrs[span], + write_meta.sizes[span], + dst_dev, ) - - # NOTE: TransferRequest moves (not copies) src/dst MemoryDescs internally. - # After this call, src_memory_descs and dst_memory_descs are in a moved-from - # state and must NOT be accessed again. - return TransferRequest( - TransferOp.WRITE, src_memory_descs, dst_memory_descs, write_meta.peer_name, None - ) + # TransferRequest moves the descriptors; do not access them afterwards. + requests.append( + TransferRequest( + TransferOp.WRITE, src_memory_descs, dst_memory_descs, write_meta.peer_name, None + ) + ) + return requests[0] if len(requests) == 1 else _TransferRequestGroup(tuple(requests)) @nvtx_range("_deliver_kv_to_agent") def _deliver_kv_to_agent(self, write_meta: WriteMeta): @@ -1384,13 +1472,13 @@ def _build_kv_write_meta(self, task: KVSendTask, req_info: RecvReqInfo) -> Write # Aggregate fragment pointers from all matching pool pairs. # Each pool pair produces one or more region pairs (rp), each containing # a numpy array of src/dst pointers and a uniform bytes_per_region. - src_frag_parts: list[np.ndarray] = [] - dst_frag_parts: list[np.ndarray] = [] + src_frag_parts: dict[tuple[str, str], list[np.ndarray]] = defaultdict(list) + dst_frag_parts: dict[tuple[str, str], list[np.ndarray]] = defaultdict(list) # Instead of calling np.full() per region pair to build a size array and # then np.concatenate() all of them, we record (count, bytes_per_region) # tuples and construct the final sizes array with a single np.repeat(). # For 48k+ items this avoids many small allocations in the hot loop. - size_specs: list[tuple[int, int]] = [] + size_specs: dict[tuple[str, str], list[tuple[int, int]]] = defaultdict(list) dst_device_id = peer_ri.device_id extractor = self._registrar.self_extractor peer_extractor = self._registrar.peer_extractor( @@ -1450,20 +1538,38 @@ def _build_kv_write_meta(self, task: KVSendTask, req_info: RecvReqInfo) -> Write mapper = self._registrar.get_kv_map(peer_ri, (self_lg, self_pi), (peer_lg, peer_pi)) region_pair = mapper.map(src_region, dst_region) region_pairs = region_pair if isinstance(region_pair, list) else [region_pair] + src_view = lg_info.pool_views[self_pi] + dst_view = peer_extractor.page_table.layer_groups[peer_lg].pool_views[peer_pi] + memory_types = ( + get_physical_pool(extractor.page_table, self_lg, src_view.pool_idx).memory_type, + get_physical_pool( + peer_extractor.page_table, peer_lg, dst_view.pool_idx + ).memory_type, + ) for rp in region_pairs: - src_frag_parts.append(rp.src.memory.ptrs) - dst_frag_parts.append(rp.dst.memory.ptrs) - size_specs.append((rp.src.memory.ptrs.size, rp.src.memory.bytes_per_region)) + if rp.src.memory.ptrs.size == 0: + continue + src_frag_parts[memory_types].append(rp.src.memory.ptrs) + dst_frag_parts[memory_types].append(rp.dst.memory.ptrs) + size_specs[memory_types].append( + (rp.src.memory.ptrs.size, rp.src.memory.bytes_per_region) + ) + transfer_groups = [] if src_frag_parts: - src_frags = np.concatenate(src_frag_parts) - dst_frags = np.concatenate(dst_frag_parts) + src_frags = np.concatenate([p for parts in src_frag_parts.values() for p in parts]) + dst_frags = np.concatenate([p for parts in dst_frag_parts.values() for p in parts]) # Build the kv_sizes array in one shot: np.repeat expands each # bytes_per_region value by its count, e.g.: # values=[4096, 8192], counts=[100, 200] # → [4096]*100 ++ [8192]*200 - counts, values = zip(*size_specs) + counts, values = zip(*(spec for specs in size_specs.values() for spec in specs)) kv_sizes = np.repeat(np.array(values, dtype=np.int64), counts) + start = 0 + for (src_type, dst_type), specs in size_specs.items(): + stop = start + sum(count for count, _ in specs) + transfer_groups.append(_TransferGroup(src_type, dst_type, start, stop)) + start = stop else: src_frags = np.array([], dtype=np.int64) dst_frags = np.array([], dtype=np.int64) @@ -1487,7 +1593,12 @@ def _build_kv_write_meta(self, task: KVSendTask, req_info: RecvReqInfo) -> Write slice_id=task.slice_id, receiver_slice_id=req_info.slice_id if req_info.slice_id is not None else 0, is_last_slice=task._chunk.is_last, - bounce_dst_base=req_info.bounce_dst_base, + bounce_dst_base=( + req_info.bounce_dst_base + if all(g.src_memory_type == g.dst_memory_type == "VRAM" for g in transfer_groups) + else None + ), + transfer_groups=tuple(transfer_groups), ) def _build_aux_write_meta(self, task: AuxSendTask, req_info: RecvReqInfo) -> WriteMeta: @@ -2928,6 +3039,12 @@ def dispatch_task(self, task: KVRecvTask) -> None: receiver_page_table=self._registrar.self_extractor.page_table, dst_slot=mamba_dst_slot, ) + # Bounce gather/scatter operates on device allocations. A mixed-memory + # handoff writes directly to the published host and device pools. + allow_bounce = allow_bounce and not any( + page_table is not None and page_table.has_host_pools + for page_table in (self._registrar.self_extractor.page_table, peer_infos.page_table) + ) bounced = allow_bounce and self._bounce.reserve( receiver_req, task.expected_transfers, extra_bytes=extra_bytes ) @@ -4303,6 +4420,7 @@ def _setup_transfer_engine(self): bounce=self._bounce, enforce_physical_ownership=self._config.enforce_physical_ownership, retirement_watchdog=self._retirement_watchdog, + request_info_timeout_s=self._config.tx_overall_timeout_s, ) self._receiver = Receiver( self._peer_registrar, @@ -4323,14 +4441,15 @@ def _setup_transfer_engine(self): def _register_kv_cache(self): assert self._rank_info.page_table is not None - memory_descs = get_unique_pool_memory_descs( - self._rank_info.page_table, self._rank_info.device_id - ) - if memory_descs: - reg_memory_desc = RegMemoryDescs("VRAM", memory_descs) - self._agent.register_memory(reg_memory_desc) - logger.debug(f"Registered KV cache memory with transfer agent: {memory_descs}") - self._registered_mem.append(reg_memory_desc) + for memory_type in ("VRAM", "DRAM"): + memory_descs = get_unique_pool_memory_descs( + self._rank_info.page_table, self._rank_info.device_id, memory_type + ) + if memory_descs: + reg_memory_desc = RegMemoryDescs(memory_type, memory_descs) + self._agent.register_memory(reg_memory_desc) + logger.debug(f"Registered {memory_type} KV cache memory: {memory_descs}") + self._registered_mem.append(reg_memory_desc) def _register_aux_buffer(self): assert self._aux_buffer is not None diff --git a/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py b/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py index b1132c10ddc4..5fc48b9adb56 100644 --- a/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py +++ b/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py @@ -397,10 +397,30 @@ def _window_size_for_layer(internal_layer_id: int): pool_groups: List[PhysicalPoolGroup] = [] storage_pg_to_list_idx: Dict[int, int] = {} layer_groups_by_id: List[LayerGroup | None] = [None] * len(manager.impl.layer_grouping) + # Native sparse pool descriptors expose their host hot-tier allocations. + # Dense buffers in the same layer group still have device-backed pools. + sparse_buffers = { + (int(layer.layer_id), str(buffer.role)) + for layer in config.layers + for buffer in getattr(layer, "buffers", ()) + if getattr(buffer, "is_sparse", False) + } for pg_desc in pool_group_descs: storage_pg_idx = int(pg_desc.pool_group_index) storage_pg_to_list_idx[storage_pg_idx] = len(pool_groups) + pool_memory_types = [] + for pool_idx in range(len(pg_desc.pools)): + backing = { + (int(buffer_id.layer_id), str(buffer_id.role)) in sparse_buffers + for variant in pg_desc.slot_desc.variants + for buffer_id in variant.coalesced_buffers[pool_idx].buffer_ids + } + if len(backing) > 1: + raise ValueError( + f"Physical pool {storage_pg_idx}/{pool_idx} mixes host and device buffers" + ) + pool_memory_types.append("DRAM" if True in backing else "VRAM") pool_groups.append( PhysicalPoolGroup( pools=[ @@ -408,8 +428,9 @@ def _window_size_for_layer(internal_layer_id: int): base_address=int(pool.base_address), slot_bytes=int(pool.slot_bytes), num_slots=int(pg_desc.num_slots), + memory_type=pool_memory_types[pool_idx], ) - for pool in pg_desc.pools + for pool_idx, pool in enumerate(pg_desc.pools) ] ) ) diff --git a/tensorrt_llm/_torch/disaggregation/resource/page.py b/tensorrt_llm/_torch/disaggregation/resource/page.py index 1254c06001f6..4acc3b357bfa 100644 --- a/tensorrt_llm/_torch/disaggregation/resource/page.py +++ b/tensorrt_llm/_torch/disaggregation/resource/page.py @@ -17,7 +17,7 @@ from dataclasses import dataclass, field from enum import IntEnum -from typing import FrozenSet, List, Optional, Tuple +from typing import FrozenSet, List, Literal, Optional, Tuple import numpy as np @@ -103,14 +103,19 @@ class PhysicalPool: ``slot_bytes`` spans all layers and roles coalesced into one slot. The slot address is ``base_address + slot * slot_stride_bytes``; per-layer and per-role offsets are carried by each PoolView. + ``memory_type`` describes the backing allocation, including the host hot + tier for sparse buffers, independently of the logical layer layout. """ base_address: int # uint64 slot_bytes: int num_slots: int slot_stride_bytes: Optional[int] = None + memory_type: Literal["VRAM", "DRAM"] = "VRAM" def __post_init__(self) -> None: + if self.memory_type not in ("VRAM", "DRAM"): + raise ValueError(f"Unsupported KV pool memory type: {self.memory_type!r}") if self.slot_stride_bytes is None: self.slot_stride_bytes = self.slot_bytes if self.slot_stride_bytes < self.slot_bytes: @@ -122,6 +127,7 @@ def to_dict(self) -> dict: "slot_bytes": int(self.slot_bytes), "num_slots": int(self.num_slots), "slot_stride_bytes": int(self.slot_stride_bytes), + "memory_type": self.memory_type, } @staticmethod @@ -135,6 +141,7 @@ def from_dict(data: dict) -> "PhysicalPool": if data.get("slot_stride_bytes") is not None else None ), + memory_type=data.get("memory_type", "VRAM"), ) @@ -371,6 +378,11 @@ class KVCachePageTable: layer_groups: List[LayerGroup] pool_groups: List[PhysicalPoolGroup] # indexed by LayerGroup.pool_group_idx + @property + def has_host_pools(self) -> bool: + """Whether any transferable allocation is backed by host memory.""" + return any(pool.memory_type == "DRAM" for group in self.pool_groups for pool in group.pools) + def to_dict(self) -> dict: return { "tokens_per_block": int(self.tokens_per_block), diff --git a/tensorrt_llm/_torch/disaggregation/resource/utils.py b/tensorrt_llm/_torch/disaggregation/resource/utils.py index 579cc42d1d13..81d70ce5d9a4 100644 --- a/tensorrt_llm/_torch/disaggregation/resource/utils.py +++ b/tensorrt_llm/_torch/disaggregation/resource/utils.py @@ -239,10 +239,13 @@ def get_physical_pool(page_table: KVCachePageTable, lg_idx: int, pool_idx: int) def get_unique_pool_memory_descs( - page_table: KVCachePageTable, device_id: int + page_table: KVCachePageTable, device_id: int, memory_type: str | None = None ) -> list[tuple[int, int, int, str]]: - """Return deduplicated (ptr, size, device_id, name) tuples for all physical pools.""" - unique_pools: dict[tuple[int, int], int] = {} # (ptr, size) -> index + """Return deduplicated pool descriptors, optionally filtered by memory type. + + Host descriptors use device ID zero; device descriptors use the rank's GPU. + """ + unique_pools: dict[tuple[int, int, str], int] = {} pool_counter = 0 for lg_idx, lg in enumerate(page_table.layer_groups): if lg.kind == CacheKind.STATE: @@ -257,9 +260,12 @@ def get_unique_pool_memory_descs( ) for pool in pools: pool_size = pool.num_slots * pool.slot_stride_bytes - if any(b <= pool.base_address < b + s for (b, s) in unique_pools): + if any( + kind == pool.memory_type and b <= pool.base_address < b + s + for b, s, kind in unique_pools + ): continue - pool_key = (pool.base_address, pool_size) + pool_key = (pool.base_address, pool_size, pool.memory_type) if pool_key not in unique_pools: unique_pools[pool_key] = pool_counter pool_counter += 1 @@ -267,13 +273,18 @@ def get_unique_pool_memory_descs( # PAGED (attention): each pool view is an independent allocation for pv in lg.pool_views: pool = get_physical_pool(page_table, lg_idx, pv.pool_idx) - pool_key = (pool.base_address, pool.num_slots * pool.slot_stride_bytes) + pool_key = ( + pool.base_address, + pool.num_slots * pool.slot_stride_bytes, + pool.memory_type, + ) if pool_key not in unique_pools: unique_pools[pool_key] = pool_counter pool_counter += 1 return [ - (pool_ptr, pool_size, device_id, f"kv_cache_memory_pool{idx}") - for (pool_ptr, pool_size), idx in unique_pools.items() + (pool_ptr, pool_size, 0 if kind == "DRAM" else device_id, f"kv_cache_memory_pool{idx}") + for (pool_ptr, pool_size, kind), idx in unique_pools.items() + if memory_type is None or memory_type == kind ] diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py index a29eae1fa19e..2ceccf2e94a7 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/tensorrt_llm/_torch/model_config.py @@ -1237,8 +1237,15 @@ def update_sparse_attention_indexer_config(pretrained_config, kwargs): enable_heuristic_topk = sparse_attention_config.enable_heuristic_topk use_self_sampling_topk = sparse_attention_config.use_self_sampling_topk use_gvr_emission = sparse_attention_config.use_gvr_emission + gvr_deterministic_topk = sparse_attention_config.gvr_deterministic_topk indexer_k_dtype = sparse_attention_config.indexer_k_dtype index_share_for_mtp_iteration = sparse_attention_config.index_share_for_mtp_iteration + offload_kwargs = { + 'enable_kv_cache_offload': + sparse_attention_config.enable_kv_cache_offload, + 'kv_cache_offload_config': + sparse_attention_config.kv_cache_offload_config, + } else: index_n_heads = pretrained_config.index_n_heads index_head_dim = pretrained_config.index_head_dim @@ -1252,8 +1259,10 @@ def update_sparse_attention_indexer_config(pretrained_config, kwargs): enable_heuristic_topk = False use_self_sampling_topk = True use_gvr_emission = False + gvr_deterministic_topk = False indexer_k_dtype = "fp8" index_share_for_mtp_iteration = None + offload_kwargs = {} kwargs[ 'sparse_attention_config'] = DeepSeekSparseAttentionConfig( index_n_heads=index_n_heads, @@ -1270,9 +1279,11 @@ def update_sparse_attention_indexer_config(pretrained_config, kwargs): enable_heuristic_topk=enable_heuristic_topk, use_self_sampling_topk=use_self_sampling_topk, use_gvr_emission=use_gvr_emission, + gvr_deterministic_topk=gvr_deterministic_topk, indexer_k_dtype=indexer_k_dtype, index_share_for_mtp_iteration= - index_share_for_mtp_iteration) + index_share_for_mtp_iteration, + **offload_kwargs) elif pretrained_config.architectures[ 0] in _DEEPSEEK_V41_ARCHITECTURES: sparse_config = kwargs.get('sparse_attention_config') @@ -1369,13 +1380,16 @@ def update_sparse_attention_indexer_config(pretrained_config, kwargs): compress_ratios = list(compress_ratios) + [1] * ( total_layers - len(compress_ratios)) - indexer_k_dtype = indexer_config.pop('indexer_k_dtype') + resolved_sparse_config = ( + sparse_attention_config.model_dump() + if sparse_attention_config is not None else {}) + resolved_sparse_config.update( + compress_ratios=compress_ratios, + window_size=window_size, + **indexer_config) kwargs[ 'sparse_attention_config'] = DeepSeekV4SparseAttentionConfig( - compress_ratios=compress_ratios, - window_size=window_size, - indexer_k_dtype=indexer_k_dtype, - **indexer_config) + **resolved_sparse_config) else: raise ValueError( "checkpoint_dir is None. Cannot load model config without a valid checkpoint directory." diff --git a/tensorrt_llm/_torch/modules/mhc/mhc_cuda.py b/tensorrt_llm/_torch/modules/mhc/mhc_cuda.py index a7ff8409fcc6..a91567d8a1a8 100644 --- a/tensorrt_llm/_torch/modules/mhc/mhc_cuda.py +++ b/tensorrt_llm/_torch/modules/mhc/mhc_cuda.py @@ -826,11 +826,17 @@ def _alloc_fused_hc_scratch( _FUSED_HC_FALLBACK_TACTIC_FMA = ("fused_half_fma", 2, 1, 256, 1) -def _get_fused_hc_fallback_tactic(hidden_size: int | None = None): +def _get_fused_hc_fallback_tactic( + hidden_size: int | None = None, *, deterministic: bool = False +) -> tuple[str, int, int, int, int]: mma_ok = _fused_hc_mma_supported() if hidden_size is not None: mma_ok = mma_ok and hidden_size in _FUSED_HC_MMA_SUPPORTED_HIDDEN_SIZES - return _FUSED_HC_FALLBACK_TACTIC_MMA if mma_ok else _FUSED_HC_FALLBACK_TACTIC_FMA + if mma_ok: + if deterministic: + return ("fused_half_mma", 0, 1, 0, 1) + return _FUSED_HC_FALLBACK_TACTIC_MMA + return _FUSED_HC_FALLBACK_TACTIC_FMA class MhcFusedHcRunner(TunableRunner): @@ -854,6 +860,11 @@ class MhcFusedHcRunner(TunableRunner): Fallback (-1): delegates to the C++ heuristic (fused_half_mma + auto ks + auto bs). + + ``FORCE_DETERMINISTIC=1`` or ``FORCE_MHC_KERNEL_DETERMINISTIC=1`` uses a + fixed half-fused tactic with one K split for tuned and fallback execution, + except that the FMA fallback tactic is kept: ``mhc_fused_hc`` passes it under + ``FORCE_DETERMINISTIC=1`` for its row-count-independent reduction order. """ tuning_config = TuningConfig( @@ -887,6 +898,10 @@ def __init__( ): self.n = n self.hidden_size = hidden_size + self._force_deterministic = ( + os.getenv("FORCE_DETERMINISTIC", "0") == "1" + or os.getenv("FORCE_MHC_KERNEL_DETERMINISTIC", "0") == "1" + ) self.rms_eps = rms_eps self.hc_pre_eps = hc_pre_eps self.hc_sinkhorn_eps = hc_sinkhorn_eps @@ -894,9 +909,13 @@ def __init__( self.sinkhorn_repeat = sinkhorn_repeat def unique_id(self): + if self._force_deterministic: + return (self.n, self.hidden_size, "deterministic") return (self.n, self.hidden_size) def get_valid_tactics(self, inputs, profile: OptimizationProfile, **kwargs): + if self._force_deterministic: + return [_get_fused_hc_fallback_tactic(self.hidden_size, deterministic=True)] M = inputs[0].shape[0] tactics = [] seen = set() @@ -982,8 +1001,15 @@ def forward(self, inputs, *, tactic=-1, **kwargs): hc_base_cur = hc_base_cur.to(torch.float32).contiguous() pre_mix_ext = kwargs.get("pre_mix_ext") - if tactic == -1: - tactic = _get_fused_hc_fallback_tactic(self.hidden_size) + if tactic == -1 or (self._force_deterministic and tactic != _FUSED_HC_FALLBACK_TACTIC_FMA): + # One fixed tactic avoids both atomic split-K reductions and + # arithmetic changes from independently tuned rank/process choices. + # The FMA fallback stays: it has neither, and mhc_fused_hc passes it + # under FORCE_DETERMINISTIC because its reduction order does not + # depend on the number of rows. + tactic = _get_fused_hc_fallback_tactic( + self.hidden_size, deterministic=self._force_deterministic + ) backend, tile_n, num_k_splits, bigfuse_bs, tile_m = tactic backend_code = _FUSED_HC_BACKEND_CODE[backend] diff --git a/tensorrt_llm/_torch/modules/top_k.py b/tensorrt_llm/_torch/modules/top_k.py index 2a2fda1b87d3..4ab65437b64a 100644 --- a/tensorrt_llm/_torch/modules/top_k.py +++ b/tensorrt_llm/_torch/modules/top_k.py @@ -4,6 +4,7 @@ from __future__ import annotations +import os from enum import Enum import torch @@ -36,6 +37,10 @@ class TopK(nn.Module): request metadata and retain a stable address across CUDA Graph replays. Native CUDA scratch is supplied by the caller so its lifetime follows the caller's eager or CUDA-graph metadata. + + ``FORCE_DETERMINISTIC=1`` or ``FORCE_TOPK_KERNEL_DETERMINISTIC=1`` selects + the PyTorch implementation for repeatable indices, including their ordering + and selection among tied scores. """ def __init__( @@ -51,6 +56,13 @@ def __init__( ) -> None: super().__init__() self.top_k = top_k + if ( + os.getenv("FORCE_DETERMINISTIC", "0") == "1" + or os.getenv("FORCE_TOPK_KERNEL_DETERMINISTIC", "0") == "1" + ): + # Native selection uses atomic appends, including for ties at the + # cutoff. Fix both the selected set and its order for repeatability. + prefill_implementation = decode_implementation = TopKImplementation.TORCH self.prefill_implementation = TopKImplementation( prefill_implementation or TopKImplementation.CUDA_RADIX ) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 151c4d84ddcb..45984f2a47d0 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -57,10 +57,10 @@ from . import config_utils from .config_utils import (MambaKVCacheParams, _is_sliding_attention_layer, extract_mamba_kv_cache_params, - get_layer_attention_window, is_hybrid_linear, - is_kimi_linear, is_mla, is_nemotron_hybrid, - is_qwen3_hybrid, is_qwen4_exp, resolve_vocab_size, - uses_vswa_kv_cache_layout) + get_layer_attention_window, is_decode_only, + is_hybrid_linear, is_kimi_linear, is_mla, + is_nemotron_hybrid, is_qwen3_hybrid, is_qwen4_exp, + resolve_vocab_size, uses_vswa_kv_cache_layout) from .connectors.kv_cache_connector import KvCacheConnectorManager from .dwdp import DwdpManager from .engine.lora import make_lora_model_config @@ -1977,10 +1977,12 @@ def configure_kv_cache_capacity(self, # A completed chunk-aware profiling run already prices the cached-KV # staging buffers in the measured peak. Do not subtract an additional # workspace reserve or install an attended-KV admission cap for it. + # A decode-only worker never runs the context-MLA gather at all. profiled_mla_chunks = (py_executor is not None and not self._skip_est and self._mla_chunked_profile_length is not None) - w_bytes_per_token = (0 if profiled_mla_chunks else - get_attention_workspace_bytes_per_token( + w_bytes_per_token = (0 if profiled_mla_chunks + or is_decode_only(self._sparse_attention_config) + else get_attention_workspace_bytes_per_token( self._model_engine.model.model_config, self._mapping)) workspace_is_chunked_prefill_bounded = True diff --git a/tensorrt_llm/_torch/pyexecutor/config_utils.py b/tensorrt_llm/_torch/pyexecutor/config_utils.py index bab8948bd2e8..27920f6aa519 100644 --- a/tensorrt_llm/_torch/pyexecutor/config_utils.py +++ b/tensorrt_llm/_torch/pyexecutor/config_utils.py @@ -10,6 +10,7 @@ from tensorrt_llm._utils import str_dtype_to_torch from tensorrt_llm.llmapi.llm_args import (CacheTransceiverConfig, + SparseAttentionConfig, is_speculators_dspark_config, translate_speculators_dspark_config) from tensorrt_llm.logger import logger @@ -423,6 +424,18 @@ def is_minimax_m3(sparse_attention_config): return sparse_attention_config is not None and sparse_attention_config.algorithm == "minimax_m3" +def is_decode_only( + sparse_attention_config: Optional[SparseAttentionConfig]) -> bool: + """True when the sparse attention config restricts the worker to decode. + + Sparse KV cache offload keeps attention history in host memory and admits + one query token per request per step, so an offloading worker serves the + decode phase of disaggregated serving and never runs a context batch. + """ + return (sparse_attention_config is not None and + sparse_attention_config.get_kv_cache_offload_config() is not None) + + def is_qwen3_next(config): return hasattr( config, 'architectures' diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/graphs.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/graphs.py index a85547af6ea8..a46203b0007c 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/graphs.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/graphs.py @@ -25,6 +25,7 @@ ) from tensorrt_llm._torch.nccl_window_tensor_scope import discard_nccl_window_tensor_outputs from tensorrt_llm._torch.pyexecutor.breakable_cuda_graph_runner import BreakableCUDAGraphRunner +from tensorrt_llm._torch.pyexecutor.config_utils import is_decode_only from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import CUDAGraphRunner, KeyType from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, get_draft_token_length @@ -934,6 +935,12 @@ def _capture_prefill_cuda_graphs(self, resource_manager: ResourceManager) -> Non and not self._config.torch_compile_enabled ): return + if is_decode_only(self.sparse_attention_config): + logger.warning( + "Skipping prefill CUDA graph capture: this worker serves " + "decode only and never runs a context batch." + ) + return logger.info("Running prefill CUDA graph warmup...") prefill_cuda_graph_num_tokens = sorted( diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/warmup.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/warmup.py index 4b5f99651907..113b7519a845 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/warmup.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/warmup.py @@ -21,6 +21,7 @@ from tensorrt_llm._torch.models.modeling_utils import timing_metric from tensorrt_llm._torch.modules.mamba.mamba2_metadata import Mamba2Metadata from tensorrt_llm._torch.nccl_window_tensor_scope import discard_nccl_window_tensor_outputs +from tensorrt_llm._torch.pyexecutor.config_utils import is_decode_only from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManager from tensorrt_llm._torch.pyexecutor.resource_manager import ResourceManager, ResourceManagerType from tensorrt_llm._torch.pyexecutor.sampler.ops.flashinfer import warmup_sampling_module @@ -341,6 +342,20 @@ def _get_full_general_warmup_requests( # Deduplicate the warmup_configs while keeping the order. return list(dict.fromkeys(warmup_configs)) + def _generation_warmup_shapes( + self, warmup_requests_configs: list[tuple[int, int]] + ) -> list[tuple[int, int]]: + """Keeps the warmup shapes whose tokens all belong to generation requests. + + Each shape is ``(num_tokens, num_gen_requests)``. Order is preserved. + """ + tokens_per_gen_request = 1 + self._config.max_total_draft_tokens + return [ + (num_tokens, num_gen_requests) + for num_tokens, num_gen_requests in warmup_requests_configs + if num_tokens == num_gen_requests * tokens_per_gen_request + ] + def _warmup_scheduled( self, resource_manager: ResourceManager, kv_cache_manager ) -> EagerWorkspaceReclaimer | None: @@ -387,6 +402,10 @@ def _warmup_scheduled( and self.guided_decoder is None and not isinstance(kv_cache_manager, MambaHybridCacheManager) ) + # A decode-only worker never runs a context batch, so every phase below + # warms up pure-generation shapes only. The flag comes from the + # configuration, so every rank takes the same branch. + decode_only = is_decode_only(self.sparse_attention_config) log_mem_snapshot("warmup/before_warmup") # Compile the DSv4 indexer-Q CuTe DSL kernels before the first @@ -402,16 +421,21 @@ def _warmup_scheduled( ), self._torch_compile_bypass(), ): - self._run_attention_warmup(resource_manager, can_run_general_warmup) + self._run_attention_warmup( + resource_manager, can_run_general_warmup and not decode_only + ) if can_run_general_warmup: # Specialize torch.compile graphs across the key input shapes before CUDA graph capture. with self._warmup_timer.phase( "general", metrics=self._metrics, metric_name="general_warmup_seconds" ): - warmup_requests_configs = self._agreement.agree_shapes( - self._get_full_general_warmup_requests(resource_manager) - ) + warmup_requests_configs = self._get_full_general_warmup_requests(resource_manager) + if decode_only: + warmup_requests_configs = self._generation_warmup_shapes( + warmup_requests_configs + ) + warmup_requests_configs = self._agreement.agree_shapes(warmup_requests_configs) # Currently graph has not been captured, disable cuda graph for this warmup. with self._graphs.no_cuda_graph(): self._general_warmup(resource_manager, warmup_requests_configs) @@ -444,7 +468,7 @@ def _warmup_scheduled( ), self._torch_compile_bypass(), ): - self._run_autotuner_warmup(resource_manager) + self._run_autotuner_warmup(resource_manager, generation_only=decode_only) log_mem_snapshot("warmup/after_autotuner") # Pre-JIT Mamba SSD multi-seq + HAS_INITSTATES=True Triton kernels # for Mamba hybrid models. Runs regardless of enable_autotuner, @@ -501,6 +525,10 @@ def _warmup_scheduled( self._torch_compile_bypass(), ): warmup_requests_configs = self._get_max_shape_warmup_requests(resource_manager) + if decode_only: + warmup_requests_configs = self._generation_warmup_shapes( + warmup_requests_configs + ) self._general_warmup(resource_manager, warmup_requests_configs) log_mem_snapshot("warmup/after_memory_pool_prepop") 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 fc2164f91cc8..b15f22e167fb 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 @@ -60,9 +60,10 @@ copy_batch_block_offsets_to_device, gather_base_page_rows, ) -from tensorrt_llm.llmapi.llm_args import KvCacheConfig, KVEventsConfig +from tensorrt_llm.llmapi.llm_args import KvCacheConfig, KVEventsConfig, SparseKvCacheOffloadConfig from tensorrt_llm.logger import logger from tensorrt_llm.mapping import Mapping +from tensorrt_llm.math_utils import ceil_div, pad_up from tensorrt_llm.runtime.kv_cache_hash import get_effective_kv_cache_event_hash_algo from tensorrt_llm.runtime.kv_cache_manager_v2 import ( _KV_CACHE_ITERATION_STATS_DELTA_FIELDS, @@ -74,7 +75,6 @@ AttentionLayerConfig, AttentionReusePolicy, AttnLifeCycle, - Batch, BatchDesc, BeamIndex, BufferConfig, @@ -137,6 +137,17 @@ scan_kv_buffer_for_invalid_values, ) from ..scheduler import ScheduledRequests +from .sparse_offload_memory import ( + HostHistoryReservations, + life_cycle_pages, + plan_host_history_bytes, + solve_sparse_offload_capacity, + sparse_arena_bytes, + sparse_arena_shapes, + sparse_batch_shape, + sparse_history_bytes_per_block, + sparse_offload_pool_ratio, +) if TYPE_CHECKING: from tensorrt_llm._torch.attention.backends.interface import AttentionMetadata @@ -498,12 +509,80 @@ def _estimate_full_attn_size_per_token( ) +# Shares of a rank's available host memory and of RLIMIT_MEMLOCK that automatic +# host sizing may pin. +AUTO_HOST_MEMORY_SHARE = 0.5 +AUTO_HOST_MEMLOCK_SHARE = 0.8 + + +class HostMemoryLimits(NamedTuple): + """The host memory a rank may plan pinned allocations against. + + Attributes: + local_ranks: Ranks sharing this node. + mem_available: The node's available memory in bytes; ``inf`` when unknown. + memlock_limit: The RLIMIT_MEMLOCK soft limit in bytes; ``inf`` when unlimited. + """ + + local_ranks: int + mem_available: float + memlock_limit: float + + @property + def rank_share(self) -> float: + """This rank's share of the node's available memory.""" + return self.mem_available / self.local_ranks + + @property + def auto_pinnable_bytes(self) -> float: + """Host bytes automatic sizing may pin on this rank; ``inf`` when unbounded.""" + return min( + self.rank_share * AUTO_HOST_MEMORY_SHARE, + self.memlock_limit * AUTO_HOST_MEMLOCK_SHARE, + ) + + +def get_host_memory_limits(mapping: Mapping) -> HostMemoryLimits: + """Return the host memory a rank may plan pinned allocations against. + + Args: + mapping: The rank's mapping, which gives the ranks sharing its node. + """ + import resource + + local_ranks = max(1, Distributed.get(mapping).local_world_size) + # MemAvailable counts the page cache that pinning reclaims. Weights read at startup + # fill it, so free pages alone can understate usable memory by hundreds of GiB. + mem_available = None + try: + with open("/proc/meminfo") as meminfo: + for line in meminfo: + if line.startswith("MemAvailable:"): + mem_available = int(line.split()[1]) * 1024 + break + except (OSError, IndexError, ValueError): + pass + if mem_available is None: + # Kernels before 3.14 and non-Linux hosts do not publish MemAvailable. + try: + mem_available = os.sysconf("SC_PAGE_SIZE") * os.sysconf("SC_AVPHYS_PAGES") + except (ValueError, OSError): + mem_available = float("inf") + try: + soft_limit, _ = resource.getrlimit(resource.RLIMIT_MEMLOCK) + memlock_limit = soft_limit if soft_limit != resource.RLIM_INFINITY else float("inf") + except (ValueError, OSError): + memlock_limit = float("inf") + return HostMemoryLimits(local_ranks, mem_available, memlock_limit) + + def _compute_auto_host_tier_quota( quota: int, local_ranks: int, mem_available: float, memlock_limit: float, shares_device_memory: bool = False, + pins_offloaded_history: bool = False, ) -> int: """Compute the auto-provisioned host cache tier quota for a single rank. @@ -528,18 +607,23 @@ def _compute_auto_host_tier_quota( shares_device_memory: Whether host and device memory are one physical pool (integrated GPU). Charges ``quota`` against ``mem_available`` and allows a zero result. + pins_offloaded_history: Whether sparse KV cache offload pins host + memory beside the tier. ``mem_available`` and ``memlock_limit`` + already exclude it, so falling back to the device quota could pin + past them; allows a zero result instead. Returns: - Host cache tier quota in bytes. Positive, except on integrated GPUs - with no headroom left, where 0 means "provision no host tier". + Host cache tier quota in bytes. Positive, except on integrated GPUs or + beside offloaded history with no headroom left, where 0 means + "provision no host tier". """ if shares_device_memory and mem_available != float("inf"): mem_available = max(0.0, mem_available - quota) candidates = [quota] if mem_available != float("inf"): - candidates.append(int(mem_available / local_ranks * 0.5)) + candidates.append(int(mem_available / local_ranks * AUTO_HOST_MEMORY_SHARE)) if memlock_limit != float("inf"): - candidates.append(int(memlock_limit * 0.8)) + candidates.append(int(memlock_limit * AUTO_HOST_MEMLOCK_SHARE)) host_quota = min(candidates) if host_quota <= 0: if shares_device_memory: @@ -555,6 +639,8 @@ def _compute_auto_host_tier_quota( "kv_cache_config.host_cache_size." ) return 0 + if pins_offloaded_history: + return 0 logger.warning( f"KV cache manager v2 auto host tier sizing computed a " f"non-positive quota ({host_quota}); falling back to the " @@ -607,6 +693,34 @@ def _sync_kv_cache_manager_init_status( return local_status +def _raise_on_every_rank(error: Optional[Exception], mapping: Mapping, what: str) -> None: + """Raise on every rank if a check failed on any rank, and on none otherwise. + + A check that reads rank-local state, such as free memory, memory limits or + the rank's own layers, can fail on some ranks only. A rank that raised alone + would leave the others waiting in the next collective, so every rank calls + this at the same point with its own outcome. + + Args: + error: This rank's error, or None if its check passed. + mapping: Parallelism mapping; ``world_size`` selects whether a + collective is needed. + what: What was checked, for the error the other ranks raise. + + Raises: + Exception: ``error``, on a rank whose check failed. + RuntimeError: On the other ranks, when any rank's check failed. + """ + failed_rank = mapping.rank if error is not None else mapping.world_size + if mapping.world_size > 1: + # The lowest failing rank, or world_size when every rank passed. + failed_rank = Distributed.get(mapping).allreduce(failed_rank, op=ReduceOp.MIN) + if error is not None: + raise error + if failed_rank < mapping.world_size: + raise RuntimeError(f"{what} failed on rank {failed_rank}; see that rank's log.") + + def _estimate_swa_cache_size( layer_sizes: Sequence[int], attention_windows: Sequence[Optional[int]], @@ -1756,8 +1870,6 @@ def copy_block_offsets_to( class KVCacheManagerV2(BaseResourceManager): - # Sparse managers attach a metadata batch during initialization. - sparse_metadata_batch: Batch | None = None # Filled lazily by _cold_pool_group_membership(); the grouping is fixed after construction. # Declared on the class so it is present even when an instance is built without running __init__. _cold_pool_group_membership_cache: Optional[tuple[tuple[int, frozenset[int]], ...]] = None @@ -1799,6 +1911,25 @@ class KVCacheManagerV2(BaseResourceManager): # need snapshotted (request_prefix_snapshot); kept apart from the request's # own expect_snapshot_points, which other code rewrites. Created on first use. _waiter_snapshot_points: Optional[Dict[int, set[int]]] = None + # Sizing of sparse attention history offloaded to host memory; None without + # offload. Declared on the class so it is present even when an instance is + # built without running __init__. + _sparse_kv_offload: Optional[SparseKvCacheOffloadConfig] = None + # Host history pages admitted requests hold through their last token. + _host_history_reservations: Optional[HostHistoryReservations] = None + # Host history pages the smallest rank's pool has for requests once dummy + # requests are set aside, the most one request can be given; None when no + # rank offloads sparse history. + _host_history_capacity_blocks: Optional[int] = None + # Host bytes the host-backed sparse pools may take; 0 without offload. + _host_history_budget: int = 0 + # Device bytes offload allocates once for max_batch_size rows, outside the + # pools and the arenas, and the pinned host bytes of the Batch's staging + # copy; 0 without offload. + _sparse_batch_buffer_bytes: int = 0 + _sparse_batch_staging_bytes: int = 0 + # The device quota's split when the host budget holds the host-backed life + # cycles to less than their share; None to split by decode bytes. def __init__( self, @@ -1835,6 +1966,7 @@ def __init__( joint_kv_cache_reuse: bool = False, max_cuda_graph_batch_size: Optional[int] = None, layer_specs: Optional[Sequence[KVCacheLayerSpec]] = None, + sparse_kv_offload_config: Optional[SparseKvCacheOffloadConfig] = None, **kwargs, ) -> None: self.mapping = mapping @@ -1857,6 +1989,13 @@ def __init__( self._init_mixed_precision( kv_cache_config, kwargs.get("pretrained_config"), max_seq_len, dtype ) + # A subclass that declares sparse buffers whose history lives in host + # memory passes their sizing; the manager then accounts for that memory. + if sparse_kv_offload_config is not None: + self._check_sparse_kv_offload_config( + sparse_kv_offload_config, kv_cache_config, mapping, spec_config=spec_config + ) + self._sparse_kv_offload = sparse_kv_offload_config self._validate_speculative_config(spec_config) self.is_disagg = is_disagg self.kv_connector_manager = kv_connector_manager @@ -2043,7 +2182,6 @@ def __init__( self._stream = ( execution_stream if execution_stream is not None else torch.cuda.current_stream() ) - self.sparse_metadata_batch: Batch | None = None logger.info(f"[KVCacheManager] execution_stream: {self._stream}") # Materialize an exact per-local-layer vector for cache and attention consumers. @@ -2184,14 +2322,98 @@ def append_to_kv_heads_per_layer( self.is_vswa = uses_vswa_kv_cache_layout(self.max_attention_window_vec) + # With pipeline parallelism, multiple microbatches can be in-flight + # simultaneously, so we need slots for all concurrent sequences. + # Reserve stable slots for persistent request IDs such as CUDA-graph + # padding requests. The default preserves the main-branch allocation. + # In disaggregated mode, use a coefficient of 2: at any moment up to + # `max_num_sequences` requests can be actively generating while another + # up to `max_num_sequences` requests are still in KV transfer + # (TRANS_IN_PROGRESS) and continue to hold their index slots. The 2x + # capacity lets the next batch of active requests acquire slots without + # waiting for the previous batch's transfers to finish. With the overlap + # scheduler on (non-PP), a retiring request holds its index slot one extra + # iteration and needs the same coefficient. Sparse cache leases do not + # follow this count: only requests seated in the attention Batch hold one. + max_num_sequences = max_batch_size * mapping.pp_size + needs_extra_index_slots = is_disagg or ( + not disable_overlap_scheduler and not mapping.has_pp() + ) + assert num_reserved_index_slots >= 0, "num_reserved_index_slots must be non-negative" + # Both diagnostics below are off unless their environment variable is + # set; with neither set nothing here changes any allocation, any page + # content or any reported count. + self._guard_page_fill = os.environ.get("TRTLLM_KV_GUARD_PAGE", "").strip().lower() + self._guard_page_value = _parse_kv_fill_value(self._guard_page_fill, "TRTLLM_KV_GUARD_PAGE") + self._guard_page_by_layer: dict[int, int] = {} + self._fresh_page_fill = _parse_kv_fill_value( + os.environ.get("TRTLLM_KV_FRESH_PAGE_FILL", "").strip().lower(), + "TRTLLM_KV_FRESH_PAGE_FILL", + ) + # request id -> pool id -> the pool's last-seen base page-index array + # (indexed by block ordinal). Comparing it against the current array is + # both the early-out (unchanged -> nothing handed out since the previous + # step) and the per-ordinal freshness key (see ``_fill_fresh_kv_pages``). + self._fresh_pages_filled: Dict[int, Dict[int, np.ndarray]] = {} + self._fresh_fill_announced = False + self._fresh_fill_unavailable_announced = False + self.max_admissible_sequences = max_num_sequences * (2 if needs_extra_index_slots else 1) + # The guard page is held by a permanent sequence, so it needs an index + # slot of its own. Taking one of the scheduler's would change which + # requests get admitted, so a run with the diagnostic on would no + # longer be comparable with the run it is being read against. + index_mapper_capacity = ( + self.max_admissible_sequences + + num_reserved_index_slots + + (1 if self._guard_page_value is not None else 0) + ) + max_util_for_resume = kv_cache_config.max_util_for_resume quota = sys.maxsize + # With sparse KV cache offload, the sparse cache arenas and the buffers + # sized for max_batch_size rows are allocated beside the pools, outside + # the device tier's quota, so they come out of the device byte budget. + # A lease is held only while a cache is seated in the attention Batch, so + # the arena needs one per row rather than one per admitted sequence. + reserved_arena_bytes = self._get_sparse_arena_bytes( + kv_cache_config, tokens_per_block=tokens_per_block, max_leases=self.max_batch_size + ) + self._sparse_batch_buffer_bytes, self._sparse_batch_staging_bytes = ( + self._get_sparse_batch_buffer_bytes( + kv_cache_config, tokens_per_block=tokens_per_block, max_leases=self.max_batch_size + ) + ) + reserved_bytes = reserved_arena_bytes + self._sparse_batch_buffer_bytes + offload_error: Optional[ValueError] = None if ( kv_cache_config.max_gpu_total_bytes is not None and kv_cache_config.max_gpu_total_bytes > 0 ): quota = int(kv_cache_config.max_gpu_total_bytes) logger.info(f"max_gpu_total_bytes is provided. New quota is {quota / (1 << 30)}GiB") + + # The arenas and buffers sit beside the pools, so they come out of the device + # memory budget. max_tokens below sizes the pools alone and does not pay for them. + if reserved_bytes > 0 and quota < sys.maxsize: + reserved = ( + f"{reserved_arena_bytes / (1 << 30):.2f}GiB of sparse cache arena for " + f"{self.max_batch_size} leases and " + f"{self._sparse_batch_buffer_bytes / (1 << 30):.2f}GiB of offload buffers for " + f"max_batch_size={self.max_batch_size}" + ) + if reserved_bytes >= quota: + # Raised below, once every rank has sized its offload. + offload_error = ValueError( + f"Sparse KV cache offload needs {reserved}, more than the " + f"{quota / (1 << 30):.2f}GiB device KV budget. Lower max_batch_size or " + "the sparse cache size, or raise the device KV budget." + ) + else: + quota -= reserved_bytes + logger.info( + f"Reserved {reserved} from the device KV budget. " + f"New quota is {quota / (1 << 30)}GiB" + ) if kv_cache_config.max_tokens is not None: quota_from_max_tokens = int( math.ceil( @@ -2209,6 +2431,20 @@ def append_to_kv_heads_per_layer( assert quota < sys.maxsize, ( "Quota not set. Check kv_cache_config.max_tokens or kv_cache_config.max_gpu_total_bytes" ) + if self._sparse_kv_offload is not None: + if offload_error is None: + try: + self._plan_host_history_budget( + kv_cache_config, tokens_per_block=tokens_per_block + ) + except ValueError as error: + offload_error = error + # Each rank sizes offload from its own device budget, host memory and + # layers, so sizing can fail on some ranks only. A rank that raised + # alone would leave the rest waiting in the token sync below. + _raise_on_every_rank(offload_error, mapping, "Sizing sparse KV cache offload") + elif offload_error is not None: + raise offload_error # Sync resumable token capacity across ranks so the scheduler produces # identical batches. Normalize to tokens because cache costs vary @@ -2263,31 +2499,26 @@ def append_to_kv_heads_per_layer( # Automatically provision a host tier matching the GPU quota so # suspend/resume works out of the box. Cap at available host # memory and pinnable memory limit to avoid allocation failures. - import resource - + # # Rank-aware host tier sizing: divide the per-node memory budget # by the number of ranks sharing this physical node. Each rank # independently provisions a host tier; without this, N # co-located ranks each reserving a device-quota-sized block can # OOM the host (observed on GB300 NVL72 with 4 ranks/node and # ~170GiB device quota each on a 975GiB node). - local_ranks = max(1, Distributed.get(mapping).local_world_size) - - try: - mem_available = os.sysconf("SC_PAGE_SIZE") * os.sysconf("SC_AVPHYS_PAGES") - except (ValueError, OSError): - mem_available = float("inf") - try: - _soft, _hard = resource.getrlimit(resource.RLIMIT_MEMLOCK) - memlock_limit = _soft if _soft != resource.RLIM_INFINITY else float("inf") - except (ValueError, OSError): - memlock_limit = float("inf") + local_ranks, mem_available, memlock_limit = get_host_memory_limits(mapping) + # Every rank on the node also pins its host-backed sparse pools and + # its sparse Batch's staging copy. + offload_pinned_bytes = self._host_history_budget + self._sparse_batch_staging_bytes + mem_available = max(0.0, mem_available - local_ranks * offload_pinned_bytes) + memlock_limit = max(0.0, memlock_limit - offload_pinned_bytes) host_quota = _compute_auto_host_tier_quota( quota, local_ranks, mem_available, memlock_limit, shares_device_memory=is_device_integrated(), + pins_offloaded_history=self._sparse_kv_offload is not None, ) logger.info( f"KV cache manager v2 auto host tier sizing: " @@ -2297,6 +2528,22 @@ def append_to_kv_heads_per_layer( # Reduce the rank-local auto quota to the fleet minimum; see # _sync_host_tier_quota for why divergence wedges collectives. host_quota = _sync_host_tier_quota(host_quota, mapping) + disk_cache_size = kv_cache_config.disk_cache_size + if ( + host_quota <= 0 + and self._sparse_kv_offload is not None + and not (disk_cache_size is not None and disk_cache_size > 0) + ): + # Every rank holds the same synced quota, so every rank raises. + raise ValueError( + "Sparse KV cache offload pins " + f"{offload_pinned_bytes / (1 << 30):.2f}GiB of host memory on this rank, " + "which leaves none for the host cache tier. Without a tier to evict to, " + "suspending a request cannot free its KV cache pages, and a decode-only " + "worker cannot recompute them, so scheduling can stall. Lower " + "sparse_attention_config.kv_cache_offload_config.host_history_size, or set " + "kv_cache_config.host_cache_size (0 runs without a host tier)." + ) if host_quota > 0: cache_tiers.append(HostCacheTierConfig(quota=int(host_quota))) logger.info( @@ -2315,14 +2562,20 @@ def append_to_kv_heads_per_layer( self.vocab_size = vocab_size - config = self._build_base_config( + config = self._build_manager_config( kv_cache_config, tokens_per_block=tokens_per_block, cache_tiers=cache_tiers, ) - config = self._build_cache_config(config) - config = self._remove_zero_size_buffers(config) config = self._add_mixed_precision_layers(config) + arena_bytes = sparse_arena_bytes(config, self.max_batch_size) + if arena_bytes > reserved_arena_bytes: + raise RuntimeError( + f"The sparse cache arenas take {arena_bytes} bytes of device memory, but only " + f"{reserved_arena_bytes} bytes were reserved for them. A manager whose layers " + "declare a sparse_cache must pass sparse_kv_offload_config, and its layer " + "configs must not depend on the cache tier quotas." + ) has_host_cache_tier = any( isinstance(tier, HostCacheTierConfig) for tier in config.cache_tiers ) @@ -2394,6 +2647,13 @@ def create_cold_page_codec(cache_config: object) -> Optional[object]: self.kv_cache_manager_py_config = config self.impl = candidate self.can_evict = len(config.cache_tiers) > 1 + if self._sparse_kv_offload is not None and not self.can_evict: + logger.warning( + "Sparse KV cache offload runs without a host or disk cache tier: suspending a " + "request cannot free its KV cache pages, and a decode-only worker cannot " + "recompute them, so scheduling can stall once the device KV cache is full. Set " + "kv_cache_config.host_cache_size to give evicted pages a host tier." + ) if self.event_manager is not None: self.event_manager.set_layer_group_window_sizes( self._get_event_window_sizes_by_layer_group( @@ -2473,14 +2733,7 @@ def create_cold_page_codec(cache_config: object) -> Optional[object]: # stay int. self.max_seq_len = int(max_num_tokens) - # Pad max_blocks_per_seq to next multiple of 4 (copy_block_offsets kernel). - # Include the same maximum generation lead used by allocation and sizing. - max_seq_capacity = self.max_seq_len + self._generation_kv_capacity_headroom - self.max_blocks_per_seq = ( - max_seq_capacity + self._ledger_tokens_per_block - 1 - ) // self._ledger_tokens_per_block - if self.max_blocks_per_seq % 4 != 0: - self.max_blocks_per_seq = ((self.max_blocks_per_seq + 3) // 4) * 4 + self.max_blocks_per_seq = self._get_max_blocks_per_seq(self.max_seq_len) self.enable_block_reuse = kv_cache_config.enable_block_reuse self.disk_prefetch_num_reqs = kv_cache_config.disk_prefetch_num_reqs @@ -2519,48 +2772,6 @@ def create_cold_page_codec(cache_config: object) -> Optional[object]: ) self.can_share_inflight_prefix = self._can_share_inflight_prefix() - # With pipeline parallelism, multiple microbatches can be in-flight - # simultaneously, so we need slots for all concurrent sequences. - # Reserve stable slots for persistent request IDs such as CUDA-graph - # padding requests. The default preserves the main-branch allocation. - # In disaggregated mode, use a coefficient of 2: at any moment up to - # `max_num_sequences` requests can be actively generating while another - # up to `max_num_sequences` requests are still in KV transfer - # (TRANS_IN_PROGRESS) and continue to hold their index slots. The 2x - # capacity lets the next batch of active requests acquire slots without - # waiting for the previous batch's transfers to finish. With the overlap - # scheduler on (non-PP), a retiring request holds its lease one extra - # iteration and needs the same coefficient. - max_num_sequences = max_batch_size * mapping.pp_size - needs_extra_leases = is_disagg or (not disable_overlap_scheduler and not mapping.has_pp()) - assert num_reserved_index_slots >= 0, "num_reserved_index_slots must be non-negative" - # Both diagnostics below are off unless their environment variable is - # set; with neither set nothing here changes any allocation, any page - # content or any reported count. - self._guard_page_fill = os.environ.get("TRTLLM_KV_GUARD_PAGE", "").strip().lower() - self._guard_page_value = _parse_kv_fill_value(self._guard_page_fill, "TRTLLM_KV_GUARD_PAGE") - self._guard_page_by_layer: dict[int, int] = {} - self._fresh_page_fill = _parse_kv_fill_value( - os.environ.get("TRTLLM_KV_FRESH_PAGE_FILL", "").strip().lower(), - "TRTLLM_KV_FRESH_PAGE_FILL", - ) - # request id -> pool id -> the pool's last-seen base page-index array - # (indexed by block ordinal). Comparing it against the current array is - # both the early-out (unchanged -> nothing handed out since the previous - # step) and the per-ordinal freshness key (see ``_fill_fresh_kv_pages``). - self._fresh_pages_filled: Dict[int, Dict[int, np.ndarray]] = {} - self._fresh_fill_announced = False - self._fresh_fill_unavailable_announced = False - self.max_admissible_sequences = max_num_sequences * (2 if needs_extra_leases else 1) - # The guard page is held by a permanent sequence, so it needs an index - # slot of its own. Taking one of the scheduler's would change which - # requests get admitted, so a run with the diagnostic on would no - # longer be comparable with the run it is being read against. - index_mapper_capacity = ( - self.max_admissible_sequences - + num_reserved_index_slots - + (1 if self._guard_page_value is not None else 0) - ) logger.info( f"KVCacheManagerV2: IndexMapper capacity={index_mapper_capacity} " f"(max_num_sequences={max_num_sequences}, is_disagg={is_disagg}, " @@ -2579,24 +2790,11 @@ def create_cold_page_codec(cache_config: object) -> Optional[object]: self._pending_preemption: Dict[int, LlmRequest] = {} self._use_per_layer_page_tables = False self._prepare_page_table_tensor(index_mapper_capacity) - self._sparse_layer_group_ids = tuple( - sorted( - { - self.layer_to_pool_mapping_dict[buf.layer_id] - for buf in self.impl.all_buffer_ids - if self.impl.is_sparse(buf.layer_id, buf.role) - } - ) - ) - # 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() + if self._sparse_kv_offload is not None: + self._setup_sparse_kv_offload(kv_cache_config, max_seq_len, num_reserved_index_slots) # Last: bind the publisher socket and start its thread only once every check # above has passed. Constructing the manager is side-effect free, so a failure @@ -3108,11 +3306,22 @@ def _prepare_page_table_tensor(self, index_mapper_capacity: int) -> None: layer_id = self._pool_layer_ids_by_role[(pool_id, role_a)] self.index_scales[pool_id] = self.impl.get_page_index_scale(layer_id, role_a) if role_b is not None: - self.kv_offset[pool_id] = exact_div( + offset = exact_div( self.impl.get_mem_pool_base_address(layer_id, role_b, PageIndexMode.SHARED) - self.impl.get_mem_pool_base_address(layer_id, role_a, PageIndexMode.SHARED), self.impl.get_page_stride(layer_id, role_a), ) + # This page table reaches the second role by its distance from the first, + # which holds only while both sit in one medium. A host-backed pool beside + # a device one is further apart than a page index can say, and the managers + # whose layouts address them directly override this method. + if not -(2**31) <= offset <= 2**31 - 1: + raise ValueError( + f"pool {pool_id} puts {role_b} {offset} pages from {role_a}, which a " + f"page index cannot carry; a layer whose roles live in different media " + f"needs a manager that addresses them directly" + ) + self.kv_offset[pool_id] = offset else: self.kv_offset[pool_id] = 0 # Plain-int mirror of index_scales so the per-step block-table build @@ -3476,12 +3685,10 @@ def _prepare_swa_scratch_copy_tensors(self, index_mapper_capacity: int) -> None: device=device, ) - min_scale = int(scales.min().item()) if scales.numel() > 0 else 1 - max_scratch_pages = int(scratch_pages.max().item()) if scratch_pages.numel() > 0 else 1 - self._max_scratch_slots = max( - 1, - (self.max_blocks_per_seq * max_scratch_pages + min_scale - 1) // min_scale, - ) + # Asked rather than derived: a block still in the window, or held as a sink, is + # never stale and so never scratch, and that reasoning belongs with the window and + # rewind margins the manager already holds. + self._max_scratch_slots = max(1, self.impl.max_scratch_slots(self.max_blocks_per_seq)) scratch_slots_shape = ( self.num_pools, staging_capacity, @@ -3543,17 +3750,18 @@ def _copy_scratch_metadata_to_device( for pool_id in range(self.num_pools): for context_idx, request_id in enumerate(request_ids[:num_contexts]): - desc = self.kv_cache_map[request_id].get_scratch_desc(pool_id) - if desc is None: + cache = self.kv_cache_map[request_id] + scratch_range = cache.get_swa_scratch_range(pool_id) + if not scratch_range: continue - slot_ids = desc.slot_ids + slot_ids = cache.get_scratch_slots(pool_id) if len(slot_ids) > self._max_scratch_slots: raise RuntimeError( f"Scratch slot count {len(slot_ids)} exceeds staging capacity " f"{self._max_scratch_slots}" ) - host_begs[pool_id, context_idx] = int(desc.range.beg) - host_ends[pool_id, context_idx] = int(desc.range.end) + host_begs[pool_id, context_idx] = int(scratch_range.beg) + host_ends[pool_id, context_idx] = int(scratch_range.end) for slot_idx, slot_id in enumerate(slot_ids): host_slots[pool_id, context_idx, slot_idx] = int(slot_id) @@ -3686,17 +3894,20 @@ def _build_base_config( self.max_num_tokens if self.max_num_tokens is not None else typical_seq_len ) + self.num_extra_kv_tokens generation_history_length = max(0, typical_seq_len - self.max_draft_len - 1) - typical_step = BatchDesc( - [KVCacheDesc(capacity=context_capacity, history_length=0)] - + [ - KVCacheDesc( - capacity=typical_seq_len, - history_length=generation_history_length, - beam_width=beam_width, - ) - ] - * (generation_request_capacity - 1) + generation = KVCacheDesc( + capacity=typical_seq_len, + history_length=generation_history_length, + beam_width=beam_width, ) + if self._sparse_kv_offload is not None: + # A worker that offloads sparse history only decodes, so every + # resident request is a generation request. + typical_step = BatchDesc([generation] * generation_request_capacity) + else: + typical_step = BatchDesc( + [KVCacheDesc(capacity=context_capacity, history_length=0)] + + [generation] * (generation_request_capacity - 1) + ) # General and chunked-prefill warmup uses one fresh context request # at the per-iteration token budget. @@ -3878,6 +4089,150 @@ def _build_cache_config(self, config: KVCacheManagerConfigPy) -> KVCacheManagerC """Customize the general cache config for a specialized cache manager.""" return config + def _build_manager_config( + self, + kv_cache_config: KvCacheConfig, + *, + tokens_per_block: int, + cache_tiers: List[CacheTierConfig], + ) -> KVCacheManagerConfigPy: + """Build the config the native manager is constructed from. + + Args: + kv_cache_config: The user's KV cache config. + tokens_per_block: Tokens per block. + cache_tiers: The cache tiers and their quotas. + """ + config = self._build_base_config( + kv_cache_config, + tokens_per_block=tokens_per_block, + cache_tiers=cache_tiers, + ) + config = self._build_cache_config(config) + config = self._remove_zero_size_buffers(config) + if self._sparse_kv_offload is not None and self._host_history_budget > 0: + # The manager divides the device quota and this budget together, holding the + # host-backed pools inside it rather than being handed a quota already lowered. + config = replace(config, host_sparse_pool_quota=self._host_history_budget) + if self._sparse_kv_offload is not None and config.initial_pool_ratio is None: + # Split the quota the way decoding requests use it, so the host-backed + # pools it sizes hold as many requests as the device-backed ones. + config = replace( + config, + initial_pool_ratio=sparse_offload_pool_ratio( + life_cycle_pages(config.layers, config.tokens_per_block), + capacity=self._get_sparse_offload_seq_len(kv_cache_config), + capacity_headroom=self._generation_kv_capacity_headroom, + tokens_per_block=config.tokens_per_block, + ), + ) + return config + + def _get_sparse_arena_bytes( + self, kv_cache_config: KvCacheConfig, *, tokens_per_block: int, max_leases: int + ) -> int: + """Return the device bytes of the sparse cache arenas the manager will allocate. + + The device quota must leave room for the arenas before the cache config + is built. Layer configs do not depend on the tier quotas, so a config + built with a placeholder quota has the same arenas as the real one. + """ + if self._sparse_kv_offload is None: + return 0 + config = self._build_manager_config( + kv_cache_config, + tokens_per_block=tokens_per_block, + cache_tiers=[GpuCacheTierConfig(quota=sys.maxsize)], + ) + return sparse_arena_bytes(config, max_leases) + + def _get_sparse_batch_buffer_bytes( + self, kv_cache_config: KvCacheConfig, *, tokens_per_block: int, max_leases: int + ) -> tuple[int, int]: + """Return the bytes offload allocates once for ``max_batch_size`` rows. + + These are the native Batch the sparse fetch reads and the model's offload + scratch, each sized for every row whatever a step holds. The Batch sizes + its page tables for the configured ``max_seq_len``, an upper bound on the + length the manager settles on. + + Returns: + The device bytes, and the pinned host bytes of the Batch's staging copy. + """ + if self._sparse_kv_offload is None: + return 0, 0 + config = self._build_manager_config( + kv_cache_config, + tokens_per_block=tokens_per_block, + cache_tiers=[GpuCacheTierConfig(quota=sys.maxsize)], + ) + batch = sparse_batch_shape( + config.layers, + tokens_per_block=config.tokens_per_block, + max_batch_size=self.max_batch_size, + max_blocks_per_seq=self._get_max_blocks_per_seq(self.max_seq_len), + ) + device_bytes = batch.device_bytes + self.get_sparse_offload_scratch_bytes(config) + return device_bytes, batch.staging_bytes + + def get_sparse_offload_scratch_bytes(self, config: KVCacheManagerConfigPy) -> int: + """Return the device bytes the model allocates for sparse KV cache offload. + + These are buffers outside the cache manager sized for ``max_batch_size`` + rows, such as the attention metadata's selection, fetch and + dequantization scratch, counted once for every copy the model holds. A + manager whose model allocates any overrides this. + + Args: + config: The manager config, built before the device quota is known. + """ + return 0 + + def _get_max_blocks_per_seq(self, max_seq_len: int) -> int: + """Return the page-table width for requests of up to ``max_seq_len`` tokens. + + It includes the generation lead that allocation and sizing reserve, and + is padded to a multiple of 4 for the copy_block_offsets kernel. + """ + max_seq_capacity = max_seq_len + self._generation_kv_capacity_headroom + return pad_up(ceil_div(max_seq_capacity, self._ledger_tokens_per_block), 4) + + def _get_sparse_offload_seq_len(self, kv_cache_config: KvCacheConfig) -> int: + """Return the request length sparse KV cache offload is sized for.""" + typical_seq_len = self._get_typical_seq_len(kv_cache_config) + return typical_seq_len if typical_seq_len is not None else self.max_seq_len + + def _plan_host_history_budget( + self, kv_cache_config: KvCacheConfig, *, tokens_per_block: int + ) -> None: + """Settle the host memory the offloaded history may pin. + + The manager is given this budget and holds its host-backed pools inside it while + dividing the device quota, so nothing is lowered here: a bound applied before the + manager and again inside it would size the level twice. + """ + config = self._build_manager_config( + kv_cache_config, + tokens_per_block=tokens_per_block, + cache_tiers=[GpuCacheTierConfig(quota=sys.maxsize)], + ) + num_tokens = self._get_sparse_offload_seq_len(kv_cache_config) + # An explicit host cache tier is pinned too, beside the offloaded history. + host_tier_bytes = max(0, kv_cache_config.host_cache_size or 0) + self._host_history_budget = plan_host_history_bytes( + host_history_size=self._sparse_kv_offload.host_history_size, + max_batch_size=self.max_batch_size, + host_bytes_per_request=sparse_history_bytes_per_block(config.layers) + * ceil_div(num_tokens, tokens_per_block), + pinnable_host_bytes=get_host_memory_limits(self.mapping).auto_pinnable_bytes + - host_tier_bytes, + staging_bytes=self._sparse_batch_staging_bytes, + ) + logger.info( + f"Sparse KV cache offload may pin {self._host_history_budget / (1 << 30):.2f}GiB of " + "host memory for its history; the cache manager holds its host-backed pools to it." + ) + def _remove_zero_size_buffers(self, config: KVCacheManagerConfigPy) -> KVCacheManagerConfigPy: """Exclude empty buffers before creating the runtime storage pools.""" if config.layers and all( @@ -3903,6 +4258,10 @@ def _remove_zero_size_buffers(self, config: KVCacheManagerConfigPy) -> KVCacheMa sliding_window_size=layer.sliding_window_size, num_sink_tokens=layer.num_sink_tokens, reuse_policy=layer.reuse_policy, + # The fetch path stays while a sparse buffer remains to use it. + sparse_cache=layer.sparse_cache + if any(buffer.is_sparse for buffer in buffers) + else None, ) ) else: @@ -4730,11 +5089,9 @@ def try_allocate_generation(self, req: LlmRequest) -> bool: return False if not kv_cache.is_active: - if not kv_cache.resume(self._stream.cuda_stream, is_decoding=True): + if not kv_cache.resume(self._stream.cuda_stream): return False self._restore_page_index_bufs(req.py_request_id, kv_cache) - elif not kv_cache.enter_decode(): - return False if not self._ensure_generation_beam_width(req, kv_cache): return False @@ -4835,6 +5192,12 @@ def revert_allocate_context(self, req: LlmRequest) -> bool: def _set_page_index_bufs(self, request_id: int, kv_cache: _KVCache) -> None: assert kv_cache.beam_width <= self.max_beam_width + if kv_cache.batch_row is not None: + # A member's page indices belong to its Batch row; binding a second buffer + # would take the writes with it and strand that row at whatever it held when + # the cache joined. Adapters that need separate host tables must snapshot + # the Batch's rows before those tables are consumed. + return index = self.index_mapper.get_index(request_id) for beam_idx in range(int(kv_cache.beam_width)): for pool_idx in range(self.num_pools): @@ -4880,10 +5243,6 @@ def _restore_page_index_bufs(self, request_id: int, kv_cache: _KVCache) -> None: """ self._set_page_index_bufs(request_id, kv_cache) - if self.sparse_metadata_batch is not None: - index = self.index_mapper.get_index(request_id) - self.sparse_metadata_batch.add(kv_cache, index) - def _resume_and_restore(self, req_id: int, kv_cache, optional_reuse_groups=None) -> bool: """Resume a suspended KV cache and restore its page index buffers. @@ -5284,7 +5643,38 @@ def prepare_disagg_gen_init(self, req: LlmRequest) -> bool: ``kv_cache.history_length`` to ``prompt_len``. Returns True on success, False if preparation or resize failed (cache is suspended on resize failure). + + With sparse KV cache offload, the request is admitted only if the host + history pool holds its history through its last output token, so no + running request runs out of host pages. A request larger than the pool + is rejected when it arrives; see ``get_request_kv_block_budget``. """ + reservations = self._host_history_reservations + if reservations is None: + return self._prepare_disagg_gen_init_cache(req) + num_blocks = self._get_host_history_blocks(req) + if not reservations.can_reserve(req.py_request_id, num_blocks): + if num_blocks > reservations.capacity_blocks: + logger.warning( + f"Request {req.py_request_id} needs {num_blocks} blocks of host history, " + f"more than the {reservations.capacity_blocks} the host history pool " + "holds; it cannot be admitted." + ) + return False + if not self._prepare_disagg_gen_init_cache(req): + return False + reservations.reserve(req.py_request_id, num_blocks) + return True + + def _get_host_history_blocks(self, req: LlmRequest) -> int: + """Return the host history pages ``req`` reaches by its last output token.""" + return ceil_div( + req.prompt_len + max(0, req.max_new_tokens) + self.num_extra_kv_tokens, + self.tokens_per_block, + ) + + def _prepare_disagg_gen_init_cache(self, req: LlmRequest) -> bool: + """Allocate a disagg generation init request's cache; see ``prepare_disagg_gen_init``.""" reused = self.prepare_context_cache(req) if reused is None: return False @@ -5597,7 +5987,6 @@ def prepare_resources(self, scheduled_batch: ScheduledRequests): # Mirror the main manager. Under one-model spec decoding the # scheduler may already have created the context cache. self._prepare_draft_resources(scheduled_batch) - self._publish_sparse_metadata() return # KV pages are allocated in `KVCacheV2Scheduler`, so by this point every @@ -5608,18 +5997,6 @@ def prepare_resources(self, scheduled_batch: ScheduledRequests): # would cost a page index gather per context request for nothing. if self.kv_connector_manager is not None and not self.kv_connector_manager.capacity_only: self._run_kv_connector_hooks(scheduled_batch) - self._publish_sparse_metadata() - - def _publish_sparse_metadata(self) -> None: - """Refresh stable GPU rows on the execution stream before model work. - - Batch rows match IndexMapper slots, including holes. Sparse consumers use - ``sparse_metadata_batch`` for raw tables and eligible-history counts. - Readers on another stream must use Batch.wait_ready/record_read. - """ - if self.sparse_metadata_batch is not None: - self.sparse_metadata_batch.record_read(self._stream.cuda_stream) - self.sparse_metadata_batch.publish(self._stream.cuda_stream) def _run_kv_connector_hooks(self, scheduled_batch: ScheduledRequests) -> None: """Serve final-batch queries for connectors without source reservations.""" @@ -5676,8 +6053,6 @@ def report_batch_to_connector( self._accept_connector_prefix_reservations(scheduled_batch) self._adopt_allocation_stream(scheduled_batch) self.kv_connector_manager.build_scheduler_output(scheduled_batch, self) - # Connector acceptance can resize requests after prepare_resources. - self._publish_sparse_metadata() # ---- KV connector prefix ---- @@ -6621,6 +6996,267 @@ def get_primary_block_counts(self) -> tuple[int, int]: sum(pool.total for pool in stats), ) + @staticmethod + def _check_sparse_kv_offload_config( + offload_config: SparseKvCacheOffloadConfig, + kv_cache_config: KvCacheConfig, + mapping: Mapping, + *, + spec_config: Optional["DecodingBaseConfig"], + ) -> None: + """Reject sparse KV cache offload settings any rank cannot honor. + + Host memory differs by node, so every rank calls this at the same point, + and all of them raise when any rank cannot honor the settings. + + Raises: + ValueError: KV pool rebalancing or speculative decoding is enabled, or + the host history size, with an explicit host cache tier, exceeds + this rank's share of available host memory or the pinned-memory + limit. + RuntimeError: Another rank cannot honor the settings. + """ + error: Optional[ValueError] = None + host_history_size = offload_config.host_history_size + if kv_cache_config.enable_kv_pool_rebalance: + error = ValueError( + "Sparse KV cache offload is incompatible with " + "kv_cache_config.enable_kv_pool_rebalance=True: host-backed history pools " + "take their pool group's slot count, so a rebalance would resize pinned host " + "memory outside the host budget." + ) + elif spec_config is not None: + error = ValueError( + "Sparse KV cache offload does not support speculative decoding: a sparse cache " + "takes one input token per decode step. Remove speculative_config or disable " + "sparse_attention_config.enable_kv_cache_offload." + ) + elif kv_cache_config.enable_block_reuse: + error = ValueError( + "Sparse KV cache offload does not support block reuse: offloaded history lives " + "in host-backed pools that reuse cannot share. Set " + "kv_cache_config.enable_block_reuse=False." + ) + elif host_history_size is not None: + limits = get_host_memory_limits(mapping) + # An explicit host cache tier is pinned beside the offloaded history; + # an automatic one is sized from what the history leaves. + host_tier_size = max(0, kv_cache_config.host_cache_size or 0) + pinned = ( + f"The KV cache offload host_history_size ({host_history_size / (1 << 30):.2f}GiB)" + ) + if host_tier_size: + pinned += ( + f" plus kv_cache_config.host_cache_size ({host_tier_size / (1 << 30):.2f}GiB)" + ) + if host_history_size + host_tier_size > limits.rank_share: + error = ValueError( + f"{pinned} exceeds this rank's {limits.rank_share / (1 << 30):.2f}GiB share of " + f"available host memory ({limits.local_ranks} rank(s) on this node)." + ) + elif host_history_size + host_tier_size > limits.memlock_limit: + error = ValueError( + f"{pinned} exceeds the {limits.memlock_limit / (1 << 30):.2f}GiB " + "RLIMIT_MEMLOCK limit." + ) + _raise_on_every_rank(error, mapping, "The sparse KV cache offload config check") + + def _setup_sparse_kv_offload( + self, kv_cache_config: KvCacheConfig, max_seq_len: int, num_reserved_index_slots: int + ) -> None: + """Check the host-backed pools, start host admission and report the capacity. + + A rank whose layers have no sparse buffers holds no host history. Every + rank calls this at the same point. + + Args: + kv_cache_config: The KV cache configuration. + max_seq_len: The longest request the executor accepts, before this + manager clamps ``max_seq_len`` to what its pools hold. + num_reserved_index_slots: Index slots reserved for persistent dummy + requests, such as CUDA-graph padding. + + Raises: + ValueError: This rank's host-backed pools exceed their host budget. + RuntimeError: Another rank's do. + """ + host_pool_groups = self.get_host_backed_pool_groups() + host_pool_bytes = self._get_host_backed_pool_bytes() if host_pool_groups else 0 + sizing_error: Optional[ValueError] = None + if host_pool_bytes > self._host_history_budget: + # The quota leaves room for how the native manager rounds its split, + # so only its slot floors can take the pools past the budget. + sizing_error = ValueError( + f"The host-backed sparse history pools take {host_pool_bytes / (1 << 30):.2f}GiB, " + f"more than their {self._host_history_budget / (1 << 30):.2f}GiB host budget: the " + "KV cache manager keeps a minimum number of slots in every pool group, and here " + "that minimum needs more host memory than the budget holds. Raise " + "sparse_attention_config.kv_cache_offload_config.host_history_size." + ) + _raise_on_every_rank(sizing_error, self.mapping, "Sizing the host-backed history pools") + local_capacity_blocks = sys.maxsize + if host_pool_groups: + # Dummy requests hold history pages without a reservation, so + # admission keeps their pages out of what it can promise. + local_capacity_blocks = max( + 0, + min(int(pool_group.num_slots) for pool_group in host_pool_groups) + - self._get_dummy_history_blocks(num_reserved_index_slots), + ) + # The executor rejects a request on arrival when its history can never + # fit (get_request_kv_block_budget). Every rank must reject the same + # requests, so all of them judge against the smallest pool. + capacity_blocks = local_capacity_blocks + if self.mapping.world_size > 1: + capacity_blocks = Distributed.get(self.mapping).allreduce( + capacity_blocks, op=ReduceOp.MIN + ) + if capacity_blocks < sys.maxsize: + self._host_history_capacity_blocks = capacity_blocks + if not self.is_estimating_kv_cache: + self._warn_if_host_history_rejects_requests(capacity_blocks, max_seq_len) + if not host_pool_groups: + return + self._host_history_reservations = HostHistoryReservations(local_capacity_blocks) + if not self.is_estimating_kv_cache: + self._log_sparse_kv_offload_capacity(kv_cache_config) + + def _get_dummy_history_blocks(self, num_reserved_index_slots: int) -> int: + """Return the history pages dummy requests may hold at once, outside admission. + + These are the CUDA-graph padding dummies, one per reserved index slot, + an attention-DP dummy when attention DP is on, and the guard page when + it is held. Each holds one token plus its draft and extra tokens. + Offload rejects speculative decoding (``_check_sparse_kv_offload_config``), + so padding uses one draft length and needs one dummy; the count relies + on that. + """ + num_dummies = ( + num_reserved_index_slots + + int(self.mapping.enable_attention_dp) + + int(_GUARD_PAGE_REQUEST_ID in self.kv_cache_map) + ) + tokens_per_dummy = 1 + 2 * self.max_total_draft_tokens + self.num_extra_kv_tokens + return num_dummies * ceil_div(tokens_per_dummy, self.tokens_per_block) + + def _warn_if_host_history_rejects_requests( + self, capacity_blocks: int, max_seq_len: int + ) -> None: + """Warn when the smallest host history pool cannot hold a ``max_seq_len`` request. + + Requests of up to ``max_seq_len`` tokens reach the executor, which + rejects on arrival those whose history does not fit. + """ + needed_blocks = ceil_div(max_seq_len + self.num_extra_kv_tokens, self.tokens_per_block) + if capacity_blocks >= needed_blocks: + return + max_request_tokens = capacity_blocks * self.tokens_per_block - self.num_extra_kv_tokens + logger.warning( + f"Sparse KV cache offload rejects requests whose prompt and max_tokens exceed " + f"{max_request_tokens} tokens: the smallest host history pool has room for " + f"{capacity_blocks} blocks of {self.tokens_per_block} tokens, and a request of " + f"max_seq_len={max_seq_len} tokens needs {needed_blocks}. Raise " + "sparse_attention_config.kv_cache_offload_config.host_history_size or the device KV " + "cache budget, or lower max_seq_len." + ) + + def _log_sparse_kv_offload_capacity(self, kv_cache_config: KvCacheConfig) -> None: + """Log how many requests fit, which limit binds, and the batch size that fits.""" + config = self.kv_cache_manager_py_config + num_tokens = self._get_sparse_offload_seq_len(kv_cache_config) + host_bytes_per_request = sparse_history_bytes_per_block(config.layers) * ceil_div( + num_tokens, self.tokens_per_block + ) + capacity = solve_sparse_offload_capacity( + num_tokens=num_tokens, + max_batch_size=self.max_batch_size, + # The device quota counts only device-backed pools. + gpu_pool_bytes=self.impl.get_quota(GPU_LEVEL), + host_pool_bytes=self._get_host_backed_pool_bytes(), + host_budget_bytes=self._host_history_budget, + gpu_bytes_per_request=self.get_request_device_bytes(num_tokens), + fixed_gpu_bytes_per_request=self.get_request_device_bytes(0), + host_bytes_per_request=host_bytes_per_request, + arenas=sparse_arena_shapes(config.layers, tokens_per_block=config.tokens_per_block), + num_leases=self.max_batch_size, + batch_buffer_bytes=self._sparse_batch_buffer_bytes, + ) + logger.info(capacity.describe()) + if capacity.fits < self.max_batch_size: + logger.warning( + f"Only {capacity.fits} of max_batch_size={self.max_batch_size} requests of " + f"{num_tokens} tokens fit, limited by {capacity.limited_by} memory; the rest " + "wait for capacity." + ) + + def get_request_device_bytes(self, num_tokens: int) -> int: + """Return the device pool bytes one decoding request of ``num_tokens`` tokens holds. + + Excludes the request's sparse cache lease. With sparse KV cache offload, + it must also leave out the history that lives in host memory, as the + quota estimators do; a manager with its own estimators overrides this. + """ + layer_sizes, attention_windows = self._get_runtime_cache_size_layer_components() + _, size_per_token, size_per_request = _estimate_cache_size_components( + layer_sizes, + attention_windows, + self.tokens_per_block, + scratch=False, + generation_capacity_headroom=self._generation_kv_capacity_headroom, + ) + return int(num_tokens * size_per_token + size_per_request) + + def _get_host_backed_pool_bytes(self) -> int: + """Return the bytes of the hot pools whose pages are pinned host memory.""" + sparse_buffers = self._get_sparse_buffer_ids() + total = 0 + for pool_group in self.get_host_backed_pool_groups(): + # A pool holds the coalesced buffer at its index in every variant, and + # variants of one group agree on sparsity position by position. + coalesced_buffers = pool_group.slot_desc.variants[0].coalesced_buffers + for pool in pool_group.pools: + buffer_ids = coalesced_buffers[int(pool.pool_index)].buffer_ids + if any( + (int(buffer_id.layer_id), str(buffer_id.role)) in sparse_buffers + for buffer_id in buffer_ids + ): + total += pool.slot_bytes * pool_group.num_slots + return total + + def _get_sparse_buffer_ids(self) -> set[tuple[int, str]]: + """Return the ``(layer_id, role)`` of every sparse buffer the manager allocates.""" + return { + (int(layer.layer_id), str(buffer.role)) + for layer in self.kv_cache_manager_py_config.layers + if isinstance(layer, AttentionLayerConfig) + for buffer in layer.buffers + if buffer.is_sparse + } + + def get_host_backed_pool_groups(self) -> list: + """Return the hot pool groups that hold host-backed pools. + + A sparse buffer's pool is host memory mapped into the device address + space. It shares its pool group, and so the group's slot count, with the + device-backed pools of the same life cycle. + + Returns: + The ``PoolGroupDesc`` of each hot pool group with a host-backed pool. + """ + sparse_buffers = self._get_sparse_buffer_ids() + if not sparse_buffers: + return [] + return [ + pool_group + for pool_group in self.impl.pool_group_descs + if any( + (int(buffer_id.layer_id), str(buffer_id.role)) in sparse_buffers + for variant in pool_group.slot_desc.variants + for coalesced in variant.coalesced_buffers + for buffer_id in coalesced.buffer_ids + ) + ] + def flush_iteration_events(self): event_manager = self.event_manager if event_manager is not None: @@ -7081,9 +7717,6 @@ def release_index_slot(self, request_id: int) -> None: # mirrored, and the target may release the same request twice. return if kv_cache is not None: - if self.sparse_metadata_batch is not None: - self.sparse_metadata_batch.record_read(self._stream.cuda_stream) - self.sparse_metadata_batch.remove(kv_cache) for beam_idx in range(int(kv_cache.beam_width)): for pool_idx in range(self.num_pools): kv_cache.set_base_page_index_buf(BeamIndex(beam_idx), pool_idx, None) @@ -7144,6 +7777,8 @@ def free_resources(self, request: LlmRequest, pin_on_release: bool = False): self._disagg_receive_ready.pop(request.py_request_id, None) self._on_allocation_stream.discard(request.py_request_id) self._allocation_ready.pop(request.py_request_id, None) + if self._host_history_reservations is not None: + self._host_history_reservations.release(request.py_request_id) self._remove_kv_cache(request.py_request_id) def _after_shrink(self, request_id: int, kv_cache: _KVCache) -> None: @@ -7386,8 +8021,17 @@ def calculate_scaling_factor_size_bytes( return get_size_in_bytes(cache_size // quant_vector_size, scaling_factor_dtype) def _iter_cache_buffers_for_invalid_check(self) -> Iterable[Tuple[int, torch.Tensor]]: + # Host-backed sparse history is left out: nothing reads a row of it + # before a transfer or a write-back fills it, and walking it would cross + # PCIe for every page of the host budget. + sparse_buffers = ( + self._get_sparse_buffer_ids() if self._sparse_kv_offload is not None else set() + ) pool_handled = set() for layer_id, layer_offset in self.layer_offsets.items(): + pool_id = self.layer_to_pool_mapping_dict[layer_offset] + if pool_id in pool_handled or (int(layer_offset), str(Role.KEY)) in sparse_buffers: + continue buffer = self.get_buffers(layer_id) if buffer is None: continue @@ -7447,12 +8091,10 @@ def check_invalid_values_in_kv_cache(self, fill_with_zero: bool = False) -> bool return bool(has_invalid_values) def shutdown(self): + # Before the caches close: a lender waiting on a lease has to be told to stop while + # the caches it lent against are still alive. if self._sharing is not None: self._sharing._on_shutdown() - if self.sparse_metadata_batch is not None: - self.sparse_metadata_batch.record_read(self._stream.cuda_stream) - self.sparse_metadata_batch.close() - self.sparse_metadata_batch = None for kv_cache in self.kv_cache_map.values(): kv_cache.close() self.kv_cache_map.clear() @@ -7494,6 +8136,23 @@ def get_needed_resource_to_completion(self, request: LlmRequest) -> int: # return need_blocks return 0 + def kv_block_budget_applies(self) -> bool: + # Only offloaded sparse history has a request-level budget here. + return self._host_history_capacity_blocks is not None + + def get_request_kv_block_budget(self, request: LlmRequest) -> Optional[Tuple[int, int]]: + """Return the host history pages ``request`` needs and the most it can be given. + + With sparse KV cache offload, a request is admitted only once its host + history pages through its last output token can be reserved + (``prepare_disagg_gen_init``). A request that needs more pages than the + smallest rank's pool holds would wait forever, so the executor rejects + it when it arrives. Returns ``None`` without offload. + """ + if self._host_history_capacity_blocks is None: + return None + return self._get_host_history_blocks(request), self._host_history_capacity_blocks + # TODO: refactor get_cache_size_per_token and get_cache_bytes_per_token to use the same logic @classmethod def get_cache_size_per_token( @@ -7887,6 +8546,19 @@ def get_cache_quota_per_request( quota = sum(math.ceil(bytes_ / granularity) * granularity for bytes_ in layer_bytes) return quota if quota > 0 else None + def update_forward_resources( + self, scheduled_batch: ScheduledRequests, *, wait_for_input_copy: Callable[[], None] + ) -> None: + """Update resources after a successful model forward has been enqueued. + + The caller's stream must wait for model execution before this hook. + It runs outside CUDA graph capture, before CPU sampling-result processing + and the next scheduling pass. Managers whose history depends on executed inputs + can override it; ordinary caches update history after sampling. + Overrides must call wait_for_input_copy before changing host page tables + that the model's input preparation may still be copying to the device. + """ + def update_context_resources(self, scheduled_batch: ScheduledRequests): """Update KV cache for context requests in the current batch. @@ -8040,21 +8712,6 @@ def copy_batch_block_offsets( num_seqs: int, max_blocks: Optional[int] = None, ): - self._publish_sparse_metadata() - # Sparse history can remain on GPU while offload is deferred. Check all - # mapped pages after publication, which can retry the deferred offload. - if self.sparse_metadata_batch is not None and any( - level is not None and level != GPU_LEVEL - for req_id in request_ids - for layer_group_id in self._sparse_layer_group_ids - for level in self.kv_cache_map[req_id] - .get_page_storage_snapshot(layer_group_id) - .cache_levels - ): - raise RuntimeError( - "Offloaded sparse history requires Batch metadata and sparse fetch; " - "dense attention offsets cannot address host slots" - ) # max_blocks is accepted for signature parity with KVCacheManager; the # device-side copy op here already scales with allocated blocks only. assert beam_width <= self.max_copy_beam_width @@ -8188,9 +8845,7 @@ def _create_kv_cache( kv_cache.stop_committing() self.impl.mark_stats_excluded(request_id) kv_cache.discard_pending_stats() - index = self.index_mapper.add_new_sequence(request_id) - if self.sparse_metadata_batch is not None: - self.sparse_metadata_batch.add(kv_cache, index) + self.index_mapper.add_new_sequence(request_id) self._set_page_index_bufs(request_id, kv_cache) return kv_cache diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_layout.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_layout.py index 26d98f8a6f68..55fc1d365184 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_layout.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_layout.py @@ -568,7 +568,11 @@ def derive_layout(manager: KVCacheManagerV2) -> ManagerLayout: windows=tuple(desc.window for desc in layer_groups), sink_blocks=tuple(desc.sink_blocks for desc in layer_groups), recurrent=tuple(desc.kind == "state" for desc in layer_groups), - sparse=tuple(_sparse(layer_cfg[layer_ids[0]]) for layer_ids in grouping), + # Every layer of the group, not just its first: grouping follows the life cycle, + # which sparsity does not take part in, so a group can hold both kinds. + sparse=tuple( + any(_sparse(layer_cfg[layer_id]) for layer_id in layer_ids) for layer_ids in grouping + ), layers=tuple(desc.layers for desc in layer_groups), shards=tuple((desc.shard.count, desc.shard.index) for desc in layer_groups), page_bytes={g: sum(p.slot_bytes for p in pools) for g, pools in device_pools.items()}, diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/sparse_offload_memory.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/sparse_offload_memory.py new file mode 100644 index 000000000000..d6d910a1bc68 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/sparse_offload_memory.py @@ -0,0 +1,839 @@ +# 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. + +"""Memory accounting for sparse attention KV cache offload. + +With offload, the KV history a sparse attention layer selects from lives in +host-backed pages, and each request leases a GPU cache of selected units from a +sparse cell arena. + +The native manager puts a sparse buffer's pages in a host-backed pool of the +same pool group as its life cycle's dense buffers. A slot spans every pool of +its group, so the host-backed pools get the slot count the device quota buys +for the dense ones, and nothing else bounds them. The host memory offload takes +is therefore set by the device quota and by how that quota is split across life +cycles. The arena is a separate device allocation outside every tier's quota. + +Everything here is derived from the layer configs a cache manager declares, so +it is the same for every model. A model marks each offloaded buffer +``is_sparse``, gives its layer a fetch path from :func:`sparse_cache_config`, +and reports the device bytes of its other buffers through +``KVCacheManagerV2.get_request_device_bytes``. + +The module has three parts: mirrors of the native layout, which must change +when ``cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2`` does; the sizing +policy; and the capacity report and host admission. +""" + +from __future__ import annotations + +import math +from collections import Counter, defaultdict +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Sequence + +from tensorrt_llm.math_utils import ceil_div, pad_up +from tensorrt_llm.runtime.kv_cache_manager_v2 import AttentionLayerConfig, SparseCacheConfig + +if TYPE_CHECKING: + from tensorrt_llm.runtime.kv_cache_manager_v2 import ( + BufferConfig, + KVCacheManagerConfig, + LayerConfig, + ) + +_GIB = float(1 << 30) +_MIB = float(1 << 20) + +# --------------------------------------------------------------------------- +# Native layout: SparseCellArena and leaseLayout in sparseCache.cpp, the Batch in +# batch.cpp, pool and life-cycle formation in storage/config.cpp and +# lifeCycleRegistry.cpp, cacheTierGranularity in storageManager.cpp, and HostMem +# in utils/hostMem.cpp. +# --------------------------------------------------------------------------- + +# SparseCellArena rounds every region to this many bytes. +SPARSE_ARENA_ALIGNMENT = 256 +# Every arena and Batch index array holds 32-bit entries (int32_t or uint32_t). +_INDEX_BYTES = 4 +# A Batch row holds two watermarks, history length and capacity. +_BATCH_WATERMARKS_PER_ROW = 2 +# Pinned host memory is mapped in whole pages of this many bytes. +_HOST_PAGE_BYTES = 4096 +# A decode step writes one query token per request, which completes at most one +# unit, so one unit per layer is in flight (_sparseCacheMaxInFlightUnits). +_SPARSE_MAX_FLUSH_UNITS = 1 + + +def _index_array_bytes(num_entries: int) -> int: + """Return the aligned bytes of an arena array of ``num_entries`` 32-bit entries.""" + return pad_up(_INDEX_BYTES * num_entries, SPARSE_ARENA_ALIGNMENT) + + +def sparse_units_per_page(unit_tokens: int | None, tokens_per_block: int) -> int: + """Return the selection units one page holds. + + Args: + unit_tokens: Tokens one unit covers, or None when a unit is a whole page. + tokens_per_block: Tokens per page. + """ + return tokens_per_block // (unit_tokens or tokens_per_block) + + +@dataclass(frozen=True, slots=True) +class SparseArenaShape: + """GPU cell arena of one sparse life cycle. + + Attributes: + num_layers: Layers with a sparse fetch path. + num_selections: Selections those layers resolve between them, which is how + many metadata rows a lease carries. One per layer unless layers name a + shared ``selection_id``, in which case they share a row. + cells_per_buffer: Cells each fetched sparse buffer gets per lease. + units_per_cell: Selection units in one cell. + cell_bytes: Bytes of one cell of each fetched sparse buffer, in layer order. + """ + + num_layers: int + num_selections: int + cells_per_buffer: int + units_per_cell: int + cell_bytes: tuple[int, ...] + + @property + def entries(self) -> int: + """Cache entries of one lease: one per cached unit of each selection. + + Per selection, not per layer: layers sharing a selection hold the same units + in mirrored cells, so one resident set serves them all. + """ + return self.num_selections * self.cells_per_buffer * self.units_per_cell + + @property + def lease_metadata_bytes(self) -> int: + """Bytes of one lease's residency and eviction state. + + One word per entry: the unit it holds and when that unit was last used. + + The reservation only. A step's work lists belong to the Batch, not to a lease: + they are sized by the batch and reused by whichever requests are seated. + """ + return _index_array_bytes(self.entries) + + @property + def lease_bytes(self) -> int: + """GPU bytes one lease adds to the arena.""" + return self.cells_per_buffer * sum(self.cell_bytes) + self.lease_metadata_bytes + + def arena_bytes(self, max_leases: int) -> int: + """Return the GPU bytes of an arena sized for ``max_leases`` leases.""" + # One region per pool. Buffers of one cell size share a pool, so a region + # holds every lease's cells of every buffer of that size. + buffers_by_cell_bytes: dict[int, int] = defaultdict(int) + for cell_bytes in self.cell_bytes: + buffers_by_cell_bytes[cell_bytes] += 1 + cell_regions = sum( + pad_up( + max_leases * num_buffers * self.cells_per_buffer * cell_bytes, + SPARSE_ARENA_ALIGNMENT, + ) + for cell_bytes, num_buffers in buffers_by_cell_bytes.items() + ) + num_buffers = len(self.cell_bytes) + shared_arrays = ( + num_buffers, # layer of each buffer + num_buffers, # slot of each buffer in its pool + self.num_layers, # metadata row each layer resolves within + ) + shared = sum(_index_array_bytes(n) for n in shared_arrays) + return cell_regions + shared + max_leases * self.lease_metadata_bytes + + +def _fetched_sparse_buffers(layer: LayerConfig) -> list[BufferConfig]: + """Return the sparse buffers of ``layer`` that are read through a sparse cache.""" + if getattr(layer, "sparse_cache", None) is None: + return [] + return [buffer for buffer in layer.buffers if buffer.is_sparse] + + +def _fetch_path_life_cycles( + layers: Sequence[LayerConfig], tokens_per_block: int +) -> list[list[AttentionLayerConfig]]: + """Return the layers with a sparse fetch path, grouped by life cycle.""" + # Layers share a life cycle when their window and sink pages match. + groups: dict[tuple[int | None, int], list[AttentionLayerConfig]] = defaultdict(list) + for layer in layers: + if _fetched_sparse_buffers(layer): + sink_blocks = ceil_div(layer.num_sink_tokens or 0, tokens_per_block) + groups[layer.sliding_window_size, sink_blocks].append(layer) + return list(groups.values()) + + +def sparse_arena_shapes( + layers: Sequence[LayerConfig], *, tokens_per_block: int +) -> list[SparseArenaShape]: + """Return the arena of every life cycle whose layers have a sparse fetch path. + + Args: + layers: Every layer config the manager allocates. + tokens_per_block: Tokens per page. + + Raises: + ValueError: Layers sharing a life cycle declare different + ``sparse_cache`` values. They share one arena, so the native manager + rejects them too. + """ + shapes = [] + for grouped_layers in _fetch_path_life_cycles(layers, tokens_per_block): + geometries = { + ( + layer.sparse_cache.num_selected_blocks, + layer.sparse_cache.max_cached_units, + sparse_units_per_page(layer.sparse_cache.unit_tokens, tokens_per_block), + ) + for layer in grouped_layers + } + if len(geometries) > 1: + raise ValueError( + "Layers sharing a life cycle must declare the same sparse_cache; they share " + f"one arena. Got (selected, cached, units per page) {sorted(geometries)}." + ) + ((_, max_cached_units, units_per_page),) = geometries + # An unset selection_id reads as the layer's own, so a layer that names + # nothing keeps a row to itself and only an explicit id merges two. + selections = { + layer.sparse_cache.selection_id + if layer.sparse_cache.selection_id is not None + else layer.layer_id + for layer in grouped_layers + } + shapes.append( + SparseArenaShape( + num_layers=len(grouped_layers), + num_selections=len(selections), + cells_per_buffer=ceil_div(max_cached_units, units_per_page), + units_per_cell=units_per_page, + # A cell holds one page of one buffer, and a buffer's size is one page. + cell_bytes=tuple( + buffer.size + for layer in grouped_layers + for buffer in _fetched_sparse_buffers(layer) + ), + ) + ) + return shapes + + +def sparse_arena_bytes(config: KVCacheManagerConfig, max_leases: int) -> int: + """Return the GPU bytes of every sparse arena the Batch allocates. + + ``max_leases`` is the Batch's row count, which is what the arenas are sized + for: a lease is taken when a cache joins and returned when it leaves. + + The layers alone decide whether there is an arena at all, so nothing else is + read until they say there is one: a config with no sparse fetch path allocates + nothing, and need not describe its block size. + """ + if not any( + getattr(layer, "sparse_cache", None) is not None + and any(buffer.is_sparse for buffer in layer.buffers) + for layer in config.layers + ): + return 0 + shapes = sparse_arena_shapes(config.layers, tokens_per_block=config.tokens_per_block) + if not shapes or max_leases <= 0: + return 0 + return sum(shape.arena_bytes(max_leases) for shape in shapes) + + +def sparse_history_bytes_per_block(layers: Sequence[LayerConfig]) -> int: + """Return the host bytes one request's offloaded history takes per page of tokens. + + Every sparse buffer keeps its history in host-backed pages, and a buffer's + size is one page, so this is the summed size of the sparse buffers. A + windowed sparse layer is counted as if it kept its whole history. + """ + return sum(buffer.size for layer in layers for buffer in layer.buffers if buffer.is_sparse) + + +@dataclass(frozen=True) +class LifeCyclePage: + """What one page of a life cycle costs, split by where the native manager puts it. + + Attributes: + window_size: Sliding window in tokens, or None when the life cycle keeps + all of its history. + num_sink_tokens: Tokens kept from the start of the sequence beside the + window. + is_state: Whether the life cycle holds one slot of state per request, + as SSM layers do, rather than pages of tokens. + device_bytes: Bytes of its dense buffers, which the device quota pays for. + host_bytes: Bytes of its sparse buffers, which live in pinned host memory. + device_pool_bytes: Slot bytes of each device-backed pool, largest first. + The native manager gives a life cycle's dense buffers of one page + size a pool of their own, so these sum to ``device_bytes``. + """ + + window_size: int | None + num_sink_tokens: int + is_state: bool + device_bytes: int + host_bytes: int + device_pool_bytes: tuple[int, ...] + + def pages_per_request( + self, *, capacity: int, capacity_headroom: int, tokens_per_block: int + ) -> int: + """Return the most pages one decoding request holds. + + Args: + capacity: Tokens the request's cache holds, its history included. + capacity_headroom: Tokens a decode step reserves above the history. + tokens_per_block: Tokens per page. + """ + if self.is_state: + return 1 + pages = ceil_div(capacity, tokens_per_block) + if self.window_size is None: + return pages + # The live interval spans window + headroom - 1 tokens, at any page offset. + window_pages = ceil_div(self.window_size + capacity_headroom - 2, tokens_per_block) + 1 + sink_pages = ceil_div(self.num_sink_tokens, tokens_per_block) + return min(pages, window_pages + sink_pages) + + +def life_cycle_pages(layers: Sequence[LayerConfig], tokens_per_block: int) -> list[LifeCyclePage]: + """Return the life cycles ``layers`` form, in the order the native manager numbers them. + + The native manager numbers life cycles by their first layer, which is also + the order of ``initial_pool_ratio``. Attention layers share a life cycle + when their window and sink pages match; all SSM layers share one. + + Raises: + ValueError: A life cycle has sparse buffers but no dense one. The native + manager sizes a pool group's slots from its device-backed pools, so it + rejects such a life cycle too. + """ + keys: list[tuple] = [] + device_pages: dict[tuple, Counter[int]] = {} + host_bytes: dict[tuple, int] = {} + pages: dict[tuple, tuple[int | None, int, bool]] = {} + for layer in layers: + if isinstance(layer, AttentionLayerConfig): + window = layer.sliding_window_size + sink_tokens = layer.num_sink_tokens or 0 + key = ("attention", window, ceil_div(sink_tokens, tokens_per_block)) + shape = (window, sink_tokens, False) + else: + key = ("state",) + shape = (None, 0, True) + if key not in pages: + keys.append(key) + pages[key] = shape + device_pages[key] = Counter() + host_bytes[key] = 0 + for buffer in layer.buffers: + if buffer.is_sparse: + host_bytes[key] += buffer.size + elif buffer.size: + override = buffer.tokens_per_block_override or tokens_per_block + device_pages[key][buffer.size * (tokens_per_block // override)] += 1 + result = [] + for key in keys: + device_pool_bytes = tuple( + sorted((size * count for size, count in device_pages[key].items()), reverse=True) + ) + if host_bytes[key] and not device_pool_bytes: + raise ValueError( + "A life cycle with sparse buffers must also hold a dense buffer: the native " + "KV cache manager sizes a pool group's slots from its device-backed pools." + ) + window, sink_tokens, is_state = pages[key] + result.append( + LifeCyclePage( + window_size=window, + num_sink_tokens=sink_tokens, + is_state=is_state, + device_bytes=sum(device_pool_bytes), + host_bytes=host_bytes[key], + device_pool_bytes=device_pool_bytes, + ) + ) + return result + + +@dataclass(frozen=True, slots=True) +class SparseBatchShape: + """The native Batch the sparse fetch reads, sized once for every step. + + It holds a page table per life cycle, and two watermarks and one arena slot per row, + in a device mirror, a pinned host staging copy and a device staging copy, and the + fetch's work arrays on the device. Every step uploads the host staging copy to the + device one, then copies the mirror's part of it across. + + Attributes: + num_life_cycles: Life cycles of the manager; each has a page table. + max_batch_size: Rows. + max_blocks_per_seq: Columns of each page table. + selection_widths: Units one step selects, for each life cycle with a sparse + fetch path. + layer_counts: Layers of each of those life cycles. The write-back keeps a work + list per layer, so this scales that term; a zero would drop it and under-size + the allocation, which is why it is required rather than defaulted. + flush_capacities: Units one step can write back per layer, for each of those life + cycles, from the step's token budget. + """ + + num_life_cycles: int + max_batch_size: int + max_blocks_per_seq: int + selection_widths: tuple[int, ...] + layer_counts: tuple[int, ...] + flush_capacities: tuple[int, ...] + + @property + def block_bytes(self) -> int: + """Device bytes of the staged block: the page tables, watermarks and leases.""" + entries = self.max_batch_size * ( + self.num_life_cycles * self.max_blocks_per_seq + _BATCH_WATERMARKS_PER_ROW + 1 + ) + return _INDEX_BYTES * entries + + @property + def mirror_bytes(self) -> int: + """Device bytes of the mirror: the block's tail, which a captured kernel reads. + + The page tables leading the block are read from its device staging copy by the + index conversion, so they are not carried on. + """ + return _INDEX_BYTES * self.max_batch_size * (_BATCH_WATERMARKS_PER_ROW + 1) + + @property + def staging_bytes(self) -> int: + """Pinned host bytes of the staging copy of the block.""" + return pad_up(self.block_bytes, _HOST_PAGE_BYTES) + + @property + def device_bytes(self) -> int: + """Device bytes of the Batch: the staged block, the mirror of its tail and the + fetch's work arrays.""" + return self.block_bytes + self.mirror_bytes + self.fetch_scratch_bytes + + @property + def fetch_scratch_bytes(self) -> int: + """Device bytes of a step's work arrays, for the fetch and the write-back both. + + Per life cycle and row: one entry per selected unit, a miss count and the unit + in flight. Then the two compacted work lists, which carry a row, an ordinal and + a slot per entry plus one total per list: one list for the fetch, sized for every + row missing its whole selection, and one per layer for the write-back, sized by + the step's token budget. + """ + rows = self.max_batch_size + total = 0 + for width, num_layers, flush_capacity in zip( + self.selection_widths, self.layer_counts, self.flush_capacities, strict=True + ): + total += rows * (width + 1 + _SPARSE_MAX_FLUSH_UNITS) + total += 3 * rows * width + 1 + total += num_layers * (3 * flush_capacity + 1) + total += rows * num_layers + return _INDEX_BYTES * total + + +def sparse_batch_shape( + layers: Sequence[LayerConfig], + *, + tokens_per_block: int, + max_batch_size: int, + max_blocks_per_seq: int, + max_total_input_tokens: int | None = None, +) -> SparseBatchShape: + """Return the Batch a manager with ``layers`` steps its requests in. + + Args: + layers: Every layer config the manager allocates. + tokens_per_block: Tokens per page. + max_batch_size: Rows of the Batch. + max_blocks_per_seq: Pages the longest request can hold. + max_total_input_tokens: Tokens one step may ingest across every row. It sizes + the write-back's work lists: a row advancing ``t`` tokens crosses at most + ``t / unit_tokens + 1`` units. Left unset, a decode step is assumed, where + each row advances one token and so crosses at most one unit boundary. + """ + fetch_groups = list(_fetch_path_life_cycles(layers, tokens_per_block)) + + def flush_capacity(grouped_layers: Sequence[LayerConfig]) -> int: + if max_total_input_tokens is None: + return max_batch_size * _SPARSE_MAX_FLUSH_UNITS + unit_tokens = grouped_layers[0].sparse_cache.unit_tokens or tokens_per_block + return max_total_input_tokens // unit_tokens + max_batch_size + + return SparseBatchShape( + num_life_cycles=len(life_cycle_pages(layers, tokens_per_block)), + max_batch_size=max_batch_size, + max_blocks_per_seq=max_blocks_per_seq, + selection_widths=tuple( + grouped_layers[0].sparse_cache.num_selected_blocks for grouped_layers in fetch_groups + ), + layer_counts=tuple(len(grouped_layers) for grouped_layers in fetch_groups), + flush_capacities=tuple(flush_capacity(grouped_layers) for grouped_layers in fetch_groups), + ) + + +# --------------------------------------------------------------------------- +# Sizing policy: the GPU cache per request, the quota split, the host budget and +# the device quota it bounds. +# --------------------------------------------------------------------------- + + +def sparse_cache_cells(cache_factor: float, num_selected_units: int, units_per_cell: int) -> int: + """Return the cells one sparse buffer gets in each request's device cache. + + The cache holds ``cache_factor`` selections plus the unit being written, + rounded up to whole cells. A request has at most one unit in flight, so that + unit adds one whatever the factor. A factor of 1 is the smallest cache the + native manager accepts. + + Args: + cache_factor: Selections the cache holds; at least 1. + num_selected_units: Units one step selects. + units_per_cell: Units in one cell, i.e. in one page. + """ + if cache_factor < 1: + raise ValueError(f"cache_factor must be at least 1, got {cache_factor}") + if num_selected_units <= 0 or units_per_cell <= 0: + raise ValueError( + "num_selected_units and units_per_cell must be positive, got " + f"{num_selected_units} and {units_per_cell}" + ) + # Round before taking the ceiling so float error cannot add a unit. + selected_units = math.ceil(round(cache_factor * num_selected_units, 6)) + return ceil_div(selected_units + 1, units_per_cell) + + +def sparse_cache_config( + *, + num_selected_units: int, + unit_tokens: int | None, + tokens_per_block: int, + gpu_cache_factor: float, + selection_id: int | None = None, +) -> SparseCacheConfig: + """Return a layer's sparse fetch path whose cache holds ``gpu_cache_factor`` selections. + + A selection unit is what one selection names: a token, a compressed row or + pool of tokens, or a whole page. Each request caches selected units in GPU + cells of one page each. The cache holds ``gpu_cache_factor`` selections plus + the unit being written, rounded up to whole cells. + + Args: + num_selected_units: Units one decode step selects for one request. When + heads select separately, this is the size of their union. + unit_tokens: Tokens one unit covers, or None when a unit is a whole page. + Must divide ``tokens_per_block``. + tokens_per_block: Tokens per page. + gpu_cache_factor: Selections the cache holds; at least 1. + selection_id: Names the selection this layer resolves, where several layers + are handed one. Layers naming the same id share a resident set and + resolve once between them, so the id must be shared by exactly the + layers an indexer selects for, and those layers must agree on every + other argument here. None, the default, means the layer selects for + itself; a layer that names nothing never shares. + """ + if unit_tokens is not None and (unit_tokens <= 0 or tokens_per_block % unit_tokens): + raise ValueError( + f"unit_tokens ({unit_tokens}) must divide tokens_per_block ({tokens_per_block})" + ) + units_per_page = sparse_units_per_page(unit_tokens, tokens_per_block) + cells = sparse_cache_cells(gpu_cache_factor, num_selected_units, units_per_page) + return SparseCacheConfig( + num_selected_blocks=num_selected_units, + max_cached_units=cells * units_per_page, + # Resolved rather than passed through: the config names a unit size, and a whole + # page is one of them. None is this helper's convenience, not the config's. + unit_tokens=tokens_per_block // units_per_page, + selection_id=selection_id, + ) + + +def sparse_offload_pool_ratio( + life_cycles: Sequence[LifeCyclePage], + *, + capacity: int, + capacity_headroom: int, + tokens_per_block: int, +) -> list[float]: + """Return the device quota's split across life cycles for decoding requests. + + Each life cycle's share is the device bytes its pages take in a request of + ``capacity`` tokens, so the quota buys the same number of such requests in + every life cycle. Host-backed pools cost the quota nothing and get no share. + + Args: + life_cycles: The manager's life cycles, in native order. + capacity: Tokens one request's cache holds. + capacity_headroom: Tokens a decode step reserves above the history. + tokens_per_block: Tokens per page. + + Returns: + One positive share per life cycle, summing to 1; the order and meaning of + ``KVCacheManagerConfig.initial_pool_ratio``. + """ + # Device and host bytes together. The manager's pool ratio counts every pool, host-backed + # ones included -- it is projected onto cold levels, where one pool holds dense and sparse + # buffers together, and a rebalance compares it against a layout measured the same way. + # Weighing device bytes alone here would hand the manager a ratio in one currency and + # measure the result in another. + weights = [ + (life_cycle.device_bytes + life_cycle.host_bytes) + * life_cycle.pages_per_request( + capacity=capacity, + capacity_headroom=capacity_headroom, + tokens_per_block=tokens_per_block, + ) + for life_cycle in life_cycles + ] + if any(weight <= 0 for weight in weights): + raise ValueError(f"Every life cycle needs bytes to share the quota, got {weights}") + total = sum(weights) + return [weight / total for weight in weights] + + +def plan_host_history_bytes( + *, + host_history_size: int | None, + max_batch_size: int, + host_bytes_per_request: int, + pinnable_host_bytes: float, + staging_bytes: int, +) -> int: + """Return the host bytes the host-backed history pools may take. + + Offload also pins the sparse Batch's staging buffers, so they come out of + the same host memory. An explicit ``host_history_size`` covers both. + Otherwise the pools hold ``max_batch_size`` requests, capped so that pools + and staging fit in the host memory the rank may pin. + + Args: + host_history_size: The configured size, or None to size automatically. + max_batch_size: Requests the automatic budget holds. + host_bytes_per_request: Host history bytes one request holds. + pinnable_host_bytes: Host bytes automatic sizing may pin for offload; + ``inf`` when unknown. + staging_bytes: Pinned host bytes of the sparse Batch's staging buffers. + + Raises: + ValueError: No host memory is left for the pools. + """ + if host_history_size is not None: + budget = host_history_size - staging_bytes + if budget <= 0: + raise ValueError( + "sparse_attention_config.kv_cache_offload_config.host_history_size " + f"({host_history_size / _MIB:.1f}MiB) leaves no host memory for offloaded " + f"history: the sparse Batch's staging buffers take {staging_bytes / _MIB:.1f}MiB " + "of it." + ) + return budget + if not host_bytes_per_request: + # This rank's layers offload no history. + return 0 + budget = min(max_batch_size * host_bytes_per_request, pinnable_host_bytes - staging_bytes) + if budget <= 0: + raise ValueError( + "Sizing sparse KV cache offload automatically leaves no host memory for offloaded " + f"history: this rank may pin {max(0.0, pinnable_host_bytes) / _MIB:.1f}MiB, and the " + f"sparse Batch's staging buffers take {staging_bytes / _MIB:.1f}MiB. Set " + "sparse_attention_config.kv_cache_offload_config.host_history_size, or lower " + "kv_cache_config.host_cache_size." + ) + return int(budget) + + +# --------------------------------------------------------------------------- +# Capacity report and host admission. +# --------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class SparseOffloadCapacity: + """How many requests of one length a worker with sparse KV cache offload holds. + + Attributes: + num_tokens: Tokens per request. + max_batch_size: The configured batch limit. + gpu_pool_bytes: Device bytes of the KV pools. + host_pool_bytes: Host bytes of the host-backed pools the manager allocated. + host_budget_bytes: Host bytes offloaded history may take. + gpu_bytes_per_request: Device pool bytes one request holds. + fixed_gpu_bytes_per_request: The part of ``gpu_bytes_per_request`` that + does not grow with length, such as sliding windows. + host_bytes_per_request: Host history bytes one request holds. + arena_bytes: Device bytes of the sparse cache arenas. + num_leases: Leases each arena holds. + lease_bytes: Device bytes of one lease, summed over the arenas. + batch_buffer_bytes: Device bytes of the buffers sized for + ``max_batch_size`` rows: the Batch and the model's offload scratch. + by_limit: Requests that fit by each limit: ``max_batch_size``, ``host`` + and ``gpu``, in the order that breaks ties. + recommended_max_batch_size: The largest batch the device budget and the + host budget hold. + """ + + num_tokens: int + max_batch_size: int + gpu_pool_bytes: int + host_pool_bytes: int + host_budget_bytes: int + gpu_bytes_per_request: int + fixed_gpu_bytes_per_request: int + host_bytes_per_request: int + arena_bytes: int + num_leases: int + lease_bytes: int + batch_buffer_bytes: int + by_limit: dict[str, int] = field(compare=False) + recommended_max_batch_size: int + + @property + def limited_by(self) -> str: + """The limit that admits the fewest requests.""" + return min(self.by_limit, key=self.by_limit.__getitem__) + + @property + def fits(self) -> int: + """Requests that fit.""" + return self.by_limit[self.limited_by] + + def describe(self) -> str: + """Return a one-line report for the log.""" + return ( + f"Sparse KV cache offload capacity at {self.num_tokens} tokens per request: " + f"GPU pools {self.gpu_pool_bytes / _GIB:.2f}GiB at " + f"{self.gpu_bytes_per_request / _MIB:.1f}MiB per request " + f"({self.fixed_gpu_bytes_per_request / _MIB:.1f}MiB fixed); arena " + f"{self.arena_bytes / _GIB:.2f}GiB for {self.num_leases} leases of " + f"{self.lease_bytes / _MIB:.2f}MiB; offload buffers " + f"{self.batch_buffer_bytes / _MIB:.1f}MiB for {self.max_batch_size} rows; " + f"host-backed history pools " + f"{self.host_pool_bytes / _GIB:.2f}GiB at " + f"{self.host_bytes_per_request / _MIB:.1f}MiB per request, host budget " + f"{self.host_budget_bytes / _GIB:.2f}GiB. Requests: {self.by_limit['gpu']} by " + f"GPU, {self.by_limit['host']} by host, max_batch_size {self.max_batch_size}; " + f"{self.fits} fit, limited by {self.limited_by}. A max_batch_size of " + f"{self.recommended_max_batch_size} fits the GPU budget and the host budget." + ) + + +def solve_sparse_offload_capacity( + *, + num_tokens: int, + max_batch_size: int, + gpu_pool_bytes: int, + host_pool_bytes: int, + host_budget_bytes: int, + gpu_bytes_per_request: int, + fixed_gpu_bytes_per_request: int, + host_bytes_per_request: int, + arenas: Sequence[SparseArenaShape], + num_leases: int, + batch_buffer_bytes: int, +) -> SparseOffloadCapacity: + """Return how many requests of ``num_tokens`` tokens fit, and the batch size that fits. + + Args: + num_tokens: Tokens per request. + max_batch_size: The configured batch limit, the rows of the Batch. + gpu_pool_bytes: Device bytes of the KV pools. + host_pool_bytes: Host bytes of the host-backed pools the manager allocated. + host_budget_bytes: Host bytes offloaded history may take. + gpu_bytes_per_request: Device pool bytes one request holds. + fixed_gpu_bytes_per_request: The part of ``gpu_bytes_per_request`` that + does not grow with length. + host_bytes_per_request: Host history bytes one request holds. + arenas: The arena of each sparse life cycle. + num_leases: Leases each arena holds. A cache holds a lease only while it + is seated in the Batch, so this is ``max_batch_size``; any beyond it + are padding. + batch_buffer_bytes: Device bytes of the buffers sized for + ``max_batch_size`` rows, reserved apart from ``gpu_pool_bytes``. + """ + gpu_bytes_per_request = max(1, gpu_bytes_per_request) + host_bytes_per_request = max(1, host_bytes_per_request) + arena_bytes = sum(arena.arena_bytes(num_leases) for arena in arenas) + lease_bytes = sum(arena.lease_bytes for arena in arenas) + by_limit = { + "max_batch_size": max_batch_size, + "host": host_pool_bytes // host_bytes_per_request, + "gpu": gpu_pool_bytes // gpu_bytes_per_request, + } + # The device budget pays for the pools, the arena and the buffers sized for + # max_batch_size rows. The arena's shared tables and padding leases are + # fixed; every batch slot adds a request's pages, the one lease its seated + # cache holds and its rows of those buffers. + padding_leases = max(0, num_leases - max_batch_size) + fixed_arena_bytes = sum(arena.arena_bytes(padding_leases) for arena in arenas) + batch_bytes_per_slot = ceil_div(batch_buffer_bytes, max(1, max_batch_size)) + recommended = min( + max(0, gpu_pool_bytes + arena_bytes + batch_buffer_bytes - fixed_arena_bytes) + // (gpu_bytes_per_request + lease_bytes + batch_bytes_per_slot), + host_budget_bytes // host_bytes_per_request, + ) + return SparseOffloadCapacity( + num_tokens=num_tokens, + max_batch_size=max_batch_size, + gpu_pool_bytes=gpu_pool_bytes, + host_pool_bytes=host_pool_bytes, + host_budget_bytes=host_budget_bytes, + gpu_bytes_per_request=gpu_bytes_per_request, + fixed_gpu_bytes_per_request=fixed_gpu_bytes_per_request, + host_bytes_per_request=host_bytes_per_request, + arena_bytes=arena_bytes, + num_leases=num_leases, + lease_bytes=lease_bytes, + batch_buffer_bytes=batch_buffer_bytes, + by_limit=by_limit, + recommended_max_batch_size=recommended, + ) + + +class HostHistoryReservations: + """Host history pages reserved for admitted requests through their last token. + + A request whose history outgrows the host-backed pool while it decodes has + nowhere to go, so admission reserves every page a request can reach. + """ + + def __init__(self, capacity_blocks: int) -> None: + """Track reservations against a pool of ``capacity_blocks`` pages.""" + self.capacity_blocks = capacity_blocks + self.reserved_blocks = 0 + self._blocks_by_request: dict[int, int] = {} + + def can_reserve(self, request_id: int, num_blocks: int) -> bool: + """Return whether ``num_blocks`` pages fit for ``request_id``, replacing its reservation.""" + held = self._blocks_by_request.get(request_id, 0) + return self.reserved_blocks - held + num_blocks <= self.capacity_blocks + + def reserve(self, request_id: int, num_blocks: int) -> None: + """Reserve ``num_blocks`` pages for ``request_id``, replacing its reservation.""" + held = self._blocks_by_request.get(request_id, 0) + self._blocks_by_request[request_id] = num_blocks + self.reserved_blocks += num_blocks - held + + def release(self, request_id: int) -> None: + """Release the pages reserved for ``request_id``, if any.""" + self.reserved_blocks -= self._blocks_by_request.pop(request_id, 0) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index d319f73e439d..77be7db6a6a7 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -3397,20 +3397,22 @@ def _executor_loop_pp(self): self._handle_dynamic_draft_len(scheduled_batch) - self._maybe_record_hang_diagnostic_phase( - "preparing_resources", scheduled_batch) - self.resource_manager.prepare_resources(scheduled_batch) - # The generation requests that do not have batch_idx # need to be in front of the batch due to the assumptions # made in model_engine.py::_forward_step. This is only important # for disaggregated serving. For non-disaggregated serving, # the generation requests always have batch_idx. + # Sort before preparing resources: a resource manager can + # assign per-request rows in this order. scheduled_batch.generation_requests = sorted( # stable sort scheduled_batch.generation_requests, key=lambda req: int(req.py_batch_idx is not None), ) + self._maybe_record_hang_diagnostic_phase( + "preparing_resources", scheduled_batch) + self.resource_manager.prepare_resources(scheduled_batch) + if self.kv_cache_transceiver: # Return the first token to the client self._handle_first_token_response(scheduled_batch) @@ -6334,6 +6336,18 @@ def _executor_loop_overlap(self) -> None: self._handle_dynamic_draft_len(scheduled_batch) + # The generation requests that do not have batch_idx + # need to be in front of the batch due to the assumptions + # made in model_engine.py::_forward_step. This is only important + # for disaggregated serving. For non-disaggregated serving, + # the generation requests always have batch_idx. + # Sort before preparing resources: a resource manager can + # assign per-request rows in this order. + scheduled_batch.generation_requests = sorted( # stable sort + scheduled_batch.generation_requests, + key=lambda req: int(req.py_batch_idx is not None), + ) + self._maybe_record_hang_diagnostic_phase( "preparing_resources", scheduled_batch) self.resource_manager.prepare_resources(scheduled_batch) @@ -6362,17 +6376,6 @@ def _executor_loop_overlap(self) -> None: # and let the later iteration to update it. should_process_previous_batch = can_queue or not can_queue_this_rank if can_queue: - - # The generation requests that do not have batch_idx - # need to be in front of the batch due to the assumptions - # made in model_engine.py::_forward_step. This is only important - # for disaggregated serving. For non-disaggregated serving, - # the generation requests always have batch_idx. - scheduled_batch.generation_requests = sorted( # stable sort - scheduled_batch.generation_requests, - key=lambda req: int(req.py_batch_idx is not None), - ) - if self.kv_cache_transceiver: # Return the first token to the client self._handle_first_token_response(scheduled_batch) @@ -9034,6 +9037,11 @@ def forward(scheduled_requests, resource_manager, new_tensors_device, if not self._loop_on_execution_stream: torch.cuda.current_stream().wait_stream(self.execution_stream) + if self._is_kv_manager_v2: + self.kv_cache_manager.update_forward_resources( + scheduled_requests, + wait_for_input_copy=self._wait_for_model_engine_input_copy) + self._kv_connector_wait_for_save() if route_capture is not None: route_capture.finish_forward() # R3: disarm between forwards diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index f0e465944c5a..aed446d3c33b 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -49,7 +49,7 @@ from ._util import (KvCacheCreator, _adjust_torch_mem_fraction, create_py_executor_instance, instantiate_sampler, is_disagg_enabled, is_mla) -from .config_utils import (is_hybrid_linear, is_minimax_m3, +from .config_utils import (is_decode_only, is_hybrid_linear, is_minimax_m3, resolve_cache_transceiver_config, uses_fp4_mla_attention, uses_vswa_kv_cache_layout) from .connectors.kv_cache_connector import KvCacheConnectorManager @@ -687,6 +687,13 @@ def _create_py_executor( "generation servers; aggregated serving with Helix is not " "supported yet (set TRTLLM_HELIX_AGG_EXPERIMENTAL=1 for " "development only).") + if (is_decode_only(llm_args.sparse_attention_config) + and not is_disagg_enabled(cache_transceiver_config)): + raise NotImplementedError( + "Sparse KV cache offload runs only on disaggregated generation " + "servers: it admits one input token per request per step, so it " + "cannot serve context requests. Set cache_transceiver_config, or " + "disable sparse_attention_config.enable_kv_cache_offload.") has_spec_drafter = False if spec_config is not None: @@ -1159,6 +1166,14 @@ def allocation_scope(current_stage: ExecutorMemoryType): # Get draft config for one-engine speculative decoding if available draft_config = getattr(model_engine.model, 'draft_config', None) + if not skip_est and is_decode_only(sparse_attention_config): + # The estimation dry run prefills dummy prompts, which a + # decode-only worker cannot execute. Size the KV cache from free + # memory instead, as with TRTLLM_SKIP_KV_CACHE_ESTIMATION=1. + logger.info("KV cache size estimation is not supported for a " + "decode-only worker, disable it.") + skip_est = True + kv_cache_creator = KvCacheCreator( model_engine=model_engine, mapping=mapping, diff --git a/tensorrt_llm/llmapi/__init__.py b/tensorrt_llm/llmapi/__init__.py index 8d1dfd245852..63f004f54b10 100644 --- a/tensorrt_llm/llmapi/__init__.py +++ b/tensorrt_llm/llmapi/__init__.py @@ -30,8 +30,8 @@ RuntimeMemoryStatus, SADecodingConfig, SAEnhancerConfig, SaveHiddenStatesDecodingConfig, SchedulerConfig, SkipSoftmaxAttentionConfig, SleepConfig, - TorchCompileConfig, TorchLlmArgs, - TriAttentionKvCacheCompressionConfig, + SparseKvCacheOffloadConfig, TorchCompileConfig, + TorchLlmArgs, TriAttentionKvCacheCompressionConfig, UserProvidedDecodingConfig) from .llm_utils import KvCacheRetentionConfig, QuantAlgo, QuantConfig from .mm_encoder import MultimodalEncoder @@ -99,6 +99,7 @@ 'DeepSeekSparseAttentionConfig', 'DeepSeekV4SparseAttentionConfig', 'MiniMaxM3SparseAttentionConfig', + 'SparseKvCacheOffloadConfig', 'SchedulingParams', 'SkipSoftmaxAttentionConfig', 'ColdPageQuantizationCompressionConfig', diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index d174a47f010e..85df2786ccc3 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -915,6 +915,40 @@ class GuidedDecodingBackend(Enum): description="The stop token ids for guided decoding config.") +class SparseKvCacheOffloadConfig(StrictBaseModel): + """Memory sizing for sparse attention KV cache offload. + + With offload, the KV history a sparse attention layer selects from lives in + host memory, and each request keeps a GPU cache of the units its selections + name. + """ + gpu_cache_factor: float = Field( + default=2.0, + ge=1.0, + status="prototype", + description= + "Size of each request's GPU cache for offloaded attention history, in " + "selections: the cache holds this many times the units one decode step " + "selects, plus the unit being written, rounded up to whole pages. A " + "larger cache keeps more history on the GPU between decode steps, so " + "fewer units are fetched from host memory, at the cost of GPU memory per " + "request.") + host_history_size: Optional[NonNegativeInt] = Field( + default=None, + status="prototype", + description= + "Pinned host memory in bytes, per rank, for offloaded attention history " + "and the staging buffers that publish it to the GPU. History pages share " + "slots with the GPU KV cache pools, so the GPU KV cache quota is lowered " + "when needed to stay within it. It is separate from " + "kv_cache_config.host_cache_size, the host tier that evicted KV cache " + "pages move to, and the two together must fit this rank's share of " + "available host memory and the pinned-memory limit. When unset, it " + "holds max_batch_size requests at kv_cache_config.avg_seq_len, or " + "max_seq_len when that is unset, capped by the host memory automatic " + "sizing may pin.") + + class BaseSparseAttentionConfig(StrictBaseModel): """Configuration for sparse attention.""" algorithm: str @@ -925,6 +959,14 @@ def supports_backend(self, backend: str) -> bool: """ return True + def get_kv_cache_offload_config( + self) -> Optional[SparseKvCacheOffloadConfig]: + """Return the KV cache offload sizing, or None without offload. + + Override in algorithms that can offload attention history to host memory. + """ + return None + def to_sparse_params(self, **kwargs): """Lower user-facing config into SparseParams.""" raise NotImplementedError( @@ -1340,7 +1382,9 @@ class DeepSeekSparseAttentionConfig(SeqLenAwareSparseAttentionConfig): default=True, description= "Whether to use the CuTe DSL paged MQA logits kernel on SM100-family " - "GPUs instead of C++ DeepGEMM. Falls back to DeepGEMM off the SM100 " + "GPUs instead of C++ DeepGEMM. Defaults to False with KV cache " + "offload enabled and True otherwise; an explicit value overrides " + "this default. Falls back to DeepGEMM off the SM100 " "family or when the cutlass-dsl package is missing.") q_split_threshold: int = Field( default=8192, @@ -1404,6 +1448,43 @@ class DeepSeekSparseAttentionConfig(SeqLenAwareSparseAttentionConfig): description= "Reuse the indexer Top-K across MTP draft steps instead of recomputing " "it each step. Defaults to the model's HF config value.") + enable_kv_cache_offload: bool = Field( + default=False, + status="prototype", + description= + "Offload the sparse attention KV history to host memory with KV cache " + "manager v2. Each decode step fetches the units the indexer selects " + "into a per-request GPU cache; the indexer cache remains on GPU. Only " + "a decode-only worker, such as the generation worker of a " + "disaggregated deployment, can offload.") + kv_cache_offload_config: SparseKvCacheOffloadConfig = Field( + default_factory=SparseKvCacheOffloadConfig, + status="prototype", + description="Memory sizing for KV cache offload. Applies only with " + "enable_kv_cache_offload.") + + @model_validator(mode="after") + def _validate_kv_cache_offload(self): + """Validate offload sizing and resolve the default indexer backend. + + ``index_topk`` may still be None here; the checkpoint fills it later. + """ + if self.enable_kv_cache_offload: + if self.index_topk is not None and self.index_topk <= 0: + raise ValueError( + "KV cache offload requires a positive index_topk.") + if "use_cute_dsl_paged_mqa_logits" not in self.model_fields_set: + self.use_cute_dsl_paged_mqa_logits = False + elif self.kv_cache_offload_config != SparseKvCacheOffloadConfig(): + raise ValueError("kv_cache_offload_config requires " + "enable_kv_cache_offload=True.") + return self + + def get_kv_cache_offload_config( + self) -> Optional[SparseKvCacheOffloadConfig]: + if not self.enable_kv_cache_offload: + return None + return self.kv_cache_offload_config @model_validator(mode="after") def _validate_indexer_k_dtype(self): @@ -1537,6 +1618,7 @@ def _value(name: str, default=None): pretrained_config, kwargs.get("layer_idx")), mtp_index_share=bool(_value("index_share_for_mtp_iteration", False)), + enable_kv_cache_offload=self.enable_kv_cache_offload, ) def to_sparse_metadata_params(self, **kwargs): @@ -1573,6 +1655,7 @@ def _value(name: str, default=None): has_shared_indexer_layers=has_shared_indexer_layers, mtp_index_share=bool(_value("index_share_for_mtp_iteration", False)), + enable_kv_cache_offload=self.enable_kv_cache_offload, ) @@ -1611,8 +1694,8 @@ class DeepSeekV4SparseAttentionConfig(DeepSeekSparseAttentionConfig): description= "Offload ratio-4 compressed attention KV history to host memory with " "KV cache manager v2. The indexer and sliding-window caches remain on " - "GPU. This feature is under development and currently cannot be " - "enabled for inference.") + "GPU. Only a decode-only worker, such as the generation worker of a " + "disaggregated deployment, can offload.") @field_validator("index_head_dim") @classmethod @@ -1632,7 +1715,8 @@ def normalize_compress_ratios(cls, compress_ratios): return [1 if ratio == 0 else ratio for ratio in compress_ratios] @model_validator(mode="after") - def validate_kv_cache_offload(self) -> "DeepSeekV4SparseAttentionConfig": + def _validate_deepseek_v4_kv_cache_offload( + self) -> "DeepSeekV4SparseAttentionConfig": if self.enable_kv_cache_offload: if 4 not in self.compress_ratios: raise ValueError( diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py index 3c710cb115a2..d907ac5fb823 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py @@ -164,7 +164,6 @@ class _KVCacheManagerConfigFieldSpec: ReusedBlocksByLevel = _cpp.ReusedBlocksByLevel Batch = _cpp.Batch IndexTableSpec = _cpp.IndexTableSpec -ScratchDesc = _cpp.ScratchDesc SlotDesc = _cpp.SlotDesc SlotDescVariant = _cpp.SlotDescVariant SsmLayerConfig = _cpp.SsmLayerConfig @@ -273,7 +272,6 @@ def typed_range(*args: int) -> range: "Priority", "ReuseGroupStatus", "ReuseScope", - "ScratchDesc", "Batch", "IndexTableSpec", "KVCacheIterationStatsDelta", diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi index 1211ecfd3d89..fefaadf158b4 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi @@ -798,7 +798,6 @@ class _KVCache: def resume( self, cuda_stream: CudaStream | None = None, - is_decoding: bool | None = None, optional_reuse_groups: Sequence[LayerGroupId] | None = None, ) -> bool: """Lock this cache's pages at their required levels. @@ -809,10 +808,6 @@ class _KVCache: Takes no sparse cache arena lease: that follows membership, so a request receiving its history costs the GPU cache nothing until it is seated. - - ``is_decoding`` is accepted and ignored. Resuming no longer distinguishes a decoding - request; the parameter is kept so callers that still pass it keep working, and goes - with the full sparse attention support. """ def prefetch(self, target: CacheLevel) -> bool: ... def get_swa_scratch_range(self, layer_group_id: LayerGroupId) -> HalfOpenRange: @@ -927,11 +922,6 @@ class HalfOpenRange: def __len__(self) -> int: ... def __eq__(self, other: object) -> bool: ... -class ScratchDesc: - range: HalfOpenRange - slot_ids: Sequence[int] - def __bool__(self) -> bool: ... - @dataclass(slots=True, frozen=True) class PageIndexConverter: scale: int @@ -943,7 +933,7 @@ class PageIndexConverter: self, base_indices: Sequence[int], index_mode: PageIndexMode | None = None, - scratch_range: HalfOpenRange | ScratchDesc | None = None, + scratch_range: HalfOpenRange | None = None, scratch_slots: Sequence[int] = (), ) -> list[int]: ... @@ -1289,10 +1279,17 @@ class Batch: 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. + Only flush a history that is final -- one the row is going to take up. A history + past the row's capacity is refused: those units were never written, so they have + no cell to write back from. + + Bounded by what was in flight, 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. The bound is aggregate: the units + crossed, summed over rows, must fit the per-layer write-back list, which is sized + from ``max_total_input_tokens``. One prefilling row may therefore use the budget + several decoding rows leave unspent, and a batch that exceeds it is refused before + any work is issued. Reads the index table as of the last :meth:`publish`, which covers the blocks in range because they were allocated by an earlier step. @@ -1508,10 +1505,12 @@ class Batch: 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. + The miss list and per-slot cell mapping this leaves behind are read by every + :meth:`fetch_resolved_pages` that follows, and they belong to the life cycle rather + than to the selection: the next resolve of *any* selection in the same life cycle + overwrites them. Fetch every member layer before resolving another selection there, + or those fetches copy the later selection's units. Both run on ``stream``, so that + ordering is the caller's only across streams. """ def fetch_resolved_pages( diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_introspection.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_introspection.py index 9dc70bf7797e..c4b564aaedca 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_introspection.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_introspection.py @@ -347,7 +347,11 @@ class SparseResolveResult(NamedTuple): """State and outputs of one sparse resolve launch.""" resident_ordinal: list[int] - lru_order: list[int] + #: Steps since each entry's last use, and whether that use found the entry already + #: resident. Eviction takes the oldest first and, within one age, what was fetched + #: before what was hit. ``-1`` and ``0`` where the entry is free. + entry_age: list[int] + entry_hit: 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. @@ -460,7 +464,7 @@ def sparse_flush_list( def sparse_resolve( resident_ordinal: list[int], - lru_order: list[int], + entry_age: list[int], num_layers: int, entries_per_layer: int, selected: list[int], @@ -477,6 +481,9 @@ def sparse_resolve( 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. + ``entry_age`` is the age in steps each entry starts at, or an empty list to leave every + entry at zero -- which leaves the eviction order among them arbitrary. + 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 @@ -492,7 +499,8 @@ def sparse_resolve( ( resident, - lru, + ages, + hits, out, miss_row, miss_ord, @@ -502,7 +510,7 @@ def sparse_resolve( in_flight, ) = _cpp().sparse_resolve( resident_ordinal, - lru_order, + entry_age, num_layers, entries_per_layer, selected, @@ -514,7 +522,8 @@ def sparse_resolve( ) return SparseResolveResult( list(resident), - list(lru), + list(ages), + list(hits), list(out), list(miss_row), list(miss_ord), diff --git a/tensorrt_llm/schemas/trtllm-serve-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-config.schema.json index 30578b5e3d0a..c8caf572fd8d 100644 --- a/tensorrt_llm/schemas/trtllm-serve-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-config.schema.json @@ -1302,6 +1302,34 @@ }, "DeepSeekSparseAttentionConfig": { "additionalProperties": false, + "allOf": [ + { + "else": { + "properties": { + "use_cute_dsl_paged_mqa_logits": { + "default": true + } + } + }, + "if": { + "properties": { + "enable_kv_cache_offload": { + "const": true + } + }, + "required": [ + "enable_kv_cache_offload" + ] + }, + "then": { + "properties": { + "use_cute_dsl_paged_mqa_logits": { + "default": false + } + } + } + } + ], "description": "Configuration for DeepSeek Sparse Attention.", "properties": { "algorithm": { @@ -1316,6 +1344,13 @@ "title": "Enable Heuristic Topk", "type": "boolean" }, + "enable_kv_cache_offload": { + "default": false, + "description": "Offload the sparse attention KV history to host memory with KV cache manager v2. Each decode step fetches the units the indexer selects into a per-request GPU cache; the indexer cache remains on GPU. Only a decode-only worker, such as the generation worker of a disaggregated deployment, can offload.", + "status": "prototype", + "title": "Enable Kv Cache Offload", + "type": "boolean" + }, "gvr_deterministic_topk": { "default": false, "description": "Make the self-sampling GVR Top-K deterministic: the selected index set is {score > tau} plus the smallest indices among the ties at tau, and the output order is a fixed function of the input (decode and prefill). Costs about 0.2-3 us per Top-K launch. Only takes effect with enable_heuristic_topk=True and use_self_sampling_topk=True; ignored otherwise.", @@ -1404,6 +1439,11 @@ "title": "Indexer Rope Interleave", "type": "boolean" }, + "kv_cache_offload_config": { + "$ref": "#/$defs/SparseKvCacheOffloadConfig", + "description": "Memory sizing for KV cache offload. Applies only with enable_kv_cache_offload.", + "status": "prototype" + }, "q_split_threshold": { "default": 8192, "description": "If number of packed tokens in prefill chunk exceeds this threshold, q tokens will be evenly distributed across ranks for indexer computation. If negative, q split will always be disabled.", @@ -1430,8 +1470,7 @@ "type": "boolean" }, "use_cute_dsl_paged_mqa_logits": { - "default": true, - "description": "Whether to use the CuTe DSL paged MQA logits kernel on SM100-family GPUs instead of C++ DeepGEMM. Falls back to DeepGEMM off the SM100 family or when the cutlass-dsl package is missing.", + "description": "Whether to use the CuTe DSL paged MQA logits kernel on SM100-family GPUs instead of C++ DeepGEMM. Defaults to False with KV cache offload enabled and True otherwise; an explicit value overrides this default. Falls back to DeepGEMM off the SM100 family or when the cutlass-dsl package is missing.", "title": "Use Cute Dsl Paged Mqa Logits", "type": "boolean" }, @@ -1459,6 +1498,34 @@ }, "DeepSeekV4SparseAttentionConfig": { "additionalProperties": false, + "allOf": [ + { + "else": { + "properties": { + "use_cute_dsl_paged_mqa_logits": { + "default": true + } + } + }, + "if": { + "properties": { + "enable_kv_cache_offload": { + "const": true + } + }, + "required": [ + "enable_kv_cache_offload" + ] + }, + "then": { + "properties": { + "use_cute_dsl_paged_mqa_logits": { + "default": false + } + } + } + } + ], "description": "Configuration for DeepSeek-V4 Sparse Attention.", "properties": { "algorithm": { @@ -1483,7 +1550,7 @@ }, "enable_kv_cache_offload": { "default": false, - "description": "Offload ratio-4 compressed attention KV history to host memory with KV cache manager v2. The indexer and sliding-window caches remain on GPU. This feature is under development and currently cannot be enabled for inference.", + "description": "Offload ratio-4 compressed attention KV history to host memory with KV cache manager v2. The indexer and sliding-window caches remain on GPU. Only a decode-only worker, such as the generation worker of a disaggregated deployment, can offload.", "status": "prototype", "title": "Enable Kv Cache Offload", "type": "boolean" @@ -1576,6 +1643,11 @@ "title": "Indexer Rope Interleave", "type": "boolean" }, + "kv_cache_offload_config": { + "$ref": "#/$defs/SparseKvCacheOffloadConfig", + "description": "Memory sizing for KV cache offload. Applies only with enable_kv_cache_offload.", + "status": "prototype" + }, "q_split_threshold": { "default": 8192, "description": "If number of packed tokens in prefill chunk exceeds this threshold, q tokens will be evenly distributed across ranks for indexer computation. If negative, q split will always be disabled.", @@ -1602,8 +1674,7 @@ "type": "boolean" }, "use_cute_dsl_paged_mqa_logits": { - "default": true, - "description": "Whether to use the CuTe DSL paged MQA logits kernel on SM100-family GPUs instead of C++ DeepGEMM. Falls back to DeepGEMM off the SM100 family or when the cutlass-dsl package is missing.", + "description": "Whether to use the CuTe DSL paged MQA logits kernel on SM100-family GPUs instead of C++ DeepGEMM. Defaults to False with KV cache offload enabled and True otherwise; an explicit value overrides this default. Falls back to DeepGEMM off the SM100 family or when the cutlass-dsl package is missing.", "title": "Use Cute Dsl Paged Mqa Logits", "type": "boolean" }, @@ -6001,6 +6072,37 @@ "title": "SleepConfig", "type": "object" }, + "SparseKvCacheOffloadConfig": { + "additionalProperties": false, + "description": "Memory sizing for sparse attention KV cache offload.\n\nWith offload, the KV history a sparse attention layer selects from lives in\nhost memory, and each request keeps a GPU cache of the units its selections\nname.", + "properties": { + "gpu_cache_factor": { + "default": 2.0, + "description": "Size of each request's GPU cache for offloaded attention history, in selections: the cache holds this many times the units one decode step selects, plus the unit being written, rounded up to whole pages. A larger cache keeps more history on the GPU between decode steps, so fewer units are fetched from host memory, at the cost of GPU memory per request.", + "minimum": 1.0, + "status": "prototype", + "title": "Gpu Cache Factor", + "type": "number" + }, + "host_history_size": { + "anyOf": [ + { + "minimum": 0, + "type": "integer" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Pinned host memory in bytes, per rank, for offloaded attention history and the staging buffers that publish it to the GPU. History pages share slots with the GPU KV cache pools, so the GPU KV cache quota is lowered when needed to stay within it. It is separate from kv_cache_config.host_cache_size, the host tier that evicted KV cache pages move to, and the two together must fit this rank's share of available host memory and the pinned-memory limit. When unset, it holds max_batch_size requests at kv_cache_config.avg_seq_len, or max_seq_len when that is unset, capped by the host memory automatic sizing may pin.", + "status": "prototype", + "title": "Host History Size" + } + }, + "title": "SparseKvCacheOffloadConfig", + "type": "object" + }, "TelemetryConfig": { "additionalProperties": false, "description": "Telemetry configuration for usage data collection.\n\nControls opt-out behavior and tracks which entry point invoked TRT-LLM.", diff --git a/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json index ea3619438431..0aff994712e2 100644 --- a/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json @@ -1302,6 +1302,34 @@ }, "DeepSeekSparseAttentionConfig": { "additionalProperties": false, + "allOf": [ + { + "else": { + "properties": { + "use_cute_dsl_paged_mqa_logits": { + "default": true + } + } + }, + "if": { + "properties": { + "enable_kv_cache_offload": { + "const": true + } + }, + "required": [ + "enable_kv_cache_offload" + ] + }, + "then": { + "properties": { + "use_cute_dsl_paged_mqa_logits": { + "default": false + } + } + } + } + ], "description": "Configuration for DeepSeek Sparse Attention.", "properties": { "algorithm": { @@ -1316,6 +1344,13 @@ "title": "Enable Heuristic Topk", "type": "boolean" }, + "enable_kv_cache_offload": { + "default": false, + "description": "Offload the sparse attention KV history to host memory with KV cache manager v2. Each decode step fetches the units the indexer selects into a per-request GPU cache; the indexer cache remains on GPU. Only a decode-only worker, such as the generation worker of a disaggregated deployment, can offload.", + "status": "prototype", + "title": "Enable Kv Cache Offload", + "type": "boolean" + }, "gvr_deterministic_topk": { "default": false, "description": "Make the self-sampling GVR Top-K deterministic: the selected index set is {score > tau} plus the smallest indices among the ties at tau, and the output order is a fixed function of the input (decode and prefill). Costs about 0.2-3 us per Top-K launch. Only takes effect with enable_heuristic_topk=True and use_self_sampling_topk=True; ignored otherwise.", @@ -1404,6 +1439,11 @@ "title": "Indexer Rope Interleave", "type": "boolean" }, + "kv_cache_offload_config": { + "$ref": "#/$defs/SparseKvCacheOffloadConfig", + "description": "Memory sizing for KV cache offload. Applies only with enable_kv_cache_offload.", + "status": "prototype" + }, "q_split_threshold": { "default": 8192, "description": "If number of packed tokens in prefill chunk exceeds this threshold, q tokens will be evenly distributed across ranks for indexer computation. If negative, q split will always be disabled.", @@ -1430,8 +1470,7 @@ "type": "boolean" }, "use_cute_dsl_paged_mqa_logits": { - "default": true, - "description": "Whether to use the CuTe DSL paged MQA logits kernel on SM100-family GPUs instead of C++ DeepGEMM. Falls back to DeepGEMM off the SM100 family or when the cutlass-dsl package is missing.", + "description": "Whether to use the CuTe DSL paged MQA logits kernel on SM100-family GPUs instead of C++ DeepGEMM. Defaults to False with KV cache offload enabled and True otherwise; an explicit value overrides this default. Falls back to DeepGEMM off the SM100 family or when the cutlass-dsl package is missing.", "title": "Use Cute Dsl Paged Mqa Logits", "type": "boolean" }, @@ -1459,6 +1498,34 @@ }, "DeepSeekV4SparseAttentionConfig": { "additionalProperties": false, + "allOf": [ + { + "else": { + "properties": { + "use_cute_dsl_paged_mqa_logits": { + "default": true + } + } + }, + "if": { + "properties": { + "enable_kv_cache_offload": { + "const": true + } + }, + "required": [ + "enable_kv_cache_offload" + ] + }, + "then": { + "properties": { + "use_cute_dsl_paged_mqa_logits": { + "default": false + } + } + } + } + ], "description": "Configuration for DeepSeek-V4 Sparse Attention.", "properties": { "algorithm": { @@ -1483,7 +1550,7 @@ }, "enable_kv_cache_offload": { "default": false, - "description": "Offload ratio-4 compressed attention KV history to host memory with KV cache manager v2. The indexer and sliding-window caches remain on GPU. This feature is under development and currently cannot be enabled for inference.", + "description": "Offload ratio-4 compressed attention KV history to host memory with KV cache manager v2. The indexer and sliding-window caches remain on GPU. Only a decode-only worker, such as the generation worker of a disaggregated deployment, can offload.", "status": "prototype", "title": "Enable Kv Cache Offload", "type": "boolean" @@ -1576,6 +1643,11 @@ "title": "Indexer Rope Interleave", "type": "boolean" }, + "kv_cache_offload_config": { + "$ref": "#/$defs/SparseKvCacheOffloadConfig", + "description": "Memory sizing for KV cache offload. Applies only with enable_kv_cache_offload.", + "status": "prototype" + }, "q_split_threshold": { "default": 8192, "description": "If number of packed tokens in prefill chunk exceeds this threshold, q tokens will be evenly distributed across ranks for indexer computation. If negative, q split will always be disabled.", @@ -1602,8 +1674,7 @@ "type": "boolean" }, "use_cute_dsl_paged_mqa_logits": { - "default": true, - "description": "Whether to use the CuTe DSL paged MQA logits kernel on SM100-family GPUs instead of C++ DeepGEMM. Falls back to DeepGEMM off the SM100 family or when the cutlass-dsl package is missing.", + "description": "Whether to use the CuTe DSL paged MQA logits kernel on SM100-family GPUs instead of C++ DeepGEMM. Defaults to False with KV cache offload enabled and True otherwise; an explicit value overrides this default. Falls back to DeepGEMM off the SM100 family or when the cutlass-dsl package is missing.", "title": "Use Cute Dsl Paged Mqa Logits", "type": "boolean" }, @@ -7576,6 +7647,37 @@ "title": "SleepConfig", "type": "object" }, + "SparseKvCacheOffloadConfig": { + "additionalProperties": false, + "description": "Memory sizing for sparse attention KV cache offload.\n\nWith offload, the KV history a sparse attention layer selects from lives in\nhost memory, and each request keeps a GPU cache of the units its selections\nname.", + "properties": { + "gpu_cache_factor": { + "default": 2.0, + "description": "Size of each request's GPU cache for offloaded attention history, in selections: the cache holds this many times the units one decode step selects, plus the unit being written, rounded up to whole pages. A larger cache keeps more history on the GPU between decode steps, so fewer units are fetched from host memory, at the cost of GPU memory per request.", + "minimum": 1.0, + "status": "prototype", + "title": "Gpu Cache Factor", + "type": "number" + }, + "host_history_size": { + "anyOf": [ + { + "minimum": 0, + "type": "integer" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Pinned host memory in bytes, per rank, for offloaded attention history and the staging buffers that publish it to the GPU. History pages share slots with the GPU KV cache pools, so the GPU KV cache quota is lowered when needed to stay within it. It is separate from kv_cache_config.host_cache_size, the host tier that evicted KV cache pages move to, and the two together must fit this rank's share of available host memory and the pinned-memory limit. When unset, it holds max_batch_size requests at kv_cache_config.avg_seq_len, or max_seq_len when that is unset, capped by the host memory automatic sizing may pin.", + "status": "prototype", + "title": "Host History Size" + } + }, + "title": "SparseKvCacheOffloadConfig", + "type": "object" + }, "TelemetryConfig": { "additionalProperties": false, "description": "Telemetry configuration for usage data collection.\n\nControls opt-out behavior and tracks which entry point invoked TRT-LLM.", diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index 89f32273d0ef..bf96168f596b 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -1527,6 +1527,16 @@ "kind": "categorical", "path": "sparse_attention_config.kt_cache_dtype" }, + { + "capture_policy": "float", + "kind": "value", + "path": "sparse_attention_config.kv_cache_offload_config.gpu_cache_factor" + }, + { + "capture_policy": "int|none", + "kind": "value", + "path": "sparse_attention_config.kv_cache_offload_config.host_history_size" + }, { "capture_policy": "int|none", "kind": "value", diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 6ce9b590f6d2..64eb2abf2c28 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -138,6 +138,7 @@ l0_b200: - unittest/_torch/modeling/test_glm5_next_contracts.py - unittest/llmapi/test_deepseek_v4_tokenizer.py - unittest/_torch/modules/test_mhc.py + - unittest/_torch/modules/test_top_k.py - unittest/_torch/modules/test_engram.py - unittest/_torch/modules/test_gate_up_swiglu_quack.py - unittest/_torch/modules/test_multi_stream_utils.py diff --git a/tests/integration/test_lists/test-db/l0_gb300.yml b/tests/integration/test_lists/test-db/l0_gb300.yml index ce21a01c9ddc..ddc455f8a4fe 100644 --- a/tests/integration/test_lists/test-db/l0_gb300.yml +++ b/tests/integration/test_lists/test-db/l0_gb300.yml @@ -43,6 +43,7 @@ l0_gb300: - accuracy/test_llm_api_pytorch.py::TestQwen3_8_Flash_Next::test_nvfp4_adp4_mtp3_trtllm_ple_offload - accuracy/test_llm_api_pytorch.py::TestQwen3_8_Flash_Next::test_nvfp4_adp4_megamoe_cutedsl_ple_offload - accuracy/test_disaggregated_serving.py::TestQwen3_8_Flash_Next::test_fp8_nixl_python[prefix_cache] + - unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py::TestSparseResolve::test_colliding_selections_preserve_slots_and_eviction - unittest/_torch/moe/kernels/parallel - unittest/_torch/kernels/parallel TIMEOUT (90) # Kimi K3 KDA unit tests (single GPU, random weights, no checkpoint). diff --git a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_kernel.py b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_kernel.py index 7f680d440d9e..99f5b71ecd80 100644 --- a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_kernel.py +++ b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_kernel.py @@ -11,6 +11,8 @@ import triton from utils.util import skip_pre_blackwell +from tensorrt_llm._utils import TensorWrapper, convert_to_torch_tensor + _CUDA_SUPPORTED_HEAD_DIMS = (128, 512) @@ -1874,6 +1876,181 @@ def test_chunked_prefill(compress_ratio, head_dim, overlap, batch_size, start_po _HAS_POSTPROCESS_SCATTER = False +@pytest.mark.parametrize( + "cache_dtype,residual_dim", + [ + (0, 0), + (1, 0), + pytest.param(4, 0, marks=skip_pre_blackwell), + pytest.param(4, 64, marks=skip_pre_blackwell), + ], +) +@pytest.mark.parametrize("graph", [False, True]) +@pytest.mark.skipif(not _HAS_POSTPROCESS_SCATTER, reason="Compressor CUDA ops are required") +def test_sparse_compressor_scatter_matches_paged_rows(cache_dtype, residual_dim, graph): + """One-row page adaptation writes paired arena rows and masks inactive requests.""" + device = "cuda" + head_dim, batch = 512, 4 + dtype = {0: torch.bfloat16, 1: torch.float8_e4m3fn, 4: torch.uint8}[cache_dtype] + row_size = (head_dim + residual_dim) // 2 if cache_dtype == 4 else head_dim + x = torch.randn(batch, head_dim, device=device, dtype=torch.bfloat16) + weights = torch.ones(head_dim, device=device, dtype=torch.bfloat16) + cos_sin = torch.randn(32, 2, 32, device=device) + positions = torch.arange(batch, device=device, dtype=torch.int32) + num_outputs = torch.ones(batch, device=device, dtype=torch.int32) + cu = torch.arange(batch + 1, device=device, dtype=torch.int32) + mask = torch.tensor([True, True, True, False], device=device) + reference = torch.zeros(batch, 1, row_size, dtype=dtype, device=device) + arena = torch.zeros(11, 1, row_size, dtype=dtype, device=device) + scale_reference = ( + torch.zeros( + batch, 1, (head_dim + residual_dim) // 16, dtype=torch.float8_e4m3fn, device=device + ) + if cache_dtype == 4 + else None + ) + scale_arena = ( + torch.zeros( + 11, 1, (head_dim + residual_dim) // 16, dtype=torch.float8_e4m3fn, device=device + ) + if cache_dtype == 4 + else None + ) + # Strided columns from a native output whose capacity exceeds this graph's batch. + write_mapping = torch.tensor( + [[7, 7], [-1, -1], [3, 3], [0, 0], [10, 10], [-1, -1]], dtype=torch.int32, device=device + ) + starts = torch.zeros(batch, dtype=torch.int32, device=device) + blocks = torch.arange(batch, dtype=torch.int32, device=device).view(batch, 1) + + def scatter(data, scales, table, counts, storage_positions): + torch.ops.trtllm.compressor_postprocess_scatter( + x, + None, + weights, + 1e-5, + cos_sin, + positions, + 448, + 64, + data, + scales, + counts, + cu, + storage_positions, + table, + mask, + 1, + cache_dtype, + 1.0, + residual_dim, + True, + None, + None, + ) + + scatter(reference, scale_reference, blocks, num_outputs, starts) + starts.fill_(1000) # Sparse writes must not use logical positions as physical addresses. + sparse_starts = torch.zeros_like(starts) + data_view = convert_to_torch_tensor(TensorWrapper(arena.data_ptr(), dtype, arena.shape)) + scale_view = ( + convert_to_torch_tensor( + TensorWrapper(scale_arena.data_ptr(), scale_arena.dtype, scale_arena.shape) + ) + if scale_arena is not None + else None + ) + + def sparse_scatter(): + indices = write_mapping[:batch, 0] + scatter( + data_view, + scale_view, + indices.contiguous().view(batch, 1), + torch.where(indices >= 0, num_outputs, 0), + sparse_starts, + ) + + sparse_scatter() + if graph: + capture = torch.cuda.CUDAGraph() + with torch.cuda.graph(capture): + sparse_scatter() + arena.zero_() + if scale_arena is not None: + scale_arena.zero_() + # Replay changes the physical cells without changing pointer arguments. + write_mapping[0] = 8 + write_mapping[1] = 4 + write_mapping[2] = -1 + capture.replay() + expected = torch.zeros_like(arena) + expected[8 if graph else 7] = reference[0] + expected[4 if graph else 3] = reference[1 if graph else 2] + torch.testing.assert_close(arena.view(torch.uint8), expected.view(torch.uint8)) + if scale_arena is not None: + expected_scales = torch.zeros_like(scale_arena) + expected_scales[8 if graph else 7] = scale_reference[0] + expected_scales[4 if graph else 3] = scale_reference[1 if graph else 2] + torch.testing.assert_close(scale_arena.view(torch.uint8), expected_scales.view(torch.uint8)) + + +@pytest.mark.parametrize("residual_dim", [0, 64]) +@pytest.mark.parametrize("global_scale", [1.0, 0.25]) +@pytest.mark.parametrize("graph", [False, True]) +@skip_pre_blackwell +def test_sparse_nvfp4_gather_uses_shared_pool_indices(residual_dim, global_scale, graph): + """One index addresses both arena pools through the existing dequantization kernel.""" + head_dim = 512 + row_bytes = (head_dim + residual_dim) // 2 + scale_bytes = (head_dim + residual_dim) // 16 + data = torch.randint(0, 256, (11, row_bytes), dtype=torch.uint8, device="cuda") + scales = (torch.rand(11, scale_bytes, device="cuda") + 0.1).to(torch.float8_e4m3fn) + data_indices = torch.tensor([[7, 2, -1, 0], [3, 8, 4, 11]], dtype=torch.int32, device="cuda") + original_indices = data_indices.clone() + data_view = convert_to_torch_tensor( + TensorWrapper(data.data_ptr(), data.dtype, (11, 1, row_bytes)) + ) + scale_view = convert_to_torch_tensor( + TensorWrapper(scales.data_ptr(), scales.dtype, (11, 1, scale_bytes)) + ) + out = torch.zeros(2, 4, head_dim, dtype=torch.float8_e4m3fn, device="cuda") + scale = torch.tensor([global_scale], device="cuda") + + def gather(): + torch.ops.trtllm.nvfp4_mla_kv_cache_gather_direct( + data_view, scale_view, data_indices, out, scale, residual_dim, data_view.shape[0] + ) + + gather() + if graph: + data_indices.copy_(original_indices) + capture = torch.cuda.CUDAGraph() + with torch.cuda.graph(capture): + gather() + # The gather leaves rows whose index is out of range untouched, so the + # replay's expectation of zeros there only holds from a cleared output. + out.zero_() + original_indices.copy_( + torch.tensor([[10, -1, 1, 4], [2, 11, 5, 0]], dtype=torch.int32, device="cuda") + ) + data_indices.copy_(original_indices) + capture.replay() + valid = (original_indices >= 0) & (original_indices < data_view.shape[0]) + safe_indices = torch.where(valid, original_indices, 0).flatten().long() + paired_data = data[safe_indices].contiguous() + paired_scales = scales.view(torch.uint8)[safe_indices].contiguous().view(torch.float8_e4m3fn) + expected_indices = ( + torch.arange(8, device="cuda", dtype=torch.int32).view(2, 4).masked_fill(~valid, -1) + ) + expected_out = torch.zeros_like(out) + torch.ops.trtllm.nvfp4_mla_kv_cache_gather_direct( + paired_data, paired_scales, expected_indices, expected_out, scale, residual_dim, 8 + ) + torch.testing.assert_close(data_indices, expected_indices) + torch.testing.assert_close(out.view(torch.uint8), expected_out.view(torch.uint8)) + + def _scatter_reference_to_paged_cache( kv: torch.Tensor, num_comp_tokens: torch.Tensor, diff --git a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_module.py b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_module.py index 85efe8539f60..38a0852633b4 100644 --- a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_module.py +++ b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_module.py @@ -90,6 +90,8 @@ def rotate_activation(x: torch.Tensor) -> torch.Tensor: class DummyAttentionMetadata: """Dummy attention metadata for testing Compressor.forward.""" + sparse_offload_state = None + def __init__( self, num_contexts: int, @@ -1947,10 +1949,66 @@ def test_compressor_wkv_gate_uses_checkpoint_dtype(): assert compressor.wkv_gate.weight.dtype == DTYPE -def _run_compressor_with_fake_postprocess(monkeypatch, kv_cache_dtype: str, is_indexer: bool): +def _run_compressor_with_fake_postprocess( + monkeypatch, kv_cache_dtype: str, is_indexer: bool, sparse=False, sparse_write_index=7 +): compressor = _create_small_compressor(kv_cache_dtype, is_indexer) metadata = _create_minimal_metadata(compressor) seen = {} + if sparse: + from types import SimpleNamespace + from unittest.mock import Mock + + metadata.num_contexts = 0 + metadata.num_generations = 1 + metadata.num_ctx_tokens = 0 + metadata.num_tokens = 1 + metadata.num_gen_tokens_per_seq = 1 + metadata.cu_seq_lens_cuda = torch.tensor([0, 1], device=DEVICE, dtype=torch.int32) + # Write rows by Batch row: the one request sits in row 2, not at its input position. + writes = torch.full((4, 2), -1, device=DEVICE, dtype=torch.int32) + writes[2] = sparse_write_index + is_nvfp4 = kv_cache_dtype == "nvfp4" + descriptor = SimpleNamespace( + data_buffer_index=1, + scale_buffer_index=0 if is_nvfp4 else None, + ) + metadata.sparse_offload_state = SimpleNamespace( + layers={0: descriptor}, + # Neither buffer shares a slice here, so each reads at its own index. + out_slots={0: (0, 1)}, + batch_rows=torch.tensor([2, -1, -1, -1], device=DEVICE, dtype=torch.int32), + scatter_start_positions=torch.zeros(4, device=DEVICE, dtype=torch.int32), + ) + data = torch.empty( + ( + 12, + 1, + (compressor.head_dim + compressor.nvfp4_residual_dim) // 2 + if is_nvfp4 + else compressor.head_dim, + ), + device=DEVICE, + dtype=torch.uint8 + if is_nvfp4 + else (torch.float8_e4m3fn if kv_cache_dtype == "fp8_pertensor" else DTYPE), + ) + scales = ( + torch.empty( + (12, 1, (compressor.head_dim + compressor.nvfp4_residual_dim) // 16), + device=DEVICE, + dtype=torch.float8_e4m3fn, + ) + if is_nvfp4 + else None + ) + metadata.kv_cache_manager.get_sparse_compress_buffers = Mock(return_value=(data, scales)) + seen["arena"] = (data, scales) + seen["zero_positions"] = metadata.sparse_offload_state.scatter_start_positions + metadata.kv_cache_manager.prepare_sparse_compressor_write_indices = Mock( + return_value=writes + ) + seen["prepare"] = metadata.kv_cache_manager.prepare_sparse_compressor_write_indices def fake_prefill_reduction(*args): seen["kv_score_dtype"] = args[0].dtype @@ -1958,7 +2016,8 @@ def fake_prefill_reduction(*args): kv_comp.fill_(0.25) def fake_paged_kv_compress(*args): - raise AssertionError("generation compression path should not run") + assert sparse + args[6].fill_(0.25) def fake_postprocess_scatter( kv_comp, @@ -1991,6 +2050,11 @@ def fake_postprocess_scatter( seen["kv_cache_scale"] = kv_cache_scale seen["nvfp4_global_scale"] = nvfp4_global_scale seen["nvfp4_residual_dim"] = nvfp4_residual_dim + seen["kv_cache"] = kv_cache + seen["block_table"] = block_table + seen["tokens_per_block"] = tokens_per_block + seen["start_pos"] = start_pos + seen["num_outputs"] = num_comp_tokens if kv_out is not None: kv_out.fill_(0.5) if quant_output is not None: @@ -2019,12 +2083,46 @@ def fake_postprocess_scatter( raising=False, ) - x = torch.randn(compressor.compress_ratio, 16, device=DEVICE, dtype=DTYPE) + x = torch.randn(metadata.num_tokens, 16, device=DEVICE, dtype=DTYPE) with torch.no_grad(): output = compressor(x, metadata) return output, seen +@pytest.mark.parametrize( + "kv_cache_dtype,is_indexer", + [ + ("default", False), + ("fp8_pertensor", False), + ("nvfp4", False), + ("default", True), + ("fp8_blockwise", True), + ("mxfp4", True), + ], +) +@pytest.mark.parametrize("write_index", [7, -1]) +def test_compressor_routes_only_main_writes_to_arena( + monkeypatch, kv_cache_dtype, is_indexer, write_index +): + _, seen = _run_compressor_with_fake_postprocess( + monkeypatch, kv_cache_dtype, is_indexer, sparse=True, sparse_write_index=write_index + ) + if is_indexer: + seen["prepare"].assert_not_called() + assert seen["kv_cache"] is not None + assert seen["block_table"] is not None + return + seen["prepare"].assert_called_once() + assert seen["kv_cache"] is seen["arena"][0] + assert seen["kv_cache_scale"] is seen["arena"][1] + assert seen["tokens_per_block"] == 1 + assert seen["block_table"].shape == (1, 1) and seen["block_table"].is_contiguous() + assert seen["block_table"].tolist() == [[write_index]] + assert seen["start_pos"].data_ptr() == seen["zero_positions"].data_ptr() + assert (seen["start_pos"] == 0).all() + assert (seen["num_outputs"] == int(write_index >= 0)).all() + + def test_main_compressor_does_not_materialize_postprocess_output(monkeypatch): """The fused kernel writes main compressor output directly to paged cache. 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 43820e811229..b97ff998088d 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 @@ -21,6 +21,7 @@ import torch from utils.util import skip_pre_blackwell +from tensorrt_llm._torch.attention.backends.sparse import offload as sparse_offload from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4 import DeepseekV4CacheManager from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4 import ( cache_manager as deepseek_v4_cache, @@ -30,6 +31,7 @@ DeepseekV4AttentionType, compress_ratio_has_attention, ) +from tensorrt_llm._torch.attention.backends.sparse.offload import SparseOffloadPaddingRegion from tensorrt_llm._torch.disaggregation.native.peer import PeerRegistrar from tensorrt_llm._torch.disaggregation.native.rank_info import RankInfo from tensorrt_llm._torch.disaggregation.resource.kv_extractor import ( @@ -41,10 +43,17 @@ Nvfp4ColdPageQuantizationCompression, ) from tensorrt_llm._torch.pyexecutor._util import CacheCost +from tensorrt_llm._torch.pyexecutor.kv_cache import kv_cache_manager_v2 as kv_cache_v2_module from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import ( + CUDA_GRAPH_DUMMY_REQUEST_ID, BlockReusePolicy, KVCacheManagerV2, ) +from tensorrt_llm._torch.pyexecutor.kv_cache.sparse_offload_memory import ( + sparse_arena_bytes, + sparse_arena_shapes, + sparse_history_bytes_per_block, +) from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, LlmRequestState from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests from tensorrt_llm._utils import binding_to_torch_dtype @@ -56,18 +65,23 @@ DeepSeekV4SparseAttentionConfig, DraftTargetDecodingConfig, KvCacheConfig, + MTPDecodingConfig, + SparseKvCacheOffloadConfig, ) from tensorrt_llm.mapping import Mapping from tensorrt_llm.runtime.kv_cache_manager_v2 import ( BAD_PAGE_INDEX, + AttentionLayerConfig, BatchDesc, BufferConfig, + BufferId, CacheTier, GpuCacheTierConfig, - HostCacheTierConfig, + IndexTableSpec, KVCacheDesc, KVCacheManagerConfig, PageIndexMode, + SparseCacheConfig, _introspection, ) @@ -81,18 +95,61 @@ class _SparseRuntime: - """Test double for the declared APIs; no sparse transfer implementation.""" + """Expose native buffer configuration without allocating a cache manager.""" - def __init__(self): + def __init__(self, tokens_per_block=128): self.groups = {41: 13, 19: 13, 42: 9, 20: 9, 43: 5, 21: 5, 7: 14} self.sparse = {41, 19} - self.converters = { - 41: SimpleNamespace(scale=5, expansion=1, layer_offset=17), - 19: SimpleNamespace(scale=5, expansion=1, layer_offset=23), - } - self.pool_bases = {41: 100000, 19: 100000} - self.calls = [] - self.uploads = [] + self.expansions = {41: tokens_per_block // 4, 19: tokens_per_block // 4} + self.pool_indices = {} + self.layers = {} + for layer_id in (41, 19): + self.configure_buffers(layer_id, [(DeepseekV4AttentionType.COMPRESS.role, 0)]) + + def sparse_index_table_specs(self): + """One request per (layer, pool) run, as the native manager derives them. + + The fake's sparse buffers are one per layer, so each is its own run. + """ + return [ + IndexTableSpec([BufferId(layer_id, role)]) + for layer_id in sorted(self.sparse) + for role in (DeepseekV4AttentionType.COMPRESS.role,) + ] + + def configure_buffers(self, layer_id, roles_and_pools): + self.layers[layer_id] = AttentionLayerConfig( + layer_id=layer_id, + buffers=[ + BufferConfig(role=role, size=1024, is_sparse=True) for role, _ in roles_and_pools + ], + sparse_cache=SparseCacheConfig( + num_selected_blocks=8, max_cached_units=16, unit_tokens=4 + ), + ) + for role, pool in roles_and_pools: + self.pool_indices[BufferId(layer_id, role)] = pool + + @property + def init_config(self): + return SimpleNamespace( + layers=list(self.layers.values()), + constraints=[ + SimpleNamespace(kv_caches=[None] * 2), + SimpleNamespace(kv_caches=[None] * 6), + ], + ) + + @property + def pool_group_descs(self): + count = max(self.pool_indices.values()) + 1 + coalesced = [SimpleNamespace(buffer_ids=[]) for _ in range(count)] + # Storage's buffer order need not be declaration order. The adapter + # must use this only to locate the pool, then preserve declaration order. + for buffer_id, pool in reversed(list(self.pool_indices.items())): + coalesced[pool].buffer_ids.append(buffer_id) + variant = SimpleNamespace(coalesced_buffers=coalesced) + return [SimpleNamespace(slot_desc=SimpleNamespace(variants=[variant]))] def is_sparse(self, layer_id, role): return layer_id in self.sparse @@ -100,35 +157,22 @@ def is_sparse(self, layer_id, role): def get_layer_group_id(self, layer_id): return self.groups[layer_id] - def get_page_index_converter(self, layer_id, role): - return self.converters[layer_id] - - def get_mem_pool_base_address(self, layer_id, role, index_mode): - if index_mode == PageIndexMode.SHARED: - return self.pool_bases[layer_id] + self.converters[layer_id].layer_offset * 4096 - assert index_mode == PageIndexMode.PER_LAYER - return self.pool_bases[layer_id] - - def get_page_index_upper_bound(self, layer_id, role): - return 320 - self.converters[layer_id].layer_offset + def get_sparse_cache_units_per_page(self, layer_id): + return self.expansions[layer_id] - def copy_base_page_indices_to_device(self, kv_caches, layer_group_id, out, stream): - assert stream == torch.cuda.current_stream().cuda_stream - assert out.shape[0] == len(kv_caches) - self.calls.append( - ([cache.id for cache in kv_caches], layer_group_id, out.data_ptr(), stream) - ) - table = torch.full(out.shape, -1, dtype=torch.int32, pin_memory=True) - for row, cache in enumerate(kv_caches): - table[row, : cache.num_blocks] = torch.tensor(cache.pages, dtype=torch.int32) - self.uploads.append(table) # Keep each immutable pinned source alive through its upload. - out.copy_(table, non_blocking=True) + def get_sparse_cache_base_address(self, layer_id, role): + return 100000 + layer_id * 4096 + self.pool_indices[BufferId(layer_id, role)] * 1024 def _manager(tokens_per_block=128): manager = object.__new__(DeepseekV4CacheManager) manager.kv_cache_type = CacheTypeCpp.SELFKONLY manager._enable_kv_cache_offload = True + manager._sparse_batch = None + manager._sparse_batch_members = {} + manager._sparse_batch_prepared = False + manager.max_batch_size = 4 + manager._sparse_offload_topk = 8 manager.max_beam_width = 1 manager.max_draft_len = 0 manager.max_total_draft_tokens = 0 @@ -137,6 +181,10 @@ def _manager(tokens_per_block=128): manager._has_cp_helix = False manager.is_draft = False manager._use_nvfp4_compress = False + manager._nvfp4_residual_dim = 0 + manager._sparse_arena_buffers = {} + manager.dtype = DataType.BF16 + manager.head_dim = 512 manager.use_fp8_ds_mla = False manager.pp_layers = [2, 3, 4] manager._compress_ratios = [1, 4, 4, 128, 4] @@ -151,7 +199,11 @@ def _manager(tokens_per_block=128): (4, DeepseekV4AttentionType.SWA): 21, (3, DeepseekV4AttentionType.COMPRESS): 7, } - manager.impl = _SparseRuntime() + manager.impl = _SparseRuntime(tokens_per_block) + # The arenas are the Batch's, and this manager has no native one; answer for it. + manager._sparse_arena_base = lambda layer_id, role: manager.impl.get_sparse_cache_base_address( + layer_id, role + ) manager.kv_cache_map = { 10: SimpleNamespace(id=10, history_length=tokens_per_block - 1, pages=[0, 9, -1]), 20: SimpleNamespace(id=20, history_length=tokens_per_block, pages=[0, 0, 12]), @@ -165,46 +217,432 @@ def _manager(tokens_per_block=128): @pytest.mark.cpu_only -def test_sparse_offload_descriptors_use_runtime_ids(): - manager = _manager() +@pytest.mark.parametrize("tokens_per_block", [128, 256]) +def test_sparse_offload_descriptors_use_runtime_ids(tokens_per_block): + manager = _manager(tokens_per_block) descriptors = manager.get_sparse_offload_descriptors() assert list(descriptors) == [2, 4] - assert [d.buffer_id.layer_id for d in descriptors.values()] == [41, 19] + assert [d.buffer_ids[0].layer_id for d in descriptors.values()] == [41, 19] for descriptor in descriptors.values(): - assert descriptor.buffer_id.role == DeepseekV4AttentionType.COMPRESS.role + assert descriptor.buffer_ids[0].role == DeepseekV4AttentionType.COMPRESS.role assert descriptor.group_id == 13 - assert descriptor.page_scale == 5 + assert descriptor.unit_tokens == 4 + assert descriptor.data_buffer_index == 0 + assert descriptor.arena_row_counts == ( + manager.max_batch_size * 2 * (tokens_per_block // 4), + ) + assert descriptor.arena_base_addresses == ( + manager.impl.get_sparse_cache_base_address(*descriptor.buffer_ids[0]), + ) @pytest.mark.cpu_only @pytest.mark.parametrize( "invalid", - ["group", "scale", "pool", "expansion", "ordinary", "not_sparse", "nvfp4", "footer", "api"], + ["expansion", "not_sparse", "unconfigured", "width", "footer"], ) def test_sparse_offload_rejects_incompatible_layout(invalid): manager = _manager() - if invalid == "group": - manager.impl.groups[19] = 100 - elif invalid == "scale": - manager.impl.converters[19].scale = 3 - elif invalid == "pool": - manager.impl.pool_bases[19] += 128 - elif invalid == "expansion": - manager.impl.converters[41].expansion = 2 - elif invalid == "ordinary": - manager.impl.groups[42] = 13 + if invalid == "expansion": + manager.impl.expansions[41] = 1 elif invalid == "not_sparse": manager.impl.sparse.remove(41) - elif invalid == "nvfp4": - manager._use_nvfp4_compress = True elif invalid == "footer": manager.use_fp8_ds_mla = True else: - manager.impl = SimpleNamespace() + config = manager.impl.layers[41] + if invalid == "unconfigured": + manager.impl.layers[41] = AttentionLayerConfig( + layer_id=config.layer_id, buffers=config.buffers + ) + elif invalid == "width": + cache = config.sparse_cache + cache.num_selected_blocks += 1 + config.sparse_cache = cache with pytest.raises((ValueError, NotImplementedError)): manager.get_sparse_offload_descriptors() +@pytest.mark.cpu_only +def test_sparse_offload_accepts_a_dense_buffer_in_a_sparse_layer(): + """An indexer sits beside the data it selects over; only the latter is fetched.""" + manager = _manager() + config = manager.impl.layers[41] + buffers = config.buffers + buffers.append(BufferConfig(role="ordinary", size=16)) + config.buffers = buffers + descriptors = manager.get_sparse_offload_descriptors() + assert descriptors + roles = {b.role for d in descriptors.values() for b in d.buffer_ids} + assert "ordinary" not in roles + + +@pytest.mark.cpu_only +def test_sparse_offload_accepts_a_dense_layer_in_a_sparse_group(): + manager = _manager() + manager.impl.groups[42] = 13 + role = DeepseekV4AttentionType.INDEXER_COMPRESS.role + manager.impl.layers[42] = AttentionLayerConfig( + layer_id=42, buffers=[BufferConfig(role=role, size=1024)] + ) + manager.impl.pool_indices[BufferId(42, role)] = 1 + + descriptors = manager.get_sparse_offload_descriptors() + + assert set(descriptors) == {2, 4} + assert all(buffer.layer_id != 42 for d in descriptors.values() for buffer in d.buffer_ids) + + +@pytest.mark.cpu_only +def test_sparse_offload_descriptors_include_all_roles_in_native_order(): + manager = _manager() + manager._use_nvfp4_compress = True + main_role = DeepseekV4AttentionType.COMPRESS.role + scale_role = deepseek_v4_cache.COMPRESS_BLOCK_SCALE_ROLE + manager.impl.configure_buffers( + 41, [(main_role, 1), (scale_role, 0), ("scale_padding", 0), ("data_padding", 1)] + ) + manager.impl.configure_buffers(19, [(main_role, 1), (scale_role, 0)]) + descriptor = manager.get_sparse_offload_descriptors()[2] + assert [buffer.role for buffer in descriptor.buffer_ids] == [ + scale_role, + "scale_padding", + main_role, + "data_padding", + ] + assert all(buffer.layer_id == 41 for buffer in descriptor.buffer_ids) + assert descriptor.data_buffer_index == 2 + assert descriptor.scale_buffer_index == 0 + + +@pytest.mark.cpu_only +def test_sparse_offload_padding_regions_hold_units_not_pages(): + manager = _manager() + manager._use_nvfp4_compress = True + main = DeepseekV4AttentionType.COMPRESS.role + scale = deepseek_v4_cache.COMPRESS_BLOCK_SCALE_ROLE + manager.impl.configure_buffers( + 41, [(main, 1), (scale, 0), ("scale_padding", 0), ("data_padding", 1)] + ) + manager.impl.configure_buffers(19, [(main, 1), (scale, 0)]) + units_per_page = manager.tokens_per_block // 4 + for layer in manager.impl.layers.values(): + buffers = layer.buffers + for buffer in buffers: + # A sparse buffer's size is a whole page of 288-byte rows, or 36-byte scales. + buffer.size = (36 if "scale" in buffer.role else 288) * units_per_page + layer.buffers = buffers + + descriptor = manager.get_sparse_offload_descriptors()[2] + + # Sixteen cached units round up to one page of 32, and every lease holds one + # region per buffer of the pool: three here, the padding role second. + assert descriptor.padding_regions == tuple( + SparseOffloadPaddingRegion( + base_address=manager.impl.get_sparse_cache_base_address(41, role) + 32 * row_bytes, + bytes_per_lease=32 * row_bytes, + lease_stride=3 * 32 * row_bytes, + ) + for role, row_bytes in (("scale_padding", 36), ("data_padding", 288)) + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +@pytest.mark.parametrize("row_sizes", [(256, 32), (288, 36)]) +@pytest.mark.parametrize("max_units", [16, 65]) +def test_sparse_arena_padding_initialization_covers_all_leases(row_sizes, max_units): + manager = _manager() + manager._use_nvfp4_compress = True + manager._stream = torch.cuda.Stream() + manager._nvfp4_residual_dim = row_sizes[0] * 2 - manager.head_dim + main = DeepseekV4AttentionType.COMPRESS.role + scale = deepseek_v4_cache.COMPRESS_BLOCK_SCALE_ROLE + manager.impl.configure_buffers( + 41, [(main, 1), (scale, 0), ("scale_padding", 0), ("data_padding", 1)] + ) + manager.impl.configure_buffers(19, [(main, 1), (scale, 0)]) + expansion = manager.tokens_per_block // 4 + for layer in manager.impl.layers.values(): + buffers = layer.buffers + for buffer in buffers: + # A sparse buffer's size is a whole page of rows. + buffer.size = row_sizes["scale" in buffer.role] * expansion + layer.buffers = buffers + cache = layer.sparse_cache + cache.max_cached_units = max_units if layer.layer_id == 19 else 16 + layer.sparse_cache = cache + rows = (max_units + expansion - 1) // expansion * expansion + # The Batch owns the arenas and leases one row of them per batch row, so the + # fixture's arena is as wide as its batch limit. Two model buffers and one + # padding buffer share each pool. + leases = manager.max_batch_size + arenas = [ + torch.full((leases, 3, rows, size), 37, dtype=torch.uint8, device="cuda") + for size in row_sizes + ] + manager.impl.get_sparse_cache_base_address = lambda layer, role: arenas[ + "scale" in role + ].data_ptr() + descriptors = manager.get_sparse_offload_descriptors() + manager._initialize_sparse_arena_buffers(descriptors) + data_view, scale_view = manager.get_sparse_compress_buffers(2) + assert manager.get_sparse_compress_buffers(4)[0] is data_view + assert manager.get_sparse_compress_buffers(4)[1] is scale_view + for view, arena, size in zip((data_view, scale_view), arenas, row_sizes): + assert view.shape == (leases * 3 * rows, 1, size) + assert view.is_contiguous() and view.data_ptr() == arena.data_ptr() + torch.testing.assert_close(view.view(torch.uint8).reshape_as(arena), arena) + manager._stream.wait_stream(torch.cuda.current_stream()) + manager._initialize_sparse_arena_padding(descriptors) + for arena in arenas: + expected = torch.full_like(arena, 37) + expected[:, 1] = 0 + torch.testing.assert_close(arena, expected) + # Rebuilding graph metadata describes the same regions without reinitializing + # an arena that may already contain live request data. + for arena in arenas: + arena.fill_(19) + manager.get_sparse_offload_descriptors() + assert all(torch.all(arena == 19) for arena in arenas) + assert torch.all(data_view == 19) + assert torch.all(scale_view.view(torch.uint8) == 19) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("cache_dtype", [DataType.BF16, DataType.FP8, DataType.UINT8]) +@pytest.mark.parametrize("max_units", [16, 65]) +def test_sparse_arena_views_use_pool_capacity_and_share_storage(cache_dtype, max_units): + manager = _manager() + manager.dtype = cache_dtype + manager._use_nvfp4_compress = cache_dtype == DataType.UINT8 + manager._nvfp4_residual_dim = 64 + manager._stream = SimpleNamespace(device="cuda:0") + main = DeepseekV4AttentionType.COMPRESS.role + scale = deepseek_v4_cache.COMPRESS_BLOCK_SCALE_ROLE + roles = [(main, 1), ("data_padding", 1)] + if manager._use_nvfp4_compress: + roles += [(scale, 0), ("scale_padding", 0)] + manager.impl.configure_buffers(41, roles) + manager.impl.configure_buffers( + 19, [(main, 1)] + ([(scale, 0)] if manager._use_nvfp4_compress else []) + ) + config = manager.impl.layers[19] + cache = config.sparse_cache + cache.max_cached_units = max_units + config.sparse_cache = cache + manager.impl.get_sparse_cache_base_address = lambda layer, role: ( + 8192 if "scale" in role else 4096 + ) + descriptors = manager.get_sparse_offload_descriptors() + rows = manager.max_batch_size * 3 * ((max_units + 31) // 32 * 32) + with ( + patch.object(torch.cuda, "device"), + patch.object(torch.cuda, "is_current_stream_capturing", return_value=False), + patch.object( + deepseek_v4_cache, + "convert_to_torch_tensor", + side_effect=lambda wrapper: torch.empty(wrapper.shape, dtype=wrapper.dtype), + ) as convert, + ): + manager._initialize_sparse_arena_buffers(descriptors) + data, scales = manager.get_sparse_compress_buffers(2) + assert data.shape == (rows, 1, 288 if manager._use_nvfp4_compress else 512) + assert data.dtype == binding_to_torch_dtype(cache_dtype) + assert data.is_contiguous() + assert manager.get_sparse_compress_buffers(4)[0] is data + if manager._use_nvfp4_compress: + assert scales.shape == (rows, 1, 36) + assert scales.dtype == torch.float8_e4m3fn + assert manager.get_sparse_compress_buffers(4)[1] is scales + else: + assert scales is None + assert convert.call_count == (2 if manager._use_nvfp4_compress else 1) + assert convert.call_args_list[0].args[0].data_ptr() == 4096 + + with ( + patch.object(torch.cuda, "device"), + patch.object(torch.cuda, "is_current_stream_capturing", return_value=True), + patch.object(deepseek_v4_cache, "convert_to_torch_tensor") as convert, + ): + with pytest.raises(RuntimeError, match="before graph capture"): + manager._initialize_sparse_arena_buffers(descriptors) + convert.assert_not_called() + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("tokens_per_block", [16, 128, 256]) +@pytest.mark.parametrize("sparse", [False, True]) +@pytest.mark.parametrize("residual_dim", [0, 64]) +def test_compress_storage_views_match_allocation_size( + tokens_per_block: int, sparse: bool, residual_dim: int +) -> None: + manager = _manager(tokens_per_block) + manager._use_nvfp4_compress = True + manager._nvfp4_residual_dim = residual_dim + manager.index_head_dim = 128 + manager.compressed_block_sizes = [ + tokens_per_block // ratio for ratio in manager._compress_ratios + ] + if not sparse: + manager.impl.sparse.clear() + main = DeepseekV4AttentionType.COMPRESS.role + scale = deepseek_v4_cache.COMPRESS_BLOCK_SCALE_ROLE + addresses = {main: 4096, scale: 8192} + manager.impl.get_mem_pool_base_address = Mock( + side_effect=lambda layer, role, mode: addresses[role] + ) + # A sparse buffer's page index space is in pages, and a page holds + # compressed_block_size rows. + manager.impl.get_page_index_upper_bound = Mock(return_value=3) + manager.compress_pool_ptrs = {4: addresses[main]} + manager.compress_scale_pool_ptrs = {4: addresses[scale]} + with patch.object( + deepseek_v4_cache, + "convert_to_torch_tensor", + side_effect=lambda wrapper: torch.empty(wrapper.shape, dtype=wrapper.dtype), + ) as convert: + data = manager.get_buffers(2, DeepseekV4AttentionType.COMPRESS) + scales = manager.get_compress_scale_buffers(2) + flat_data, flat_scales = manager.get_compress_pool_buffers(4) + + rows_per_page = tokens_per_block // 4 + storage_dim = manager.head_dim + residual_dim + assert data.shape == (3, rows_per_page, storage_dim // 2) + assert scales.shape == (3, rows_per_page, storage_dim // 16) + assert data.dtype == torch.uint8 + assert scales.dtype == torch.float8_e4m3fn + assert data.numel() == flat_data.numel() == 3 * (tokens_per_block // 4) * storage_dim // 2 + assert scales.numel() == flat_scales.numel() == 3 * (tokens_per_block // 4) * storage_dim // 16 + assert [call.args[0].data_ptr() for call in convert.call_args_list] == [4096, 8192, 4096, 8192] + + +@pytest.mark.parametrize( + "device", + [ + pytest.param("cpu", marks=pytest.mark.cpu_only), + pytest.param( + "cuda", + marks=pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required"), + ), + ], +) +@pytest.mark.parametrize("nvfp4", [False, True]) +@pytest.mark.parametrize("fill_with_zero", [False, True]) +def test_sparse_warmup_cleanup_covers_arenas_and_leaves_host_history( + monkeypatch: pytest.MonkeyPatch, device: str, nvfp4: bool, fill_with_zero: bool +) -> None: + manager = _manager() + manager._use_nvfp4_compress = nvfp4 + manager._nvfp4_residual_dim = 64 + manager.index_head_dim = 128 + manager.compressed_block_sizes = [ + manager.tokens_per_block // ratio for ratio in manager._compress_ratios + ] + manager._layer_attn_to_layer_id = { + (2, DeepseekV4AttentionType.COMPRESS): 41, + (4, DeepseekV4AttentionType.COMPRESS): 19, + } + main = DeepseekV4AttentionType.COMPRESS.role + scale = deepseek_v4_cache.COMPRESS_BLOCK_SCALE_ROLE + manager._guard_page_by_layer = {(41, main): 5} + manager._guard_page_value = 17 + backings = [] + + def allocate(rows: int, width: int, dtype: torch.dtype, *, host: bool) -> torch.Tensor: + storage = torch.full( + (rows + 2, 1, width), + 37, + dtype=torch.uint8 if dtype == torch.float8_e4m3fn else dtype, + device="cpu" if host else device, + pin_memory=host and device == "cuda", + ) + backings.append((storage, storage[[0, -1]].clone())) + return storage[1:-1].view(dtype) + + data_dtype = torch.uint8 if nvfp4 else torch.bfloat16 + data_width = (manager.head_dim + manager.nvfp4_residual_dim) // 2 if nvfp4 else manager.head_dim + host_rows = 3 * (manager.tokens_per_block // 4) + storage = {} + for layer in (41, 19): + storage[layer, main] = allocate(host_rows, data_width, data_dtype, host=True) + if nvfp4: + storage[layer, scale] = allocate(host_rows, 36, torch.float8_e4m3fn, host=True) + arena_data = allocate(11, data_width, data_dtype, host=False) + arena_scales = allocate(11, 36, torch.float8_e4m3fn, host=False) if nvfp4 else None + manager._sparse_arena_buffers = {2: (arena_data, arena_scales), 4: (arena_data, arena_scales)} + if not nvfp4: + arena_data[0, 0, 0] = torch.nan + buffers = list(storage.values()) + [arena_data] + ([arena_scales] if nvfp4 else []) + originals = [buffer.view(torch.uint8).clone() for buffer in buffers] + manager.impl.get_mem_pool_base_address = Mock( + side_effect=lambda layer, role, mode: storage[layer, role].data_ptr() + ) + manager.impl.get_page_index_upper_bound = Mock( + return_value=host_rows // manager.compressed_block_sizes[2] + ) + by_address = {buffer.data_ptr(): buffer for buffer in storage.values()} + convert = deepseek_v4_cache.convert_to_torch_tensor + + def convert_storage(wrapper: deepseek_v4_cache.TensorWrapper) -> torch.Tensor: + buffer = by_address[wrapper.data_ptr()] + # Check the native allocation bound before exposing a raw pointer to CUDA. + assert wrapper.numel() == buffer.numel() + assert wrapper.dtype == buffer.dtype + return convert(wrapper) if device == "cuda" else buffer.reshape(wrapper.shape) + + monkeypatch.setattr(deepseek_v4_cache, "convert_to_torch_tensor", convert_storage) + if device == "cpu": + manager._stream = SimpleNamespace(device="cpu") + cleanup_stream = Mock() + monkeypatch.setattr(torch.cuda, "current_stream", lambda device: cleanup_stream) + monkeypatch.setattr(torch.cuda, "current_device", lambda: "cpu") + monkeypatch.setattr(torch.cuda, "synchronize", Mock()) + else: + manager._stream = torch.cuda.Stream() + manager._stream.wait_stream(torch.cuda.current_stream()) + + with patch.object(torch, "isnan", wraps=torch.isnan) as isnan: + assert manager.check_invalid_values_in_kv_cache(fill_with_zero) == (not nvfp4) + if device == "cpu": + cleanup_stream.wait_stream.assert_called_once_with(manager._stream) + scanned = [call.args[0].data_ptr() for call in isnan.call_args_list] + for arena in (arena_data, arena_scales): + if arena is not None: + assert scanned.count(arena.data_ptr()) == 1 + # Nothing reads host history before a transfer or a write-back fills it, so + # the scan neither reads nor clears it, guard page included. + assert not {buffer.data_ptr() for buffer in storage.values()} & set(scanned) + for buffer, original in zip(buffers, originals): + cleared = fill_with_zero and all(buffer is not host for host in storage.values()) + expected = torch.zeros_like(original) if cleared else original + torch.testing.assert_close(buffer.view(torch.uint8), expected) + for backing, edges in backings: + torch.testing.assert_close(backing[[0, -1]].view(torch.uint8), edges.view(torch.uint8)) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("invalid", ["count", "position", "missing_scales"]) +def test_sparse_offload_descriptors_reject_unpaired_nvfp4_arenas(invalid): + manager = _manager() + manager._use_nvfp4_compress = True + main = DeepseekV4AttentionType.COMPRESS.role + scale = deepseek_v4_cache.COMPRESS_BLOCK_SCALE_ROLE + manager.impl.configure_buffers(41, [(main, 1), (scale, 0)]) + manager.impl.configure_buffers(19, [(main, 1), (scale, 0)]) + if invalid == "count": + # An extra buffer in a later layer changes every request's lease stride. + manager.impl.configure_buffers(19, [(main, 1), (scale, 0), ("scale_padding", 0)]) + elif invalid == "position": + # Equal pool sizes alone do not align the corresponding buffer positions. + manager.impl.configure_buffers( + 41, [("data_padding", 1), (main, 1), (scale, 0), ("scale_padding", 0)] + ) + else: + manager.impl.configure_buffers(41, [(main, 1)]) + with pytest.raises(ValueError, match="block-scale buffer|data and scale arenas"): + manager.get_sparse_offload_descriptors() + + class _HistoryCache: def __init__(self, tokens_per_block: int, cuda_stream: int) -> None: self.id, self.beam_width, self.num_blocks = 20, 1, 3 @@ -241,6 +679,7 @@ def _history_case( context_current_position=tokens_per_block - 1, context_remaining_length=0, is_dummy_request=False, + is_disagg_generation_init_state=False, py_num_accepted_draft_tokens=0, py_rewind_len=0, state=LlmRequestState.GENERATION_IN_PROGRESS, @@ -263,20 +702,23 @@ def test_completed_history_advances_without_reuse_or_capacity_growth(tokens_per_ assert cache.history_length == tokens_per_block - 1 assert cache.calls == [(None, tokens_per_block - 1)] batch.context_requests, batch.generation_requests = [], [request] - for completed in ( - tokens_per_block, - tokens_per_block + 1, - 2 * tokens_per_block, - 2 * tokens_per_block + 1, - ): + for completed in range(tokens_per_block, 2 * tokens_per_block + 2): + # What scheduling records for the step, as try_allocate_generation does. + manager._allocated_draft_lens[request.py_request_id] = 0 + manager._inflight_draft_lens.setdefault(request.py_request_id, []).append(0) + manager.update_forward_resources(batch, wait_for_input_copy=Mock()) + assert cache.history_length == completed + assert cache.calls[-1] == (None, completed) + assert request.py_request_id not in manager._allocated_draft_lens + assert request.py_request_id not in manager._inflight_draft_lens # Sampling has appended one output token; that token has no KV yet. request.max_beam_num_tokens = completed + 1 manager.update_resources(batch) assert cache.history_length == completed - assert cache.calls[-1] == (3 * tokens_per_block, completed) - assert manager._stream.wait_stream.call_count == 5 + assert cache.calls[-1] == (None, completed) + assert manager._stream.wait_stream.call_count == tokens_per_block + 3 manager._stream.wait_stream.assert_called_with(caller_stream) - assert len(cache.calls) == 5 + assert len(cache.calls) == tokens_per_block + 3 @pytest.mark.cpu_only @@ -286,7 +728,8 @@ def test_completed_history_advances_without_reuse_or_capacity_growth(tokens_per_ ) def test_history_update_rejects_failures_and_skips_inactive(context: bool, problem: str) -> None: manager, cache, request, batch = _history_case() - update = manager.update_context_resources if context else manager.update_resources + update = manager.update_context_resources if context else manager.update_forward_resources + update_kwargs = {} if context else {"wait_for_input_copy": Mock()} if not context: batch.context_requests, batch.generation_requests = [], [request] if problem == "resize_failure": @@ -302,109 +745,924 @@ def test_history_update_rejects_failures_and_skips_inactive(context: bool, probl patch("torch.cuda.is_current_stream_capturing", return_value=problem == "capture"), ): if problem in ("suspended", "removed"): - update(batch) + update(batch, **update_kwargs) manager._stream.wait_stream.assert_not_called() else: with pytest.raises((ValueError, RuntimeError)): - update(batch) + update(batch, **update_kwargs) assert cache.history_length == 0 assert len(cache.calls) == (1 if problem == "resize_failure" else 0) @pytest.mark.cpu_only -def test_descriptor_uses_common_pool_origin_and_layer_relative_bound() -> None: +def test_sparse_descriptors_allow_different_groups_and_pool_geometry() -> None: manager = _manager() + manager.impl.groups[19] = 100 + manager.impl.configure_buffers(19, [(DeepseekV4AttentionType.COMPRESS.role, 1)]) descriptors = manager.get_sparse_offload_descriptors() - first, second = descriptors[2], descriptors[4] - assert first.page_index_upper_bound == 303 - assert second.page_index_upper_bound == 297 - role = first.buffer_id.role - assert manager.impl.get_mem_pool_base_address( - 41, role, PageIndexMode.SHARED - ) != manager.impl.get_mem_pool_base_address(19, role, PageIndexMode.SHARED) - manager.impl.get_page_index_upper_bound = Mock(return_value=0) - with pytest.raises(ValueError, match="physical-page bound"): - manager.get_sparse_offload_descriptors() + assert descriptors[2].group_id == 13 + assert descriptors[4].group_id == 100 @pytest.mark.cpu_only -@pytest.mark.parametrize( - ("first", "last"), [(True, False), (False, False), (False, True), (True, True)] -) +@pytest.mark.parametrize("gen_init", [False, True]) @pytest.mark.parametrize("enabled", [False, True]) -def test_chunked_prefill_checked_from_scheduled_request(first, last, enabled): +def test_prefill_rejected_from_scheduled_request(gen_init, enabled): + """Offload serves generation only, so a context request that runs a forward is refused. + + A disaggregated gen-init request is the exception: it is resized and its KV + transferred before it joins the Batch, so it runs no forward this step. + """ manager = _manager() manager._enable_kv_cache_offload = enabled batch = SimpleNamespace( - context_requests=[SimpleNamespace(is_first_context_chunk=first, is_last_context_chunk=last)] + context_requests=[ + SimpleNamespace( + py_request_id=20, + is_first_context_chunk=True, + is_last_context_chunk=True, + is_disagg_generation_init_state=gen_init, + ) + ], + generation_requests=[], ) - with patch.object(KVCacheManagerV2, "prepare_resources") as prepare: - if enabled and not (first and last): - with pytest.raises(NotImplementedError, match="chunked prefill"): + with ( + patch.object(KVCacheManagerV2, "prepare_resources") as prepare, + patch.object(manager, "_prepare_sparse_batch") as publish, + ): + if enabled and not gen_init: + with pytest.raises(NotImplementedError, match="generation only"): manager.prepare_resources(batch) prepare.assert_not_called() + publish.assert_not_called() else: manager.prepare_resources(batch) prepare.assert_called_once_with(batch) + # Nothing runs a forward here, so no batch is published either way. + publish.assert_not_called() + + +class _SchedulingBatch: + """Model native Batch's exclusive page-index ownership and published snapshots.""" + + def __init__(self, manager, cuda_stream, max_batch_size, max_blocks_per_seq, index_tables=()): + self.cuda_stream = cuda_stream + self.max_batch_size = max_batch_size + self.max_blocks_per_seq = max_blocks_per_seq + self.index_tables = list(index_tables) + self.rows = [None] * max_batch_size + self.snapshots = [] + self.upload_pending = False + self.flushes = [] + self.joins = [] + self.leaves = [] + + def row_of(self, cache): + return self.rows.index(cache) if cache in self.rows else -1 + + def flush_sparse_units(self, new_history_lengths): + # Recorded rather than performed: what matters here is that the write-back is asked + # for while the members are still seated, and with each row's own next history. + assert len(new_history_lengths) == self.max_batch_size + for row, cache in enumerate(self.rows): + if cache is not None: + assert new_history_lengths[row] >= cache.history_length + self.flushes.append(list(new_history_lengths)) + + def add_kv_cache(self, cache): + assert not self.upload_pending + assert cache.page_index_owner == "runtime" + row = self.rows.index(None) + self.rows[row] = cache + cache.page_index_owner = "batch" + self.joins.append(cache.id) + return row + + def remove_kv_cache(self, cache): + assert not self.upload_pending + assert cache.page_index_owner == "batch" + self.rows[self.rows.index(cache)] = None + cache.page_index_owner = "cache" + self.leaves.append(cache.id) + + def coalesced_page_table(self, pool_id: int) -> torch.Tensor: + return torch.tensor( + [ + list(cache.pages) + [BAD_PAGE_INDEX] * (self.max_blocks_per_seq - len(cache.pages)) + if cache is not None + else [BAD_PAGE_INDEX] * self.max_blocks_per_seq + for cache in self.rows + ], + dtype=torch.int32, + ) + + def publish(self): + assert not self.upload_pending + self.snapshots.append( + [ + (cache.id, cache.history_length, cache.capacity, tuple(cache.pages)) + if cache is not None + else (None, 0, 0, (BAD_PAGE_INDEX,) * self.max_blocks_per_seq) + for cache in self.rows + ] + ) + self.upload_pending = True + + def wait_published(self): + self.upload_pending = False + + +@pytest.fixture +def sparse_scheduling_case(): + manager, cache, request, _ = _history_case() + manager.max_batch_size = 4 + cache.page_index_owner = "runtime" + second = _HistoryCache(manager.tokens_per_block, manager._stream.cuda_stream) + second.id = 10 + second.history_length = manager.tokens_per_block + second.page_index_owner = "runtime" + manager.kv_cache_map[10] = second + request.is_first_context_chunk = request.is_last_context_chunk = True + other_request = SimpleNamespace(py_request_id=10, is_disagg_generation_init_state=False) + scheduled = SimpleNamespace(context_requests=[], generation_requests=[request, other_request]) + manager.num_pools = 2 + manager.index_mapper = Mock() + slots = {20: 3, 10: 1} + manager.index_mapper.get_index.side_effect = slots.__getitem__ + manager.host_kv_cache_block_offsets = torch.full((2, 4, 1, 6), -99, dtype=torch.int32) + + def restore(request_id, request_cache): + assert request_cache.page_index_owner == "cache" + request_cache.page_index_owner = "runtime" + + with ( + patch.object(sparse_offload, "Batch", side_effect=_SchedulingBatch) as create_batch, + patch.object(KVCacheManagerV2, "prepare_resources"), + patch.object(manager, "_set_page_index_bufs", side_effect=restore) as restore_buffers, + patch("torch.cuda.current_stream", return_value=Mock(cuda_stream=23)) as current_stream, + patch("torch.cuda.is_current_stream_capturing", return_value=False) as capturing, + ): + yield SimpleNamespace( + manager=manager, + cache=cache, + request=request, + scheduled=scheduled, + slots=slots, + create_batch=create_batch, + restore_buffers=restore_buffers, + current_stream=current_stream, + capturing=capturing, + ) @pytest.mark.cpu_only -@pytest.mark.parametrize("enabled", [False, True]) -@pytest.mark.parametrize("tokens_per_block", [128, 256]) -@pytest.mark.parametrize("pp_layers", [[0, 1, 2, 3], [2, 3]]) -@pytest.mark.parametrize(("dtype", "bytes_per_element"), [(DataType.BF16, 2), (DataType.FP8, 1)]) -def test_sparse_offload_cache_roles( - enabled: bool, - tokens_per_block: int, - pp_layers: list[int], - dtype: DataType, - bytes_per_element: int, +def test_sparse_batch_publishes_schedule_and_reuses_allocation(sparse_scheduling_case): + case = sparse_scheduling_case + manager = case.manager + manager.prepare_resources(case.scheduled) + batch = manager.get_sparse_batch([20, 10, CUDA_GRAPH_DUMMY_REQUEST_ID]) + assert [row[0] for row in batch.snapshots[-1]] == [20, 10, None, None] + assert batch.snapshots[-1][0][1:] == (0, 384, (0, 7, 8)) + assert batch.snapshots[-1][1][1:] == (128, 384, (0, 7, 8)) + assert batch.snapshots[-1][2] == (None, 0, 0, (-1,) * manager.max_blocks_per_seq) + case.current_stream.return_value.wait_stream.assert_called_with(manager._stream) + + # A new schedule order moves nothing: each request keeps its row and lease. + case.scheduled.generation_requests.reverse() + manager.prepare_resources(case.scheduled) + assert manager.get_sparse_batch([10, 20]) is batch + assert [row[0] for row in batch.snapshots[-1]] == [20, 10, None, None] + call = case.create_batch.call_args + assert call.args == (manager.impl, manager._stream.cuda_stream) + assert call.kwargs["max_batch_size"] == 4 + assert call.kwargs["max_blocks_per_seq"] == 6 + # By content: IndexTableSpec has no value equality, so comparing against a freshly + # derived list would only ever compare object identity and never hold. + assert [ + [(buffer.layer_id, buffer.role) for buffer in request.buffers] + for request in call.kwargs["index_tables"] + ] == [ + [(layer_id, DeepseekV4AttentionType.COMPRESS.role)] + for layer_id in sorted(manager.impl.sparse) + ] + case.restore_buffers.assert_not_called() + assert batch.joins == [20, 10] + assert not batch.leaves + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("max_beam_width", [1, 2]) +def test_sparse_batch_retains_continuing_members_and_refreshes_legacy_tables( + sparse_scheduling_case: SimpleNamespace, max_beam_width: int ) -> None: + case = sparse_scheduling_case + manager = case.manager + manager.max_beam_width = max_beam_width + manager.host_kv_cache_block_offsets = torch.full( + (2, 4 * max_beam_width, 1, 6), -99, dtype=torch.int32 + ) + manager.prepare_resources(case.scheduled) + batch = manager.get_sparse_batch([20, 10]) + case.cache.pages = [4, 8, 9] + case.scheduled.generation_requests.reverse() + manager.prepare_resources(case.scheduled) + + assert batch.joins == [20, 10] + assert not batch.leaves + assert manager.get_sparse_batch([10, 20]) is batch + # Only the columns below the members' block counts are copied; past them the + # runtime's binding keeps its own BAD_PAGE_INDEX, here the -99 the test filled. + for pool_id in range(manager.num_pools): + assert manager.host_kv_cache_block_offsets[pool_id, 3 * max_beam_width, 0].tolist() == [ + 4, + 8, + 9, + -99, + -99, + -99, + ] + assert manager.host_kv_cache_block_offsets[pool_id, max_beam_width, 0].tolist() == [ + 0, + 7, + 8, + -99, + -99, + -99, + ] + + case.scheduled.generation_requests = [case.request] + manager.prepare_resources(case.scheduled) + assert manager.get_sparse_batch([20]) is batch + assert batch.joins == [20, 10] + assert batch.leaves == [10] + assert batch.rows == [case.cache, None, None, None] + assert case.cache.page_index_owner == "batch" + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("warmup", [False, True]) +def test_sparse_batch_metadata_joins_scheduler_or_prepares_warmup(sparse_scheduling_case, warmup): + case = sparse_scheduling_case + if not warmup: + case.manager.prepare_resources(case.scheduled) + case.manager.prepare_sparse_batch([20, 10, CUDA_GRAPH_DUMMY_REQUEST_ID]) + batch = case.manager.get_sparse_batch([20, 10]) + case.manager.prepare_sparse_batch([20, 10]) + # Metadata may order the published requests differently; nothing is republished. + case.manager.prepare_sparse_batch([10, 20]) + assert len(batch.snapshots) == 1 + assert [row[0] for row in batch.snapshots[-1]] == [20, 10, None, None] + case.current_stream.return_value.wait_stream.assert_called_with(case.manager._stream) + with pytest.raises(ValueError, match="do not match"): + case.manager.prepare_sparse_batch([20]) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("dummy_count", [0, 1, 4]) +def test_sparse_batch_clears_rows_for_empty_or_padded_schedule(sparse_scheduling_case, dummy_count): + case = sparse_scheduling_case + case.manager.prepare_resources(case.scheduled) + case.scheduled.generation_requests = [ + SimpleNamespace( + py_request_id=CUDA_GRAPH_DUMMY_REQUEST_ID, is_disagg_generation_init_state=False + ) + ] * dummy_count + case.manager.prepare_resources(case.scheduled) + batch = case.manager.get_sparse_batch([CUDA_GRAPH_DUMMY_REQUEST_ID] * dummy_count) + assert all(row == (None, 0, 0, (-1,) * 6) for row in batch.snapshots[-1]) + assert all(cache.page_index_owner == "runtime" for cache in case.manager.kv_cache_map.values()) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + "problem", + [ + "missing", + "inactive", + "beam", + "stream", + "history", + "blocks", + "duplicate", + "padding", + "size", + "capture", + ], +) +def test_sparse_batch_rejects_invalid_schedule_before_binding(sparse_scheduling_case, problem): + case = sparse_scheduling_case + if problem == "missing": + del case.manager.kv_cache_map[20] + elif problem == "inactive": + case.cache.is_active = False + elif problem == "beam": + case.cache.beam_width = 2 + elif problem == "stream": + case.cache.cuda_stream += 1 + elif problem == "history": + case.cache.history_length = case.cache.capacity + 1 + elif problem == "blocks": + case.cache.num_blocks = case.manager.max_blocks_per_seq + 1 + elif problem == "duplicate": + case.scheduled.generation_requests = [case.request, case.request] + elif problem == "padding": + case.scheduled.generation_requests.insert( + 0, + SimpleNamespace( + py_request_id=CUDA_GRAPH_DUMMY_REQUEST_ID, is_disagg_generation_init_state=False + ), + ) + elif problem == "size": + case.manager.max_batch_size = 1 + else: + case.capturing.return_value = True + with pytest.raises((ValueError, RuntimeError)): + case.manager.prepare_resources(case.scheduled) + case.create_batch.assert_not_called() + case.restore_buffers.assert_not_called() + + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + "request_ids", + [ + [20], + [20, 20], + [20, 10, 30], + [20, -1, 10], + [20, 10] + [-1] * 3, + [CUDA_GRAPH_DUMMY_REQUEST_ID, 20, 10], + [20, 10] + [CUDA_GRAPH_DUMMY_REQUEST_ID] * 3, + ], +) +def test_sparse_batch_checks_consumer_requests(sparse_scheduling_case, request_ids): + case = sparse_scheduling_case + case.manager.prepare_resources(case.scheduled) + with pytest.raises(ValueError, match="do not match"): + case.manager.get_sparse_batch(request_ids) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("request_ids", [[20, 10], [10, 20], [10, 20, CUDA_GRAPH_DUMMY_REQUEST_ID]]) +def test_sparse_batch_accepts_members_in_any_order(sparse_scheduling_case, request_ids): + case = sparse_scheduling_case + case.manager.prepare_resources(case.scheduled) + assert case.manager.get_sparse_batch(request_ids) is case.manager._sparse_batch + + +def _add_scheduled_request(case: SimpleNamespace, request_id: int, slot: int) -> _HistoryCache: + cache = _HistoryCache(case.manager.tokens_per_block, case.manager._stream.cuda_stream) + cache.id = request_id + cache.page_index_owner = "runtime" + case.manager.kv_cache_map[request_id] = cache + case.slots[request_id] = slot + return cache + + +@pytest.mark.cpu_only +def test_sparse_batch_rows_are_slots_that_arrivals_refill(sparse_scheduling_case): + """A departure leaves a free row, which the next arrival takes; nobody else moves.""" + case = sparse_scheduling_case + manager = case.manager + _add_scheduled_request(case, 30, slot=0) + third = SimpleNamespace(py_request_id=30, is_disagg_generation_init_state=False) + other = case.scheduled.generation_requests[1] + + def schedule(*requests): + case.scheduled.generation_requests = list(requests) + manager.prepare_resources(case.scheduled) + return {rid: member.row for rid, member in manager._sparse_batch_members.items()} + + assert schedule(case.request, other) == {20: 0, 10: 1} + batch = manager._sparse_batch + assert schedule(other) == {10: 1} + assert batch.rows == [None, manager.kv_cache_map[10], None, None] + # The arrival runs first in the input but takes the lowest free row. + assert schedule(third, other) == {10: 1, 30: 0} + assert manager.get_sparse_batch([30, 10]) is batch + assert schedule(other, third) == {10: 1, 30: 0} + assert batch.joins == [20, 10, 30] + assert batch.leaves == [20] + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("hole", [False, True]) +def test_sparse_batch_snapshot_follows_rows_with_or_without_holes(sparse_scheduling_case, hole): + case = sparse_scheduling_case + manager = case.manager + third = _add_scheduled_request(case, 30, slot=0) + third.pages = [5, 6, 7] + scheduled = case.scheduled.generation_requests + [ + SimpleNamespace(py_request_id=30, is_disagg_generation_init_state=False) + ] + case.scheduled.generation_requests = scheduled + manager.prepare_resources(case.scheduled) + assert {rid: member.row for rid, member in manager._sparse_batch_members.items()} == { + 20: 0, + 10: 1, + 30: 2, + } + if hole: + # Request 10 leaves row 1 free, so rows 0 and 2 are gathered. + case.scheduled.generation_requests = [scheduled[2], scheduled[0]] + # A member that grew writes its new pages through its Batch row. + case.cache.pages = [4, 8, 9, 11] + case.cache.num_blocks = 4 + third.pages = [5, 6, 12] + manager.prepare_resources(case.scheduled) + host = manager.host_kv_cache_block_offsets + for pool_id in range(manager.num_pools): + assert host[pool_id, 3, 0].tolist() == [4, 8, 9, 11, -99, -99] + assert host[pool_id, 0, 0].tolist() == [5, 6, 12, -1, -99, -99] + if not hole: + assert host[pool_id, 1, 0].tolist() == [0, 7, 8, -1, -99, -99] + + +def _gen_init_request(request_id: int) -> SimpleNamespace: + return SimpleNamespace( + py_request_id=request_id, + is_first_context_chunk=True, + is_last_context_chunk=True, + is_disagg_generation_init_state=True, + ) + + +@pytest.mark.cpu_only +def test_sparse_batch_ignores_disagg_gen_init_preparation(sparse_scheduling_case): + """Gen-init requests are prepared before their KV transfer and run no forward. + + Their admission bypasses the batch budget, so a burst can exceed the Batch rows. + """ + case = sparse_scheduling_case + burst = case.manager.max_batch_size + 2 + case.scheduled.context_requests = [_gen_init_request(100 + i) for i in range(burst)] + case.scheduled.generation_requests = [] + + case.manager.prepare_resources(case.scheduled) + + case.create_batch.assert_not_called() + with pytest.raises(RuntimeError, match="must be prepared"): + case.manager.get_sparse_batch([]) + + +@pytest.mark.cpu_only +def test_disagg_gen_init_retains_continuing_sparse_cache_leases( + sparse_scheduling_case: SimpleNamespace, +) -> None: + case = sparse_scheduling_case + case.manager.prepare_resources(case.scheduled) + batch = case.manager.get_sparse_batch([20, 10]) + case.manager.prepare_resources( + SimpleNamespace(context_requests=[_gen_init_request(100)], generation_requests=[]) + ) + assert [cache.id for cache in batch.rows if cache is not None] == [20, 10] + assert not batch.leaves + case.manager.prepare_resources(case.scheduled) + assert case.manager.get_sparse_batch([20, 10]) is batch + assert batch.joins == [20, 10] + + +@pytest.mark.cpu_only +def test_releasing_unseated_request_preserves_the_published_batch( + sparse_scheduling_case: SimpleNamespace, +) -> None: + case = sparse_scheduling_case + case.manager.prepare_resources(case.scheduled) + batch = case.manager.get_sparse_batch([20, 10]) + with patch.object(KVCacheManagerV2, "free_resources"): + case.manager.free_resources(_gen_init_request(100)) + assert case.manager.get_sparse_batch([20, 10]) is batch + assert not batch.leaves + + +@pytest.mark.cpu_only +def test_sparse_batch_publishes_only_requests_that_run_a_forward(sparse_scheduling_case): + case = sparse_scheduling_case + case.scheduled.context_requests = [_gen_init_request(100)] + + case.manager.prepare_resources(case.scheduled) + + batch = case.manager.get_sparse_batch([20, 10]) + assert [row[0] for row in batch.snapshots[-1]] == [20, 10, None, None] + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("failure", ["add", "publish"]) +def test_sparse_batch_restores_ownership_after_publication_failure(sparse_scheduling_case, failure): + case = sparse_scheduling_case + case.manager.prepare_resources(case.scheduled) + batch = case.manager.get_sparse_batch([20, 10]) + case.manager._release_sparse_batch() + if failure == "add": + add = batch.add_kv_cache + + def fail_second(cache): + if cache.id == 10: + raise RuntimeError("injected add failure") + return add(cache) + + batch.add_kv_cache = fail_second + else: + batch.publish = Mock(side_effect=RuntimeError("injected publish failure")) + with pytest.raises(RuntimeError, match="injected"): + case.manager.prepare_resources(case.scheduled) + assert batch.rows == [None] * 4 + assert all(cache.page_index_owner == "runtime" for cache in case.manager.kv_cache_map.values()) + with pytest.raises(RuntimeError, match="must be prepared"): + case.manager.get_sparse_batch([20, 10]) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + "operation", + [ + "update_context_resources", + "suspend_request", + "free_resources", + "release_index_slot", + "shutdown", + ], +) +def test_sparse_batch_detaches_before_cache_lifecycle_mutations(sparse_scheduling_case, operation): + case = sparse_scheduling_case + manager = case.manager + manager.prepare_resources(case.scheduled) + batch = manager.get_sparse_batch([20, 10]) + + def mutate(*args): + if operation == "shutdown": + assert batch.rows == [None] * 4 + assert all( + cache.page_index_owner == "runtime" for cache in manager.kv_cache_map.values() + ) + else: + assert batch.rows == [None, manager.kv_cache_map[10], None, None] + assert case.cache.page_index_owner == "runtime" + assert manager.kv_cache_map[10].page_index_owner == "batch" + + with patch.object(KVCacheManagerV2, operation, side_effect=mutate) as parent: + if operation.startswith("update_"): + case.scheduled.context_requests = [case.request] + getattr(manager, operation)(case.scheduled) + elif operation == "release_index_slot": + manager.release_index_slot(20) + elif operation == "shutdown": + manager.shutdown() + else: + getattr(manager, operation)(case.request) + parent.assert_called_once() + with pytest.raises(RuntimeError, match="must be prepared"): + manager.get_sparse_batch([20, 10]) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("context", [False, True]) +def test_sparse_batch_publishes_updated_history_on_next_step(sparse_scheduling_case, context): + case = sparse_scheduling_case + manager, cache, request = case.manager, case.cache, case.request + cache.history_length, cache.capacity = 7, 8 + request.context_current_position = 8 + request.max_beam_num_tokens = 9 # The sampled output token does not yet have KV. + # Offload serves generation only, so the context side of this is reached solely by a + # disaggregated gen-init request: resized and transferred before it joins the Batch, it + # runs no forward and nothing is published for it until it becomes a generation request. + request.is_disagg_generation_init_state = context + case.scheduled.context_requests = [request] if context else [] + case.scheduled.generation_requests = [] if context else [request] + manager.prepare_resources(case.scheduled) + batch = None + if not context: + batch = manager.get_sparse_batch([20]) + assert batch.snapshots[-1][0][1:3] == (7, 8) + + if context: + manager.update_context_resources(case.scheduled) + else: + manager.update_forward_resources(case.scheduled, wait_for_input_copy=Mock()) + manager.update_resources(case.scheduled) + assert cache.history_length == cache.capacity == 8 + assert cache.page_index_owner == ("runtime" if context else "batch") + if batch is not None: + assert not batch.upload_pending + + # Scheduling reserves the next input token after the completed history update. + assert cache.resize(9, None) + request.is_disagg_generation_init_state = False + case.scheduled.context_requests = [] + case.scheduled.generation_requests = [request] + manager.prepare_resources(case.scheduled) + stepped = manager.get_sparse_batch([20]) + if batch is not None: + assert stepped is batch + assert stepped.snapshots[-1][0][1:3] == (8, 9) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("history_length", [7, 127, 255]) +def test_sparse_overlap_history_advances_before_delayed_samples( + sparse_scheduling_case: SimpleNamespace, history_length: int +) -> None: + case = sparse_scheduling_case + manager, cache, request = case.manager, case.cache, case.request + case.scheduled.generation_requests = [request] + cache.history_length, cache.capacity = history_length, history_length + 1 + request.max_beam_num_tokens = history_length + 1 + manager.prepare_resources(case.scheduled) + batch = manager.get_sparse_batch([20]) + + # Forward N produces KV before the executor consumes sampling result N. + manager.update_forward_resources(case.scheduled, wait_for_input_copy=Mock()) + assert cache.history_length == history_length + 1 + assert cache.page_index_owner == "batch" + + assert cache.resize(history_length + 2, None) + manager.prepare_resources(case.scheduled) + assert batch.snapshots[-1][0][1:3] == (history_length + 1, history_length + 2) + calls = list(cache.calls) + # A delayed update must not detach the newly published Batch. + manager.update_resources(case.scheduled) + assert manager.get_sparse_batch([20]) is batch + assert cache.calls == calls + + manager.update_forward_resources(case.scheduled, wait_for_input_copy=Mock()) + assert cache.history_length == history_length + 2 + calls = list(cache.calls) + # The overlap executor now consumes N's sample after forward N+1. + # Its CPU token count describes an older KV frontier. + request.max_beam_num_tokens = history_length + 2 + manager.update_resources(case.scheduled) + assert cache.history_length == history_length + 2 + assert cache.calls == calls + + assert cache.resize(history_length + 3, None) + manager.prepare_resources(case.scheduled) + assert batch.snapshots[-1][0][1:3] == (history_length + 2, history_length + 3) + + +@pytest.mark.cpu_only +def test_sparse_history_does_not_advance_without_forward( + sparse_scheduling_case: SimpleNamespace, +) -> None: + case = sparse_scheduling_case + manager, cache, request = case.manager, case.cache, case.request + case.scheduled.generation_requests = [request] + cache.history_length, cache.capacity = 7, 8 + manager.prepare_resources(case.scheduled) + batch = manager.get_sparse_batch([20]) + # Resource preparation alone can be followed by can_queue=False. + manager.update_resources(case.scheduled) + manager.prepare_resources(case.scheduled) + assert cache.history_length == 7 + assert not cache.calls + assert batch.snapshots[-1][0][1:3] == (7, 8) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("copy_fails", [False, True]) +def test_sparse_forward_waits_for_input_copy_before_history_update( + sparse_scheduling_case: SimpleNamespace, copy_fails: bool +) -> None: + case = sparse_scheduling_case + manager, cache, request = case.manager, case.cache, case.request + case.scheduled.generation_requests = [request] + cache.history_length, cache.capacity = 7, 8 + manager.prepare_resources(case.scheduled) + batch = manager.get_sparse_batch([20]) + + def wait_for_input_copy() -> None: + assert manager.get_sparse_batch([20]) is batch + assert cache.page_index_owner == "batch" + assert not cache.calls + if copy_fails: + raise RuntimeError("input copy failed") + + wait = Mock(side_effect=wait_for_input_copy) + if copy_fails: + with pytest.raises(RuntimeError, match="input copy failed"): + manager.update_forward_resources(case.scheduled, wait_for_input_copy=wait) + assert manager.get_sparse_batch([20]) is batch + assert cache.history_length == 7 + assert not cache.calls + else: + manager.update_forward_resources(case.scheduled, wait_for_input_copy=wait) + assert cache.history_length == 8 + assert cache.page_index_owner == "batch" + wait.assert_called_once_with() + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("history_length", [3, 127]) +def test_sparse_forward_keeps_flush_proof_attached_until_history_advances( + sparse_scheduling_case: SimpleNamespace, history_length: int +) -> None: + case = sparse_scheduling_case + manager, cache, request = case.manager, case.cache, case.request + case.scheduled.generation_requests = [request] + cache.history_length, cache.capacity = history_length, history_length + 1 + manager.prepare_resources(case.scheduled) + batch = manager.get_sparse_batch([20]) + manager._stream.wait_stream.reset_mock() + flush = batch.flush_sparse_units + resize = cache.resize + + def flush_after_forward(histories): + manager._stream.wait_stream.assert_called_once_with(case.current_stream.return_value) + flush(histories) + + def advance_attached(capacity, history): + row = batch.row_of(cache) + assert row >= 0, "The native flush watermark belongs to the attached Batch row" + assert batch.flushes[-1][row] >= history + assert not batch.upload_pending, "History mutation must not race the staged table upload" + return resize(capacity, history) + + with ( + patch.object(batch, "flush_sparse_units", side_effect=flush_after_forward), + patch.object(cache, "resize", side_effect=advance_attached), + ): + manager.update_forward_resources(case.scheduled, wait_for_input_copy=Mock()) + + assert cache.history_length == history_length + 1 + assert cache.page_index_owner == "batch" + assert batch.row_of(cache) == 0 + assert batch.joins == [20] + assert not batch.leaves + + +@pytest.mark.cpu_only +def test_sparse_forward_requires_reserved_input_capacity() -> None: + manager, cache, request, batch = _history_case() + batch.context_requests, batch.generation_requests = [], [request] + cache.history_length = cache.capacity + with ( + patch("torch.cuda.current_stream"), + patch("torch.cuda.is_current_stream_capturing", return_value=False), + pytest.raises(ValueError, match="no reserved input token"), + ): + manager.update_forward_resources(batch, wait_for_input_copy=Mock()) + assert not cache.calls + + +@pytest.mark.cpu_only +def test_sparse_scheduling_keeps_a_reserved_input_token() -> None: + """An attention-DP dummy arrives holding its input token; the scheduler must not grow it.""" + manager, cache, request, batch = _history_case() + batch.context_requests, batch.generation_requests = [], [request] + # add_dummy_requests(is_gen=True) reserves the step's input token up front. + cache.history_length, cache.capacity = 0, 1 + + with patch.object(KVCacheManagerV2, "try_allocate_generation") as grow: + assert manager.try_allocate_generation(request) + grow.assert_not_called() + assert not cache.calls + assert request.py_request_id not in manager._allocated_draft_lens + + # A fleet-wide skipped step (can_queue=False) has no growth to give back. + manager.revert_allocate_generation(request) + assert not cache.calls + + # The forward consumes exactly the reserved token. + with ( + patch("torch.cuda.current_stream"), + patch("torch.cuda.is_current_stream_capturing", return_value=False), + ): + manager.update_forward_resources(batch, wait_for_input_copy=Mock()) + assert cache.history_length == cache.capacity == 1 + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("offload", [False, True]) +@pytest.mark.parametrize("reserved", [False, True]) +def test_generation_growth_defers_to_the_base_policy(offload: bool, reserved: bool) -> None: + """Only an offload cache that already holds its input token skips growth. + + Without offload, history follows sampling and lags one step under the overlap + scheduler, so capacity beyond history is normal there and must still grow. + """ + manager, cache, request, _ = _history_case() + manager._enable_kv_cache_offload = offload + cache.history_length = 7 + cache.capacity = 8 if reserved else 7 + + with patch.object(KVCacheManagerV2, "try_allocate_generation", return_value=True) as grow: + assert manager.try_allocate_generation(request) + assert grow.call_count == (0 if offload and reserved else 1) + + +@pytest.mark.cpu_only +def test_non_offload_history_still_follows_sampling() -> None: + manager, cache, request, batch = _history_case() + manager._enable_kv_cache_offload = False + batch.context_requests, batch.generation_requests = [], [request] + wait = Mock() + manager.update_forward_resources(batch, wait_for_input_copy=wait) + wait.assert_not_called() + assert not cache.calls + manager.update_resources(batch) + assert cache.history_length == request.max_beam_num_tokens - 1 + manager._stream.wait_stream.assert_not_called() + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("offload", [False, True]) +def test_sparse_batch_not_allocated_without_local_sparse_layers(sparse_scheduling_case, offload): + case = sparse_scheduling_case + case.manager._enable_kv_cache_offload = offload + if offload: + case.manager.pp_layers = [3] + case.manager.prepare_resources(case.scheduled) + case.create_batch.assert_not_called() + case.current_stream.assert_not_called() + + +def _cache_config_case( + *, + enabled: bool, + tokens_per_block: int = 128, + pp_layers: list[int] | None = None, + dtype: DataType = DataType.BF16, + nvfp4_residual_dim: int = 0, + index_topk: int | None = 512, + indexer_k_dtype: str = "fp8", + cache_factor: float = 2.0, +) -> tuple[DeepseekV4CacheManager, KVCacheManagerConfig]: manager = object.__new__(DeepseekV4CacheManager) manager._enable_kv_cache_offload = enabled - manager.pp_layers = pp_layers + manager._sparse_offload_topk = index_topk + manager._sparse_kv_offload = ( + SparseKvCacheOffloadConfig(gpu_cache_factor=cache_factor) if enabled else None + ) + manager.pp_layers = [0, 1, 2, 3] if pp_layers is None else pp_layers manager._compress_ratios = [1, 4, 128, 4] manager.compressed_block_sizes = [tokens_per_block // r for r in manager._compress_ratios] manager.tokens_per_block = tokens_per_block manager.head_dim = 512 manager.index_head_dim = 128 - manager.dtype = dtype - manager._indexer_k_dtype = "fp8" - manager._use_nvfp4_compress = False + manager.dtype = DataType.FP8 if dtype == DataType.NVFP4 else dtype + manager._indexer_k_dtype = indexer_k_dtype + manager._use_nvfp4_compress = dtype == DataType.NVFP4 + manager._nvfp4_residual_dim = nvfp4_residual_dim manager.use_fp8_ds_mla = False manager._swa_window_size = 128 manager.max_seq_len = 1024 manager.max_batch_size = 2 manager.is_estimating_kv_cache = False manager.max_cuda_graph_batch_size = None - manager.max_attention_window_vec = [None] * len(pp_layers) + manager.max_attention_window_vec = [None] * len(manager.pp_layers) manager.max_draft_len = manager._max_draft_len = 0 manager.num_extra_kv_tokens = 0 config = KVCacheManagerConfig( tokens_per_block=tokens_per_block, - cache_tiers=[ - GpuCacheTierConfig(quota=1 << 20), - HostCacheTierConfig(quota=1 << 20), - ], + cache_tiers=[GpuCacheTierConfig(quota=1 << 20)], layers=[], ) + return manager, config - def make_buffer_config(*, role: str, size: int, is_sparse: bool = False) -> BufferConfig: - """Keep native buffers while the mock records the unsupported sparse flag.""" - return BufferConfig(role=role, size=size) - with patch.object( - deepseek_v4_cache, "BufferConfig", side_effect=make_buffer_config - ) as create_buffer: - result = manager._build_cache_config(config) +@pytest.mark.cpu_only +@pytest.mark.parametrize("enabled", [False, True]) +@pytest.mark.parametrize("tokens_per_block", [128, 256]) +@pytest.mark.parametrize("pp_layers", [[0, 1, 2, 3], [2, 3], [0, 2]]) +@pytest.mark.parametrize("index_topk", [1, 512]) +@pytest.mark.parametrize("indexer_k_dtype", ["fp8", "fp4"]) +@pytest.mark.parametrize( + ("dtype", "nvfp4_residual_dim", "data_row_bytes", "scale_row_bytes"), + [ + (DataType.BF16, 0, 1024, None), + (DataType.FP8, 0, 512, None), + (DataType.NVFP4, 0, 256, 32), + (DataType.NVFP4, 64, 288, 36), + ], +) +def test_sparse_offload_cache_roles( + enabled: bool, + tokens_per_block: int, + pp_layers: list[int], + index_topk: int, + indexer_k_dtype: str, + dtype: DataType, + nvfp4_residual_dim: int, + data_row_bytes: int, + scale_row_bytes: int | None, +) -> None: + manager, config = _cache_config_case( + enabled=enabled, + tokens_per_block=tokens_per_block, + pp_layers=pp_layers, + dtype=dtype, + nvfp4_residual_dim=nvfp4_residual_dim, + index_topk=index_topk, + indexer_k_dtype=indexer_k_dtype, + ) + result = manager._build_cache_config(config) assert result.tokens_per_block == tokens_per_block assert len(result.layers) == manager._num_manager_layers - sparse_flags = {} - buffers = [(layer.layer_id, buffer) for layer in result.layers for buffer in layer.buffers] - for (layer_id, buffer), call in zip(buffers, create_buffer.call_args_list, strict=True): - assert (buffer.role, buffer.size) == (call.kwargs["role"], call.kwargs["size"]) - sparse_flags[layer_id, buffer.role] = call.kwargs.get("is_sparse", False) + for layer in result.layers: + for buffer in layer.buffers: + if not buffer.is_sparse: + assert buffer.tokens_per_block_override is None + assert buffer.sparse_unit_chunks == 1 layers = {layer.layer_id: layer for layer in result.layers} for (model_layer, attn_type), layer_id in manager._layer_attn_to_layer_id.items(): assert model_layer in pp_layers @@ -413,42 +1671,123 @@ def make_buffer_config(*, role: str, size: int, is_sparse: bool = False) -> Buff attn_type, ) buffer = next(b for b in layers[layer_id].buffers if b.role == attn_type.role) - assert buffer.tokens_per_block_override is None ratio = manager._compress_ratios[model_layer] - assert sparse_flags[layer_id, buffer.role] == ( + assert buffer.is_sparse == ( enabled and ratio == 4 and attn_type == DeepseekV4AttentionType.COMPRESS ) if attn_type == DeepseekV4AttentionType.COMPRESS: - assert buffer.size == tokens_per_block // ratio * manager.head_dim * bytes_per_element + row_sizes = {attn_type.role: data_row_bytes} + if scale_row_bytes is not None: + row_sizes[deepseek_v4_cache.COMPRESS_BLOCK_SCALE_ROLE] = scale_row_bytes + compressed_rows = tokens_per_block // ratio + for role, row_bytes in row_sizes.items(): + role_buffer = next(b for b in layers[layer_id].buffers if b.role == role) + if enabled and ratio == 4: + assert role_buffer.is_sparse + # A sparse buffer's size is its whole page, as for any other buffer; + # the layer's sparse_cache says how many units it divides into. + assert role_buffer.tokens_per_block_override is None + assert role_buffer.sparse_unit_chunks == 1 + assert role_buffer.size == compressed_rows * row_bytes + cache = layers[layer_id].sparse_cache + assert cache is not None + assert cache.num_selected_blocks == index_topk + # Two selections plus the row being written, in whole pages of + # compressed rows. + cells = -(-(2 * index_topk + 1) // compressed_rows) + assert cache.max_cached_units == cells * compressed_rows + assert cache.unit_tokens == 4 + else: + assert role_buffer.size == compressed_rows * row_bytes if ratio == 4: indexer_layer_id = manager._layer_attn_to_layer_id[ model_layer, DeepseekV4AttentionType.INDEXER_COMPRESS ] - assert (layer_id != indexer_layer_id) is enabled + # COMPRESS and INDEXER_COMPRESS are one layer either way: they are one + # attention computation, and only the first is fetched through the cache. + assert layer_id == indexer_layer_id @pytest.mark.cpu_only -def test_sparse_offload_rejects_inference_before_cache_allocation() -> None: - sparse_config = DeepSeekV4SparseAttentionConfig(enable_kv_cache_offload=True) - with ( - patch.object(deepseek_v4_cache, "get_sm_version") as get_sm_version, - patch.object(deepseek_v4_cache.KVCacheManagerV2, "__init__") as initialize_cache, - pytest.raises(NotImplementedError, match="KVCM v2 sparse runtime"), - ): - DeepseekV4CacheManager( - kv_cache_config=KvCacheConfig(enable_block_reuse=False, host_cache_size=1 << 20), - kv_cache_type=CacheTypeCpp.SELFKONLY, - num_layers=len(sparse_config.compress_ratios), - max_batch_size=1, - tokens_per_block=128, - max_seq_len=256, - vocab_size=128, - mapping=Mapping(), - sparse_attn_config=sparse_config, - ) +@pytest.mark.parametrize("enabled", [False, True]) +@pytest.mark.parametrize("pp_layers", [[0, 1], [0, 2]]) +def test_sparse_offload_cache_config_footer_layout(enabled: bool, pp_layers: list[int]) -> None: + manager, config = _cache_config_case(enabled=enabled, pp_layers=pp_layers, dtype=DataType.FP8) + manager.use_fp8_ds_mla = True + if enabled and 1 in pp_layers: + with pytest.raises(NotImplementedError, match="separate data and footer-scale regions"): + manager._build_cache_config(config) + else: + result = manager._build_cache_config(config) + assert all(not buffer.is_sparse for layer in result.layers for buffer in layer.buffers) - get_sm_version.assert_not_called() - initialize_cache.assert_not_called() + +@pytest.mark.cpu_only +@pytest.mark.parametrize(("cache_factor", "rows_per_layer"), [(1.0, 544), (1.5, 800), (3.0, 1568)]) +def test_sparse_offload_cache_factor_scales_whole_cells( + cache_factor: float, rows_per_layer: int +) -> None: + manager, config = _cache_config_case(enabled=True, cache_factor=cache_factor) + result = manager._build_cache_config(config) + + caches = [layer.sparse_cache for layer in result.layers if layer.sparse_cache is not None] + assert caches + # The factor scales the 512 selected rows, the row being written adds one, and + # the total rounds up to cells of 32 rows. + assert {cache.max_cached_units for cache in caches} == {rows_per_layer} + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("index_topk", [None, 0, -1]) +def test_sparse_offload_cache_config_requires_positive_topk(index_topk: int | None) -> None: + manager, config = _cache_config_case(enabled=True, index_topk=index_topk) + with pytest.raises(ValueError, match="positive index_topk"): + manager._build_cache_config(config) + + +@skip_pre_blackwell +@pytest.mark.parametrize("tokens_per_block", [128, 256]) +@pytest.mark.parametrize("dtype", [DataType.BF16, DataType.FP8, DataType.NVFP4]) +def test_sparse_offload_initializes_host_history_and_gpu_arena( + tokens_per_block: int, dtype: DataType +) -> None: + sparse_config = DeepSeekV4SparseAttentionConfig( + enable_kv_cache_offload=True, + compress_ratios=[4, 128], + index_topk=8, + kv_cache_offload_config={"host_history_size": 64 << 20}, + ) + manager = DeepseekV4CacheManager( + kv_cache_config=KvCacheConfig(enable_block_reuse=False, max_tokens=1024, host_cache_size=0), + kv_cache_type=CacheTypeCpp.SELFKONLY, + num_layers=len(sparse_config.compress_ratios), + head_dim=512, + max_batch_size=1, + tokens_per_block=tokens_per_block, + max_seq_len=512, + max_num_tokens=512, + vocab_size=128, + mapping=Mapping(), + dtype=dtype, + sparse_attn_config=sparse_config, + ) + try: + assert manager.get_host_backed_pool_groups() + page_table = build_page_table_from_manager(manager) + assert {pool.memory_type for group in page_table.pool_groups for pool in group.pools} == { + "DRAM", + "VRAM", + } + assert set(manager.get_sparse_offload_descriptors()) == {0} + assert set(manager._sparse_arena_buffers) == {0} + data, scales = manager._sparse_arena_buffers[0] + assert data.is_cuda + if dtype == DataType.NVFP4: + assert scales is not None and scales.is_cuda + else: + assert scales is None + finally: + manager.shutdown() @pytest.mark.parametrize(("avg_seq_len", "expected"), [(None, 1024), (256, 256)]) @@ -532,6 +1871,7 @@ class FakeModelConfig: compress_ratios=[1, 4, 1, 128], indexer_k_dtype="fp8", window_size=128, + enable_kv_cache_offload=False, ) pretrained_config = SimpleNamespace( kv_lora_rank=512, @@ -556,6 +1896,7 @@ def get_num_attention_layers(self) -> int: def test_quota_from_max_tokens_models_context_swa_scratch(): manager = object.__new__(DeepseekV4CacheManager) + manager._enable_kv_cache_offload = False manager.pp_layers = [0, 1] manager._compress_ratios = [4, 4] manager.dtype = DataType.BF16 @@ -592,6 +1933,7 @@ def test_quota_from_max_tokens_models_context_swa_scratch(): def test_needed_resource_uses_context_swa_scratch_slope(): manager = object.__new__(DeepseekV4CacheManager) + manager._enable_kv_cache_offload = False manager.pp_layers = [0, 1] manager._compress_ratios = [4, 4] manager.dtype = DataType.BF16 @@ -782,6 +2124,7 @@ def _create_deepseek_v4_cache_manager( enable_swa_scratch_reuse: bool = True, cold_page_codec_provider: object | None = None, host_cache_size: int | None = None, + max_gpu_total_bytes: int = 0, ) -> Tuple[DeepseekV4CacheManager, DeepSeekV4SparseAttentionConfig]: """Helper to create a DeepseekV4CacheManager for testing.""" @@ -805,6 +2148,7 @@ def _create_deepseek_v4_cache_manager( event_buffer_max_size=0, enable_swa_scratch_reuse=enable_swa_scratch_reuse, host_cache_size=host_cache_size, + max_gpu_total_bytes=max_gpu_total_bytes, ) # Create mapping (single GPU, no parallelism) @@ -946,6 +2290,10 @@ def test_nvfp4_cold_page_codec_migrates_real_csa_hca_through_host( enable_swa_scratch_reuse=False, cold_page_codec_provider=provider, host_cache_size=64 << 20, + # Below what the constraints need, so the manager sizes every pool group at + # its floor: room for one max-length request, which the pressure request + # then has to evict. + max_gpu_total_bytes=1 << 20, ) try: @@ -1413,10 +2761,13 @@ def _reference_copy_batch_page_indices( int(copy_idx[row]), 0, ].tolist() - scratch = None + scratch_range = None + scratch_slots = [] if row < num_contexts and page_index_mode == PageIndexMode.PER_LAYER: - scratch = cache_manager.kv_cache_map[request_id].get_scratch_desc(pool_id) - converted = converter(base_indices, page_index_mode, scratch) + cache = cache_manager.kv_cache_map[request_id] + scratch_range = cache.get_swa_scratch_range(pool_id) or None + scratch_slots = cache.get_scratch_slots(pool_id) + converted = converter(base_indices, page_index_mode, scratch_range, scratch_slots) expected[row, : len(converted)] = torch.tensor(converted, dtype=torch.int32) return expected @@ -2643,6 +3994,221 @@ def test_check_invalid_values_in_kv_cache( cache_manager.shutdown() +# Offload memory accounting. + +V4_FLASH_COMPRESS_RATIOS = [1, 1] + [4, 128] * 20 + [4] + + +def _v4_flash_offload_manager( + *, max_batch_size: int = 256, cache_factor: float = 2.0 +) -> DeepseekV4CacheManager: + manager, _ = _cache_config_case( + enabled=True, + pp_layers=list(range(len(V4_FLASH_COMPRESS_RATIOS))), + dtype=DataType.NVFP4, + nvfp4_residual_dim=64, + indexer_k_dtype="fp4", + cache_factor=cache_factor, + ) + manager._compress_ratios = V4_FLASH_COMPRESS_RATIOS + manager.compressed_block_sizes = [ + manager.tokens_per_block // r for r in V4_FLASH_COMPRESS_RATIOS + ] + manager.max_batch_size = max_batch_size + manager._max_num_tokens = 8192 + manager.enable_swa_scratch_reuse = False + manager._generation_kv_capacity_headroom = 1 + return manager + + +def _v4_flash_offload_config(manager: DeepseekV4CacheManager) -> KVCacheManagerConfig: + """Return the manager's cache config.""" + return manager._build_cache_config( + KVCacheManagerConfig( + tokens_per_block=manager.tokens_per_block, + cache_tiers=[GpuCacheTierConfig(quota=1 << 30)], + layers=[], + constraints=[BatchDesc([KVCacheDesc(capacity=0, history_length=0)] * 513)], + ) + ) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize(("cache_factor", "lease_bytes"), [(1.0, 3_747_200), (2.0, 7_273_856)]) +def test_offload_splits_v4_flash_nvfp4_history_by_memory_kind( + cache_factor: float, lease_bytes: int +) -> None: + manager = _v4_flash_offload_manager(cache_factor=cache_factor) + config = _v4_flash_offload_config(manager) + + # 21 CSA layers keep 81 B per token in host memory: 288 B of NVFP4 data and + # 36 B of scales per row of four tokens. + assert sparse_history_bytes_per_block(config.layers) == 128 * 21 * 81 + # The GPU keeps the MXFP4 indexer keys (17 B per token) and the HCA history + # (324 B per 128 tokens): 407.625 B per token, reported in whole bytes. + assert manager.get_cache_bytes_per_token() == 408 + (arena,) = sparse_arena_shapes(config.layers, tokens_per_block=manager.tokens_per_block) + assert arena.lease_bytes == lease_bytes + assert sparse_arena_bytes(config, 513) == arena.arena_bytes(513) + + manager._enable_kv_cache_offload = False + config = _v4_flash_offload_config(manager) + assert manager.get_cache_bytes_per_token() == 2098 + assert sparse_history_bytes_per_block(config.layers) == 0 + assert sparse_arena_bytes(config, 513) == 0 + + +@pytest.mark.cpu_only +def test_offload_scratch_counts_both_attention_metadata_copies() -> None: + manager = _v4_flash_offload_manager() + config = _v4_flash_offload_config(manager) + + # A copy holds 64 MiB of FP8 rows dequantized from NVFP4, [256, 512, 512], with a + # zeroed tail of one FMHA KV tile (128 rows of 512), and int32 indices: the + # selections [256, 512], the fetch outputs [256, 2, 512] and write destinations + # [256, 2] of the data and scale buffers, the scatter start positions [256], the + # Batch row map [256] and the active count. + per_copy = (256 * 512 + 128) * 512 + 4 * (256 * 512 + 256 * 2 * 512 + 256 * 2 + 256 + 256 + 1) + assert per_copy == 68_751_364 + # The eager attention metadata holds one, and the CUDA graph metadata share one. + assert manager.get_sparse_offload_scratch_bytes(config) == 2 * per_copy + + manager._enable_kv_cache_offload = False + assert manager.get_sparse_offload_scratch_bytes(_v4_flash_offload_config(manager)) == 0 + + +@pytest.mark.cpu_only +def test_offload_counts_gpu_history_and_window_pages_exactly() -> None: + manager = _v4_flash_offload_manager() + max_tokens = 1 << 20 + + # Every window can span two pages: 43 SWA layers, 21 CSA and 20 HCA + # compressor states. + assert manager._get_cache_bytes_for_tokens(0, context=False) == 81_657_856 + # 52,176 GPU bytes per 128 tokens, with no row rounded down. + assert manager._get_cache_bytes_for_tokens(max_tokens, context=False) == ( + max_tokens * 52_176 // 128 + 81_657_856 + ) + assert manager.get_request_device_bytes(max_tokens) == ( + manager._get_cache_bytes_for_tokens(max_tokens, context=False) + ) + quota = manager._get_quota_from_max_tokens(max_tokens) + assert manager._get_max_tokens_from_quota(quota) == pytest.approx(max_tokens) + + # Without offload the estimates are unchanged: whole bytes per token and one + # page per window. + manager._enable_kv_cache_offload = False + assert manager.get_cache_bytes_per_token() == 2098 + assert manager._get_cache_bytes_for_tokens(0, context=False) == 40_828_928 + assert manager._get_cache_bytes_for_tokens(max_tokens, context=False) == ( + max_tokens * 2098 + 40_828_928 + ) + + +@pytest.mark.cpu_only +def test_sparse_offload_rejects_kv_pool_rebalance() -> None: + sparse_config = DeepSeekV4SparseAttentionConfig(enable_kv_cache_offload=True) + with ( + patch.object(deepseek_v4_cache.KVCacheManagerV2, "__init__") as initialize_cache, + pytest.raises(ValueError, match="enable_kv_pool_rebalance"), + ): + DeepseekV4CacheManager( + kv_cache_config=KvCacheConfig(enable_block_reuse=False, enable_kv_pool_rebalance=True), + kv_cache_type=CacheTypeCpp.SELFKONLY, + num_layers=len(sparse_config.compress_ratios), + max_batch_size=1, + tokens_per_block=128, + max_seq_len=256, + vocab_size=128, + mapping=Mapping(), + sparse_attn_config=sparse_config, + ) + + initialize_cache.assert_not_called() + + +@pytest.mark.cpu_only +def test_sparse_offload_rejects_speculative_decoding() -> None: + sparse_config = DeepSeekV4SparseAttentionConfig(enable_kv_cache_offload=True) + with ( + patch.object(deepseek_v4_cache.KVCacheManagerV2, "__init__") as initialize_cache, + pytest.raises(ValueError, match="does not support speculative decoding"), + ): + DeepseekV4CacheManager( + kv_cache_config=KvCacheConfig(enable_block_reuse=False), + kv_cache_type=CacheTypeCpp.SELFKONLY, + num_layers=len(sparse_config.compress_ratios), + max_batch_size=1, + tokens_per_block=128, + max_seq_len=256, + vocab_size=128, + mapping=Mapping(), + sparse_attn_config=sparse_config, + spec_config=MTPDecodingConfig(max_draft_len=1), + ) + + initialize_cache.assert_not_called() + + +@pytest.mark.cpu_only +def test_offload_capacity_report_names_the_binding_limit() -> None: + manager = _v4_flash_offload_manager() + manager.max_seq_len = 1 << 17 + manager.mapping = Mapping() + manager.kv_cache_manager_py_config = _v4_flash_offload_config(manager) + # Half of a 100 GiB share of host memory, the automatic budget on a node of + # four ranks with 400 GiB available. + manager._host_history_budget = 50 << 30 + # The native group holds 4,096 slots: four requests of 128K tokens. Its pool of + # sparse pages holds 217,728 host bytes per slot; its indexer pool is on the GPU. + sparse_layer = next( + layer for layer in manager.kv_cache_manager_py_config.layers if layer.sparse_cache + ) + sparse_buffer = next(buffer for buffer in sparse_layer.buffers if buffer.is_sparse) + dense_buffer = next(buffer for buffer in sparse_layer.buffers if not buffer.is_sparse) + host_pool = SimpleNamespace( + num_slots=4096, + slot_desc=SimpleNamespace( + variants=[ + SimpleNamespace( + coalesced_buffers=[ + SimpleNamespace( + buffer_ids=[ + SimpleNamespace(layer_id=sparse_layer.layer_id, role=buffer.role) + ] + ) + for buffer in (sparse_buffer, dense_buffer) + ] + ) + ] + ), + pools=[ + SimpleNamespace(pool_index=0, slot_bytes=217_728), + SimpleNamespace(pool_index=1, slot_bytes=52_176), + ], + ) + manager.get_host_backed_pool_groups = lambda: [host_pool] + # The device quota counts only device-backed pools. + manager.impl = SimpleNamespace(get_quota=lambda level: 100 << 30) + + with ( + patch.object(kv_cache_v2_module.logger, "info") as info, + patch.object(kv_cache_v2_module.logger, "warning") as warning, + ): + manager._log_sparse_kv_offload_capacity(KvCacheConfig()) + + report = info.call_args.args[0] + # The DeepSeek-V4 estimator sizes each 128K-token request at 128.8 MiB of GPU pools. + assert "128.8MiB per request" in report + assert "4 by host" in report + assert "limited by host" in report + # One lease per Batch row, not per constrained request: max_batch_size, not 513. + assert "256 leases" in report + # The 50 GiB host budget holds 240 requests of 128K tokens. + assert "A max_batch_size of 240 fits" in report + assert "Only 4 of max_batch_size=256" in warning.call_args.args[0] + + if __name__ == "__main__": tester = TestDeepseekV4CacheManager() print("=== FP8, prompt_lens=[1024, 2048, 4096], steps=100 ===") diff --git a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_kernels.py b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_kernels.py index 5548c96c65c3..75c161302ce4 100644 --- a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_kernels.py +++ b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_kernels.py @@ -13,23 +13,141 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Tests for DeepSeek-V4 sparse offload kernels.""" +"""Tests for compact sparse history selection and DeepSeek-V4 index conversion.""" import pytest import torch from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.kernels import ( - check_sparse_read_table, deepseek_v4_local_to_global_indices, - merge_sparse_read_table, - prepare_sparse_write_table, - select_sparse_history_pages, +) +from tensorrt_llm._torch.attention.backends.sparse.offload_kernels import ( + select_sparse_history_units, ) _requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") -# Sparse history selection and read-table kernels. +@_requires_cuda +@pytest.mark.parametrize("strided", [False, True]) +@pytest.mark.parametrize("active", [0, 2, 4]) +@pytest.mark.parametrize("split_extra", [False, True]) +@pytest.mark.parametrize("batch_rows", [[0, 1, 2, 3], [2, 0, 3, 1]]) +def test_sparse_index_preparation_merges_history_and_writes( + strided, active, split_extra, batch_rows +): + # Request order differs from query order, and Batch rows need not follow + # either. Some rows complete a unit, one has no newly completed row, and + # another has no valid write destination. + topk = _cuda_tensor([[1, -1, 0, 99], [32, 31, 0, -2], [0, -1, -1, -1], [0, 1, 2, -1]], strided) + requests = _cuda_tensor([1, 0, 3, 2], strided) + lengths = _cuda_tensor([33, 2, 3, 1], strided) + # Watermarks, fetches and writes of requests 0..3, stored at their Batch rows. + by_row = torch.tensor(batch_rows).argsort() + history = _cuda_tensor(torch.tensor([131, 8, 11, 3])[by_row], strided) + fetched = _cuda_tensor((torch.arange(4 * 3 * 4).reshape(4, 3, 4) + 100)[by_row], strided)[:, 1] + writes = _cuda_tensor( + torch.tensor([[7, 7, 50], [-1, -1, -1], [5, 5, 42], [-1, -1, -1]])[by_row], strided + )[:, 1] + kwargs = dict( + req_id=requests, + block_table_swa=_cuda_tensor([[2], [3], [5], [9]]), + swa_local_indices=_cuda_tensor([[0, 1, -1, 3]] * 4), + swa_pool_base_ptr=4096, + swa_buffer_ptr=4096, + tokens_per_block=4, + token_stride=1024, + compress_ratio=4, + num_compressed_indices=4, + compressed_local_indices=topk, + compressed_fetched_indices=fetched, + compressed_write_indices=writes, + compressed_history_lengths=history, + compressed_lengths=lengths, + compressed_batch_rows=_cuda_tensor(batch_rows, strided), + active_request_count=_cuda_tensor([active]), + mask_invalid_pages=True, + split_extra=split_extra, + ) + out = deepseek_v4_local_to_global_indices(**kwargs) + expected = torch.full((4, 8), -1, dtype=torch.int32) + for query, request in enumerate(requests.cpu().tolist()): + if query >= active or request >= active: + continue + # The SWA table and the lengths follow the request; sparse state, its row. + page = [2, 3, 5, 9][request] + row = batch_rows[request] + expected[query, :4] = torch.tensor([page * 4, page * 4 + 1, -1, page * 4 + 3]) + for slot, unit in enumerate(topk.cpu().tolist()[query]): + if not 0 <= unit < lengths[request].item(): + continue + if unit < history[row].item() // 4: + expected[query, 4 + slot] = fetched[row, slot].cpu() + elif unit == history[row].item() // 4: + expected[query, 4 + slot] = writes[row].cpu() + if split_extra: + torch.testing.assert_close(out[0].cpu(), expected[:, :4]) + torch.testing.assert_close(out[1].cpu(), expected[:, 4:]) + else: + torch.testing.assert_close(out.cpu(), expected) + + +@_requires_cuda +def test_sparse_index_preparation_replays_with_new_watermarks_and_padding(): + # The graph has two query rows; native Batch retains four rows. Invalid + # request IDs in graph padding must never index either cache's metadata. + requests = _cuda_tensor([0, 999]) + topk = _cuda_tensor([[0, 1, 2, -1], [0, 1, 2, 3]]) + fetched = _cuda_tensor([[71, 29, 88, -1]] * 4) + writes = _cuda_tensor([7, -1, -1, -1], strided=True) + history = _cuda_tensor([7, 0, 0, 0]) + lengths = _cuda_tensor([2, 0]) + batch_rows = _cuda_tensor([0, -1, -1, -1]) + active = _cuda_tensor([1]) + kwargs = dict( + req_id=requests, + block_table_swa=_cuda_tensor([[2], [3]]), + swa_local_indices=_cuda_tensor([[0, 1, -1, 3]] * 2), + swa_pool_base_ptr=4096, + swa_buffer_ptr=4096, + tokens_per_block=4, + token_stride=1024, + compress_ratio=4, + num_compressed_indices=4, + compressed_local_indices=topk, + compressed_fetched_indices=fetched, + compressed_write_indices=writes, + compressed_history_lengths=history, + compressed_lengths=lengths, + compressed_batch_rows=batch_rows, + active_request_count=active, + mask_invalid_pages=True, + ) + deepseek_v4_local_to_global_indices(**kwargs) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = deepseek_v4_local_to_global_indices(**kwargs) + for row, history_tokens, length, destination, active_count, compressed in ( + (0, 7, 2, 7, 1, [71, 7, -1, -1]), + (0, 8, 2, -1, 1, [71, 29, -1, -1]), + (0, 11, 3, 5, 1, [71, 29, 5, -1]), + (0, 11, 3, -1, 1, [71, 29, -1, -1]), + (0, 11, 3, 5, 0, [-1] * 4), + # The request now sits in Batch row 2; the replay follows the new map. + (2, 11, 3, 5, 1, [71, 29, 5, -1]), + ): + history.zero_() + writes.fill_(-1) + batch_rows[0] = row + history[row] = history_tokens + lengths[0] = length + writes[row] = destination + active[0] = active_count + graph.replay() + swa = [8, 9, -1, 11] if active_count else [-1] * 4 + torch.testing.assert_close( + out.cpu(), torch.tensor([swa + compressed, [-1] * 8], dtype=torch.int32) + ) def _cuda_tensor(value: torch.Tensor | list, strided: bool = False) -> torch.Tensor: @@ -44,298 +162,177 @@ def _cuda_tensor(value: torch.Tensor | list, strided: bool = False) -> torch.Ten return view -def _selection_reference( - topk: torch.Tensor, - request_rows: torch.Tensor, - active_count: int, - lengths: torch.Tensor, - raw: torch.Tensor, - history: torch.Tensor, - tokens_per_page: int, - width: int, -) -> torch.Tensor: - result = torch.full((raw.shape[0], width), -1, dtype=torch.int32) - for query, request in enumerate(request_rows.tolist()): - if query >= active_count or not 0 <= request < min(active_count, raw.shape[0]): +def _selection_reference(topk, requests, batch_rows, active, lengths, history): + out = torch.full((history.numel(), topk.shape[1]), -1, dtype=torch.int32) + for query, request in enumerate(requests.tolist()): + if query >= active or not 0 <= request < min(active, history.numel(), lengths.numel()): continue - pages = set() - for token in topk[query].tolist(): - if 0 <= token < lengths[request]: - ordinal = token // tokens_per_page - if ordinal < min(history[request], raw.shape[1]) and raw[request, ordinal] >= 0: - pages.add(ordinal) - ordered = sorted(pages) - result[request, : len(ordered)] = torch.tensor(ordered, dtype=torch.int32) - return result - - -def _merge_reference( - fetched: torch.Tensor, - write: torch.Tensor, - history: torch.Tensor, - active_count: int, - scale: int, -) -> torch.Tensor: - result = torch.full_like(write, -1) - for request in range(min(active_count, write.shape[0])): - for ordinal in range(write.shape[1]): - if ordinal < history[request]: - page = int(fetched[request, ordinal]) * scale - else: - page = int(write[request, ordinal]) - if 0 <= page <= torch.iinfo(torch.int32).max: - result[request, ordinal] = page - return result + row = batch_rows[request].item() + if not 0 <= row < history.numel(): + continue + for k, unit in enumerate(topk[query].tolist()): + if 0 <= unit < min(lengths[request], history[row] // 4): + out[row, k] = unit + return out @_requires_cuda -@pytest.mark.parametrize("tokens_per_block", [128, 256]) @pytest.mark.parametrize("strided", [False, True]) -def test_select_history_pages_boundaries(tokens_per_block: int, strided: bool) -> None: - p = tokens_per_block // 4 - original_history = torch.tensor( - [ - tokens_per_block - 1, - tokens_per_block, - tokens_per_block + 1, - 3 * tokens_per_block + 7, - 2 * tokens_per_block, - 4 * tokens_per_block, - 10000, - 10000, - ], - dtype=torch.int32, - ) - history = original_history // tokens_per_block - lengths = (original_history + 1) // 4 - raw = torch.arange(8 * 6, dtype=torch.int32).reshape(8, 6) + 17 - raw[1, 0] = 0 # Host slot zero must not be mistaken for an absent page. - raw[2, 0] = -1 - raw[3, 1] = -1 - by_request = [ - [0, p - 1, p, -1], - [p - 1, 0, 0, p, -1], - [0, p - 1, p, -1], - [2 * p + 5, 0, p + 1, p + 1, 3 * p, 3 * p + 2, -1], - [p - 1, p, 2 * p - 1, 2 * p, 0, -2], - [3 * p + 1, 2 * p, p, 0, 0, 3 * p - 1, 4 * p, -1], - ] - request_rows = torch.tensor([5, 2, 0, 4, 1, 3, 0, 999], dtype=torch.int32) - topk = torch.full((8, 16), -1, dtype=torch.int32) - for query, request in enumerate(request_rows[:6].tolist()): - tokens = by_request[request] - topk[query, : len(tokens)] = torch.tensor(tokens, dtype=torch.int32) - topk[6:] = 0 # Stale graph padding, including a duplicate live request ID. - topk_gpu = _cuda_tensor(topk, strided) - raw_gpu = _cuda_tensor(raw, strided) - out = _cuda_tensor(torch.full((8, 9), 999, dtype=torch.int32), strided) - active_count = _cuda_tensor([6]) - args = ( - topk_gpu, - _cuda_tensor(request_rows, strided), - active_count, +@pytest.mark.parametrize("permuted", [False, True]) +def test_select_history_units_uses_compressed_boundaries_and_preserves_slots(strided, permuted): + # Requests 0..7 in Batch rows; when permuted, rows 4 and 6 are free. + rows = [9, 2, 5, 0, 7, 1, 3, 8] if permuted else list(range(8)) + batch_rows = torch.tensor(rows + [-1, -1], dtype=torch.int32) + request_history = torch.tensor([0, 3, 4, 7, 8, 127, 128, 129], dtype=torch.int32) + history = torch.full((10,), 999, dtype=torch.int32) + history[rows] = request_history + lengths = (request_history + 1) // 4 + requests = torch.tensor([7, 4, 1, 6, 0, 5, 3, 2], dtype=torch.int32) + topk = torch.tensor([[31, -1, 0, 32, 1, 30, -2, 1000]] * 8, dtype=torch.int32) + out = _cuda_tensor(torch.full((10, 8), 999, dtype=torch.int32), strided) + gpu_topk = _cuda_tensor(topk, strided) + select_sparse_history_units( + gpu_topk, + _cuda_tensor(requests, strided), + _cuda_tensor(batch_rows, strided), + _cuda_tensor([8]), _cuda_tensor(lengths, strided), - raw_gpu, _cuda_tensor(history, strided), - p, out, + unit_tokens=4, ) - - select_sparse_history_pages(*args) - - expected = _selection_reference(topk, request_rows, 6, lengths, raw, history, p, 9) - torch.testing.assert_close(out.cpu(), expected, rtol=0, atol=0) - assert out[1, 0].item() == 0 - assert (out[0] == -1).all() # A newly completed page is still the resident tail. - torch.testing.assert_close(topk_gpu.cpu(), topk, rtol=0, atol=0) - torch.testing.assert_close(raw_gpu.cpu(), raw, rtol=0, atol=0) - - active_count.zero_() - select_sparse_history_pages(*args) - assert (out == -1).all() + expected = _selection_reference(topk, requests, batch_rows, 8, lengths, history) + torch.testing.assert_close(out.cpu(), expected) + torch.testing.assert_close(gpu_topk.cpu(), topk) + # History token 7 already makes compressed unit 0 eligible. + assert out[rows[3], 2].item() == 0 + # Unit 1 is still in flight at history 7 / capacity 8. + assert out[rows[3], 4].item() == -1 + assert out[rows[5], 0].item() == -1 and out[rows[6], 0].item() == 31 + # Full native Batch capacity is cleared, free rows included. + free = sorted(set(range(10)) - set(rows)) + assert (out[free] == -1).all() @_requires_cuda -@pytest.mark.parametrize("tokens_per_page", [32, 64]) -@pytest.mark.parametrize(("topk_size", "max_blocks"), [(1, 1), (7, 3), (513, 1100), (1024, 257)]) -def test_select_history_pages_scattered( - tokens_per_page: int, topk_size: int, max_blocks: int -) -> None: +@pytest.mark.parametrize("topk_size", [1, 7, 513, 2048]) +def test_select_history_units_scattered(topk_size): generator = torch.Generator().manual_seed(42) - tokens = (torch.arange(topk_size) % max_blocks) * tokens_per_page - tokens += torch.arange(topk_size) % tokens_per_page - topk = torch.stack([tokens[torch.randperm(topk_size, generator=generator)] for _ in range(3)]) - topk = topk.to(torch.int32) - raw = torch.arange(3 * max_blocks, dtype=torch.int32).reshape(3, max_blocks) - history = torch.tensor([max_blocks, max_blocks // 2, 0], dtype=torch.int32) - lengths = torch.full((3,), max_blocks * tokens_per_page, dtype=torch.int32) + topk = torch.stack( + [torch.randperm(3000, generator=generator)[:topk_size] for _ in range(3)] + ).to(torch.int32) + topk[:, 1::5] = -1 requests = torch.tensor([2, 0, 1], dtype=torch.int32) - width = min(topk_size, max_blocks) - out = torch.empty((3, width), dtype=torch.int32, device="cuda") - args = ( + batch_rows = torch.tensor([3, 0, 4, -1, -1], dtype=torch.int32) + history = torch.tensor([4127, 0, 0, 12000, 0], dtype=torch.int32) + lengths = torch.tensor([3000, 1032, 1], dtype=torch.int32) + out = torch.empty((5, topk_size), dtype=torch.int32, device="cuda") + select_sparse_history_units( _cuda_tensor(topk), _cuda_tensor(requests), + _cuda_tensor(batch_rows), _cuda_tensor([3]), _cuda_tensor(lengths), - _cuda_tensor(raw), _cuda_tensor(history), - tokens_per_page, + out, + unit_tokens=4, ) - - select_sparse_history_pages(*args, out) - - expected = _selection_reference( - topk, requests, 3, lengths, raw, history, tokens_per_page, width + torch.testing.assert_close( + out.cpu(), _selection_reference(topk, requests, batch_rows, 3, lengths, history) ) - torch.testing.assert_close(out.cpu(), expected, rtol=0, atol=0) - with pytest.raises(ValueError, match="selection must not be truncated"): - select_sparse_history_pages(*args, out[:, : width - 1]) @_requires_cuda -def test_select_history_pages_masks_invalid_request_ids() -> None: - out = torch.full((4, 1), 999, dtype=torch.int32, device="cuda") - select_sparse_history_pages( - _cuda_tensor([[0], [0], [0], [0]]), - _cuda_tensor([-1, 999, 1, 0]), - _cuda_tensor([4]), - _cuda_tensor([32] * 4), - _cuda_tensor([[0], [1], [2], [3]]), - _cuda_tensor([1] * 4), - 32, +def test_select_history_units_masks_invalid_request_ids_rows_and_padding(): + # Queries 0 and 1 carry invalid request IDs, query 3 is graph padding, and + # request 1 maps to no Batch row. + out = torch.full((6, 2), 99, dtype=torch.int32, device="cuda") + select_sparse_history_units( + _cuda_tensor([[0, 1]] * 4), + _cuda_tensor([-1, 999, 0, 0]), + _cuda_tensor([5, -1, 2, -1, -1, -1]), + _cuda_tensor([3]), + _cuda_tensor([2] * 4), + _cuda_tensor([8] * 6), out, + unit_tokens=4, ) - torch.testing.assert_close(out.cpu(), torch.tensor([[0], [0], [-1], [-1]], dtype=torch.int32)) - - -@_requires_cuda -@pytest.mark.parametrize(("batch", "max_blocks", "topk"), [(0, 3, 7), (2, 0, 7), (2, 3, 0)]) -def test_sparse_offload_empty_shapes(batch: int, max_blocks: int, topk: int) -> None: - raw = torch.empty((batch, max_blocks), dtype=torch.int32, device="cuda") - out = torch.full((batch, max(1, min(topk, max_blocks))), 999, dtype=torch.int32, device="cuda") - history = _cuda_tensor([0] * batch) - active_count = _cuda_tensor([0]) - select_sparse_history_pages( - torch.empty((batch, topk), dtype=torch.int32, device="cuda"), - _cuda_tensor(list(range(batch))), - active_count, - _cuda_tensor([0] * batch), - raw, - history, - 32, + torch.testing.assert_close(out[5], _cuda_tensor([0, 1])) + assert (out[:5] == -1).all() + out.fill_(99) + select_sparse_history_units( + _cuda_tensor([[0, 1]] * 2), + _cuda_tensor([1, 0]), + _cuda_tensor([5, -1, 2, -1, -1, -1]), + _cuda_tensor([2]), + _cuda_tensor([2] * 2), + _cuda_tensor([8] * 6), out, + unit_tokens=4, ) - assert (out == -1).all() - read = torch.full_like(raw, 999) - merge_sparse_read_table(raw, raw, history, active_count, read, fetched_page_scale=1) - assert (read == -1).all() + torch.testing.assert_close(out[5], _cuda_tensor([0, 1])) + assert (out[:5] == -1).all() @_requires_cuda -@pytest.mark.parametrize("scale", [1, 4]) -@pytest.mark.parametrize("strided", [False, True]) -def test_merge_sparse_read_table(scale: int, strided: bool) -> None: - generator = torch.Generator().manual_seed(17) - fetched = torch.randint(0, 4000, (5, 513), generator=generator, dtype=torch.int32) - write = torch.randint(0, 4000, (5, 513), generator=generator, dtype=torch.int32) - fetched[:, ::17] = -1 - fetched[1, 0] = 0 - fetched[2, 1] = 1 << 30 # Scaling must not wrap this to a valid page zero. - write[:, ::11] = -1 - write[0, 0] = 0 - history = torch.tensor([0, 1, 257, 513, 2], dtype=torch.int32) - fetched_gpu = _cuda_tensor(fetched, strided) - write_gpu = _cuda_tensor(write, strided) - out = _cuda_tensor(torch.full_like(write, 999), strided) - - merge_sparse_read_table( - fetched_gpu, - write_gpu, - _cuda_tensor(history, strided), - _cuda_tensor([4]), +@pytest.mark.parametrize(("batch", "queries", "topk"), [(4, 0, 8), (0, 0, 8), (4, 2, 0)]) +def test_select_history_units_empty_shapes(batch, queries, topk): + out = torch.full((batch, topk), 99, dtype=torch.int32, device="cuda") + select_sparse_history_units( + torch.empty((queries, topk), dtype=torch.int32, device="cuda"), + torch.zeros(queries, dtype=torch.int32, device="cuda"), + torch.full((batch,), -1, dtype=torch.int32, device="cuda"), + _cuda_tensor([0]), + torch.zeros(queries, dtype=torch.int32, device="cuda"), + torch.zeros(batch, dtype=torch.int32, device="cuda"), out, - fetched_page_scale=scale, + unit_tokens=4, ) - - expected = _merge_reference(fetched, write, history, 4, scale) - torch.testing.assert_close(out.cpu(), expected, rtol=0, atol=0) - torch.testing.assert_close(fetched_gpu.cpu(), fetched, rtol=0, atol=0) - torch.testing.assert_close(write_gpu.cpu(), write, rtol=0, atol=0) + assert (out == -1).all() @_requires_cuda -def test_sparse_offload_cuda_graph_replay() -> None: - batch, max_blocks, topk_size, p = 3, 6, 8, 32 - topk = torch.zeros((batch, topk_size), dtype=torch.int32, device="cuda") - requests = _cuda_tensor([2, 0, 1]) - active = _cuda_tensor([3]) - lengths = _cuda_tensor([128] * batch) - history = _cuda_tensor([2, 0, 1]) - raw = torch.zeros((batch, max_blocks), dtype=torch.int32, device="cuda") - fetched = torch.zeros_like(raw) - write = torch.zeros_like(raw) - selected = torch.empty_like(raw) - read = torch.empty_like(raw) - - def run() -> None: - select_sparse_history_pages(topk, requests, active, lengths, raw, history, p, selected) - merge_sparse_read_table(fetched, write, history, active, read, fetched_page_scale=4) - - stream = torch.cuda.Stream() - stream.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(stream): - run() - stream.synchronize() - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph, stream=stream): - run() - pointers = selected.data_ptr(), read.data_ptr() - - for n_active, request_ids, history_counts in ( - (3, [2, 0, 1], [2, 0, 1]), - (1, [0, 0, 0], [3, 2, 1]), - (0, [0, 0, 0], [3, 3, 3]), - (3, [1, 2, 0], [0, 2, 3]), - ): - tokens_cpu = torch.tensor([[64, 0, 32, 95, 0, -1, 96, 127]] * batch, dtype=torch.int32) - requests_cpu = torch.tensor(request_ids, dtype=torch.int32) - history_cpu = torch.tensor(history_counts, dtype=torch.int32) - lengths_cpu = torch.tensor([128, 97, 64], dtype=torch.int32) - raw_cpu = torch.arange(batch * max_blocks, dtype=torch.int32).reshape(batch, max_blocks) - raw_cpu[0, 1] = -1 - expected_selected = _selection_reference( - tokens_cpu, requests_cpu, n_active, lengths_cpu, raw_cpu, history_cpu, p, max_blocks - ) - # Emulate KVCM's mapping input to the merger. No transfer is implemented - # by this test: the independently tested selector and merger are captured. - fetched_cpu = torch.full_like(raw_cpu, -1) - for request in range(n_active): - for position, ordinal in enumerate(expected_selected[request].tolist()): - if ordinal >= 0: - fetched_cpu[request, ordinal] = 100 + request * max_blocks + position - write_cpu = raw_cpu + 500 # Deliberately populated history catches accidental fallback. - for target, source in ( - (topk, tokens_cpu), - (requests, requests_cpu), - (active, torch.tensor([n_active])), - (lengths, lengths_cpu), - (history, history_cpu), - (raw, raw_cpu), - (fetched, fetched_cpu), - (write, write_cpu), - ): - target.copy_(source) - graph.replay() - torch.testing.assert_close(selected.cpu(), expected_selected, rtol=0, atol=0) - torch.testing.assert_close( - read.cpu(), - _merge_reference(fetched_cpu, write_cpu, history_cpu, n_active, 4), - rtol=0, - atol=0, +@pytest.mark.parametrize( + "invalid", + ["dtype", "cpu", "width", "requests", "batch_rows", "history", "lengths", "unit"], +) +def test_select_history_units_rejects_invalid_layout(invalid): + topk = _cuda_tensor([[0, 1], [1, 0]]) + requests, batch_rows, active, lengths, history = ( + _cuda_tensor([0, 1]), + _cuda_tensor([0, 1, -1, -1]), + _cuda_tensor([2]), + _cuda_tensor([2, 2]), + _cuda_tensor([8] * 4), + ) + unit = 4 + if invalid == "dtype": + topk = topk.to(torch.int64) + elif invalid == "cpu": + topk = topk.cpu() + elif invalid == "width": + topk = topk[:, :1] + elif invalid == "requests": + requests = requests[:1] + elif invalid == "batch_rows": + batch_rows = batch_rows[:2] + elif invalid == "history": + history = history[:2] + elif invalid == "lengths": + lengths = lengths[:1] + else: + unit = 0 + with pytest.raises(ValueError): + select_sparse_history_units( + topk, + requests, + batch_rows, + active, + lengths, + history, + torch.empty((4, 2), dtype=torch.int32, device="cuda"), + unit_tokens=unit, ) - torch.testing.assert_close(topk.cpu(), tokens_cpu, rtol=0, atol=0) - torch.testing.assert_close(write.cpu(), write_cpu, rtol=0, atol=0) - assert (selected.data_ptr(), read.data_ptr()) == pointers @_requires_cuda @@ -406,60 +403,8 @@ def reference( torch.testing.assert_close(result.cpu(), expected, rtol=0, atol=0) -@_requires_cuda -@pytest.mark.parametrize("scale", [1, 5]) -def test_prepare_write_table_preserves_invalids_and_strides(scale): - raw = torch.tensor( - [[0, 2, -1, 1 << 30], [3, 0, 2, -9], [7, 8, 9, 10]], dtype=torch.int32, device="cuda" - ) - storage = torch.full((6, 8), 999, dtype=torch.int32, device="cuda") - out = storage[::2, ::2] - history = torch.tensor([1, 2, 0], dtype=torch.int32, device="cuda") - active = torch.tensor([2], dtype=torch.int32, device="cuda") - prepare_sparse_write_table(raw, history, active, out, page_scale=scale) - expected = [ - [-1, 2 * scale, -1, (1 << 30) if scale == 1 else -1], - [-1, -1, 2 * scale, -1], - [-1] * 4, - ] - torch.testing.assert_close(out.cpu(), torch.tensor(expected, dtype=torch.int32)) - assert (storage[1::2] == 999).all() and (storage[:, 1::2] == 999).all() - - -def _cuda(value: list[int] | list[list[int]]) -> torch.Tensor: - return torch.tensor(value, dtype=torch.int32, device="cuda") - - -@_requires_cuda -@pytest.mark.parametrize( - "problem", ["none", "missing", "bound", "length", "ordinal", "request", "active", "padding"] -) -def test_read_coverage_detects_required_pages_without_host_reads(problem: str) -> None: - topk = _cuda([[0, 31, -1, -1], [0, 31, 32, 63], [0] * 4, [0] * 4]) - requests, active, lengths = _cuda([1, 0, 2, 999]), _cuda([2]), _cuda([64] * 4) - read = _cuda([[6, 8, 10], [2, -1, 4], [-1] * 3, [-1] * 3]) - if problem == "missing": - read[0, 1] = -1 - elif problem == "bound": - read[0, 1] = 16 - elif problem == "length": - topk[1, 0] = 64 - elif problem == "ordinal": - topk[1, 0], lengths[0] = 128, 160 - elif problem == "request": - requests[0] = 999 - elif problem == "active": - active.fill_(5) - elif problem == "padding": - active.zero_() # Stale, invalid queries cannot invalidate an empty replay. - valid = _cuda([77]) - check_sparse_read_table(topk, requests, active, lengths, read, 32, 16, valid) - assert valid.item() == int(problem in ("none", "padding")) - # Status is rewritten, rather than sticky from the previous failed step. - active.zero_() - check_sparse_read_table(topk, requests, active, lengths, read, 32, 16, valid) - torch._assert_async(valid) - assert valid.item() == 1 +def _cuda(value): + return _cuda_tensor(value) @_requires_cuda diff --git a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_offload.py b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_offload.py index 7998f8750142..1215ea98c988 100644 --- a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_offload.py +++ b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_offload.py @@ -13,48 +13,42 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Sparse offload metadata, attention integration, and CUDA graph replay tests.""" +"""Native Batch fetch adaptation and persistent DeepSeek-V4 offload metadata.""" -import math -from dataclasses import replace from types import MethodType, SimpleNamespace from unittest.mock import Mock, patch import pytest import torch -import triton -import triton.language as tl -from tensorrt_llm._torch.attention.backends.interface import ( - AttentionForwardArgs, - AttentionInputType, -) -from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.backend import ( - DeepseekV4TrtllmAttention, +from tensorrt_llm._torch.attention.backends.interface import AttentionInputType +from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4 import backend as deepseek_v4_backend +from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4 import ( + cache_manager as deepseek_v4_cache, ) from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.kernels import ( - check_sparse_read_table, - merge_sparse_read_table, - select_sparse_history_pages, + deepseek_v4_local_to_global_indices, ) from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.metadata import ( DeepseekV4TrtllmAttentionMetadata, ) -from tensorrt_llm._torch.attention.backends.sparse.hooks import prepare_sparse_runtime_params -from tensorrt_llm._torch.attention.backends.sparse.params import SparseBackendForwardArgs +from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.params import DeepseekV4AttentionType +from tensorrt_llm._torch.attention.backends.sparse.offload import ( + NVFP4_SCRATCH_TAIL_ROWS, + SparseOffloadState, + _SparseBatchMember, + allocate_sparse_offload_state, +) from tensorrt_llm._torch.memory_buffer_utils import Buffers from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import CUDA_GRAPH_DUMMY_REQUEST_ID from tensorrt_llm.bindings import DataType from tensorrt_llm.llmapi.llm_args import DeepSeekV4SparseAttentionConfig -from .test_deepseek_v4_cache_manager import _history_case, _manager +from .test_deepseek_v4_cache_manager import _manager _requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") -# Offload metadata allocation and per-batch preparation. - - def _metadata(manager, *, enabled=True, graph=False, buffers=None, initialize=True): # Exercise the production offload initialization without constructing the # unrelated indexer/compressor/native attention workspaces. @@ -124,124 +118,218 @@ def test_sparse_offload_stage_without_local_csa_allocates_nothing(): assert _metadata(manager).sparse_offload_state is None +@pytest.mark.cpu_only +def test_sparse_metadata_uses_native_capacity_and_configured_width(): + manager = _manager() + metadata = _metadata(manager, initialize=False) + metadata.max_num_sequences = 2 + with patch.object( + metadata, + "get_empty", + side_effect=lambda buffers, shape, **kwargs: torch.empty(shape, dtype=torch.int32), + ): + metadata._init_sparse_offload_state() + state = metadata.sparse_offload_state + assert state.selected_history_units.shape == (4, 8) + assert state.fetched_indices[1].shape == (4, 1, 8) + assert state.scatter_start_positions.shape == (4,) + assert (state.scatter_start_positions == 0).all() + assert state.batch_rows.shape == (4,) and (state.batch_rows == -1).all() + assert state.history_lengths is None and state.batch is None + assert (state.selected_history_units == -1).all() + + +@pytest.mark.cpu_only +def test_nvfp4_scratch_keeps_a_zeroed_tail_past_its_rows(): + """The sparse MLA FMHA can read one KV tile past the last row its indices name.""" + allocated = {} + + def allocate(name, shape, dtype): + if dtype.itemsize == 1: + tensor = torch.full(shape, 0x7F, dtype=torch.uint8).view(dtype) + else: + tensor = torch.empty(shape, dtype=dtype) + allocated[name] = tensor + return tensor + + layers = {2: SimpleNamespace(buffer_ids=(None, None))} + state = allocate_sparse_offload_state( + allocate, layers, max_batch_size=4, topk=8, nvfp4_head_dim=512 + ) + rows = 4 * 8 + flat = allocated["nvfp4_scratch"] + assert flat.shape == (rows + NVFP4_SCRATCH_TAIL_ROWS, 512) + assert state.nvfp4_scratch.shape == (4, 8, 512) + assert state.nvfp4_scratch.data_ptr() == flat.data_ptr() + assert (flat[:rows].view(torch.uint8) == 0x7F).all() + assert (flat[rows:].view(torch.uint8) == 0).all() + assert SparseOffloadState.device_bytes( + max_batch_size=4, topk=8, buffer_counts={2}, nvfp4_head_dim=512 + ) == sum(tensor.numel() * tensor.element_size() for tensor in allocated.values()) + + +@pytest.mark.cpu_only +def test_sparse_metadata_rejects_different_selection_width(): + metadata = _metadata(_manager(), initialize=False) + metadata.sparse_mla_topk = 4 + with patch.object(metadata, "get_empty") as allocate: + with pytest.raises(ValueError, match="configured positive top-k"): + metadata._init_sparse_offload_state() + allocate.assert_not_called() + + @_requires_cuda @pytest.mark.parametrize("graph", [False, True]) def test_sparse_offload_buffers_are_persistent_and_separate(graph): manager = _manager() + scale = deepseek_v4_cache.COMPRESS_BLOCK_SCALE_ROLE + main = next(iter(manager.impl.layers[19].buffers)).role + manager.impl.configure_buffers(41, [(main, 1), (scale, 0)]) buffers = Buffers() if graph else None metadata = _metadata(manager, graph=graph, buffers=buffers) state = metadata.sparse_offload_state - assert state.history_blocks_host.is_pinned() - assert state.selected_history_pages.shape == (4, 6) - tables = [*state.base_page_tables.values(), state.fetched_page_table, state.compress_read_table] - assert all(table.shape == (4, 6) for table in tables) - device_tensors = tables + [ - state.selected_history_pages, - state.history_blocks, + tensors = [ + state.selected_history_units, state.active_request_count, + state.batch_rows, + state.scatter_start_positions, + *state.fetched_indices.values(), + *state.write_indices.values(), ] - assert all(tensor.is_cuda and tensor.dtype == torch.int32 for tensor in device_tensors) - assert len({tensor.data_ptr() for tensor in device_tensors}) == len(device_tensors) - assert all((table == -1).all() for table in tables) - assert (state.history_blocks == 0).all() and state.active_request_count.item() == 0 + assert state.selected_history_units.shape == (4, 8) + assert state.fetched_indices[1].shape == (4, 1, 8) + assert state.fetched_indices[2].shape == (4, 2, 8) + assert all(t.is_cuda and t.dtype == torch.int32 and t.is_contiguous() for t in tensors) + assert len({t.data_ptr() for t in tensors}) == len(tensors) if graph: - # Same-name graph buffers can be reused for serial graph buckets. - reused = _metadata(manager, graph=True, buffers=buffers).sparse_offload_state - assert reused.base_page_tables[13].data_ptr() == state.base_page_tables[13].data_ptr() - assert reused.compress_read_table.data_ptr() == state.compress_read_table.data_ptr() - assert reused.history_blocks_host.data_ptr() != state.history_blocks_host.data_ptr() + smaller = _metadata(manager, graph=True, buffers=buffers, initialize=False) + smaller.max_num_sequences = 2 + smaller._init_sparse_offload_state() + reused = smaller.sparse_offload_state + assert reused.selected_history_units.data_ptr() == state.selected_history_units.data_ptr() + assert reused.fetched_indices[2].data_ptr() == state.fetched_indices[2].data_ptr() + assert reused.write_indices[2].data_ptr() == state.write_indices[2].data_ptr() + assert reused.scatter_start_positions.data_ptr() == state.scatter_start_positions.data_ptr() + assert reused.batch_rows.data_ptr() == state.batch_rows.data_ptr() + + +class _PublishedBatch: + """Device watermarks with the same stable DLPack-view ownership as native Batch.""" + + def __init__(self, capacity): + self.max_batch_size = capacity + self.history = torch.zeros(capacity, dtype=torch.int32, device="cuda") + + def history_lengths(self): + return self.history + + def sparse_out_slots(self, buffers): + """One slice per buffer: this double's pools are distinct, so none coincide.""" + return list(range(len(buffers))) + + +def _publish(manager, request_ids, rows=None): + """Seat the live requests in ``rows``, by default in input order from row 0.""" + if manager._sparse_batch is None: + manager._sparse_batch = _PublishedBatch(manager.max_batch_size) + live = [request for request in request_ids if request != CUDA_GRAPH_DUMMY_REQUEST_ID] + rows = range(len(live)) if rows is None else rows + manager._sparse_batch_members = { + request: _SparseBatchMember( + cache=manager.kv_cache_map[request], + row=row, + blocks=manager.kv_cache_map[request].num_blocks, + ) + for request, row in zip(live, rows, strict=True) + } + manager._sparse_batch_prepared = True + values = [0] * manager.max_batch_size + for member in manager._sparse_batch_members.values(): + values[member.row] = member.cache.history_length + with torch.cuda.stream(manager._stream): + manager._sparse_batch.history.copy_(torch.tensor(values, dtype=torch.int32, device="cuda")) @_requires_cuda -@pytest.mark.parametrize("tokens_per_block", [128, 256]) -def test_prepare_sparse_offload_request_order_and_boundaries(tokens_per_block): - manager = _manager(tokens_per_block) +def test_prepare_sparse_offload_binds_batch_without_copying_watermarks(): + manager = _manager() manager._stream = torch.cuda.current_stream() - metadata = _metadata(manager) - state = metadata.sparse_offload_state - write = torch.empty((4, 6), device="cuda", dtype=torch.int32) - pointers = [t.data_ptr() for t in (state.history_blocks, state.base_page_tables[13], write)] - prepare_stream = torch.cuda.Stream() - for ids in ( - [30, 10, 20], - [20, CUDA_GRAPH_DUMMY_REQUEST_ID, CUDA_GRAPH_DUMMY_REQUEST_ID], - [], - [10], + state = _metadata(manager).sparse_offload_state + batch = None + for request_ids in ([20, 10], [10, CUDA_GRAPH_DUMMY_REQUEST_ID], [], [20]): + _publish(manager, request_ids) + manager.prepare_sparse_offload(state, request_ids, beam_width=1) + assert state.prepared + assert state.batch is manager._sparse_batch + if batch is not None: + assert state.batch is batch + batch = state.batch + assert state.history_lengths.data_ptr() == batch.history.data_ptr() + live = len(manager._sparse_batch_members) + assert state.active_request_count.item() == live + assert state.batch_rows.tolist() == list(range(live)) + [-1] * (4 - live) + assert (state.selected_history_units == -1).all() + assert (state.fetched_indices[1] == -1).all() + batch.history[0] = 77 + assert state.history_lengths[0].item() == 77 + + +@_requires_cuda +def test_prepare_sparse_offload_maps_input_positions_to_batch_rows(): + """Rows are slots: requests keep theirs whatever their input order.""" + manager = _manager() + manager._stream = torch.cuda.current_stream() + state = _metadata(manager).sparse_offload_state + _publish(manager, [20, 10, 30], rows=[3, 0, 2]) + for request_ids, expected in ( + ([20, 10, 30], [3, 0, 2, -1]), + ([30, 20, 10], [2, 3, 0, -1]), + ([10, 30, 20, CUDA_GRAPH_DUMMY_REQUEST_ID], [0, 2, 3, -1]), ): - state.fetched_page_table.fill_(999) - state.compress_read_table.fill_(999) - state.selected_history_pages.fill_(999) - prepare_stream.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(prepare_stream): - manager.prepare_sparse_offload(state, ids, write, beam_width=1) - torch.cuda.current_stream().wait_stream(prepare_stream) - live_ids = [req for req in ids if req != CUDA_GRAPH_DUMMY_REQUEST_ID] - expected_raw = torch.full((4, 6), -1, dtype=torch.int32) - expected_write = torch.full_like(expected_raw, -1) - expected_history = torch.zeros(4, dtype=torch.int32) - for row, req in enumerate(live_ids): - cache = manager.kv_cache_map[req] - count = cache.history_length // tokens_per_block - expected_history[row] = count - expected_raw[row, : cache.num_blocks] = torch.tensor(cache.pages) - for ordinal, page in enumerate(cache.pages): - if ordinal >= count and page >= 0: - expected_write[row, ordinal] = page * 5 - torch.testing.assert_close(state.base_page_tables[13].cpu(), expected_raw) - torch.testing.assert_close(write.cpu(), expected_write) - torch.testing.assert_close(state.history_blocks.cpu(), expected_history) - assert state.active_request_count.item() == len(live_ids) - assert all( - (t == -1).all() - for t in ( - state.selected_history_pages, - state.fetched_page_table, - state.compress_read_table, - ) - ) - if live_ids: - assert manager.impl.calls[-1][:2] == (live_ids, 13) - assert manager.impl.calls[-1][3] == prepare_stream.cuda_stream - assert pointers == [ - t.data_ptr() for t in (state.history_blocks, state.base_page_tables[13], write) - ] - # Change the frontier and physical assignment without replacing tensors. - manager.kv_cache_map[10].history_length = 2 * tokens_per_block - manager.kv_cache_map[10].pages = [17, 19, 0] - assert len(manager.impl.calls) == 3 # The empty batch does not call KVCM. + manager.prepare_sparse_offload(state, request_ids, beam_width=1) + assert state.batch_rows.tolist() == expected + assert state.active_request_count.item() == 3 @_requires_cuda @pytest.mark.parametrize( "invalid", - ["batch", "duplicates", "dummy_order", "unknown", "width", "frontier", "beam", "capture"], + [ + "unpublished", + "missing", + "stranger", + "duplicate", + "beam", + "capacity", + "replacement", + "capture", + ], ) def test_sparse_offload_rejects_invalid_preparation(invalid): manager = _manager() manager._stream = torch.cuda.current_stream() state = _metadata(manager).sparse_offload_state - write = torch.empty((4, 6), device="cuda", dtype=torch.int32) - ids = [10] - beam = 1 - if invalid == "batch": - ids = [10] * 5 - elif invalid == "duplicates": - ids = [10, 10] - elif invalid == "dummy_order": - ids = [CUDA_GRAPH_DUMMY_REQUEST_ID, 10] - elif invalid == "unknown": - ids = [12345] - elif invalid == "width": - manager.kv_cache_map[10].num_blocks = 7 - elif invalid == "frontier": - manager.kv_cache_map[10].history_length = 99999 + _publish(manager, [20, 10]) + ids, beam = [20, 10], 1 + if invalid == "unpublished": + manager._sparse_batch_prepared = False + elif invalid == "missing": + ids = [20, CUDA_GRAPH_DUMMY_REQUEST_ID] + elif invalid == "stranger": + ids = [20, 30] + elif invalid == "duplicate": + ids = [20, 20] elif invalid == "beam": beam = 2 + elif invalid == "capacity": + manager._sparse_batch.max_batch_size = 5 + elif invalid == "replacement": + manager.prepare_sparse_offload(state, ids, beam_width=1) + manager._sparse_batch = _PublishedBatch(manager.max_batch_size) with patch("torch.cuda.is_current_stream_capturing", return_value=invalid == "capture"): - with pytest.raises((ValueError, KeyError, RuntimeError)): - manager.prepare_sparse_offload(state, ids, write, beam_width=beam) - assert not manager.impl.calls - assert state.active_request_count.item() == 0 + with pytest.raises((ValueError, RuntimeError)): + manager.prepare_sparse_offload(state, ids, beam_width=beam) + assert not state.prepared @_requires_cuda @@ -253,7 +341,7 @@ def test_metadata_routes_only_sparse_main_table(enabled): metadata = _metadata(manager, enabled=enabled) metadata.request_ids = [20, 10] metadata._seq_lens = torch.ones(2, dtype=torch.int32) - metadata._num_contexts = 1 + metadata._num_contexts = 0 metadata.sliding_block_tables = torch.empty((1, 1, 4, 6), dtype=torch.int32, device="cuda") metadata.compress_block_tables = { ratio: torch.empty((4, 6), dtype=torch.int32, device="cuda") for ratio in (4, 128) @@ -262,6 +350,8 @@ def test_metadata_routes_only_sparse_main_table(enabled): manager.copy_batch_compress_block_tables = Mock( side_effect=lambda out, *args, **kwargs: out.fill_(29) ) + if enabled: + _publish(manager, metadata.request_ids) metadata.prepare_for_block_tables() assert (metadata.sliding_block_tables == 71).all() assert (metadata.compress_block_tables[128] == 29).all() @@ -271,84 +361,393 @@ def test_metadata_routes_only_sparse_main_table(enabled): ] assert ratios == ([128] if enabled else [4, 128]) if enabled: - assert manager.impl.calls[-1][0] == [20, 10] - assert metadata.compress_block_tables[4][0, 0].item() == -1 - assert metadata.compress_block_tables[4][0, 1].item() == 0 + assert metadata.sparse_offload_state.batch is manager._sparse_batch + assert (metadata.compress_block_tables[4] == -1).all() else: assert metadata.sparse_offload_state is None - assert not manager.impl.calls -@_requires_cuda -def test_sparse_metadata_refresh_before_graph_replay(): +class _FetchDouble: + """Compact, per-role physical indices; deliberately accepts only the native ABI.""" + + def __init__(self, batch): + # Attached as an attribute of the Batch, which does not bind a receiver, so the + # batch the call is standing in for is supplied here instead. + self.batch = batch + self.calls = [] + + def __call__(self, *, buffers, selected, out, stream): + assert selected.shape == (self.batch.max_batch_size, 8) + assert out.shape == (self.batch.max_batch_size, len(buffers), 8) + assert selected.is_contiguous() and out.is_contiguous() + assert stream == torch.cuda.current_stream().cuda_stream + self.calls.append((self.batch, buffers, selected, out, stream)) + for i, buffer in enumerate(buffers): + # Paired data/scale roles share indices; padding has separate cells. + offset = (0 if buffer.layer_id == 41 else 1000) + ( + 101 if "padding" in buffer.role else 0 + ) + out[:, i, :].copy_(torch.where(selected >= 0, selected + offset, -1)) + + +def _fetch_case(multiple_roles=False, rows=None): manager = _manager() manager._stream = torch.cuda.current_stream() - state = _metadata(manager, graph=True, buffers=Buffers()).sparse_offload_state - write = torch.empty((4, 6), dtype=torch.int32, device="cuda") - topk = torch.tensor([[0, 31, 32, 63, 64, -1, 32, 0]] * 4, dtype=torch.int32, device="cuda") - requests = torch.arange(4, dtype=torch.int32, device="cuda") - lengths = torch.full((4,), 96, dtype=torch.int32, device="cuda") - - def consume(): - select_sparse_history_pages( - topk, - requests, - state.active_request_count, - lengths, - state.base_page_tables[13], - state.history_blocks, - 32, - state.selected_history_pages, + if multiple_roles: + manager._use_nvfp4_compress = True + main = manager.impl.layers[41].buffers[0].role + manager.impl.configure_buffers( + 41, + [ + (main, 1), + (deepseek_v4_cache.COMPRESS_BLOCK_SCALE_ROLE, 0), + ("padding", 1), + ("scale_padding", 0), + ], ) - # No fetch is simulated here. History remains absent, and only the - # prepared resident table can supply valid addresses to the merger. - merge_sparse_read_table( - state.fetched_page_table, - write, - state.history_blocks, - state.active_request_count, - state.compress_read_table, - fetched_page_scale=1, + manager.impl.configure_buffers( + 19, [(main, 1), (deepseek_v4_cache.COMPRESS_BLOCK_SCALE_ROLE, 0)] ) + manager.kv_cache_map[20].history_length = 7 + manager.kv_cache_map[10].history_length = 12 + _publish(manager, [20, 10], rows) + state = _metadata(manager).sparse_offload_state + manager.prepare_sparse_offload(state, [20, 10], beam_width=1) + manager._sparse_batch.fetch_sparse_pages = _FetchDouble(manager._sparse_batch) + topk = torch.tensor( + [[2, -1, 0, 1, -2, 3, 99, -1], [1, 0, -1, 2, -1, -1, -1, -1]], + dtype=torch.int32, + device="cuda", + ) + # Query 0 is request 10 at input position 1; query 1 is request 20 at position 0. + requests = torch.tensor([1, 0], dtype=torch.int32, device="cuda") + lengths = torch.tensor([2, 3], dtype=torch.int32, device="cuda") + return SimpleNamespace( + manager=manager, state=state, topk=topk, requests=requests, lengths=lengths + ) - manager.prepare_sparse_offload(state, [20, 10], write, beam_width=1) - capture_stream = torch.cuda.Stream() - capture_stream.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(capture_stream): - consume() - capture_stream.synchronize() - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph, stream=capture_stream): - consume() - pointers = state.selected_history_pages.data_ptr(), state.compress_read_table.data_ptr() - for ids in ([20, 10], [30, CUDA_GRAPH_DUMMY_REQUEST_ID], [], [10, 30]): - manager.prepare_sparse_offload(state, ids, write, beam_width=1) - graph.replay() - expected = torch.full((4, 6), -1, dtype=torch.int32) - live = [req for req in ids if req != CUDA_GRAPH_DUMMY_REQUEST_ID] - for row, req in enumerate(live): - cache = manager.kv_cache_map[req] - pages = sorted( - { - token // 32 - for token in topk[row].cpu().tolist() - if 0 <= token < 96 - and token // 32 < cache.history_length // 128 - and cache.pages[token // 32] >= 0 - } + +def _fetch(case, layer=2): + return case.manager.fetch_sparse_indices( + case.state, layer, case.topk, case.requests, case.lengths + ) + + +def _expected_selection(manager, by_request): + """Each member's expected selection, in its Batch row; free rows select nothing.""" + expected = torch.full((manager.max_batch_size, 8), -1, dtype=torch.int32, device="cuda") + for request, member in manager._sparse_batch_members.items(): + expected[member.row] = torch.tensor(by_request[request], dtype=torch.int32) + return expected + + +@_requires_cuda +@pytest.mark.parametrize("rows", [None, [3, 1]]) +def test_prepare_writes_and_attention_merge_on_ordered_streams(rows): + case = _fetch_case(rows=rows) + row_20 = case.manager._sparse_batch_members[20].row + producer = torch.cuda.Stream() + seen = [] + + def prepare(*, buffers, out, stream): + assert out.shape == (4, 1) and out.is_contiguous() + assert stream == producer.cuda_stream + assert buffers == case.state.layers[2].buffer_ids + seen.append(out.data_ptr()) + out.fill_(-1) + # Request 20 has history 7 and capacity 8, so its unit 1 is newly complete. + out[row_20, 0] = 77 + + case.manager._sparse_batch.prepare_sparse_write_indices = prepare + producer.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(producer): + writes = case.manager.prepare_sparse_compressor_write_indices(case.state, 2) + torch.cuda.current_stream().wait_stream(producer) + with ( + patch.object(torch.Tensor, "item", side_effect=AssertionError("host read")), + patch.object(torch.Tensor, "cpu", side_effect=AssertionError("host read")), + ): + fetched = _fetch(case) + out = deepseek_v4_local_to_global_indices( + req_id=case.requests, + block_table_swa=torch.zeros((2, 1), dtype=torch.int32, device="cuda"), + swa_local_indices=torch.full((2, 4), -1, dtype=torch.int32, device="cuda"), + swa_pool_base_ptr=4096, + swa_buffer_ptr=4096, + tokens_per_block=128, + token_stride=1024, + compress_ratio=4, + num_compressed_indices=8, + compressed_local_indices=case.topk, + compressed_fetched_indices=fetched[:, 0], + compressed_write_indices=writes[:, 0], + compressed_history_lengths=case.state.history_lengths, + compressed_lengths=case.lengths, + compressed_batch_rows=case.state.batch_rows, + active_request_count=case.state.active_request_count, + mask_invalid_pages=True, + ) + assert seen == [writes.data_ptr()] + expected = torch.tensor( + [[[2, -1, 0, 1, -1, -1, -1, -1]], [[77, 0, -1, -1, -1, -1, -1, -1]]], + dtype=torch.int32, + device="cuda", + ) + torch.testing.assert_close(out[:, 4:], expected[:, 0]) + case.manager.prepare_sparse_offload(case.state, [20, 10], beam_width=1) + assert case.state.write_layer is None + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("nvfp4", [False, True]) +def test_attention_uses_arena_base_and_physical_indices_without_host_pages(nvfp4): + manager = _manager() + manager.dtype = DataType.BF16 + manager.layer_offsets = {2: 0} + physical = torch.tensor([[[71, -1, 9, 0, -1, -1, -1, -1]]], dtype=torch.int32) + if nvfp4: + manager._use_nvfp4_compress = True + manager._nvfp4_residual_dim = 0 + for layer_id in (41, 19): + manager.impl.configure_buffers( + layer_id, + [ + (DeepseekV4AttentionType.COMPRESS.role, 1), + (deepseek_v4_cache.COMPRESS_BLOCK_SCALE_ROLE, 0), + ], ) - expected[row, : len(pages)] = torch.tensor(pages, dtype=torch.int32) - torch.testing.assert_close(state.selected_history_pages.cpu(), expected) - torch.testing.assert_close(state.compress_read_table, write) - assert pointers == ( - state.selected_history_pages.data_ptr(), - state.compress_read_table.data_ptr(), + physical = physical.expand(-1, 2, -1).clone() + physical = physical.expand(4, -1, -1).clone() + manager.fetch_sparse_indices = Mock(return_value=physical) + writes = torch.full((4, physical.shape[1], 1), 77, dtype=torch.int32) + state = SimpleNamespace( + layers=manager.get_sparse_offload_descriptors(), + active_request_count=torch.tensor([1], dtype=torch.int32), + batch_rows=torch.tensor([0, -1, -1, -1], dtype=torch.int32), + write_indices={physical.shape[1]: writes}, + history_lengths=torch.tensor([4, 0, 0, 0], dtype=torch.int32), + write_layer=2, + # The mocked fetch cannot fill this, so stand in for it: one slice per buffer. + out_slots={2: tuple(range(physical.shape[1]))}, + nvfp4_scratch=torch.empty(1, 8, 512, dtype=torch.float8_e4m3fn) if nvfp4 else None, + ) + descriptor = state.layers[2] + if nvfp4: + pool_rows = descriptor.arena_row_counts[descriptor.data_buffer_index] + manager._sparse_arena_buffers[2] = ( + torch.empty((pool_rows, 1, 256), dtype=torch.uint8), + torch.empty((pool_rows, 1, 32), dtype=torch.float8_e4m3fn), ) - manager.kv_cache_map[30].history_length = 2 * 128 - manager.kv_cache_map[30].pages = [0, 16, 3] + metadata = SimpleNamespace( + kv_cache_manager=manager, + sparse_offload_state=state, + num_ctx_tokens=0, + num_tokens=1, + num_seqs=1, + sparse_mla_base_ptrs={1: 4096, 4: 999999}, + swa_buffer_ptrs={2: 4096}, + compressed_buffer_ptrs={2: 888888}, + req_idx_per_token=torch.tensor([0], dtype=torch.int32), + swa_local_indices_cuda=torch.zeros(1, 4, dtype=torch.int32), + sliding_block_tables=torch.zeros(1, 4, 1, 1, dtype=torch.int32), + compress_block_tables={4: None}, + compressed_kv_lens_cuda={4: torch.tensor([2], dtype=torch.int32)}, + max_compressed_indices={4: 8}, + ) + backend = SimpleNamespace( + layer_idx=2, + head_dim=512, + compress_ratio=4, + use_fp8_ds_mla=False, + _uses_nvfp4_compress=nvfp4, + _nvfp4_compress_scale_quant_orig=torch.ones(1), + sparse_attention_config=SimpleNamespace(index_head_dim=128, enable_kv_cache_offload=True), + ) + args = SimpleNamespace( + attention_input_type=AttentionInputType.generation_only, + sparse_backend_args=SimpleNamespace(topk_indices=torch.zeros(1, 8, dtype=torch.int32)), + sparse_runtime_params=SimpleNamespace(aux_kv_cache_pool_ptr=0), + ) + expected = ( + torch.tensor([[0, -1, 2, 3, -1, -1, -1, -1]], dtype=torch.int32) + if nvfp4 + else physical[:1, 0] + ) + prepared = torch.cat((torch.tensor([[10, 11, 12, 13]]), physical[:1, 0]), dim=1).int() + def gather_rows( + data_pool, scale_pool, indices, output, global_scale, residual_dim, num_pool_tokens + ): + assert data_pool is manager._sparse_arena_buffers[2][0] + assert scale_pool is manager._sparse_arena_buffers[2][1] + assert num_pool_tokens == pool_rows + torch.testing.assert_close(indices, physical[:1, descriptor.data_buffer_index]) + assert output.data_ptr() == state.nvfp4_scratch.data_ptr() + indices.copy_(expected) -# Sparse attention fetch and consumption. + with ( + patch.object( + deepseek_v4_backend, "deepseek_v4_local_to_global_indices", return_value=prepared + ) as convert, + patch.object( + torch.ops.trtllm, + "nvfp4_mla_kv_cache_gather_direct", + side_effect=gather_rows, + create=True, + ) as gather, + ): + result = deepseek_v4_backend.DeepseekV4TrtllmAttention.sparse_attn_predict( + backend, None, None, metadata, args + ) + manager.fetch_sparse_indices.assert_called_once() + assert args.sparse_runtime_params.aux_kv_cache_pool_ptr == ( + state.nvfp4_scratch.data_ptr() + if nvfp4 + else descriptor.arena_base_addresses[descriptor.data_buffer_index] + ) + converted = convert.call_args.kwargs + torch.testing.assert_close( + converted["compressed_fetched_indices"], physical[:, descriptor.data_buffer_index] + ) + torch.testing.assert_close( + converted["compressed_write_indices"], writes[:, descriptor.data_buffer_index] + ) + assert converted["compressed_history_lengths"] is state.history_lengths + assert converted["compressed_batch_rows"] is state.batch_rows + torch.testing.assert_close( + result[0][:, :4], torch.tensor([[10, 11, 12, 13]], dtype=torch.int32) + ) + torch.testing.assert_close(result[0][:, 4:], expected) + assert gather.call_count == int(nvfp4) + state.write_layer = None + with pytest.raises(RuntimeError, match="Prepare this layer"): + deepseek_v4_backend.DeepseekV4TrtllmAttention.sparse_attn_predict( + backend, None, None, metadata, args + ) + + +@_requires_cuda +@pytest.mark.parametrize("multiple_roles", [False, True]) +@pytest.mark.parametrize("rows", [None, [2, 0], [3, 1]]) +def test_fetch_uses_batch_rows_and_preserves_topk_slots(multiple_roles, rows): + case = _fetch_case(multiple_roles, rows) + original = case.topk.clone() + with ( + patch.object(torch.Tensor, "item", side_effect=AssertionError("host read")), + patch.object(torch.Tensor, "cpu", side_effect=AssertionError("host read")), + ): + out = _fetch(case) + expected = _expected_selection( + case.manager, + {20: [-1, 0, -1, -1, -1, -1, -1, -1], 10: [2, -1, 0, 1, -1, -1, -1, -1]}, + ) + torch.testing.assert_close(case.state.selected_history_units, expected) + torch.testing.assert_close(case.topk, original) + for i in range(out.shape[1]): + torch.testing.assert_close( + out[:, i], torch.where(expected >= 0, expected + (i % 2) * 101, -1) + ) + call = case.manager._sparse_batch.fetch_sparse_pages.calls[-1] + assert call[0] is case.manager._sparse_batch + assert call[1] == case.state.layers[2].buffer_ids + assert call[2] is case.state.selected_history_units and call[3] is out + second = _fetch(case, layer=4) + torch.testing.assert_close(second[:, 0], torch.where(expected >= 0, expected + 1000, -1)) + assert second is case.state.fetched_indices[2 if multiple_roles else 1] + if not multiple_roles: + assert second is out + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("invalid", ["disabled", "unprepared", "prefill", "batch", "watermarks"]) +def test_fetch_rejects_unprepared_or_prefill_metadata(invalid): + manager = _manager() + state = SimpleNamespace( + prepared=True, is_prefill=False, batch=object(), history_lengths=object() + ) + if invalid == "disabled": + manager._enable_kv_cache_offload = False + elif invalid == "unprepared": + state.prepared = False + elif invalid == "prefill": + state.is_prefill = True + elif invalid == "batch": + state.batch = None + else: + state.history_lengths = None + with pytest.raises(RuntimeError, match="prepared decode metadata"): + manager.fetch_sparse_indices(state, 2, None, None, None) + + +@_requires_cuda +def test_fetch_propagates_native_failure(): + case = _fetch_case() + case.manager._sparse_batch.fetch_sparse_pages = Mock(side_effect=RuntimeError("fetch failed")) + with pytest.raises(RuntimeError, match="fetch failed"): + _fetch(case) + + +@_requires_cuda +@pytest.mark.parametrize("empty", ["topk", "history", "batch"]) +def test_fetch_refreshes_all_outputs_for_empty_selection(empty): + case = _fetch_case(multiple_roles=True) + _fetch(case) + if empty == "topk": + case.topk.fill_(-1) + elif empty == "history": + case.manager._sparse_batch.history.zero_() + else: + _publish(case.manager, []) + case.manager.prepare_sparse_offload( + case.state, [CUDA_GRAPH_DUMMY_REQUEST_ID] * 2, beam_width=1 + ) + out = _fetch(case) + assert (out == -1).all() + assert len(case.manager._sparse_batch.fetch_sparse_pages.calls) == 2 + + +@_requires_cuda +def test_fetch_graph_replay_refreshes_batch_watermarks_and_padding(): + case = _fetch_case(multiple_roles=True) + # Warm up before capture; the graph has two query rows but Batch has four. + _fetch(case) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = _fetch(case) + ptr = out.data_ptr() + for ids, rows, history in ( + ([10, 20], None, 8), + # Both requests moved rows, and row 0 is free: the replay follows the map. + ([10, 20], [3, 1], 8), + ([20, CUDA_GRAPH_DUMMY_REQUEST_ID], [2], 12), + ([], None, 0), + ): + case.manager.kv_cache_map[20].history_length = history + case.manager.kv_cache_map[10].history_length = history + _publish(case.manager, ids, rows) + case.manager.prepare_sparse_offload(case.state, ids, beam_width=1) + case.requests.copy_( + torch.tensor([0, 0], dtype=torch.int32, device="cuda") + if len(case.manager._sparse_batch_members) < 2 + else torch.tensor([1, 0], dtype=torch.int32, device="cuda") + ) + case.topk.copy_( + torch.tensor([[0, -1, 1, 2, -1, -1, -1, -1]] * 2, dtype=torch.int32, device="cuda") + ) + case.lengths.fill_(4) + graph.replay() + expected = torch.full_like(case.state.selected_history_units, -1) + for member in case.manager._sparse_batch_members.values(): + expected[member.row, 0] = 0 + if history // 4 > 1: + expected[member.row, 2] = 1 + if history // 4 > 2: + expected[member.row, 3] = 2 + for role in range(out.shape[1]): + torch.testing.assert_close( + out[:, role], torch.where(expected >= 0, expected + (role % 2) * 101, -1) + ) + assert out.data_ptr() == ptr def _host_batch(*, prefill=False): @@ -421,522 +820,3 @@ def test_offload_rejects_batch_before_table_preparation(unsupported): DeepseekV4TrtllmAttentionMetadata.prepare(metadata) metadata.kv_cache_manager.compute_sliding_block_tables.assert_not_called() assert not metadata.sparse_offload_state.prepared - - -@triton.jit -def _fetch_double_kernel( - raw, - selected, - history, - out, - source, - pool, - M: tl.constexpr, - S: tl.constexpr, - P: tl.constexpr, - D: tl.constexpr, - OFFSET: tl.constexpr, - FETCH_SCALE: tl.constexpr, -): - row, j = tl.program_id(0), tl.program_id(1) - ordinal = tl.load(selected + row * S + j) - valid = (ordinal >= 0) & (ordinal < tl.load(history + row)) - host_slot = tl.load(raw + row * M + ordinal, mask=valid, other=0) - # Positional scratch lives above all resident slots; pages interleave - # layers at scale 5. Keep the original row position within every page. - physical_page = (16 + row * S + j) * 5 - x = tl.arange(0, P * D) - value = tl.load(source + host_slot * P * D + x, mask=valid, other=0) - tl.store(pool + (OFFSET + physical_page * P) * D + x, value, mask=valid) - tl.store(out + row * M + ordinal, physical_page // FETCH_SCALE, mask=valid) - - -class _FetchDouble: - def __init__(self, state, pool, sources, offsets, page_tokens, calls): - self.state, self.pool, self.sources = state, pool, sources - self.offsets, self.page_tokens, self.calls = offsets, page_tokens, calls - - def __call__(self, *, buffers, page_table, selected, num_blocks, out, stream): - assert len(buffers) == 1 - descriptor = next(d for d in self.state.layers.values() if d.buffer_id == buffers[0]) - assert page_table is self.state.base_page_tables[descriptor.group_id] - assert selected is self.state.selected_history_pages - assert num_blocks is self.state.history_blocks - assert out is self.state.fetched_page_table - assert stream == torch.cuda.current_stream().cuda_stream - self.calls.append((buffers[0], stream)) - out.fill_(-1) - _fetch_double_kernel[(selected.shape[0], selected.shape[1])]( - page_table, - selected, - num_blocks, - out, - self.sources[buffers[0].layer_id], - self.pool, - page_table.shape[1], - selected.shape[1], - self.page_tokens, - self.pool.shape[1], - self.offsets[buffers[0].layer_id], - descriptor.fetched_page_scale, - ) - - -def _attention_case(*, tokens_per_block=128, fetched_page_scale=1, fp8=False): - manager = _manager(tokens_per_block) - manager._stream = torch.cuda.current_stream() - manager.dtype = DataType.FP8 if fp8 else DataType.BF16 - manager.layer_offsets = {2: 0, 3: 1, 4: 2} - metadata = _metadata(manager) - state = metadata.sparse_offload_state - state.layers = { - layer: replace(d, fetched_page_scale=fetched_page_scale) - for layer, d in state.layers.items() - } - p, dim = tokens_per_block // 4, 16 - manager.kv_cache_map[20].history_length = 2 * tokens_per_block - 1 - manager.kv_cache_map[20].pages = [0, 3, -1] - manager.kv_cache_map[30].history_length = 2 * tokens_per_block + 3 - manager.kv_cache_map[30].pages = [8, 7, 0] - manager.kv_cache_map[10].pages = [4, -1, -1] - metadata.request_ids = [20, 30, 10, CUDA_GRAPH_DUMMY_REQUEST_ID] - metadata._num_contexts = metadata._num_ctx_tokens = 0 - metadata._num_generations = metadata._num_tokens = 4 - metadata.req_idx_per_token = torch.tensor([1, 0, 2, 3], dtype=torch.int32, device="cuda") - metadata.compress_block_tables = {4: torch.empty((4, 6), dtype=torch.int32, device="cuda")} - metadata.compressed_kv_lens_cuda = { - 4: torch.tensor([2 * p, 2 * p + 1, p, 0], dtype=torch.int32, device="cuda") - } - metadata.max_compressed_indices = {1: 0, 4: 8, 128: 8} - metadata.sparse_mla_topk_lens = { - r: torch.full((4,), 136, dtype=torch.int32, device="cuda") for r in (1, 4, 128) - } - metadata.swa_local_indices_cuda = torch.full((4, 128), -1, dtype=torch.int32, device="cuda") - metadata.swa_local_indices_cuda[:3, :2] = torch.tensor([0, 3], dtype=torch.int32, device="cuda") - metadata.sliding_block_tables = torch.zeros((3, 1, 4, 6), dtype=torch.int32, device="cuda") - # This pool is only read through returned addresses; the double supplies - # device source pages to stand in for host history, not an H2D transport. - pool = ( - torch.arange(64 * 5 * p * dim, device="cuda", dtype=torch.float32).reshape(-1, dim) / 10000 - ) - offsets = {41: p, 19: 3 * p} - sources = { - bid: torch.arange(16 * p * dim, device="cuda", dtype=torch.float32).reshape(-1, dim) / 10000 - + bid - for bid in offsets - } - calls = [] - manager.impl.fetch_sparse_pages = _FetchDouble(state, pool, sources, offsets, p, calls) - stride = dim * (1 if fp8 else 2) - metadata.sparse_mla_base_ptrs = {1: 100000, 4: pool.data_ptr(), 128: 200000} - metadata.swa_buffer_ptrs = { - 2: 100000 + 7 * stride, - 3: 100000 + 9 * stride, - 4: 100000 + 11 * stride, - } - metadata.compressed_buffer_ptrs = { - layer: pool.data_ptr() + offsets[d.buffer_id.layer_id] * stride - for layer, d in state.layers.items() - } - config = DeepSeekV4SparseAttentionConfig(enable_kv_cache_offload=True, index_topk=8) - backends = [] - for layer in state.layers: - backend = object.__new__(DeepseekV4TrtllmAttention) - backend.layer_idx, backend.compress_ratio, backend.head_dim = layer, 4, dim - backend.sparse_attention_config = config - backend.sparse_params = config.to_sparse_params() - backend.use_fp8_ds_mla = backend._uses_nvfp4_compress = False - backend.q_scaling, backend.qk_nope_head_dim, backend.qk_rope_head_dim = 1.0, dim, 0 - backend.kv_scale_quant_orig = torch.tensor([0.25], device="cuda") - backends.append(backend) - topk = torch.tensor( - [ - [2 * p, p + 1, 0, p + 1, 2 * p - 1, p - 1, -1, -1], - [2 * p - 1, p - 1, 0, p, p - 1, -1, -1, -1], - [p - 1, 0, 0, -1, -1, -1, -1, -1], - [0] * 8, - ], - dtype=torch.int32, - device="cuda", - ) - forward_args = AttentionForwardArgs( - attention_input_type=AttentionInputType.generation_only, - sparse_backend_args=SparseBackendForwardArgs(topk_indices=topk), - ) - if fp8: - forward_args.fmha_scheduler_counter = torch.full((1,), 42, dtype=torch.int32, device="cuda") - forward_args.mla_bmm1_scale = torch.empty(2, device="cuda") - forward_args.mla_bmm2_scale = torch.empty(1, device="cuda") - manager.prepare_sparse_offload( - state, metadata.request_ids, metadata.compress_block_tables[4], beam_width=1 - ) - return SimpleNamespace( - manager=manager, - metadata=metadata, - state=state, - backends=backends, - args=forward_args, - pool=pool, - sources=sources, - offsets=offsets, - calls=calls, - p=p, - ) - - -def _consume(case, backend): - backend._prepare_sparse_forward_args(case.metadata, case.args) - return prepare_sparse_runtime_params( - backend, case.pool[:4], None, None, case.metadata, case.args - ) - - -def _check_consumption(case, backend, runtime): - metadata, state = case.metadata, case.state - descriptor = state.layers[backend.layer_idx] - offset = case.offsets[descriptor.buffer_id.layer_id] - raw = state.base_page_tables[descriptor.group_id].cpu() - history = state.history_blocks.cpu() - topk = case.args.sparse_backend_args.topk_indices.cpu() - requests = metadata.req_idx_per_token.cpu().tolist() - active = state.active_request_count.item() - indices = runtime.sparse_attn_indices.cpu() - fetched = state.fetched_page_table.cpu() - assert runtime.sparse_attn_offsets is None - assert runtime.aux_kv_cache_pool_ptr == metadata.sparse_mla_base_ptrs[4] - assert runtime.sparse_attn_kv_lens.data_ptr() == metadata.sparse_mla_topk_lens[4].data_ptr() - assert runtime.sparse_attn_indices_block_size == backend.sparse_params.indices_block_size - assert indices.shape == (4, 136) - expected_indices = torch.full_like(topk, -1) - for query, row in enumerate(requests): - if query >= active: - continue - for j, token in enumerate(topk[query].tolist()): - if token < 0: - continue - ordinal, position = divmod(token, case.p) - if ordinal < history[row]: - physical_page = int(fetched[row, ordinal]) * descriptor.fetched_page_scale - assert physical_page >= 0 # Every referenced history page must be staged. - expected_value = case.sources[descriptor.buffer_id.layer_id][ - int(raw[row, ordinal]) * case.p + position - ] - else: - physical_page = int(raw[row, ordinal]) * descriptor.page_scale - expected_value = case.pool[offset + physical_page * case.p + position] - expected_index = offset + physical_page * case.p + position - expected_indices[query, j] = expected_index - torch.testing.assert_close( - case.pool[int(indices[query, 128 + j])], expected_value, rtol=0, atol=0 - ) - torch.testing.assert_close(indices[:, 128:], expected_indices, rtol=0, atol=0) - # Existing SWA addressing is unaffected by selecting/fetching compressed KV. - swa_offset = {2: 7, 4: 11}[backend.layer_idx] - torch.testing.assert_close( - indices[:active, :2], - torch.tensor([[swa_offset, swa_offset + 3]] * active, dtype=torch.int32).reshape(active, 2), - ) - assert (indices[active:] == -1).all() - assert (indices[:, 2:128] == -1).all() - - -@_requires_cuda -@pytest.mark.parametrize("tokens_per_block", [128, 256]) -@pytest.mark.parametrize("fetched_page_scale", [1, 5]) -def test_per_layer_fetch_reaches_sparse_attention_inputs(tokens_per_block, fetched_page_scale): - case = _attention_case(tokens_per_block=tokens_per_block, fetched_page_scale=fetched_page_scale) - original_topk = case.args.sparse_backend_args.topk_indices.clone() - original_write = case.metadata.compress_block_tables[4].clone() - for backend in case.backends: - runtime = _consume(case, backend) - _check_consumption(case, backend, runtime) - assert [b.layer_id for b, _ in case.calls] == [41, 19] - torch.testing.assert_close(case.args.sparse_backend_args.topk_indices, original_topk) - torch.testing.assert_close(case.metadata.compress_block_tables[4], original_write) - - -@_requires_cuda -def test_fetch_preserves_fp8_scheduler_prologue(): - case = _attention_case(fp8=True) - runtime = _consume(case, case.backends[0]) - _check_consumption(case, case.backends[0], runtime) - assert case.args.fmha_scheduler_counter.item() == 0 - scale = 0.25**2 / math.sqrt(16) - torch.testing.assert_close( - case.args.mla_bmm1_scale, torch.tensor([scale, scale * math.log2(math.e)], device="cuda") - ) - torch.testing.assert_close(case.args.mla_bmm2_scale, torch.tensor([0.25], device="cuda")) - - -@_requires_cuda -@pytest.mark.parametrize("failure", ["units", "scale", "api", "prepared", "prefill", "transfer"]) -def test_sparse_fetch_fails_without_resident_fallback(failure): - case = _attention_case() - if failure in ("units", "scale"): - case.state.layers[2] = replace( - case.state.layers[2], fetched_page_scale=None if failure == "units" else 0 - ) - elif failure == "api": - case.manager.impl.fetch_sparse_pages = None - elif failure == "prepared": - case.state.prepared = False - elif failure == "prefill": - case.state.is_prefill = True - else: - case.manager.impl.fetch_sparse_pages = Mock(side_effect=RuntimeError("fetch failed")) - with pytest.raises((NotImplementedError, RuntimeError, ValueError)): - _consume(case, case.backends[0]) - assert not case.calls - - -@_requires_cuda -def test_decode_fetches_even_for_zero_history_or_all_invalid_selection(): - case = _attention_case() - for cache in case.manager.kv_cache_map.values(): - cache.history_length = 0 - case.args.sparse_backend_args.topk_indices[:3].copy_( - torch.tensor([[0, 1, -1, -1, -1, -1, -1, -1]] * 3, dtype=torch.int32, device="cuda") - ) - case.manager.prepare_sparse_offload( - case.state, case.metadata.request_ids, case.metadata.compress_block_tables[4], beam_width=1 - ) - _consume(case, case.backends[0]) - assert (case.state.selected_history_pages == -1).all() - assert (case.state.fetched_page_table == -1).all() - torch.testing.assert_close( - case.state.compress_read_table, case.metadata.compress_block_tables[4] - ) - case.args.sparse_backend_args.topk_indices.fill_(-1) - _consume(case, case.backends[1]) - assert len(case.calls) == 2 - - -@_requires_cuda -def test_fresh_prefill_consumes_write_table_without_fetch(): - case = _attention_case() - for cache in case.manager.kv_cache_map.values(): - cache.history_length = 0 - case.manager.prepare_sparse_offload( - case.state, case.metadata.request_ids, case.metadata.compress_block_tables[4], beam_width=1 - ) - case.state.is_prefill = True - case.metadata._num_contexts, case.metadata._num_generations = 2, 0 - case.metadata._num_ctx_tokens = 4 - case.args.attention_input_type = AttentionInputType.context_only - case.args.sparse_backend_args.topk_indices.copy_( - torch.tensor([[0, 1, 31, -1, -1, -1, -1, -1]] * 4, dtype=torch.int32, device="cuda") - ) - # Multiple queries/request are legal only for fresh prefill, where no - # selection union is required. Repeated ordinals retain their positions. - case.metadata.req_idx_per_token.copy_( - torch.tensor([0, 0, 1, 1], dtype=torch.int32, device="cuda") - ) - runtime = _consume(case, case.backends[0]) - assert not case.calls - assert ( - runtime.sparse_attn_indices[:, 128:][case.args.sparse_backend_args.topk_indices >= 0] >= 0 - ).all() - - -@_requires_cuda -def test_offload_disabled_and_other_ratios_do_not_fetch(): - case = _attention_case() - backend = case.backends[0] - backend.sparse_attention_config.enable_kv_cache_offload = False - _consume(case, backend) - backend.sparse_attention_config.enable_kv_cache_offload = True - backend.compress_ratio = 1 - runtime = _consume(case, backend) - assert runtime.sparse_attn_indices.shape == (4, 128) - backend.compress_ratio = 128 - case.metadata.compress_block_tables[128] = case.metadata.compress_block_tables[4] - case.metadata.compressed_local_indices_cuda = torch.zeros( - (4, 8), dtype=torch.int32, device="cuda" - ) - _consume(case, backend) - assert not case.calls - - -@_requires_cuda -@pytest.mark.parametrize("debug_assert", [False, True]) -def test_fetch_graph_replay_refreshes_requests_topk_and_padding(monkeypatch, debug_assert): - monkeypatch.setenv("TLLM_DSV4_OFFLOAD_DEBUG_ASSERT", "1" if debug_assert else "0") - case = _attention_case(fetched_page_scale=5) - backend = case.backends[0] - consumer = torch.cuda.Stream() - consumer.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(consumer): - _consume(case, backend) # Compile before capture on the consuming stream. - torch.cuda.current_stream().wait_stream(consumer) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph, stream=consumer): - runtime = _consume(case, backend) - pointers = [ - tensor.data_ptr() - for tensor in ( - case.state.selected_history_pages, - case.state.fetched_page_table, - case.state.compress_read_table, - ) - ] - for ids, requests, topk in [ - ( - [30, 20, 10], - [0, 1, 2, 3], - [ - [0, 32, 64, -1, -1, -1, -1, -1], - [0, 63, -1, -1, -1, -1, -1, -1], - [31, -1, -1, -1, -1, -1, -1, -1], - [0] * 8, - ], - ), - ([10], [0, 1, 2, 3], [[0, 31, -1, -1, -1, -1, -1, -1]] + [[0] * 8] * 3), - ([], [0, 1, 2, 3], [[0] * 8] * 4), - ]: - case.metadata.request_ids = ids + [CUDA_GRAPH_DUMMY_REQUEST_ID] * (4 - len(ids)) - case.manager.prepare_sparse_offload( - case.state, - case.metadata.request_ids, - case.metadata.compress_block_tables[4], - beam_width=1, - ) - lengths = [(case.manager.kv_cache_map[i].history_length + 1) // 4 for i in ids] + [0] * ( - 4 - len(ids) - ) - case.metadata.compressed_kv_lens_cuda[4].copy_( - torch.tensor(lengths, dtype=torch.int32, device="cuda") - ) - case.metadata.req_idx_per_token.copy_( - torch.tensor(requests, dtype=torch.int32, device="cuda") - ) - case.args.sparse_backend_args.topk_indices.copy_( - torch.tensor(topk, dtype=torch.int32, device="cuda") - ) - graph.replay() - _check_consumption(case, backend, runtime) - assert pointers == [ - tensor.data_ptr() - for tensor in ( - case.state.selected_history_pages, - case.state.fetched_page_table, - case.state.compress_read_table, - ) - ] - assert (case.state.fetched_page_table == -1).all() - - -# History policy validation and stream ordering. - - -@pytest.mark.cpu_only -@pytest.mark.parametrize( - "policy", ["enable_block_reuse", "kv_compression_manages_history", "_has_cp_helix", "is_draft"] -) -def test_offload_rejects_other_history_policies_during_initialization(policy: str) -> None: - manager = _manager() - setattr(manager, policy, True) - with pytest.raises(NotImplementedError, match="Sparse offload requires"): - _metadata(manager) - - -@_requires_cuda -@pytest.mark.parametrize("debug_assert", [False, True]) -def test_missing_fetch_result_assertion_is_debug_only(monkeypatch, debug_assert) -> None: - monkeypatch.delenv("TLLM_DSV4_OFFLOAD_DEBUG_ASSERT", raising=False) - if debug_assert: - monkeypatch.setenv("TLLM_DSV4_OFFLOAD_DEBUG_ASSERT", "1") - case = _attention_case() - assert case.state.debug_assert == debug_assert - case.manager.impl.fetch_sparse_pages = lambda **kwargs: kwargs["out"].fill_(-1) - - # Avoid poisoning the suite's CUDA context with a deliberate device assert. - def check(valid: torch.Tensor, message: str) -> None: - if not valid.item(): - raise RuntimeError(message) - - with ( - patch("torch._assert_async", side_effect=check) as device_assert, - patch( - "tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.cache_manager.check_sparse_read_table", - wraps=check_sparse_read_table, - ) as validate, - patch( - "tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.backend.deepseek_v4_local_to_global_indices" - ) as convert, - ): - if debug_assert: - with pytest.raises(RuntimeError, match="missing/out-of-bounds KV page"): - _consume(case, case.backends[0]) - device_assert.assert_called_once() - convert.assert_not_called() - else: - _consume(case, case.backends[0]) - device_assert.assert_not_called() - convert.assert_called_once() - validate.assert_called_once() - assert case.state.read_table_valid.item() == 0 - - -@_requires_cuda -@pytest.mark.parametrize("tokens_per_block", [128, 256]) -def test_history_demotion_waits_for_all_layer_writers_and_readers(tokens_per_block: int) -> None: - cache_stream, model_stream = torch.cuda.Stream(), torch.cuda.Stream() - manager, cache, request, batch = _history_case(tokens_per_block, stream=cache_stream) - state = _metadata(manager).sparse_offload_state - write = torch.empty((4, 6), dtype=torch.int32, device="cuda") - p = tokens_per_block // 4 - # Two CSA layers write different data to the same logical history page. - source = torch.zeros((2, p, 16), dtype=torch.bfloat16, device="cuda") - host = torch.empty(source.shape, dtype=source.dtype, pin_memory=True) - read_before_recycle = torch.empty_like(source) - fetched = torch.empty_like(source) - model_stream.wait_stream(torch.cuda.current_stream()) - resize = cache.resize - - def demote(capacity: int | None, history: int | None) -> bool: - previous = cache.history_length - result = resize(capacity, history) - if previous < tokens_per_block <= cache.history_length: - with torch.cuda.stream(cache_stream): - host.copy_(source, non_blocking=True) - source.fill_(-99) # Simulate reusing the released resident slots. - return result - - cache.resize = demote - with torch.cuda.stream(model_stream): - source[0, :-1].fill_(1) - source[1, :-1].fill_(2) - manager.update_context_resources(batch) - manager.prepare_sparse_offload(state, [20], write, beam_width=1) - model_stream.synchronize() - assert cache.history_length == tokens_per_block - 1 - assert state.history_blocks[0].item() == 0 - assert write[0, 0].item() == 0 - - events = [] - producers = [torch.cuda.Stream(), torch.cuda.Stream()] - for layer, producer in enumerate(producers): - producer.wait_stream(model_stream) - with torch.cuda.stream(producer): - torch.cuda._sleep(2_000_000 * (layer + 1)) - source[layer, -1].fill_(layer + 3) - events.append(producer.record_event()) - batch.context_requests, batch.generation_requests = [], [request] - request.max_beam_num_tokens = tokens_per_block + 1 - with torch.cuda.stream(model_stream): - for event in events: - model_stream.wait_event(event) - read_before_recycle.copy_(source) - manager.update_resources(batch) - manager.prepare_sparse_offload(state, [20], write, beam_width=1) - fetched.copy_(host, non_blocking=True) - model_stream.synchronize() - assert cache.history_length == tokens_per_block - assert state.history_blocks[0].item() == 1 - assert write[0, 0].item() == -1 - assert (source == -99).all() - expected = torch.empty_like(source) - expected[0, :-1], expected[1, :-1] = 1, 2 - expected[0, -1], expected[1, -1] = 3, 4 - torch.testing.assert_close(read_before_recycle, expected, rtol=0, atol=0) - torch.testing.assert_close(fetched, expected, rtol=0, atol=0) diff --git a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_prims_ts_sparse_mla_fmha.py b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_prims_ts_sparse_mla_fmha.py index 5beda56866d2..977906c4b996 100644 --- a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_prims_ts_sparse_mla_fmha.py +++ b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_prims_ts_sparse_mla_fmha.py @@ -14,6 +14,7 @@ PrimsTSSparseMlaFmha, get_prims_ts_toolchain_unsupported_reason, ) +from tensorrt_llm._torch.attention.backends.interface import AttentionInputType from tensorrt_llm._utils import get_sm_version from ..test_sparse_mla_forward import HAS_FLASH_MLA @@ -217,6 +218,40 @@ class _Attn: # weak-referenceable stand-in for TrtllmAttention assert len(planned) == 3 and len(cache) == 3 +@pytest.mark.parametrize( + "compress_ratio, offload, offloaded_reason", + [(4, True, True), (4, False, False), (128, True, False)], +) +def test_offloaded_compressed_rows_stay_on_the_fallback( + compress_ratio, offload, offloaded_reason +) -> None: + """Offloaded ratio-4 rows sit in the sparse cache arena, which the kernel cannot read.""" + + class _Attn: + pass + + attn = _Attn() + attn.compress_ratio = compress_ratio + library = PrimsTSSparseMlaFmha(attn) + metadata = SimpleNamespace( + kv_cache_manager=SimpleNamespace(swa_pool_ptr=0), + sparse_offload_state=object() if offload else None, + # Rejected by the next check, so a layer that passes the offload check stops there. + beam_width=2, + ) + reason = library._unsupported_reason( + SimpleNamespace(device=torch.device("cuda"), ndim=2), + None, + None, + metadata, + SimpleNamespace(attention_input_type=AttentionInputType.generation_only), + ) + if offloaded_reason: + assert "sparse cache arena" in reason + else: + assert reason == "beam search is not supported." + + @pytest.mark.parametrize("offload, heuristic_topk", [(True, False), (False, True)]) def test_plan_keeps_the_hole_checking_path(monkeypatch, offload, heuristic_topk) -> None: """Offload (pages masked to -1 in place) and the heuristic GVR top-k (pads from its diff --git a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py index 5d09cdc3d680..a3febd26607d 100644 --- a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py +++ b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py @@ -584,15 +584,20 @@ def test_shared_topk_lifecycle(monkeypatch): context_topk = torch.tensor([[0, 1, 2], [1, 2, 3]], dtype=torch.int32) generation_topk = torch.tensor([[2, 1, 0], [3, 2, 1]], dtype=torch.int32) + # Offload off: sparse_attn_predict routes decode batches to the arena when it is on, + # which is a different path from the topk sharing under test. + sparse_params = SimpleNamespace(enable_kv_cache_offload=False) full_backend = SimpleNamespace( indexer=SimpleNamespace( forward_from_projected=Mock(side_effect=[context_topk, generation_topk]) ), get_local_layer_idx=Mock(return_value=0), + sparse_params=sparse_params, ) shared_backend = SimpleNamespace( indexer=None, get_local_layer_idx=Mock(return_value=1), + sparse_params=sparse_params, ) # sparse_attn_predict dispatches the local->global remap through the # backend's own _remap_topk_to_global; bind the real method onto the stubs. diff --git a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_offload_cache_manager.py b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_offload_cache_manager.py new file mode 100644 index 000000000000..c6ca0c9626a7 --- /dev/null +++ b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_offload_cache_manager.py @@ -0,0 +1,789 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""DSACacheManagerV2 with sparse KV cache offload, on the native KV cache manager.""" + +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch.attention.backends.sparse.dsa import DSACacheManagerV2 +from tensorrt_llm._torch.attention.backends.sparse.dsa import custom_ops as dsa_custom_ops +from tensorrt_llm._torch.attention.backends.sparse.dsa.backend import DSATrtllmAttention +from tensorrt_llm._torch.attention.backends.sparse.dsa.metadata import DSAtrtllmAttentionMetadata +from tensorrt_llm._torch.attention.backends.sparse.dsa.module import forward_sparse_attn_custom_op +from tensorrt_llm._torch.attention.backends.sparse.offload import ( + SPARSE_OFFLOAD_STATE_COPIES, + SparseOffloadState, + allocate_sparse_offload_state, +) +from tensorrt_llm._torch.attention.backends.sparse.offload_kernels import merge_sparse_arena_indices +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention +from tensorrt_llm._torch.modules.multi_stream_utils import with_multi_stream +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import Role +from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest +from tensorrt_llm._utils import TensorWrapper, convert_to_torch_tensor, get_sm_version +from tensorrt_llm.bindings import DataType, SamplingConfig +from tensorrt_llm.bindings.internal.batch_manager import CacheType as CacheTypeCpp +from tensorrt_llm.llmapi.llm_args import DeepSeekSparseAttentionConfig, KvCacheConfig +from tensorrt_llm.mapping import Mapping + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or get_sm_version() < 100, + reason="DSA KV cache offload needs an SM100+ GPU", +) + +NUM_LAYERS = 7 +FULL_INDEXER_LAYERS = (0, 1, 2, 6) +TOPK = 64 +TOKENS_PER_BLOCK = 64 +MAX_BATCH_SIZE = 4 +HEAD_DIM = 576 +INDEXER_BYTES_PER_TOKEN = 132 +# An NVFP4 latent token: 576 values and a 64-value RoPE residual, two values per +# byte and one E4M3 block scale per 16 values. +NVFP4_RESIDUAL_DIM = 64 +NVFP4_DATA_BYTES = (HEAD_DIM + NVFP4_RESIDUAL_DIM) // 2 +NVFP4_SCALE_BYTES = (HEAD_DIM + NVFP4_RESIDUAL_DIM) // 16 + + +def _pretrained_config(num_layers: int = NUM_LAYERS) -> SimpleNamespace: + # GLM-5.2's indexer pattern: a full layer every fourth from layer 3 on. + return SimpleNamespace( + num_hidden_layers=num_layers, + index_topk_freq=4, + index_skip_topk_offset=3, + index_topk=TOPK, + index_head_dim=128, + index_n_heads=32, + qk_rope_head_dim=64, + ) + + +def _pretrained_config_all_full(num_layers: int = NUM_LAYERS) -> SimpleNamespace: + """Every layer owns an indexer, so no layer reuses another's top-k.""" + config = _pretrained_config(num_layers) + config.index_topk_freq = 1 + config.index_skip_topk_offset = 0 + return config + + +def _manager( + offload: bool = True, + num_layers: int = NUM_LAYERS, + dtype: DataType = DataType.FP8, + pretrained_config: SimpleNamespace | None = None, +) -> DSACacheManagerV2: + sparse_config = DeepSeekSparseAttentionConfig( + index_head_dim=128, + index_n_heads=32, + index_topk=TOPK, + enable_kv_cache_offload=offload, + kv_cache_offload_config=( + {"gpu_cache_factor": 2.0, "host_history_size": 1 << 30} if offload else {} + ), + ) + return DSACacheManagerV2( + kv_cache_config=KvCacheConfig( + enable_block_reuse=False, max_tokens=16384, host_cache_size=0 + ), + kv_cache_type=CacheTypeCpp.SELFKONLY, + num_layers=num_layers, + num_kv_heads=1, + head_dim=HEAD_DIM, + tokens_per_block=TOKENS_PER_BLOCK, + max_seq_len=1024, + max_batch_size=MAX_BATCH_SIZE, + mapping=Mapping(world_size=1, rank=0, tp_size=1, pp_size=1), + dtype=dtype, + sparse_attention_config=sparse_config, + pretrained_config=pretrained_config or _pretrained_config(num_layers), + ) + + +def _request(request_id: int, prompt_len: int) -> LlmRequest: + return LlmRequest( + request_id=request_id, + max_new_tokens=32, + input_tokens=list(range(prompt_len)), + sampling_config=SamplingConfig(), + is_streaming=False, + ) + + +@pytest.fixture +def manager(): + manager = _manager() + yield manager + manager.shutdown() + + +def test_offload_puts_every_latent_in_host_memory_and_keeps_the_indexer(manager): + layers = manager.kv_cache_manager_py_config.layers + assert len(layers) == NUM_LAYERS + for layer in layers: + roles = {str(buffer.role): buffer for buffer in layer.buffers} + assert roles[str(Role.KEY)].is_sparse + assert layer.sparse_cache.num_selected_blocks == TOPK + assert layer.sparse_cache.unit_tokens == 1 + has_indexer = str(Role.INDEX_KEY) in roles + assert has_indexer == (int(layer.layer_id) in FULL_INDEXER_LAYERS) + if has_indexer: + assert not roles[str(Role.INDEX_KEY)].is_sparse + for local_layer in range(NUM_LAYERS): + expected = INDEXER_BYTES_PER_TOKEN if local_layer in FULL_INDEXER_LAYERS else 0 + assert manager.get_layer_bytes_per_token(local_layer, Role.ALL) == expected + assert manager.get_layer_bytes_per_token(local_layer, Role.KEY) == HEAD_DIM + descriptors = manager.get_sparse_offload_descriptors() + assert sorted(descriptors) == list(range(NUM_LAYERS)) + for layer, descriptor in descriptors.items(): + assert descriptor.unit_tokens == 1 + assert descriptor.scale_buffer_index is None + assert len(descriptor.buffer_ids) == 1 + pool_pointers, pool_mapping = manager.get_sparse_latent_write_pool(layer) + assert pool_pointers.tolist() == [[descriptor.arena_base_addresses[0], 0]] + assert pool_mapping.tolist() == [[0, 0]] + assert manager._sparse_offload_nvfp4_head_dim() is None + + +def test_nvfp4_offload_puts_block_scales_in_host_memory_beside_the_latent(): + manager = _manager(dtype=DataType.NVFP4) + try: + for layer in manager.kv_cache_manager_py_config.layers: + roles = {str(buffer.role): buffer for buffer in layer.buffers} + assert roles[str(Role.KEY)].is_sparse + assert roles[str(Role.KEY_BLOCK_SCALE)].is_sparse + for local_layer in range(NUM_LAYERS): + expected = INDEXER_BYTES_PER_TOKEN if local_layer in FULL_INDEXER_LAYERS else 0 + assert manager.get_layer_bytes_per_token(local_layer, Role.ALL) == expected + assert manager.get_layer_bytes_per_token(local_layer, Role.KEY) == NVFP4_DATA_BYTES + assert ( + manager.get_layer_bytes_per_token(local_layer, Role.KEY_BLOCK_SCALE) + == NVFP4_SCALE_BYTES + ) + for layer, descriptor in manager.get_sparse_offload_descriptors().items(): + data, scales = descriptor.data_buffer_index, descriptor.scale_buffer_index + assert len(descriptor.buffer_ids) == 2 and {data, scales} == {0, 1} + # The writer and the gather address both arenas with the data's rows. + assert descriptor.arena_row_counts[data] == descriptor.arena_row_counts[scales] + data_base = descriptor.arena_base_addresses[data] + scale_base = descriptor.arena_base_addresses[scales] + pool_pointers, pool_mapping = manager.get_sparse_latent_write_pool(layer) + assert pool_pointers.tolist() == [[[data_base, scale_base], [0, 0]]] + assert pool_mapping.tolist() == [[0, 0]] + data_rows, scale_rows = manager.get_sparse_latent_nvfp4_arenas(layer) + assert (data_rows.data_ptr(), scale_rows.data_ptr()) == (data_base, scale_base) + assert (data_rows.dtype, scale_rows.dtype) == (torch.uint8, torch.float8_e4m3fn) + rows = descriptor.arena_row_counts[data] + assert data_rows.shape == (rows, 1, NVFP4_DATA_BYTES) + assert scale_rows.shape == (rows, 1, NVFP4_SCALE_BYTES) + # Attention reads the dequantized FP8 rows of each selection from offload scratch. + assert manager._sparse_offload_nvfp4_head_dim() == HEAD_DIM + state_bytes = super(DSACacheManagerV2, manager).get_sparse_offload_scratch_bytes( + manager.kv_cache_manager_py_config + ) + assert state_bytes == SPARSE_OFFLOAD_STATE_COPIES * SparseOffloadState.device_bytes( + max_batch_size=MAX_BATCH_SIZE, topk=TOPK, buffer_counts={2}, nvfp4_head_dim=HEAD_DIM + ) + finally: + manager.shutdown() + + +def test_shared_indexer_layers_name_one_selection(manager): + """A run of layers reusing one indexer resolves once and shares one resident set. + + The groups here are GLM-5.2's: full indexers at 0, 1, 2 and 6, so layers 3-5 + reuse layer 2's top-k and nothing else is shared. A layer that selects for + itself names no selection, which is what keeps the fused fetch path for every + model without indexer sharing. + """ + expected_owner = {0: None, 1: None, 2: 2, 3: 2, 4: 2, 5: 2, 6: None} + assert manager._selection_leaders() == expected_owner + + layers = manager.kv_cache_manager_py_config.layers + ids = {int(layer.layer_id): layer.sparse_cache.selection_id for layer in layers} + # Members agree; the layers selecting for themselves name nothing. The value is + # negative so it cannot be read as a layer id, which is a space an unset id + # already occupies. + assert ids == {0: None, 1: None, 2: -3, 3: -3, 4: -3, 5: -3, 6: None} + + descriptors = manager.get_sparse_offload_descriptors() + assert {layer: d.selection_owner for layer, d in descriptors.items()} == expected_owner + + +def test_a_model_without_indexer_sharing_shares_no_selection(): + """Every layer owning an indexer leaves every selection a layer's own.""" + manager = _manager(pretrained_config=_pretrained_config_all_full()) + try: + assert set(manager._selection_leaders().values()) == {None} + for layer in manager.kv_cache_manager_py_config.layers: + assert layer.sparse_cache.selection_id is None + for descriptor in manager.get_sparse_offload_descriptors().values(): + assert descriptor.selection_owner is None + finally: + manager.shutdown() + + +def test_offload_refuses_the_unified_pool_view(manager): + with pytest.raises(RuntimeError, match="host memory"): + manager.get_unique_primary_pool() + + +def test_decode_step_writes_fetches_and_advances_history(manager): + prompt_len = 3 * TOKENS_PER_BLOCK + 5 + request = _request(7, prompt_len) + assert manager.prepare_disagg_gen_init(request) + assert manager.try_allocate_generation(request) + cache = manager.kv_cache_map[request.py_request_id] + assert (cache.history_length, cache.capacity) == (prompt_len, prompt_len + 1) + + descriptors = manager.get_sparse_offload_descriptors() + state = allocate_sparse_offload_state( + lambda name, shape, dtype: torch.empty(shape, dtype=dtype, device="cuda"), + descriptors, + max_batch_size=MAX_BATCH_SIZE, + topk=TOPK, + ) + manager.prepare_sparse_batch([request.py_request_id]) + manager.prepare_sparse_offload(state, [request.py_request_id], beam_width=1) + assert state.history_lengths[0].item() == prompt_len + + # The last history token and the token written this step, then older history. + topk = torch.full((1, TOPK), -1, dtype=torch.int32, device="cuda") + topk[0, : TOPK - 2] = torch.arange(TOPK - 2, dtype=torch.int32) + topk[0, TOPK - 2] = prompt_len - 1 + topk[0, TOPK - 1] = prompt_len + request_indices = torch.zeros(1, dtype=torch.int32, device="cuda") + kv_lengths = torch.tensor([prompt_len + 1], dtype=torch.int32, device="cuda") + out = torch.empty((1, TOPK), dtype=torch.int32, device="cuda") + assert state.batch_rows.tolist() == [0] + [-1] * (MAX_BATCH_SIZE - 1) + for layer in range(NUM_LAYERS): + writes = manager.prepare_sparse_write_indices(state, layer) + assert writes[0, 0].item() >= 0 + assert (writes[1:, 0] == -1).all() + fetched = manager.fetch_sparse_indices(state, layer, topk, request_indices, kv_lengths) + merge_sparse_arena_indices( + topk, + request_indices, + state.batch_rows, + state.active_request_count, + kv_lengths, + state.history_lengths, + fetched[:, 0, :], + writes[:, 0], + out, + unit_tokens=1, + ) + merged = out[0].tolist() + assert all(index >= 0 for index in merged) + # Every selected token has its own arena row, the written one included. + assert len(set(merged)) == TOPK + assert merged[-1] == writes[0, 0].item() + assert fetched[0, 0, -1].item() == -1 + rows_in_arena = descriptors[layer].arena_row_counts[0] + assert max(merged) < rows_in_arena + + scheduled = SimpleNamespace(context_requests=[], generation_requests=[request]) + manager.update_forward_resources(scheduled, wait_for_input_copy=lambda: None) + torch.cuda.synchronize() + assert cache.history_length == prompt_len + 1 + manager.free_resources(request) + + +def _arena_rows(descriptor, index: int, row_bytes: int) -> torch.Tensor: + """One buffer's arena as int32 words per row.""" + return convert_to_torch_tensor( + TensorWrapper( + descriptor.arena_base_addresses[index], + torch.int32, + (descriptor.arena_row_counts[index], row_bytes // 4), + ) + ) + + +@pytest.mark.parametrize( + ("num_layers", "num_requests", "dtype"), + [ + (1, 1, DataType.FP8), + (2, 1, DataType.FP8), + (NUM_LAYERS, 1, DataType.FP8), + (NUM_LAYERS, 3, DataType.FP8), + (NUM_LAYERS, 1, DataType.NVFP4), + (NUM_LAYERS, 3, DataType.NVFP4), + ], + ids=["1layer", "2layers", "7layers", "7layers-3requests", "nvfp4", "nvfp4-3requests"], +) +@pytest.mark.parametrize("seed", [0, 1]) +def test_fetched_rows_hold_written_history_across_evictions(seed, num_layers, num_requests, dtype): + """Every step writes each request's token row; later fetches must return it. + + The GPU cache holds a few selections, so a long decode evicts rows and fetches + them back from the host pages the write-back filled. Requests join at different + steps, so the Batch's rows carry different histories. Input runs newest first, + as on a disaggregated decode worker, while each request keeps the Batch row it + joined in, so input positions and Batch rows disagree. An NVFP4 latent's block + scales travel in their own buffer and must come back in the data's rows. + """ + generator = torch.Generator().manual_seed(seed) + manager = _manager(num_layers=num_layers, dtype=dtype) + requests = [_request(9 + index, 1) for index in range(num_requests)] + join_step = {request.py_request_id: 40 * index for index, request in enumerate(requests)} + descriptors = manager.get_sparse_offload_descriptors() + nvfp4 = dtype == DataType.NVFP4 + arenas = { + layer: [_arena_rows(descriptor, descriptor.data_buffer_index, HEAD_DIM)] + for layer, descriptor in descriptors.items() + } + if nvfp4: + arenas = { + layer: [ + _arena_rows(descriptor, descriptor.data_buffer_index, NVFP4_DATA_BYTES), + _arena_rows(descriptor, descriptor.scale_buffer_index, NVFP4_SCALE_BYTES), + ] + for layer, descriptor in descriptors.items() + } + state = allocate_sparse_offload_state( + lambda name, shape, dtype: torch.empty(shape, dtype=dtype, device="cuda"), + descriptors, + max_batch_size=MAX_BATCH_SIZE, + topk=TOPK, + ) + active: list[LlmRequest] = [] + permuted = False + for step in range(6 * TOKENS_PER_BLOCK): + for request in requests: + if join_step[request.py_request_id] == step: + assert manager.prepare_disagg_gen_init(request) + active.append(request) + for request in active: + assert manager.try_allocate_generation(request) + scheduled = SimpleNamespace(context_requests=[], generation_requests=active[::-1]) + manager.prepare_resources(scheduled) + request_ids = [request.py_request_id for request in scheduled.generation_requests] + manager.prepare_sparse_offload(state, request_ids, beam_width=1) + # The rows the manager recorded at add_kv_cache are still the native ones. + for member in manager._sparse_batch_members.values(): + assert manager._sparse_batch.row_of(member.cache) == member.row + batch_rows = state.batch_rows.tolist()[: len(active)] + permuted |= batch_rows != list(range(len(active))) + histories = [manager.kv_cache_map[rid].history_length for rid in request_ids] + topk = torch.full((len(active), TOPK), -1, dtype=torch.int32) + for position, history in enumerate(histories): + # Unit 0 came with the prompt and was never written; select from the rest, + # always including the token written this step. + candidates = torch.randperm(history, generator=generator)[: TOPK - 1] + candidates = candidates[candidates > 0][: TOPK - 1].to(torch.int32) + topk[position, : candidates.numel()] = candidates + topk[position, -1] = history + topk = topk.cuda() + request_indices = torch.arange(len(active), dtype=torch.int32, device="cuda") + kv_lengths = torch.tensor([h + 1 for h in histories], dtype=torch.int32, device="cuda") + out = torch.empty((len(active), TOPK), dtype=torch.int32, device="cuda") + for layer in range(num_layers): + descriptor = descriptors[layer] + data = descriptor.data_buffer_index + writes = manager.prepare_sparse_write_indices(state, layer) + # At the buffer's slice, not its position: buffers whose arena indices + # coincide share one, and the slices no buffer maps to are never written. + data = state.out_slots[layer][data] + for row, rid, history in zip(batch_rows, request_ids, histories): + write_row = writes[row, data].item() + assert write_row >= 0, (rid, history, layer) + # The token's row: its position, its layer, its request. + for arena in arenas[layer]: + arena[write_row, :3] = torch.tensor( + [history, layer, rid], dtype=torch.int32, device="cuda" + ) + free_rows = [row for row in range(MAX_BATCH_SIZE) if row not in batch_rows] + assert (writes[free_rows, data] == -1).all() + fetched = manager.fetch_sparse_indices(state, layer, topk, request_indices, kv_lengths) + if nvfp4: + # The scale mirrors the data pool, so it resolves to the same index and + # reads the data's slice rather than one of its own. One slice written + # once is what replaces the two equal rows this used to compare. + assert state.out_slots[layer][descriptor.scale_buffer_index] == data + merge_sparse_arena_indices( + topk, + request_indices, + state.batch_rows, + state.active_request_count, + kv_lengths, + state.history_lengths, + fetched[:, data, :], + writes[:, data], + out, + unit_tokens=1, + ) + for position, rid in enumerate(request_ids): + valid = topk[position] >= 0 + units = topk[position][valid] + expected = torch.stack( + (units, torch.full_like(units, layer), torch.full_like(units, rid)), dim=1 + ) + for arena in arenas[layer]: + got = arena[out[position][valid].long(), :3] + wrong = (got != expected).any(dim=1) + assert not wrong.any(), ( + f"request {rid}, history {histories[position]}, layer {layer}: units " + f"{units[wrong].tolist()} read {got[wrong].tolist()}" + ) + manager.update_forward_resources(scheduled, wait_for_input_copy=lambda: None) + for rid, history in zip(request_ids, histories): + assert manager.kv_cache_map[rid].history_length == history + 1 + assert permuted == (num_requests > 1) + for request in active: + manager.free_resources(request) + manager.shutdown() + + +def test_nvfp4_attention_reads_the_selected_rows_dequantized(): + """Attention reads FP8 scratch rows dequantized from the arena rows the merge picked. + + Arena rows hold random NVFP4 bytes and scales. The reference dequantizes the + merged rows after copying them out of the arena, so it checks which rows the + gather read, from which views, and that each index became its scratch row. + """ + manager = _manager(dtype=DataType.NVFP4) + generator = torch.Generator(device="cuda").manual_seed(0) + layer = 2 + try: + prompt_len = 3 * TOKENS_PER_BLOCK + 5 + request = _request(7, prompt_len) + assert manager.prepare_disagg_gen_init(request) + assert manager.try_allocate_generation(request) + state = allocate_sparse_offload_state( + lambda name, shape, dtype: torch.empty(shape, dtype=dtype, device="cuda"), + manager.get_sparse_offload_descriptors(), + max_batch_size=MAX_BATCH_SIZE, + topk=TOPK, + nvfp4_head_dim=manager._sparse_offload_nvfp4_head_dim(), + ) + manager.prepare_sparse_batch([request.py_request_id]) + manager.prepare_sparse_offload(state, [request.py_request_id], beam_width=1) + data_rows, scale_rows = manager.get_sparse_latent_nvfp4_arenas(layer) + data_rows.copy_( + torch.randint( + 0, 256, data_rows.shape, dtype=torch.uint8, device="cuda", generator=generator + ) + ) + scale_rows.copy_( + torch.empty(scale_rows.shape, device="cuda") + .uniform_(0.25, 4.0, generator=generator) + .to(torch.float8_e4m3fn) + ) + writes = manager.prepare_sparse_write_indices(state, layer) + + # Older history, the last history token and the token written this step, + # with unselected slots between them. + topk = torch.full((1, TOPK), -1, dtype=torch.int32, device="cuda") + topk[0, : TOPK - 8] = torch.arange(TOPK - 8, dtype=torch.int32) * 3 + topk[0, TOPK - 2] = prompt_len - 1 + topk[0, TOPK - 1] = prompt_len + metadata = SimpleNamespace( + kv_cache_manager=manager, + sparse_offload_state=state, + num_contexts=0, + num_seqs=1, + sparse_mla_topk=TOPK, + req_idx_per_token=torch.zeros(1, dtype=torch.int32, device="cuda"), + kv_lens_cuda=torch.tensor([prompt_len + 1], dtype=torch.int32, device="cuda"), + sparse_offload_attn_indices=torch.empty( + (MAX_BATCH_SIZE, TOPK), dtype=torch.int32, device="cuda" + ), + ) + backend = SimpleNamespace( + layer_idx=layer, + _nvfp4_mla_kv_scale_quant_orig=torch.tensor([0.5], device="cuda"), + ) + args = SimpleNamespace(sparse_runtime_params=SimpleNamespace(aux_kv_cache_pool_ptr=0)) + indices = DSATrtllmAttention._sparse_offload_arena_indices(backend, topk, metadata, args) + + descriptor = state.layers[layer] + data = descriptor.data_buffer_index + rows = torch.empty_like(topk) + merge_sparse_arena_indices( + topk, + metadata.req_idx_per_token, + state.batch_rows, + state.active_request_count, + metadata.kv_lens_cuda, + state.history_lengths, + state.fetched_indices[2][:, data, :], + writes[:, data], + rows, + unit_tokens=1, + ) + valid = rows[0] >= 0 + assert valid.sum().item() == TOPK - 6 + assert rows[0, -1].item() == writes[0, data].item() + flat = rows.view(-1).clamp(min=0).long() + reference = torch.empty((1, TOPK, HEAD_DIM), dtype=torch.float8_e4m3fn, device="cuda") + compact = torch.where(valid, torch.arange(TOPK, device="cuda"), -1).to(torch.int32)[None] + torch.ops.trtllm.nvfp4_mla_kv_cache_gather_direct( + data_rows[flat].contiguous(), + scale_rows.view(torch.uint8)[flat].contiguous().view(torch.float8_e4m3fn), + compact, + reference, + backend._nvfp4_mla_kv_scale_quant_orig, + NVFP4_RESIDUAL_DIM, + TOPK, + ) + assert args.sparse_runtime_params.aux_kv_cache_pool_ptr == state.nvfp4_scratch.data_ptr() + torch.testing.assert_close(indices, compact) + scratch = state.nvfp4_scratch[0].view(torch.uint8) + assert torch.equal(scratch[valid], reference[0].view(torch.uint8)[valid]) + manager.free_resources(request) + finally: + manager.shutdown() + + +def test_offload_scratch_counts_the_latent_writer(manager): + with_writer = manager.get_sparse_offload_scratch_bytes(manager.kv_cache_manager_py_config) + state_only = super(DSACacheManagerV2, manager).get_sparse_offload_scratch_bytes( + manager.kv_cache_manager_py_config + ) + assert with_writer - state_only == 2 * 3 * MAX_BATCH_SIZE * 4 + + +@pytest.mark.parametrize("dtype", [DataType.FP8, DataType.NVFP4]) +@pytest.mark.parametrize("capture", [False, True], ids=["eager", "cuda-graph"]) +@pytest.mark.parametrize("custom_op", [False, True], ids=["direct", "registered"]) +def test_overlapped_write_mapping_tracks_reordered_requests(monkeypatch, dtype, capture, custom_op): + """The writer joins reservation and reads current Batch rows on every replay.""" + manager = _manager(num_layers=2, dtype=dtype) + requests = [_request(17, 1), _request(18, 1)] + descriptors = manager.get_sparse_offload_descriptors() + state = allocate_sparse_offload_state( + lambda name, shape, dtype: torch.empty(shape, dtype=dtype, device="cuda"), + descriptors, + max_batch_size=MAX_BATCH_SIZE, + topk=TOPK, + ) + metadata = SimpleNamespace( + kv_cache_manager=manager, + sparse_offload_state=state, + sparse_offload_stream=torch.cuda.Stream(), + sparse_offload_start_event=torch.cuda.Event(), + sparse_offload_write_event=torch.cuda.Event(), + sparse_offload_write_pending=None, + sparse_latent_block_offsets=torch.full( + (1, MAX_BATCH_SIZE, 2, 1), -1, dtype=torch.int32, device="cuda" + ), + sparse_latent_write_slots=torch.full( + (MAX_BATCH_SIZE,), -1, dtype=torch.int32, device="cuda" + ), + ) + observed = torch.empty((2, MAX_BATCH_SIZE), dtype=torch.int32, device="cuda") + observed_slots = torch.empty_like(observed) + latent = torch.empty((len(requests), HEAD_DIM), dtype=torch.bfloat16, device="cuda") + backends = [] + for layer in range(2): + backend = object.__new__(DSATrtllmAttention) + backend.layer_idx = layer + backend.sparse_params = SimpleNamespace(enable_kv_cache_offload=True) + backends.append(backend) + + def consume_write_mapping(self, fused_q, q_pe, latent_cache, metadata, **kwargs): + target = kwargs["latent_write_target"] + observed[self.layer_idx].copy_(target.block_offsets[0, :, 0, 0]) + observed_slots[self.layer_idx].copy_(metadata.sparse_latent_write_slots) + + monkeypatch.setattr(TrtllmAttention, "mla_rope_generation", consume_write_mapping) + + def project(hidden_states, position_ids, layer_idx): + assert metadata.sparse_offload_write_pending == int(layer_idx) + return [latent] * 4 + + def attend(q, compressed_kv, k_pe, latent_cache, intermediates, positions, layer_idx, output): + backend = backends[int(layer_idx)] + backend.prepare_sparse_write(metadata) + backend.mla_rope_generation(None, None, latent, metadata) + + monkeypatch.setattr(dsa_custom_ops, "_extract_extra_attrs", lambda layer: (metadata, None)) + monkeypatch.setattr(torch.ops.trtllm, "mla_dsa_proj", project) + monkeypatch.setattr(dsa_custom_ops, "maybe_bcg_mla_dsa_attn_inplace", attend) + + def forward(): + for backend in backends: + if custom_op: + mla = SimpleNamespace(mqa=backend, layer_idx_str=str(backend.layer_idx)) + forward_sparse_attn_custom_op(mla, latent, None, [latent], None) + continue + backend.prepare_sparse_write(metadata) + # Entry through the attention custom op must reuse the pending mapping. + backend.prepare_sparse_write(metadata) + backend.mla_rope_generation(None, None, latent, metadata) + + graph = None + try: + for request in requests: + assert manager.prepare_disagg_gen_init(request) + for step in range(6): + for request in requests: + assert manager.try_allocate_generation(request) + ordered = requests if step % 2 == 0 else requests[::-1] + scheduled = SimpleNamespace(context_requests=[], generation_requests=ordered) + manager.prepare_resources(scheduled) + manager.prepare_sparse_offload( + state, [request.py_request_id for request in ordered], beam_width=1 + ) + with with_multi_stream(True): + if capture and graph is None: + forward() + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + forward() + if graph is None: + forward() + else: + graph.replay() + assert metadata.sparse_offload_write_pending is None + # The serial native API is idempotent for the same in-flight token. + # It independently checks the destinations consumed by the writer. + for layer in range(2): + writes = manager.prepare_sparse_write_indices(state, layer) + data = descriptors[layer].data_buffer_index + expected = torch.full_like(observed[layer], -1) + expected[: len(requests)] = writes[state.batch_rows[: len(requests)], data] + torch.testing.assert_close(observed[layer], expected) + torch.testing.assert_close( + observed_slots[layer], (expected >= 0).to(torch.int32) - 1 + ) + manager.update_forward_resources(scheduled, wait_for_input_copy=lambda: None) + for request in requests: + manager.free_resources(request) + finally: + torch.cuda.synchronize() + manager.shutdown() + + +@pytest.mark.parametrize("dtype", [DataType.FP8, DataType.NVFP4]) +@pytest.mark.parametrize("capture", [False, True], ids=["eager", "cuda-graph"]) +@pytest.mark.parametrize("explicit_pattern", [False, True], ids=["glm52", "glm53"]) +def test_shared_layer_prefetch_preserves_history(monkeypatch, dtype, capture, explicit_pattern): + """Lookahead fetch preserves rows through eviction, request reordering and replay.""" + config = _pretrained_config() + if explicit_pattern: + config.index_topk_pattern = ["F", "S", "S", "F", "S", "F", "S"] + manager = _manager(dtype=dtype, pretrained_config=config) + requests = [_request(17, 1), _request(18, 1)] + metadata = object.__new__(DSAtrtllmAttentionMetadata) + metadata.kv_cache_manager = manager + metadata.sparse_metadata_params = SimpleNamespace(enable_kv_cache_offload=True) + metadata.beam_width = 1 + metadata.max_draft_tokens = 0 + metadata.draft_kv_cache_manager = None + metadata.sparse_mla_topk = TOPK + metadata.is_cuda_graph = False + metadata.cuda_graph_buffers = {} + metadata.get_empty = lambda buffers, shape, dtype, **kwargs: torch.empty( + shape, dtype=dtype, device="cuda" + ) + metadata._init_dsa_sparse_offload_state() + expected_pairs = {0: 1, 1: 2, 3: 4, 5: 6} if explicit_pattern else {2: 3, 3: 4, 4: 5} + assert metadata.sparse_offload_next_shared_layer == expected_pairs + state = metadata.sparse_offload_state + # Inspect raw data and scale bytes separately; dequantization has its own test. + state.nvfp4_scratch = None + num_requests = len(requests) + metadata._num_tokens = num_requests + metadata._seq_lens = torch.ones(num_requests, dtype=torch.int32) + metadata._num_contexts = 0 + metadata.shared_topk_indices = torch.full( + (num_requests, TOPK), -1, dtype=torch.int32, device="cuda" + ) + metadata.req_idx_per_token = torch.arange(num_requests, dtype=torch.int32, device="cuda") + metadata.kv_lens_cuda = torch.empty(num_requests, dtype=torch.int32, device="cuda") + latent = torch.empty((num_requests, HEAD_DIM), dtype=torch.bfloat16, device="cuda") + token_values = torch.empty((num_requests, 3), dtype=torch.int32, device="cuda") + arenas = {} + backends = [] + for layer, descriptor in state.layers.items(): + if dtype == DataType.NVFP4: + arenas[layer] = [ + _arena_rows(descriptor, descriptor.data_buffer_index, NVFP4_DATA_BYTES), + _arena_rows(descriptor, descriptor.scale_buffer_index, NVFP4_SCALE_BYTES), + ] + else: + arenas[layer] = [_arena_rows(descriptor, descriptor.data_buffer_index, HEAD_DIM)] + backend = object.__new__(DSATrtllmAttention) + backend.layer_idx = layer + backend.sparse_params = SimpleNamespace(enable_kv_cache_offload=True) + backends.append(backend) + observed = torch.empty( + (NUM_LAYERS, len(arenas[0]), num_requests, TOPK, 3), dtype=torch.int32, device="cuda" + ) + args = SimpleNamespace(sparse_runtime_params=SimpleNamespace(aux_kv_cache_pool_ptr=0)) + + def write_token(self, fused_q, q_pe, latent_cache, metadata, **kwargs): + rows = kwargs["latent_write_target"].block_offsets[0, :num_requests, 0, 0].long() + token_values[:, 1] = self.layer_idx + for arena in arenas[self.layer_idx]: + arena[rows, :3] = token_values + + monkeypatch.setattr(TrtllmAttention, "mla_rope_generation", write_token) + + def forward(): + for backend in backends: + backend.prepare_sparse_write(metadata) + backend.mla_rope_generation(None, None, latent, metadata) + indices = backend._sparse_offload_arena_indices( + metadata.shared_topk_indices, metadata, args + ) + for buffer, arena in enumerate(arenas[backend.layer_idx]): + observed[backend.layer_idx, buffer].copy_(arena[indices.clamp_min(0).long(), :3]) + backend.prefetch_next_sparse_history(metadata) + + graph = None + generator = torch.Generator().manual_seed(17) + try: + for request in requests: + assert manager.prepare_disagg_gen_init(request) + for step in range(3 * TOKENS_PER_BLOCK): + for request in requests: + assert manager.try_allocate_generation(request) + ordered = requests if step % 2 == 0 else requests[::-1] + scheduled = SimpleNamespace(context_requests=[], generation_requests=ordered) + manager.prepare_resources(scheduled) + ids = [request.py_request_id for request in ordered] + manager.prepare_sparse_offload(state, ids, beam_width=1) + history = step + 1 + selected = torch.full((num_requests, TOPK), -1, dtype=torch.int32) + for row in range(num_requests): + candidates = torch.randperm(history - 1, generator=generator)[: TOPK - 1] + 1 + selected[row, : candidates.numel()] = candidates + selected[row, -1] = history + metadata.shared_topk_indices.copy_(selected) + metadata.kv_lens_cuda.fill_(history + 1) + token_values.copy_(torch.tensor([[history, 0, rid] for rid in ids], dtype=torch.int32)) + with with_multi_stream(True): + if capture and graph is None: + forward() + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + forward() + if graph is None: + forward() + else: + graph.replay() + assert metadata.sparse_offload_write_pending is None + assert metadata.sparse_offload_fetch_pending is None + actual = observed.cpu() + for layer in range(NUM_LAYERS): + for row, rid in enumerate(ids): + units = selected[row][selected[row] >= 0] + expected = torch.stack( + (units, torch.full_like(units, layer), torch.full_like(units, rid)), dim=1 + ) + for buffer in range(len(arenas[layer])): + torch.testing.assert_close( + actual[layer, buffer, row][selected[row] >= 0], expected + ) + manager.update_forward_resources(scheduled, wait_for_input_copy=lambda: None) + for request in requests: + manager.free_resources(request) + finally: + torch.cuda.synchronize() + manager.shutdown() diff --git a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_offload_config.py b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_offload_config.py new file mode 100644 index 000000000000..6e13ee322fc5 --- /dev/null +++ b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_offload_config.py @@ -0,0 +1,134 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""DSA sparse KV cache offload configuration and its lowering.""" + +import pytest +from transformers import PretrainedConfig + +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm.llmapi.llm_args import ( + DeepSeekSparseAttentionConfig, + DeepSeekV4SparseAttentionConfig, + SparseKvCacheOffloadConfig, +) + +_OFFLOAD = {"gpu_cache_factor": 3.0, "host_history_size": 2 * 1024**3} + + +def _glm_config() -> PretrainedConfig: + return PretrainedConfig( + architectures=["GlmMoeDsaForCausalLM"], + num_hidden_layers=7, + index_n_heads=32, + index_head_dim=128, + index_topk=2048, + index_topk_freq=4, + index_skip_topk_offset=3, + indexer_rope_interleave=True, + ) + + +@pytest.mark.parametrize("enabled", [False, True]) +def test_dsa_offload_lowers_into_sparse_params(enabled): + config = DeepSeekSparseAttentionConfig( + index_topk=2048, + enable_kv_cache_offload=enabled, + kv_cache_offload_config=_OFFLOAD if enabled else {}, + ) + expected = SparseKvCacheOffloadConfig(**_OFFLOAD) if enabled else None + assert config.get_kv_cache_offload_config() == expected + assert config.to_sparse_params().enable_kv_cache_offload is enabled + assert config.to_sparse_metadata_params().enable_kv_cache_offload is enabled + + +def test_dsa_offload_sizing_requires_offload(): + with pytest.raises(ValueError, match="requires enable_kv_cache_offload"): + DeepSeekSparseAttentionConfig(kv_cache_offload_config=_OFFLOAD) + + +def test_dsa_offload_defers_the_top_k_check_to_the_checkpoint(): + # index_topk comes from the checkpoint when the user leaves it unset. + assert DeepSeekSparseAttentionConfig(enable_kv_cache_offload=True).index_topk is None + with pytest.raises(ValueError, match="positive index_topk"): + DeepSeekSparseAttentionConfig(enable_kv_cache_offload=True, index_topk=0) + + +def test_deepseek_v4_offload_keeps_its_ratio_check(): + with pytest.raises(ValueError, match="ratio-4"): + DeepSeekV4SparseAttentionConfig(compress_ratios=[1, 128], enable_kv_cache_offload=True) + with pytest.raises(ValueError, match="requires enable_kv_cache_offload"): + DeepSeekV4SparseAttentionConfig(kv_cache_offload_config=_OFFLOAD) + + +@pytest.mark.parametrize("enabled", [False, True]) +def test_glm_model_config_preserves_sparse_offload(tmp_path, monkeypatch, enabled): + monkeypatch.setattr( + "tensorrt_llm._torch.model_config.load_pretrained_config", + lambda *args, **kwargs: _glm_config(), + ) + sparse_config = DeepSeekSparseAttentionConfig( + enable_kv_cache_offload=enabled, + kv_cache_offload_config=_OFFLOAD if enabled else {}, + ) + + resolved = ModelConfig.from_pretrained( + str(tmp_path), + sparse_attention_config=sparse_config, + attn_backend="TRTLLM", + moe_backend="TRTLLM", + ).sparse_attention_config + + assert resolved.enable_kv_cache_offload is enabled + assert resolved.kv_cache_offload_config == sparse_config.kv_cache_offload_config + # The checkpoint fills the selection width the offload cache is sized by. + assert resolved.index_topk == 2048 + + +@pytest.mark.parametrize( + ("indexer_mask", "pp_layers", "expected"), + [ + # GLM-5.2's pattern over seven layers: full indexers at 0, 1, 2 and 6, so + # layers 3-5 reuse layer 2's top-k and nothing else is shared. + ( + [True, True, True, False, False, False, True], + None, + {0: None, 1: None, 2: 2, 3: 2, 4: 2, 5: 2, 6: None}, + ), + # Every layer owning an indexer shares nothing, which is what keeps the + # fused fetch path for models without indexer sharing. + ([True] * 4, None, {0: None, 1: None, 2: None, 3: None}), + # A pipeline rank holding global layers 8-13, full at 8 and 11. Leaders are + # model layers, not local offsets. + ( + [False] * 8 + [True, False, False, True, False, False], + list(range(8, 14)), + {8: 8, 9: 8, 10: 8, 11: 11, 12: 11, 13: 11}, + ), + # A rank whose first layer owns no indexer: its top-k comes from a layer on + # another rank, so the run has to start here anyway. The remap path's own + # scan starts a group at local zero for the same reason. + ( + [False] * 10 + [True, False], + list(range(8, 12)), + {8: 8, 9: 8, 10: 10, 11: 10}, + ), + ], + ids=["glm-5.2", "no-sharing", "pipelined", "rank-starts-shared"], +) +def test_selection_leaders_follow_the_indexer_runs(indexer_mask, pp_layers, expected): + """Which layers are handed one selection, and so share one resident set. + + A run starts at a layer owning an indexer; the layers after it reuse its top-k. + A run of one names no selection -- a layer that selects for itself never shares. + """ + from types import SimpleNamespace + + from tensorrt_llm._torch.attention.backends.sparse.dsa.cache_manager import DSACacheManagerV2 + + manager = SimpleNamespace( + pp_layers=pp_layers if pp_layers is not None else list(range(len(indexer_mask))), + _indexer_k_cache_layer_mask=indexer_mask, + _selection_leader_cache=None, + ) + + assert DSACacheManagerV2._selection_leaders(manager) == expected diff --git a/tests/unittest/_torch/attention/sparse/test_sparse_offload_kernels.py b/tests/unittest/_torch/attention/sparse/test_sparse_offload_kernels.py new file mode 100644 index 000000000000..a7e2ea409989 --- /dev/null +++ b/tests/unittest/_torch/attention/sparse/test_sparse_offload_kernels.py @@ -0,0 +1,262 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Tests for the kernels that assemble sparse offload arena rows.""" + +import pytest +import torch + +from tensorrt_llm._torch.attention.backends.sparse.offload_kernels import ( + gather_sparse_batch_rows, + merge_sparse_arena_indices, + select_sparse_history_units, +) + +_requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + + +def _cuda_tensor(value, strided: bool = False) -> torch.Tensor: + tensor = torch.as_tensor(value, dtype=torch.int32, device="cuda") + if not strided: + return tensor + storage = torch.full( + tuple(2 * n for n in tensor.shape), 12345, dtype=torch.int32, device="cuda" + ) + view = storage[tuple(slice(None, None, 2) for _ in tensor.shape)] + view.copy_(tensor) + return view + + +def _merge_reference( + topk, requests, batch_rows, active, lengths, history, fetched, writes, unit_tokens +): + out = torch.full(topk.shape, -1, dtype=torch.int32) + for query, request in enumerate(requests.tolist()): + if query >= active or not 0 <= request < min(active, lengths.numel()): + continue + row = batch_rows[request].item() + if not 0 <= row < history.numel(): + continue + history_units = history[row].item() // unit_tokens + for k, unit in enumerate(topk[query].tolist()): + if not 0 <= unit < lengths[request].item(): + continue + if unit < history_units: + out[query, k] = fetched[row, k] + elif unit == history_units: + out[query, k] = writes[row] + return out + + +@_requires_cuda +@pytest.mark.parametrize("strided", [False, True]) +@pytest.mark.parametrize("active", [0, 2, 3]) +def test_merge_takes_fetched_history_and_the_write_row(strided, active): + # One-token units: a decode step's kv length is its history plus the token + # being written, and that token's ordinal equals the history. Batch rows are + # slots: input positions 0, 1 and 2 sit in rows 2, 3 and 0, and row 1 is free. + history = torch.tensor([7, 0, 12, 5], dtype=torch.int32) + batch_rows = torch.tensor([2, 3, 0, -1], dtype=torch.int32) + lengths = torch.tensor([13, 6, 8], dtype=torch.int32) + requests = torch.tensor([2, 0, 1], dtype=torch.int32) + topk = torch.tensor( + [[7, 6, 0, 8, -1], [12, 0, -1, 11, 13], [0, 1, -1, -1, -1]], dtype=torch.int32 + ) + fetched = torch.arange(4 * 5, dtype=torch.int32).reshape(4, 5) + 100 + writes = torch.tensor([40, 41, 42, 43], dtype=torch.int32) + out = torch.full((3, 5), 999, dtype=torch.int32, device="cuda") + merge_sparse_arena_indices( + _cuda_tensor(topk, strided), + _cuda_tensor(requests, strided), + _cuda_tensor(batch_rows, strided), + _cuda_tensor([active]), + _cuda_tensor(lengths, strided), + _cuda_tensor(history, strided), + _cuda_tensor(fetched, strided), + _cuda_tensor(writes, strided), + out, + unit_tokens=1, + ) + expected = _merge_reference( + topk, requests, batch_rows, active, lengths, history, fetched, writes, 1 + ) + torch.testing.assert_close(out.cpu(), expected) + if active == 3: + # Query 0 is input position 2, in row 0 with history 7: ordinal 7 is the + # token being written, and older ordinals come from row 0's fetch. + assert out[0].tolist() == [40, 101, 102, -1, -1] + # Query 1 is input position 0, in row 2 with history 12; 13 is past its length. + assert out[1].tolist() == [42, 111, -1, 113, -1] + + +@_requires_cuda +@pytest.mark.parametrize("topk_size", [1, 7, 2048]) +def test_merge_matches_selection_and_fetch_order(topk_size): + generator = torch.Generator().manual_seed(7) + # Rows 1 and 3 are free; positions 0, 1 and 2 sit in rows 4, 0 and 2. + history = torch.tensor([41, 0, 1, 0, 3000], dtype=torch.int32) + batch_rows = torch.tensor([4, 0, 2, -1, -1], dtype=torch.int32) + lengths = torch.tensor([3001, 42, 2], dtype=torch.int32) + requests = torch.tensor([1, 2, 0], dtype=torch.int32) + topk = torch.stack( + [torch.randperm(3001, generator=generator)[:topk_size] for _ in range(3)] + ).to(torch.int32) + topk[:, 2::7] = -1 + selected = torch.empty((5, topk_size), dtype=torch.int32, device="cuda") + select_sparse_history_units( + _cuda_tensor(topk), + _cuda_tensor(requests), + _cuda_tensor(batch_rows), + _cuda_tensor([3]), + _cuda_tensor(lengths), + _cuda_tensor(history), + selected, + unit_tokens=1, + ) + assert (selected[[1, 3]] == -1).all(), "Free rows select nothing" + # A fetch returns an arena row for every selected history unit, in slot order. + fetched = torch.where(selected >= 0, selected * 3 + 1, -1) + writes = _cuda_tensor([5, 6, 7, 8, 9]) + out = torch.empty((3, topk_size), dtype=torch.int32, device="cuda") + merge_sparse_arena_indices( + _cuda_tensor(topk), + _cuda_tensor(requests), + _cuda_tensor(batch_rows), + _cuda_tensor([3]), + _cuda_tensor(lengths), + _cuda_tensor(history), + fetched, + writes, + out, + unit_tokens=1, + ) + expected = _merge_reference( + topk, requests, batch_rows, 3, lengths, history, fetched.cpu(), writes.cpu(), 1 + ) + torch.testing.assert_close(out.cpu(), expected) + + +@_requires_cuda +def test_merge_replays_in_a_cuda_graph_with_new_watermarks_and_rows(): + topk = _cuda_tensor([[0, 1, 2, 3], [0, 1, -1, -1]]) + requests = _cuda_tensor([0, 1]) + batch_rows = _cuda_tensor([0, 1, -1]) + active = _cuda_tensor([2]) + history = _cuda_tensor([2, 1, 0]) + lengths = _cuda_tensor([3, 2]) + fetched = _cuda_tensor([[10, 11, 12, 13], [20, 21, 22, 23], [30, 31, 32, 33]]) + writes = _cuda_tensor([7, 8, 9]) + out = torch.empty((2, 4), dtype=torch.int32, device="cuda") + args = (topk, requests, batch_rows, active, lengths, history, fetched, writes, out) + merge_sparse_arena_indices(*args, unit_tokens=1) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + merge_sparse_arena_indices(*args, unit_tokens=1) + for row0, hist0, length0, active_count, expected0 in ( + (0, 2, 3, 2, [10, 11, 7, -1]), + (0, 3, 4, 2, [10, 11, 12, 7]), + (0, 0, 1, 2, [7, -1, -1, -1]), + (0, 3, 4, 0, [-1, -1, -1, -1]), + # The request moved to row 2 between steps: the replay follows the map. + (2, 1, 3, 2, [30, 9, -1, -1]), + (-1, 1, 3, 2, [-1, -1, -1, -1]), + ): + batch_rows[0] = row0 + history[max(row0, 0)] = hist0 + lengths[0] = length0 + active[0] = active_count + graph.replay() + torch.testing.assert_close(out[0].cpu(), torch.tensor(expected0, dtype=torch.int32)) + + +@_requires_cuda +@pytest.mark.parametrize( + "invalid", ["dtype", "width", "requests", "batch_rows", "writes", "lengths", "unit"] +) +def test_merge_rejects_invalid_layout(invalid): + topk = _cuda_tensor([[0, 1], [1, 0]]) + requests, batch_rows, active, lengths, history = ( + _cuda_tensor([0, 1]), + _cuda_tensor([0, 1, -1, -1]), + _cuda_tensor([2]), + _cuda_tensor([2, 2]), + _cuda_tensor([1] * 4), + ) + fetched = _cuda_tensor([[0, 0]] * 4) + writes = _cuda_tensor([0] * 4) + unit = 1 + if invalid == "dtype": + topk = topk.to(torch.int64) + elif invalid == "width": + fetched = _cuda_tensor([[0]] * 4) + elif invalid == "requests": + requests = requests[:1] + elif invalid == "batch_rows": + batch_rows = batch_rows[:2] + elif invalid == "writes": + writes = writes[:2] + elif invalid == "lengths": + lengths = lengths[:1] + else: + unit = 0 + with pytest.raises(ValueError): + merge_sparse_arena_indices( + topk, + requests, + batch_rows, + active, + lengths, + history, + fetched, + writes, + torch.empty((2, 2), dtype=torch.int32, device="cuda"), + unit_tokens=unit, + ) + + +@_requires_cuda +@pytest.mark.parametrize("strided", [False, True]) +@pytest.mark.parametrize("count", [1, 3, 300]) +def test_gather_reads_each_request_from_its_batch_row(strided, count): + generator = torch.Generator().manual_seed(count) + num_rows = max(2 * count, 4) + values = torch.randint(0, 1 << 20, (num_rows,), generator=generator, dtype=torch.int32) + rows = torch.randperm(num_rows, generator=generator)[:count].to(torch.int32) + rows[::3] = -1 + columns = _cuda_tensor(torch.stack((values + 1, values), dim=1), strided) + out = _cuda_tensor(torch.full((count,), 999, dtype=torch.int32), strided) + result = gather_sparse_batch_rows(columns[:, 1], _cuda_tensor(rows, strided), out) + assert result is out + expected = torch.where(rows >= 0, values[rows.clamp(min=0).long()], -1) + torch.testing.assert_close(out.cpu(), expected) + allocated = gather_sparse_batch_rows(columns[:, 1], _cuda_tensor(rows)) + assert allocated.is_contiguous() and allocated.dtype == torch.int32 + torch.testing.assert_close(allocated.cpu(), expected) + + +@_requires_cuda +def test_gather_replays_with_new_rows_and_ignores_rows_past_the_batch(): + values = _cuda_tensor([10, 11, 12]) + rows = _cuda_tensor([2, 0]) + out = torch.empty(2, dtype=torch.int32, device="cuda") + gather_sparse_batch_rows(values, rows, out) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + gather_sparse_batch_rows(values, rows, out) + for new_rows, expected in (([2, 0], [12, 10]), ([1, -1], [11, -1]), ([3, 2], [-1, 12])): + rows.copy_(_cuda_tensor(new_rows)) + graph.replay() + assert out.tolist() == expected + + +@_requires_cuda +@pytest.mark.parametrize("invalid", ["dtype", "ndim", "out"]) +def test_gather_rejects_invalid_layout(invalid): + values, rows, out = _cuda_tensor([1, 2]), _cuda_tensor([0, 1]), None + if invalid == "dtype": + values = values.to(torch.int64) + elif invalid == "ndim": + values = values.view(1, 2) + else: + out = torch.empty(3, dtype=torch.int32, device="cuda") + with pytest.raises(ValueError): + gather_sparse_batch_rows(values, rows, out) diff --git a/tests/unittest/_torch/executor/engine/runners/decoder/test_warmup.py b/tests/unittest/_torch/executor/engine/runners/decoder/test_warmup.py index 2ca5fbd39a61..d596a1b580f5 100644 --- a/tests/unittest/_torch/executor/engine/runners/decoder/test_warmup.py +++ b/tests/unittest/_torch/executor/engine/runners/decoder/test_warmup.py @@ -31,14 +31,17 @@ ) import tensorrt_llm._torch.pyexecutor.engine.model_call as model_call_module +import tensorrt_llm._torch.pyexecutor.engine.runners.decoder.graphs as graphs_module import tensorrt_llm._torch.pyexecutor.engine.runners.decoder.warmup as decoder_warmup_module from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata from tensorrt_llm._torch.custom_ops.torch_custom_ops import MXFP8GemmRunner from tensorrt_llm._torch.models.modeling_utils import remove_weights from tensorrt_llm._torch.modules.linear import MXFP8LinearMethod +from tensorrt_llm._torch.pyexecutor.config_utils import is_decode_only from tensorrt_llm._torch.pyexecutor.engine.model_call import ModelCaller from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import DecoderRunner from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import warmup as warmup_module +from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.graphs import DecoderCudaGraphs from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.metadata import DecoderMetadata from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.warmup import DecoderWarmup from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 @@ -46,7 +49,14 @@ from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests from tensorrt_llm._torch.pyexecutor.warmup_timer import _WarmupTimer from tensorrt_llm._torch.utils import is_torch_compiling, torch_compiling -from tensorrt_llm.llmapi.llm_args import DraftTargetDecodingConfig, PrefillCudaGraphBackend +from tensorrt_llm.llmapi.llm_args import ( + BaseSparseAttentionConfig, + DeepSeekSparseAttentionConfig, + DeepSeekV4SparseAttentionConfig, + DraftTargetDecodingConfig, + PrefillCudaGraphBackend, + SparseKvCacheOffloadConfig, +) if TYPE_CHECKING: from collections.abc import Callable @@ -1446,3 +1456,86 @@ def record_metric_scope( any("Skipping warm up as no KV Cache manager allocated." in log for log in logs) ) self.assertGreaterEqual(warmup._metrics["sampling_warmup_seconds"], 0) + + +class _OffloadingSparseAttentionConfig(BaseSparseAttentionConfig): + """Any sparse attention algorithm that offloads KV history to host memory.""" + + algorithm: str = "offloading" + + def get_kv_cache_offload_config(self) -> SparseKvCacheOffloadConfig: + return SparseKvCacheOffloadConfig() + + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + ("config", "expected"), + [ + (None, False), + (DeepSeekSparseAttentionConfig(), False), + (DeepSeekV4SparseAttentionConfig(), False), + (DeepSeekV4SparseAttentionConfig(enable_kv_cache_offload=True), True), + (_OffloadingSparseAttentionConfig(), True), + ], + ids=["none", "dsa", "dsv4", "dsv4_offload", "any_offload"], +) +def test_only_kv_cache_offload_makes_a_worker_decode_only(config, expected: bool) -> None: + assert is_decode_only(config) is expected + + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + ("max_total_draft_tokens", "expected"), + [(0, [(1, 1), (13, 13)]), (2, [(3, 1), (39, 13)])], +) +def test_generation_warmup_shapes_drop_context_tokens( + max_total_draft_tokens: int, expected: list[tuple[int, int]] +) -> None: + warmup = SimpleNamespace(_config=SimpleNamespace(max_total_draft_tokens=max_total_draft_tokens)) + configs = [(1, 0), (1, 1), (3, 1), (258, 0), (13, 13), (39, 13), (2, 0), (2, 1)] + + assert DecoderWarmup._generation_warmup_shapes(warmup, configs) == expected + + +@pytest.mark.cpu_only +def test_decode_only_skips_prefill_cuda_graph_capture() -> None: + """A decode-only worker never runs a context batch, so prefill graphs are dead weight.""" + graphs = SimpleNamespace( + _config=SimpleNamespace( + prefill_cuda_graph_backend=PrefillCudaGraphBackend.BREAKABLE, + torch_compile_enabled=True, + prefill_cuda_graph_num_tokens=[16], + ), + sparse_attention_config=_OffloadingSparseAttentionConfig(), + _dummy_requests=SimpleNamespace(create_warmup_request=Mock()), + ) + + DecoderCudaGraphs._capture_prefill_cuda_graphs(graphs, Mock()) + + graphs._dummy_requests.create_warmup_request.assert_not_called() + + +@pytest.mark.parametrize("decode_only", [False, True], ids=["default", "decode_only"]) +def test_decode_only_warmup_requests_generation_shapes_only(decode_only: bool) -> None: + """Every warmup phase of a decode-only worker builds pure-generation batches.""" + with _warmup_and_resources() as (warmup, resource_manager): + shapes: list[tuple[int, int]] = [] + + def record(resources, num_tokens, num_gen_requests, **kwargs): + shapes.append((num_tokens, num_gen_requests)) + return None # No batch, so each phase skips its forward. + + with ( + patch.object(decoder_warmup_module, "is_decode_only", return_value=decode_only), + patch.object(graphs_module, "is_decode_only", return_value=decode_only), + patch.object(warmup._dummy_requests, "create_warmup_request", side_effect=record), + patch.object(warmup._graphs, "_capture_generation_cuda_graphs"), + ): + warmup.run(resource_manager) + + assert shapes + context_shapes = [shape for shape in shapes if shape[0] != shape[1]] + if decode_only: + assert context_shapes == [], shapes + else: + assert context_shapes, shapes 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 79717b94f868..fe3a2cdba1fb 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 @@ -184,11 +184,7 @@ def sparse_layers(monkeypatch, indices): 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, - SparseCacheConfig, - ) + from tensorrt_llm.runtime.kv_cache_manager_v2 import AttentionLayerConfig, BufferConfig build = KVCacheManagerV2._build_cache_config @@ -206,9 +202,9 @@ def with_sparse_layers(self, config): 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 - ), + # No sparse_cache: declaring one obliges the manager to be given a + # sparse_kv_offload_config to reserve the arenas from, and the refusal + # under test is about sparse buffers, which a layer has either way. ) return build(self, replace(config, layers=layers)) @@ -232,7 +228,8 @@ def test_a_layer_group_of_sparse_buffers_is_refused(kit, host_tier_manager, monk with host_tier_manager(**keywords) as mgr: layers = {int(layer.layer_id): layer for layer in mgr.impl.init_config.layers} flags = [ - any(b.is_sparse for b in layers[int(ids[0])].buffers) for ids in mgr.impl.layer_grouping + any(b.is_sparse for lid in ids for b in layers[int(lid)].buffers) + for ids in mgr.impl.layer_grouping ] sparse = [lg for lg, flag in enumerate(flags) if flag] assert sparse and (0 in sparse) == (case == "first"), ( diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_scheduler.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_scheduler.py index ae062f0983c5..7195f4e61f55 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_scheduler.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_scheduler.py @@ -28,7 +28,6 @@ KVCacheManagerV2, ) from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState -from tensorrt_llm._torch.pyexecutor.resource_manager import CacheTypeCpp from tensorrt_llm._torch.pyexecutor.scheduler.prefix_wait_planner import ( MAX_WAIT_PASSES, MIN_GAIN_BLOCKS, @@ -39,293 +38,6 @@ pytestmark = pytest.mark.cpu_only -@pytest.mark.parametrize("active", [False, True]) -@pytest.mark.parametrize("admitted", [False, True]) -@pytest.mark.parametrize("beam_admitted", [False, True]) -def test_generation_admits_decode_before_capacity_growth( - active: bool, admitted: bool, beam_admitted: bool -) -> None: - manager = object.__new__(KVCacheManagerV2) - cache = Mock(is_active=active, capacity=8) - cache.enter_decode.return_value = admitted - cache.resume.return_value = admitted - cache.resize.return_value = True - manager.kv_cache_map = {1: cache} - manager._stream = Mock(cuda_stream=123) - manager._restore_page_index_bufs = Mock() - manager._ensure_generation_beam_width = Mock(return_value=beam_admitted) - manager._generation_draft_slots = Mock(return_value=0) - manager._allocated_draft_lens = {} - manager._inflight_draft_lens = {} - manager._has_cp_helix = False - manager._fill_fresh_kv_pages = Mock() - manager._log_window_crossing = Mock() - req = Mock(py_request_id=1) - - allocated = admitted and beam_admitted - assert manager.try_allocate_generation(req) == allocated - admission = call.enter_decode() if active else call.resume(123, is_decoding=True) - assert cache.mock_calls == [admission] + ([call.resize(9)] if allocated else []) - if admitted: - manager._ensure_generation_beam_width.assert_called_once_with(req, cache) - else: - manager._ensure_generation_beam_width.assert_not_called() - if not active and admitted: - manager._restore_page_index_bufs.assert_called_once_with(1, cache) - else: - manager._restore_page_index_bufs.assert_not_called() - assert manager._allocated_draft_lens == ({1: 0} if allocated else {}) - - -@pytest.mark.parametrize("sparse", [False, True]) -def test_resume_restores_page_buffers_and_sparse_metadata_row(sparse: bool) -> None: - manager = object.__new__(KVCacheManagerV2) - manager.index_mapper = Mock() - manager.index_mapper.get_index.return_value = 3 - manager._set_page_index_bufs = Mock() - manager.sparse_metadata_batch = Mock() if sparse else None - cache = Mock() - - manager._restore_page_index_bufs(7, cache) - - manager._set_page_index_bufs.assert_called_once_with(7, cache) - if sparse: - manager.index_mapper.get_index.assert_called_once_with(7) - manager.sparse_metadata_batch.add.assert_called_once_with(cache, 3) - - -@pytest.mark.parametrize("is_draft", [False, True]) -def test_sparse_metadata_publishes_after_preparation(is_draft: bool) -> None: - manager = object.__new__(KVCacheManagerV2) - manager.impl = Mock() - order = Mock() - manager._disagg_receive_ready = {} - manager.is_draft = is_draft - manager._stream = Mock(cuda_stream=123) - manager.sparse_metadata_batch = Mock() - manager.kv_connector_manager = Mock(capacity_only=False) - manager._prepare_draft_resources = order.prepare_draft - manager._run_kv_connector_hooks = order.connector - order.attach_mock(manager.sparse_metadata_batch.record_read, "record_read") - order.attach_mock(manager.sparse_metadata_batch.publish, "publish") - scheduled = Mock(context_requests=[]) - manager.prepare_resources(scheduled) - prepare = call.prepare_draft(scheduled) if is_draft else call.connector(scheduled) - assert order.mock_calls == [prepare, call.record_read(123), call.publish(123)] - - -def test_sparse_metadata_republishes_after_connector_acceptance() -> None: - manager = object.__new__(KVCacheManagerV2) - order = Mock() - manager.is_draft = False - manager._stream = Mock(cuda_stream=123) - manager.sparse_metadata_batch = Mock() - manager.kv_connector_manager = Mock() - manager._connector_reservations_enabled = Mock(return_value=True) - manager._accept_connector_prefix_reservations = order.accept - manager._on_allocation_stream = set() - order.attach_mock(manager.kv_connector_manager.build_scheduler_output, "report") - order.attach_mock(manager.sparse_metadata_batch.record_read, "record_read") - order.attach_mock(manager.sparse_metadata_batch.publish, "publish") - scheduled = Mock(context_requests=[]) - manager.report_batch_to_connector(scheduled) - assert order.mock_calls == [ - call.accept(scheduled), - call.report(scheduled, manager), - call.record_read(123), - call.publish(123), - ] - - -@pytest.fixture -def sparse_offset_manager() -> KVCacheManagerV2: - manager = object.__new__(KVCacheManagerV2) - manager._stream = Mock(cuda_stream=123) - manager.sparse_metadata_batch = Mock() - manager._sparse_layer_group_ids = (1, 3) - manager.max_copy_beam_width = 1 - manager.kv_cache_type = CacheTypeCpp.SELF - manager.tokens_per_block = 4 - manager.kv_cache_map = {req_id: Mock(is_decoding=True, history_length=8) for req_id in (7, 8)} - for cache in manager.kv_cache_map.values(): - cache.get_page_storage_snapshot.return_value = Mock( - cache_levels=[0, 0], eligible_history_blocks=0 - ) - manager._use_per_layer_page_tables = False - manager._copy_batch_block_offsets_per_layer = Mock() - manager._page_table_materializer = Mock(uses_device_expansion=False, uses_device_staging=False) - manager.index_mapper = Mock() - manager.index_mapper.get_copy_index.return_value = Mock(shape=(2,)) - manager.host_kv_cache_block_offsets = Mock() - manager.index_scales = Mock() - manager.kv_offset = Mock() - return manager - - -@pytest.mark.parametrize("per_layer", [False, True]) -@pytest.mark.parametrize( - "cache_levels", - [pytest.param([0, 0], id="gpu"), pytest.param([None, 0, None], id="invalid-slots")], -) -def test_sparse_gpu_resident_history_uses_dense_attention_offsets( - sparse_offset_manager: KVCacheManagerV2, cache_levels: list[int | None], per_layer: bool -) -> None: - manager = sparse_offset_manager - manager._use_per_layer_page_tables = per_layer - manager.kv_cache_map[8].get_page_storage_snapshot.return_value.cache_levels = cache_levels - unscheduled_cache = Mock(is_decoding=True, history_length=8) - unscheduled_cache.get_page_storage_snapshot.return_value = Mock(cache_levels=[1, 1]) - manager.kv_cache_map[9] = unscheduled_cache - destination = Mock() - with patch( - "tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2." - "copy_batch_block_offsets_to_device" - ) as dense_copy: - manager.copy_batch_block_offsets(destination, [7, 8], 1, 0, 2) - if per_layer: - manager._copy_batch_block_offsets_per_layer.assert_called_once_with( - destination, [7, 8], manager.index_mapper.get_copy_index.return_value, 0, 2 - ) - dense_copy.assert_not_called() - else: - dense_copy.assert_called_once_with( - manager.host_kv_cache_block_offsets, - destination, - manager.index_mapper.get_copy_index.return_value, - manager.index_scales, - manager.kv_offset, - 123, - ) - manager._copy_batch_block_offsets_per_layer.assert_not_called() - for req_id in (7, 8): - assert manager.kv_cache_map[req_id].get_page_storage_snapshot.call_args_list == [ - call(1), - call(3), - ] - unscheduled_cache.get_page_storage_snapshot.assert_not_called() - manager.sparse_metadata_batch.publish.assert_called_once_with(123) - - -@pytest.mark.parametrize( - ("cache_levels", "eligible_history_blocks"), - [ - pytest.param([1, 1], 2, id="host"), - pytest.param([0, 1], 0, id="host-after-gpu"), - pytest.param([None, 1], 0, id="host-after-invalid-slot"), - ], -) -def test_sparse_host_indices_cannot_reach_dense_attention_offsets( - sparse_offset_manager: KVCacheManagerV2, - cache_levels: list[int | None], - eligible_history_blocks: int, -) -> None: - manager = sparse_offset_manager - snapshots = { - 1: Mock(cache_levels=[0, 0], eligible_history_blocks=0), - 3: Mock(cache_levels=cache_levels, eligible_history_blocks=eligible_history_blocks), - } - manager.kv_cache_map[8].get_page_storage_snapshot.side_effect = snapshots.__getitem__ - with patch( - "tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2." - "copy_batch_block_offsets_to_device" - ) as dense_copy: - with pytest.raises(RuntimeError, match="Offloaded sparse history"): - manager.copy_batch_block_offsets(Mock(), [7, 8], 1, 0, 2) - dense_copy.assert_not_called() - manager._copy_batch_block_offsets_per_layer.assert_not_called() - manager.sparse_metadata_batch.publish.assert_called_once_with(123) - - -def test_sparse_dense_offsets_check_residency_after_publication( - sparse_offset_manager: KVCacheManagerV2, -) -> None: - manager = sparse_offset_manager - snapshot = manager.kv_cache_map[8].get_page_storage_snapshot.return_value - - def offload_history(stream: int) -> None: - snapshot.cache_levels = [1, 1] - snapshot.eligible_history_blocks = 2 - - manager.sparse_metadata_batch.publish.side_effect = offload_history - with patch( - "tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2." - "copy_batch_block_offsets_to_device" - ) as dense_copy: - with pytest.raises(RuntimeError, match="Offloaded sparse history"): - manager.copy_batch_block_offsets(Mock(), [7, 8], 1, 0, 2) - dense_copy.assert_not_called() - manager._copy_batch_block_offsets_per_layer.assert_not_called() - manager.sparse_metadata_batch.publish.assert_called_once_with(123) - - -def test_sparse_publication_failure_stops_metadata_preparation() -> None: - manager = object.__new__(KVCacheManagerV2) - manager._stream = Mock(cuda_stream=123) - manager.sparse_metadata_batch = Mock() - manager.sparse_metadata_batch.publish.side_effect = RuntimeError("upload failed") - with patch( - "tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2." - "copy_batch_block_offsets_to_device" - ) as dense_copy: - with pytest.raises(RuntimeError, match="upload failed"): - manager.copy_batch_block_offsets(Mock(), [7], 1, 0, 1) - dense_copy.assert_not_called() - - -def test_dense_metadata_preparation_uses_existing_offsets() -> None: - manager = object.__new__(KVCacheManagerV2) - manager._stream = Mock(cuda_stream=123) - manager.max_copy_beam_width = 1 - manager.kv_cache_type = CacheTypeCpp.SELF - manager._use_per_layer_page_tables = False - manager._page_table_materializer = Mock(uses_device_expansion=False, uses_device_staging=False) - manager.index_mapper = Mock() - manager.index_mapper.get_copy_index.return_value = Mock(shape=(1,)) - manager.host_kv_cache_block_offsets = Mock() - manager.index_scales = Mock() - manager.kv_offset = Mock() - destination = Mock() - with patch( - "tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2." - "copy_batch_block_offsets_to_device" - ) as dense_copy: - manager.copy_batch_block_offsets(destination, [7], 1, 0, 1) - dense_copy.assert_called_once_with( - manager.host_kv_cache_block_offsets, - destination, - manager.index_mapper.get_copy_index.return_value, - manager.index_scales, - manager.kv_offset, - 123, - ) - - -def test_sparse_index_slot_release_detaches_batch_before_reuse() -> None: - manager = object.__new__(KVCacheManagerV2) - manager.is_draft = False - manager._stream = Mock(cuda_stream=123) - manager.max_beam_width = 2 - manager.num_pools = 1 - manager._early_freed_index_requests = set() - cache = Mock(beam_width=1) - manager.kv_cache_map = {7: cache} - manager.sparse_metadata_batch = Mock() - manager.index_mapper = Mock() - order = Mock() - order.attach_mock(manager.sparse_metadata_batch.record_read, "record_read") - order.attach_mock(manager.sparse_metadata_batch.remove, "remove") - order.attach_mock(cache.set_base_page_index_buf, "detach_buffer") - order.attach_mock(manager.index_mapper.remove_sequence, "release_slot") - manager.release_index_slot(7) - assert order.mock_calls == [ - call.record_read(123), - call.remove(cache), - call.detach_buffer(0, 0, None), - call.release_slot(7), - ] - assert manager._early_freed_index_requests == {7} - - # --------------------------------------------------------------------------- # State value constants # --------------------------------------------------------------------------- @@ -5058,7 +4770,6 @@ def test_recompute_paused_request_does_not_defer_duplicate(self): def test_connector_credit_admits_another_request(chunked: bool) -> None: from types import SimpleNamespace - from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, SamplingConfig from tensorrt_llm._torch.pyexecutor.scheduler.scheduler_v2 import BudgetTracker diff --git a/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py b/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py index addb326f0a91..f220b6361411 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py @@ -15,9 +15,10 @@ import array import os +import sys from dataclasses import dataclass, field, replace from types import SimpleNamespace -from unittest.mock import Mock, call, patch +from unittest.mock import ANY, Mock, call, patch import numpy as np import pytest @@ -28,18 +29,28 @@ from tensorrt_llm._torch.pyexecutor.kv_cache import kv_cache_manager_v2 as kv_cache_v2_module from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import ( BlockReusePolicy, + HostMemoryLimits, KVCacheLayerOwner, KVCacheLayerSpec, KVCacheManagerV2, Role, _extend_swa_windows_for_reuse, _KVCacheManagerInitStatus, + _raise_on_every_rank, _swa_branch_snapshot_point, _swa_endpoint_priority, _sync_kv_cache_manager_init_status, _update_kv_cache_draft_token_location, ) +from tensorrt_llm._torch.pyexecutor.kv_cache.sparse_offload_memory import ( + SPARSE_ARENA_ALIGNMENT, + HostHistoryReservations, + sparse_arena_bytes, + sparse_arena_shapes, + sparse_batch_shape, +) from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, LlmRequestState +from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor from tensorrt_llm._torch.pyexecutor.resource_manager import ResourceManager, ResourceManagerType from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests from tensorrt_llm.bindings import DataType, SamplingConfig @@ -55,6 +66,7 @@ KvCacheConfig, MTPDecodingConfig, PARDDecodingConfig, + SparseKvCacheOffloadConfig, ) from tensorrt_llm.mapping import Mapping from tensorrt_llm.quantization import QuantMode @@ -63,9 +75,9 @@ DEFAULT_BEAM_INDEX, AttentionLayerConfig, AttnLifeCycle, + Batch, BatchDesc, BufferConfig, - BufferId, CacheLevel, CudaStream, DataRole, @@ -78,6 +90,7 @@ LayerId, PageIndexMode, ReuseScope, + SparseCacheConfig, SsmLayerConfig, _introspection, ) @@ -150,6 +163,8 @@ def test_create_kv_cache_swa_endpoint_priority( manager.num_pools = 0 manager.impl = Mock() manager.impl.create_kv_cache.return_value.beam_width = 1 + # Belongs to no Batch, so it still owns its page index buffers. + manager.impl.create_kv_cache.return_value.batch_row = None manager._create_kv_cache(0, None, None, is_dummy=is_dummy, expected_prompt_length=prompt) @@ -325,6 +340,7 @@ def _make_cache_config_for_test( pp_layers: list[int] | None = None, dtype: DataType = DataType.HALF, kv_cache_type: CacheType = CacheType.SELFKONLY, + sparse_kv_offload: SparseKvCacheOffloadConfig | None = None, ) -> KVCacheManagerConfig: if max_attention_window_vec is None: max_attention_window_vec = [None] @@ -334,6 +350,7 @@ def _make_cache_config_for_test( cache_manager = object.__new__(KVCacheManagerV2) cache_manager.layer_specs = None + cache_manager._sparse_kv_offload = sparse_kv_offload cache_manager.max_beam_width = 1 cache_manager.kv_cache_type = kv_cache_type cache_manager.dtype = dtype @@ -423,8 +440,6 @@ def build_cache_config( fake_impl.layer_grouping = [[0]] fake_impl.pool_group_descs = [] fake_impl.get_layer_group_id.side_effect = lambda _: 0 - fake_impl.all_buffer_ids = [BufferId(LayerId(0), Role.KEY)] - fake_impl.is_sparse.return_value = False module = "tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2" with ( @@ -946,6 +961,50 @@ def test_zero_size_filter_preserves_ssm_and_nonempty_extra_buffers() -> None: assert filtered.initial_pool_ratio == pytest.approx([0.25, 0.75]) +def test_zero_size_filter_keeps_a_sparse_fetch_path() -> None: + manager = object.__new__(KVCacheManagerV2) + manager.pp_layers = [0, 1] + manager.num_layers = 2 + manager.num_local_layers = 2 + manager.layer_offsets = {0: 0, 1: 1} + manager.num_kv_heads_per_layer = [1, 1] + manager.head_dim_per_layer = [64, 64] + manager.max_attention_window_vec = [None, None] + cache = SparseCacheConfig( + num_selected_blocks=1, max_cached_units=2, unit_tokens=TOKENS_PER_BLOCK + ) + config = KVCacheManagerConfig( + tokens_per_block=TOKENS_PER_BLOCK, + cache_tiers=[GpuCacheTierConfig(quota=1 << 20)], + layers=[ + AttentionLayerConfig( + layer_id=LayerId(0), + buffers=[ + BufferConfig(role=Role.KEY, size=1024, is_sparse=True), + BufferConfig(role=Role.VALUE, size=256), + BufferConfig(role=Role.INDEX_KEY, size=0), + ], + sparse_cache=cache, + ), + # Its only sparse buffer is empty, so nothing is left to fetch. + AttentionLayerConfig( + layer_id=LayerId(1), + buffers=[ + BufferConfig(role=Role.KEY, size=0, is_sparse=True), + BufferConfig(role=Role.VALUE, size=256), + ], + sparse_cache=cache, + ), + ], + ) + + kept, emptied = manager._remove_zero_size_buffers(config).layers + + assert kept.sparse_cache is not None + assert kept.sparse_cache.max_cached_units == 2 + assert emptied.sparse_cache is None + + def test_zero_size_filter_rejects_empty_local_cache() -> None: manager = object.__new__(KVCacheManagerV2) config = KVCacheManagerConfig( @@ -1328,6 +1387,21 @@ def test_avg_seq_len_builds_context_warmup_constraint() -> None: assert config.constraints == [BatchDesc([KVCacheDesc(capacity=2048, history_length=0)])] +def test_sparse_kv_offload_typical_step_models_only_decoding() -> None: + config = _make_cache_config_for_test( + KvCacheConfig(host_cache_size=0, avg_seq_len=1024), + max_batch_size=3, + max_seq_len=1024, + max_num_tokens=2048, + max_draft_len=2, + sparse_kv_offload=SparseKvCacheOffloadConfig(), + ) + + # A worker that offloads sparse history only decodes. + assert config.typical_step == BatchDesc([KVCacheDesc(capacity=1024, history_length=1021)] * 3) + assert config.constraints == [BatchDesc([KVCacheDesc(capacity=2048, history_length=0)])] + + def test_avg_seq_len_updates_typical_step() -> None: config = _make_cache_config_for_test( KvCacheConfig(avg_seq_len=256), @@ -3114,6 +3188,8 @@ def _index_mapper_capacity_for( is_disagg: bool = False, num_reserved_index_slots: int = 1, disable_overlap_scheduler: bool = True, + sparse: bool = False, + guard_page: bool = False, ) -> tuple[int, int, int]: """Construct a manager and return the three sizes it derives. @@ -3130,8 +3206,6 @@ def _index_mapper_capacity_for( fake_impl.layer_grouping = [[0]] fake_impl.pool_group_descs = [] fake_impl.get_layer_group_id.side_effect = lambda _: 0 - fake_impl.all_buffer_ids = [BufferId(LayerId(0), Role.KEY)] - fake_impl.is_sparse.return_value = False def build_base_config( self: KVCacheManagerV2, @@ -3139,27 +3213,44 @@ def build_base_config( *, tokens_per_block: int, cache_tiers: list[object], - ) -> _FakeManagerConfig: - del self, config, tokens_per_block - return _FakeManagerConfig(cache_tiers=cache_tiers) + ) -> KVCacheManagerConfig: + del self, config + return KVCacheManagerConfig( + tokens_per_block=tokens_per_block, + cache_tiers=cache_tiers, + layers=[ + AttentionLayerConfig( + layer_id=LayerId(0), + buffers=[ + BufferConfig(role=Role.KEY, size=128, is_sparse=sparse), + BufferConfig(role=Role.VALUE, size=128), + ], + sparse_cache=SparseCacheConfig( + num_selected_blocks=1, max_cached_units=2, unit_tokens=tokens_per_block + ) + if sparse + else None, + ) + ], + ) with ( - # Pin both diagnostics off: the guard page would add an index slot. patch.dict( os.environ, - {"TRTLLM_KV_GUARD_PAGE": "", "TRTLLM_KV_FRESH_PAGE_FILL": ""}, + {"TRTLLM_KV_GUARD_PAGE": "1" if guard_page else "", "TRTLLM_KV_FRESH_PAGE_FILL": ""}, ), patch(f"{module}.IndexMapper") as index_mapper_cls, - patch(f"{module}.KVCacheManagerPy", Mock(return_value=fake_impl)), + patch(f"{module}.KVCacheManagerPy", return_value=fake_impl) as create_impl, patch.object(KVCacheManagerV2, "_build_base_config", build_base_config), patch.object(KVCacheManagerV2, "_build_cache_config", lambda self, config: config), patch.object(KVCacheManagerV2, "get_num_available_tokens", return_value=MAX_SEQ_LEN), patch.object(KVCacheManagerV2, "_prepare_page_table_tensor") as page_table, patch.object(KVCacheManagerV2, "_log_kv_cache_pool_lifecycle_mapping"), + patch.object(KVCacheManagerV2, "_reserve_guard_page"), ): manager = KVCacheManagerV2( # A quota must be set or __init__ asserts; the value is irrelevant here. - KvCacheConfig(max_gpu_total_bytes=16 << 20), + KvCacheConfig(max_gpu_total_bytes=16 << 20, enable_block_reuse=False), CacheType.SELFKONLY, num_layers=1, num_kv_heads=1, @@ -3179,9 +3270,14 @@ def build_base_config( is_disagg=is_disagg, num_reserved_index_slots=num_reserved_index_slots, disable_overlap_scheduler=disable_overlap_scheduler, + sparse_kv_offload_config=SparseKvCacheOffloadConfig() if sparse else None, ) index_mapper_cls.assert_called_once() page_table.assert_called_once() + config = create_impl.call_args.args[0] + # The sparse arenas are the Batch's and take their lease count from its rows, so + # nothing encodes a request count as a constraint any more. + assert not config.constraints return ( index_mapper_cls.call_args.args[0], page_table.call_args.args[0], @@ -3207,10 +3303,13 @@ def build_base_config( pytest.param(2, 4, True, True, 1, 17, id="pp4_disagg_unchanged"), # Reserved slots are still added on top of the widened pool. pytest.param(2, 1, False, False, 5, 9, id="reserved_slots_still_added"), + pytest.param(1, 1, True, False, 1, 2, id="single_live_request_and_padding"), + pytest.param(2, 1, True, False, 0, 2, id="no_reserved_slots"), ] @pytest.mark.cpu_only +@pytest.mark.parametrize("sparse", [False, True]) @pytest.mark.parametrize( "max_batch_size,pp_size,disable_overlap_scheduler,is_disagg,reserved,expected", _INDEX_MAPPER_CAPACITY_CASES, @@ -3222,6 +3321,7 @@ def test_index_mapper_capacity_covers_the_overlapping_cohorts( is_disagg: bool, reserved: int, expected: int, + sparse: bool, ) -> None: capacity, page_table_capacity, max_admissible_sequences = _index_mapper_capacity_for( max_batch_size=max_batch_size, @@ -3229,144 +3329,1334 @@ def test_index_mapper_capacity_covers_the_overlapping_cohorts( is_disagg=is_disagg, num_reserved_index_slots=reserved, disable_overlap_scheduler=disable_overlap_scheduler, + sparse=sparse, ) assert capacity == expected assert page_table_capacity == expected assert max_admissible_sequences == expected - reserved -@pytest.mark.parametrize( - "setting, expected", - [ - ("", None), - ("1", 0.0), - ("on", 0.0), - ("true", 0.0), - ("zero", 0.0), - ("0", 0.0), - ("-2.5", -2.5), - ("garbage", None), - ], -) -def test_kv_fill_value_parsing(setting, expected): - """The knobs are off when unset and ignore anything they cannot parse.""" - assert kv_cache_v2_module._parse_kv_fill_value(setting, "TEST_VAR") == expected +@pytest.mark.cpu_only +def test_sparse_arena_also_reserves_guard_request() -> None: + assert _index_mapper_capacity_for( + max_batch_size=2, num_reserved_index_slots=1, sparse=True, guard_page=True + ) == (4, 4, 2) -def test_kv_fill_value_parsing_nan(): - value = kv_cache_v2_module._parse_kv_fill_value("nan", "TEST_VAR") - assert value is not None and np.isnan(value) +def _sparse_arena_config( + max_batch_size: int, + explicit_ratio: bool, + dense_size: int = 4096, + *, + num_layers: int = 1, + shared_selection: bool = False, +) -> KVCacheManagerConfig: + return KVCacheManagerConfig( + tokens_per_block=16, + cache_tiers=[GpuCacheTierConfig(4 << 20), HostCacheTierConfig(4 << 20)], + layers=[ + AttentionLayerConfig( + layer_id=LayerId(layer), + buffers=[ + # Sparse pages of four 4-token units, beside the dense buffer + # that sizes the pool group's slots. + BufferConfig(role=Role.KEY, size=4 * 4096, is_sparse=True), + BufferConfig(role=Role.KEY_BLOCK_SCALE, size=4 * 128, is_sparse=True), + BufferConfig(role=Role.VALUE, size=dense_size), + ], + sparse_cache=SparseCacheConfig( + num_selected_blocks=2, + max_cached_units=4, + unit_tokens=4, + selection_id=0 if shared_selection else None, + ), + ) + for layer in range(num_layers) + ], + constraints=[] + if explicit_ratio + else [BatchDesc([KVCacheDesc(capacity=16, history_length=0)] * max_batch_size)], + initial_pool_ratio=[1.0] if explicit_ratio else None, + ) -def test_fill_kv_pages_float_buffer(): - buffer = torch.ones(4, 3, 2) - kv_cache_v2_module._fill_kv_pages(buffer, [1, 3], 0.0) - assert torch.equal(buffer[0], torch.ones(3, 2)) - assert torch.equal(buffer[1], torch.zeros(3, 2)) - assert torch.equal(buffer[2], torch.ones(3, 2)) - assert torch.equal(buffer[3], torch.zeros(3, 2)) +@pytest.mark.parametrize("explicit_ratio", [False, True]) +def test_sparse_arena_leases_follow_the_batch_rows(explicit_ratio: bool) -> None: + """The arena is the Batch's, so its rows alone decide how many leases exist.""" + torch.cuda.synchronize() + config = _sparse_arena_config(2, explicit_ratio) + core = KVCacheManager(config) + stream = CudaStream(torch.cuda.current_stream().cuda_stream) + try: + (shape,) = sparse_arena_shapes(config.layers, tokens_per_block=config.tokens_per_block) + for rows in (1, 2, 5): + batch = Batch( + core, + stream, + max_batch_size=rows, + max_blocks_per_seq=1, + index_tables=core.sparse_index_table_specs(), + ) + try: + assert batch.sparse_arena_bytes == shape.arena_bytes(rows) + finally: + del batch + finally: + core.shutdown() -def test_fill_kv_pages_packed_buffer(): - """A packed sub-byte pool is int8: zero stays zero, anything else poisons.""" - buffer = torch.ones(3, 4, dtype=torch.int8) - kv_cache_v2_module._fill_kv_pages(buffer, [0], 0.0) - kv_cache_v2_module._fill_kv_pages(buffer, [2], float("nan")) - assert torch.equal(buffer[0], torch.zeros(4, dtype=torch.int8)) - assert torch.equal(buffer[1], torch.ones(4, dtype=torch.int8)) - assert torch.equal(buffer[2], torch.full((4,), 0x7F, dtype=torch.int8)) +@pytest.mark.parametrize("max_batch_size,reserved", [(1, 1), (3, 2)]) +@pytest.mark.parametrize("explicit_ratio", [False, True]) +def test_sparse_arena_retains_padding_with_full_live_batch( + max_batch_size: int, reserved: int, explicit_ratio: bool +) -> None: + torch.cuda.synchronize() + config = _sparse_arena_config(max_batch_size, explicit_ratio) + core = KVCacheManager(config) + caches = [] + stream = CudaStream(torch.cuda.current_stream().cuda_stream) + batch = Batch( + core, + stream, + max_batch_size=max_batch_size, + max_blocks_per_seq=1, + index_tables=core.sparse_index_table_specs(), + ) + try: + # A lease follows membership, so padding caches outside the Batch hold none and a + # resumed request receiving its history costs the arena nothing. + lc_id = core.get_layer_group_id(0) + for index in range(reserved + max_batch_size): + cache = core.create_kv_cache() + caches.append(cache) + assert cache.resume(stream) + assert cache.resize(1) + assert _introspection.sparse_cache_base_cell(cache, lc_id) is None + if index >= reserved: + batch.add_kv_cache(cache) + assert _introspection.sparse_cache_base_cell(cache, lc_id) is not None + assert batch.max_batch_size == max_batch_size + batch.publish() + batch.wait_published() + + # The rows bound membership, not the arena: a free row always has a free lease, + # so a full Batch refuses on its rows rather than running out of cells. + extra = core.create_kv_cache() + caches.append(extra) + assert extra.resume(stream) + assert extra.resize(1) + with pytest.raises(Exception, match="full"): + batch.add_kv_cache(extra) + batch.remove_kv_cache(caches[-2]) + assert _introspection.sparse_cache_base_cell(caches[-2], lc_id) is None + batch.add_kv_cache(extra) + assert _introspection.sparse_cache_base_cell(extra, lc_id) is not None + finally: + for cache in caches: + if batch.row_of(cache) is not None: + batch.remove_kv_cache(cache) + cache.close() + core.shutdown() -def test_fill_kv_pages_empty_is_a_noop(): - buffer = torch.ones(2, 2) - assert kv_cache_v2_module._fill_kv_pages(buffer, [], 0.0) is True - assert torch.equal(buffer, torch.ones(2, 2)) +def test_sparse_arena_shape_matches_native_layout() -> None: + torch.cuda.synchronize() + leases = 5 + config = _sparse_arena_config(2, False) + core = KVCacheManager(config) + try: + (shape,) = sparse_arena_shapes(config.layers, tokens_per_block=config.tokens_per_block) + layer = LayerId(0) + units_per_page = core.get_sparse_cache_units_per_page(layer) + assert shape.units_per_cell == units_per_page + assert shape.cells_per_buffer == -( + -core.get_sparse_cache_max_units(layer) // units_per_page + ) + assert len(shape.cell_bytes) == core.get_sparse_cache_buffer_count(layer) + # Pool regions are consecutive, each holding every lease's cells for one pool. + regions = { + -(-(leases * shape.cells_per_buffer * cell_bytes) // SPARSE_ARENA_ALIGNMENT) + * SPARSE_ARENA_ALIGNMENT + for cell_bytes in shape.cell_bytes + } + batch = Batch( + core, + CudaStream(torch.cuda.current_stream().cuda_stream), + max_batch_size=leases, + max_blocks_per_seq=1, + index_tables=core.sparse_index_table_specs(), + ) + key_base = batch.get_sparse_cache_base_address(layer, Role.KEY) + scale_base = batch.get_sparse_cache_base_address(layer, Role.KEY_BLOCK_SCALE) + assert abs(scale_base - key_base) in regions + + # The sparse buffers' pools share one pool group with the dense buffer, and + # their slots hold one page of every sparse buffer: a cell's worth each. + sparse_buffers = {(0, str(Role.KEY)), (0, str(Role.KEY_BLOCK_SCALE))} + (group,) = core.pool_group_descs + coalesced_buffers = group.slot_desc.variants[0].coalesced_buffers + sparse_pool_bytes = [ + pool.slot_bytes + for pool in group.pools + if any( + (int(buffer_id.layer_id), str(buffer_id.role)) in sparse_buffers + for buffer_id in coalesced_buffers[int(pool.pool_index)].buffer_ids + ) + ] + assert sorted(sparse_pool_bytes) == sorted(shape.cell_bytes) + assert len(group.pools) == 3 + finally: + core.shutdown() -def test_fill_kv_pages_reports_applied(): - buffer = torch.ones(4, 2) - assert kv_cache_v2_module._fill_kv_pages(buffer, [1], 0.0) is True +# A dense buffer that sorts between the sparse ones leaves a pool index no sparse buffer +# uses; one that sorts last leaves none. The two exercise different arena layouts, and +# until both were checked a missing placeholder region and a phantom shared array cancelled. +@pytest.mark.parametrize("dense_size", [4096, 256], ids=["dense-between", "dense-last"]) +def test_sparse_arena_bytes_match_the_native_allocation(dense_size: int) -> None: + """Python predicts the arena's bytes; this is the only thing that checks the prediction. + The reservation has to run before a manager exists, so the bytes are modelled in + Python and compared against a figure Python also computed -- which cannot catch the + model drifting from the allocation. Reading the allocation back closes that loop. + """ + torch.cuda.synchronize() + leases = 5 + config = _sparse_arena_config(2, False, dense_size) + core = KVCacheManager(config) + try: + (shape,) = sparse_arena_shapes(config.layers, tokens_per_block=config.tokens_per_block) + batch = Batch( + core, + CudaStream(torch.cuda.current_stream().cuda_stream), + max_batch_size=leases, + max_blocks_per_seq=1, + index_tables=core.sparse_index_table_specs(), + ) + assert batch.sparse_arena_bytes == shape.arena_bytes(leases) + assert batch.sparse_arena_bytes > 0 + finally: + core.shutdown() -def test_fill_kv_pages_skips_dtype_without_index_fill(): - """FP8 has no ``index_fill_`` kernel; the fill is skipped, not fatal.""" + +@pytest.mark.parametrize("shared_selection", [False, True], ids=["apart", "shared"]) +def test_sparse_arena_bytes_follow_the_selection_count(shared_selection: bool) -> None: + """The metadata is sized by selections, so sharing one has to shrink the allocation. + + The same loop as above, closed over the dimension a selection group changes: the + Python model counts metadata rows itself, and nothing else would catch it counting + layers where the arena counts selections. + """ + torch.cuda.synchronize() + leases = 5 + config = _sparse_arena_config(2, False, num_layers=2, shared_selection=shared_selection) + core = KVCacheManager(config) try: - torch.zeros(2, dtype=torch.float8_e4m3fn).index_fill_(0, torch.tensor([0]), float("nan")) - except NotImplementedError: - pass - else: - pytest.skip("index_fill_ is supported for float8 on this platform") + (shape,) = sparse_arena_shapes(config.layers, tokens_per_block=config.tokens_per_block) + assert shape.num_layers == 2 + assert shape.num_selections == (1 if shared_selection else 2) + batch = Batch( + core, + CudaStream(torch.cuda.current_stream().cuda_stream), + max_batch_size=leases, + max_blocks_per_seq=1, + index_tables=core.sparse_index_table_specs(), + ) + assert batch.sparse_arena_bytes == shape.arena_bytes(leases) + finally: + core.shutdown() - buffer = torch.zeros(4, 2, dtype=torch.float8_e4m3fn) - assert kv_cache_v2_module._fill_kv_pages(buffer, [1], float("nan")) is False +def _native_config_with_sparse_arena( + *, + arena_bytes: int, + budget: int, + sparse_layers: bool = False, + batch_buffer_bytes: int = 0, + mapping: Mapping | None = None, + sparse_kv_offload_config: SparseKvCacheOffloadConfig | None = None, + kv_cache_config: KvCacheConfig | None = None, + host_history_budget: int = 0, +) -> KVCacheManagerConfig: + """Construct a manager that reserves ``arena_bytes`` for its arena and + ``batch_buffer_bytes`` for its offload buffers, and plans + ``host_history_budget`` bytes of host history; return the config it gives + the native manager.""" + module = "tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2" + layers = ( + _sparse_arena_config(2, False).layers + if sparse_layers + else [ + AttentionLayerConfig( + layer_id=LayerId(0), buffers=[BufferConfig(role=Role.KEY, size=128)] + ) + ] + ) + fake_impl = Mock() + fake_impl.layer_grouping = [[0]] + fake_impl.pool_group_descs = [] + fake_impl.get_layer_group_id.side_effect = lambda _: 0 -def test_guard_page_accessors_report_no_guard_by_default(): - """With the guard page off, backends see ``None`` and keep page 0.""" - manager = SimpleNamespace(_guard_page_by_layer={}) - assert KVCacheManagerV2.guard_page_index(manager, 0) is None - assert KVCacheManagerV2.guard_page_indices(manager) == frozenset() + def build_base_config( + self: KVCacheManagerV2, + config: KvCacheConfig, + *, + tokens_per_block: int, + cache_tiers: list[object], + ) -> KVCacheManagerConfig: + del self, config + return KVCacheManagerConfig( + tokens_per_block=tokens_per_block, cache_tiers=cache_tiers, layers=layers + ) + def plan_host_history_budget( + self: KVCacheManagerV2, kv_cache_config: KvCacheConfig, **kwargs: object + ) -> None: + # Settles the budget only; the manager divides the device quota against it. + self._host_history_budget = host_history_budget -def test_guard_page_accessors_report_reserved_pages(): - # Layers in one pool can carry different page-index scales, so one - # reservation can surface as more than one index. - manager = SimpleNamespace(_guard_page_by_layer={0: 7, 1: 14, 2: 7}) - assert KVCacheManagerV2.guard_page_index(manager, 1) == 14 - assert KVCacheManagerV2.guard_page_index(manager, 9) is None - assert KVCacheManagerV2.guard_page_indices(manager) == frozenset({7, 14}) + with ( + patch.dict(os.environ, {"TRTLLM_KV_GUARD_PAGE": "", "TRTLLM_KV_FRESH_PAGE_FILL": ""}), + patch(f"{module}.IndexMapper"), + patch(f"{module}.KVCacheManagerPy", return_value=fake_impl) as create_impl, + patch.object(KVCacheManagerV2, "_build_base_config", build_base_config), + patch.object(KVCacheManagerV2, "_build_cache_config", lambda self, config: config), + patch.object(KVCacheManagerV2, "get_num_available_tokens", return_value=MAX_SEQ_LEN), + patch.object(KVCacheManagerV2, "_prepare_page_table_tensor"), + patch.object(KVCacheManagerV2, "_log_kv_cache_pool_lifecycle_mapping"), + patch.object(KVCacheManagerV2, "_reserve_guard_page"), + patch.object( + KVCacheManagerV2, "_get_sparse_arena_bytes", return_value=arena_bytes + ) as get_arena_bytes, + patch.object( + KVCacheManagerV2, + "_get_sparse_batch_buffer_bytes", + return_value=(batch_buffer_bytes, 4096 if batch_buffer_bytes else 0), + ) as get_batch_buffer_bytes, + patch.object(KVCacheManagerV2, "_plan_host_history_budget", plan_host_history_budget), + ): + KVCacheManagerV2( + kv_cache_config or KvCacheConfig(max_gpu_total_bytes=budget, enable_block_reuse=False), + CacheType.SELFKONLY, + num_layers=1, + num_kv_heads=1, + head_dim=1, + tokens_per_block=TOKENS_PER_BLOCK, + max_seq_len=MAX_SEQ_LEN, + max_batch_size=2, + mapping=mapping or Mapping(world_size=1, rank=0, tp_size=1, pp_size=1), + dtype=DataType.HALF, + vocab_size=16, + execution_stream=Mock(), + is_disagg=True, + sparse_kv_offload_config=sparse_kv_offload_config, + ) + # A lease follows Batch membership, so the arena is sized by the rows rather than by + # the 5 indices the mapper can hand out -- disagg doubling and padding included. + get_arena_bytes.assert_called_once_with(ANY, tokens_per_block=TOKENS_PER_BLOCK, max_leases=2) + get_batch_buffer_bytes.assert_called_once_with( + ANY, tokens_per_block=TOKENS_PER_BLOCK, max_leases=2 + ) + return create_impl.call_args.args[0] -class _FreshFillKVCache: - """Just enough of a KV cache for ``_fill_fresh_kv_pages``: a mutable page - list standing in for the pool's base page indices.""" +def _device_quota_with_sparse_arena(**kwargs: object) -> int: + """Return the device quota of ``_native_config_with_sparse_arena``.""" + return _native_config_with_sparse_arena(**kwargs).cache_tiers[0].quota - def __init__(self, pages: list[int], num_committed_tokens: int = 0) -> None: - self.pages = list(pages) - self.num_committed_tokens = num_committed_tokens - @property - def num_blocks(self) -> int: - return len(self.pages) +@pytest.mark.cpu_only +@pytest.mark.parametrize( + ("arena_bytes", "batch_buffer_bytes"), [(0, 0), (3 << 20, 0), (3 << 20, 1 << 20)] +) +def test_device_quota_leaves_room_for_the_offload_allocations( + arena_bytes: int, batch_buffer_bytes: int +) -> None: + quota = _device_quota_with_sparse_arena( + arena_bytes=arena_bytes, budget=16 << 20, batch_buffer_bytes=batch_buffer_bytes + ) - def get_base_page_indices(self, pool_id: int) -> np.ndarray: - del pool_id - return np.asarray(self.pages, dtype=np.int32) + assert quota == (16 << 20) - arena_bytes - batch_buffer_bytes -_FRESH_FILL = 3.0 +@pytest.mark.cpu_only +@pytest.mark.parametrize( + ("arena_bytes", "batch_buffer_bytes"), [(16 << 20, 0), (12 << 20, 4 << 20)] +) +def test_offload_allocations_larger_than_the_budget_are_rejected( + arena_bytes: int, batch_buffer_bytes: int +) -> None: + with pytest.raises(ValueError, match="Sparse KV cache offload needs"): + _device_quota_with_sparse_arena( + arena_bytes=arena_bytes, budget=16 << 20, batch_buffer_bytes=batch_buffer_bytes + ) -def _make_fresh_fill_manager( - buffer: torch.Tensor, - kv_cache: _FreshFillKVCache, - monkeypatch: pytest.MonkeyPatch, - fill_value: float | None = _FRESH_FILL, -) -> KVCacheManagerV2: - # These tests use CPU tensors and an explicit stream double; a device-wide wait is forbidden. - monkeypatch.setattr( - torch.cuda, "synchronize", Mock(side_effect=AssertionError("device-wide fill sync")) - ) - monkeypatch.setattr(torch.cuda, "current_stream", Mock(return_value=Mock())) - manager = object.__new__(KVCacheManagerV2) - manager._stream = SimpleNamespace(device=torch.device("cuda:0")) - manager._fresh_page_fill = fill_value - manager.layer_specs = None - manager._fresh_pages_filled = {} - manager._fresh_fill_announced = True - manager.kv_cache_map = {1: kv_cache} - manager.tokens_per_block = TOKENS_PER_BLOCK - manager.num_pools = 1 - manager.pp_layers = [0] - manager.layer_offsets = {0: 0} - manager.layer_to_pool_mapping_dict = {0: 0} - manager.kv_factor = 1 - manager.get_buffers = lambda layer_idx: buffer - manager.get_layer_page_index_scale = lambda layer_idx: 1 - return manager +@pytest.mark.cpu_only +@pytest.mark.parametrize("max_tokens", [4 * TOKENS_PER_BLOCK, 1 << 30]) +def test_max_tokens_sizes_the_pools_beside_the_offload_allocations(max_tokens: int) -> None: + # The arena and buffers come out of the device memory budget. max_tokens sizes the + # pools alone, so a pool for a few tokens still fits beside a large arena. + def quota(arena_bytes: int, batch_buffer_bytes: int) -> int: + return _device_quota_with_sparse_arena( + arena_bytes=arena_bytes, + budget=16 << 20, + batch_buffer_bytes=batch_buffer_bytes, + kv_cache_config=KvCacheConfig( + max_gpu_total_bytes=16 << 20, max_tokens=max_tokens, enable_block_reuse=False + ), + ) + pools = quota(0, 0) + assert quota(12 << 20, 1 << 20) == min(pools, (16 << 20) - (13 << 20)) -def test_fresh_fill_covers_pages_on_allocation(monkeypatch: pytest.MonkeyPatch) -> None: + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + ("arena_bytes", "sizing_result", "error", "match"), + [ + # Rank 1 fits its budget, but rank 0 does not. + (0, 0, RuntimeError, "Sizing sparse KV cache offload failed on rank 0"), + # Rank 1 itself does not fit, and raises its own error. + (16 << 20, 1, ValueError, "Sparse KV cache offload needs"), + ], +) +def test_offload_sizing_fails_on_every_rank_before_the_token_sync( + arena_bytes: int, sizing_result: int, error: type, match: str +) -> None: + dist = Mock() + # The config check passes everywhere; then sizing reports the lowest failing rank. + dist.allreduce.side_effect = [2, sizing_result] + + with ( + patch.object(Distributed, "get", return_value=dist), + pytest.raises(error, match=match), + ): + _device_quota_with_sparse_arena( + arena_bytes=arena_bytes, + budget=16 << 20, + mapping=Mapping(world_size=2, rank=1, tp_size=2), + sparse_kv_offload_config=SparseKvCacheOffloadConfig(), + ) + + # Every rank stops after agreeing; none is left in the token sync. + this_rank = 2 if arena_bytes == 0 else 1 + assert dist.allreduce.call_args_list == [ + call(2, op=ReduceOp.MIN), + call(this_rank, op=ReduceOp.MIN), + ] + + +@pytest.mark.cpu_only +def test_automatic_host_tier_gets_what_offload_leaves(monkeypatch) -> None: + # The 8 GiB of history and 4 KiB of staging leave 10 MiB of the memlock limit. + monkeypatch.setattr( + kv_cache_v2_module, + "get_host_memory_limits", + lambda mapping: HostMemoryLimits(1, float(400 << 30), float((8 << 30) + 4096 + (10 << 20))), + ) + + with patch.object(kv_cache_v2_module.logger, "warning") as warning: + config = _native_config_with_sparse_arena( + arena_bytes=0, + budget=16 << 20, + batch_buffer_bytes=1 << 20, + sparse_kv_offload_config=SparseKvCacheOffloadConfig(), + host_history_budget=8 << 30, + ) + + _, host_tier = config.cache_tiers + assert isinstance(host_tier, HostCacheTierConfig) + assert host_tier.quota == int((10 << 20) * 0.8) + assert not any("cache tier" in c.args[0] for c in warning.call_args_list) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("disk_tier", [False, True]) +def test_offload_refuses_to_leave_no_room_for_the_automatic_host_tier( + monkeypatch, tmp_path, disk_tier: bool +) -> None: + # The offloaded history takes all of the memlock limit. + monkeypatch.setattr( + kv_cache_v2_module, + "get_host_memory_limits", + lambda mapping: HostMemoryLimits(1, float(400 << 30), float(8 << 30)), + ) + kwargs = dict( + arena_bytes=0, + budget=16 << 20, + sparse_kv_offload_config=SparseKvCacheOffloadConfig(), + host_history_budget=8 << 30, + ) + + if not disk_tier: + # Falling back to the device quota would pin past the limit. + with pytest.raises(ValueError, match="leaves none for the host cache tier"): + _native_config_with_sparse_arena(**kwargs) + return + # A disk tier still gives evicted pages somewhere to go. + config = _native_config_with_sparse_arena( + **kwargs, + kv_cache_config=KvCacheConfig( + max_gpu_total_bytes=16 << 20, + disk_cache_size=16 << 20, + disk_cache_path=str(tmp_path), + enable_block_reuse=False, + ), + ) + assert [type(tier) for tier in config.cache_tiers] == [ + GpuCacheTierConfig, + DiskCacheTierConfig, + ] + + +@pytest.mark.cpu_only +def test_offload_warns_when_it_runs_without_a_cold_tier() -> None: + with patch.object(kv_cache_v2_module.logger, "warning") as warning: + config = _native_config_with_sparse_arena( + arena_bytes=0, + budget=16 << 20, + sparse_kv_offload_config=SparseKvCacheOffloadConfig(), + kv_cache_config=KvCacheConfig( + max_gpu_total_bytes=16 << 20, host_cache_size=0, enable_block_reuse=False + ), + ) + + assert [type(tier) for tier in config.cache_tiers] == [GpuCacheTierConfig] + assert any( + "runs without a host or disk cache tier" in c.args[0] for c in warning.call_args_list + ) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + ("error", "lowest_failing_rank", "raised", "match"), + [ + (None, 4, None, None), + (None, 1, RuntimeError, "Checking failed on rank 1"), + (ValueError("no room"), 1, ValueError, "no room"), + ], +) +def test_raise_on_every_rank_agrees_across_ranks( + error: Exception | None, lowest_failing_rank: int, raised: type | None, match: str | None +) -> None: + mapping = SimpleNamespace(world_size=4, rank=2) + dist = Mock() + dist.allreduce.return_value = lowest_failing_rank + + with patch.object(Distributed, "get", return_value=dist): + if raised is None: + _raise_on_every_rank(error, mapping, "Checking") + else: + with pytest.raises(raised, match=match): + _raise_on_every_rank(error, mapping, "Checking") + + # A passing rank offers world_size, so the minimum names the lowest failing rank. + dist.allreduce.assert_called_once_with(2 if error else 4, op=ReduceOp.MIN) + + +@pytest.mark.cpu_only +def test_raise_on_every_rank_needs_no_collective_on_one_rank() -> None: + mapping = SimpleNamespace(world_size=1, rank=0) + + with patch.object(Distributed, "get") as get: + _raise_on_every_rank(None, mapping, "Checking") + with pytest.raises(ValueError, match="no room"): + _raise_on_every_rank(ValueError("no room"), mapping, "Checking") + + get.assert_not_called() + + +@pytest.mark.cpu_only +def test_sparse_arena_without_a_reservation_is_rejected() -> None: + # Cached sparse buffers declared without sparse_kv_offload_config reserve no + # device memory for their arena. + with pytest.raises(RuntimeError, match="sparse_kv_offload_config"): + _device_quota_with_sparse_arena(arena_bytes=0, budget=16 << 20, sparse_layers=True) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("offload", [False, True]) +def test_sparse_arena_bytes_come_from_a_provisional_config(offload: bool) -> None: + manager = object.__new__(KVCacheManagerV2) + if offload: + manager._sparse_kv_offload = SparseKvCacheOffloadConfig() + config = _sparse_arena_config(2, False) + + with patch.object(KVCacheManagerV2, "_build_manager_config", return_value=config) as build: + arena_bytes = manager._get_sparse_arena_bytes( + KvCacheConfig(), tokens_per_block=16, max_leases=5 + ) + + if not offload: + assert arena_bytes == 0 + build.assert_not_called() + return + assert arena_bytes == sparse_arena_bytes(config, 5) > 0 + # Layer configs do not depend on tier quotas, so any quota gives the same arena. + (gpu_tier,) = build.call_args.kwargs["cache_tiers"] + assert gpu_tier.quota == sys.maxsize + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("offload", [False, True]) +def test_sparse_batch_buffers_come_from_a_provisional_config(offload: bool) -> None: + manager = object.__new__(KVCacheManagerV2) + manager.max_seq_len = MAX_SEQ_LEN + manager.max_batch_size = 2 + manager._generation_kv_capacity_headroom = 1 + manager._ledger_tokens_per_block = TOKENS_PER_BLOCK + if offload: + manager._sparse_kv_offload = SparseKvCacheOffloadConfig() + config = KVCacheManagerConfig( + tokens_per_block=TOKENS_PER_BLOCK, + cache_tiers=[GpuCacheTierConfig(quota=sys.maxsize)], + layers=_offload_layers(), + ) + + with ( + patch.object(KVCacheManagerV2, "_build_manager_config", return_value=config) as build, + patch.object( + KVCacheManagerV2, "get_sparse_offload_scratch_bytes", return_value=1000 + ) as get_scratch_bytes, + ): + device_bytes, staging_bytes = manager._get_sparse_batch_buffer_bytes( + KvCacheConfig(), tokens_per_block=TOKENS_PER_BLOCK, max_leases=5 + ) + + if not offload: + assert (device_bytes, staging_bytes) == (0, 0) + build.assert_not_called() + return + # 17 tokens of 16 and the decode step's one span five pages, padded to eight. + batch = sparse_batch_shape( + config.layers, tokens_per_block=TOKENS_PER_BLOCK, max_batch_size=2, max_blocks_per_seq=8 + ) + assert batch.num_life_cycles == 2 + assert batch.selection_widths == (1,) + assert device_bytes == batch.block_bytes + batch.mirror_bytes + batch.fetch_scratch_bytes + 1000 + assert staging_bytes == batch.staging_bytes == 4096 + get_scratch_bytes.assert_called_once_with(config) + (gpu_tier,) = build.call_args.kwargs["cache_tiers"] + assert gpu_tier.quota == sys.maxsize + + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + ("max_seq_len", "headroom", "max_blocks_per_seq"), [(16, 1, 8), (15, 1, 4), (16, 0, 4)] +) +def test_page_table_holds_the_decode_step_and_pads_to_four_pages( + max_seq_len: int, headroom: int, max_blocks_per_seq: int +) -> None: + manager = object.__new__(KVCacheManagerV2) + manager._generation_kv_capacity_headroom = headroom + manager._ledger_tokens_per_block = TOKENS_PER_BLOCK + + assert manager._get_max_blocks_per_seq(max_seq_len) == max_blocks_per_seq + + +def _gen_init_manager(host_capacity_blocks: int | None) -> KVCacheManagerV2: + manager = object.__new__(KVCacheManagerV2) + if host_capacity_blocks is not None: + manager._host_history_reservations = HostHistoryReservations(host_capacity_blocks) + manager.num_extra_kv_tokens = 0 + manager.tokens_per_block = 128 + manager._on_allocation_stream = set() + manager._allocation_ready = {} + return manager + + +def _gen_init_request(request_id: int, prompt_len: int, max_new_tokens: int): + return SimpleNamespace( + py_request_id=request_id, + prompt_len=prompt_len, + orig_prompt_len=prompt_len, + max_new_tokens=max_new_tokens, + py_beam_width=1, + ) + + +@pytest.mark.cpu_only +def test_gen_init_reserves_host_history_to_completion() -> None: + manager = _gen_init_manager(host_capacity_blocks=10) + reservations = manager._host_history_reservations + + with patch.object( + KVCacheManagerV2, "_prepare_disagg_gen_init_cache", return_value=True + ) as prepare: + # 641 tokens through the last output token span six blocks. + assert manager.prepare_disagg_gen_init(_gen_init_request(1, 512, 129)) + assert reservations.reserved_blocks == 6 + # Five more blocks do not fit beside it, and the cache is left untouched. + assert not manager.prepare_disagg_gen_init(_gen_init_request(2, 512, 128)) + assert prepare.call_count == 1 + # Preparing a request again replaces its reservation. + assert manager.prepare_disagg_gen_init(_gen_init_request(1, 256, 0)) + assert manager.prepare_disagg_gen_init(_gen_init_request(2, 512, 128)) + assert reservations.reserved_blocks == 7 + # A request larger than the whole pool is never admitted. + assert not manager.prepare_disagg_gen_init(_gen_init_request(3, 128 * 11, 0)) + with patch.object(KVCacheManagerV2, "_prepare_disagg_gen_init_cache", return_value=False): + # A request whose cache could not be prepared reserves nothing. + assert not manager.prepare_disagg_gen_init(_gen_init_request(4, 128, 0)) + assert reservations.reserved_blocks == 7 + + +@pytest.mark.cpu_only +def test_gen_init_without_sparse_offload_reserves_nothing() -> None: + manager = _gen_init_manager(host_capacity_blocks=None) + + with patch.object( + KVCacheManagerV2, "_prepare_disagg_gen_init_cache", return_value=True + ) as prepare: + assert manager.prepare_disagg_gen_init(_gen_init_request(1, 1 << 20, 0)) + + prepare.assert_called_once() + assert not manager.kv_block_budget_applies() + assert manager.get_request_kv_block_budget(_gen_init_request(1, 1 << 20, 0)) is None + + +@pytest.mark.cpu_only +def test_request_larger_than_host_history_is_rejected_on_arrival() -> None: + manager = _gen_init_manager(host_capacity_blocks=10) + manager._host_history_capacity_blocks = 10 + executor = PyExecutor.__new__(PyExecutor) + executor.resource_manager = ResourceManager({ResourceManagerType.KV_CACHE_MANAGER: manager}) + + assert manager.kv_block_budget_applies() + # 1281 tokens through the last output token span 11 blocks, which no rank can + # reserve, so the request fails instead of waiting for admission forever. + with pytest.raises(ValueError, match="requires 11 KV cache blocks.*capacity of 10 blocks"): + executor._validate_request_budget(_gen_init_request(1, 1024, 257)) + # Ten blocks fit, and admission reserves exactly what the check counted. + request = _gen_init_request(2, 1024, 256) + executor._validate_request_budget(request) + with patch.object(KVCacheManagerV2, "_prepare_disagg_gen_init_cache", return_value=True): + assert manager.prepare_disagg_gen_init(request) + assert manager._host_history_reservations.reserved_blocks == 10 + + +def _offload_setup_manager( + local_slots: list[int], + *, + world_size: int = 1, + attention_dp: bool = False, + guard_page: bool = False, + max_total_draft_tokens: int = 0, +) -> KVCacheManagerV2: + """A manager stand-in whose host-backed pool groups hold ``local_slots`` slots.""" + manager = object.__new__(KVCacheManagerV2) + manager.mapping = SimpleNamespace( + world_size=world_size, rank=0, enable_attention_dp=attention_dp + ) + manager.is_estimating_kv_cache = False + manager.num_extra_kv_tokens = 0 + manager.max_total_draft_tokens = max_total_draft_tokens + manager.tokens_per_block = 128 + manager._host_history_budget = 1 << 40 + manager.kv_cache_map = {kv_cache_v2_module._GUARD_PAGE_REQUEST_ID: Mock()} if guard_page else {} + manager.get_host_backed_pool_groups = lambda: [ + SimpleNamespace(num_slots=n) for n in local_slots + ] + manager._get_host_backed_pool_bytes = lambda: 0 + return manager + + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + ("world_size", "local_slots", "peer_slots", "capacity_blocks"), + [ + (1, [12, 10], None, 10), + # Every rank judges arriving requests by the smallest pool, including a + # rank whose layers hold no host history, but reserves from its own. + (2, [12, 10], 8, 8), + (2, [], 8, 8), + (2, [], sys.maxsize, None), + ], +) +def test_sparse_offload_setup_judges_requests_by_the_smallest_host_pool( + world_size: int, local_slots: list[int], peer_slots: int | None, capacity_blocks: int | None +) -> None: + manager = _offload_setup_manager(local_slots, world_size=world_size) + dist = Mock() + dist.allreduce.side_effect = lambda value, op: min(value, peer_slots) + + with ( + patch.object(kv_cache_v2_module.Distributed, "get", return_value=dist), + patch.object(KVCacheManagerV2, "_log_sparse_kv_offload_capacity"), + patch.object(kv_cache_v2_module.logger, "warning") as warning, + ): + manager._setup_sparse_kv_offload( + KvCacheConfig(), max_seq_len=1024, num_reserved_index_slots=0 + ) + + if world_size > 1: + # The ranks agree that their pools fit, then on the smallest pool. + assert dist.allreduce.call_args_list == [ + call(world_size, op=ReduceOp.MIN), + call(min(local_slots, default=sys.maxsize), op=ReduceOp.MIN), + ] + else: + dist.allreduce.assert_not_called() + assert manager._host_history_capacity_blocks == capacity_blocks + assert manager.kv_block_budget_applies() == (capacity_blocks is not None) + if local_slots: + assert manager._host_history_reservations.capacity_blocks == min(local_slots) + else: + assert manager._host_history_reservations is None + # Eight blocks of 128 tokens hold a max_seq_len request of 1024 tokens. + warning.assert_not_called() + + +@pytest.mark.cpu_only +def test_sparse_offload_setup_warns_when_max_seq_len_outgrows_host_history() -> None: + manager = _offload_setup_manager([8]) + + with ( + patch.object(KVCacheManagerV2, "_log_sparse_kv_offload_capacity"), + patch.object(kv_cache_v2_module.logger, "warning") as warning, + ): + manager._setup_sparse_kv_offload( + KvCacheConfig(), max_seq_len=1025, num_reserved_index_slots=0 + ) + + message = warning.call_args.args[0] + assert "rejects requests whose prompt and max_tokens exceed 1024 tokens" in message + assert "max_seq_len=1025 tokens needs 9" in message + assert "kv_cache_offload_config.host_history_size" in message + + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + ("world_size", "host_pool_bytes", "lowest_failing_rank", "error", "match"), + [ + # Only the native slot floors can take the pools past the budget. + (1, 2 << 30, None, ValueError, "take 2.00GiB, more than their 1.00GiB host budget"), + # This rank's pools are too large; it raises its own error once every + # rank has agreed. + (2, 2 << 30, 1, ValueError, "take 2.00GiB, more than their 1.00GiB host budget"), + # Rank 0's pools are too large; this rank's fit. + (2, 1 << 30, 0, RuntimeError, "host-backed history pools failed on rank 0"), + ], +) +def test_sparse_offload_setup_rejects_host_pools_over_their_budget( + world_size: int, + host_pool_bytes: int, + lowest_failing_rank: int | None, + error: type, + match: str, +) -> None: + manager = _offload_setup_manager([8], world_size=world_size) + manager.mapping.rank = world_size - 1 + manager._host_history_budget = 1 << 30 + manager._get_host_backed_pool_bytes = lambda: host_pool_bytes + dist = Mock() + dist.allreduce.return_value = lowest_failing_rank + + with ( + patch.object(kv_cache_v2_module.Distributed, "get", return_value=dist), + pytest.raises(error, match=match), + ): + manager._setup_sparse_kv_offload( + KvCacheConfig(), max_seq_len=1024, num_reserved_index_slots=0 + ) + + # No rank reaches the capacity sync, and none starts admitting requests. + if world_size > 1: + this_rank = world_size - 1 if host_pool_bytes > 1 << 30 else world_size + dist.allreduce.assert_called_once_with(this_rank, op=ReduceOp.MIN) + else: + dist.allreduce.assert_not_called() + assert manager._host_history_reservations is None + + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + ("num_reserved_index_slots", "attention_dp", "guard_page", "draft_tokens", "set_aside"), + [ + # The CUDA-graph padding dummy. + (1, False, False, 0, 1), + # Plus the attention-DP dummy and the guard page. + (1, True, False, 0, 2), + (1, True, True, 0, 3), + # A dummy with 64 draft tokens holds 1 + 2 x 64 tokens, two blocks of 128. + (2, False, False, 64, 4), + ], +) +def test_sparse_offload_setup_sets_dummy_history_aside( + num_reserved_index_slots: int, + attention_dp: bool, + guard_page: bool, + draft_tokens: int, + set_aside: int, +) -> None: + manager = _offload_setup_manager( + [12, 10], + attention_dp=attention_dp, + guard_page=guard_page, + max_total_draft_tokens=draft_tokens, + ) + + with patch.object(KVCacheManagerV2, "_log_sparse_kv_offload_capacity"): + manager._setup_sparse_kv_offload( + KvCacheConfig(), max_seq_len=128, num_reserved_index_slots=num_reserved_index_slots + ) + + # Dummy requests hold history pages without a reservation, so neither + # admission nor the check on arrival may promise those pages to a request. + assert manager._host_history_reservations.capacity_blocks == 10 - set_aside + assert manager._host_history_capacity_blocks == 10 - set_aside + whole_pool = _gen_init_request(1, 128 * 10, 0) + assert manager.get_request_kv_block_budget(whole_pool) == (10, 10 - set_aside) + + +@pytest.mark.cpu_only +def test_free_resources_releases_host_history() -> None: + manager = _gen_init_manager(host_capacity_blocks=10) + manager._host_history_reservations.reserve(1, 6) + manager.kv_connector_manager = None + manager.conversation_manager = None + manager._allocated_draft_lens = {} + manager._inflight_draft_lens = {} + manager._request_stats_enabled_ids = set() + manager._fresh_pages_filled = {} + manager._disagg_receive_ready = {} + manager._pending_preemption = {} + manager._on_allocation_stream = set() + manager._allocation_ready = {} + manager.kv_cache_map = {} + manager.impl = Mock() + + manager.free_resources(SimpleNamespace(py_request_id=1, request_id=1)) + + assert manager._host_history_reservations.reserved_blocks == 0 + + +@pytest.mark.cpu_only +def test_host_memory_limits_cap_automatic_pinning() -> None: + assert HostMemoryLimits(4, float(400 << 30), float("inf")).auto_pinnable_bytes == 50 << 30 + assert HostMemoryLimits(1, float("inf"), float(10 << 30)).auto_pinnable_bytes == 8 << 30 + assert HostMemoryLimits(1, float("inf"), float("inf")).auto_pinnable_bytes == float("inf") + + +@pytest.mark.cpu_only +def test_sparse_kv_offload_config_is_checked_against_this_rank(monkeypatch) -> None: + limits = {"value": HostMemoryLimits(4, float(400 << 30), float("inf"))} + monkeypatch.setattr( + kv_cache_v2_module, "get_host_memory_limits", lambda mapping: limits["value"] + ) + + def check(offload_config, kv_cache_config, mapping): + KVCacheManagerV2._check_sparse_kv_offload_config( + offload_config, kv_cache_config, mapping, spec_config=None + ) + + kv_cache_config = KvCacheConfig(enable_block_reuse=False) + + check(SparseKvCacheOffloadConfig(host_history_size=50 << 30), kv_cache_config, Mapping()) + # A quarter of 400 GiB is this rank's share. + with pytest.raises(ValueError, match="share of available host memory"): + check(SparseKvCacheOffloadConfig(host_history_size=101 << 30), kv_cache_config, Mapping()) + # An explicit host cache tier is pinned beside the history. + check( + SparseKvCacheOffloadConfig(host_history_size=50 << 30), + KvCacheConfig(enable_block_reuse=False, host_cache_size=50 << 30), + Mapping(), + ) + with pytest.raises( + ValueError, match=r"plus kv_cache_config.host_cache_size \(51.00GiB\) exceeds this rank" + ): + check( + SparseKvCacheOffloadConfig(host_history_size=50 << 30), + KvCacheConfig(enable_block_reuse=False, host_cache_size=51 << 30), + Mapping(), + ) + limits["value"] = HostMemoryLimits(1, float("inf"), float(10 << 30)) + with pytest.raises(ValueError, match="RLIMIT_MEMLOCK"): + check(SparseKvCacheOffloadConfig(host_history_size=50 << 30), kv_cache_config, Mapping()) + with pytest.raises(ValueError, match=r"plus kv_cache_config.host_cache_size .* RLIMIT_MEMLOCK"): + check( + SparseKvCacheOffloadConfig(host_history_size=5 << 30), + KvCacheConfig(enable_block_reuse=False, host_cache_size=6 << 30), + Mapping(), + ) + with pytest.raises(ValueError, match="enable_kv_pool_rebalance"): + check( + SparseKvCacheOffloadConfig(), + KvCacheConfig(enable_block_reuse=False, enable_kv_pool_rebalance=True), + Mapping(), + ) + + +@pytest.mark.cpu_only +def test_sparse_kv_offload_rejects_speculative_decoding() -> None: + # A sparse cache takes one input token per decode step, and the history set + # aside for dummy requests counts a single CUDA-graph padding dummy. + with pytest.raises(ValueError, match="does not support speculative decoding"): + KVCacheManagerV2._check_sparse_kv_offload_config( + SparseKvCacheOffloadConfig(), + KvCacheConfig(enable_block_reuse=False), + Mapping(), + spec_config=MTPDecodingConfig(max_draft_len=1), + ) + + +@pytest.mark.cpu_only +def test_sparse_kv_offload_rejects_block_reuse() -> None: + # Reuse is on by default, so the check runs before any pool or arena is sized. + with pytest.raises(ValueError, match="does not support block reuse"): + KVCacheManagerV2._check_sparse_kv_offload_config( + SparseKvCacheOffloadConfig(), KvCacheConfig(), Mapping(), spec_config=None + ) + + +@pytest.mark.cpu_only +def test_sparse_kv_offload_config_check_fails_on_every_rank(monkeypatch) -> None: + # This rank's node has room, but rank 0's does not. + monkeypatch.setattr( + kv_cache_v2_module, + "get_host_memory_limits", + lambda mapping: HostMemoryLimits(4, float(400 << 30), float("inf")), + ) + dist = Mock() + dist.allreduce.return_value = 0 + + with ( + patch.object(Distributed, "get", return_value=dist), + pytest.raises(RuntimeError, match="config check failed on rank 0"), + ): + KVCacheManagerV2._check_sparse_kv_offload_config( + SparseKvCacheOffloadConfig(host_history_size=50 << 30), + KvCacheConfig(enable_block_reuse=False), + Mapping(world_size=2, rank=1, tp_size=2), + spec_config=None, + ) + + dist.allreduce.assert_called_once_with(2, op=ReduceOp.MIN) + + +def _offload_layers() -> list[AttentionLayerConfig]: + """A layer whose 1 KiB sparse pages sit beside 256-byte dense ones, and a windowed layer.""" + return [ + AttentionLayerConfig( + layer_id=LayerId(0), + buffers=[ + BufferConfig(role=Role.KEY, size=1024, is_sparse=True), + BufferConfig(role=Role.VALUE, size=256), + ], + sparse_cache=SparseCacheConfig( + num_selected_blocks=1, max_cached_units=2, unit_tokens=TOKENS_PER_BLOCK + ), + ), + AttentionLayerConfig( + layer_id=LayerId(1), + buffers=[BufferConfig(role=Role.KEY, size=4096)], + sliding_window_size=8, + ), + ] + + +def _offload_manager(host_history_size: int | None = None) -> KVCacheManagerV2: + manager = object.__new__(KVCacheManagerV2) + manager._sparse_kv_offload = SparseKvCacheOffloadConfig(host_history_size=host_history_size) + manager.max_seq_len = MAX_SEQ_LEN + manager.max_batch_size = 2 + manager.mapping = Mapping() + manager._generation_kv_capacity_headroom = 1 + return manager + + +def _base_config(layers: list, initial_pool_ratio: list[float] | None = None): + config = KVCacheManagerConfig( + tokens_per_block=TOKENS_PER_BLOCK, + cache_tiers=[GpuCacheTierConfig(quota=1 << 20)], + layers=layers, + initial_pool_ratio=initial_pool_ratio, + ) + return ( + patch.object(KVCacheManagerV2, "_build_base_config", return_value=config), + patch.object(KVCacheManagerV2, "_build_cache_config", lambda self, config: config), + ) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("pinned", [False, True]) +def test_offload_splits_the_ratio_by_all_bytes_unless_pinned(pinned: bool) -> None: + manager = _offload_manager() + build_base, build_cache = _base_config(_offload_layers(), [0.25, 0.75] if pinned else None) + + with build_base, build_cache: + config = manager._build_manager_config( + KvCacheConfig(), + tokens_per_block=TOKENS_PER_BLOCK, + cache_tiers=[GpuCacheTierConfig(quota=1 << 20)], + ) + + if pinned: + assert config.initial_pool_ratio == [0.25, 0.75] + return + # A 16-token request holds four pages in the no-window life cycle and three pages of + # 4 KiB in the window-8 one. A no-window page is 1280 bytes: 1024 host-backed sparse + # beside 256 device. The ratio weighs both media -- it is projected onto cold levels, + # where one pool holds them together, and compared against a layout measured the same + # way. That host bytes cost the GPU quota nothing is a fact about the quota alone. + no_window = 4 * (1024 + 256) + windowed = 3 * 4096 + total = no_window + windowed + assert config.initial_pool_ratio == pytest.approx([no_window / total, windowed / total]) + + +@pytest.mark.cpu_only +def test_host_backed_bytes_count_only_the_sparse_pools() -> None: + manager = object.__new__(KVCacheManagerV2) + manager.kv_cache_manager_py_config = KVCacheManagerConfig( + tokens_per_block=TOKENS_PER_BLOCK, + cache_tiers=[GpuCacheTierConfig(quota=1 << 20)], + layers=_offload_layers(), + ) + + def group(num_slots: int, *pools: tuple[int, int, str]) -> SimpleNamespace: + return SimpleNamespace( + num_slots=num_slots, + slot_desc=SimpleNamespace( + variants=[ + SimpleNamespace( + coalesced_buffers=[ + SimpleNamespace( + buffer_ids=[SimpleNamespace(layer_id=LayerId(layer), role=role)] + ) + for layer, _, role in pools + ] + ) + ] + ), + pools=[ + SimpleNamespace(pool_index=index, slot_bytes=slot_bytes) + for index, (_, slot_bytes, _) in enumerate(pools) + ], + ) + + mixed = group(10, (0, 1024, Role.KEY), (0, 256, Role.VALUE)) + dense = group(7, (1, 4096, Role.KEY)) + manager.impl = SimpleNamespace(pool_group_descs=[mixed, dense]) + + assert manager.get_host_backed_pool_groups() == [mixed] + # The mixed group's dense pool is GPU memory; only its sparse pool is host memory. + assert manager._get_host_backed_pool_bytes() == 10 * 1024 + + +@pytest.mark.cpu_only +def test_offload_hands_its_host_budget_to_the_manager(monkeypatch) -> None: + """The device quota is not lowered to fit the host budget; the manager divides both. + + Bounding the quota beforehand shrank every pool group, device-only ones included, to + hold back host memory only some of them spend. The manager is given the budget instead + and holds the host-backed pools inside it while the rest keep the device bytes. + """ + + module = "tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2" + fake_impl = Mock() + fake_impl.layer_grouping = [[0]] + fake_impl.pool_group_descs = [] + fake_impl.get_layer_group_id.side_effect = lambda _: 0 + + def build_base_config( + self: KVCacheManagerV2, + config: KvCacheConfig, + *, + tokens_per_block: int, + cache_tiers: list[object], + ) -> KVCacheManagerConfig: + del self, config + return KVCacheManagerConfig( + tokens_per_block=tokens_per_block, + cache_tiers=cache_tiers, + layers=_offload_layers()[:1], + ) + + # One rank on a node with 250 MiB of host memory available. + monkeypatch.setattr( + kv_cache_v2_module, + "get_host_memory_limits", + lambda mapping: HostMemoryLimits(1, float(250 << 20), float("inf")), + ) + with ( + patch.dict(os.environ, {"TRTLLM_KV_GUARD_PAGE": "", "TRTLLM_KV_FRESH_PAGE_FILL": ""}), + patch(f"{module}.IndexMapper"), + patch(f"{module}.KVCacheManagerPy", return_value=fake_impl) as create_impl, + patch(f"{module}.is_device_integrated", return_value=False), + patch.object(KVCacheManagerV2, "_build_base_config", build_base_config), + patch.object(KVCacheManagerV2, "_build_cache_config", lambda self, config: config), + patch.object(KVCacheManagerV2, "get_num_available_tokens", return_value=MAX_SEQ_LEN), + patch.object(KVCacheManagerV2, "_prepare_page_table_tensor"), + patch.object(KVCacheManagerV2, "_log_kv_cache_pool_lifecycle_mapping"), + patch.object(KVCacheManagerV2, "_reserve_guard_page"), + ): + KVCacheManagerV2( + KvCacheConfig(max_gpu_total_bytes=64 << 20, enable_block_reuse=False), + CacheType.SELFKONLY, + num_layers=1, + num_kv_heads=1, + head_dim=1, + tokens_per_block=TOKENS_PER_BLOCK, + max_seq_len=MAX_SEQ_LEN, + max_batch_size=2, + mapping=Mapping(world_size=1, rank=0, tp_size=1, pp_size=1), + dtype=DataType.HALF, + vocab_size=16, + execution_stream=Mock(), + is_disagg=True, + sparse_kv_offload_config=SparseKvCacheOffloadConfig(host_history_size=200 << 20), + ) + + config = create_impl.call_args.args[0] + gpu_tier, host_tier = config.cache_tiers + assert config.initial_pool_ratio == [1.0] + # Only the offload allocations come off the 64 MiB budget now. Under the old + # bound this was about 46 MiB, the host budget divided by the 4 host bytes each + # device byte carries -- a reduction every pool group paid, not just the ones + # holding host-backed pools. + assert gpu_tier.quota > (63 << 20) + # The 200 MiB of host history less the sparse Batch's one-page staging copy. No + # granule holdback: the manager counts host bytes in their own grain, so there is + # no device-side rounding to guess at from out here. + assert config.host_sparse_pool_quota == (200 << 20) - 4096 + # The automatic host tier still takes half of what the 200 MiB leaves of the 250 MiB. + assert host_tier.quota == (50 << 20) // 2 + + +@pytest.mark.parametrize( + "setting, expected", + [ + ("", None), + ("1", 0.0), + ("on", 0.0), + ("true", 0.0), + ("zero", 0.0), + ("0", 0.0), + ("-2.5", -2.5), + ("garbage", None), + ], +) +def test_kv_fill_value_parsing(setting, expected): + """The knobs are off when unset and ignore anything they cannot parse.""" + assert kv_cache_v2_module._parse_kv_fill_value(setting, "TEST_VAR") == expected + + +def test_kv_fill_value_parsing_nan(): + value = kv_cache_v2_module._parse_kv_fill_value("nan", "TEST_VAR") + assert value is not None and np.isnan(value) + + +def test_fill_kv_pages_float_buffer(): + buffer = torch.ones(4, 3, 2) + kv_cache_v2_module._fill_kv_pages(buffer, [1, 3], 0.0) + assert torch.equal(buffer[0], torch.ones(3, 2)) + assert torch.equal(buffer[1], torch.zeros(3, 2)) + assert torch.equal(buffer[2], torch.ones(3, 2)) + assert torch.equal(buffer[3], torch.zeros(3, 2)) + + +def test_fill_kv_pages_packed_buffer(): + """A packed sub-byte pool is int8: zero stays zero, anything else poisons.""" + buffer = torch.ones(3, 4, dtype=torch.int8) + kv_cache_v2_module._fill_kv_pages(buffer, [0], 0.0) + kv_cache_v2_module._fill_kv_pages(buffer, [2], float("nan")) + assert torch.equal(buffer[0], torch.zeros(4, dtype=torch.int8)) + assert torch.equal(buffer[1], torch.ones(4, dtype=torch.int8)) + assert torch.equal(buffer[2], torch.full((4,), 0x7F, dtype=torch.int8)) + + +def test_fill_kv_pages_empty_is_a_noop(): + buffer = torch.ones(2, 2) + assert kv_cache_v2_module._fill_kv_pages(buffer, [], 0.0) is True + assert torch.equal(buffer, torch.ones(2, 2)) + + +def test_fill_kv_pages_reports_applied(): + buffer = torch.ones(4, 2) + assert kv_cache_v2_module._fill_kv_pages(buffer, [1], 0.0) is True + + +def test_fill_kv_pages_skips_dtype_without_index_fill(): + """FP8 has no ``index_fill_`` kernel; the fill is skipped, not fatal.""" + try: + torch.zeros(2, dtype=torch.float8_e4m3fn).index_fill_(0, torch.tensor([0]), float("nan")) + except NotImplementedError: + pass + else: + pytest.skip("index_fill_ is supported for float8 on this platform") + + buffer = torch.zeros(4, 2, dtype=torch.float8_e4m3fn) + assert kv_cache_v2_module._fill_kv_pages(buffer, [1], float("nan")) is False + + +def test_guard_page_accessors_report_no_guard_by_default(): + """With the guard page off, backends see ``None`` and keep page 0.""" + manager = SimpleNamespace(_guard_page_by_layer={}) + assert KVCacheManagerV2.guard_page_index(manager, 0) is None + assert KVCacheManagerV2.guard_page_indices(manager) == frozenset() + + +def test_guard_page_accessors_report_reserved_pages(): + # Layers in one pool can carry different page-index scales, so one + # reservation can surface as more than one index. + manager = SimpleNamespace(_guard_page_by_layer={0: 7, 1: 14, 2: 7}) + assert KVCacheManagerV2.guard_page_index(manager, 1) == 14 + assert KVCacheManagerV2.guard_page_index(manager, 9) is None + assert KVCacheManagerV2.guard_page_indices(manager) == frozenset({7, 14}) + + +class _FreshFillKVCache: + """Just enough of a KV cache for ``_fill_fresh_kv_pages``: a mutable page + list standing in for the pool's base page indices.""" + + def __init__(self, pages: list[int], num_committed_tokens: int = 0) -> None: + self.pages = list(pages) + self.num_committed_tokens = num_committed_tokens + + @property + def num_blocks(self) -> int: + return len(self.pages) + + def get_base_page_indices(self, pool_id: int) -> np.ndarray: + del pool_id + return np.asarray(self.pages, dtype=np.int32) + + +_FRESH_FILL = 3.0 + + +def _make_fresh_fill_manager( + buffer: torch.Tensor, + kv_cache: _FreshFillKVCache, + monkeypatch: pytest.MonkeyPatch, + fill_value: float | None = _FRESH_FILL, +) -> KVCacheManagerV2: + # These tests use CPU tensors and an explicit stream double; a device-wide wait is forbidden. + monkeypatch.setattr( + torch.cuda, "synchronize", Mock(side_effect=AssertionError("device-wide fill sync")) + ) + monkeypatch.setattr(torch.cuda, "current_stream", Mock(return_value=Mock())) + manager = object.__new__(KVCacheManagerV2) + manager._stream = SimpleNamespace(device=torch.device("cuda:0")) + manager._fresh_page_fill = fill_value + manager.layer_specs = None + manager._fresh_pages_filled = {} + manager._fresh_fill_announced = True + manager.kv_cache_map = {1: kv_cache} + manager.tokens_per_block = TOKENS_PER_BLOCK + manager.num_pools = 1 + manager.pp_layers = [0] + manager.layer_offsets = {0: 0} + manager.layer_to_pool_mapping_dict = {0: 0} + manager.kv_factor = 1 + manager.get_buffers = lambda layer_idx: buffer + manager.get_layer_page_index_scale = lambda layer_idx: 1 + return manager + + +def test_fresh_fill_covers_pages_on_allocation(monkeypatch: pytest.MonkeyPatch) -> None: """Newly handed-out pages get the pattern; pages the request keeps do not.""" buffer = torch.zeros(8, 2) kv_cache = _FreshFillKVCache([2, 5]) @@ -3685,6 +4975,29 @@ def test_warmup_zeroing_preserves_guard_page(monkeypatch: pytest.MonkeyPatch) -> assert torch.equal(buffer[5], torch.zeros(2)) +@pytest.mark.cpu_only +@pytest.mark.parametrize("offload", [False, True]) +def test_warmup_scan_leaves_host_backed_history_alone(offload: bool) -> None: + manager = object.__new__(KVCacheManagerV2) + manager.layer_offsets = {0: 0, 1: 1} + manager.layer_to_pool_mapping_dict = {0: 0, 1: 1} + manager.layer_specs = None + manager.get_buffers = lambda layer_idx, kv_layout="NHD": torch.ones(2) + # Layer 0's KEY pages are offloaded history; layer 1's are device memory. + manager.kv_cache_manager_py_config = KVCacheManagerConfig( + tokens_per_block=TOKENS_PER_BLOCK, + cache_tiers=[GpuCacheTierConfig(quota=1 << 20)], + layers=_offload_layers(), + ) + if offload: + manager._sparse_kv_offload = SparseKvCacheOffloadConfig() + + scanned = [layer_id for layer_id, _ in manager._iter_cache_buffers_for_invalid_check()] + + # Nothing reads host history before a transfer or a write-back fills it. + assert scanned == ([1] if offload else [0, 1]) + + class _FakeLogger: """Stand-in for tensorrt_llm.logger with a settable level and captured messages.""" diff --git a/tests/unittest/_torch/executor/kv_cache/test_kvv2_host_tier_sizing.py b/tests/unittest/_torch/executor/kv_cache/test_kvv2_host_tier_sizing.py index 174886442d7f..2855e8413b0b 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kvv2_host_tier_sizing.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kvv2_host_tier_sizing.py @@ -123,6 +123,22 @@ def test_non_positive_result_falls_back_to_device_quota(self, memlock_limit): == 173 * GiB ) + @pytest.mark.parametrize("memlock_limit", [0.0, float(1)]) + def test_offloaded_history_without_headroom_provisions_no_tier(self, memlock_limit): + # Sparse KV cache offload has already pinned what the limits left, so the + # device quota would pin past them; the manager decides what to do + # without a tier instead. + assert ( + _compute_auto_host_tier_quota( + quota=173 * GiB, + local_ranks=4, + mem_available=float(440 * GiB), + memlock_limit=memlock_limit, + pins_offloaded_history=True, + ) + == 0 + ) + def test_result_is_always_positive(self): # Exhausted node memory reading must not produce a non-positive tier. assert ( diff --git a/tests/unittest/_torch/executor/kv_cache/test_mixed_precision_kv_cache.py b/tests/unittest/_torch/executor/kv_cache/test_mixed_precision_kv_cache.py index f1ba0fb27ddc..3b24cabae9c2 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_mixed_precision_kv_cache.py +++ b/tests/unittest/_torch/executor/kv_cache/test_mixed_precision_kv_cache.py @@ -238,7 +238,7 @@ def test_native_common_reuse_and_window_retention(make_manager, policy, reused_t assert source.resume(stream) and source.resize(len(tokens)) source.commit(tokens) source.suspend() - assert source.resume(stream, is_decoding=True) + assert source.resume(stream) pages = list(source.get_base_page_indices(high_group)) assert pages[0] >= 0 and pages[4] >= 0 assert pages[1:4] == [kv.BAD_PAGE_INDEX] * 3 diff --git a/tests/unittest/_torch/executor/kv_cache/test_sparse_offload_memory.py b/tests/unittest/_torch/executor/kv_cache/test_sparse_offload_memory.py new file mode 100644 index 000000000000..5ab7f2574d50 --- /dev/null +++ b/tests/unittest/_torch/executor/kv_cache/test_sparse_offload_memory.py @@ -0,0 +1,767 @@ +# 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. + +import pytest + +from tensorrt_llm._torch.pyexecutor.kv_cache.sparse_offload_memory import ( + HostHistoryReservations, + LifeCyclePage, + SparseArenaShape, + SparseBatchShape, + life_cycle_pages, + plan_host_history_bytes, + solve_sparse_offload_capacity, + sparse_arena_bytes, + sparse_arena_shapes, + sparse_batch_shape, + sparse_cache_cells, + sparse_cache_config, + sparse_history_bytes_per_block, + sparse_offload_pool_ratio, + sparse_units_per_page, +) +from tensorrt_llm.math_utils import ceil_div +from tensorrt_llm.runtime.kv_cache_manager_v2 import ( + AttentionLayerConfig, + BatchDesc, + BufferConfig, + CacheTier, + DataRole, + GpuCacheTierConfig, + KVCacheDesc, + KVCacheManagerConfig, + LayerId, + SparseCacheConfig, + SsmLayerConfig, + _introspection, +) + +TOKENS_PER_BLOCK = 128 +# A V4-Flash NVFP4 compressed row covers four tokens: 288 data bytes and 36 scale bytes. +V4_FLASH_HOST_BYTES_PER_BLOCK = 21 * 32 * (288 + 36) +V4_FLASH_ARENA = SparseArenaShape( + num_layers=21, + num_selections=21, + cells_per_buffer=33, + units_per_cell=32, + cell_bytes=(32 * 288, 32 * 36) * 21, +) +# 256 rows of up to 1M tokens: a table of 8,196 pages for each of three life +# cycles, and one fetch of 512 rows. +V4_FLASH_BATCH = SparseBatchShape( + num_life_cycles=3, + max_batch_size=256, + max_blocks_per_seq=8196, + selection_widths=(512,), + layer_counts=(2,), + flush_capacities=(256,), +) +# The Batch and DeepSeek-V4's attention metadata offload state, of which the eager +# and the CUDA graph metadata each hold one: 64 MiB of dequantized NVFP4 rows and +# 1.5 MiB of indices. +V4_FLASH_BATCH_BUFFER_BYTES = V4_FLASH_BATCH.device_bytes + 2 * 68_685_828 + + +def _layer( + layer_id: int, + *buffers: tuple[str, int, bool], + window: int | None = None, + sparse_cache: SparseCacheConfig | None = None, +) -> AttentionLayerConfig: + return AttentionLayerConfig( + layer_id=LayerId(layer_id), + buffers=[ + BufferConfig(role=DataRole(role), size=size, is_sparse=is_sparse) + for role, size, is_sparse in buffers + ], + sliding_window_size=window, + sparse_cache=sparse_cache, + ) + + +def _sparse_layer( + layer_id: int, + max_cached_units: int = 33 * 32, + *, + fetched: bool = True, + selection_id: int | None = None, +) -> AttentionLayerConfig: + """A V4-Flash CSA layer: offloaded NVFP4 rows beside the GPU-resident fp4 indexer. + + One compressed row covers four tokens, 32 rows to a page. + """ + return _layer( + layer_id, + ("data", 32 * 288, True), + ("scale", 32 * 36, True), + ("indexer", 32 * 68, False), + sparse_cache=SparseCacheConfig( + num_selected_blocks=512, + max_cached_units=max_cached_units, + unit_tokens=4, + selection_id=selection_id, + ) + if fetched + else None, + ) + + +def _config(layers, constraints) -> KVCacheManagerConfig: + return KVCacheManagerConfig( + tokens_per_block=TOKENS_PER_BLOCK, + cache_tiers=[GpuCacheTierConfig(quota=1 << 30)], + layers=layers, + constraints=constraints, + ) + + +def _v4_flash_like_layers() -> list: + return [ + _layer(0, ("swa", 65_536, False), window=128), + _sparse_layer(1), + # A ratio-128 layer keeps its one compressed row per page on the GPU. + _layer(2, ("data", 288, False), ("scale", 36, False)), + AttentionLayerConfig( + layer_id=LayerId(3), + buffers=[ + BufferConfig(role=DataRole("state"), size=16_384), + # A page of eight-token blocks expands to 16 of them. + BufferConfig(role=DataRole("score"), size=25, tokens_per_block_override=8), + ], + sliding_window_size=8, + ), + _sparse_layer(4), + ] + + +# Native layout + + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + ("unit_tokens", "units_per_page"), [(4, 32), (1, 128), (TOKENS_PER_BLOCK, 1), (None, 1)] +) +def test_units_per_page(unit_tokens: int | None, units_per_page: int) -> None: + assert sparse_units_per_page(unit_tokens, TOKENS_PER_BLOCK) == units_per_page + + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + ("cells", "lease_bytes"), + [ + # 21 layers x cells x (9216 B data + 1152 B scale), plus one int32 word per + # cached entry. A lease holds the reservation and nothing of the step: the + # work lists belong to the Batch, sized by the batch rather than by whoever + # is seated. + (17, 3_747_200), + (32, 7_053_312), + (33, 7_273_856), + ], +) +def test_lease_matches_native_layout(cells: int, lease_bytes: int) -> None: + shape = SparseArenaShape( + num_layers=21, + num_selections=21, + cells_per_buffer=cells, + units_per_cell=32, + cell_bytes=(9216, 1152) * 21, + ) + + assert shape.lease_bytes == lease_bytes + # Shared tables and region alignment are paid once, not per lease. + assert shape.arena_bytes(513) - shape.arena_bytes(512) == pytest.approx(lease_bytes, abs=512) + + +@pytest.mark.cpu_only +def test_arena_follows_the_batch_rows() -> None: + layers = [_layer(0, ("key", 65_536, False), window=128)] + layers += [_sparse_layer(i) for i in range(1, 22)] + config = _config( + layers, + [ + BatchDesc([KVCacheDesc(capacity=4096, history_length=4095)] * 6), + BatchDesc([KVCacheDesc(capacity=0, history_length=0)] * 513), + ], + ) + + (shape,) = sparse_arena_shapes(config.layers, tokens_per_block=TOKENS_PER_BLOCK) + + assert shape == V4_FLASH_ARENA + assert shape.lease_bytes == 7_273_856 + assert sparse_arena_bytes(config, 513) == shape.arena_bytes(513) + + +@pytest.mark.cpu_only +def test_layers_of_one_arena_must_agree_on_the_cache() -> None: + layers = [_sparse_layer(0, 17 * 32), _sparse_layer(1, 33 * 32)] + + with pytest.raises(ValueError, match="same sparse_cache"): + sparse_arena_shapes(layers, tokens_per_block=TOKENS_PER_BLOCK) + + +@pytest.mark.cpu_only +def test_layers_without_a_fetch_path_join_no_arena() -> None: + # A layer whose life cycle has an arena, but which reads nothing through it. + layers = [_sparse_layer(0), _sparse_layer(1, 17 * 32, fetched=False)] + + (shape,) = sparse_arena_shapes(layers, tokens_per_block=TOKENS_PER_BLOCK) + + assert shape.num_layers == 1 + assert shape.cell_bytes == (32 * 288, 32 * 36) + + +@pytest.mark.cpu_only +def test_layers_sharing_a_selection_share_a_metadata_row() -> None: + """The arena's metadata is sized by selections, so sharing one shrinks it.""" + apart = sparse_arena_shapes( + [_sparse_layer(0), _sparse_layer(1)], tokens_per_block=TOKENS_PER_BLOCK + ) + together = sparse_arena_shapes( + [_sparse_layer(0), _sparse_layer(1, selection_id=0)], tokens_per_block=TOKENS_PER_BLOCK + ) + + (apart_shape,) = apart + (shared_shape,) = together + assert (apart_shape.num_layers, apart_shape.num_selections) == (2, 2) + assert (shared_shape.num_layers, shared_shape.num_selections) == (2, 1) + # The cells are per layer and unchanged; only the residency and recency rows merge. + assert shared_shape.cell_bytes == apart_shape.cell_bytes + assert shared_shape.entries == apart_shape.entries // 2 + assert shared_shape.lease_bytes < apart_shape.lease_bytes + + +@pytest.mark.cpu_only +def test_layers_with_different_windows_get_separate_arenas() -> None: + windowed = _sparse_layer(1) + windowed.sliding_window_size = 4096 + + shapes = sparse_arena_shapes([_sparse_layer(0), windowed], tokens_per_block=TOKENS_PER_BLOCK) + + assert [shape.num_layers for shape in shapes] == [1, 1] + + +@pytest.mark.cpu_only +def test_no_arena_without_a_fetch_path_or_leases() -> None: + lease_constraint = [BatchDesc([KVCacheDesc(capacity=0, history_length=0)] * 4)] + + assert ( + sparse_arena_bytes(_config([_layer(0, ("key", 65_536, False))], lease_constraint), 4) == 0 + ) + # A sparse buffer without a fetch path has no cache to lease. + unfetched = _sparse_layer(0, fetched=False) + assert sparse_arena_bytes(_config([unfetched], lease_constraint), 4) == 0 + # No rows to seat means no leases, whatever the layers declare. + assert sparse_arena_bytes(_config([_sparse_layer(0)], []), 0) == 0 + + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + ("quota", "granule_bytes"), + [ + (512 << 20, 2 << 20), + (1 << 30, 2 << 20), + (3 << 30, 4 << 20), + (16 << 30, 32 << 20), + (1 << 40, 32 << 20), + ], +) +def test_gpu_pool_granule_follows_the_quota(quota: int, granule_bytes: int) -> None: + assert _introspection.tier_allocation_granularity(CacheTier.GPU_MEM, quota) == granule_bytes + + +@pytest.mark.cpu_only +def test_host_history_holds_a_page_of_every_sparse_buffer() -> None: + layers = [_sparse_layer(layer_id) for layer_id in range(21)] + layers.append(_layer(21, ("key", 8704, False))) + + assert sparse_history_bytes_per_block(layers) == V4_FLASH_HOST_BYTES_PER_BLOCK + # A sparse buffer without a fetch path still keeps its history on the host. + assert ( + sparse_history_bytes_per_block([_layer(0, ("key", 4096, True), ("value", 4096, False))]) + == 4096 + ) + + +@pytest.mark.cpu_only +def test_life_cycles_follow_native_numbering_and_split_bytes_by_memory() -> None: + life_cycles = life_cycle_pages(_v4_flash_like_layers(), TOKENS_PER_BLOCK) + + assert life_cycles == [ + LifeCyclePage( + window_size=128, + num_sink_tokens=0, + is_state=False, + device_bytes=65_536, + host_bytes=0, + device_pool_bytes=(65_536,), + ), + # Both compressed layers and the ratio-128 layer share the no-window life + # cycle, and the two indexers of one page size share a pool. + LifeCyclePage( + window_size=None, + num_sink_tokens=0, + is_state=False, + device_bytes=2 * 2176 + 288 + 36, + host_bytes=2 * (9216 + 1152), + device_pool_bytes=(2 * 2176, 288, 36), + ), + LifeCyclePage( + window_size=8, + num_sink_tokens=0, + is_state=False, + device_bytes=16_384 + 400, + host_bytes=0, + device_pool_bytes=(16_384, 400), + ), + ] + ssm = SsmLayerConfig( + layer_id=LayerId(5), buffers=[BufferConfig(role=DataRole("ssm"), size=4096)] + ) + (state,) = life_cycle_pages([ssm], TOKENS_PER_BLOCK) + assert state.is_state and state.device_bytes == 4096 + assert state.device_pool_bytes == (4096,) + + +@pytest.mark.cpu_only +def test_life_cycle_with_only_sparse_buffers_is_rejected() -> None: + with pytest.raises(ValueError, match="must also hold a dense buffer"): + life_cycle_pages([_layer(0, ("history", 4096, True))], TOKENS_PER_BLOCK) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + ("window", "sink_tokens", "capacity", "headroom", "pages"), + [ + (None, 0, 1 << 17, 1, 1024), + # A 128-token window and the token being written can straddle two pages. + (128, 0, 1 << 17, 1, 2), + (128, 0, 1 << 17, 4, 3), + (128, 200, 1 << 17, 1, 4), + # A request never holds more pages than its capacity spans. + (4096, 0, 300, 1, 3), + ], +) +def test_pages_a_decoding_request_holds( + window: int | None, sink_tokens: int, capacity: int, headroom: int, pages: int +) -> None: + life_cycle = LifeCyclePage( + window_size=window, + num_sink_tokens=sink_tokens, + is_state=False, + device_bytes=1, + host_bytes=0, + device_pool_bytes=(1,), + ) + assert ( + life_cycle.pages_per_request( + capacity=capacity, capacity_headroom=headroom, tokens_per_block=TOKENS_PER_BLOCK + ) + == pages + ) + + +@pytest.mark.cpu_only +def test_batch_matches_native_layout() -> None: + # Three int32 page tables of 8,196 pages, two watermarks per row, and one arena slot + # per row, which serves every life cycle. + assert V4_FLASH_BATCH.block_bytes == 4 * 256 * (3 * 8196 + 2 + 1) + assert V4_FLASH_BATCH.block_bytes == 25_181_184 + # The mirror carries only the block's tail: what a captured kernel reads. + assert V4_FLASH_BATCH.mirror_bytes == 4 * 256 * (2 + 1) == 3_072 + # The staging copy is mapped in whole 4 KiB pages. + assert V4_FLASH_BATCH.staging_bytes == 25_182_208 + # Per row: a slot per selected unit, a miss count and the unit in flight. Then the + # two compacted work lists -- a row, an ordinal and a slot per entry, plus a total + # per list -- one for the fetch, sized for every row missing its whole selection, + # and one per layer for the write-back, sized by the step's token budget. + assert V4_FLASH_BATCH.fetch_scratch_bytes == 4 * ( + 256 * (512 + 2) + (3 * 256 * 512 + 1) + 2 * (3 * 256 + 1) + 256 * 2 + ) + assert V4_FLASH_BATCH.fetch_scratch_bytes == 2_107_404 + # The device holds the staged block a publish fills ahead of the step, the mirror of + # its tail that a captured kernel reads, and the work arrays. + assert V4_FLASH_BATCH.device_bytes == 25_181_184 + 3_072 + 2_107_404 == 27_291_660 + + +@pytest.mark.cpu_only +def test_the_token_budget_sizes_the_write_back_list() -> None: + """A row advancing t tokens crosses at most t/unit_tokens + 1 units.""" + layers = _v4_flash_like_layers() + shape = sparse_batch_shape( + layers, + tokens_per_block=TOKENS_PER_BLOCK, + max_batch_size=256, + max_blocks_per_seq=8196, + max_total_input_tokens=8192, + ) + + # Four tokens to a unit, and one partial unit per row at the start of its move. + assert shape.flush_capacities == (8192 // 4 + 256,) + # The bound is an aggregate: it does not multiply by rows, which is what makes it + # affordable. Per row it would have to allow for one row taking the whole budget. + assert shape.flush_capacities[0] * 2 < 256 * (8192 // 4 + 1) + assert shape.fetch_scratch_bytes > V4_FLASH_BATCH.fetch_scratch_bytes + + # Left unset it assumes a decode step: one token per row, so one unit boundary. + assert V4_FLASH_BATCH.flush_capacities == (256,) + + +@pytest.mark.cpu_only +def test_batch_has_a_page_table_per_life_cycle_and_a_fetch_per_sparse_one() -> None: + layers = _v4_flash_like_layers() + + assert ( + sparse_batch_shape( + layers, tokens_per_block=TOKENS_PER_BLOCK, max_batch_size=256, max_blocks_per_seq=8196 + ) + == V4_FLASH_BATCH + ) + # Layers without a fetch path still get page tables, but no fetch work arrays. + layers = [_layer(0, ("swa", 4096, False), window=128), _sparse_layer(1, fetched=False)] + batch = sparse_batch_shape( + layers, tokens_per_block=TOKENS_PER_BLOCK, max_batch_size=4, max_blocks_per_seq=8 + ) + assert batch.num_life_cycles == 2 + assert batch.selection_widths == () + assert batch.fetch_scratch_bytes == 0 + + +# Sizing policy + + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + ("cache_factor", "num_selected_blocks", "units_per_cell", "cells"), + [ + # The smallest cache holds 512 selected rows plus the row being written. + (1.0, 512, 32, 17), + # The factor scales the selection; the row being written stays one row. + (1.5, 512, 32, 25), + (2.0, 512, 32, 33), + (3.0, 512, 32, 49), + # A selection that does not fill whole cells rounds up once, at the end. + (2.0, 500, 32, 32), + (1.0, 512, 64, 9), + (1.0, 1, 32, 1), + # Float error from a computed factor does not add a unit. + (0.1 * 3 / 0.3, 512, 32, 17), + ], +) +def test_sparse_cache_cells_scale_the_selection( + cache_factor: float, num_selected_blocks: int, units_per_cell: int, cells: int +) -> None: + assert sparse_cache_cells(cache_factor, num_selected_blocks, units_per_cell) == cells + + +@pytest.mark.cpu_only +def test_sparse_cache_cells_reject_less_than_one_selection() -> None: + with pytest.raises(ValueError, match="at least 1"): + sparse_cache_cells(0.5, 512, 32) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + ("unit_tokens", "num_selected_units", "units_per_page", "cells"), + [ + # Two selections of 512 rows plus the one being written fill 33 cells of 32 rows. + (4, 512, 32, 33), + # Two selections of 2048 tokens plus one fill 33 cells of 128 tokens. + (1, 2048, 128, 33), + # Two selections of 128 pools of 16 tokens plus one fill 33 cells of 8 pools. + (16, 128, 8, 33), + # Two selections of 64 pages plus one, a unit given as a page or left unset. + (TOKENS_PER_BLOCK, 64, 1, 129), + (None, 64, 1, 129), + ], + ids=["compressed_row", "token", "token_pool", "page", "unset"], +) +def test_sparse_cache_config_sizes_the_cache_for_any_selection_unit( + unit_tokens: int | None, num_selected_units: int, units_per_page: int, cells: int +) -> None: + cache = sparse_cache_config( + num_selected_units=num_selected_units, + unit_tokens=unit_tokens, + tokens_per_block=TOKENS_PER_BLOCK, + gpu_cache_factor=2.0, + ) + + assert cache.num_selected_blocks == num_selected_units + # The config names a definite unit size: a caller passing none means a whole page, + # and the helper resolves that rather than storing the absence. + assert cache.unit_tokens == (TOKENS_PER_BLOCK if unit_tokens is None else unit_tokens) + # The default cache holds two selections plus the unit being written, in whole cells. + assert cache.max_cached_units == cells * units_per_page + (arena,) = sparse_arena_shapes( + [_layer(0, ("history", 64 * units_per_page, True), ("key", 64, False), sparse_cache=cache)], + tokens_per_block=TOKENS_PER_BLOCK, + ) + assert arena.units_per_cell == units_per_page + assert arena.cells_per_buffer == cells + assert arena.cell_bytes == (64 * units_per_page,) + + +@pytest.mark.cpu_only +def test_sparse_cache_config_rejects_a_unit_that_does_not_divide_a_page() -> None: + with pytest.raises(ValueError, match="must divide tokens_per_block"): + sparse_cache_config( + num_selected_units=8, + unit_tokens=48, + tokens_per_block=TOKENS_PER_BLOCK, + gpu_cache_factor=1.0, + ) + + +@pytest.mark.cpu_only +def test_pool_ratio_buys_the_same_requests_in_every_life_cycle() -> None: + life_cycles = life_cycle_pages(_v4_flash_like_layers(), TOKENS_PER_BLOCK) + + ratio = sparse_offload_pool_ratio( + life_cycles, capacity=1 << 17, capacity_headroom=1, tokens_per_block=TOKENS_PER_BLOCK + ) + + # Every pool counts, host-backed ones included: the ratio is compared against a layout + # measured the same way, and is projected onto cold levels where one pool holds dense and + # sparse buffers together. Host pages costing the GPU quota nothing is a fact about the + # quota, not about the ratio. + request_bytes = [ + (life_cycle.device_bytes + life_cycle.host_bytes) + * life_cycle.pages_per_request( + capacity=1 << 17, capacity_headroom=1, tokens_per_block=TOKENS_PER_BLOCK + ) + for life_cycle in life_cycles + ] + assert ratio == pytest.approx([b / sum(request_bytes) for b in request_bytes]) + # The history life cycle is the one with host-backed pages, so counting them moves it. + assert ratio[1] > 0.9, "the offloaded history dominates once its host pages are counted" + assert sum(ratio) == pytest.approx(1.0, abs=1e-9) + + +def _history_beside_five_window_pools() -> list[LifeCyclePage]: + """The history's one dense pool wastes little of its granules on whole slots, + while the windowed life cycles have five pools that can each leave one.""" + layers = [ + _layer( + 0, ("key", 131_072, False), ("value", 65_536, False), ("x", 16_384, False), window=128 + ), + _sparse_layer(1), + _layer(2, ("state", 65_536, False), ("score", 800, False), window=8), + ] + return life_cycle_pages(layers, TOKENS_PER_BLOCK) + + +def _native_slots(life_cycles: list[LifeCyclePage], ratio: list[float], quota: int) -> list[int]: + """Slots the native manager gives each life cycle, one pool group apiece. + + It rounds the quota up to a whole granule, then splits it by ``ratio``. + """ + granule = _introspection.tier_allocation_granularity(CacheTier.GPU_MEM, quota) + return _introspection.ratio_to_slot_count_list( + ceil_div(quota, granule) * granule, + [list(life_cycle.device_pool_bytes) for life_cycle in life_cycles], + ratio, + granule, + [1] * len(life_cycles), + ) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + ("host_history_size", "pinnable_host_bytes", "budget"), + [ + # An explicit size covers the staging buffers too, whatever automatic + # sizing may pin. + (1 << 30, 0.0, (1 << 30) - 4096), + # Otherwise the pools hold max_batch_size requests... + (None, float("inf"), 256 * 1000), + # ...unless they and the staging buffers outgrow what automatic sizing + # may pin. + (None, 50_000.0, 50_000 - 4096), + ], +) +def test_host_budget( + host_history_size: int | None, pinnable_host_bytes: float, budget: int +) -> None: + assert ( + plan_host_history_bytes( + host_history_size=host_history_size, + max_batch_size=256, + host_bytes_per_request=1000, + pinnable_host_bytes=pinnable_host_bytes, + staging_bytes=4096, + ) + == budget + ) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + ("host_history_size", "pinnable_host_bytes", "match"), + [ + (4096, float("inf"), "host_history_size .* leaves no host memory for offloaded history"), + (None, 4096.0, "automatically leaves no host memory for offloaded history"), + ], +) +def test_host_budget_must_leave_the_pools_room( + host_history_size: int | None, pinnable_host_bytes: float, match: str +) -> None: + with pytest.raises(ValueError, match=match): + plan_host_history_bytes( + host_history_size=host_history_size, + max_batch_size=256, + host_bytes_per_request=1000, + pinnable_host_bytes=pinnable_host_bytes, + staging_bytes=4096, + ) + + +@pytest.mark.cpu_only +def test_host_budget_is_empty_on_a_rank_without_offloaded_history() -> None: + assert ( + plan_host_history_bytes( + host_history_size=None, + max_batch_size=256, + host_bytes_per_request=0, + pinnable_host_bytes=0.0, + staging_bytes=0, + ) + == 0 + ) + + +# Capacity report and host admission + + +def _v4_flash_capacity(**overrides): + # 128K-token requests on a disaggregated decode worker with max_batch_size 256. + kwargs = dict( + num_tokens=1 << 17, + max_batch_size=256, + gpu_pool_bytes=100 << 30, + # The host pool holds 4,096 blocks: four requests of 128K tokens. + host_pool_bytes=4096 * V4_FLASH_HOST_BYTES_PER_BLOCK, + host_budget_bytes=50 << 30, + # Every window spans two pages, plus 52,176 GPU bytes per 128 tokens. + gpu_bytes_per_request=81_657_856 + 1024 * 52_176, + fixed_gpu_bytes_per_request=81_657_856, + host_bytes_per_request=1024 * V4_FLASH_HOST_BYTES_PER_BLOCK, + arenas=[V4_FLASH_ARENA], + # One lease per Batch row. + num_leases=256, + batch_buffer_bytes=V4_FLASH_BATCH_BUFFER_BYTES, + ) + kwargs.update(overrides) + return solve_sparse_offload_capacity(**kwargs) + + +@pytest.mark.cpu_only +def test_capacity_names_the_binding_limit_and_the_batch_that_fits() -> None: + capacity = _v4_flash_capacity() + + assert capacity.by_limit == {"max_batch_size": 256, "host": 4, "gpu": 794} + assert capacity.limited_by == "host" + assert capacity.fits == 4 + assert capacity.lease_bytes == 7_273_856 + assert capacity.arena_bytes == V4_FLASH_ARENA.arena_bytes(256) + # The 50 GiB host budget holds 240 requests of 128K tokens; the GPU budget + # would hold more. + assert capacity.recommended_max_batch_size == 240 + report = capacity.describe() + assert "4 by host" in report + assert "limited by host" in report + assert "256 leases" in report + assert "offload buffers 157.0MiB for 256 rows" in report + assert "A max_batch_size of 240 fits" in report + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("padding_leases", [0, 1]) +def test_capacity_recommendation_charges_a_batch_slot_one_lease(padding_leases: int) -> None: + # A cache holds a lease only while it is seated in the Batch, so each slot + # holds one, even on a disaggregated or overlap worker where the index mapper + # admits two requests per slot. + lease_bytes = V4_FLASH_ARENA.lease_bytes + gpu_bytes_per_request = 81_657_856 + 1024 * 52_176 + # Ten batch slots of one lease each, plus any padding leases. + budget = ( + 10 * (gpu_bytes_per_request + lease_bytes) + + V4_FLASH_ARENA.arena_bytes(padding_leases) + - V4_FLASH_ARENA.arena_bytes(10 + padding_leases) + ) + capacity = _v4_flash_capacity( + max_batch_size=10, + gpu_pool_bytes=budget, + host_budget_bytes=1 << 50, + num_leases=10 + padding_leases, + batch_buffer_bytes=0, + ) + assert capacity.recommended_max_batch_size == 10 + + short = _v4_flash_capacity( + max_batch_size=10, + gpu_pool_bytes=budget - 1, + host_budget_bytes=1 << 50, + num_leases=10 + padding_leases, + batch_buffer_bytes=0, + ) + assert short.recommended_max_batch_size == 9 + + +@pytest.mark.cpu_only +def test_capacity_recommendation_charges_every_batch_slot_its_rows_of_offload_buffers() -> None: + gpu_bytes_per_request = 81_657_856 + 1024 * 52_176 + # 1 MiB of offload buffers per row, sized for 16 rows. + slot_bytes = gpu_bytes_per_request + V4_FLASH_ARENA.lease_bytes + (1 << 20) + # Device memory for ten batch slots, the arena's shared tables and the + # offload buffers, of which the pools get what the rest leaves. + budget = 10 * slot_bytes + V4_FLASH_ARENA.arena_bytes(0) + gpu_pool_bytes = budget - V4_FLASH_ARENA.arena_bytes(16) - (16 << 20) + capacity = _v4_flash_capacity( + max_batch_size=16, + gpu_pool_bytes=gpu_pool_bytes, + host_budget_bytes=1 << 50, + num_leases=16, + batch_buffer_bytes=16 << 20, + ) + + assert capacity.recommended_max_batch_size == 10 + + +@pytest.mark.cpu_only +def test_capacity_ties_go_to_the_batch_limit_then_host() -> None: + capacity = _v4_flash_capacity(max_batch_size=4, gpu_pool_bytes=4 * (81_657_856 + 1024 * 52_176)) + assert capacity.by_limit == {"max_batch_size": 4, "host": 4, "gpu": 4} + assert capacity.limited_by == "max_batch_size" + + capacity = _v4_flash_capacity(gpu_pool_bytes=4 * (81_657_856 + 1024 * 52_176)) + assert capacity.limited_by == "host" + + +@pytest.mark.cpu_only +def test_host_reservations_replace_and_release() -> None: + reservations = HostHistoryReservations(capacity_blocks=10) + + assert reservations.can_reserve(1, 6) + reservations.reserve(1, 6) + assert not reservations.can_reserve(2, 5) + # A request's new reservation replaces the one it holds. + assert reservations.can_reserve(1, 10) + reservations.reserve(1, 2) + assert reservations.can_reserve(2, 5) + reservations.reserve(2, 5) + assert reservations.reserved_blocks == 7 + + reservations.release(1) + reservations.release(9) + assert reservations.reserved_blocks == 5 diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py index 08592600e4eb..495a64275b33 100644 --- a/tests/unittest/_torch/executor/test_py_executor.py +++ b/tests/unittest/_torch/executor/test_py_executor.py @@ -15,6 +15,7 @@ import threading import time import types +from collections.abc import Callable from contextlib import contextmanager, nullcontext from unittest.mock import MagicMock, Mock, patch @@ -84,6 +85,7 @@ def _forward_step_executor( sampler=Mock(get_cache_indirection=Mock(return_value=object())), execution_stream=Mock(), _loop_on_execution_stream=False, + _is_kv_manager_v2=False, _iter_adp_dummy_ctx_tokens=0, _iter_adp_dummy_gen_tokens=0, _compute_adp_dummy_tokens=PyExecutor._compute_adp_dummy_tokens, @@ -861,6 +863,75 @@ def test_free_request_resources_releases_target_and_draft_model_state() -> None: assert not executor._prefetched_request_ids +@pytest.mark.parametrize("is_v2", [False, True]) +@pytest.mark.parametrize("forward_fails", [False, True]) +def test_forward_step_updates_kv_after_model_stream_join(is_v2: bool, forward_fails: bool) -> None: + batch = ScheduledRequests() + events = [] + logits = object() + + def forward(*args: object, **kwargs: object) -> object: + events.append("forward") + if forward_fails: + raise RuntimeError("forward failed") + return logits + + caller_stream = Mock() + caller_stream.wait_stream.side_effect = lambda stream: events.append("join") + kv_cache_manager = Mock() + + def update_forward_resources( + scheduled_batch: ScheduledRequests, *, wait_for_input_copy: Callable[[], None] + ) -> None: + wait_for_input_copy() + events.append("history") + + kv_cache_manager.update_forward_resources.side_effect = update_forward_resources + executor = types.SimpleNamespace( + model_engine=types.SimpleNamespace(route_capture=None, forward=forward), + resource_manager=object(), + kv_cache_manager=kv_cache_manager, + _is_kv_manager_v2=is_v2, + _loop_on_execution_stream=False, + enable_attention_dp=False, + enable_spec_decode=False, + runtime_draft_len=0, + iter_counter=0, + sampler=Mock(), + execution_stream=Mock(), + _iter_adp_dummy_ctx_tokens=0, + _iter_adp_dummy_gen_tokens=0, + _compute_adp_dummy_tokens=PyExecutor._compute_adp_dummy_tokens, + _maybe_record_hang_diagnostic_phase=Mock(), + _attach_encoder_output_to_execution_stream=Mock(), + _mark_cross_kv_projection_consumed=Mock(), + _kv_connector_wait_for_save=Mock(), + _wait_for_model_engine_input_copy=Mock(side_effect=lambda: events.append("input_copy")), + _handle_errors=Mock(), + ) + with ( + patch("torch.cuda.current_stream", return_value=caller_stream), + patch("torch.cuda.stream", return_value=nullcontext()), + patch("tensorrt_llm._torch.pyexecutor.py_executor.ExpertStatistic.set_iter"), + ): + result = PyExecutor._forward_step(executor, batch) + if forward_fails: + assert result is None + assert events == ["forward"] + executor._handle_errors.assert_called_once_with("forward failed") + else: + assert result is logits + assert events == ["forward", "join"] + (["input_copy", "history"] if is_v2 else []) + executor._handle_errors.assert_not_called() + if is_v2 and not forward_fails: + kv_cache_manager.update_forward_resources.assert_called_once_with( + batch, wait_for_input_copy=executor._wait_for_model_engine_input_copy + ) + else: + kv_cache_manager.update_forward_resources.assert_not_called() + executor._wait_for_model_engine_input_copy.assert_not_called() + + class _InflightRequestIds: def __init__(self): self.ids = set() diff --git a/tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py b/tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py index 86b15e37d6b2..b093f85b1396 100644 --- a/tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py +++ b/tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py @@ -725,3 +725,27 @@ def teardown_managers(self, resources): assert allocation_phases == ["final_kv_cache_allocation", "final_executor_creation"] assert phases[-1] == "executor_start_worker" assert timers[0].depth == 0 + + +@pytest.mark.parametrize("decode_only", [False, True], ids=["default", "decode_only"]) +def test_decode_only_worker_skips_kv_cache_estimation(monkeypatch, decode_only): + """A decode-only worker cannot run the dummy-prompt profiling forward.""" + skip_est_args = [] + + class RecordingCacheCreator(_DummyKvCacheCreator): + """Record whether estimation was disabled at construction.""" + + def __init__(self, **kwargs): + super().__init__(**kwargs) + skip_est_args.append(kwargs["skip_est"]) + + monkeypatch.delenv("TRTLLM_SKIP_KV_CACHE_ESTIMATION", raising=False) + monkeypatch.setattr(py_executor_creator, "is_decode_only", lambda config: decode_only) + _run_create_py_executor( + monkeypatch, + sm_version=100, + kv_cache_quant_algo=QuantAlgo.NO_QUANT, + kv_cache_creator_cls=RecordingCacheCreator, + ) + + assert skip_est_args == [decode_only] diff --git a/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py b/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py index dee28f4734a1..f53b0cc5ca31 100644 --- a/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py +++ b/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py @@ -1146,6 +1146,37 @@ def test_deepseek_v4_model_config_defaults_to_fp4_indexer(tmp_path, monkeypatch) assert model_config.sparse_attention_config.indexer_k_dtype == "fp4" +@pytest.mark.parametrize("enable_offload", [False, True]) +def test_deepseek_v4_model_config_preserves_sparse_offload(tmp_path, monkeypatch, enable_offload): + config = DeepseekV4Config( + architectures=["DeepseekV4ForCausalLM"], + num_hidden_layers=8, + compress_ratios=[0, 0, 4, 128, 4, 128, 4, 128], + ) + monkeypatch.setattr( + "tensorrt_llm._torch.model_config.load_pretrained_config", lambda *args, **kwargs: config + ) + offload_options = ( + {"gpu_cache_factor": 3.0, "host_history_size": 2 * 1024**3} if enable_offload else {} + ) + sparse_config = DeepSeekV4SparseAttentionConfig( + enable_kv_cache_offload=enable_offload, + kv_cache_offload_config=offload_options, + ) + + resolved = ModelConfig.from_pretrained( + str(tmp_path), + sparse_attention_config=sparse_config, + attn_backend="TRTLLM", + moe_backend="TRTLLM", + ).sparse_attention_config + + assert resolved.enable_kv_cache_offload is enable_offload + assert resolved.kv_cache_offload_config == sparse_config.kv_cache_offload_config + assert resolved.compress_ratios == [1, 1, 4, 128, 4, 128, 4, 128] + assert sparse_config.compress_ratios == [1, 1, 4, 128, 4, 128, 4] + + def test_deepseek_v4_model_config_defaults_to_fp8_before_blackwell(tmp_path, monkeypatch): checkpoint_ratios = [128, 128, 4, 128, 4, 128, 0, 4] config = DeepseekV4Config( diff --git a/tests/unittest/_torch/modules/test_mhc.py b/tests/unittest/_torch/modules/test_mhc.py index 0a072bcfb4b1..154daaefff88 100644 --- a/tests/unittest/_torch/modules/test_mhc.py +++ b/tests/unittest/_torch/modules/test_mhc.py @@ -22,6 +22,46 @@ timing_stats = defaultdict(dict) +@skip_pre_blackwell +@pytest.mark.parametrize("num_tokens", [1, 129]) +@pytest.mark.parametrize("force_fma", [False, True]) +@pytest.mark.parametrize( + "deterministic_env", ["FORCE_DETERMINISTIC", "FORCE_MHC_KERNEL_DETERMINISTIC"] +) +def test_mhc_deterministic_mode_repeats_with_cached_tactic( + monkeypatch: pytest.MonkeyPatch, num_tokens: int, force_fma: bool, deterministic_env: str +) -> None: + from tensorrt_llm._torch.modules.mhc import mhc_cuda + + monkeypatch.delenv("FORCE_DETERMINISTIC", raising=False) + monkeypatch.delenv("FORCE_MHC_KERNEL_DETERMINISTIC", raising=False) + monkeypatch.setenv(deterministic_env, "1") + if force_fma: + monkeypatch.setattr(mhc_cuda, "_fused_hc_mma_supported", lambda: False) + runner = mhc_cuda.MhcFusedHcRunner(4, 4096, 1e-6, 1e-6, 1e-6, 2.0, 20) + generator = torch.Generator(device="cuda").manual_seed(0) + + def random_tensor(*shape: int, dtype: torch.dtype = torch.float32) -> torch.Tensor: + return torch.randn(shape, generator=generator, dtype=dtype, device="cuda") + + inputs = [ + random_tensor(num_tokens, 4096, dtype=torch.bfloat16), + random_tensor(num_tokens, 4, 4096, dtype=torch.bfloat16), + random_tensor(num_tokens, 4).sigmoid(), + random_tensor(num_tokens, 4, 4).softmax(-1), + random_tensor(24, 16384) * 0.01, + torch.ones(3, device="cuda"), + torch.zeros(24, device="cuda"), + ] + assert len(runner.get_valid_tactics(inputs, None)) == 1 + reference = runner(inputs) + for _ in range(20): + actual = runner(inputs, tactic=("fused_half_mma", 0, 16, 0, 1)) + for expected, value in zip(reference, actual): + assert torch.isfinite(value).all() + assert torch.equal(value, expected) + + def _mhc_fused_hc_mma_available() -> bool: try: from tensorrt_llm._torch.modules.mhc.mhc_cuda import _fused_hc_mma_supported @@ -1198,7 +1238,9 @@ def _make_fused_hc_runner_case(n: int, hidden_size: int, hc_mult: int, seed: int return runner, inputs -@pytest.mark.parametrize("num_tokens", [1, 129]) +# 300 rows take a different automatic BigFuse block size than one row, which the +# row-invariance check below must not see. +@pytest.mark.parametrize("num_tokens", [1, 129, 300]) @pytest.mark.parametrize("hidden_size", [4096, 7168]) @pytest.mark.parametrize("fuse_norm", [False, True]) @pytest.mark.parametrize("lagged", [False, True]) @@ -1260,6 +1302,15 @@ def run( single_pre_mix = pre_mix_ext[row : row + 1].clone() if pre_mix_ext is not None else None for actual, reference in zip(run(single_inputs, single_pre_mix), expected, strict=True): torch.testing.assert_close(actual, reference[row : row + 1], rtol=0, atol=0) + # Nor may splitting the batch, which checks every row. + half = (num_tokens + 1) // 2 + for start, stop in ((0, half), (half, num_tokens)): + if start == stop: + continue + part_inputs = [v[start:stop].clone() for v in inputs[:4]] + inputs[4:] + part_pre_mix = pre_mix_ext[start:stop].clone() if pre_mix_ext is not None else None + for actual, reference in zip(run(part_inputs, part_pre_mix), expected, strict=True): + torch.testing.assert_close(actual, reference[start:stop], rtol=0, atol=0) stream = torch.cuda.Stream() stream.wait_stream(torch.cuda.current_stream()) diff --git a/tests/unittest/_torch/modules/test_top_k.py b/tests/unittest/_torch/modules/test_top_k.py index 0c9b7522e027..a2828160a525 100644 --- a/tests/unittest/_torch/modules/test_top_k.py +++ b/tests/unittest/_torch/modules/test_top_k.py @@ -13,6 +13,47 @@ from tensorrt_llm._torch.modules.top_k import _CUTE_DSL_PREFILL_COPY_BITS, TopK, TopKImplementation +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +@pytest.mark.parametrize("is_prefill", [False, True]) +@pytest.mark.parametrize("implementation", list(TopKImplementation)) +@pytest.mark.parametrize( + "deterministic_env", ["FORCE_DETERMINISTIC", "FORCE_TOPK_KERNEL_DETERMINISTIC"] +) +def test_deterministic_topk_repeats_tied_selection( + monkeypatch: pytest.MonkeyPatch, + is_prefill: bool, + implementation: TopKImplementation, + deterministic_env: str, +) -> None: + monkeypatch.delenv("FORCE_DETERMINISTIC", raising=False) + monkeypatch.delenv("FORCE_TOPK_KERNEL_DETERMINISTIC", raising=False) + monkeypatch.setenv(deterministic_env, "1") + generator = torch.Generator(device="cuda").manual_seed(0) + scores = torch.randn(3, 2048, generator=generator, device="cuda").round() + output = torch.empty(3, 512, dtype=torch.int32, device="cuda") + lengths = torch.tensor([2048, 3, 0], dtype=torch.int32, device="cuda") + top_k = TopK( + 512, + prefill_implementation=implementation, + decode_implementation=implementation, + compress_ratio=4, + ) + kwargs = dict( + is_prefill=is_prefill, + row_starts=torch.zeros_like(lengths), + row_ends=lengths, + sequence_lengths=lengths * 4, + scan_lengths=lengths, + ) + reference = top_k(scores, output, **kwargs).clone() + assert (reference >= 0).sum(dim=1).tolist() == [512, 3, 0] + selected_scores = scores[0, reference[0].long()].sort().values + torch.testing.assert_close(selected_scores, scores[0].topk(512).values.sort().values) + for _ in range(20): + actual = top_k(scores, output, **kwargs) + assert torch.equal(actual, reference) + + def test_prefill_torch_masks_dirty_scores_and_pads_output() -> None: scores = torch.tensor( [ diff --git a/tests/unittest/disaggregated/test_chunked_transfer.py b/tests/unittest/disaggregated/test_chunked_transfer.py index 4db5a7e30ecd..93d7086873ac 100644 --- a/tests/unittest/disaggregated/test_chunked_transfer.py +++ b/tests/unittest/disaggregated/test_chunked_transfer.py @@ -152,11 +152,19 @@ def _make_projection_sender() -> Sender: ) extractor = MagicMock() + # Write metas group fragments by the memory type of each view's physical pool, so each + # group's single view resolves to a device pool. extractor.page_table = SimpleNamespace( tokens_per_block=_PROJECTION_TPB, + pool_groups=[SimpleNamespace(pools=[SimpleNamespace(memory_type="VRAM")])], layer_groups=[ - SimpleNamespace(kind=CacheKind.PAGED, sliding_window_size=_PROJECTION_PROMPT_TOKENS), - SimpleNamespace(kind=CacheKind.PAGED, sliding_window_size=_PROJECTION_PROMPT_TOKENS), + SimpleNamespace( + kind=CacheKind.PAGED, + sliding_window_size=_PROJECTION_PROMPT_TOKENS, + pool_group_idx=0, + pool_views=[SimpleNamespace(pool_idx=0)], + ) + for _ in range(2) ], ) # extract(region_ids, layer_group_id, pool_idx) - the caller passes the diff --git a/tests/unittest/disaggregated/test_deepseek_v4_kv_transfer.py b/tests/unittest/disaggregated/test_deepseek_v4_kv_transfer.py index 1daa4b016558..6521b60e2720 100644 --- a/tests/unittest/disaggregated/test_deepseek_v4_kv_transfer.py +++ b/tests/unittest/disaggregated/test_deepseek_v4_kv_transfer.py @@ -20,9 +20,13 @@ transfers across different TP/PP/DP configurations. """ +import ctypes import functools +import uuid +from types import SimpleNamespace from typing import Dict, List, Optional, Sequence, Tuple +import numpy as np import pytest import torch from kv_transfer_harness import ( @@ -513,6 +517,230 @@ def test_deepseek_v4_kv_transfer( print("PASSED") +def _native_transfer_manager( + offload: bool, dtype: DataType, tokens_per_block: int +) -> DeepseekV4CacheManager: + """Build real V4 pools without enabling the guarded model-inference path.""" + from tensorrt_llm.llmapi.llm_args import SparseKvCacheOffloadConfig + from tensorrt_llm.runtime.kv_cache_manager_v2 import ( + GpuCacheTierConfig, + KVCacheManager, + KVCacheManagerConfig, + ) + + adapter = object.__new__(DeepseekV4CacheManager) + adapter._enable_kv_cache_offload = offload + adapter._sparse_offload_topk = 8 + adapter._sparse_kv_offload = SparseKvCacheOffloadConfig() if offload else None + adapter.pp_layers = [0, 1, 2, 3] + adapter._compress_ratios = [1, 4, 128, 4] + adapter.compressed_block_sizes = [tokens_per_block // r for r in adapter._compress_ratios] + adapter.tokens_per_block = tokens_per_block + adapter.head_dim = HEAD_DIM + adapter.index_head_dim = INDEX_HEAD_DIM + adapter.dtype = DataType.FP8 if dtype == DataType.NVFP4 else dtype + adapter._indexer_k_dtype = "fp4" + adapter._use_nvfp4_compress = dtype == DataType.NVFP4 + adapter._nvfp4_residual_dim = 64 + adapter.use_fp8_ds_mla = False + adapter._swa_window_size = WINDOW_SIZE + adapter.max_seq_len = 4 * tokens_per_block + adapter.max_batch_size = 2 + adapter.is_estimating_kv_cache = False + adapter.max_cuda_graph_batch_size = None + adapter.max_attention_window_vec = [None] * 4 + adapter.max_draft_len = adapter._max_draft_len = 0 + adapter.num_extra_kv_tokens = 0 + adapter.num_kv_heads_per_layer = [1] * 4 + config = adapter._build_cache_config( + KVCacheManagerConfig( + tokens_per_block=tokens_per_block, + cache_tiers=[GpuCacheTierConfig(quota=32 << 20)], + layers=[], + ) + ) + adapter.impl = KVCacheManager(config) + return adapter + + +@pytest.mark.timeout(180) +@pytest.mark.parametrize("dtype", [DataType.BF16, DataType.FP8, DataType.NVFP4]) +@pytest.mark.parametrize("tokens_per_block", [128, 256]) +def test_deepseek_v4_native_sparse_host_transfer(dtype: DataType, tokens_per_block: int) -> None: + """NIXL writes V4 compressed KV/scales to host and all other roles to GPU. + + Uses the production V4 layout builder, native KVCM allocations, and the + sender's registration, mapping, grouping, and submission. Raw byte checks + include partial pages and repeated requests reusing the same pools. + """ + from cuda.bindings import runtime as cudart + from test_sparse_kv_transfer import _rank_info, _sender_and_meta + + from tensorrt_llm._torch.disaggregation.native.transfer import TransferWorker + from tensorrt_llm._torch.disaggregation.nixl._agent_cpp import BindingsNixlTransferAgent + from tensorrt_llm.runtime.kv_cache_manager_v2 import CudaStream + + device_id = torch.cuda.current_device() + stream = CudaStream(torch.cuda.current_stream().cuda_stream) + managers, workers, agents = [], [], [] + suffix = uuid.uuid4().hex + + def copy_bytes(dst: int, src: int, size: int, kind) -> None: + (error,) = cudart.cudaMemcpy(dst, src, size, kind) + assert int(error) == 0 + + try: + for offload, name in ((False, "ctx"), (True, "gen")): + manager = _native_transfer_manager(offload, dtype, tokens_per_block) + managers.append(manager) + table = KVRegionExtractorV1(manager).page_table + agent = BindingsNixlTransferAgent(f"{name}-{suffix}", num_threads=1) + agents.append(agent) + worker = SimpleNamespace() + worker._rank_info = _rank_info(name, table) + worker._rank_info.device_id = device_id + worker._agent = agent + worker._registered_mem = [] + TransferWorker._register_kv_cache(worker) + workers.append(worker) + src, dst = managers + assert {d.type for d in workers[0]._registered_mem} == {"VRAM"} + assert {d.type for d in workers[1]._registered_mem} == {"VRAM", "DRAM"} + agents[0].load_remote_agent(agents[1].name, agents[1].get_local_agent_desc()) + + for request_idx, prompt_len in enumerate( + (tokens_per_block - 1, tokens_per_block, tokens_per_block + 1, 3 * tokens_per_block - 1) + ): + caches = [] + transfer_active = False + try: + blocks = [] + for manager in managers: + cache = manager.impl.create_kv_cache() + caches.append(cache) + assert cache.resume(stream) + assert cache.resize(prompt_len, prompt_len) + blocks.append( + [ + np.array( + list(cache.get_aggregated_page_indices(lg, valid_only=False)), + dtype=np.int64, + ) + for lg in range(len(manager.impl.layer_grouping)) + ] + ) + torch.cuda.synchronize() + expected = [] + host_scale_buffers = 0 + for layer in dst.impl.init_config.layers: + lg = dst.impl.get_layer_group_id(layer.layer_id) + src_lg = src.impl.get_layer_group_id(layer.layer_id) + src_indices = list(caches[0].get_base_page_indices(src_lg)) + dst_indices = list(caches[1].get_base_page_indices(lg)) + for buffer_idx, buffer in enumerate(layer.buffers): + if buffer.role in dst.get_disagg_ignored_roles(): + continue + src_base = src.impl.get_mem_pool_base_address(layer.layer_id, buffer.role) + dst_base = dst.impl.get_mem_pool_base_address(layer.layer_id, buffer.role) + src_stride = src.impl.get_page_stride(layer.layer_id, buffer.role) + dst_stride = dst.impl.get_page_stride(layer.layer_id, buffer.role) + src_converter = src.impl.get_page_index_converter( + layer.layer_id, buffer.role + ) + dst_converter = dst.impl.get_page_index_converter( + layer.layer_id, buffer.role + ) + assert src_converter.expansion == dst_converter.expansion == 1 + src_role_indices = src_converter(src_indices) + dst_role_indices = dst_converter(dst_indices) + host_scale_buffers += int(buffer.is_sparse and "scale" in str(buffer.role)) + for ordinal, (src_index, dst_index) in enumerate( + zip(src_role_indices, dst_role_indices) + ): + if src_index < 0 or dst_index < 0: + continue + data = ( + ( + np.arange(buffer.size, dtype=np.int64) * (2 * buffer_idx + 1) + + int(layer.layer_id) * 23 + + ordinal * 17 + + request_idx * 11 + ) + % 251 + ).astype(np.uint8) + src_ptr = int(src_base) + int(src_index) * src_stride + dst_ptr = int(dst_base) + int(dst_index) * dst_stride + copy_bytes( + src_ptr, + data.ctypes.data, + buffer.size, + cudart.cudaMemcpyKind.cudaMemcpyHostToDevice, + ) + if buffer.is_sparse: + ctypes.memset(dst_ptr, 0, buffer.size) + else: + zeros = np.zeros(buffer.size, dtype=np.uint8) + copy_bytes( + dst_ptr, + zeros.ctypes.data, + buffer.size, + cudart.cudaMemcpyKind.cudaMemcpyHostToDevice, + ) + expected.append((dst_ptr, data, buffer.is_sparse, str(buffer.role))) + if dtype == DataType.NVFP4: + assert host_scale_buffers == 2 + sender, meta = _sender_and_meta( + workers[0]._rank_info.page_table, + workers[1]._rank_info.page_table, + blocks[0], + blocks[1], + device_id=device_id, + prompt_len=prompt_len, + ) + sender._agent = agents[0] + sender._enforce_physical_ownership = False + meta.peer_name = agents[1].name + assert {(g.src_memory_type, g.dst_memory_type) for g in meta.transfer_groups} == { + ("VRAM", "VRAM"), + ("VRAM", "DRAM"), + } + request = sender._make_agent_request(meta, device_id) + transfer_active = True + assert sender._submit_transfer(meta.task, 0, request) == (True, None) + transfer_active = False + assert meta.task.resources_drained + for address, data, on_host, role in expected: + if on_host: + actual = np.frombuffer(ctypes.string_at(address, data.size), dtype=np.uint8) + else: + actual = np.empty_like(data) + copy_bytes( + actual.ctypes.data, + address, + data.size, + cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost, + ) + np.testing.assert_array_equal(actual, data, err_msg=role) + finally: + if transfer_active: + # Quiesce any failed submission before releasing its pools. + for agent in agents: + agent.shutdown() + agents.clear() + workers.clear() + for cache in caches: + cache.close() + finally: + for worker in workers: + for registration in reversed(worker._registered_mem): + worker._agent.deregister_memory(registration) + worker._registered_mem.clear() + for agent in agents: + agent.shutdown() + for manager in managers: + manager.impl.shutdown() + + # --------------------------------------------------------------------------- # Uneven PP layer test configurations # --------------------------------------------------------------------------- diff --git a/tests/unittest/disaggregated/test_sparse_kv_transfer.py b/tests/unittest/disaggregated/test_sparse_kv_transfer.py new file mode 100644 index 000000000000..6c0974385ce5 --- /dev/null +++ b/tests/unittest/disaggregated/test_sparse_kv_transfer.py @@ -0,0 +1,341 @@ +# 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. + +"""Memory descriptors and logical completion for mixed host/device KV handoffs.""" + +import threading +from types import SimpleNamespace +from unittest.mock import Mock + +import numpy as np +import pytest + +from tensorrt_llm._torch.disaggregation.base import TokenRange +from tensorrt_llm._torch.disaggregation.base.agent import MemoryType +from tensorrt_llm._torch.disaggregation.base.transfer import SessionStatus +from tensorrt_llm._torch.disaggregation.native import transfer as transfer_mod +from tensorrt_llm._torch.disaggregation.native.extent import Chunk +from tensorrt_llm._torch.disaggregation.native.mixers.attention.spec import AttentionInfo +from tensorrt_llm._torch.disaggregation.native.peer import PeerRegistrar +from tensorrt_llm._torch.disaggregation.native.rank_info import RankInfo +from tensorrt_llm._torch.disaggregation.resource.kv_extractor import ( + KVRegionExtractorV1, + build_page_table_from_manager, +) +from tensorrt_llm._torch.disaggregation.resource.page import ( + KVCachePageTable, + MapperKind, + PhysicalPool, +) +from tensorrt_llm._torch.disaggregation.resource.utils import get_unique_pool_memory_descs +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2, Role +from tensorrt_llm.disaggregated_params import DisaggregatedParams + +pytestmark = pytest.mark.cpu_only + + +def _manager(offload: bool) -> Mock: + roles = ("compress", "compress_scale", "indexer") + sizes = (288 * 32, 36 * 32, 144 * 32) + layers = [ + SimpleNamespace( + layer_id=layer, + window_size=None, + buffers=[ + SimpleNamespace(role=role, is_sparse=offload and role != "indexer") + for role in roles + ], + ) + for layer in range(2) + ] + coalesced = [ + SimpleNamespace( + single_buffer_size=size, + buffer_ids=[SimpleNamespace(layer_id=layer, role=role) for layer in range(2)], + ) + for role, size in zip(roles, sizes) + ] + impl = SimpleNamespace( + init_config=SimpleNamespace(tokens_per_block=128, layers=layers), + layer_grouping=((0, 1),), + pool_group_descs=[ + SimpleNamespace( + pool_group_index=0, + num_slots=4, + pools=[ + SimpleNamespace(base_address=0x100000 * (i + 1), slot_bytes=2 * size) + for i, size in enumerate(sizes) + ], + slot_desc=SimpleNamespace( + variants=[SimpleNamespace(layer_group_id=0, coalesced_buffers=coalesced)] + ), + ) + ], + ) + # The page-table builder accepts only a KVCacheManagerV2. + return Mock( + spec=KVCacheManagerV2, + get_disagg_pool_view_partition=None, + get_disagg_ignored_roles=lambda: frozenset(), + impl=impl, + pp_layers=[0, 1], + num_kv_heads_per_layer=[1, 1], + get_disagg_role_mapper_kinds=lambda: {Role.ALL: MapperKind.REPLICATED}, + get_disagg_role_layouts=lambda: {}, + ) + + +def _rank_info(name: str, page_table: KVCachePageTable) -> RankInfo: + return RankInfo( + instance_name=name, + instance_rank=0, + tp_size=1, + tp_rank=0, + pp_size=1, + pp_rank=0, + layer_num_per_pp=[ + len( + { + layer.global_layer_id + for group in page_table.layer_groups + for layer in group.local_layers + } + ) + ], + sender_endpoints=["sender"], + self_endpoint="receiver", + transfer_engine_info=b"", + device_id=3, + attention=AttentionInfo(1, page_table.tokens_per_block, 512, 1, False, True), + page_table=page_table, + ) + + +def _sender_and_meta( + src: KVCachePageTable, + dst: KVCachePageTable, + src_blocks: list[np.ndarray] | None = None, + dst_blocks: list[np.ndarray] | None = None, + device_id: int = 3, + prompt_len: int | None = None, +) -> tuple[transfer_mod.Sender, transfer_mod.WriteMeta]: + sender = object.__new__(transfer_mod.Sender) + source = _rank_info("ctx", src) + source.device_id = device_id + # Exercise the metadata exchange, including memory types. + destination = RankInfo.from_bytes(_rank_info("gen", dst).to_bytes()) + destination.device_id = device_id + sender._registrar = PeerRegistrar(source, KVRegionExtractorV1(src)) + sender._registrar.register("gen", 0, destination) + src_blocks = src_blocks if src_blocks is not None else [np.array([1, 3], dtype=np.int64)] + dst_blocks = dst_blocks if dst_blocks is not None else [np.array([2, 0], dtype=np.int64)] + if prompt_len is None: + prompt_len = len(src_blocks[0]) * src.tokens_per_block + chunk = Chunk( + src_blocks, [group.kind for group in src.layer_groups], TokenRange(0, prompt_len), True + ) + task = transfer_mod.KVSendTask(chunk, DisaggregatedParams(disagg_request_id=401), 0, prompt_len) + task._unique_rid = 401 + req_info = transfer_mod.RecvReqInfo(1, "gen", 0, dst_blocks, 401, bounce_dst_base=0x900000) + meta = sender._build_kv_write_meta(task, req_info) + sender._enforce_physical_ownership = True + sender._ownership_poisoned = None + sender._ownership_poison_lock = threading.Lock() + return sender, meta + + +@pytest.mark.parametrize("memory_type", ["VRAM", "DRAM"]) +def test_pool_memory_type_roundtrip(memory_type: str) -> None: + pool = PhysicalPool(0x1000, 64, 4, memory_type=memory_type) + assert PhysicalPool.from_dict(pool.to_dict()) == pool + old_payload = pool.to_dict() + del old_payload["memory_type"] + assert PhysicalPool.from_dict(old_payload).memory_type == "VRAM" + with pytest.raises(ValueError, match="memory type"): + PhysicalPool(0x1000, 64, 4, memory_type="invalid") + + +def test_v2_extracts_host_data_and_scales_beside_device_indexer() -> None: + table = build_page_table_from_manager(_manager(True)) + assert [p.memory_type for p in table.pool_groups[0].pools] == ["DRAM", "DRAM", "VRAM"] + assert table.has_host_pools + assert not build_page_table_from_manager(_manager(False)).has_host_pools + manager = _manager(True) + manager.impl.init_config.layers[1].buffers[0].is_sparse = False + with pytest.raises(ValueError, match="mixes host and device"): + build_page_table_from_manager(manager) + + +def test_registration_separates_memory_types_and_device_ids() -> None: + table = build_page_table_from_manager(_manager(True)) + worker = SimpleNamespace(_rank_info=_rank_info("gen", table), _agent=Mock(), _registered_mem=[]) + transfer_mod.TransferWorker._register_kv_cache(worker) + assert [d.type for d in worker._registered_mem] == ["VRAM", "DRAM"] + for desc in worker._registered_mem: + assert desc.descs == get_unique_pool_memory_descs(table, 3, desc.type) + assert {entry[2] for entry in desc.descs} == ({0} if desc.type == "DRAM" else {3}) + assert sum(len(d.descs) for d in worker._registered_mem) == 3 + assert worker._agent.register_memory.call_count == 2 + + +@pytest.mark.parametrize("src_offload,dst_offload", [(False, False), (False, True), (True, False)]) +def test_receiver_bypasses_device_bounce_for_host_pools( + src_offload: bool, dst_offload: bool +) -> None: + receiver = object.__new__(transfer_mod.Receiver) + src = _rank_info("ctx", build_page_table_from_manager(_manager(src_offload))) + dst = _rank_info("gen", build_page_table_from_manager(_manager(dst_offload))) + receiver._registrar = SimpleNamespace( + self_rank_info=dst, + self_extractor=KVRegionExtractorV1(dst.page_table), + get_peer_overlap=lambda *_: SimpleNamespace(ranks=[0]), + ) + receiver._enforce_physical_ownership = False + receiver._get_sender_info = Mock(return_value=src) + receiver._bounce = Mock(reserve=Mock(return_value=False)) + receiver._get_session = Mock(return_value=Mock(_sender_endpoints=set())) + receiver._request_sender_data = Mock() + task = SimpleNamespace( + _params=DisaggregatedParams(ctx_dp_rank=0), + _unique_rid=401, + _perf_timer=None, + slice_id=0, + ) + req_info = transfer_mod.RecvReqInfo(1, "gen", 0, [np.array([0])], 401) + receiver._build_recv_req_info = Mock(return_value=req_info) + receiver.dispatch_task(task) + assert receiver._bounce.reserve.call_count == int(not src_offload and not dst_offload) + receiver._request_sender_data.assert_called_once() + + +@pytest.mark.parametrize("offload", [False, True]) +def test_transfer_groups_keep_the_same_byte_ranges(offload: bool) -> None: + src = build_page_table_from_manager(_manager(False)) + dst = build_page_table_from_manager(_manager(offload)) + sender, meta = _sender_and_meta(src, dst) + assert meta.expected_transfers == 1 + assert int(meta.sizes.sum()) == 2 * sum(p.slot_bytes for p in src.pool_groups[0].pools) + assert {(g.src_memory_type, g.dst_memory_type) for g in meta.transfer_groups} == ( + {("VRAM", "DRAM"), ("VRAM", "VRAM")} if offload else {("VRAM", "VRAM")} + ) + assert meta.bounce_dst_base == (None if offload else 0x900000) + request = sender._make_agent_request(meta, 3) + requests = request.requests if offload else (request,) + assert len(requests) == (2 if offload else 1) + assert {r.src_descs.type for r in requests} == {MemoryType.VRAM} + assert {r.dst_descs.type for r in requests} == ( + {MemoryType.DRAM, MemoryType.VRAM} if offload else {MemoryType.VRAM} + ) + + +@pytest.mark.parametrize("outcome", ["success", "wait_failure", "submission_failure"]) +def test_grouped_submission_retains_every_handle_until_settled(outcome: str) -> None: + sender, meta = _sender_and_meta( + build_page_table_from_manager(_manager(False)), + build_page_table_from_manager(_manager(True)), + ) + request = sender._make_agent_request(meta, 3) + first = Mock(wait=Mock(return_value=True), is_completed=Mock(return_value=True)) + second = Mock( + wait=Mock(return_value=outcome == "success"), + is_completed=Mock(return_value=outcome == "success"), + ) + first.last_status_str.return_value = "DONE" + second.last_status_str.return_value = "ERROR" + sender._agent = Mock() + sender._agent.submit_transfer_requests.side_effect = [ + first, + RuntimeError("second submission failed") if outcome == "submission_failure" else second, + ] + task = meta.task + assert task.begin_physical_operation(0) + completed, _ = sender._submit_transfer(task, 0, request) + assert completed == (outcome == "success") + if completed: + assert task.resources_drained + first.wait.assert_called_once() + second.wait.assert_called_once() + else: + operation = task._physical_operations[0] + assert operation.request is request + assert operation.status.statuses[0] is first + assert not task.poll_in_doubt_physical_operation(0) + assert not task.resources_drained + second.is_completed.return_value = True + assert task.poll_in_doubt_physical_operation(0) == (outcome == "wait_failure") + + +@pytest.mark.parametrize("enforce_ownership", [False, True]) +def test_one_logical_completion_after_all_grouped_writes(enforce_ownership: bool) -> None: + sender, meta = _sender_and_meta( + build_page_table_from_manager(_manager(False)), + build_page_table_from_manager(_manager(True)), + ) + task = meta.task + sender._enforce_physical_ownership = enforce_ownership + sender._sessions_lock = threading.Lock() + session = SimpleNamespace( + kv_tasks=[task], lock=threading.Lock(), status=SessionStatus.READY, set_exception=Mock() + ) + sender._sessions = {401: session} + sender._device_id = 3 + sender._instance_rank = 0 + sender._bounce = Mock() + dealer = Mock() + sender._get_result_dealer = Mock(return_value=dealer) + + def wait() -> bool: + assert task.transferred_count == 0 + dealer.send.assert_not_called() + if enforce_ownership: + assert not task.resources_drained + return True + + sender._agent = Mock() + sender._agent.submit_transfer_requests.side_effect = [Mock(wait=wait), Mock(wait=wait)] + if enforce_ownership: + assert task.begin_physical_operation(0) + sender._deliver_kv_to_agent(meta) + assert task.transferred_count == 1 + assert task.resources_drained + assert task.is_done + dealer.send.assert_called_once() + sender._bounce.build_request.assert_not_called() + + +@pytest.mark.parametrize("first_completes", [False, True]) +def test_legacy_grouped_submission_waits_before_submitting_next(first_completes: bool) -> None: + sender, meta = _sender_and_meta( + build_page_table_from_manager(_manager(False)), + build_page_table_from_manager(_manager(True)), + ) + sender._enforce_physical_ownership = False + first = Mock(wait=Mock(return_value=first_completes)) + first.last_status_str.return_value = "ERROR" + + def submit(request): + if sender._agent.submit_transfer_requests.call_count == 1: + return first + first.wait.assert_called_once() + raise RuntimeError("second submission failed") + + sender._agent = Mock(submit_transfer_requests=Mock(side_effect=submit)) + request = sender._make_agent_request(meta, 3) + if first_completes: + with pytest.raises(RuntimeError, match="second submission failed"): + sender._submit_transfer(meta.task, 0, request) + else: + assert sender._submit_transfer(meta.task, 0, request) == (False, "ERROR") + assert sender._agent.submit_transfer_requests.call_count == 1 diff --git a/tests/unittest/disaggregated/test_transfer_ownership_regressions.py b/tests/unittest/disaggregated/test_transfer_ownership_regressions.py index f415e387b063..53ab98df3ac8 100644 --- a/tests/unittest/disaggregated/test_transfer_ownership_regressions.py +++ b/tests/unittest/disaggregated/test_transfer_ownership_regressions.py @@ -849,7 +849,7 @@ def test_cancel_after_publication_cannot_overtake_request_data( protocol_order: list[str] = [] initial_cancel_outcome: str | None = None thread_results: queue.Queue[Exception | None] = queue.Queue() - result_started = threading.Event() + result_recorded = threading.Event() receiver = object.__new__(Receiver) receiver._sessions_lock = threading.Lock() @@ -921,13 +921,13 @@ def cancel() -> None: session.cancel() def finish_writer() -> None: - result_started.set() session.process_kv_agent_result( peer_rank=0, receiver_slice_id=0, is_last_slice=True, status=AgentResult.FAILED, ) + result_recorded.set() receive_thread = _start_checked_thread(receive, thread_results) cancel_thread = None @@ -937,9 +937,8 @@ def finish_writer() -> None: cancel_thread = _start_checked_thread(cancel, thread_results) initial_cancel_outcome = race_outcomes.get(timeout=10) result_thread = _start_checked_thread(finish_writer, thread_results) - assert result_started.wait(timeout=10) - assert result_thread.is_alive(), ( - "terminal writer evidence should wait until REQUEST_DATA publication finishes" + assert result_recorded.wait(timeout=5), ( + "terminal writer evidence must not wait for REQUEST_DATA publication" ) assert not session.resources_drained(), ( "terminal writer evidence authorized reuse before REQUEST_DATA publication finished" diff --git a/tests/unittest/disaggregated/test_unsubmitted_session.py b/tests/unittest/disaggregated/test_unsubmitted_session.py index fbe719015676..53c372b00a88 100644 --- a/tests/unittest/disaggregated/test_unsubmitted_session.py +++ b/tests/unittest/disaggregated/test_unsubmitted_session.py @@ -47,6 +47,7 @@ def _sender(*, ownership: bool = True) -> transfer_mod.Sender: sender._pre_cancelled_rids = {} sender._peer_requests_lock = threading.Lock() sender._peer_requests, sender._peer_requests_timestamps = {}, {} + sender._request_info_ttl_s = transfer_mod.Sender._STALE_REQ_INFO_TTL_S sender._shutdown, sender._shutdown_requested = True, False sender._retirement_watchdog = None sender._instance_rank = 7 @@ -139,6 +140,32 @@ def test_metadata_expiry_does_not_prevent_cancel_acknowledgment() -> None: assert not sender._sessions +@pytest.mark.parametrize("transfer_timeout_s", [None, 60.0, 600.0]) +def test_queued_prefill_metadata_uses_transfer_timeout( + monkeypatch: pytest.MonkeyPatch, transfer_timeout_s: float | None +) -> None: + monkeypatch.setattr(transfer_mod, "ZMQMessenger", Mock()) + monkeypatch.setattr(transfer_mod, "KV_TRANSFER_NUM_THREADS", 0) + now = 0.0 + monkeypatch.setattr(transfer_mod.time, "monotonic", lambda: now) + sender = transfer_mod.Sender(Mock(), Mock(), request_info_timeout_s=transfer_timeout_s) + info = transfer_mod.RecvReqInfo.from_bytes(_request_data()[1]) + try: + sender._add_req_info(401, 0, info) + now = 121.0 + sender.sweep_stale_req_infos() + if transfer_timeout_s == 600.0: + assert sender._get_req_info(401) == {0: info} + else: + assert sender._get_req_info(401) is None + + now = 601.0 + sender.sweep_stale_req_infos() + assert sender._get_req_info(401) is None + finally: + sender.shutdown() + + @pytest.mark.parametrize("pending_settlement", [False, True]) @pytest.mark.parametrize("ack_fails_once", [False, True]) def test_worker_orders_session_ack_after_piece_reports( 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 829aab663454..3f9966158604 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 @@ -1647,7 +1647,7 @@ def test_unconfigured_life_cycle_has_no_cache_state(self): class TestSparseResolve(unittest.TestCase): - """The sparse resolve kernel: probe index, victim choice and recency rebuild.""" + """The sparse resolve kernel: probe index, victim choice and the ages it leaves.""" def setUp(self) -> None: init_cuda_once() @@ -1662,7 +1662,7 @@ def _resolve( selected: list[int], *, num_layers: int = 1, - lru: list[int] | None = None, + ages: list[int] | None = None, history_units: int = 1 << 20, capacity_units: int | None = None, max_in_flight_units: int = 0, @@ -1671,13 +1671,11 @@ def _resolve( ): 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, + ages or [], num_layers, entries, selected, @@ -1710,29 +1708,33 @@ def test_padding_resolves_to_no_slot(self): self.assertEqual(result.out_slot[1:], [-1, -1]) def test_hits_outrank_units_fetched_this_step(self): + """Both entries are used this step, and the history flag is what separates them.""" # 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") + self.assertEqual(result.entry_age[0], 0) + self.assertEqual(result.entry_hit[0], 1, "the entry already held its unit") + self.assertEqual(result.entry_age[fetched], 0) + self.assertEqual(result.entry_hit[fetched], 0, "the unit was fetched into the entry") - def test_victims_come_from_the_lru_tail(self): - # A full cache, with entry 2 placed last in recency so it is the victim. + def test_victims_are_the_stalest_entries(self): + # A full cache, with entry 2 the oldest so it is the victim. resident = [10, 11, 12, 13] - lru = [0, 1, 3, 2] - result = self._resolve(resident, [99], lru=lru) + result = self._resolve(resident, [99], ages=[1, 1, 9, 1]) 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. + # The units that were not selected keep their data, and age by nothing: the step + # clock only advances between resolves. self.assertEqual([result.resident_ordinal[i] for i in (0, 1, 3)], [10, 11, 13]) + self.assertEqual([result.entry_age[i] for i in (0, 1, 3)], [1, 1, 1]) 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) + ages = [1, 1, 9, 1] # entry 2 (unit 12) is the oldest + result = self._resolve(resident, [99], ages=ages, 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) @@ -1840,35 +1842,20 @@ def test_layers_resolve_independently(self): 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}" - ) + def test_large_caches_evict_the_free_entries_then_the_stalest(self): + """Free entries go first, then the oldest, and every other entry is left alone.""" 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 + count = min(entries // 4, 512) + num_free = min(entries // 8, count // 2) resident = list(range(entries - in_flight)) + list( range(history, history + in_flight) ) - resident[: entries // 8] = self._free(entries // 8) + resident[:num_free] = self._free(num_free) 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 ) @@ -1880,18 +1867,22 @@ def test_large_caches_preserve_victims_and_recency_order(self): 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] - ) + free_slots = {slot for slot, ordinal in enumerate(resident) if ordinal < 0} + # Two ages, far enough apart that the cache has to take the older group + # whole before it reaches anything else. + eligible = [ + slot + for slot in range(entries) + if slot not in hit_slots and slot not in free_slots + ] + stale = set(rng.sample(eligible, count - num_free)) + ages = [9 if slot in stale else 1 for slot in range(entries)] + victims = free_slots | stale try: result = self._resolve( resident, selected, - lru=lru, + ages=ages, history_units=history, capacity_units=history + in_flight, max_in_flight_units=in_flight, @@ -1901,7 +1892,6 @@ def test_large_caches_preserve_victims_and_recency_order(self): 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: @@ -1914,40 +1904,42 @@ def test_large_caches_preserve_victims_and_recency_order(self): for slot, ordinal in enumerate(resident): if slot not in victims: self.assertEqual(result.resident_ordinal[slot], ordinal) + # Every entry the step used is now the youngest there is, and the flag + # separates the ones that were already holding their unit. + for slot in range(entries): + if slot in victims: + self.assertEqual((result.entry_age[slot], result.entry_hit[slot]), (0, 0)) + elif slot in hit_slots - free_slots: + self.assertEqual((result.entry_age[slot], result.entry_hit[slot]), (0, 1)) + else: + self.assertEqual(result.entry_age[slot], ages[slot]) 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.""" + def test_colliding_selections_preserve_slots_and_eviction(self) -> None: + """Sparse ordinals and padding preserve the same eviction 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) + history = 200000 + resident = [i * 64 for i in range(entries - 2)] + [history, history + 1] count = min(entries // 4, 32) hits = resident[:count] - misses = [(entries + i) * 2048 for i in range(count)] + misses = [(entries + i) * 64 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] - ) + eligible = [slot for slot in range(entries) if slot not in hit_slots] + victims = set(random.Random(entries).sample(eligible, count)) + ages = [9 if slot in victims else 1 for slot in range(entries)] result = self._resolve( resident, selected, - lru=lru, + ages=ages, 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]) @@ -4496,11 +4488,11 @@ def test_a_fetch_neither_moves_nor_overwrites_the_prepared_cell(self): 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 + It does not cover the separate held-entry 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. + as matched and never as evictable -- which the victim test already excludes. Dropping + the held term 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() diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index ede5add24e59..5145219b5915 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -62,9 +62,10 @@ PrefillCudaGraphBackend, PybindMirror, RayPlacementConfig, SkipSoftmaxAttentionConfig, - SleepConfig, SpeculativeConfig, - StrictBaseModel, TorchCompileConfig, - TorchLlmArgs, + SleepConfig, + SparseKvCacheOffloadConfig, + SpeculativeConfig, StrictBaseModel, + TorchCompileConfig, TorchLlmArgs, UserProvidedDecodingConfig, update_llm_args_with_extra_dict, update_llm_args_with_extra_options) @@ -4807,6 +4808,87 @@ def test_ckpt_sparse_attention_config_can_be_passed_directly(self): 100.0 * math.exp(5.0 * 0.5)) +@pytest.mark.cpu_only +@pytest.mark.parametrize("config_cls", [ + llm_args_mod.DeepSeekSparseAttentionConfig, + DeepSeekV4SparseAttentionConfig, +]) +class TestSparseOffloadIndexerDefault: + + @pytest.mark.parametrize("enabled", [False, True]) + @pytest.mark.parametrize("requested", [None, False, True]) + def test_default_override_and_round_trip( + self, + config_cls: type[llm_args_mod.DeepSeekSparseAttentionConfig], + enabled: bool, + requested: bool | None, + ) -> None: + kwargs = {"enable_kv_cache_offload": enabled} + if requested is not None: + kwargs["use_cute_dsl_paged_mqa_logits"] = requested + expected = not enabled if requested is None else requested + config = config_cls(**kwargs) + + for resolved in ( + config, + config_cls.model_validate_json(config.model_dump_json()), + config_cls.model_validate( + config.model_dump(exclude_unset=True)), + ): + assert resolved.use_cute_dsl_paged_mqa_logits is expected + assert resolved.to_sparse_params( + ).use_cute_dsl_paged_mqa_logits is expected + assert resolved.to_sparse_metadata_params( + ).use_cute_dsl_paged_mqa_logits is expected + + @pytest.mark.parametrize("enabled", [False, True]) + def test_yaml_bool_coercion( + self, + config_cls: type[llm_args_mod.DeepSeekSparseAttentionConfig], + enabled: bool, + ) -> None: + config = config_cls(enable_kv_cache_offload=str(enabled).lower()) + assert config.use_cute_dsl_paged_mqa_logits is not enabled + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("enabled", [False, True]) +@pytest.mark.parametrize("requested", [None, False, True]) +def test_glm_model_config_preserves_indexer_choice( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + enabled: bool, + requested: bool | None, +) -> None: + checkpoint = PretrainedConfig( + architectures=["GlmMoeDsaForCausalLM"], + num_hidden_layers=7, + index_n_heads=32, + index_head_dim=128, + index_topk=2048, + ) + monkeypatch.setattr( + "tensorrt_llm._torch.model_config.load_pretrained_config", + lambda *args, **kwargs: checkpoint) + kwargs = {"enable_kv_cache_offload": enabled} + if requested is not None: + kwargs["use_cute_dsl_paged_mqa_logits"] = requested + config = llm_args_mod.DeepSeekSparseAttentionConfig(**kwargs) + + resolved = ModelConfig.from_pretrained( + str(tmp_path), + sparse_attention_config=config, + attn_backend="TRTLLM", + moe_backend="TRTLLM", + ).sparse_attention_config + + expected = not enabled if requested is None else requested + assert resolved.use_cute_dsl_paged_mqa_logits is expected + assert resolved.to_sparse_params().use_cute_dsl_paged_mqa_logits is expected + assert resolved.to_sparse_metadata_params( + ).use_cute_dsl_paged_mqa_logits is expected + + @pytest.mark.cpu_only class TestDeepSeekV4SparseAttentionConfig: @@ -4867,6 +4949,54 @@ def test_kv_cache_offload_requires_positive_topk(self, index_topk): DeepSeekV4SparseAttentionConfig(index_topk=index_topk, enable_kv_cache_offload=True) + def test_kv_cache_offload_size_defaults_round_trip(self): + config = DeepSeekV4SparseAttentionConfig() + + assert config.kv_cache_offload_config == SparseKvCacheOffloadConfig() + assert config.kv_cache_offload_config.gpu_cache_factor == 2.0 + assert config.kv_cache_offload_config.host_history_size is None + assert config.get_kv_cache_offload_config() is None + # Defaults stay valid when a dump passes every field back explicitly. + DeepSeekV4SparseAttentionConfig(**config.model_dump()) + + def test_kv_cache_offload_sizes_round_trip(self): + config = DeepSeekV4SparseAttentionConfig(enable_kv_cache_offload=True, + kv_cache_offload_config={ + "gpu_cache_factor": 3.0, + "host_history_size": + 1 << 30 + }) + + restored = DeepSeekV4SparseAttentionConfig.model_validate_json( + config.model_dump_json()) + + assert restored.get_kv_cache_offload_config( + ) == SparseKvCacheOffloadConfig(gpu_cache_factor=3.0, + host_history_size=1 << 30) + + @pytest.mark.parametrize("field, value", [("gpu_cache_factor", 0.5), + ("host_history_size", -1)]) + def test_kv_cache_offload_sizes_reject_out_of_range(self, field, value): + with pytest.raises(ValidationError, match=field): + SparseKvCacheOffloadConfig(**{field: value}) + + @pytest.mark.parametrize("field, value", [("gpu_cache_factor", 3.0), + ("host_history_size", 1 << 30)]) + def test_kv_cache_offload_sizes_require_offload(self, field, value): + with pytest.raises(ValidationError, + match="requires enable_kv_cache_offload"): + DeepSeekV4SparseAttentionConfig( + kv_cache_offload_config={field: value}) + + def test_only_offloading_configs_report_kv_cache_offload(self): + assert llm_args_mod.DeepSeekSparseAttentionConfig( + ).get_kv_cache_offload_config() is None + assert SkipSoftmaxAttentionConfig().get_kv_cache_offload_config( + ) is None + config = DeepSeekV4SparseAttentionConfig(enable_kv_cache_offload=True) + assert config.get_kv_cache_offload_config( + ) is config.kv_cache_offload_config + @pytest.mark.parametrize("compress_ratios", [[], [-1, 4, 128]]) def test_invalid_compress_ratios_raise(self, compress_ratios): with pytest.raises(ValidationError, match="compress_ratios"): diff --git a/tests/unittest/llmapi/test_trtllm_serve_config_schemas.py b/tests/unittest/llmapi/test_trtllm_serve_config_schemas.py index 5b9dbaa91d93..5bae7af2801d 100644 --- a/tests/unittest/llmapi/test_trtllm_serve_config_schemas.py +++ b/tests/unittest/llmapi/test_trtllm_serve_config_schemas.py @@ -318,6 +318,48 @@ def test_llm_model_fields_have_schema(schemas: dict[str, dict]) -> None: _assert_model_fields_in_schema(TorchLlmArgs, schema["properties"]) +@pytest.mark.parametrize("filename", [generator.SERVE_SCHEMA, generator.DISAGG_SCHEMA]) +@pytest.mark.parametrize( + ("config_name", "algorithm"), + [("DeepSeekSparseAttentionConfig", "dsa"), ("DeepSeekV4SparseAttentionConfig", "deepseek_v4")], +) +@pytest.mark.parametrize("enabled", [None, False, True]) +def test_sparse_offload_indexer_default( + schemas: dict[str, dict], filename: str, config_name: str, algorithm: str, enabled: bool | None +) -> None: + from tensorrt_llm.llmapi.llm_args import ( + DeepSeekSparseAttentionConfig, + DeepSeekV4SparseAttentionConfig, + ) + + config_cls = { + "dsa": DeepSeekSparseAttentionConfig, + "deepseek_v4": DeepSeekV4SparseAttentionConfig, + }[algorithm] + sparse_config: dict[str, str | bool] = {"algorithm": algorithm} + if enabled is not None: + sparse_config["enable_kv_cache_offload"] = enabled + schema = schemas[filename] + definition = schema["$defs"][config_name] + field = "use_cute_dsl_paged_mqa_logits" + assert "default" not in definition["properties"][field] + condition = definition["allOf"][0] + branch = "then" if Draft202012Validator(condition["if"]).is_valid(sparse_config) else "else" + expected = enabled is not True + assert condition[branch]["properties"][field]["default"] is expected + assert config_cls.model_validate(sparse_config).use_cute_dsl_paged_mqa_logits is expected + validator = Draft202012Validator(schema) + assert validator.is_valid({"sparse_attention_config": sparse_config}) + for value in (False, True): + explicit = {**sparse_config, field: value} + assert validator.is_valid({"sparse_attention_config": explicit}) + assert config_cls.model_validate(explicit).use_cute_dsl_paged_mqa_logits is value + for invalid in (None, "invalid", 2): + assert not validator.is_valid( + {"sparse_attention_config": {**sparse_config, field: invalid}} + ) + + def test_model_field_guard_detects_missing_fields_and_aliases() -> None: class Config(BaseModel): plain: int diff --git a/token_parity.md b/token_parity.md new file mode 100644 index 000000000000..215d7b17c2c9 --- /dev/null +++ b/token_parity.md @@ -0,0 +1,295 @@ + + +The DeepSeek V4 Flash token-parity fixes make **TopK, MHC, and FMHA honor +`FORCE_DETERMINISTIC=1`**. Repeated resident-cache runs were already producing +different tokens, so the resident/offload comparison needed a repeatable baseline. +The implementation is split into a disaggregated sparse-offload commit for +transport and configuration, followed by a general determinism commit for +TopK, MHC, FMHA, kernel-specific controls, tests, and this report. The validation +below was performed before splitting the commits. + +All five cumulative feature stages pass exact resident-versus-CPU-offload token +parity. Including the fixed baseline and profiling probes, the successful runs +cover **98 requests and 15,680 generated tokens**, with **36 exact sequential +pairs and 8 exact concurrent pairs**. These results compare matching feature +configurations; they do not establish identical tokens across every configuration. +Block reuse is exercised on the context worker only. +[Final validation report](../sparse-disagg-e2e-artifacts/sequential-20261017/final-report.json). + +Tracing repeated resident runs found identical TopK scores with different index +orders. A 50-repeat kernel check reproduced native ordering differences; tied +scores also changed the selected set. The PyTorch implementation had no ordering +or selection mismatches in those checks. Separately, MHC split-K reductions +changed intermediate mixing values by approximately `1e-7` to `3e-7`, enough to +propagate into different generated tokens. + +- **TopK:** Native atomic appends changed index ordering and, with tied scores, + sometimes changed which indices were selected. Deterministic mode forces the + PyTorch implementation for both prefill and decode, which produced repeatable + selections and ordering in the tested configuration. + [Implementation](tensorrt_llm/_torch/modules/top_k.py). +- **MHC:** Atomic split-K reductions introduced small floating-point differences + that propagated into token divergence. Deterministic mode forces a fixed, + single-split tactic, gives deterministic execution a separate tuning-cache + identity, and overrides cached or explicitly supplied tactics. It uses the + half-fused MMA implementation where supported, with a single-split FMA fallback. + [Implementation](tensorrt_llm/_torch/modules/mhc/mhc_cuda.py). +- **FMHA:** Deterministic mode fixes the KV reduction to one partition so eager + execution and CUDA graph capture select the same reduction layout. + [Implementation](cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h). + +Sparse-offload settings are also preserved during checkpoint loading so the +offload run actually uses the CPU hot tier. +[Configuration fix](tensorrt_llm/_torch/model_config.py). + +Validation passed: **43 kernel regression tests** and **6/6 exact resident/offload +token matches**, covering 12 requests and 1,920 generated tokens on the full +43-layer DeepSeek V4 Flash model. The comparison used four B200 GPUs, NVFP4 KV +cache, context and generation TP2/EP2 with attention TP, the NIXL Python +transceiver, `FORCE_DETERMINISTIC=1`, and disabled autotuning. Prompt lengths were +129, 8,193, 127, 128, 2,049, and 8,193 tokens, with 160 output tokens per request. + +That initial comparison used eager execution with CUDA graphs, overlap scheduling, +chunked prefill, and block reuse disabled. Attention DP had encountered a CUDA +illegal access during startup memory profiling; the subsequent validation below +covers its current status. +The passing run used the production implementation without diagnostic overrides. +Local evidence is available in the +[root-cause report](../sparse-disagg-e2e-artifacts/parity/root-cause-evidence.json) +and the [validation report](../sparse-disagg-e2e-artifacts/refresh-20261017T173731Z/final-report.json). + +Sequential validation used the original combined commit `432a1ba93d` plus the +fixes described below, before splitting the changes into separate commits. +The following features were enabled cumulatively. Each stage must match resident +versus offloaded tokens with the same enabled features before the next feature +is enabled. +Cross-configuration comparisons are also recorded; exact TP/DP token identity is +a separate criterion because their low-precision reduction layouts differ. +The same six prompts, 160 output tokens, full checkpoint, NVFP4 format, and +deterministic kernel choices are retained. Runtime profiles and worker statistics +verify that the enabled features execute. + +| Stage | Newly enabled feature | Status | +| --- | --- | --- | +| 1 | CUDA graphs, decode batch sizes 1 and 2 | Passed: 6/6 resident/offload and eager matches | +| 2 | Generation attention DP | Passed: 6/6 sequential and 2/2 concurrent resident/offload matches | +| 3 | Chunked prefill, 2,048-token scheduling budget | Passed: 6/6 sequential and 2/2 concurrent matches; schedules identical | +| 4 | Context-worker block reuse, including the repeated 8,193-token prompt | Passed: 6/6 sequential and 2/2 concurrent matches; reuse verified | +| 5 | Overlap scheduler | Passed: 6/6 sequential and 2/2 concurrent matches; all sequential outputs also match stage 4 | + +The initial CUDA-graph resident run matched both 8,193-token requests and the +2,049-token request, but differed for the 127-, 128-, and 129-token prompts. +An eager resident/offload control on the same allocation matched all six prior +reference outputs and all six resident/offload pairs. A focused deterministic +TopK eager/graph comparison passed 22 cases. + +Full-model tracing located the first difference at generation layer 2's attention +output. Earlier layer outputs, queries, and TopK selections matched. Profiles +showed that graph capture selected a CGA multi-CTA KV reduction despite +`FORCE_DETERMINISTIC=1`; eager short requests selected a single-CTA reduction. +Running eager execution with multiple CUDA streams still matched the reference. +The FMHA selector now forces `MultiCtasKvMode::Disabled` and +`mMaxNumCtasPerSeqKv = 1` in deterministic mode. This prevents its heuristics from +overriding the single-partition choice when graph capture supplies a longer KV +length than the live request. +[Implementation](cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h). + +Both new native regressions failed before the fix and pass after it. They check +kernel selection with 32 and 64 query heads and exact eager/graph output equality +at seven live KV lengths from 127 to 8,193 tokens, with a graph captured at 9,216. +[Native tests](cpp/tests/unit_tests/kernels/fmhaDeterministicTest.cpp). +The rebuilt libraries pass a fresh eager resident/offload comparison: **6/6 exact +token matches**, with repeatable outputs for the repeated long prompt in both +modes. **Stage 1 passes:** all six CUDA-graph outputs match the corresponding eager +outputs in both resident and offload modes, and all six resident/offload pairs +match. Each mode's profiled repeat also matches its earlier long-prompt output. +Profiles show three graph launches on each generation rank, using only the +single-partition FMHA kernel in the captured decode steps. +[Diagnostic evidence](../sparse-disagg-e2e-artifacts/sequential-20261017/cuda-graph-root-cause.json). + +The fix also changes eager reduction choices, so five of the six eager outputs +differ from the old build. The 129-token prompt remains identical; the other +short prompts diverge late in generation. Subsequent feature comparisons use this +fixed implementation in both runs. +[Reference changes](../sparse-disagg-e2e-artifacts/sequential-20261017/fmha-reference-change.json). + +Stage 2 completes attention-DP startup without the earlier CUDA illegal access, +but all six resident outputs differ from stage 1. The first differing output-token +indices are 28, 18, 24, 24, 32, and 18, respectively. Repeated outputs are stable. +Tracing found exact equality in the first layer's input, quantized queries, and +native attention output. The first difference is its output projection: TP2 +rounds each partial GEMM result to BF16 before reducing the two results, while DP +computes the full projection locally. Replaying captured tensors reproduces both +paths exactly: 1,476 of 4,096 output elements differ, with a maximum absolute +difference of 0.03125. Graph and eager DP runs produce identical tokens. +[Projection evidence](../sparse-disagg-e2e-artifacts/sequential-20261017/attention-dp-projection.json). + +The current E2E criterion compares resident and offloaded execution within each +feature configuration. The TP/DP numerical difference is retained in the report. +Sequential requests ran on DP rank 0. Two concurrent long requests exercised both +ranks; one matched the sequential output and one differed starting at output +token 27. The sequential outputs and concurrent output pair both repeated exactly +across two independent resident runs. The harness records this sensitivity +separately and requires the pair +of concurrent outputs to match between resident and offload modes, allowing the +two identical requests to exchange rank assignments. **Stage 2 passes:** all six +sequential outputs and both concurrent outputs match exactly between resident and +offload modes. Worker statistics confirm active requests on DP ranks 0 and 1 in +both modes, and profiles confirm graph replay on both ranks. +[Attention-DP validation](../sparse-disagg-e2e-artifacts/sequential-20261017/attention-dp-root-cause.json). + +The stage 3 resident run verifies chunked execution: 29 context iterations, at +most 2,048 scheduled context tokens per iteration, and 21 iterations with cached +context KV. Five of the six sequential outputs match stage 2; the 2,049-token +prompt differs at exactly one output token, index 28. All six sequential +resident/offload pairs match exactly. The initial concurrent pair differed, and +worker statistics showed different prefill chunk schedules: the second resident +request started with 1,920 tokens, while the corresponding offload request +started with 2,048. The harness now submits the second request after the first +streaming token. This preserves concurrent decode on both DP ranks while keeping +prefill scheduling serial and comparable. **Stage 3 passes on rerun:** all six +sequential pairs and both concurrent outputs match exactly, with identical +resident/offload chunk schedules: 30 context iterations, at most 2,048 tokens per +iteration, and 21 iterations with cached context KV. Both DP ranks execute real +requests, and graph replay remains verified. +[Scheduling evidence](../sparse-disagg-e2e-artifacts/sequential-20261017/chunked-prefill-probe-schedule.json). + +Stage 4's first attempt did not exercise reuse: `FORCE_DETERMINISTIC=1` +overrides the requested cache configuration and disables it. The test correctly +failed its runtime feature check despite stable tokens. The rerun uses the +existing attention, MoE, and all-reduce determinism switches plus new +`FORCE_TOPK_KERNEL_DETERMINISTIC=1` and `FORCE_MHC_KERNEL_DETERMINISTIC=1` +switches, with global `FORCE_DETERMINISTIC=0`. This preserves the deterministic +kernel choices while allowing prefix reuse. The global flag's behavior is +unchanged. Stages 4 and 5 use: + +```bash +export FORCE_DETERMINISTIC=0 +export FORCE_ATTENTION_KERNEL_DETERMINISTIC=1 +export FORCE_MOE_KERNEL_DETERMINISTIC=1 +export FORCE_ALL_REDUCE_DETERMINISTIC=1 +export FORCE_TOPK_KERNEL_DETERMINISTIC=1 +export FORCE_MHC_KERNEL_DETERMINISTIC=1 +``` + +Reuse is enabled on the context worker in both modes; generation +receives prompt KV through disaggregation. Sparse generation's explicit +restriction on local prefix reuse remains in place, so this test does not +establish support for generation-worker reuse of CPU-backed prefixes. +The 24 focused TopK/MHC regressions pass with the global and kernel-specific +controls, including tied-score selection and cached split-K tactic overrides. +**Stage 4 passes:** all six sequential pairs and both concurrent outputs match +exactly between resident and offload. Both runs verify actual reuse: 276 reused +blocks and 12 context iterations, down from 30 without reuse, with identical +chunk schedules. Repeated long prompts reuse 8,192 cached tokens and compute one +prompt token; their outputs and the profiled repeat remain identical. Five +sequential outputs match stage 3; the 127-token prompt differs beginning at output +token 24. This cross-configuration difference remains recorded separately. +[Reuse-control evidence](../sparse-disagg-e2e-artifacts/sequential-20261017/block-reuse-controls.json). + +**Stage 5 passes:** all six sequential pairs and both concurrent outputs match +exactly between resident and offload with overlap scheduling enabled. All six +sequential outputs also match stage 4 in each mode; the repeated prompt and +profiled repeat remain identical. Both modes report the overlap scheduler, +active requests on both attention-DP ranks, 276 reused blocks, and identical +context schedules. CUDA traces show three graph launches on each generation rank. +No further production change was needed for overlap in this configuration. +[Overlap validation](../sparse-disagg-e2e-artifacts/sequential-20261017/overlap-scheduler-validation.json). + +Run artifacts are stored locally in +`../sparse-disagg-e2e-artifacts/sequential-20261017/` relative to this worktree. + +**32K sparse-offload capacity validation, October 17–18, 2026 (UTC).** +The full 43-layer DeepSeek V4 Flash model sustained **1,134 simultaneous decode +requests, 567 per attention-DP rank**, with exactly **32,768 input tokens and +160 output tokens per request** on four B200 GPUs. This is the measured capacity +with 32 GiB of CPU sparse-history memory per generation rank, +`free_gpu_memory_fraction=0.9`, and generation `max_batch_size=1024` per rank. +It is not an unrestricted hardware maximum or a throughput benchmark. + +The two offload trials at 1,134 both completed and produced identical output +multisets. Offering 1,136, 1,140, or 1,182 requests still peaked at 1,134 active +decodes; all offered requests eventually returned 160 tokens. The clean 1,136 +probe released its first batch after 1,134 receives. The exploratory 1,140 and +1,182 probes required recorded manual barrier adjustments. Startup's aggregate +estimate of 591 requests per rank (1,182 total) exceeded the observed usable +capacity; the larger gated batches encountered attention-DP dummy-allocation +pressure. +[Capacity evidence](../sparse-disagg-e2e-artifacts/long-context-20261017T231142Z/offload-capacity-summary.json). + +The prompts are **synthetic and identical**: a repeated astronomy paragraph and +a fixed question, with BOS included in the 32,768-token count. Context prefix +reuse is enabled; generation reuse is disabled, so each decode request receives +its own KV through disaggregation. Ratio-4 sparse layers have 8,192 compressed +history rows at this input length. Their GPU fetch cache holds 1,056 rows with +TopK 512, `gpu_cache_factor=2`, and 128-token pages. The complete history is +therefore larger than the GPU fetch cache and uses the CPU hot tier. Dense pools +and ratio-128 history remain on the GPU. + +The topology is context TP2/EP2 and generation TP2/EP2 with attention DP2. Both +runs use NVFP4 KV, NIXL Python transfer, CUDA graphs, overlap scheduling, and +2,048-token chunked prefill. Greedy sampling requests exactly 160 tokens and +ignores EOS. The kernel-specific determinism flags listed above remain enabled, +with global `FORCE_DETERMINISTIC=0` and autotuning disabled. + +Exact parity compares output-token **multisets**, allowing identical prompts to +exchange attention-DP rank assignments. A test-only receive-publication barrier +starts matching decode batches, and context `max_batch_size=1` fixes prefill +shapes. The barrier leaves transfers, cache operations, model kernels, and +sampling unchanged. Each passing trial must also sustain its full requested +decode concurrency for all 159 decode iterations; the first output token comes +from context execution. + +| Offered requests | Exact resident/offload matches | Active decode requests | +| --- | --- | --- | +| 1, cold prefix | 1/1 | 1 | +| 1, reused prefix | 1/1 | 1 | +| 512 | 512/512 | 512 | +| 1,024 | 1,024/1,024 | 1,024 | +| 1,134, first trial | 1,134/1,134 | 1,134 | +| 1,134, repeat | 1,134/1,134 | 1,134 | + +All six comparisons pass: **3,806 matched request pairs and 608,960 exactly +matched output tokens per mode**. Both maximum-size repeats sustain 567 decodes +per rank for all 159 decode iterations with zero paused requests. The repeated +output multisets also match within each mode. + +The GPU-resident reference needs different allocation settings to hold the same +batch: `free_gpu_memory_fraction=0.99`, `max_util_for_resume=1.0`, +`fp8_context_mla_kv_len_cap=33024`, and `avg_seq_len=32770`. These change memory +allocation, while model computation, input/output lengths, graph sizes, and +sampling remain matched. The last setting sizes windowed pools for decode +crossing a page boundary. The default typical sequence ends exactly on a page +boundary and models one window page where this workload can need two. Earlier +resident attempts paused requests despite spare full-history blocks, changing +the decode batch and failing the comparison. Those attempts are preserved. +[Pool-sizing evidence](../sparse-disagg-e2e-artifacts/long-context-20261017T231142Z/resident-pool-sizing-diagnostic.json). + +**Ordinary asynchronous high-concurrency parity remains unresolved.** Initial +runs without the barrier showed token differences as their batch schedules +changed. The controlled results establish resident/offload parity for matching +batches, not invariance to arbitrary request arrival or batching. Warm prefix +reuse also makes this a capacity test with repeated prompts rather than a test +of sustained throughput on distinct 32K documents. +[Asynchronous comparison](../sparse-disagg-e2e-artifacts/long-context-20261017T231142Z/initial-comparison.json). + +The larger test exposed one additional **general disaggregation fix**. In +generation-first serving, receive descriptors could expire after a fixed 120 +seconds while prefill remained queued, even with a 600-second transfer timeout. +`Sender` now retains them for `max(120 seconds, configured transfer timeout)`. +The regression checks the default lifetime, a shorter configured timeout, and +retention plus eventual expiry for a longer timeout. Validation passed +**90 tests, with one skipped**, across the unsubmitted-session and bounded-polling +suites. This fix concerns transfer metadata lifetime and is not specific to +sparse KV offload. +[Implementation](tensorrt_llm/_torch/disaggregation/native/transfer.py), +[regression tests](tests/unittest/disaggregated/test_unsubmitted_session.py), +[validation evidence](../sparse-disagg-e2e-artifacts/long-context-20261017T231142Z/regression-validation.json). + +Complete results, configuration files, source hashes, failed attempts, and +iteration metrics are linked from the +[32K final report](../sparse-disagg-e2e-artifacts/long-context-20261017T231142Z/final-report.json) +and [exact comparison](../sparse-disagg-e2e-artifacts/long-context-20261017T231142Z/verified-comparison.json). From a9862f2aeea145c4f82c435187c98330faaffb92 Mon Sep 17 00:00:00 2001 From: Guiju Zhang <7135567+cascade812@users.noreply.github.com> Date: Thu, 29 Oct 2026 14:55:23 -0700 Subject: [PATCH 02/11] [TRTLLM-16282][test] Stop passing is_decoding to _KVCache.resume in the mixed-precision tests _KVCache.resume no longer takes is_decoding; #21129 kept it only as an ignored keyword until the full sparse attention support landed. The cached context and decode tests still passed it, so both would raise TypeError. Dropping the argument leaves their behaviour unchanged. Signed-off-by: Guiju Zhang <7135567+cascade812@users.noreply.github.com> --- .../_torch/attention/test_mixed_precision_backends.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/tests/unittest/_torch/attention/test_mixed_precision_backends.py b/tests/unittest/_torch/attention/test_mixed_precision_backends.py index 1c22dc3bd18e..f589afb0afee 100644 --- a/tests/unittest/_torch/attention/test_mixed_precision_backends.py +++ b/tests/unittest/_torch/attention/test_mixed_precision_backends.py @@ -267,9 +267,7 @@ def token_major(pool): cache = manager.kv_cache_map[request_id] assert cache.resize(None, begin) cache.suspend() - assert cache.resume( - torch.cuda.current_stream().cuda_stream, is_decoding=not context - ) + assert cache.resume(torch.cuda.current_stream().cuda_stream) assert cache.resize(end) manager._restore_page_index_bufs(request_id, cache) history_before = token_major(manager.get_buffers(0, kv_layout=kv_layout)).clone() @@ -468,9 +466,7 @@ def test_mla_cached_context_and_decode(backend): cache = manager.kv_cache_map[0] assert cache.resize(None, begin) cache.suspend() - assert cache.resume( - torch.cuda.current_stream().cuda_stream, is_decoding=not context - ) + assert cache.resume(torch.cuda.current_stream().cuda_stream) assert cache.resize(end) manager._restore_page_index_bufs(0, cache) metadata = backend.Metadata( From 84b8c005747f01531d1c6fca41f27b957b960548 Mon Sep 17 00:00:00 2001 From: Guiju Zhang <7135567+cascade812@users.noreply.github.com> Date: Thu, 29 Oct 2026 20:20:27 -0700 Subject: [PATCH 03/11] [TRTLLM-16282][fix] Address the review findings on sparse KV offload Runtime: - Hold the host history pools to the native manager's floor: the pages of its largest constraint batch at max_util_for_resume, plus one pool granule per host pool. An explicit host_history_size below the floor is rejected with what to change; the automatic budget is raised to it. An offloading worker's typical step is a decode step, so it adds no context-warmup constraint. - Reserve host history against the capacity synced across ranks, and read host memory limits from effective_available_host_memory(). - Reject pipeline parallelism: the attention Batch seats one micro-batch's requests at a time. - Count the converted index tables and DSA's attention indices in the device reservation. - Disaggregated transfer: only generation-first REQUEST_DATA waits for the transfer timeout. Context-first REQUEST_DATA is failed after the base TTL, and the warning reports how long it waited. The page table takes each pool's memory type from the native storage_medium. - Size host-backed pools from the native sparse_backing_medium() and get_num_bytes() instead of a Python mirror of the layer config. - PyExecutor declares _is_kv_manager_v2 = False on the class, so the forward step reads it on executors the tests build without __init__. Consolidation: - One sparse_arena_rows Triton function and input check serve the DSA and DeepSeek-V4 index merges. TLLM_SPARSE_OFFLOAD_DEBUG_ASSERT=1 checks that the merged rows cover every selected block. - get_scratch_entry() replaces three copies of the scratch lookup, and DeepSeek-V4 takes its scratch slot count from the native manager. - DeepSeek-V4 offload uses DSA's batch validation (decode only), and its config check runs once, in the base manager. SparseOffloadState drops is_prefill. - Remove the Python initial_pool_ratio override, sparse_offload_pool_ratio and LifeCyclePage.pages_per_request. - Warmup uses one generation_only flag for Helix CP and decode-only workers. - _raise_on_every_rank uses _sync_kv_cache_manager_init_status. - DeepSeekV4SparseAttentionConfig inherits enable_kv_cache_offload. Tests and CI: - Fix the phase-selective forward stub and the transceiver config of the decode-only creator test. - Add enable_kv_pool_rebalance to the KvCacheConfigV2 stubs: the offload config check reads it off kv_cache_config. - Build fmhaDeterministicTest with add_tllm_gtest and FULL_STACK. - The DSA offload gate checks for the SM100 family first. - Remove token_parity.md; the PR description summarizes it. Docs: the llm_args descriptions (schemas regenerated), the sparse offload memory docstrings and the Python surface in SPARSE_ATTENTION_DESIGN.md. Signed-off-by: Guiju Zhang <7135567+cascade812@users.noreply.github.com> # Conflicts: # tests/unittest/_torch/executor/test_py_executor.py --- .../SPARSE_ATTENTION_DESIGN.md | 16 +- cpp/tests/unit_tests/kernels/CMakeLists.txt | 13 +- .../run_cache_transceiver_test.py | 1 + .../backends/sparse/csa2/cache_manager.py | 16 +- .../backends/sparse/deepseek_v4/backend.py | 19 +- .../sparse/deepseek_v4/cache_manager.py | 52 +-- .../backends/sparse/deepseek_v4/kernels.py | 104 +++--- .../backends/sparse/deepseek_v4/metadata.py | 52 +-- .../attention/backends/sparse/dsa/backend.py | 1 + .../backends/sparse/dsa/cache_manager.py | 31 +- .../attention/backends/sparse/dsa/metadata.py | 6 +- .../attention/backends/sparse/offload.py | 39 ++- .../backends/sparse/offload_kernels.py | 176 ++++++++--- .../_torch/disaggregation/native/transfer.py | 43 ++- .../disaggregation/resource/kv_extractor.py | 30 +- .../engine/runners/decoder/warmup.py | 70 ++--- .../kv_cache/kv_cache_manager_v2.py | 282 ++++++++--------- .../kv_cache/sparse_offload_memory.py | 187 ++++++----- tensorrt_llm/_torch/pyexecutor/py_executor.py | 14 +- tensorrt_llm/llmapi/llm_args.py | 39 +-- .../schemas/trtllm-serve-config.schema.json | 6 +- .../trtllm-serve-disagg-config.schema.json | 6 +- .../test_deepseek_v4_cache_manager.py | 75 ++--- .../deepseek_v4/test_deepseek_v4_offload.py | 27 +- .../dsa/test_dsa_offload_cache_manager.py | 221 ++++++------- .../sparse/test_sparse_offload_kernels.py | 40 +++ .../test_phase_selective_forward.py | 3 +- .../engine/runners/decoder/test_warmup.py | 10 + .../kv_cache/sharing/test_staging_lender.py | 8 +- .../kv_cache/test_kvcm2_integration.py | 266 +++++++--------- .../kv_cache/test_sparse_offload_memory.py | 181 ++++++----- .../_torch/executor/test_py_executor.py | 26 +- ...y_executor_creator_mla_cache_reuse_sync.py | 7 + .../_torch/executor/test_resource_manager.py | 2 +- .../test_cache_transceiver_single_process.py | 1 + .../unittest/disaggregated/test_extractor.py | 19 +- .../disaggregated/test_kv_transfer.py | 1 + .../disaggregated/test_sparse_kv_transfer.py | 19 +- .../disaggregated/test_unsubmitted_session.py | 31 +- token_parity.md | 295 ------------------ 40 files changed, 1141 insertions(+), 1294 deletions(-) delete mode 100644 token_parity.md 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 index 8803aecf43d5..8c90b78f589a 100644 --- 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 @@ -988,8 +988,8 @@ never run. ## Sizing -Per sparse buffer: **`num_selected_blocks`** (the selection width `k`) and **N** (the -reservation upper bound, in the same units). +Per sparse cache (`SparseCacheConfig`): **`num_selected_blocks`** (the selection width +`k`) and **N**, `max_cached_units` (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 @@ -1010,15 +1010,17 @@ so `N` is sized by one demand rather than two: retention wants ## Python surface -Declared or planned in `tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi`: +Declared 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. +- `BufferConfig.is_sparse: bool` and `BufferConfig.sparse_unit_chunks: int` — + construction-time, per buffer. +- `AttentionLayerConfig.sparse_cache: SparseCacheConfig | None` — the layer's fetch path: + `num_selected_blocks` (the selection width `k`), `max_cached_units` (the reservation + bound), `unit_tokens` and `selection_id`; construction-time, per layer. - `KVCacheManager.is_sparse(layer_id, data_role) -> bool`. - `Batch` — membership, row order, device arrays, `resize()` and `publish()`, which transfers the whole table; dirty-row tracking is planned, not declared. -- `KVCacheManager.fetch_sparse_pages(batch, buffers, selected, out, stream)`. +- `Batch.fetch_sparse_pages(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. diff --git a/cpp/tests/unit_tests/kernels/CMakeLists.txt b/cpp/tests/unit_tests/kernels/CMakeLists.txt index 5086a317b1e2..21f67b7685d0 100644 --- a/cpp/tests/unit_tests/kernels/CMakeLists.txt +++ b/cpp/tests/unit_tests/kernels/CMakeLists.txt @@ -228,7 +228,18 @@ add_tllm_gtest( LABELS "unit;kernels;attention-mask;gpu") add_gtest(fmhaDsv4EpilogueFusionTest fmhaDsv4EpilogueFusionTest.cpp) -add_gtest(fmhaDeterministicTest fmhaDeterministicTest.cpp NO_GTEST_MAIN) +# The trtllm-gen FMHA runner has no component target, so this links the full +# stack. +add_tllm_gtest( + fmhaDeterministicTest + OWNER + kernel + SOURCES + fmhaDeterministicTest.cpp + FULL_STACK + NO_GTEST_MAIN + LABELS + "unit;kernels;fmha;full-stack;gpu") add_tllm_gtest( nvfp4ColdPageKernelsTest diff --git a/examples/disaggregated/slurm/cache_transceiver_test/run_cache_transceiver_test.py b/examples/disaggregated/slurm/cache_transceiver_test/run_cache_transceiver_test.py index e5f33a18ce00..9e653edc7cb9 100644 --- a/examples/disaggregated/slurm/cache_transceiver_test/run_cache_transceiver_test.py +++ b/examples/disaggregated/slurm/cache_transceiver_test/run_cache_transceiver_test.py @@ -120,6 +120,7 @@ class KvCacheConfigV2: # tensorrt_llm/llmapi/llm_args.py; KVCacheManagerV2._build_base_config() # reads it unconditionally. kv_pool_rebalance_config: KvPoolRebalanceConfig = field(default_factory=KvPoolRebalanceConfig) + enable_kv_pool_rebalance: bool = False class _Timeout(Exception): diff --git a/tensorrt_llm/_torch/attention/backends/sparse/csa2/cache_manager.py b/tensorrt_llm/_torch/attention/backends/sparse/csa2/cache_manager.py index 49f49e56c173..eb1081914ebc 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/csa2/cache_manager.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/csa2/cache_manager.py @@ -24,6 +24,7 @@ _estimate_full_attn_size_per_token, _estimate_swa_cache_size, _fill_kv_pages, + get_scratch_entry, ) from tensorrt_llm._utils import TensorWrapper, convert_to_torch_tensor, prefer_pinned from tensorrt_llm.bindings import DataType @@ -51,14 +52,6 @@ ) -def _scratch_entry(cache, pool_id): - """The scratch range and the slots serving it, or None when none of the blocks do.""" - scratch_range = cache.get_swa_scratch_range(pool_id) - if not scratch_range: - return None - return scratch_range, cache.get_scratch_slots(pool_id) - - def _validate_context_swa_layer_limit(limit: int | None, num_layers: int) -> None: if limit is not None and (type(limit) is not int or not 0 < limit < num_layers): raise ValueError("CSA2 context SWA layer limit must be an integer inside the model") @@ -797,11 +790,12 @@ def get_cache_indices(self, request_id: int, layer_idx: int, role: CSA2CacheRole pool_id = self.layer_to_pool_mapping_dict[layer_id] cache = self.kv_cache_map[request_id] converter = self._get_page_index_converter(layer_id, role.role) + scratch_range, scratch_slots = get_scratch_entry(cache, pool_id) or (None, ()) return converter( cache.get_base_page_indices(pool_id).tolist(), role.index_mode, - cache.get_swa_scratch_range(pool_id) or None, - cache.get_scratch_slots(pool_id), + scratch_range, + scratch_slots, ) def compute_batch_page_tables(self, request_ids, num_contexts) -> None: @@ -868,7 +862,7 @@ def _scratch_descriptors(self, request_ids): (index, row, entry) for row, cache in holders for index, pool in enumerate(self._scratch_pools.tolist()) - if (entry := _scratch_entry(cache, pool)) + if (entry := get_scratch_entry(cache, pool)) ] ranges = np.zeros((len(self._scratch_pools), len(request_ids), 2), dtype=np.int32) width = max((len(slot_ids) for _, _, (_, slot_ids) in descs), default=1) diff --git a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/backend.py b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/backend.py index 7cf4c96d1f68..e4188fda4eb1 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/backend.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/backend.py @@ -33,6 +33,7 @@ from tensorrt_llm.quantization import QuantMode from ..dsa.backend import _get_nvfp4_mla_kv_cache_amax +from ..offload import NVFP4_SCRATCH_TAIL_ROWS from .cache_manager import get_token_bytes from .compressor import Compressor from .indexer import DeepseekV4Indexer @@ -43,25 +44,20 @@ if TYPE_CHECKING: from tensorrt_llm.llmapi.llm_args import SparseAttentionConfig -# The trtllm-gen DSv4 dynamic sparse MLA FMHA can read its aux KV pool up to one KV tile -# (tileSizeKv, 128 rows) past the rows named in the sparse indices. With FP8 KV the aux pool -# is the whole compressed KV pool, so those reads stay inside it; with NVFP4 KV it is the -# per-call FP8 scratch below, which is sized to the batch. -_SPARSE_MLA_FMHA_KV_TILE_SIZE = 128 - def _alloc_nvfp4_compress_fp8_scratch( num_rows: int, head_dim: int, device: torch.device ) -> torch.Tensor: """Allocate the FP8 scratch the sparse MLA FMHA reads as its aux KV pool. - Rows [0, num_rows) receive the dequantized compressed KV. A zeroed tail of one FMHA KV tile - keeps the kernel's reads past the last indexed row inside the allocation; without it the - FMHA faults whenever the caching allocator places the scratch at the end of a mapped - segment (https://nvbugs/6796296). + Rows [0, num_rows) receive the dequantized compressed KV. With FP8 KV the aux pool is the + whole compressed KV pool, so the FMHA's reads past the last indexed row stay inside it; + this per-call scratch is sized to the batch, so it carries the same zeroed tail of one + FMHA KV tile as the offload scratch (``NVFP4_SCRATCH_TAIL_ROWS``). Without it the FMHA + faults whenever the caching allocator places the scratch at the end of a mapped segment. """ scratch = torch.empty( - (num_rows + _SPARSE_MLA_FMHA_KV_TILE_SIZE, head_dim), + (num_rows + NVFP4_SCRATCH_TAIL_ROWS, head_dim), dtype=torch.float8_e4m3fn, device=device, ) @@ -357,6 +353,7 @@ def sparse_attn_predict( compressed_lengths=compressed_lengths, compressed_batch_rows=sparse_state.batch_rows, active_request_count=sparse_state.active_request_count, + compressed_coverage=sparse_state.coverage, ) compress_pool_base_ptr = descriptor.arena_base_addresses[data_index] forward_args.sparse_runtime_params.aux_kv_cache_pool_ptr = ( diff --git a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/cache_manager.py b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/cache_manager.py index 14a3644dd661..d1676e352c78 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/cache_manager.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/cache_manager.py @@ -28,6 +28,7 @@ KVCacheManagerV2, _fill_kv_pages, _get_generation_kv_capacity, + get_scratch_entry, ) from tensorrt_llm._torch.pyexecutor.kv_cache.sparse_offload_memory import sparse_cache_config from tensorrt_llm._torch.pyexecutor.resource_manager import scan_kv_buffer_for_invalid_values @@ -307,14 +308,6 @@ def _get_attn_bytes_per_token( return token_bytes -def _scratch_entry(kv_cache, pool_id) -> tuple[HalfOpenRange, list[int]] | None: - """The scratch range and the slots serving it, or None when none of the blocks do.""" - scratch_range = kv_cache.get_swa_scratch_range(pool_id) - if not scratch_range: - return None - return scratch_range, kv_cache.get_scratch_slots(pool_id) - - def _get_index_mode(attn_type: DeepseekV4AttentionType) -> PageIndexMode: if attn_type in DEEPSEEK_V4_SLIDING_ATTENTION: return PageIndexMode.PER_LAYER @@ -379,14 +372,8 @@ def __init__( enabled=sparse_attn_config.enable_kv_cache_offload, topk=sparse_attn_config.index_topk, ) + # The base manager checks what offload can honor before it sizes anything. sparse_kv_offload_config = sparse_attn_config.get_kv_cache_offload_config() - if sparse_kv_offload_config is not None: - self._check_sparse_kv_offload_config( - sparse_kv_offload_config, - kv_cache_config, - mapping, - spec_config=kwargs.get("spec_config"), - ) # DeepSeek-V4 specific attributes initialization assert kv_cache_type == CacheTypeCpp.SELFKONLY, "DeepSeek-V4 only supports SELFKONLY" @@ -896,13 +883,9 @@ def _prepare_page_table_tensor(self, index_mapper_capacity: int) -> None: ) if self.enable_swa_scratch_reuse: - valid_scales = self._layer_attn_scales[self._layer_attn_pool_ids >= 0] - min_scale = int(valid_scales.min().item()) if valid_scales.numel() > 0 else 1 - max_scratch_pages = int(self._scratch_pages.max().item()) - self._max_scratch_slots = max( - 1, - (self.max_blocks_per_seq * max_scratch_pages + min_scale - 1) // min_scale, - ) + # Asked of the manager, as the base does: only it knows which blocks the + # window and the rewind margin keep from going stale. + self._max_scratch_slots = max(1, self.impl.max_scratch_slots(self.max_blocks_per_seq)) scratch_slots_shape = ( self.num_pools, staging_capacity, @@ -998,12 +981,8 @@ def get_cache_indices( base_indices = kv_cache.get_base_page_indices(pool_id).tolist() converter = self.impl.get_page_index_converter(layer_id, data_role) page_index_mode = _get_index_mode(attn_type) - return converter( - base_indices, - page_index_mode, - kv_cache.get_swa_scratch_range(pool_id) or None, - kv_cache.get_scratch_slots(pool_id), - ) + scratch_range, scratch_slots = get_scratch_entry(kv_cache, pool_id) or (None, ()) + return converter(base_indices, page_index_mode, scratch_range, scratch_slots) def _get_window_page_headroom(self) -> int | None: """Headroom for counting the pages a window holds with offload, else None.""" @@ -1699,7 +1678,7 @@ def compute_sliding_block_tables( if self.enable_swa_scratch_reuse and num_contexts > 0: scratch_descs_by_pool = [ [ - _scratch_entry(self.kv_cache_map[req], pool_id) + get_scratch_entry(self.kv_cache_map[req], pool_id) for req in request_ids[:num_contexts] ] for pool_id in range(self.num_pools) @@ -1875,15 +1854,12 @@ def prepare_sparse_compressor_write_indices( ) -> torch.Tensor: """Reserve this layer's pinned destinations on the compressor's stream. - Return int32[B, F] in Batch and descriptor buffer order. The caller - must join Batch publication before this call, and join the compressor - stream before fetch/attention or a history update. No KV is fetched here. - Padding roles were initialized during manager construction. The fused - compressor uses the data column as a one-row page table for data and - NVFP4 scales, whose matching geometry is checked during descriptor construction - and, once per layer outside graph capture, by ``prepare_sparse_write_indices`` - against the indices native returns. Each buffer still uses its own arena base - and byte stride. + Returns :meth:`prepare_sparse_write_indices`'s int32[B, F], in Batch and + descriptor buffer order, under the same stream contract. The fused compressor + uses the data column as a one-row page table for both the data and the NVFP4 + scales; that is sound because their arena geometry matches, which building the + sparse descriptors checks. Each buffer still uses its own arena base and byte + stride. """ return self.prepare_sparse_write_indices(state, layer_idx) diff --git a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/kernels.py b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/kernels.py index a4fc581cb8f1..88ff0224a929 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/kernels.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/kernels.py @@ -9,7 +9,12 @@ import triton import triton.language as tl -from ..offload_kernels import check_sparse_offload_tensor +from ..offload_kernels import ( + SPARSE_ARENA_COVERAGE_MESSAGE, + check_sparse_arena_merge_inputs, + check_sparse_offload_tensor, + sparse_arena_rows, +) ######################################################## # Index gather kernel @@ -75,6 +80,8 @@ def _deepseek_v4_local_to_global_kernel( SPARSE_BATCH_SIZE: tl.constexpr, sparse_num_lengths, UNIT_TOKENS: tl.constexpr, + coverage_ptr, + CHECK_COVERAGE: tl.constexpr, ): """ Triton kernel for converting local indices to global KV cache pool indices. @@ -162,17 +169,20 @@ def _deepseek_v4_local_to_global_kernel( if MERGE_SPARSE_INDICES: history = tl.load(history_lengths_ptr + row * history_stride, mask=live, other=0) - history_units = history // UNIT_TOKENS length = tl.load(compressed_lengths_ptr + req * length_stride, mask=live, other=0) - valid = live & (compressed_local_idx >= 0) & (compressed_local_idx < length) - fetched = tl.load( - fetched_indices_ptr + row * fetch_stride0 + compressed_ids * fetch_stride1, - mask=valid & (compressed_local_idx < history_units), - other=-1, - ) write = tl.load(write_indices_ptr + row * write_stride, mask=live, other=-1) - compressed_global_index = tl.where( - valid & (compressed_local_idx == history_units), write, fetched + compressed_global_index = sparse_arena_rows( + compressed_local_idx, + compressed_ids < num_compressed_indices, + live, + length, + history // UNIT_TOKENS, + fetched_indices_ptr + row * fetch_stride0, + fetch_stride1, + compressed_ids, + write, + coverage_ptr, + CHECK_COVERAGE, ) else: compressed_valid_mask = compressed_local_idx >= 0 @@ -245,6 +255,7 @@ def deepseek_v4_local_to_global_indices( compressed_history_lengths: torch.Tensor | None = None, compressed_lengths: torch.Tensor | None = None, compressed_batch_rows: torch.Tensor | None = None, + compressed_coverage: torch.Tensor | None = None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor | None]: """ Convert local token indices to global KV cache pool indices. @@ -281,15 +292,17 @@ def deepseek_v4_local_to_global_indices( graph dummy requests, refreshed in place before replay. Single-token decode only: mask both pools for padded queries without reading their page-table rows. Do not pass this for multi-query context inputs. - mask_invalid_pages: Mask negative physical page mappings in both pools - when consuming an offload read table. Compiled out by default to - preserve the original non-offload indexing behavior. + mask_invalid_pages: Mask negative physical page mappings in both pools. + A sparse offload Batch publishes its page tables with + ``BAD_PAGE_INDEX`` for the pages it does not hold. Compiled out by + default to preserve the original non-offload indexing behavior. compressed_fetched_indices: CUDA int32 [B, K] physical history indices, in Batch row order and original top-k column order. Supply all five sparse tensors and active_request_count together for single-token decode. The kernel merges history and in-flight destinations into - query order while preparing SWA indices; no compressed page table - or additional address conversion is used. + query order while preparing SWA indices, by the rule of + ``offload_kernels.sparse_arena_rows``; no compressed page table or + additional address conversion is used. compressed_write_indices: CUDA int32 [B] prepared in-flight destinations, by Batch row. May be a strided data-column view shared by paired NVFP4 pools. @@ -302,6 +315,10 @@ def deepseek_v4_local_to_global_indices( compressed_batch_rows: CUDA int32 [B] native Batch row of each request, in req_id's numbering, or -1 for no live request. A request whose row is -1 gets no indices in either pool. + compressed_coverage: Optional CUDA int32 [1] diagnostic flag for the + merge, as ``offload_kernels.merge_sparse_arena_indices`` takes it: + reset, cleared when a live selection gets no arena row, and asserted + on the device. Pass it only for debugging. Returns: A combined index tensor, or separate SWA and compressed tensors when @@ -334,13 +351,15 @@ def deepseek_v4_local_to_global_indices( swa_buffer_offset_in_tokens = (swa_buffer_ptr - swa_pool_base_ptr) // token_stride sparse_tensors = ( - (compressed_fetched_indices, "compressed_fetched_indices", 2), - (compressed_write_indices, "compressed_write_indices", 1), - (compressed_history_lengths, "compressed_history_lengths", 1), - (compressed_lengths, "compressed_lengths", 1), - (compressed_batch_rows, "compressed_batch_rows", 1), + compressed_fetched_indices, + compressed_write_indices, + compressed_history_lengths, + compressed_lengths, + compressed_batch_rows, ) - merge_sparse_indices = any(tensor is not None for tensor, _, _ in sparse_tensors) + merge_sparse_indices = any(tensor is not None for tensor in sparse_tensors) + if compressed_coverage is not None and not merge_sparse_indices: + raise ValueError("compressed_coverage checks a sparse merge; pass the sparse mappings") sparse_batch_size = sparse_num_lengths = 0 fetch_stride0 = fetch_stride1 = write_stride = history_stride = length_stride = 0 batch_rows_stride = 0 @@ -348,32 +367,27 @@ def deepseek_v4_local_to_global_indices( if ( not has_compressed or active_request_count is None - or any(tensor is None for tensor, _, _ in sparse_tensors) + or compressed_local_indices is None + or any(tensor is None for tensor in sparse_tensors) ): raise ValueError( - "Sparse indices require a compressed layer, all mappings and an active count" + "Sparse indices require a compressed layer, the original top-k selection, " + "all mappings and an active count" ) - for tensor, name, ndim in sparse_tensors: - check_sparse_offload_tensor(tensor, name, ndim, req_id.device) - sparse_batch_size, width = compressed_fetched_indices.shape - sparse_num_lengths = compressed_lengths.numel() - if ( - width != num_compressed_indices - or compressed_write_indices.shape != (sparse_batch_size,) - or compressed_history_lengths.shape != (sparse_batch_size,) - or compressed_batch_rows.shape != (sparse_batch_size,) - or not num_tokens <= sparse_num_lengths <= sparse_batch_size - ): - raise ValueError( - "Sparse mappings must cover the native Batch and configured top-k width" - ) - if compressed_local_indices is None: - raise ValueError("Sparse index preparation requires the original top-k selection") - check_sparse_offload_tensor( - compressed_local_indices, "compressed_local_indices", 2, req_id.device + check_sparse_arena_merge_inputs( + compressed_local_indices, + compressed_lengths, + compressed_history_lengths, + compressed_batch_rows, + compressed_fetched_indices, + compressed_write_indices, + active_request_count, + compressed_coverage, ) if compressed_local_indices.shape != (num_tokens, num_compressed_indices): - raise ValueError("Sparse top-k indices must be [Q, K]") + raise ValueError("Sparse top-k indices must be [Q, K] at the configured width") + sparse_batch_size = compressed_fetched_indices.shape[0] + sparse_num_lengths = compressed_lengths.numel() fetch_stride0, fetch_stride1 = compressed_fetched_indices.stride() write_stride = compressed_write_indices.stride(0) history_stride = compressed_history_lengths.stride(0) @@ -474,6 +488,8 @@ def deepseek_v4_local_to_global_indices( "(bmm1_scale, bmm2_scale, quant_scale_o, dequant_scale_q, dequant_scale_kv)" ) + if compressed_coverage is not None: + compressed_coverage.fill_(1) # Launch kernel _deepseek_v4_local_to_global_kernel[grid]( req_id_c, @@ -533,8 +549,12 @@ def deepseek_v4_local_to_global_indices( SPARSE_BATCH_SIZE=sparse_batch_size, sparse_num_lengths=sparse_num_lengths, UNIT_TOKENS=compress_ratio, + coverage_ptr=compressed_coverage if compressed_coverage is not None else out, + CHECK_COVERAGE=compressed_coverage is not None, launch_pdl=launch_with_pdl, ) + if compressed_coverage is not None: + torch._assert_async(compressed_coverage, SPARSE_ARENA_COVERAGE_MESSAGE) if split_extra: return out, out_extra diff --git a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/metadata.py b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/metadata.py index 414962a4b0de..6f06e7bba21b 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/metadata.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/metadata.py @@ -391,7 +391,8 @@ def _init_sparse_offload_state(self) -> None: layers, max_batch_size=manager.max_batch_size, topk=self.sparse_mla_topk, - nvfp4_head_dim=manager.head_dim if manager._use_nvfp4_compress else None, + # The hook the device reservation counts the scratch by. + nvfp4_head_dim=manager._sparse_offload_nvfp4_head_dim(), ) def mla_prepare_fused_q_gen_cu_seqlens(self) -> Optional[torch.Tensor]: @@ -601,7 +602,7 @@ def _build_cache_buffer_data_pointers( swa_buffer_ptrs = {layer_idx: manager.swa_pool_ptr for layer_idx in manager.pp_layers} compressed_buffer_ptrs = { layer_idx: ( - manager._ensure_sparse_batch().get_sparse_cache_base_address( + manager._sparse_arena_base( manager._layer_attn_to_layer_id[layer_idx, DeepseekV4AttentionType.COMPRESS], DeepseekV4AttentionType.COMPRESS.role, ) @@ -685,48 +686,15 @@ def restore_after_draft_forward(self, saved_state: dict | None) -> None: return None def validate_sparse_offload_batch(self) -> None: - """Reject unsupported layouts on the host before any model writes. + """Reject batches offload cannot run, on the host before any model writes. - Only a fresh, unchunked context batch or one query per decode request - is supported. This runs during prepare(), outside capture/replay; - the per-layer path consumes the resulting phase and device metadata. - Fixed configuration is validated during offload metadata initialization. + An offloading worker only decodes, one query token per request, as DSA's + does, so the check is DSA's. It runs during prepare(), outside capture and + replay; fixed configuration is validated when the offload metadata is + initialized. """ - if not self.sparse_metadata_params.enable_kv_cache_offload: - return - state = self.sparse_offload_state - if state is not None: - state.prepared = False - state.is_prefill = False - if self.num_contexts and self.num_generations: - raise NotImplementedError( - "Sparse offload does not support mixed context/decode batches" - ) - if self.seq_lens.device.type != "cpu" or self.seq_lens_kv.device.type != "cpu": - raise ValueError("Sparse offload batch validation requires host sequence lengths") - if not torch.equal(self.seq_lens, self.seq_lens_kv): - raise NotImplementedError("Sparse offload requires matching query and new-KV lengths") - if not self.num_contexts: - if not torch.all(self.seq_lens == 1): - raise NotImplementedError("Sparse offload decode requires one query per request") - return - cached = self.kv_cache_params.num_cached_tokens_per_seq - if ( - self.prompt_lens is None - or len(self.prompt_lens) < self.num_contexts - or len(cached) < self.num_contexts - or any(cached[: self.num_contexts]) - or self.seq_lens.tolist() != self.prompt_lens[: self.num_contexts] - or any( - self.kv_cache_manager.kv_cache_map[request_id].history_length != 0 - for request_id in self.request_ids - ) - ): - raise NotImplementedError( - "Sparse offload prefill must be fresh, uncached, and unchunked" - ) - if state is not None: - state.is_prefill = True + if self.sparse_offload_state is not None: + self._validate_dsa_sparse_offload_batch() def prepare(self) -> None: self._validate_ragged_verification() diff --git a/tensorrt_llm/_torch/attention/backends/sparse/dsa/backend.py b/tensorrt_llm/_torch/attention/backends/sparse/dsa/backend.py index 62f21e1ed0e4..c2cd8d8da5fe 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/dsa/backend.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/dsa/backend.py @@ -324,6 +324,7 @@ def _sparse_offload_arena_indices( state.write_indices[len(descriptor.buffer_ids)][:, data_slot], metadata.sparse_offload_attn_indices[:num_tokens], unit_tokens=descriptor.unit_tokens, + coverage=state.coverage, ) if state.nvfp4_scratch is None: forward_args.sparse_runtime_params.aux_kv_cache_pool_ptr = ( diff --git a/tensorrt_llm/_torch/attention/backends/sparse/dsa/cache_manager.py b/tensorrt_llm/_torch/attention/backends/sparse/dsa/cache_manager.py index 23f012eb8cc2..3a84c6f11508 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/dsa/cache_manager.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/dsa/cache_manager.py @@ -22,6 +22,7 @@ convert_to_torch_tensor, get_size_in_bytes, get_sm_version, + is_sm_100f, ) from tensorrt_llm.bindings import DataType from tensorrt_llm.bindings.executor import KvCacheConfig @@ -751,7 +752,16 @@ def get_layer_bytes_per_token(self, local_layer_idx: int, data_role: DataRole) - return cache_bytes def _check_dsa_offload_support(self, dtype: DataType) -> None: - """Refuse offload on the paths that have no arena read or write yet.""" + """Refuse offload on the paths that have no arena read or write yet. + + The GPU is checked first: elsewhere it decides which KV cache dtypes are + allowed at all, so a dtype error on an unsupported GPU would mislead. + """ + if not is_sm_100f(): + raise NotImplementedError( + "DSA KV cache offload reads the latent through the trtllm-gen sparse MLA " + f"kernels, which run on the SM100 family only; this GPU is SM{get_sm_version()}." + ) if self.use_fp8_ds_mla: raise NotImplementedError( "DSA KV cache offload does not support kv_cache_config.dtype='fp8_ds_mla'." @@ -760,11 +770,6 @@ def _check_dsa_offload_support(self, dtype: DataType) -> None: raise NotImplementedError( f"DSA KV cache offload supports an FP8, BF16 or NVFP4 MLA latent, got {dtype}." ) - if get_sm_version() < 100: - raise NotImplementedError( - "DSA KV cache offload reads the latent through the trtllm-gen sparse MLA " - f"kernels, which need SM100 or newer; this GPU is SM{get_sm_version()}." - ) if self._sparse_offload_topk is None or self._sparse_offload_topk <= 0: raise ValueError("DSA KV cache offload requires a positive index_topk") @@ -951,16 +956,22 @@ def _sparse_offload_nvfp4_head_dim(self) -> Optional[int]: return self.head_dim if self.dtype == DataType.NVFP4 else None def get_sparse_offload_scratch_bytes(self, config: KVCacheManagerConfigPy) -> int: - """Count the offload state and the latent writer's one-token page table. + """Count the offload state, the latent writer's one-token page table and the + indices the sparse attention reads. Per Batch row, the page table holds a K and a V offset and the write slots - one entry, since an offloading step has one query token per request. - Both metadata copies hold them. + one entry, since an offloading step has one query token per request, and the + attention reads ``topk`` merged indices. Both metadata copies hold them. """ scratch = super().get_sparse_offload_scratch_bytes(config) if scratch: writer_entries = 3 * self.max_batch_size - scratch += SPARSE_OFFLOAD_STATE_COPIES * writer_entries * torch.int32.itemsize + attn_index_entries = self.max_batch_size * self._sparse_offload_topk + scratch += ( + SPARSE_OFFLOAD_STATE_COPIES + * (writer_entries + attn_index_entries) + * torch.int32.itemsize + ) return scratch def get_cache_bytes_per_token(self) -> int: diff --git a/tensorrt_llm/_torch/attention/backends/sparse/dsa/metadata.py b/tensorrt_llm/_torch/attention/backends/sparse/dsa/metadata.py index 01c9912833f5..2dfebe4a4554 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/dsa/metadata.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/dsa/metadata.py @@ -457,12 +457,10 @@ def _validate_dsa_sparse_offload_batch(self) -> None: An offloading worker only decodes, one query token per request. """ - state = self.sparse_offload_state - state.prepared = False - state.is_prefill = False + self.sparse_offload_state.prepared = False if self.num_contexts: raise NotImplementedError( - "DSA KV cache offload runs decode batches only; context requests reach " + "Sparse KV cache offload runs decode batches only; context requests reach " "an offloading worker through disaggregated KV transfer" ) if self.seq_lens.device.type != "cpu" or self.seq_lens_kv.device.type != "cpu": diff --git a/tensorrt_llm/_torch/attention/backends/sparse/offload.py b/tensorrt_llm/_torch/attention/backends/sparse/offload.py index ed15576c24d3..afdd91770e1f 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/offload.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/offload.py @@ -18,6 +18,7 @@ from __future__ import annotations +import os from collections import defaultdict from dataclasses import dataclass, field from typing import TYPE_CHECKING, Callable, Collection, Iterable @@ -55,6 +56,12 @@ # graph metadata, which share theirs through the graph buffer pool. SPARSE_OFFLOAD_STATE_COPIES = 2 +# With TLLM_SPARSE_OFFLOAD_DEBUG_ASSERT=1, every merge of fetched arena rows asserts on +# the device that each live selection of a history or in-flight unit got a row. Read +# at import, before any graph capture. A failed assertion makes the CUDA context +# unusable, so it is opt-in. +SPARSE_OFFLOAD_DEBUG_ASSERT = os.environ.get("TLLM_SPARSE_OFFLOAD_DEBUG_ASSERT", "0") == "1" + # The sparse MLA FMHA reads the NVFP4 scratch as its aux KV pool and can read up to # one KV tile (128 rows) past the last row its indices name (https://nvbugs/6796296), # so the scratch carries a zeroed tail of that many rows. @@ -131,15 +138,19 @@ class SparseOffloadState: Batch row; per-request inputs and attention outputs by input position. ``prepared`` means preparation was enqueued successfully; stream ordering - establishes device readiness. ``is_prefill`` is set by host batch validation - before execution, never inferred from a device count inside a layer. - ``write_layer`` tracks which layer enqueued the shared write mapping during - eager execution or capture. Graph replay uses the captured GPU operations; - this host marker does not determine device readiness. + establishes device readiness. Host batch validation clears it first, and + rejects any batch that is not pure decode. ``write_layer`` tracks which + layer enqueued the shared write mapping during eager execution or capture. + Graph replay uses the captured GPU operations; this host marker does not + determine device readiness. An in-flight selection is masked out of the history fetch. Its physical index must come from the separately prepared write mapping when assembling attention inputs. + + ``coverage`` is an int32 [1] diagnostic flag the merges of fetched rows + assert on (``offload_kernels.merge_sparse_arena_indices``), allocated only + with ``SPARSE_OFFLOAD_DEBUG_ASSERT``. """ layers: dict[int, SparseOffloadLayerDescriptor] @@ -155,10 +166,10 @@ class SparseOffloadState: #: own position. Filled on first use, the earliest the native Batch exists to ask. out_slots: dict[int, tuple[int, ...]] = field(default_factory=dict) nvfp4_scratch: torch.Tensor | None = None + coverage: torch.Tensor | None = None batch: Batch | None = None history_lengths: torch.Tensor | None = None prepared: bool = False - is_prefill: bool = False write_layer: int | None = None @staticmethod @@ -179,10 +190,11 @@ def device_bytes( buffer_counts: The distinct sparse buffer counts F of the layers. nvfp4_head_dim: Head dim of the dequantized rows with NVFP4, else None. """ - # The active count, the selections, the scatter start positions and the - # Batch row map, then the fetch outputs and write destinations of each - # buffer count. + # The active count, the selections, the scatter start positions, the + # Batch row map and any coverage flag, then the fetch outputs and write + # destinations of each buffer count. index_entries = 1 + max_batch_size * topk + 2 * max_batch_size + index_entries += int(SPARSE_OFFLOAD_DEBUG_ASSERT) index_entries += sum(max_batch_size * count * (topk + 1) for count in buffer_counts) nvfp4_bytes = 0 if nvfp4_head_dim is not None: @@ -243,6 +255,7 @@ def nvfp4_scratch(head_dim: int) -> torch.Tensor: }, scatter_start_positions=index_tensor("scatter_start_positions", (max_batch_size,), 0), nvfp4_scratch=nvfp4_scratch(nvfp4_head_dim) if nvfp4_head_dim is not None else None, + coverage=index_tensor("coverage", (1,), 1) if SPARSE_OFFLOAD_DEBUG_ASSERT else None, ) @@ -994,12 +1007,7 @@ def prepare_sparse_write_indices( Data and scale buffers use the same column when their arena indices coincide; the output-slot mapping follows the arena geometry and requires no device read. """ - if ( - not self._enable_kv_cache_offload - or not state.prepared - or state.is_prefill - or state.batch is None - ): + if not self._enable_kv_cache_offload or not state.prepared or state.batch is None: raise RuntimeError("Sparse writes require prepared decode metadata") descriptor = state.layers[layer_idx] out = state.write_indices[len(descriptor.buffer_ids)] @@ -1047,7 +1055,6 @@ def fetch_sparse_indices( if ( not self._enable_kv_cache_offload or not state.prepared - or state.is_prefill or state.batch is None or state.history_lengths is None ): diff --git a/tensorrt_llm/_torch/attention/backends/sparse/offload_kernels.py b/tensorrt_llm/_torch/attention/backends/sparse/offload_kernels.py index 8e1d9a599473..d2a6864958e9 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/offload_kernels.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/offload_kernels.py @@ -19,6 +19,11 @@ import triton import triton.language as tl +# Message of the opt-in device assertion on a merge's coverage flag. +SPARSE_ARENA_COVERAGE_MESSAGE = ( + "Sparse offload: a live selection of a fetched or in-flight unit has no arena row" +) + def check_sparse_offload_tensor( tensor: torch.Tensor, name: str, ndim: int, device: torch.device @@ -30,6 +35,95 @@ def check_sparse_offload_tensor( raise ValueError(f"{name} must be on CUDA device {device}") +def check_sparse_arena_merge_inputs( + topk_indices: torch.Tensor, + unit_lengths: torch.Tensor, + history_lengths: torch.Tensor, + batch_rows: torch.Tensor, + fetched_indices: torch.Tensor, + write_indices: torch.Tensor, + active_request_count: torch.Tensor, + coverage: torch.Tensor | None, +) -> None: + """Check the inputs every merge of fetched arena rows into a selection reads. + + Selections are [Q, K] in query order; the fetch output [B, K], write rows, + history lengths and the Batch row map [B] are by Batch row; unit lengths + cover the queries and fit the Batch; the active count and the coverage + flag are [1]. + """ + device = fetched_indices.device + for tensor, name, ndim in ( + (topk_indices, "topk_indices", 2), + (unit_lengths, "unit_lengths", 1), + (history_lengths, "history_lengths", 1), + (batch_rows, "batch_rows", 1), + (fetched_indices, "fetched_indices", 2), + (write_indices, "write_indices", 1), + (active_request_count, "active_request_count", 1), + ): + check_sparse_offload_tensor(tensor, name, ndim, device) + num_rows, width = fetched_indices.shape + num_queries, topk = topk_indices.shape + if topk != width: + raise ValueError("Selections and the fetch output must share the top-k width") + if active_request_count.shape != (1,): + raise ValueError("active_request_count must be [1]") + if ( + history_lengths.shape != (num_rows,) + or write_indices.shape != (num_rows,) + or batch_rows.shape != (num_rows,) + ): + raise ValueError( + "history_lengths, write_indices and batch_rows must cover the native Batch" + ) + if not num_queries <= unit_lengths.numel() <= num_rows: + raise ValueError("Query and unit-length rows must fit the native Batch capacity") + if coverage is not None: + check_sparse_offload_tensor(coverage, "coverage", 1, device) + if coverage.shape != (1,): + raise ValueError("coverage must be [1]") + + +@triton.jit +def sparse_arena_rows( + unit, + in_width, + live, + length, + history_units, + fetched_row_ptr, + fetched_stride, + lane, + write, + coverage_ptr, + CHECK_COVERAGE: tl.constexpr, +): + """Return the arena row each lane's selected unit is read from. + + ``unit`` holds a query's selected units in its lanes, -1 for none, and + ``in_width`` the lanes that hold a selection. A unit below the request's + ``history_units`` takes its fetched row, from ``fetched_row_ptr`` at + ``lane``; the in-flight unit takes ``write``; padding queries (``live`` + false), invalid selections and units past ``length`` take -1. + + With ``CHECK_COVERAGE``, clears the int32 flag at ``coverage_ptr`` when a + live, valid selection of a history or in-flight unit gets no row. + """ + valid = live & (unit >= 0) & (unit < length) + fetched = tl.load( + fetched_row_ptr + lane * fetched_stride, + mask=valid & in_width & (unit < history_units), + other=-1, + ) + rows = tl.where(valid & (unit == history_units), write, fetched) + if CHECK_COVERAGE: + required = valid & in_width & (unit <= history_units) + covered = tl.min((~required | (rows >= 0)).to(tl.int32), axis=0) + tl.atomic_and(coverage_ptr, covered, sem="relaxed") + return rows + + @triton.jit(do_not_specialize=["num_lengths"]) def _select_sparse_history_units_kernel( topk_ptr, @@ -179,11 +273,13 @@ def _merge_sparse_arena_indices_kernel( write_stride, out_stride0, out_stride1, + coverage_ptr, NUM_ROWS: tl.constexpr, num_lengths, TOPK: tl.constexpr, UNIT_TOKENS: tl.constexpr, BLOCK_K: tl.constexpr, + CHECK_COVERAGE: tl.constexpr, ): query = tl.program_id(0) active_count = tl.load(active_request_count_ptr) @@ -199,18 +295,21 @@ def _merge_sparse_arena_indices_kernel( k = tl.arange(0, BLOCK_K) in_width = k < TOPK unit = tl.load(topk_ptr + query * topk_stride0 + k * topk_stride1, in_width, other=-1) - valid = live & (unit >= 0) & (unit < length) - fetched = tl.load( - fetched_ptr + row * fetched_stride0 + k * fetched_stride1, - valid & in_width & (unit < history_units), - other=-1, - ) write = tl.load(write_ptr + row * write_stride, live, other=-1) - tl.store( - out_ptr + query * out_stride0 + k * out_stride1, - tl.where(valid & (unit == history_units), write, fetched), - mask=in_width, + rows = sparse_arena_rows( + unit, + in_width, + live, + length, + history_units, + fetched_ptr + row * fetched_stride0, + fetched_stride1, + k, + write, + coverage_ptr, + CHECK_COVERAGE, ) + tl.store(out_ptr + query * out_stride0 + k * out_stride1, rows, mask=in_width) def merge_sparse_arena_indices( @@ -225,6 +324,7 @@ def merge_sparse_arena_indices( out: torch.Tensor, *, unit_tokens: int, + coverage: torch.Tensor | None = None, ) -> torch.Tensor: """Return each query's selection as rows of one sparse buffer's arena. @@ -233,45 +333,37 @@ def merge_sparse_arena_indices( order; ``write_indices`` [B] is that buffer's in-flight row by Batch row, written this step. Both may be strided views of the per-buffer outputs. - ``out`` is [Q, K] in query order. A selected history unit takes its fetched - row, the in-flight unit takes the write row, and padding queries, invalid - selections and units past the unit length take -1. No tensor allocation, - host read, or synchronization is performed. + ``out`` is [Q, K] in query order, assembled by :func:`sparse_arena_rows`. + No tensor allocation, host read, or synchronization is performed. With + ``coverage``, an int32 [1] diagnostic flag, the merge also resets the flag, + clears it when a live selection gets no row, and asserts on it on the + device. A failed assertion makes the CUDA context unusable, so pass it only + for debugging. """ device = out.device - for tensor, name, ndim in ( - (topk_indices, "topk_indices", 2), - (request_indices, "request_indices", 1), - (batch_rows, "batch_rows", 1), - (active_request_count, "active_request_count", 1), - (unit_lengths, "unit_lengths", 1), - (history_lengths, "history_lengths", 1), - (fetched_indices, "fetched_indices", 2), - (write_indices, "write_indices", 1), - (out, "out", 2), - ): + check_sparse_arena_merge_inputs( + topk_indices, + unit_lengths, + history_lengths, + batch_rows, + fetched_indices, + write_indices, + active_request_count, + coverage, + ) + for tensor, name, ndim in ((request_indices, "request_indices", 1), (out, "out", 2)): check_sparse_offload_tensor(tensor, name, ndim, device) - num_rows, width = fetched_indices.shape + num_rows = fetched_indices.shape[0] num_queries, topk = topk_indices.shape if unit_tokens <= 0: raise ValueError("unit_tokens must be positive") - if topk != width or out.shape != (num_queries, topk): - raise ValueError("Selections, fetch output and out must share the top-k width") - if request_indices.shape != (num_queries,) or active_request_count.shape != (1,): - raise ValueError("request_indices must be [Q] and active_request_count must be [1]") - if ( - history_lengths.shape != (num_rows,) - or write_indices.shape != (num_rows,) - or batch_rows.shape != (num_rows,) - ): - raise ValueError( - "history_lengths, write_indices and batch_rows must cover the native Batch" - ) - if not num_queries <= unit_lengths.numel() <= num_rows: - raise ValueError("Query and unit-length rows must fit the native Batch capacity") + if out.shape != (num_queries, topk) or request_indices.shape != (num_queries,): + raise ValueError("out must be [Q, K] and request_indices [Q]") if num_queries == 0 or topk == 0: return out with torch.cuda.device(device): + if coverage is not None: + coverage.fill_(1) _merge_sparse_arena_indices_kernel[(num_queries,)]( topk_indices, request_indices, @@ -290,12 +382,16 @@ def merge_sparse_arena_indices( *fetched_indices.stride(), write_indices.stride(0), *out.stride(), + coverage if coverage is not None else out, NUM_ROWS=num_rows, num_lengths=unit_lengths.numel(), TOPK=topk, UNIT_TOKENS=unit_tokens, BLOCK_K=triton.next_power_of_2(topk), + CHECK_COVERAGE=coverage is not None, ) + if coverage is not None: + torch._assert_async(coverage, SPARSE_ARENA_COVERAGE_MESSAGE) return out diff --git a/tensorrt_llm/_torch/disaggregation/native/transfer.py b/tensorrt_llm/_torch/disaggregation/native/transfer.py index eb3b50b45d79..d3bc158524c3 100644 --- a/tensorrt_llm/_torch/disaggregation/native/transfer.py +++ b/tensorrt_llm/_torch/disaggregation/native/transfer.py @@ -610,9 +610,9 @@ def __init__( self._bounce = bounce self._enforce_physical_ownership = enforce_physical_ownership self._retirement_watchdog = retirement_watchdog - # Generation-first metadata can arrive while prefill is still queued. - # Keep it for at least the configured transfer timeout. - self._request_info_ttl_s = max(self._STALE_REQ_INFO_TTL_S, request_info_timeout_s or 0.0) + # Generation-first metadata can arrive while prefill is still queued, so it + # is kept for at least the configured transfer timeout (_request_info_ttl). + self._request_info_timeout_s = request_info_timeout_s or 0.0 self._peer_requests: dict = {} self._peer_requests_timestamps: dict[int, float] = {} # unique_rid -> insert time self._peer_requests_lock = threading.Lock() @@ -708,27 +708,40 @@ def sweep_stale_req_infos(self): """ now = time.monotonic() with self._peer_requests_lock: - stale_rids = [ - rid - for rid, ts in self._peer_requests_timestamps.items() - if now - ts > self._request_info_ttl_s - ] + stale_rids = {} + for rid, ts in self._peer_requests_timestamps.items(): + ttl = self._request_info_ttl(rid) + if now - ts > ttl: + stale_rids[rid] = ttl if not stale_rids: return - orphaned_infos: list[RecvReqInfo] = [] - for rid in stale_rids: + orphaned_infos: list[tuple[RecvReqInfo, float]] = [] + for rid, ttl in stale_rids.items(): with self._sessions_lock, self._peer_requests_lock: if rid not in self._sessions and rid in self._peer_requests: req_info_map = self._peer_requests.pop(rid) self._peer_requests_timestamps.pop(rid, None) orphaned_infos.extend( - info for info in req_info_map.values() if info.context_first + (info, ttl) for info in req_info_map.values() if info.context_first ) logger.debug(f"Swept stale RecvReqInfo for rid={rid}") - for info in orphaned_infos: - self._fail_orphaned_receiver(info) + for info, ttl in orphaned_infos: + self._fail_orphaned_receiver(info, ttl) + + def _request_info_ttl(self, unique_rid: int) -> float: + """How long REQUEST_DATA for ``unique_rid`` waits for its TxSession. + + A generation-first request can arrive while its prefill is still queued, + so it waits as long as the transfer may take. A context-first one names a + session that should already exist, so it is failed after the base TTL. + Called with ``_peer_requests_lock`` held. + """ + infos = self._peer_requests.get(unique_rid) or {} + if any(not info.context_first for info in infos.values()): + return max(self._STALE_REQ_INFO_TTL_S, self._request_info_timeout_s) + return self._STALE_REQ_INFO_TTL_S - def _fail_orphaned_receiver(self, info: RecvReqInfo) -> None: + def _fail_orphaned_receiver(self, info: RecvReqInfo, waited_s: float) -> None: """Report FAILED for a context-first REQUEST_DATA that no TxSession claimed. This resolves the receiver's task, so its RxSession fails through the @@ -738,7 +751,7 @@ def _fail_orphaned_receiver(self, info: RecvReqInfo) -> None: logger.warning( f"KV transfer request rid={info.unique_rid} from receiver " f"{info.instance_name} rank {info.instance_rank} has no TxSession on this " - f"rank after {self._STALE_REQ_INFO_TTL_S:.0f}s; failing it" + f"rank after {waited_s:.0f}s; failing it" ) self._send_failed_result_to_receiver(info, defer_to_worker=True) diff --git a/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py b/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py index 5fc48b9adb56..c6fa04e525f9 100644 --- a/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py +++ b/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py @@ -46,7 +46,7 @@ from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2, Role from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManagerV2 from tensorrt_llm._utils import nvtx_range -from tensorrt_llm.runtime.kv_cache_manager_v2 import DataRole +from tensorrt_llm.runtime.kv_cache_manager_v2 import CacheTier, DataRole # Mapper kinds a V2 manager may declare via get_disagg_role_mapper_kinds(). # A physical pool may mix kinds (V2 storage coalesces buffers purely by @@ -397,30 +397,10 @@ def _window_size_for_layer(internal_layer_id: int): pool_groups: List[PhysicalPoolGroup] = [] storage_pg_to_list_idx: Dict[int, int] = {} layer_groups_by_id: List[LayerGroup | None] = [None] * len(manager.impl.layer_grouping) - # Native sparse pool descriptors expose their host hot-tier allocations. - # Dense buffers in the same layer group still have device-backed pools. - sparse_buffers = { - (int(layer.layer_id), str(buffer.role)) - for layer in config.layers - for buffer in getattr(layer, "buffers", ()) - if getattr(buffer, "is_sparse", False) - } for pg_desc in pool_group_descs: storage_pg_idx = int(pg_desc.pool_group_index) storage_pg_to_list_idx[storage_pg_idx] = len(pool_groups) - pool_memory_types = [] - for pool_idx in range(len(pg_desc.pools)): - backing = { - (int(buffer_id.layer_id), str(buffer_id.role)) in sparse_buffers - for variant in pg_desc.slot_desc.variants - for buffer_id in variant.coalesced_buffers[pool_idx].buffer_ids - } - if len(backing) > 1: - raise ValueError( - f"Physical pool {storage_pg_idx}/{pool_idx} mixes host and device buffers" - ) - pool_memory_types.append("DRAM" if True in backing else "VRAM") pool_groups.append( PhysicalPoolGroup( pools=[ @@ -428,9 +408,13 @@ def _window_size_for_layer(internal_layer_id: int): base_address=int(pool.base_address), slot_bytes=int(pool.slot_bytes), num_slots=int(pg_desc.num_slots), - memory_type=pool_memory_types[pool_idx], + # A sparse pool of the GPU tier is host memory mapped into the + # device address space; the pool, not its group, reports that. + memory_type=( + "VRAM" if pool.storage_medium == CacheTier.GPU_MEM else "DRAM" + ), ) - for pool_idx, pool in enumerate(pg_desc.pools) + for pool in pg_desc.pools ] ) ) diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/warmup.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/warmup.py index 113b7519a845..7699d8933f3a 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/warmup.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/warmup.py @@ -345,16 +345,21 @@ def _get_full_general_warmup_requests( def _generation_warmup_shapes( self, warmup_requests_configs: list[tuple[int, int]] ) -> list[tuple[int, int]]: - """Keeps the warmup shapes whose tokens all belong to generation requests. + """Keeps the warmup shapes that build no context request. - Each shape is ``(num_tokens, num_gen_requests)``. Order is preserved. + Each shape is ``(num_tokens, num_gen_requests)``. A generation request + takes ``1 + max_total_draft_tokens`` tokens and only the tokens beyond + those become context requests, so a shape with generation requests and + no more tokens than they take is a pure-generation batch. Kept shapes + are restated as their generation tokens and deduplicated, in order. """ tokens_per_gen_request = 1 + self._config.max_total_draft_tokens - return [ - (num_tokens, num_gen_requests) + shapes = [ + (num_gen_requests * tokens_per_gen_request, num_gen_requests) for num_tokens, num_gen_requests in warmup_requests_configs - if num_tokens == num_gen_requests * tokens_per_gen_request + if num_gen_requests > 0 and num_tokens <= num_gen_requests * tokens_per_gen_request ] + return list(dict.fromkeys(shapes)) def _warmup_scheduled( self, resource_manager: ResourceManager, kv_cache_manager @@ -402,10 +407,12 @@ def _warmup_scheduled( and self.guided_decoder is None and not isinstance(kv_cache_manager, MambaHybridCacheManager) ) - # A decode-only worker never runs a context batch, so every phase below - # warms up pure-generation shapes only. The flag comes from the + # Helix CP and decode-only workers never run a context batch, so every + # phase below warms up pure-generation shapes only. Both follow from the # configuration, so every rank takes the same branch. - decode_only = is_decode_only(self.sparse_attention_config) + generation_only = self.mapping.has_cp_helix() or is_decode_only( + self.sparse_attention_config + ) log_mem_snapshot("warmup/before_warmup") # Compile the DSv4 indexer-Q CuTe DSL kernels before the first @@ -422,7 +429,7 @@ def _warmup_scheduled( self._torch_compile_bypass(), ): self._run_attention_warmup( - resource_manager, can_run_general_warmup and not decode_only + resource_manager, can_run_general_warmup and not generation_only ) if can_run_general_warmup: @@ -431,7 +438,7 @@ def _warmup_scheduled( "general", metrics=self._metrics, metric_name="general_warmup_seconds" ): warmup_requests_configs = self._get_full_general_warmup_requests(resource_manager) - if decode_only: + if generation_only: warmup_requests_configs = self._generation_warmup_shapes( warmup_requests_configs ) @@ -449,40 +456,32 @@ def _warmup_scheduled( gc.collect() torch.cuda.empty_cache() - # Helix CP is decode-only, so it tunes on generation requests. - if self._context_warmups_supported and self.mapping.has_cp_helix(): - with ( - self._warmup_timer.phase( - "autotuner", metrics=self._metrics, metric_name="autotuner_warmup_seconds" - ), - self._torch_compile_bypass(), - ): - self._run_autotuner_warmup(resource_manager, generation_only=True) - log_mem_snapshot("warmup/after_autotuner") - gc.collect() - torch.cuda.empty_cache() - elif self._context_warmups_supported: + if self._context_warmups_supported: with ( self._warmup_timer.phase( "autotuner", metrics=self._metrics, metric_name="autotuner_warmup_seconds" ), self._torch_compile_bypass(), ): - self._run_autotuner_warmup(resource_manager, generation_only=decode_only) + self._run_autotuner_warmup(resource_manager, generation_only=generation_only) log_mem_snapshot("warmup/after_autotuner") # Pre-JIT Mamba SSD multi-seq + HAS_INITSTATES=True Triton kernels # for Mamba hybrid models. Runs regardless of enable_autotuner, # since MambaHybridCacheManager skips _general_warmup and the # default autotuner shape is single-seq / no-initstates. Safe - # no-op for non-Mamba models. - with ( - self._warmup_timer.phase( - "mamba_hybrid", metrics=self._metrics, metric_name="mamba_hybrid_warmup_seconds" - ), - self._torch_compile_bypass(), - ): - self._run_mamba_hybrid_warmup(resource_manager) - log_mem_snapshot("warmup/after_mamba_hybrid") + # no-op for non-Mamba models. It runs context requests, which a + # generation-only worker never serves. + if not generation_only: + with ( + self._warmup_timer.phase( + "mamba_hybrid", + metrics=self._metrics, + metric_name="mamba_hybrid_warmup_seconds", + ), + self._torch_compile_bypass(), + ): + self._run_mamba_hybrid_warmup(resource_manager) + log_mem_snapshot("warmup/after_mamba_hybrid") # Release the autotuner's exploration-mode intermediates. The # exploration leftovers are pure waste that hide tens of GiB from # non-torch allocators (cuBLAS handle workspace, UCX/NIXL, @@ -525,7 +524,7 @@ def _warmup_scheduled( self._torch_compile_bypass(), ): warmup_requests_configs = self._get_max_shape_warmup_requests(resource_manager) - if decode_only: + if generation_only: warmup_requests_configs = self._generation_warmup_shapes( warmup_requests_configs ) @@ -983,7 +982,8 @@ def _run_autotuner_warmup( """Runs forward passes to populate the autotuner cache. ``generation_only`` tunes on one batch of ``max_batch_size`` generation - requests instead of a context batch, for Helix CP generation workers. + requests instead of a context batch, for workers that never run one + (Helix CP and decode-only workers). """ from tensorrt_llm._torch.custom_ops.torch_custom_ops import ( IS_FLASHINFER_MXFP8_CUTE_DSL_AVAILABLE, 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 b15f22e167fb..b607f21ed11b 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 @@ -84,6 +84,7 @@ DataRole, DiskCacheTierConfig, GpuCacheTierConfig, + HalfOpenRange, HostCacheTierConfig, KVCacheDesc, KVCacheEventManager, @@ -139,14 +140,13 @@ from ..scheduler import ScheduledRequests from .sparse_offload_memory import ( HostHistoryReservations, - life_cycle_pages, plan_host_history_bytes, solve_sparse_offload_capacity, sparse_arena_bytes, sparse_arena_shapes, sparse_batch_shape, sparse_history_bytes_per_block, - sparse_offload_pool_ratio, + sparse_history_floor_bytes, ) if TYPE_CHECKING: @@ -233,6 +233,21 @@ def _parse_kv_fill_value(setting: str, name: str) -> Optional[float]: return None +def get_scratch_entry( + kv_cache: _KVCache, layer_group_id: int +) -> Optional[Tuple[HalfOpenRange, Sequence[int]]]: + """Return a life cycle's scratch range and the page indices of the slots serving it. + + Returns None when the cache serves none of its blocks there from scratch, which an + empty range says; the slots are then empty too. A ``PageIndexConverter`` takes the + two as its ``scratch_range`` and ``scratch_slots``. + """ + scratch_range = kv_cache.get_swa_scratch_range(layer_group_id) + if not scratch_range: + return None + return scratch_range, kv_cache.get_scratch_slots(layer_group_id) + + def _fill_kv_pages(buffer: torch.Tensor, pages: Sequence[int], value: float) -> bool: """Write ``value`` into the named pages of one layer's KV buffer. @@ -545,29 +560,22 @@ def auto_pinnable_bytes(self) -> float: def get_host_memory_limits(mapping: Mapping) -> HostMemoryLimits: """Return the host memory a rank may plan pinned allocations against. + Available memory counts the page cache that pinning reclaims -- weights read at + startup fill it, so free pages alone can understate usable memory by hundreds of + GiB -- and is capped by the process's cgroup limits. + Args: mapping: The rank's mapping, which gives the ranks sharing its node. """ import resource + # Imported here: the models package imports this module. + from tensorrt_llm._torch.models.checkpoints.hf.rank_striped_read_ahead import ( + effective_available_host_memory, + ) + local_ranks = max(1, Distributed.get(mapping).local_world_size) - # MemAvailable counts the page cache that pinning reclaims. Weights read at startup - # fill it, so free pages alone can understate usable memory by hundreds of GiB. - mem_available = None - try: - with open("/proc/meminfo") as meminfo: - for line in meminfo: - if line.startswith("MemAvailable:"): - mem_available = int(line.split()[1]) * 1024 - break - except (OSError, IndexError, ValueError): - pass - if mem_available is None: - # Kernels before 3.14 and non-Linux hosts do not publish MemAvailable. - try: - mem_available = os.sysconf("SC_PAGE_SIZE") * os.sysconf("SC_AVPHYS_PAGES") - except (ValueError, OSError): - mem_available = float("inf") + mem_available = effective_available_host_memory() try: soft_limit, _ = resource.getrlimit(resource.RLIMIT_MEMLOCK) memlock_limit = soft_limit if soft_limit != resource.RLIM_INFINITY else float("inf") @@ -653,8 +661,8 @@ def _compute_auto_host_tier_quota( def _sync_host_tier_quota(host_quota: int, mapping: Mapping) -> int: """Reduce the auto-provisioned host cache tier quota to the fleet minimum. - ``_compute_auto_host_tier_quota`` reads rank-local host state - (``SC_AVPHYS_PAGES``, ``RLIMIT_MEMLOCK``), so co-scheduled ranks can arrive + ``_compute_auto_host_tier_quota`` reads rank-local host state (available + memory and ``RLIMIT_MEMLOCK``, see :func:`get_host_memory_limits`), so co-scheduled ranks can arrive at divergent host quotas (observed up to 10x). Divergent host-tier retention makes per-rank MAX_UTILIZATION schedulers disagree about which suspended requests can resume, which wedges collectives on @@ -711,14 +719,13 @@ def _raise_on_every_rank(error: Optional[Exception], mapping: Mapping, what: str Exception: ``error``, on a rank whose check failed. RuntimeError: On the other ranks, when any rank's check failed. """ - failed_rank = mapping.rank if error is not None else mapping.world_size - if mapping.world_size > 1: - # The lowest failing rank, or world_size when every rank passed. - failed_rank = Distributed.get(mapping).allreduce(failed_rank, op=ReduceOp.MIN) + status = _sync_kv_cache_manager_init_status( + _KVCacheManagerInitStatus.OK if error is None else _KVCacheManagerInitStatus.ABORT, mapping + ) if error is not None: raise error - if failed_rank < mapping.world_size: - raise RuntimeError(f"{what} failed on rank {failed_rank}; see that rank's log.") + if status != _KVCacheManagerInitStatus.OK: + raise RuntimeError(f"{what} failed on another rank; see the other ranks' logs.") def _estimate_swa_cache_size( @@ -1928,8 +1935,9 @@ class KVCacheManagerV2(BaseResourceManager): # copy; 0 without offload. _sparse_batch_buffer_bytes: int = 0 _sparse_batch_staging_bytes: int = 0 - # The device quota's split when the host budget holds the host-backed life - # cycles to less than their share; None to split by decode bytes. + # The manager config offload is sized from, built once with a placeholder + # device quota; None until offload sizing first asks for it. + _offload_sizing_config: Optional[KVCacheManagerConfigPy] = None def __init__( self, @@ -2379,9 +2387,7 @@ def append_to_kv_heads_per_layer( kv_cache_config, tokens_per_block=tokens_per_block, max_leases=self.max_batch_size ) self._sparse_batch_buffer_bytes, self._sparse_batch_staging_bytes = ( - self._get_sparse_batch_buffer_bytes( - kv_cache_config, tokens_per_block=tokens_per_block, max_leases=self.max_batch_size - ) + self._get_sparse_batch_buffer_bytes(kv_cache_config, tokens_per_block=tokens_per_block) ) reserved_bytes = reserved_arena_bytes + self._sparse_batch_buffer_bytes offload_error: Optional[ValueError] = None @@ -3750,11 +3756,10 @@ def _copy_scratch_metadata_to_device( for pool_id in range(self.num_pools): for context_idx, request_id in enumerate(request_ids[:num_contexts]): - cache = self.kv_cache_map[request_id] - scratch_range = cache.get_swa_scratch_range(pool_id) - if not scratch_range: + entry = get_scratch_entry(self.kv_cache_map[request_id], pool_id) + if entry is None: continue - slot_ids = cache.get_scratch_slots(pool_id) + scratch_range, slot_ids = entry if len(slot_ids) > self._max_scratch_slots: raise RuntimeError( f"Scratch slot count {len(slot_ids)} exceeds staging capacity " @@ -3885,14 +3890,18 @@ def _build_base_config( f"equal to max_seq_len ({self.max_seq_len})" ) + if self._sparse_kv_offload is not None: + # A worker that offloads sparse history only decodes: every resident + # request is a generation request, and warmup runs no context batch. + # Size its step for the length offload is sized for, which is + # max_seq_len when avg_seq_len is not set. + typical_seq_len = self._get_sparse_offload_seq_len(kv_cache_config) + if typical_seq_len is not None: - # Model one context request and enough generation requests to fill - # the resident request capacity without over-provisioning windowed - # cache pools. + # Model enough generation requests to fill the resident request + # capacity without over-provisioning windowed cache pools, beside one + # context request unless the worker only decodes. generation_request_capacity = self._get_generation_request_capacity() - context_capacity = ( - self.max_num_tokens if self.max_num_tokens is not None else typical_seq_len - ) + self.num_extra_kv_tokens generation_history_length = max(0, typical_seq_len - self.max_draft_len - 1) generation = KVCacheDesc( capacity=typical_seq_len, @@ -3900,28 +3909,28 @@ def _build_base_config( beam_width=beam_width, ) if self._sparse_kv_offload is not None: - # A worker that offloads sparse history only decodes, so every - # resident request is a generation request. typical_step = BatchDesc([generation] * generation_request_capacity) else: + context_capacity = ( + self.max_num_tokens if self.max_num_tokens is not None else typical_seq_len + ) + self.num_extra_kv_tokens typical_step = BatchDesc( [KVCacheDesc(capacity=context_capacity, history_length=0)] + [generation] * (generation_request_capacity - 1) ) - - # General and chunked-prefill warmup uses one fresh context request - # at the per-iteration token budget. - if self.max_num_tokens is not None: - constraints.append( - BatchDesc( - [ - KVCacheDesc( - capacity=self.max_num_tokens + self.num_extra_kv_tokens, - history_length=0, - ) - ] + # General and chunked-prefill warmup uses one fresh context request + # at the per-iteration token budget. + if self.max_num_tokens is not None: + constraints.append( + BatchDesc( + [ + KVCacheDesc( + capacity=self.max_num_tokens + self.num_extra_kv_tokens, + history_length=0, + ) + ] + ) ) - ) if self.layer_specs is not None: scale_roles = (Role.KEY_BLOCK_SCALE, Role.VALUE_BLOCK_SCALE) @@ -4113,41 +4122,42 @@ def _build_manager_config( if self._sparse_kv_offload is not None and self._host_history_budget > 0: # The manager divides the device quota and this budget together, holding the # host-backed pools inside it rather than being handed a quota already lowered. + # It splits both by the decode-only typical step, host-backed pools included. config = replace(config, host_sparse_pool_quota=self._host_history_budget) - if self._sparse_kv_offload is not None and config.initial_pool_ratio is None: - # Split the quota the way decoding requests use it, so the host-backed - # pools it sizes hold as many requests as the device-backed ones. - config = replace( - config, - initial_pool_ratio=sparse_offload_pool_ratio( - life_cycle_pages(config.layers, config.tokens_per_block), - capacity=self._get_sparse_offload_seq_len(kv_cache_config), - capacity_headroom=self._generation_kv_capacity_headroom, - tokens_per_block=config.tokens_per_block, - ), - ) return config + def _get_offload_sizing_config( + self, kv_cache_config: KvCacheConfig, *, tokens_per_block: int + ) -> KVCacheManagerConfigPy: + """Return the manager config sparse KV cache offload is sized from, built once. + + Offload is sized before the device quota is known. Layer configs and + constraints do not depend on the tier quotas, so a config built with a + placeholder quota has the same arenas, Batch and host floor as the real one. + """ + if self._offload_sizing_config is None: + self._offload_sizing_config = self._build_manager_config( + kv_cache_config, + tokens_per_block=tokens_per_block, + cache_tiers=[GpuCacheTierConfig(quota=sys.maxsize)], + ) + return self._offload_sizing_config + def _get_sparse_arena_bytes( self, kv_cache_config: KvCacheConfig, *, tokens_per_block: int, max_leases: int ) -> int: """Return the device bytes of the sparse cache arenas the manager will allocate. The device quota must leave room for the arenas before the cache config - is built. Layer configs do not depend on the tier quotas, so a config - built with a placeholder quota has the same arenas as the real one. + is built. """ if self._sparse_kv_offload is None: return 0 - config = self._build_manager_config( - kv_cache_config, - tokens_per_block=tokens_per_block, - cache_tiers=[GpuCacheTierConfig(quota=sys.maxsize)], - ) + config = self._get_offload_sizing_config(kv_cache_config, tokens_per_block=tokens_per_block) return sparse_arena_bytes(config, max_leases) def _get_sparse_batch_buffer_bytes( - self, kv_cache_config: KvCacheConfig, *, tokens_per_block: int, max_leases: int + self, kv_cache_config: KvCacheConfig, *, tokens_per_block: int ) -> tuple[int, int]: """Return the bytes offload allocates once for ``max_batch_size`` rows. @@ -4161,11 +4171,7 @@ def _get_sparse_batch_buffer_bytes( """ if self._sparse_kv_offload is None: return 0, 0 - config = self._build_manager_config( - kv_cache_config, - tokens_per_block=tokens_per_block, - cache_tiers=[GpuCacheTierConfig(quota=sys.maxsize)], - ) + config = self._get_offload_sizing_config(kv_cache_config, tokens_per_block=tokens_per_block) batch = sparse_batch_shape( config.layers, tokens_per_block=config.tokens_per_block, @@ -4209,13 +4215,11 @@ def _plan_host_history_budget( The manager is given this budget and holds its host-backed pools inside it while dividing the device quota, so nothing is lowered here: a bound applied before the - manager and again inside it would size the level twice. + manager and again inside it would size the level twice. The manager also keeps + the host-backed slots its constraint batches need, such as one ``max_seq_len`` + request, and refuses a budget below them, so the budget never plans less. """ - config = self._build_manager_config( - kv_cache_config, - tokens_per_block=tokens_per_block, - cache_tiers=[GpuCacheTierConfig(quota=sys.maxsize)], - ) + config = self._get_offload_sizing_config(kv_cache_config, tokens_per_block=tokens_per_block) num_tokens = self._get_sparse_offload_seq_len(kv_cache_config) # An explicit host cache tier is pinned too, beside the offloaded history. host_tier_bytes = max(0, kv_cache_config.host_cache_size or 0) @@ -4227,6 +4231,12 @@ def _plan_host_history_budget( pinnable_host_bytes=get_host_memory_limits(self.mapping).auto_pinnable_bytes - host_tier_bytes, staging_bytes=self._sparse_batch_staging_bytes, + floor_bytes=sparse_history_floor_bytes( + config.layers, + config.constraints, + tokens_per_block=config.tokens_per_block, + max_util_for_resume=config.max_util_for_resume, + ), ) logger.info( f"Sparse KV cache offload may pin {self._host_history_budget / (1 << 30):.2f}GiB of " @@ -7010,10 +7020,10 @@ def _check_sparse_kv_offload_config( and all of them raise when any rank cannot honor the settings. Raises: - ValueError: KV pool rebalancing or speculative decoding is enabled, or - the host history size, with an explicit host cache tier, exceeds - this rank's share of available host memory or the pinned-memory - limit. + ValueError: KV pool rebalancing, speculative decoding, block reuse or + pipeline parallelism is enabled, or the host history size, with an + explicit host cache tier, exceeds this rank's share of available host + memory or the pinned-memory limit. RuntimeError: Another rank cannot honor the settings. """ error: Optional[ValueError] = None @@ -7037,6 +7047,14 @@ def _check_sparse_kv_offload_config( "in host-backed pools that reuse cannot share. Set " "kv_cache_config.enable_block_reuse=False." ) + elif mapping.pp_size > 1: + error = ValueError( + "Sparse KV cache offload does not support pipeline parallelism: its attention " + "Batch seats one micro-batch's requests at a time, so each micro-batch's step " + "would release the sparse caches of the others in flight. Set " + "pipeline_parallel_size=1, or disable " + "sparse_attention_config.enable_kv_cache_offload." + ) elif host_history_size is not None: limits = get_host_memory_limits(mapping) # An explicit host cache tier is pinned beside the offloaded history; @@ -7064,10 +7082,11 @@ def _check_sparse_kv_offload_config( def _setup_sparse_kv_offload( self, kv_cache_config: KvCacheConfig, max_seq_len: int, num_reserved_index_slots: int ) -> None: - """Check the host-backed pools, start host admission and report the capacity. + """Start host admission and report the capacity. A rank whose layers have no sparse buffers holds no host history. Every - rank calls this at the same point. + rank calls this at the same point. The manager keeps its host-backed pools + within the host budget, or refuses to start, so they need no check here. Args: kv_cache_config: The KV cache configuration. @@ -7075,25 +7094,8 @@ def _setup_sparse_kv_offload( manager clamps ``max_seq_len`` to what its pools hold. num_reserved_index_slots: Index slots reserved for persistent dummy requests, such as CUDA-graph padding. - - Raises: - ValueError: This rank's host-backed pools exceed their host budget. - RuntimeError: Another rank's do. """ host_pool_groups = self.get_host_backed_pool_groups() - host_pool_bytes = self._get_host_backed_pool_bytes() if host_pool_groups else 0 - sizing_error: Optional[ValueError] = None - if host_pool_bytes > self._host_history_budget: - # The quota leaves room for how the native manager rounds its split, - # so only its slot floors can take the pools past the budget. - sizing_error = ValueError( - f"The host-backed sparse history pools take {host_pool_bytes / (1 << 30):.2f}GiB, " - f"more than their {self._host_history_budget / (1 << 30):.2f}GiB host budget: the " - "KV cache manager keeps a minimum number of slots in every pool group, and here " - "that minimum needs more host memory than the budget holds. Raise " - "sparse_attention_config.kv_cache_offload_config.host_history_size." - ) - _raise_on_every_rank(sizing_error, self.mapping, "Sizing the host-backed history pools") local_capacity_blocks = sys.maxsize if host_pool_groups: # Dummy requests hold history pages without a reservation, so @@ -7104,8 +7106,11 @@ def _setup_sparse_kv_offload( - self._get_dummy_history_blocks(num_reserved_index_slots), ) # The executor rejects a request on arrival when its history can never - # fit (get_request_kv_block_budget). Every rank must reject the same - # requests, so all of them judge against the smallest pool. + # fit (get_request_kv_block_budget), and admission reserves its history + # pages. Each rank sizes its pools from its own host memory, but every + # rank must reject and admit the same requests, or the ones that admit + # alone wait in collectives the others never enter. So all of them judge + # against the smallest pool. capacity_blocks = local_capacity_blocks if self.mapping.world_size > 1: capacity_blocks = Distributed.get(self.mapping).allreduce( @@ -7117,7 +7122,7 @@ def _setup_sparse_kv_offload( self._warn_if_host_history_rejects_requests(capacity_blocks, max_seq_len) if not host_pool_groups: return - self._host_history_reservations = HostHistoryReservations(local_capacity_blocks) + self._host_history_reservations = HostHistoryReservations(capacity_blocks) if not self.is_estimating_kv_cache: self._log_sparse_kv_offload_capacity(kv_cache_config) @@ -7207,54 +7212,28 @@ def get_request_device_bytes(self, num_tokens: int) -> int: return int(num_tokens * size_per_token + size_per_request) def _get_host_backed_pool_bytes(self) -> int: - """Return the bytes of the hot pools whose pages are pinned host memory.""" - sparse_buffers = self._get_sparse_buffer_ids() - total = 0 - for pool_group in self.get_host_backed_pool_groups(): - # A pool holds the coalesced buffer at its index in every variant, and - # variants of one group agree on sparsity position by position. - coalesced_buffers = pool_group.slot_desc.variants[0].coalesced_buffers - for pool in pool_group.pools: - buffer_ids = coalesced_buffers[int(pool.pool_index)].buffer_ids - if any( - (int(buffer_id.layer_id), str(buffer_id.role)) in sparse_buffers - for buffer_id in buffer_ids - ): - total += pool.slot_bytes * pool_group.num_slots - return total + """Return the bytes of the hot pools whose pages are pinned host memory. - def _get_sparse_buffer_ids(self) -> set[tuple[int, str]]: - """Return the ``(layer_id, role)`` of every sparse buffer the manager allocates.""" - return { - (int(layer.layer_id), str(buffer.role)) - for layer in self.kv_cache_manager_py_config.layers - if isinstance(layer, AttentionLayerConfig) - for buffer in layer.buffers - if buffer.is_sparse - } + The native count includes how each pool rounds its size to the host grain, + so it is what the host budget pays for. + """ + medium = self.impl.sparse_backing_medium() + return 0 if medium is None else int(self.impl.get_num_bytes(GPU_LEVEL, medium)) def get_host_backed_pool_groups(self) -> list: """Return the hot pool groups that hold host-backed pools. A sparse buffer's pool is host memory mapped into the device address space. It shares its pool group, and so the group's slot count, with the - device-backed pools of the same life cycle. + device-backed pools of the same life cycle; each pool reports its medium. Returns: The ``PoolGroupDesc`` of each hot pool group with a host-backed pool. """ - sparse_buffers = self._get_sparse_buffer_ids() - if not sparse_buffers: - return [] return [ pool_group for pool_group in self.impl.pool_group_descs - if any( - (int(buffer_id.layer_id), str(buffer_id.role)) in sparse_buffers - for variant in pool_group.slot_desc.variants - for coalesced in variant.coalesced_buffers - for buffer_id in coalesced.buffer_ids - ) + if any(pool.storage_medium != CacheTier.GPU_MEM for pool in pool_group.pools) ] def flush_iteration_events(self): @@ -8024,13 +8003,10 @@ def _iter_cache_buffers_for_invalid_check(self) -> Iterable[Tuple[int, torch.Ten # Host-backed sparse history is left out: nothing reads a row of it # before a transfer or a write-back fills it, and walking it would cross # PCIe for every page of the host budget. - sparse_buffers = ( - self._get_sparse_buffer_ids() if self._sparse_kv_offload is not None else set() - ) + skip_sparse = self._sparse_kv_offload is not None pool_handled = set() for layer_id, layer_offset in self.layer_offsets.items(): - pool_id = self.layer_to_pool_mapping_dict[layer_offset] - if pool_id in pool_handled or (int(layer_offset), str(Role.KEY)) in sparse_buffers: + if skip_sparse and self.impl.is_sparse(LayerId(layer_offset), Role.KEY): continue buffer = self.get_buffers(layer_id) if buffer is None: diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/sparse_offload_memory.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/sparse_offload_memory.py index d6d910a1bc68..7cf0d3efb7ff 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/sparse_offload_memory.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/sparse_offload_memory.py @@ -21,10 +21,11 @@ The native manager puts a sparse buffer's pages in a host-backed pool of the same pool group as its life cycle's dense buffers. A slot spans every pool of -its group, so the host-backed pools get the slot count the device quota buys -for the dense ones, and nothing else bounds them. The host memory offload takes -is therefore set by the device quota and by how that quota is split across life -cycles. The arena is a separate device allocation outside every tier's quota. +its group, so the host-backed pools take the slot count of their group. The +manager is handed a host budget (``host_sparse_pool_quota``) beside the device +quota, divides both by one pool ratio, and lowers a group's slot count where +its host-backed pools would exceed that budget. The arena is a separate device +allocation outside every tier's quota. Everything here is derived from the layer configs a cache manager declares, so it is the same for every model. A model marks each offloaded buffer @@ -49,6 +50,7 @@ if TYPE_CHECKING: from tensorrt_llm.runtime.kv_cache_manager_v2 import ( + BatchDesc, BufferConfig, KVCacheManagerConfig, LayerConfig, @@ -72,6 +74,9 @@ _BATCH_WATERMARKS_PER_ROW = 2 # Pinned host memory is mapped in whole pages of this many bytes. _HOST_PAGE_BYTES = 4096 +# A host-backed pool rounds its bytes up to this grain (hostMemSizeGranularity in +# utils/hostMemBacking.cpp: 2 MiB, or the commit alignment where that is larger). +_HOST_POOL_GRANULARITY = 2 << 20 # A decode step writes one query token per request, which completes at most one # unit, so one unit per layer is in flight (_sparseCacheMaxInFlightUnits). _SPARSE_MAX_FLUSH_UNITS = 1 @@ -293,26 +298,6 @@ class LifeCyclePage: host_bytes: int device_pool_bytes: tuple[int, ...] - def pages_per_request( - self, *, capacity: int, capacity_headroom: int, tokens_per_block: int - ) -> int: - """Return the most pages one decoding request holds. - - Args: - capacity: Tokens the request's cache holds, its history included. - capacity_headroom: Tokens a decode step reserves above the history. - tokens_per_block: Tokens per page. - """ - if self.is_state: - return 1 - pages = ceil_div(capacity, tokens_per_block) - if self.window_size is None: - return pages - # The live interval spans window + headroom - 1 tokens, at any page offset. - window_pages = ceil_div(self.window_size + capacity_headroom - 2, tokens_per_block) + 1 - sink_pages = ceil_div(self.num_sink_tokens, tokens_per_block) - return min(pages, window_pages + sink_pages) - def life_cycle_pages(layers: Sequence[LayerConfig], tokens_per_block: int) -> list[LifeCyclePage]: """Return the life cycles ``layers`` form, in the order the native manager numbers them. @@ -379,9 +364,10 @@ class SparseBatchShape: """The native Batch the sparse fetch reads, sized once for every step. It holds a page table per life cycle, and two watermarks and one arena slot per row, - in a device mirror, a pinned host staging copy and a device staging copy, and the - fetch's work arrays on the device. Every step uploads the host staging copy to the - device one, then copies the mirror's part of it across. + in a device mirror, a pinned host staging copy and a device staging copy, the + converted index tables the fetch and the write-back read, and the fetch's work arrays + on the device. Every step uploads the host staging copy to the device one, then + copies the mirror's part of it across. Attributes: num_life_cycles: Life cycles of the manager; each has a page table. @@ -394,6 +380,9 @@ class SparseBatchShape: the allocation, which is why it is required rather than defaulted. flush_capacities: Units one step can write back per layer, for each of those life cycles, from the step's token budget. + index_table_entries: Entries of the converted index tables, summed over the + tables; each entry is a page table's worth of indices. Required for the same + reason as ``layer_counts``. """ num_life_cycles: int @@ -402,6 +391,7 @@ class SparseBatchShape: selection_widths: tuple[int, ...] layer_counts: tuple[int, ...] flush_capacities: tuple[int, ...] + index_table_entries: int @property def block_bytes(self) -> int: @@ -425,11 +415,23 @@ def staging_bytes(self) -> int: """Pinned host bytes of the staging copy of the block.""" return pad_up(self.block_bytes, _HOST_PAGE_BYTES) + @property + def index_table_bytes(self) -> int: + """Device bytes of the converted index tables, each sized for every row. + + Their descriptors, a few bytes per table and buffer, are not counted. + """ + return ( + _INDEX_BYTES * self.index_table_entries * self.max_batch_size * self.max_blocks_per_seq + ) + @property def device_bytes(self) -> int: - """Device bytes of the Batch: the staged block, the mirror of its tail and the - fetch's work arrays.""" - return self.block_bytes + self.mirror_bytes + self.fetch_scratch_bytes + """Device bytes of the Batch: the staged block, the mirror of its tail, the + converted index tables and the fetch's work arrays.""" + return ( + self.block_bytes + self.mirror_bytes + self.index_table_bytes + self.fetch_scratch_bytes + ) @property def fetch_scratch_bytes(self) -> int: @@ -481,6 +483,17 @@ def flush_capacity(grouped_layers: Sequence[LayerConfig]) -> int: unit_tokens = grouped_layers[0].sparse_cache.unit_tokens or tokens_per_block return max_total_input_tokens // unit_tokens + max_batch_size + def index_table_entries(grouped_layers: Sequence[LayerConfig]) -> int: + # The Batch converts every sparse buffer's indices (sparseIndexTableSpecs, one spec + # per layer and pool), and merges the specs of a life cycle's layers into shared + # tables where nothing can substitute a scratch slot (_buildIndexTables). Only a + # windowed life cycle can hold scratch (supportsIndexMode); counting its layers' + # tables separately is exact with SWA scratch reuse on and an upper bound without. + entries = len(_fetched_sparse_buffers(grouped_layers[0])) + if grouped_layers[0].sliding_window_size is not None: + entries *= len(grouped_layers) + return entries + return SparseBatchShape( num_life_cycles=len(life_cycle_pages(layers, tokens_per_block)), max_batch_size=max_batch_size, @@ -490,12 +503,15 @@ def flush_capacity(grouped_layers: Sequence[LayerConfig]) -> int: ), layer_counts=tuple(len(grouped_layers) for grouped_layers in fetch_groups), flush_capacities=tuple(flush_capacity(grouped_layers) for grouped_layers in fetch_groups), + index_table_entries=sum( + index_table_entries(grouped_layers) for grouped_layers in fetch_groups + ), ) # --------------------------------------------------------------------------- -# Sizing policy: the GPU cache per request, the quota split, the host budget and -# the device quota it bounds. +# Sizing policy: the GPU cache per request and the host budget, which the native +# manager divides beside the device quota. # --------------------------------------------------------------------------- @@ -569,47 +585,47 @@ def sparse_cache_config( ) -def sparse_offload_pool_ratio( - life_cycles: Sequence[LifeCyclePage], +def sparse_history_floor_bytes( + layers: Sequence[LayerConfig], + constraints: Sequence[BatchDesc], *, - capacity: int, - capacity_headroom: int, tokens_per_block: int, -) -> list[float]: - """Return the device quota's split across life cycles for decoding requests. + max_util_for_resume: float, +) -> int: + """Return the host bytes the host-backed pools need for the batches in ``constraints``. - Each life cycle's share is the device bytes its pages take in a request of - ``capacity`` tokens, so the quota buys the same number of such requests in - every life cycle. Host-backed pools cost the quota nothing and get no share. + The native manager reserves the slots every constraint batch needs before it + divides the host budget, and refuses a budget below them + (``computePoolGroupMinSlotsFromConstraints``, ``ratioToSlotCountList``). A batch + needs a slot per page its requests' capacities cover in a life cycle that keeps + all of its history, every beam counted, and no more in a windowed one; the + manager scales that by ``1 / max_util_for_resume``. Every host-backed pool then + rounds its bytes up to the host grain, which this counts in full. Args: - life_cycles: The manager's life cycles, in native order. - capacity: Tokens one request's cache holds. - capacity_headroom: Tokens a decode step reserves above the history. + layers: Every layer config the manager allocates. + constraints: The manager config's constraint batches. tokens_per_block: Tokens per page. - - Returns: - One positive share per life cycle, summing to 1; the order and meaning of - ``KVCacheManagerConfig.initial_pool_ratio``. + max_util_for_resume: The utilization the manager keeps every pool group under. """ - # Device and host bytes together. The manager's pool ratio counts every pool, host-backed - # ones included -- it is projected onto cold levels, where one pool holds dense and sparse - # buffers together, and a rebalance compares it against a layout measured the same way. - # Weighing device bytes alone here would hand the manager a ratio in one currency and - # measure the result in another. - weights = [ - (life_cycle.device_bytes + life_cycle.host_bytes) - * life_cycle.pages_per_request( - capacity=capacity, - capacity_headroom=capacity_headroom, - tokens_per_block=tokens_per_block, + # The native manager gives each size of sparse buffer in a life cycle a pool. + host_pools = { + (layer.sliding_window_size, ceil_div(layer.num_sink_tokens or 0, tokens_per_block), size) + for layer in layers + if isinstance(layer, AttentionLayerConfig) + for size in (buffer.size for buffer in layer.buffers if buffer.is_sparse) + } + if not host_pools or not constraints: + return 0 + pages = max( + sum( + ceil_div(request.capacity, tokens_per_block) * max(1, request.beam_width) + for request in batch.kv_caches ) - for life_cycle in life_cycles - ] - if any(weight <= 0 for weight in weights): - raise ValueError(f"Every life cycle needs bytes to share the quota, got {weights}") - total = sum(weights) - return [weight / total for weight in weights] + for batch in constraints + ) + slots = math.ceil(pages / max_util_for_resume) + return slots * sparse_history_bytes_per_block(layers) + len(host_pools) * _HOST_POOL_GRANULARITY def plan_host_history_bytes( @@ -619,13 +635,15 @@ def plan_host_history_bytes( host_bytes_per_request: int, pinnable_host_bytes: float, staging_bytes: int, + floor_bytes: int = 0, ) -> int: """Return the host bytes the host-backed history pools may take. Offload also pins the sparse Batch's staging buffers, so they come out of the same host memory. An explicit ``host_history_size`` covers both. Otherwise the pools hold ``max_batch_size`` requests, capped so that pools - and staging fit in the host memory the rank may pin. + and staging fit in the host memory the rank may pin. Either way they must + reach ``floor_bytes``: the native manager refuses a smaller budget. Args: host_history_size: The configured size, or None to size automatically. @@ -634,10 +652,17 @@ def plan_host_history_bytes( pinnable_host_bytes: Host bytes automatic sizing may pin for offload; ``inf`` when unknown. staging_bytes: Pinned host bytes of the sparse Batch's staging buffers. + floor_bytes: Host bytes the manager's constraint batches need in the + host-backed pools (:func:`sparse_history_floor_bytes`). Raises: - ValueError: No host memory is left for the pools. + ValueError: No host memory is left for the pools, or less than + ``floor_bytes``. """ + floor = ( + f"the {floor_bytes / _MIB:.1f}MiB the KV cache manager reserves for the batches it must " + "always hold, such as one max_seq_len request" + ) if host_history_size is not None: budget = host_history_size - staging_bytes if budget <= 0: @@ -647,20 +672,38 @@ def plan_host_history_bytes( f"history: the sparse Batch's staging buffers take {staging_bytes / _MIB:.1f}MiB " "of it." ) + if budget < floor_bytes: + raise ValueError( + "sparse_attention_config.kv_cache_offload_config.host_history_size " + f"({host_history_size / _MIB:.1f}MiB) leaves {budget / _MIB:.1f}MiB for offloaded " + f"history beside the sparse Batch's staging buffers, less than {floor}. Raise " + "host_history_size, or lower max_seq_len." + ) return budget if not host_bytes_per_request: # This rank's layers offload no history. return 0 - budget = min(max_batch_size * host_bytes_per_request, pinnable_host_bytes - staging_bytes) - if budget <= 0: + available = pinnable_host_bytes - staging_bytes + pinnable = ( + f"this rank may pin {max(0.0, pinnable_host_bytes) / _MIB:.1f}MiB, and the sparse " + f"Batch's staging buffers take {staging_bytes / _MIB:.1f}MiB" + ) + if available <= 0: raise ValueError( "Sizing sparse KV cache offload automatically leaves no host memory for offloaded " - f"history: this rank may pin {max(0.0, pinnable_host_bytes) / _MIB:.1f}MiB, and the " - f"sparse Batch's staging buffers take {staging_bytes / _MIB:.1f}MiB. Set " + f"history: {pinnable}. Set " "sparse_attention_config.kv_cache_offload_config.host_history_size, or lower " "kv_cache_config.host_cache_size." ) - return int(budget) + if available < floor_bytes: + raise ValueError( + "Sizing sparse KV cache offload automatically leaves " + f"{available / _MIB:.1f}MiB of host memory for offloaded history, less than " + f"{floor}: {pinnable}. Set " + "sparse_attention_config.kv_cache_offload_config.host_history_size, lower " + "max_seq_len, or lower kv_cache_config.host_cache_size." + ) + return int(max(min(max_batch_size * host_bytes_per_request, available), floor_bytes)) # --------------------------------------------------------------------------- diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 77be7db6a6a7..da685eddaa4d 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -461,6 +461,8 @@ class PyExecutor: # The model's executor events (RuntimeHooks.executor_events); read once at # construction. _executor_events: Optional[ExecutorEvents] = None + # Whether the KV cache manager is a KVCacheManagerV2; set at construction. + _is_kv_manager_v2 = False # Minimum number of async micro batches for async PP execution. # This is a trade-off between memory usage and performance. @@ -3402,8 +3404,8 @@ def _executor_loop_pp(self): # made in model_engine.py::_forward_step. This is only important # for disaggregated serving. For non-disaggregated serving, # the generation requests always have batch_idx. - # Sort before preparing resources: a resource manager can - # assign per-request rows in this order. + # Sort before preparing resources, so the resource managers see + # the requests in the order the forward runs them. scheduled_batch.generation_requests = sorted( # stable sort scheduled_batch.generation_requests, key=lambda req: int(req.py_batch_idx is not None), @@ -6341,8 +6343,8 @@ def _executor_loop_overlap(self) -> None: # made in model_engine.py::_forward_step. This is only important # for disaggregated serving. For non-disaggregated serving, # the generation requests always have batch_idx. - # Sort before preparing resources: a resource manager can - # assign per-request rows in this order. + # Sort before preparing resources, so the resource managers see + # the requests in the order the forward runs them. scheduled_batch.generation_requests = sorted( # stable sort scheduled_batch.generation_requests, key=lambda req: int(req.py_batch_idx is not None), @@ -6708,10 +6710,12 @@ def _validate_request_budget(self, request: LlmRequest) -> None: for resource_type, required_blocks, primary_capacity in ( self.resource_manager.get_request_kv_block_budgets(request)): if required_blocks > primary_capacity: + # The capacity is of the pool the budget checks: the GPU primary + # pool, or with sparse KV cache offload the host history pool. raise ValueError( f"{resource_type.value} requires {required_blocks} KV cache " f"blocks to complete the request, which exceeds its " - f"GPU-primary capacity of {primary_capacity} blocks " + f"capacity of {primary_capacity} blocks " f"(prompt_len={request.orig_prompt_len}, " f"max_new_tokens={request.max_new_tokens}, " f"beam_width={request.py_beam_width}).") diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 85df2786ccc3..f5313635d2e0 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -938,15 +938,17 @@ class SparseKvCacheOffloadConfig(StrictBaseModel): status="prototype", description= "Pinned host memory in bytes, per rank, for offloaded attention history " - "and the staging buffers that publish it to the GPU. History pages share " - "slots with the GPU KV cache pools, so the GPU KV cache quota is lowered " - "when needed to stay within it. It is separate from " - "kv_cache_config.host_cache_size, the host tier that evicted KV cache " - "pages move to, and the two together must fit this rank's share of " - "available host memory and the pinned-memory limit. When unset, it " - "holds max_batch_size requests at kv_cache_config.avg_seq_len, or " - "max_seq_len when that is unset, capped by the host memory automatic " - "sizing may pin.") + "and the staging buffers that publish it to the GPU. The KV cache " + "manager holds the history pools within it, beside the GPU KV cache " + "quota, which it does not lower; a size too small for the batches the " + "manager must always hold, such as one max_seq_len request, is " + "rejected. It is separate from kv_cache_config.host_cache_size, the " + "host tier that evicted KV cache pages move to, and the two together " + "must fit this rank's share of available host memory and the " + "pinned-memory limit. When unset, it holds max_batch_size requests at " + "kv_cache_config.avg_seq_len, or max_seq_len when that is unset, " + "capped by the host memory automatic sizing may pin, and at least what " + "those batches need.") class BaseSparseAttentionConfig(StrictBaseModel): @@ -1452,11 +1454,12 @@ class DeepSeekSparseAttentionConfig(SeqLenAwareSparseAttentionConfig): default=False, status="prototype", description= - "Offload the sparse attention KV history to host memory with KV cache " - "manager v2. Each decode step fetches the units the indexer selects " - "into a per-request GPU cache; the indexer cache remains on GPU. Only " - "a decode-only worker, such as the generation worker of a " - "disaggregated deployment, can offload.") + "Offload the KV history sparse attention selects from to host memory " + "with KV cache manager v2: the MLA latent cache, or for DeepSeek-V4 " + "the ratio-4 compressed cache. Each decode step fetches the units the " + "indexer selects into a per-request GPU cache; the indexer and " + "sliding-window caches remain on GPU. Only a decode-only worker, such " + "as the generation worker of a disaggregated deployment, can offload.") kv_cache_offload_config: SparseKvCacheOffloadConfig = Field( default_factory=SparseKvCacheOffloadConfig, status="prototype", @@ -1688,14 +1691,6 @@ class DeepSeekV4SparseAttentionConfig(DeepSeekSparseAttentionConfig): description="The sliding window size in tokens for SWA layers.") index_topk: Optional[int] = Field(default=512, description="The top-k for the indexer.") - enable_kv_cache_offload: bool = Field( - default=False, - status="prototype", - description= - "Offload ratio-4 compressed attention KV history to host memory with " - "KV cache manager v2. The indexer and sliding-window caches remain on " - "GPU. Only a decode-only worker, such as the generation worker of a " - "disaggregated deployment, can offload.") @field_validator("index_head_dim") @classmethod diff --git a/tensorrt_llm/schemas/trtllm-serve-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-config.schema.json index c8caf572fd8d..dfac28e2f130 100644 --- a/tensorrt_llm/schemas/trtllm-serve-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-config.schema.json @@ -1346,7 +1346,7 @@ }, "enable_kv_cache_offload": { "default": false, - "description": "Offload the sparse attention KV history to host memory with KV cache manager v2. Each decode step fetches the units the indexer selects into a per-request GPU cache; the indexer cache remains on GPU. Only a decode-only worker, such as the generation worker of a disaggregated deployment, can offload.", + "description": "Offload the KV history sparse attention selects from to host memory with KV cache manager v2: the MLA latent cache, or for DeepSeek-V4 the ratio-4 compressed cache. Each decode step fetches the units the indexer selects into a per-request GPU cache; the indexer and sliding-window caches remain on GPU. Only a decode-only worker, such as the generation worker of a disaggregated deployment, can offload.", "status": "prototype", "title": "Enable Kv Cache Offload", "type": "boolean" @@ -1550,7 +1550,7 @@ }, "enable_kv_cache_offload": { "default": false, - "description": "Offload ratio-4 compressed attention KV history to host memory with KV cache manager v2. The indexer and sliding-window caches remain on GPU. Only a decode-only worker, such as the generation worker of a disaggregated deployment, can offload.", + "description": "Offload the KV history sparse attention selects from to host memory with KV cache manager v2: the MLA latent cache, or for DeepSeek-V4 the ratio-4 compressed cache. Each decode step fetches the units the indexer selects into a per-request GPU cache; the indexer and sliding-window caches remain on GPU. Only a decode-only worker, such as the generation worker of a disaggregated deployment, can offload.", "status": "prototype", "title": "Enable Kv Cache Offload", "type": "boolean" @@ -6095,7 +6095,7 @@ } ], "default": null, - "description": "Pinned host memory in bytes, per rank, for offloaded attention history and the staging buffers that publish it to the GPU. History pages share slots with the GPU KV cache pools, so the GPU KV cache quota is lowered when needed to stay within it. It is separate from kv_cache_config.host_cache_size, the host tier that evicted KV cache pages move to, and the two together must fit this rank's share of available host memory and the pinned-memory limit. When unset, it holds max_batch_size requests at kv_cache_config.avg_seq_len, or max_seq_len when that is unset, capped by the host memory automatic sizing may pin.", + "description": "Pinned host memory in bytes, per rank, for offloaded attention history and the staging buffers that publish it to the GPU. The KV cache manager holds the history pools within it, beside the GPU KV cache quota, which it does not lower; a size too small for the batches the manager must always hold, such as one max_seq_len request, is rejected. It is separate from kv_cache_config.host_cache_size, the host tier that evicted KV cache pages move to, and the two together must fit this rank's share of available host memory and the pinned-memory limit. When unset, it holds max_batch_size requests at kv_cache_config.avg_seq_len, or max_seq_len when that is unset, capped by the host memory automatic sizing may pin, and at least what those batches need.", "status": "prototype", "title": "Host History Size" } diff --git a/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json index 0aff994712e2..76f30a492944 100644 --- a/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json @@ -1346,7 +1346,7 @@ }, "enable_kv_cache_offload": { "default": false, - "description": "Offload the sparse attention KV history to host memory with KV cache manager v2. Each decode step fetches the units the indexer selects into a per-request GPU cache; the indexer cache remains on GPU. Only a decode-only worker, such as the generation worker of a disaggregated deployment, can offload.", + "description": "Offload the KV history sparse attention selects from to host memory with KV cache manager v2: the MLA latent cache, or for DeepSeek-V4 the ratio-4 compressed cache. Each decode step fetches the units the indexer selects into a per-request GPU cache; the indexer and sliding-window caches remain on GPU. Only a decode-only worker, such as the generation worker of a disaggregated deployment, can offload.", "status": "prototype", "title": "Enable Kv Cache Offload", "type": "boolean" @@ -1550,7 +1550,7 @@ }, "enable_kv_cache_offload": { "default": false, - "description": "Offload ratio-4 compressed attention KV history to host memory with KV cache manager v2. The indexer and sliding-window caches remain on GPU. Only a decode-only worker, such as the generation worker of a disaggregated deployment, can offload.", + "description": "Offload the KV history sparse attention selects from to host memory with KV cache manager v2: the MLA latent cache, or for DeepSeek-V4 the ratio-4 compressed cache. Each decode step fetches the units the indexer selects into a per-request GPU cache; the indexer and sliding-window caches remain on GPU. Only a decode-only worker, such as the generation worker of a disaggregated deployment, can offload.", "status": "prototype", "title": "Enable Kv Cache Offload", "type": "boolean" @@ -7670,7 +7670,7 @@ } ], "default": null, - "description": "Pinned host memory in bytes, per rank, for offloaded attention history and the staging buffers that publish it to the GPU. History pages share slots with the GPU KV cache pools, so the GPU KV cache quota is lowered when needed to stay within it. It is separate from kv_cache_config.host_cache_size, the host tier that evicted KV cache pages move to, and the two together must fit this rank's share of available host memory and the pinned-memory limit. When unset, it holds max_batch_size requests at kv_cache_config.avg_seq_len, or max_seq_len when that is unset, capped by the host memory automatic sizing may pin.", + "description": "Pinned host memory in bytes, per rank, for offloaded attention history and the staging buffers that publish it to the GPU. The KV cache manager holds the history pools within it, beside the GPU KV cache quota, which it does not lower; a size too small for the batches the manager must always hold, such as one max_seq_len request, is rejected. It is separate from kv_cache_config.host_cache_size, the host tier that evicted KV cache pages move to, and the two together must fit this rank's share of available host memory and the pinned-memory limit. When unset, it holds max_batch_size requests at kv_cache_config.avg_seq_len, or max_seq_len when that is unset, capped by the host memory automatic sizing may pin, and at least what those batches need.", "status": "prototype", "title": "Host History Size" } 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 b97ff998088d..a95cdab5df16 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 @@ -4106,33 +4106,21 @@ def test_offload_counts_gpu_history_and_window_pages_exactly() -> None: @pytest.mark.cpu_only -def test_sparse_offload_rejects_kv_pool_rebalance() -> None: +def test_sparse_offload_settings_reach_the_base_manager_check() -> None: + """The base manager rejects what offload cannot honor, such as speculative decoding + or pool rebalancing, before it sizes anything; DeepSeek-V4 hands it both configs.""" sparse_config = DeepSeekV4SparseAttentionConfig(enable_kv_cache_offload=True) - with ( - patch.object(deepseek_v4_cache.KVCacheManagerV2, "__init__") as initialize_cache, - pytest.raises(ValueError, match="enable_kv_pool_rebalance"), - ): - DeepseekV4CacheManager( - kv_cache_config=KvCacheConfig(enable_block_reuse=False, enable_kv_pool_rebalance=True), - kv_cache_type=CacheTypeCpp.SELFKONLY, - num_layers=len(sparse_config.compress_ratios), - max_batch_size=1, - tokens_per_block=128, - max_seq_len=256, - vocab_size=128, - mapping=Mapping(), - sparse_attn_config=sparse_config, - ) + spec_config = MTPDecodingConfig(max_draft_len=1) - initialize_cache.assert_not_called() + class StopAtBaseInit(Exception): + pass - -@pytest.mark.cpu_only -def test_sparse_offload_rejects_speculative_decoding() -> None: - sparse_config = DeepSeekV4SparseAttentionConfig(enable_kv_cache_offload=True) with ( - patch.object(deepseek_v4_cache.KVCacheManagerV2, "__init__") as initialize_cache, - pytest.raises(ValueError, match="does not support speculative decoding"), + patch.object(deepseek_v4_cache, "get_sm_version", return_value=100), + patch.object( + deepseek_v4_cache.KVCacheManagerV2, "__init__", side_effect=StopAtBaseInit + ) as initialize_cache, + pytest.raises(StopAtBaseInit), ): DeepseekV4CacheManager( kv_cache_config=KvCacheConfig(enable_block_reuse=False), @@ -4144,10 +4132,12 @@ def test_sparse_offload_rejects_speculative_decoding() -> None: vocab_size=128, mapping=Mapping(), sparse_attn_config=sparse_config, - spec_config=MTPDecodingConfig(max_draft_len=1), + spec_config=spec_config, ) - initialize_cache.assert_not_called() + kwargs = initialize_cache.call_args.kwargs + assert kwargs["sparse_kv_offload_config"] == sparse_config.get_kv_cache_offload_config() + assert kwargs["spec_config"] is spec_config @pytest.mark.cpu_only @@ -4160,36 +4150,13 @@ def test_offload_capacity_report_names_the_binding_limit() -> None: # four ranks with 400 GiB available. manager._host_history_budget = 50 << 30 # The native group holds 4,096 slots: four requests of 128K tokens. Its pool of - # sparse pages holds 217,728 host bytes per slot; its indexer pool is on the GPU. - sparse_layer = next( - layer for layer in manager.kv_cache_manager_py_config.layers if layer.sparse_cache - ) - sparse_buffer = next(buffer for buffer in sparse_layer.buffers if buffer.is_sparse) - dense_buffer = next(buffer for buffer in sparse_layer.buffers if not buffer.is_sparse) - host_pool = SimpleNamespace( - num_slots=4096, - slot_desc=SimpleNamespace( - variants=[ - SimpleNamespace( - coalesced_buffers=[ - SimpleNamespace( - buffer_ids=[ - SimpleNamespace(layer_id=sparse_layer.layer_id, role=buffer.role) - ] - ) - for buffer in (sparse_buffer, dense_buffer) - ] - ) - ] - ), - pools=[ - SimpleNamespace(pool_index=0, slot_bytes=217_728), - SimpleNamespace(pool_index=1, slot_bytes=52_176), - ], + # sparse pages holds 217,728 host bytes per slot; its indexer pool is on the GPU, + # and the device quota counts only device-backed pools. + manager.impl = SimpleNamespace( + get_quota=lambda level: 100 << 30, + sparse_backing_medium=lambda: CacheTier.HOST_MEM, + get_num_bytes=lambda level, medium: 4096 * 217_728 if medium == CacheTier.HOST_MEM else 0, ) - manager.get_host_backed_pool_groups = lambda: [host_pool] - # The device quota counts only device-backed pools. - manager.impl = SimpleNamespace(get_quota=lambda level: 100 << 30) with ( patch.object(kv_cache_v2_module.logger, "info") as info, diff --git a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_offload.py b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_offload.py index 1215ea98c988..f620c5ab066c 100644 --- a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_offload.py +++ b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_offload.py @@ -530,6 +530,7 @@ def test_attention_uses_arena_base_and_physical_indices_without_host_pages(nvfp4 # The mocked fetch cannot fill this, so stand in for it: one slice per buffer. out_slots={2: tuple(range(physical.shape[1]))}, nvfp4_scratch=torch.empty(1, 8, 512, dtype=torch.float8_e4m3fn) if nvfp4 else None, + coverage=None, ) descriptor = state.layers[2] if nvfp4: @@ -659,18 +660,14 @@ def test_fetch_uses_batch_rows_and_preserves_topk_slots(multiple_roles, rows): @pytest.mark.cpu_only -@pytest.mark.parametrize("invalid", ["disabled", "unprepared", "prefill", "batch", "watermarks"]) -def test_fetch_rejects_unprepared_or_prefill_metadata(invalid): +@pytest.mark.parametrize("invalid", ["disabled", "unprepared", "batch", "watermarks"]) +def test_fetch_rejects_unprepared_metadata(invalid): manager = _manager() - state = SimpleNamespace( - prepared=True, is_prefill=False, batch=object(), history_lengths=object() - ) + state = SimpleNamespace(prepared=True, batch=object(), history_lengths=object()) if invalid == "disabled": manager._enable_kv_cache_offload = False elif invalid == "unprepared": state.prepared = False - elif invalid == "prefill": - state.is_prefill = True elif invalid == "batch": state.batch = None else: @@ -757,7 +754,7 @@ def _host_batch(*, prefill=False): cache.history_length = 0 metadata = SimpleNamespace( sparse_metadata_params=SimpleNamespace(enable_kv_cache_offload=True), - sparse_offload_state=SimpleNamespace(prepared=True, is_prefill=False), + sparse_offload_state=SimpleNamespace(prepared=True), kv_cache_manager=manager, request_ids=[20, 10], beam_width=1, @@ -775,16 +772,24 @@ def _host_batch(*, prefill=False): metadata._validate_ragged_verification = MethodType( DeepseekV4TrtllmAttentionMetadata._validate_ragged_verification, metadata ) + metadata._validate_dsa_sparse_offload_batch = MethodType( + DeepseekV4TrtllmAttentionMetadata._validate_dsa_sparse_offload_batch, metadata + ) return metadata @pytest.mark.cpu_only @pytest.mark.parametrize("prefill", [False, True]) -def test_offload_admits_decode_or_fresh_prefill(prefill): +def test_offload_admits_only_decode(prefill): + # An offloading worker is a disaggregated decode worker: context requests reach it + # through the KV transfer, never as a batch of its own. metadata = _host_batch(prefill=prefill) - DeepseekV4TrtllmAttentionMetadata.validate_sparse_offload_batch(metadata) + if prefill: + with pytest.raises(NotImplementedError, match="decode batches only"): + DeepseekV4TrtllmAttentionMetadata.validate_sparse_offload_batch(metadata) + else: + DeepseekV4TrtllmAttentionMetadata.validate_sparse_offload_batch(metadata) assert not metadata.sparse_offload_state.prepared - assert metadata.sparse_offload_state.is_prefill == prefill @pytest.mark.cpu_only diff --git a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_offload_cache_manager.py b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_offload_cache_manager.py index c6ca0c9626a7..a5369e9009c7 100644 --- a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_offload_cache_manager.py +++ b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_offload_cache_manager.py @@ -22,15 +22,15 @@ from tensorrt_llm._torch.modules.multi_stream_utils import with_multi_stream from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import Role from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest -from tensorrt_llm._utils import TensorWrapper, convert_to_torch_tensor, get_sm_version +from tensorrt_llm._utils import TensorWrapper, convert_to_torch_tensor, is_sm_100f from tensorrt_llm.bindings import DataType, SamplingConfig from tensorrt_llm.bindings.internal.batch_manager import CacheType as CacheTypeCpp from tensorrt_llm.llmapi.llm_args import DeepSeekSparseAttentionConfig, KvCacheConfig from tensorrt_llm.mapping import Mapping pytestmark = pytest.mark.skipif( - not torch.cuda.is_available() or get_sm_version() < 100, - reason="DSA KV cache offload needs an SM100+ GPU", + not torch.cuda.is_available() or not is_sm_100f(), + reason="DSA KV cache offload needs an SM100-family GPU", ) NUM_LAYERS = 7 @@ -326,114 +326,117 @@ def test_fetched_rows_hold_written_history_across_evictions(seed, num_layers, nu """ generator = torch.Generator().manual_seed(seed) manager = _manager(num_layers=num_layers, dtype=dtype) - requests = [_request(9 + index, 1) for index in range(num_requests)] - join_step = {request.py_request_id: 40 * index for index, request in enumerate(requests)} - descriptors = manager.get_sparse_offload_descriptors() - nvfp4 = dtype == DataType.NVFP4 - arenas = { - layer: [_arena_rows(descriptor, descriptor.data_buffer_index, HEAD_DIM)] - for layer, descriptor in descriptors.items() - } - if nvfp4: + try: + requests = [_request(9 + index, 1) for index in range(num_requests)] + join_step = {request.py_request_id: 40 * index for index, request in enumerate(requests)} + descriptors = manager.get_sparse_offload_descriptors() + nvfp4 = dtype == DataType.NVFP4 arenas = { - layer: [ - _arena_rows(descriptor, descriptor.data_buffer_index, NVFP4_DATA_BYTES), - _arena_rows(descriptor, descriptor.scale_buffer_index, NVFP4_SCALE_BYTES), - ] + layer: [_arena_rows(descriptor, descriptor.data_buffer_index, HEAD_DIM)] for layer, descriptor in descriptors.items() } - state = allocate_sparse_offload_state( - lambda name, shape, dtype: torch.empty(shape, dtype=dtype, device="cuda"), - descriptors, - max_batch_size=MAX_BATCH_SIZE, - topk=TOPK, - ) - active: list[LlmRequest] = [] - permuted = False - for step in range(6 * TOKENS_PER_BLOCK): - for request in requests: - if join_step[request.py_request_id] == step: - assert manager.prepare_disagg_gen_init(request) - active.append(request) - for request in active: - assert manager.try_allocate_generation(request) - scheduled = SimpleNamespace(context_requests=[], generation_requests=active[::-1]) - manager.prepare_resources(scheduled) - request_ids = [request.py_request_id for request in scheduled.generation_requests] - manager.prepare_sparse_offload(state, request_ids, beam_width=1) - # The rows the manager recorded at add_kv_cache are still the native ones. - for member in manager._sparse_batch_members.values(): - assert manager._sparse_batch.row_of(member.cache) == member.row - batch_rows = state.batch_rows.tolist()[: len(active)] - permuted |= batch_rows != list(range(len(active))) - histories = [manager.kv_cache_map[rid].history_length for rid in request_ids] - topk = torch.full((len(active), TOPK), -1, dtype=torch.int32) - for position, history in enumerate(histories): - # Unit 0 came with the prompt and was never written; select from the rest, - # always including the token written this step. - candidates = torch.randperm(history, generator=generator)[: TOPK - 1] - candidates = candidates[candidates > 0][: TOPK - 1].to(torch.int32) - topk[position, : candidates.numel()] = candidates - topk[position, -1] = history - topk = topk.cuda() - request_indices = torch.arange(len(active), dtype=torch.int32, device="cuda") - kv_lengths = torch.tensor([h + 1 for h in histories], dtype=torch.int32, device="cuda") - out = torch.empty((len(active), TOPK), dtype=torch.int32, device="cuda") - for layer in range(num_layers): - descriptor = descriptors[layer] - data = descriptor.data_buffer_index - writes = manager.prepare_sparse_write_indices(state, layer) - # At the buffer's slice, not its position: buffers whose arena indices - # coincide share one, and the slices no buffer maps to are never written. - data = state.out_slots[layer][data] - for row, rid, history in zip(batch_rows, request_ids, histories): - write_row = writes[row, data].item() - assert write_row >= 0, (rid, history, layer) - # The token's row: its position, its layer, its request. - for arena in arenas[layer]: - arena[write_row, :3] = torch.tensor( - [history, layer, rid], dtype=torch.int32, device="cuda" - ) - free_rows = [row for row in range(MAX_BATCH_SIZE) if row not in batch_rows] - assert (writes[free_rows, data] == -1).all() - fetched = manager.fetch_sparse_indices(state, layer, topk, request_indices, kv_lengths) - if nvfp4: - # The scale mirrors the data pool, so it resolves to the same index and - # reads the data's slice rather than one of its own. One slice written - # once is what replaces the two equal rows this used to compare. - assert state.out_slots[layer][descriptor.scale_buffer_index] == data - merge_sparse_arena_indices( - topk, - request_indices, - state.batch_rows, - state.active_request_count, - kv_lengths, - state.history_lengths, - fetched[:, data, :], - writes[:, data], - out, - unit_tokens=1, - ) - for position, rid in enumerate(request_ids): - valid = topk[position] >= 0 - units = topk[position][valid] - expected = torch.stack( - (units, torch.full_like(units, layer), torch.full_like(units, rid)), dim=1 + if nvfp4: + arenas = { + layer: [ + _arena_rows(descriptor, descriptor.data_buffer_index, NVFP4_DATA_BYTES), + _arena_rows(descriptor, descriptor.scale_buffer_index, NVFP4_SCALE_BYTES), + ] + for layer, descriptor in descriptors.items() + } + state = allocate_sparse_offload_state( + lambda name, shape, dtype: torch.empty(shape, dtype=dtype, device="cuda"), + descriptors, + max_batch_size=MAX_BATCH_SIZE, + topk=TOPK, + ) + active: list[LlmRequest] = [] + permuted = False + for step in range(6 * TOKENS_PER_BLOCK): + for request in requests: + if join_step[request.py_request_id] == step: + assert manager.prepare_disagg_gen_init(request) + active.append(request) + for request in active: + assert manager.try_allocate_generation(request) + scheduled = SimpleNamespace(context_requests=[], generation_requests=active[::-1]) + manager.prepare_resources(scheduled) + request_ids = [request.py_request_id for request in scheduled.generation_requests] + manager.prepare_sparse_offload(state, request_ids, beam_width=1) + # The rows the manager recorded at add_kv_cache are still the native ones. + for member in manager._sparse_batch_members.values(): + assert manager._sparse_batch.row_of(member.cache) == member.row + batch_rows = state.batch_rows.tolist()[: len(active)] + permuted |= batch_rows != list(range(len(active))) + histories = [manager.kv_cache_map[rid].history_length for rid in request_ids] + topk = torch.full((len(active), TOPK), -1, dtype=torch.int32) + for position, history in enumerate(histories): + # Unit 0 came with the prompt and was never written; select from the rest, + # always including the token written this step. + candidates = torch.randperm(history, generator=generator)[: TOPK - 1] + candidates = candidates[candidates > 0][: TOPK - 1].to(torch.int32) + topk[position, : candidates.numel()] = candidates + topk[position, -1] = history + topk = topk.cuda() + request_indices = torch.arange(len(active), dtype=torch.int32, device="cuda") + kv_lengths = torch.tensor([h + 1 for h in histories], dtype=torch.int32, device="cuda") + out = torch.empty((len(active), TOPK), dtype=torch.int32, device="cuda") + for layer in range(num_layers): + descriptor = descriptors[layer] + data = descriptor.data_buffer_index + writes = manager.prepare_sparse_write_indices(state, layer) + # At the buffer's slice, not its position: buffers whose arena indices + # coincide share one, and the slices no buffer maps to are never written. + data = state.out_slots[layer][data] + for row, rid, history in zip(batch_rows, request_ids, histories): + write_row = writes[row, data].item() + assert write_row >= 0, (rid, history, layer) + # The token's row: its position, its layer, its request. + for arena in arenas[layer]: + arena[write_row, :3] = torch.tensor( + [history, layer, rid], dtype=torch.int32, device="cuda" + ) + free_rows = [row for row in range(MAX_BATCH_SIZE) if row not in batch_rows] + assert (writes[free_rows, data] == -1).all() + fetched = manager.fetch_sparse_indices( + state, layer, topk, request_indices, kv_lengths ) - for arena in arenas[layer]: - got = arena[out[position][valid].long(), :3] - wrong = (got != expected).any(dim=1) - assert not wrong.any(), ( - f"request {rid}, history {histories[position]}, layer {layer}: units " - f"{units[wrong].tolist()} read {got[wrong].tolist()}" + if nvfp4: + # The scale mirrors the data pool, so it resolves to the same index and + # reads the data's slice rather than one of its own. + assert state.out_slots[layer][descriptor.scale_buffer_index] == data + merge_sparse_arena_indices( + topk, + request_indices, + state.batch_rows, + state.active_request_count, + kv_lengths, + state.history_lengths, + fetched[:, data, :], + writes[:, data], + out, + unit_tokens=1, + ) + for position, rid in enumerate(request_ids): + valid = topk[position] >= 0 + units = topk[position][valid] + expected = torch.stack( + (units, torch.full_like(units, layer), torch.full_like(units, rid)), dim=1 ) - manager.update_forward_resources(scheduled, wait_for_input_copy=lambda: None) - for rid, history in zip(request_ids, histories): - assert manager.kv_cache_map[rid].history_length == history + 1 - assert permuted == (num_requests > 1) - for request in active: - manager.free_resources(request) - manager.shutdown() + for arena in arenas[layer]: + got = arena[out[position][valid].long(), :3] + wrong = (got != expected).any(dim=1) + assert not wrong.any(), ( + f"request {rid}, history {histories[position]}, layer {layer}: units " + f"{units[wrong].tolist()} read {got[wrong].tolist()}" + ) + manager.update_forward_resources(scheduled, wait_for_input_copy=lambda: None) + for rid, history in zip(request_ids, histories): + assert manager.kv_cache_map[rid].history_length == history + 1 + assert permuted == (num_requests > 1) + for request in active: + manager.free_resources(request) + finally: + manager.shutdown() def test_nvfp4_attention_reads_the_selected_rows_dequantized(): @@ -537,12 +540,14 @@ def test_nvfp4_attention_reads_the_selected_rows_dequantized(): manager.shutdown() -def test_offload_scratch_counts_the_latent_writer(manager): +def test_offload_scratch_counts_the_latent_writer_and_attention_indices(manager): with_writer = manager.get_sparse_offload_scratch_bytes(manager.kv_cache_manager_py_config) state_only = super(DSACacheManagerV2, manager).get_sparse_offload_scratch_bytes( manager.kv_cache_manager_py_config ) - assert with_writer - state_only == 2 * 3 * MAX_BATCH_SIZE * 4 + # Each metadata copy holds a K and a V offset and a write slot per row, and the + # top-k indices the attention reads per row. + assert with_writer - state_only == 2 * (3 + TOPK) * MAX_BATCH_SIZE * 4 @pytest.mark.parametrize("dtype", [DataType.FP8, DataType.NVFP4]) diff --git a/tests/unittest/_torch/attention/sparse/test_sparse_offload_kernels.py b/tests/unittest/_torch/attention/sparse/test_sparse_offload_kernels.py index a7e2ea409989..eafb6cc124df 100644 --- a/tests/unittest/_torch/attention/sparse/test_sparse_offload_kernels.py +++ b/tests/unittest/_torch/attention/sparse/test_sparse_offload_kernels.py @@ -2,10 +2,13 @@ # SPDX-License-Identifier: Apache-2.0 """Tests for the kernels that assemble sparse offload arena rows.""" +from unittest.mock import patch + import pytest import torch from tensorrt_llm._torch.attention.backends.sparse.offload_kernels import ( + SPARSE_ARENA_COVERAGE_MESSAGE, gather_sparse_batch_rows, merge_sparse_arena_indices, select_sparse_history_units, @@ -168,6 +171,43 @@ def test_merge_replays_in_a_cuda_graph_with_new_watermarks_and_rows(): torch.testing.assert_close(out[0].cpu(), torch.tensor(expected0, dtype=torch.int32)) +@_requires_cuda +@pytest.mark.parametrize("missing", [False, True], ids=["covered", "missing_fetch"]) +def test_merge_coverage_flags_a_live_selection_without_a_row(missing: bool): + # Request 0 holds two history units and writes its third; request 1's last + # selection is past its length, so it needs no row. + topk = _cuda_tensor([[0, 1, 2], [0, -1, 3]]) + requests = _cuda_tensor([0, 1]) + batch_rows = _cuda_tensor([0, 1]) + active = _cuda_tensor([2]) + history = _cuda_tensor([2, 1]) + lengths = _cuda_tensor([3, 2]) + fetched = _cuda_tensor([[10, -1 if missing else 11, -1], [20, -1, -1]]) + writes = _cuda_tensor([7, 8]) + out = torch.empty((2, 3), dtype=torch.int32, device="cuda") + coverage = _cuda_tensor([0]) + + # A failed device assertion would leave the CUDA context unusable, so read the flag. + with patch.object(torch, "_assert_async") as assert_async: + merge_sparse_arena_indices( + topk, + requests, + batch_rows, + active, + lengths, + history, + fetched, + writes, + out, + unit_tokens=1, + coverage=coverage, + ) + + assert coverage.item() == (0 if missing else 1) + assert_async.assert_called_once_with(coverage, SPARSE_ARENA_COVERAGE_MESSAGE) + assert out.tolist() == [[10, -1 if missing else 11, 7], [20, -1, -1]] + + @_requires_cuda @pytest.mark.parametrize( "invalid", ["dtype", "width", "requests", "batch_rows", "writes", "lengths", "unit"] diff --git a/tests/unittest/_torch/compilation/test_phase_selective_forward.py b/tests/unittest/_torch/compilation/test_phase_selective_forward.py index 80c6c69ba206..47b0b4120d10 100644 --- a/tests/unittest/_torch/compilation/test_phase_selective_forward.py +++ b/tests/unittest/_torch/compilation/test_phase_selective_forward.py @@ -236,6 +236,7 @@ def test_decoder_runner_warmup_bypasses_only_auxiliary_phases(monkeypatch, enabl graphs.preallocate_padding_dummies = Mock() warmup = object.__new__(DecoderWarmup) warmup.mapping = SimpleNamespace(cp_size=1, has_cp_helix=lambda: False) + warmup.sparse_attention_config = None warmup._inputs = SimpleNamespace(guided_decoder=None) warmup._metrics = {} warmup._warmup_timer = _WarmupTimer(rank=0) @@ -249,7 +250,7 @@ def test_decoder_runner_warmup_bypasses_only_auxiliary_phases(monkeypatch, enabl observed = [] def observe(phase): - def forward(*args): + def forward(*args, **kwargs): observed.append((phase, model.use_compiled_forward())) return forward diff --git a/tests/unittest/_torch/executor/engine/runners/decoder/test_warmup.py b/tests/unittest/_torch/executor/engine/runners/decoder/test_warmup.py index d596a1b580f5..8459bcb8dc2f 100644 --- a/tests/unittest/_torch/executor/engine/runners/decoder/test_warmup.py +++ b/tests/unittest/_torch/executor/engine/runners/decoder/test_warmup.py @@ -1497,6 +1497,16 @@ def test_generation_warmup_shapes_drop_context_tokens( assert DecoderWarmup._generation_warmup_shapes(warmup, configs) == expected +@pytest.mark.cpu_only +def test_generation_warmup_shapes_keep_generation_batches_stated_without_draft_tokens() -> None: + """The general and max-shape warmups state a generation batch as one token per + request whatever the draft length; it is still a pure-generation batch.""" + warmup = SimpleNamespace(_config=SimpleNamespace(max_total_draft_tokens=2)) + configs = [(1, 0), (1, 1), (258, 0), (13, 13), (2, 0)] + + assert DecoderWarmup._generation_warmup_shapes(warmup, configs) == [(3, 1), (39, 13)] + + @pytest.mark.cpu_only def test_decode_only_skips_prefill_cuda_graph_capture() -> None: """A decode-only worker never runs a context batch, so prefill graphs are dead weight.""" 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 fe3a2cdba1fb..23e0ac665091 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 @@ -176,11 +176,11 @@ def test_recurrent_state_is_refused(kit, hybrid_manager): def sparse_layers(monkeypatch, indices): - """Managers built from now on give their layers ``indices`` sparse attention, which no model + """Managers built from now on give their layers ``indices`` sparse buffers, 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.""" + that would select from them. ``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 diff --git a/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py b/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py index f220b6361411..19db596d6936 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py @@ -73,12 +73,14 @@ from tensorrt_llm.runtime.kv_cache_manager_v2 import ( BAD_PAGE_INDEX, DEFAULT_BEAM_INDEX, + GPU_LEVEL, AttentionLayerConfig, AttnLifeCycle, Batch, BatchDesc, BufferConfig, CacheLevel, + CacheTier, CudaStream, DataRole, DiskCacheTierConfig, @@ -1387,19 +1389,29 @@ def test_avg_seq_len_builds_context_warmup_constraint() -> None: assert config.constraints == [BatchDesc([KVCacheDesc(capacity=2048, history_length=0)])] -def test_sparse_kv_offload_typical_step_models_only_decoding() -> None: +@pytest.mark.parametrize( + ("avg_seq_len", "capacity"), [(1024, 1024), (None, 2048)], ids=["avg_seq_len", "max_seq_len"] +) +def test_sparse_kv_offload_typical_step_models_only_decoding( + avg_seq_len: int | None, capacity: int +) -> None: config = _make_cache_config_for_test( - KvCacheConfig(host_cache_size=0, avg_seq_len=1024), + KvCacheConfig(host_cache_size=0, avg_seq_len=avg_seq_len), max_batch_size=3, - max_seq_len=1024, + max_seq_len=2048, max_num_tokens=2048, max_draft_len=2, sparse_kv_offload=SparseKvCacheOffloadConfig(), ) - # A worker that offloads sparse history only decodes. - assert config.typical_step == BatchDesc([KVCacheDesc(capacity=1024, history_length=1021)] * 3) - assert config.constraints == [BatchDesc([KVCacheDesc(capacity=2048, history_length=0)])] + # A worker that offloads sparse history only decodes requests of avg_seq_len, or + # else of max_seq_len, and its warmup runs no context batch to hold room for. The + # manager splits the quota from this step, host-backed pools included. + assert config.typical_step == BatchDesc( + [KVCacheDesc(capacity=capacity, history_length=capacity - 3)] * 3 + ) + assert config.constraints == [] + assert config.initial_pool_ratio is None def test_avg_seq_len_updates_typical_step() -> None: @@ -3309,10 +3321,15 @@ def build_base_config( @pytest.mark.cpu_only -@pytest.mark.parametrize("sparse", [False, True]) @pytest.mark.parametrize( - "max_batch_size,pp_size,disable_overlap_scheduler,is_disagg,reserved,expected", - _INDEX_MAPPER_CAPACITY_CASES, + "max_batch_size,pp_size,disable_overlap_scheduler,is_disagg,reserved,expected,sparse", + [ + pytest.param(*case.values, sparse, id=f"{case.id}-{sparse}") + for case in _INDEX_MAPPER_CAPACITY_CASES + for sparse in (False, True) + # Sparse offload rejects pipeline parallelism. + if not (sparse and case.values[1] > 1) + ], ) def test_index_mapper_capacity_covers_the_overlapping_cohorts( max_batch_size: int, @@ -3654,9 +3671,7 @@ def plan_host_history_budget( # A lease follows Batch membership, so the arena is sized by the rows rather than by # the 5 indices the mapper can hand out -- disagg doubling and padding included. get_arena_bytes.assert_called_once_with(ANY, tokens_per_block=TOKENS_PER_BLOCK, max_leases=2) - get_batch_buffer_bytes.assert_called_once_with( - ANY, tokens_per_block=TOKENS_PER_BLOCK, max_leases=2 - ) + get_batch_buffer_bytes.assert_called_once_with(ANY, tokens_per_block=TOKENS_PER_BLOCK) return create_impl.call_args.args[0] @@ -3713,20 +3728,21 @@ def quota(arena_bytes: int, batch_buffer_bytes: int) -> int: @pytest.mark.cpu_only @pytest.mark.parametrize( - ("arena_bytes", "sizing_result", "error", "match"), + ("arena_bytes", "error", "match"), [ - # Rank 1 fits its budget, but rank 0 does not. - (0, 0, RuntimeError, "Sizing sparse KV cache offload failed on rank 0"), + # Rank 1 fits its budget, but another rank does not. + (0, RuntimeError, "Sizing sparse KV cache offload failed on another rank"), # Rank 1 itself does not fit, and raises its own error. - (16 << 20, 1, ValueError, "Sparse KV cache offload needs"), + (16 << 20, ValueError, "Sparse KV cache offload needs"), ], ) def test_offload_sizing_fails_on_every_rank_before_the_token_sync( - arena_bytes: int, sizing_result: int, error: type, match: str + arena_bytes: int, error: type, match: str ) -> None: + ok, abort = int(_KVCacheManagerInitStatus.OK), int(_KVCacheManagerInitStatus.ABORT) dist = Mock() - # The config check passes everywhere; then sizing reports the lowest failing rank. - dist.allreduce.side_effect = [2, sizing_result] + # The config check passes everywhere; then sizing fails on some rank. + dist.allreduce.side_effect = [ok, abort] with ( patch.object(Distributed, "get", return_value=dist), @@ -3740,10 +3756,10 @@ def test_offload_sizing_fails_on_every_rank_before_the_token_sync( ) # Every rank stops after agreeing; none is left in the token sync. - this_rank = 2 if arena_bytes == 0 else 1 + this_rank = ok if arena_bytes == 0 else abort assert dist.allreduce.call_args_list == [ - call(2, op=ReduceOp.MIN), - call(this_rank, op=ReduceOp.MIN), + call(ok, op=ReduceOp.MAX), + call(this_rank, op=ReduceOp.MAX), ] @@ -3830,19 +3846,22 @@ def test_offload_warns_when_it_runs_without_a_cold_tier() -> None: @pytest.mark.cpu_only @pytest.mark.parametrize( - ("error", "lowest_failing_rank", "raised", "match"), + ("error", "agreed_status", "raised", "match"), [ - (None, 4, None, None), - (None, 1, RuntimeError, "Checking failed on rank 1"), - (ValueError("no room"), 1, ValueError, "no room"), + (None, _KVCacheManagerInitStatus.OK, None, None), + (None, _KVCacheManagerInitStatus.ABORT, RuntimeError, "Checking failed on another rank"), + (ValueError("no room"), _KVCacheManagerInitStatus.ABORT, ValueError, "no room"), ], ) def test_raise_on_every_rank_agrees_across_ranks( - error: Exception | None, lowest_failing_rank: int, raised: type | None, match: str | None + error: Exception | None, + agreed_status: _KVCacheManagerInitStatus, + raised: type | None, + match: str | None, ) -> None: mapping = SimpleNamespace(world_size=4, rank=2) dist = Mock() - dist.allreduce.return_value = lowest_failing_rank + dist.allreduce.return_value = int(agreed_status) with patch.object(Distributed, "get", return_value=dist): if raised is None: @@ -3851,8 +3870,10 @@ def test_raise_on_every_rank_agrees_across_ranks( with pytest.raises(raised, match=match): _raise_on_every_rank(error, mapping, "Checking") - # A passing rank offers world_size, so the minimum names the lowest failing rank. - dist.allreduce.assert_called_once_with(2 if error else 4, op=ReduceOp.MIN) + # Every rank offers its own outcome through the init-status sync, and the most + # severe one is every rank's. + local_status = _KVCacheManagerInitStatus.ABORT if error else _KVCacheManagerInitStatus.OK + dist.allreduce.assert_called_once_with(int(local_status), op=ReduceOp.MAX) @pytest.mark.cpu_only @@ -3921,7 +3942,7 @@ def test_sparse_batch_buffers_come_from_a_provisional_config(offload: bool) -> N ) as get_scratch_bytes, ): device_bytes, staging_bytes = manager._get_sparse_batch_buffer_bytes( - KvCacheConfig(), tokens_per_block=TOKENS_PER_BLOCK, max_leases=5 + KvCacheConfig(), tokens_per_block=TOKENS_PER_BLOCK ) if not offload: @@ -3934,7 +3955,15 @@ def test_sparse_batch_buffers_come_from_a_provisional_config(offload: bool) -> N ) assert batch.num_life_cycles == 2 assert batch.selection_widths == (1,) - assert device_bytes == batch.block_bytes + batch.mirror_bytes + batch.fetch_scratch_bytes + 1000 + # The Batch also converts the sparse buffer's indices into one table per row. + assert batch.index_table_bytes == 4 * batch.index_table_entries * 2 * 8 > 0 + assert device_bytes == ( + batch.block_bytes + + batch.mirror_bytes + + batch.index_table_bytes + + batch.fetch_scratch_bytes + + 1000 + ) assert staging_bytes == batch.staging_bytes == 4096 get_scratch_bytes.assert_called_once_with(config) (gpu_tier,) = build.call_args.kwargs["cache_tiers"] @@ -4053,12 +4082,10 @@ def _offload_setup_manager( manager.num_extra_kv_tokens = 0 manager.max_total_draft_tokens = max_total_draft_tokens manager.tokens_per_block = 128 - manager._host_history_budget = 1 << 40 manager.kv_cache_map = {kv_cache_v2_module._GUARD_PAGE_REQUEST_ID: Mock()} if guard_page else {} manager.get_host_backed_pool_groups = lambda: [ SimpleNamespace(num_slots=n) for n in local_slots ] - manager._get_host_backed_pool_bytes = lambda: 0 return manager @@ -4067,8 +4094,8 @@ def _offload_setup_manager( ("world_size", "local_slots", "peer_slots", "capacity_blocks"), [ (1, [12, 10], None, 10), - # Every rank judges arriving requests by the smallest pool, including a - # rank whose layers hold no host history, but reserves from its own. + # Every rank judges arriving requests and admits them by the smallest pool, + # including a rank whose layers hold no host history. (2, [12, 10], 8, 8), (2, [], 8, 8), (2, [], sys.maxsize, None), @@ -4091,17 +4118,17 @@ def test_sparse_offload_setup_judges_requests_by_the_smallest_host_pool( ) if world_size > 1: - # The ranks agree that their pools fit, then on the smallest pool. - assert dist.allreduce.call_args_list == [ - call(world_size, op=ReduceOp.MIN), - call(min(local_slots, default=sys.maxsize), op=ReduceOp.MIN), - ] + dist.allreduce.assert_called_once_with( + min(local_slots, default=sys.maxsize), op=ReduceOp.MIN + ) else: dist.allreduce.assert_not_called() assert manager._host_history_capacity_blocks == capacity_blocks assert manager.kv_block_budget_applies() == (capacity_blocks is not None) if local_slots: - assert manager._host_history_reservations.capacity_blocks == min(local_slots) + # A rank that admitted against its own larger pool would serve requests + # its peers left waiting, and wait for them in collectives alone. + assert manager._host_history_reservations.capacity_blocks == capacity_blocks else: assert manager._host_history_reservations is None # Eight blocks of 128 tokens hold a max_seq_len request of 1024 tokens. @@ -4126,50 +4153,6 @@ def test_sparse_offload_setup_warns_when_max_seq_len_outgrows_host_history() -> assert "kv_cache_offload_config.host_history_size" in message -@pytest.mark.cpu_only -@pytest.mark.parametrize( - ("world_size", "host_pool_bytes", "lowest_failing_rank", "error", "match"), - [ - # Only the native slot floors can take the pools past the budget. - (1, 2 << 30, None, ValueError, "take 2.00GiB, more than their 1.00GiB host budget"), - # This rank's pools are too large; it raises its own error once every - # rank has agreed. - (2, 2 << 30, 1, ValueError, "take 2.00GiB, more than their 1.00GiB host budget"), - # Rank 0's pools are too large; this rank's fit. - (2, 1 << 30, 0, RuntimeError, "host-backed history pools failed on rank 0"), - ], -) -def test_sparse_offload_setup_rejects_host_pools_over_their_budget( - world_size: int, - host_pool_bytes: int, - lowest_failing_rank: int | None, - error: type, - match: str, -) -> None: - manager = _offload_setup_manager([8], world_size=world_size) - manager.mapping.rank = world_size - 1 - manager._host_history_budget = 1 << 30 - manager._get_host_backed_pool_bytes = lambda: host_pool_bytes - dist = Mock() - dist.allreduce.return_value = lowest_failing_rank - - with ( - patch.object(kv_cache_v2_module.Distributed, "get", return_value=dist), - pytest.raises(error, match=match), - ): - manager._setup_sparse_kv_offload( - KvCacheConfig(), max_seq_len=1024, num_reserved_index_slots=0 - ) - - # No rank reaches the capacity sync, and none starts admitting requests. - if world_size > 1: - this_rank = world_size - 1 if host_pool_bytes > 1 << 30 else world_size - dist.allreduce.assert_called_once_with(this_rank, op=ReduceOp.MIN) - else: - dist.allreduce.assert_not_called() - assert manager._host_history_reservations is None - - @pytest.mark.cpu_only @pytest.mark.parametrize( ("num_reserved_index_slots", "attention_dp", "guard_page", "draft_tokens", "set_aside"), @@ -4310,6 +4293,19 @@ def test_sparse_kv_offload_rejects_block_reuse() -> None: ) +@pytest.mark.cpu_only +def test_sparse_kv_offload_rejects_pipeline_parallelism() -> None: + # The attention Batch holds one micro-batch's members; publishing the next one would + # release the leases of the first while it is still in flight. + with pytest.raises(ValueError, match="does not support pipeline parallelism"): + KVCacheManagerV2._check_sparse_kv_offload_config( + SparseKvCacheOffloadConfig(), + KvCacheConfig(enable_block_reuse=False), + Mapping(world_size=2, rank=0, pp_size=2), + spec_config=None, + ) + + @pytest.mark.cpu_only def test_sparse_kv_offload_config_check_fails_on_every_rank(monkeypatch) -> None: # This rank's node has room, but rank 0's does not. @@ -4319,11 +4315,11 @@ def test_sparse_kv_offload_config_check_fails_on_every_rank(monkeypatch) -> None lambda mapping: HostMemoryLimits(4, float(400 << 30), float("inf")), ) dist = Mock() - dist.allreduce.return_value = 0 + dist.allreduce.return_value = int(_KVCacheManagerInitStatus.ABORT) with ( patch.object(Distributed, "get", return_value=dist), - pytest.raises(RuntimeError, match="config check failed on rank 0"), + pytest.raises(RuntimeError, match="config check failed on another rank"), ): KVCacheManagerV2._check_sparse_kv_offload_config( SparseKvCacheOffloadConfig(host_history_size=50 << 30), @@ -4332,7 +4328,7 @@ def test_sparse_kv_offload_config_check_fails_on_every_rank(monkeypatch) -> None spec_config=None, ) - dist.allreduce.assert_called_once_with(2, op=ReduceOp.MIN) + dist.allreduce.assert_called_once_with(int(_KVCacheManagerInitStatus.OK), op=ReduceOp.MAX) def _offload_layers() -> list[AttentionLayerConfig]: @@ -4381,7 +4377,7 @@ def _base_config(layers: list, initial_pool_ratio: list[float] | None = None): @pytest.mark.cpu_only @pytest.mark.parametrize("pinned", [False, True]) -def test_offload_splits_the_ratio_by_all_bytes_unless_pinned(pinned: bool) -> None: +def test_offload_leaves_the_pool_ratio_to_the_manager_unless_pinned(pinned: bool) -> None: manager = _offload_manager() build_base, build_cache = _base_config(_offload_layers(), [0.25, 0.75] if pinned else None) @@ -4392,66 +4388,51 @@ def test_offload_splits_the_ratio_by_all_bytes_unless_pinned(pinned: bool) -> No cache_tiers=[GpuCacheTierConfig(quota=1 << 20)], ) - if pinned: - assert config.initial_pool_ratio == [0.25, 0.75] - return - # A 16-token request holds four pages in the no-window life cycle and three pages of - # 4 KiB in the window-8 one. A no-window page is 1280 bytes: 1024 host-backed sparse - # beside 256 device. The ratio weighs both media -- it is projected onto cold levels, - # where one pool holds them together, and compared against a layout measured the same - # way. That host bytes cost the GPU quota nothing is a fact about the quota alone. - no_window = 4 * (1024 + 256) - windowed = 3 * 4096 - total = no_window + windowed - assert config.initial_pool_ratio == pytest.approx([no_window / total, windowed / total]) + # The manager splits the quota from the decode-only typical step itself, counting + # host-backed pools in their own grain, so offload hands it no ratio of its own. + assert config.initial_pool_ratio == ([0.25, 0.75] if pinned else None) @pytest.mark.cpu_only -def test_host_backed_bytes_count_only_the_sparse_pools() -> None: +def test_host_backed_pools_come_from_the_native_pool_descriptions() -> None: manager = object.__new__(KVCacheManagerV2) - manager.kv_cache_manager_py_config = KVCacheManagerConfig( - tokens_per_block=TOKENS_PER_BLOCK, - cache_tiers=[GpuCacheTierConfig(quota=1 << 20)], - layers=_offload_layers(), - ) - def group(num_slots: int, *pools: tuple[int, int, str]) -> SimpleNamespace: + def group(*media: CacheTier) -> SimpleNamespace: return SimpleNamespace( - num_slots=num_slots, - slot_desc=SimpleNamespace( - variants=[ - SimpleNamespace( - coalesced_buffers=[ - SimpleNamespace( - buffer_ids=[SimpleNamespace(layer_id=LayerId(layer), role=role)] - ) - for layer, _, role in pools - ] - ) - ] - ), pools=[ - SimpleNamespace(pool_index=index, slot_bytes=slot_bytes) - for index, (_, slot_bytes, _) in enumerate(pools) - ], + SimpleNamespace(pool_index=index, storage_medium=medium) + for index, medium in enumerate(media) + ] ) - mixed = group(10, (0, 1024, Role.KEY), (0, 256, Role.VALUE)) - dense = group(7, (1, 4096, Role.KEY)) - manager.impl = SimpleNamespace(pool_group_descs=[mixed, dense]) + mixed = group(CacheTier.HOST_MEM, CacheTier.GPU_MEM) + dense = group(CacheTier.GPU_MEM) + get_num_bytes = Mock(return_value=10 << 20) + manager.impl = SimpleNamespace( + pool_group_descs=[mixed, dense], + sparse_backing_medium=lambda: CacheTier.HOST_MEM, + get_num_bytes=get_num_bytes, + ) + # A group keeps its dense pools on the GPU; each pool reports its own medium. assert manager.get_host_backed_pool_groups() == [mixed] - # The mixed group's dense pool is GPU memory; only its sparse pool is host memory. - assert manager._get_host_backed_pool_bytes() == 10 * 1024 + # The native count includes how each pool rounds to the host grain. + assert manager._get_host_backed_pool_bytes() == 10 << 20 + get_num_bytes.assert_called_once_with(GPU_LEVEL, CacheTier.HOST_MEM) + + manager.impl = SimpleNamespace(pool_group_descs=[dense], sparse_backing_medium=lambda: None) + assert manager.get_host_backed_pool_groups() == [] + assert manager._get_host_backed_pool_bytes() == 0 @pytest.mark.cpu_only def test_offload_hands_its_host_budget_to_the_manager(monkeypatch) -> None: """The device quota is not lowered to fit the host budget; the manager divides both. - Bounding the quota beforehand shrank every pool group, device-only ones included, to - hold back host memory only some of them spend. The manager is given the budget instead - and holds the host-backed pools inside it while the rest keep the device bytes. + Lowering the quota to hold back host memory would shrink every pool group, device-only + ones included, though only some of them spend host memory. The manager is given the + budget instead and holds the host-backed pools inside it while the rest keep the device + bytes. """ module = "tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2" @@ -4511,15 +4492,12 @@ def build_base_config( config = create_impl.call_args.args[0] gpu_tier, host_tier = config.cache_tiers - assert config.initial_pool_ratio == [1.0] - # Only the offload allocations come off the 64 MiB budget now. Under the old - # bound this was about 46 MiB, the host budget divided by the 4 host bytes each - # device byte carries -- a reduction every pool group paid, not just the ones - # holding host-backed pools. + # The manager splits both from its typical step, host-backed pools included. + assert config.initial_pool_ratio is None + # Only the offload allocations come off the 64 MiB budget. assert gpu_tier.quota > (63 << 20) - # The 200 MiB of host history less the sparse Batch's one-page staging copy. No - # granule holdback: the manager counts host bytes in their own grain, so there is - # no device-side rounding to guess at from out here. + # The 200 MiB of host history less the sparse Batch's one-page staging copy, handed + # over whole: the manager counts host bytes in their own grain. assert config.host_sparse_pool_quota == (200 << 20) - 4096 # The automatic host tier still takes half of what the 200 MiB leaves of the 250 MiB. assert host_tier.quota == (50 << 20) // 2 @@ -4984,10 +4962,8 @@ def test_warmup_scan_leaves_host_backed_history_alone(offload: bool) -> None: manager.layer_specs = None manager.get_buffers = lambda layer_idx, kv_layout="NHD": torch.ones(2) # Layer 0's KEY pages are offloaded history; layer 1's are device memory. - manager.kv_cache_manager_py_config = KVCacheManagerConfig( - tokens_per_block=TOKENS_PER_BLOCK, - cache_tiers=[GpuCacheTierConfig(quota=1 << 20)], - layers=_offload_layers(), + manager.impl = SimpleNamespace( + is_sparse=lambda layer_id, role: layer_id == LayerId(0) and role == Role.KEY ) if offload: manager._sparse_kv_offload = SparseKvCacheOffloadConfig() diff --git a/tests/unittest/_torch/executor/kv_cache/test_sparse_offload_memory.py b/tests/unittest/_torch/executor/kv_cache/test_sparse_offload_memory.py index 5ab7f2574d50..5ec1aba4f03d 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_sparse_offload_memory.py +++ b/tests/unittest/_torch/executor/kv_cache/test_sparse_offload_memory.py @@ -29,10 +29,9 @@ sparse_cache_cells, sparse_cache_config, sparse_history_bytes_per_block, - sparse_offload_pool_ratio, + sparse_history_floor_bytes, sparse_units_per_page, ) -from tensorrt_llm.math_utils import ceil_div from tensorrt_llm.runtime.kv_cache_manager_v2 import ( AttentionLayerConfig, BatchDesc, @@ -59,7 +58,8 @@ cell_bytes=(32 * 288, 32 * 36) * 21, ) # 256 rows of up to 1M tokens: a table of 8,196 pages for each of three life -# cycles, and one fetch of 512 rows. +# cycles, a converted index table for each of the history's two sparse buffers, +# and one fetch of 512 rows. V4_FLASH_BATCH = SparseBatchShape( num_life_cycles=3, max_batch_size=256, @@ -67,6 +67,7 @@ selection_widths=(512,), layer_counts=(2,), flush_capacities=(256,), + index_table_entries=2, ) # The Batch and DeepSeek-V4's attention metadata offload state, of which the eager # and the CUDA graph metadata each hold one: 64 MiB of dequantized NVFP4 rows and @@ -340,38 +341,6 @@ def test_life_cycle_with_only_sparse_buffers_is_rejected() -> None: life_cycle_pages([_layer(0, ("history", 4096, True))], TOKENS_PER_BLOCK) -@pytest.mark.cpu_only -@pytest.mark.parametrize( - ("window", "sink_tokens", "capacity", "headroom", "pages"), - [ - (None, 0, 1 << 17, 1, 1024), - # A 128-token window and the token being written can straddle two pages. - (128, 0, 1 << 17, 1, 2), - (128, 0, 1 << 17, 4, 3), - (128, 200, 1 << 17, 1, 4), - # A request never holds more pages than its capacity spans. - (4096, 0, 300, 1, 3), - ], -) -def test_pages_a_decoding_request_holds( - window: int | None, sink_tokens: int, capacity: int, headroom: int, pages: int -) -> None: - life_cycle = LifeCyclePage( - window_size=window, - num_sink_tokens=sink_tokens, - is_state=False, - device_bytes=1, - host_bytes=0, - device_pool_bytes=(1,), - ) - assert ( - life_cycle.pages_per_request( - capacity=capacity, capacity_headroom=headroom, tokens_per_block=TOKENS_PER_BLOCK - ) - == pages - ) - - @pytest.mark.cpu_only def test_batch_matches_native_layout() -> None: # Three int32 page tables of 8,196 pages, two watermarks per row, and one arena slot @@ -390,9 +359,12 @@ def test_batch_matches_native_layout() -> None: 256 * (512 + 2) + (3 * 256 * 512 + 1) + 2 * (3 * 256 + 1) + 256 * 2 ) assert V4_FLASH_BATCH.fetch_scratch_bytes == 2_107_404 + # The fetch and the write-back read converted indices: a table per sparse buffer, + # shared by the life cycle's layers since nothing substitutes a scratch slot there. + assert V4_FLASH_BATCH.index_table_bytes == 4 * 2 * 256 * 8196 == 16_785_408 # The device holds the staged block a publish fills ahead of the step, the mirror of - # its tail that a captured kernel reads, and the work arrays. - assert V4_FLASH_BATCH.device_bytes == 25_181_184 + 3_072 + 2_107_404 == 27_291_660 + # its tail that a captured kernel reads, the converted tables and the work arrays. + assert V4_FLASH_BATCH.device_bytes == 25_181_184 + 3_072 + 16_785_408 + 2_107_404 == 44_077_068 @pytest.mark.cpu_only @@ -523,59 +495,6 @@ def test_sparse_cache_config_rejects_a_unit_that_does_not_divide_a_page() -> Non ) -@pytest.mark.cpu_only -def test_pool_ratio_buys_the_same_requests_in_every_life_cycle() -> None: - life_cycles = life_cycle_pages(_v4_flash_like_layers(), TOKENS_PER_BLOCK) - - ratio = sparse_offload_pool_ratio( - life_cycles, capacity=1 << 17, capacity_headroom=1, tokens_per_block=TOKENS_PER_BLOCK - ) - - # Every pool counts, host-backed ones included: the ratio is compared against a layout - # measured the same way, and is projected onto cold levels where one pool holds dense and - # sparse buffers together. Host pages costing the GPU quota nothing is a fact about the - # quota, not about the ratio. - request_bytes = [ - (life_cycle.device_bytes + life_cycle.host_bytes) - * life_cycle.pages_per_request( - capacity=1 << 17, capacity_headroom=1, tokens_per_block=TOKENS_PER_BLOCK - ) - for life_cycle in life_cycles - ] - assert ratio == pytest.approx([b / sum(request_bytes) for b in request_bytes]) - # The history life cycle is the one with host-backed pages, so counting them moves it. - assert ratio[1] > 0.9, "the offloaded history dominates once its host pages are counted" - assert sum(ratio) == pytest.approx(1.0, abs=1e-9) - - -def _history_beside_five_window_pools() -> list[LifeCyclePage]: - """The history's one dense pool wastes little of its granules on whole slots, - while the windowed life cycles have five pools that can each leave one.""" - layers = [ - _layer( - 0, ("key", 131_072, False), ("value", 65_536, False), ("x", 16_384, False), window=128 - ), - _sparse_layer(1), - _layer(2, ("state", 65_536, False), ("score", 800, False), window=8), - ] - return life_cycle_pages(layers, TOKENS_PER_BLOCK) - - -def _native_slots(life_cycles: list[LifeCyclePage], ratio: list[float], quota: int) -> list[int]: - """Slots the native manager gives each life cycle, one pool group apiece. - - It rounds the quota up to a whole granule, then splits it by ``ratio``. - """ - granule = _introspection.tier_allocation_granularity(CacheTier.GPU_MEM, quota) - return _introspection.ratio_to_slot_count_list( - ceil_div(quota, granule) * granule, - [list(life_cycle.device_pool_bytes) for life_cycle in life_cycles], - ratio, - granule, - [1] * len(life_cycles), - ) - - @pytest.mark.cpu_only @pytest.mark.parametrize( ("host_history_size", "pinnable_host_bytes", "budget"), @@ -605,6 +524,48 @@ def test_host_budget( ) +@pytest.mark.cpu_only +def test_automatic_host_budget_grows_to_the_floor() -> None: + # 256 requests of 1,000 bytes would take 256,000 bytes, short of the 300,000 the + # manager's constraint batches need. + assert ( + plan_host_history_bytes( + host_history_size=None, + max_batch_size=256, + host_bytes_per_request=1000, + pinnable_host_bytes=float("inf"), + staging_bytes=4096, + floor_bytes=300_000, + ) + == 300_000 + ) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + ("host_history_size", "pinnable_host_bytes", "match"), + [ + (200_000, float("inf"), "host_history_size .* less than the 0.3MiB"), + (None, 200_000.0, "automatically leaves 0.2MiB .* less than the 0.3MiB"), + ], + ids=["explicit", "automatic"], +) +def test_host_budget_below_the_floor_is_refused( + host_history_size: int | None, pinnable_host_bytes: float, match: str +) -> None: + # The native manager refuses a budget below the floor at construction; refusing it + # here says what to change. + with pytest.raises(ValueError, match=match): + plan_host_history_bytes( + host_history_size=host_history_size, + max_batch_size=256, + host_bytes_per_request=1000, + pinnable_host_bytes=pinnable_host_bytes, + staging_bytes=4096, + floor_bytes=300_000, + ) + + @pytest.mark.cpu_only @pytest.mark.parametrize( ("host_history_size", "pinnable_host_bytes", "match"), @@ -626,6 +587,44 @@ def test_host_budget_must_leave_the_pools_room( ) +@pytest.mark.cpu_only +def test_host_floor_holds_the_largest_constraint_batch() -> None: + layers = _v4_flash_like_layers() + shorter = BatchDesc([KVCacheDesc(capacity=1 << 16, history_length=0)]) + # DeepSeek-V4's constraint: one max_seq_len request beside short decode requests. + longest = BatchDesc( + [KVCacheDesc(capacity=1 << 17, history_length=(1 << 17) - 1)] + + [KVCacheDesc(capacity=1, history_length=0)] * 15 + ) + + floor = sparse_history_floor_bytes( + layers, [shorter, longest], tokens_per_block=TOKENS_PER_BLOCK, max_util_for_resume=0.95 + ) + + # The longest request spans 1,024 pages and each short one a page: 1,039 slots, or + # 1,094 under the 0.95 utilization the manager keeps. Each slot holds a page of both + # sparse layers' data and scales, and each of those two pools rounds up to a 2 MiB + # grain. + assert sparse_history_bytes_per_block(layers) == 2 * (32 * 288 + 32 * 36) + assert floor == 1094 * 20_736 + 2 * (2 << 20) == 26_879_488 + # A manager without constraints, or without host-backed pools, keeps no floor. + assert ( + sparse_history_floor_bytes( + layers, [], tokens_per_block=TOKENS_PER_BLOCK, max_util_for_resume=0.95 + ) + == 0 + ) + assert ( + sparse_history_floor_bytes( + [_layer(0, ("key", 4096, False))], + [longest], + tokens_per_block=TOKENS_PER_BLOCK, + max_util_for_resume=0.95, + ) + == 0 + ) + + @pytest.mark.cpu_only def test_host_budget_is_empty_on_a_rank_without_offloaded_history() -> None: assert ( @@ -681,7 +680,7 @@ def test_capacity_names_the_binding_limit_and_the_batch_that_fits() -> None: assert "4 by host" in report assert "limited by host" in report assert "256 leases" in report - assert "offload buffers 157.0MiB for 256 rows" in report + assert "offload buffers 173.0MiB for 256 rows" in report assert "A max_batch_size of 240 fits" in report diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py index 495a64275b33..ff81d5ebd395 100644 --- a/tests/unittest/_torch/executor/test_py_executor.py +++ b/tests/unittest/_torch/executor/test_py_executor.py @@ -887,27 +887,11 @@ def update_forward_resources( events.append("history") kv_cache_manager.update_forward_resources.side_effect = update_forward_resources - executor = types.SimpleNamespace( - model_engine=types.SimpleNamespace(route_capture=None, forward=forward), - resource_manager=object(), - kv_cache_manager=kv_cache_manager, - _is_kv_manager_v2=is_v2, - _loop_on_execution_stream=False, - enable_attention_dp=False, - enable_spec_decode=False, - runtime_draft_len=0, - iter_counter=0, - sampler=Mock(), - execution_stream=Mock(), - _iter_adp_dummy_ctx_tokens=0, - _iter_adp_dummy_gen_tokens=0, - _compute_adp_dummy_tokens=PyExecutor._compute_adp_dummy_tokens, - _maybe_record_hang_diagnostic_phase=Mock(), - _attach_encoder_output_to_execution_stream=Mock(), - _mark_cross_kv_projection_consumed=Mock(), - _kv_connector_wait_for_save=Mock(), - _wait_for_model_engine_input_copy=Mock(side_effect=lambda: events.append("input_copy")), - _handle_errors=Mock(), + executor = _forward_step_executor(types.SimpleNamespace(route_capture=None, forward=forward)) + executor.kv_cache_manager = kv_cache_manager + executor._is_kv_manager_v2 = is_v2 + executor._wait_for_model_engine_input_copy = Mock( + side_effect=lambda: events.append("input_copy") ) with ( patch("torch.cuda.current_stream", return_value=caller_stream), diff --git a/tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py b/tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py index b093f85b1396..5241b8230347 100644 --- a/tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py +++ b/tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py @@ -746,6 +746,13 @@ def __init__(self, **kwargs): sm_version=100, kv_cache_quant_algo=QuantAlgo.NO_QUANT, kv_cache_creator_cls=RecordingCacheCreator, + # A decode-only worker is a disaggregated generation server; creation refuses + # one without a cache transceiver. + cache_transceiver_config=( + CacheTransceiverConfig(backend="NIXL", transceiver_runtime="PYTHON") + if decode_only + else None + ), ) assert skip_est_args == [decode_only] diff --git a/tests/unittest/_torch/executor/test_resource_manager.py b/tests/unittest/_torch/executor/test_resource_manager.py index 2d480f8bf302..c3aa24af3120 100644 --- a/tests/unittest/_torch/executor/test_resource_manager.py +++ b/tests/unittest/_torch/executor/test_resource_manager.py @@ -1655,7 +1655,7 @@ def test_validate_request_budget_rejects_oversized_request(self): req = self._make_request(prompt_len=200, max_new_tokens=5, beam_width=4) with self.assertRaisesRegex( ValueError, "KV_CACHE_MANAGER requires 29 KV cache blocks.*" - "GPU-primary capacity of 8"): + "capacity of 8"): executor._validate_request_budget(req) # shared(2) + per_beam(2) * beam(2) = 6, so it fits. diff --git a/tests/unittest/disaggregated/test_cache_transceiver_single_process.py b/tests/unittest/disaggregated/test_cache_transceiver_single_process.py index b2c33ae83ce3..91ffa5f59e20 100644 --- a/tests/unittest/disaggregated/test_cache_transceiver_single_process.py +++ b/tests/unittest/disaggregated/test_cache_transceiver_single_process.py @@ -175,6 +175,7 @@ class KvCacheConfigV2: # tensorrt_llm/llmapi/llm_args.py; KVCacheManagerV2._build_base_config() # reads it unconditionally. kv_pool_rebalance_config: KvPoolRebalanceConfig = field(default_factory=KvPoolRebalanceConfig) + enable_kv_pool_rebalance: bool = False # --------------------------------------------------------------------------- diff --git a/tests/unittest/disaggregated/test_extractor.py b/tests/unittest/disaggregated/test_extractor.py index ec6ef42ad0f7..ff0344c58b0e 100644 --- a/tests/unittest/disaggregated/test_extractor.py +++ b/tests/unittest/disaggregated/test_extractor.py @@ -48,6 +48,7 @@ Mapping, ) from tensorrt_llm.llmapi.llm_args import KvCacheConfig as KvCacheConfigV2 +from tensorrt_llm.runtime.kv_cache_manager_v2 import CacheTier class DummyRankInfo: @@ -291,7 +292,9 @@ def _make_fake_v2_manager(attrs, role_mapper_kinds, *, num_pools=1, slot_bytes_l pool_group_index=0, num_slots=4, pools=[ - SimpleNamespace(pool_index=pi, base_address=base, slot_bytes=sb) + SimpleNamespace( + pool_index=pi, base_address=base, slot_bytes=sb, storage_medium=CacheTier.GPU_MEM + ) for pi, sb in enumerate(slot_bytes_list) ], slot_desc=SimpleNamespace( @@ -550,8 +553,12 @@ def buf(layer_id, role): pool_group_index=0, num_slots=4, pools=[ - SimpleNamespace(pool_index=0, base_address=0x10000, slot_bytes=256), - SimpleNamespace(pool_index=1, base_address=0x20000, slot_bytes=32), + SimpleNamespace( + pool_index=0, base_address=0x10000, slot_bytes=256, storage_medium=CacheTier.GPU_MEM + ), + SimpleNamespace( + pool_index=1, base_address=0x20000, slot_bytes=32, storage_medium=CacheTier.GPU_MEM + ), ], slot_desc=SimpleNamespace( variants=[ @@ -579,7 +586,11 @@ def buf(layer_id, role): attention_pg = SimpleNamespace( pool_group_index=1, num_slots=8, - pools=[SimpleNamespace(pool_index=0, base_address=0x30000, slot_bytes=128)], + pools=[ + SimpleNamespace( + pool_index=0, base_address=0x30000, slot_bytes=128, storage_medium=CacheTier.GPU_MEM + ) + ], slot_desc=SimpleNamespace( variants=[ SimpleNamespace( diff --git a/tests/unittest/disaggregated/test_kv_transfer.py b/tests/unittest/disaggregated/test_kv_transfer.py index 5f2745d50458..f8e1a053e6c9 100644 --- a/tests/unittest/disaggregated/test_kv_transfer.py +++ b/tests/unittest/disaggregated/test_kv_transfer.py @@ -113,6 +113,7 @@ class KvCacheConfigV2: # tensorrt_llm/llmapi/llm_args.py; KVCacheManagerV2._build_base_config() # reads it unconditionally. kv_pool_rebalance_config: KvPoolRebalanceConfig = field(default_factory=KvPoolRebalanceConfig) + enable_kv_pool_rebalance: bool = False _THIS_DIR = Path(__file__).resolve().parent diff --git a/tests/unittest/disaggregated/test_sparse_kv_transfer.py b/tests/unittest/disaggregated/test_sparse_kv_transfer.py index 6c0974385ce5..4036a3b86737 100644 --- a/tests/unittest/disaggregated/test_sparse_kv_transfer.py +++ b/tests/unittest/disaggregated/test_sparse_kv_transfer.py @@ -42,6 +42,7 @@ from tensorrt_llm._torch.disaggregation.resource.utils import get_unique_pool_memory_descs from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2, Role from tensorrt_llm.disaggregated_params import DisaggregatedParams +from tensorrt_llm.runtime.kv_cache_manager_v2 import CacheTier pytestmark = pytest.mark.cpu_only @@ -75,8 +76,17 @@ def _manager(offload: bool) -> Mock: pool_group_index=0, num_slots=4, pools=[ - SimpleNamespace(base_address=0x100000 * (i + 1), slot_bytes=2 * size) - for i, size in enumerate(sizes) + SimpleNamespace( + base_address=0x100000 * (i + 1), + slot_bytes=2 * size, + # Offloaded sparse buffers are host memory; the indexer stays on the GPU. + storage_medium=( + CacheTier.HOST_MEM + if offload and role != "indexer" + else CacheTier.GPU_MEM + ), + ) + for i, (role, size) in enumerate(zip(roles, sizes)) ], slot_desc=SimpleNamespace( variants=[SimpleNamespace(layer_group_id=0, coalesced_buffers=coalesced)] @@ -172,10 +182,11 @@ def test_v2_extracts_host_data_and_scales_beside_device_indexer() -> None: assert [p.memory_type for p in table.pool_groups[0].pools] == ["DRAM", "DRAM", "VRAM"] assert table.has_host_pools assert not build_page_table_from_manager(_manager(False)).has_host_pools + # Each pool's memory type is its native storage medium, not the layer config's. manager = _manager(True) manager.impl.init_config.layers[1].buffers[0].is_sparse = False - with pytest.raises(ValueError, match="mixes host and device"): - build_page_table_from_manager(manager) + table = build_page_table_from_manager(manager) + assert [p.memory_type for p in table.pool_groups[0].pools] == ["DRAM", "DRAM", "VRAM"] def test_registration_separates_memory_types_and_device_ids() -> None: diff --git a/tests/unittest/disaggregated/test_unsubmitted_session.py b/tests/unittest/disaggregated/test_unsubmitted_session.py index 53c372b00a88..5b9460418858 100644 --- a/tests/unittest/disaggregated/test_unsubmitted_session.py +++ b/tests/unittest/disaggregated/test_unsubmitted_session.py @@ -17,7 +17,7 @@ import queue import threading from types import SimpleNamespace -from unittest.mock import Mock +from unittest.mock import Mock, patch import numpy as np import pytest @@ -47,7 +47,7 @@ def _sender(*, ownership: bool = True) -> transfer_mod.Sender: sender._pre_cancelled_rids = {} sender._peer_requests_lock = threading.Lock() sender._peer_requests, sender._peer_requests_timestamps = {}, {} - sender._request_info_ttl_s = transfer_mod.Sender._STALE_REQ_INFO_TTL_S + sender._request_info_timeout_s = 0.0 sender._shutdown, sender._shutdown_requested = True, False sender._retirement_watchdog = None sender._instance_rank = 7 @@ -166,6 +166,33 @@ def test_queued_prefill_metadata_uses_transfer_timeout( sender.shutdown() +def test_orphaned_context_first_request_fails_after_the_base_ttl( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A context-first request's long transfer timeout does not keep its receiver waiting. + + Its REQUEST_DATA names a session that should already exist. + """ + monkeypatch.setattr(transfer_mod, "ZMQMessenger", Mock()) + monkeypatch.setattr(transfer_mod, "KV_TRANSFER_NUM_THREADS", 0) + now = 0.0 + monkeypatch.setattr(transfer_mod.time, "monotonic", lambda: now) + sender = transfer_mod.Sender(Mock(), Mock(), request_info_timeout_s=600.0) + info = transfer_mod.RecvReqInfo.from_bytes(_request_data()[1]) + info.context_first = True + try: + sender._add_req_info(401, 0, info) + with patch.object(sender, "_fail_orphaned_receiver") as fail: + now = 121.0 + sender.sweep_stale_req_infos() + + assert sender._get_req_info(401) is None + # The log names the TTL the entry actually waited. + fail.assert_called_once_with(info, transfer_mod.Sender._STALE_REQ_INFO_TTL_S) + finally: + sender.shutdown() + + @pytest.mark.parametrize("pending_settlement", [False, True]) @pytest.mark.parametrize("ack_fails_once", [False, True]) def test_worker_orders_session_ack_after_piece_reports( diff --git a/token_parity.md b/token_parity.md deleted file mode 100644 index 215d7b17c2c9..000000000000 --- a/token_parity.md +++ /dev/null @@ -1,295 +0,0 @@ - - -The DeepSeek V4 Flash token-parity fixes make **TopK, MHC, and FMHA honor -`FORCE_DETERMINISTIC=1`**. Repeated resident-cache runs were already producing -different tokens, so the resident/offload comparison needed a repeatable baseline. -The implementation is split into a disaggregated sparse-offload commit for -transport and configuration, followed by a general determinism commit for -TopK, MHC, FMHA, kernel-specific controls, tests, and this report. The validation -below was performed before splitting the commits. - -All five cumulative feature stages pass exact resident-versus-CPU-offload token -parity. Including the fixed baseline and profiling probes, the successful runs -cover **98 requests and 15,680 generated tokens**, with **36 exact sequential -pairs and 8 exact concurrent pairs**. These results compare matching feature -configurations; they do not establish identical tokens across every configuration. -Block reuse is exercised on the context worker only. -[Final validation report](../sparse-disagg-e2e-artifacts/sequential-20261017/final-report.json). - -Tracing repeated resident runs found identical TopK scores with different index -orders. A 50-repeat kernel check reproduced native ordering differences; tied -scores also changed the selected set. The PyTorch implementation had no ordering -or selection mismatches in those checks. Separately, MHC split-K reductions -changed intermediate mixing values by approximately `1e-7` to `3e-7`, enough to -propagate into different generated tokens. - -- **TopK:** Native atomic appends changed index ordering and, with tied scores, - sometimes changed which indices were selected. Deterministic mode forces the - PyTorch implementation for both prefill and decode, which produced repeatable - selections and ordering in the tested configuration. - [Implementation](tensorrt_llm/_torch/modules/top_k.py). -- **MHC:** Atomic split-K reductions introduced small floating-point differences - that propagated into token divergence. Deterministic mode forces a fixed, - single-split tactic, gives deterministic execution a separate tuning-cache - identity, and overrides cached or explicitly supplied tactics. It uses the - half-fused MMA implementation where supported, with a single-split FMA fallback. - [Implementation](tensorrt_llm/_torch/modules/mhc/mhc_cuda.py). -- **FMHA:** Deterministic mode fixes the KV reduction to one partition so eager - execution and CUDA graph capture select the same reduction layout. - [Implementation](cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h). - -Sparse-offload settings are also preserved during checkpoint loading so the -offload run actually uses the CPU hot tier. -[Configuration fix](tensorrt_llm/_torch/model_config.py). - -Validation passed: **43 kernel regression tests** and **6/6 exact resident/offload -token matches**, covering 12 requests and 1,920 generated tokens on the full -43-layer DeepSeek V4 Flash model. The comparison used four B200 GPUs, NVFP4 KV -cache, context and generation TP2/EP2 with attention TP, the NIXL Python -transceiver, `FORCE_DETERMINISTIC=1`, and disabled autotuning. Prompt lengths were -129, 8,193, 127, 128, 2,049, and 8,193 tokens, with 160 output tokens per request. - -That initial comparison used eager execution with CUDA graphs, overlap scheduling, -chunked prefill, and block reuse disabled. Attention DP had encountered a CUDA -illegal access during startup memory profiling; the subsequent validation below -covers its current status. -The passing run used the production implementation without diagnostic overrides. -Local evidence is available in the -[root-cause report](../sparse-disagg-e2e-artifacts/parity/root-cause-evidence.json) -and the [validation report](../sparse-disagg-e2e-artifacts/refresh-20261017T173731Z/final-report.json). - -Sequential validation used the original combined commit `432a1ba93d` plus the -fixes described below, before splitting the changes into separate commits. -The following features were enabled cumulatively. Each stage must match resident -versus offloaded tokens with the same enabled features before the next feature -is enabled. -Cross-configuration comparisons are also recorded; exact TP/DP token identity is -a separate criterion because their low-precision reduction layouts differ. -The same six prompts, 160 output tokens, full checkpoint, NVFP4 format, and -deterministic kernel choices are retained. Runtime profiles and worker statistics -verify that the enabled features execute. - -| Stage | Newly enabled feature | Status | -| --- | --- | --- | -| 1 | CUDA graphs, decode batch sizes 1 and 2 | Passed: 6/6 resident/offload and eager matches | -| 2 | Generation attention DP | Passed: 6/6 sequential and 2/2 concurrent resident/offload matches | -| 3 | Chunked prefill, 2,048-token scheduling budget | Passed: 6/6 sequential and 2/2 concurrent matches; schedules identical | -| 4 | Context-worker block reuse, including the repeated 8,193-token prompt | Passed: 6/6 sequential and 2/2 concurrent matches; reuse verified | -| 5 | Overlap scheduler | Passed: 6/6 sequential and 2/2 concurrent matches; all sequential outputs also match stage 4 | - -The initial CUDA-graph resident run matched both 8,193-token requests and the -2,049-token request, but differed for the 127-, 128-, and 129-token prompts. -An eager resident/offload control on the same allocation matched all six prior -reference outputs and all six resident/offload pairs. A focused deterministic -TopK eager/graph comparison passed 22 cases. - -Full-model tracing located the first difference at generation layer 2's attention -output. Earlier layer outputs, queries, and TopK selections matched. Profiles -showed that graph capture selected a CGA multi-CTA KV reduction despite -`FORCE_DETERMINISTIC=1`; eager short requests selected a single-CTA reduction. -Running eager execution with multiple CUDA streams still matched the reference. -The FMHA selector now forces `MultiCtasKvMode::Disabled` and -`mMaxNumCtasPerSeqKv = 1` in deterministic mode. This prevents its heuristics from -overriding the single-partition choice when graph capture supplies a longer KV -length than the live request. -[Implementation](cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h). - -Both new native regressions failed before the fix and pass after it. They check -kernel selection with 32 and 64 query heads and exact eager/graph output equality -at seven live KV lengths from 127 to 8,193 tokens, with a graph captured at 9,216. -[Native tests](cpp/tests/unit_tests/kernels/fmhaDeterministicTest.cpp). -The rebuilt libraries pass a fresh eager resident/offload comparison: **6/6 exact -token matches**, with repeatable outputs for the repeated long prompt in both -modes. **Stage 1 passes:** all six CUDA-graph outputs match the corresponding eager -outputs in both resident and offload modes, and all six resident/offload pairs -match. Each mode's profiled repeat also matches its earlier long-prompt output. -Profiles show three graph launches on each generation rank, using only the -single-partition FMHA kernel in the captured decode steps. -[Diagnostic evidence](../sparse-disagg-e2e-artifacts/sequential-20261017/cuda-graph-root-cause.json). - -The fix also changes eager reduction choices, so five of the six eager outputs -differ from the old build. The 129-token prompt remains identical; the other -short prompts diverge late in generation. Subsequent feature comparisons use this -fixed implementation in both runs. -[Reference changes](../sparse-disagg-e2e-artifacts/sequential-20261017/fmha-reference-change.json). - -Stage 2 completes attention-DP startup without the earlier CUDA illegal access, -but all six resident outputs differ from stage 1. The first differing output-token -indices are 28, 18, 24, 24, 32, and 18, respectively. Repeated outputs are stable. -Tracing found exact equality in the first layer's input, quantized queries, and -native attention output. The first difference is its output projection: TP2 -rounds each partial GEMM result to BF16 before reducing the two results, while DP -computes the full projection locally. Replaying captured tensors reproduces both -paths exactly: 1,476 of 4,096 output elements differ, with a maximum absolute -difference of 0.03125. Graph and eager DP runs produce identical tokens. -[Projection evidence](../sparse-disagg-e2e-artifacts/sequential-20261017/attention-dp-projection.json). - -The current E2E criterion compares resident and offloaded execution within each -feature configuration. The TP/DP numerical difference is retained in the report. -Sequential requests ran on DP rank 0. Two concurrent long requests exercised both -ranks; one matched the sequential output and one differed starting at output -token 27. The sequential outputs and concurrent output pair both repeated exactly -across two independent resident runs. The harness records this sensitivity -separately and requires the pair -of concurrent outputs to match between resident and offload modes, allowing the -two identical requests to exchange rank assignments. **Stage 2 passes:** all six -sequential outputs and both concurrent outputs match exactly between resident and -offload modes. Worker statistics confirm active requests on DP ranks 0 and 1 in -both modes, and profiles confirm graph replay on both ranks. -[Attention-DP validation](../sparse-disagg-e2e-artifacts/sequential-20261017/attention-dp-root-cause.json). - -The stage 3 resident run verifies chunked execution: 29 context iterations, at -most 2,048 scheduled context tokens per iteration, and 21 iterations with cached -context KV. Five of the six sequential outputs match stage 2; the 2,049-token -prompt differs at exactly one output token, index 28. All six sequential -resident/offload pairs match exactly. The initial concurrent pair differed, and -worker statistics showed different prefill chunk schedules: the second resident -request started with 1,920 tokens, while the corresponding offload request -started with 2,048. The harness now submits the second request after the first -streaming token. This preserves concurrent decode on both DP ranks while keeping -prefill scheduling serial and comparable. **Stage 3 passes on rerun:** all six -sequential pairs and both concurrent outputs match exactly, with identical -resident/offload chunk schedules: 30 context iterations, at most 2,048 tokens per -iteration, and 21 iterations with cached context KV. Both DP ranks execute real -requests, and graph replay remains verified. -[Scheduling evidence](../sparse-disagg-e2e-artifacts/sequential-20261017/chunked-prefill-probe-schedule.json). - -Stage 4's first attempt did not exercise reuse: `FORCE_DETERMINISTIC=1` -overrides the requested cache configuration and disables it. The test correctly -failed its runtime feature check despite stable tokens. The rerun uses the -existing attention, MoE, and all-reduce determinism switches plus new -`FORCE_TOPK_KERNEL_DETERMINISTIC=1` and `FORCE_MHC_KERNEL_DETERMINISTIC=1` -switches, with global `FORCE_DETERMINISTIC=0`. This preserves the deterministic -kernel choices while allowing prefix reuse. The global flag's behavior is -unchanged. Stages 4 and 5 use: - -```bash -export FORCE_DETERMINISTIC=0 -export FORCE_ATTENTION_KERNEL_DETERMINISTIC=1 -export FORCE_MOE_KERNEL_DETERMINISTIC=1 -export FORCE_ALL_REDUCE_DETERMINISTIC=1 -export FORCE_TOPK_KERNEL_DETERMINISTIC=1 -export FORCE_MHC_KERNEL_DETERMINISTIC=1 -``` - -Reuse is enabled on the context worker in both modes; generation -receives prompt KV through disaggregation. Sparse generation's explicit -restriction on local prefix reuse remains in place, so this test does not -establish support for generation-worker reuse of CPU-backed prefixes. -The 24 focused TopK/MHC regressions pass with the global and kernel-specific -controls, including tied-score selection and cached split-K tactic overrides. -**Stage 4 passes:** all six sequential pairs and both concurrent outputs match -exactly between resident and offload. Both runs verify actual reuse: 276 reused -blocks and 12 context iterations, down from 30 without reuse, with identical -chunk schedules. Repeated long prompts reuse 8,192 cached tokens and compute one -prompt token; their outputs and the profiled repeat remain identical. Five -sequential outputs match stage 3; the 127-token prompt differs beginning at output -token 24. This cross-configuration difference remains recorded separately. -[Reuse-control evidence](../sparse-disagg-e2e-artifacts/sequential-20261017/block-reuse-controls.json). - -**Stage 5 passes:** all six sequential pairs and both concurrent outputs match -exactly between resident and offload with overlap scheduling enabled. All six -sequential outputs also match stage 4 in each mode; the repeated prompt and -profiled repeat remain identical. Both modes report the overlap scheduler, -active requests on both attention-DP ranks, 276 reused blocks, and identical -context schedules. CUDA traces show three graph launches on each generation rank. -No further production change was needed for overlap in this configuration. -[Overlap validation](../sparse-disagg-e2e-artifacts/sequential-20261017/overlap-scheduler-validation.json). - -Run artifacts are stored locally in -`../sparse-disagg-e2e-artifacts/sequential-20261017/` relative to this worktree. - -**32K sparse-offload capacity validation, October 17–18, 2026 (UTC).** -The full 43-layer DeepSeek V4 Flash model sustained **1,134 simultaneous decode -requests, 567 per attention-DP rank**, with exactly **32,768 input tokens and -160 output tokens per request** on four B200 GPUs. This is the measured capacity -with 32 GiB of CPU sparse-history memory per generation rank, -`free_gpu_memory_fraction=0.9`, and generation `max_batch_size=1024` per rank. -It is not an unrestricted hardware maximum or a throughput benchmark. - -The two offload trials at 1,134 both completed and produced identical output -multisets. Offering 1,136, 1,140, or 1,182 requests still peaked at 1,134 active -decodes; all offered requests eventually returned 160 tokens. The clean 1,136 -probe released its first batch after 1,134 receives. The exploratory 1,140 and -1,182 probes required recorded manual barrier adjustments. Startup's aggregate -estimate of 591 requests per rank (1,182 total) exceeded the observed usable -capacity; the larger gated batches encountered attention-DP dummy-allocation -pressure. -[Capacity evidence](../sparse-disagg-e2e-artifacts/long-context-20261017T231142Z/offload-capacity-summary.json). - -The prompts are **synthetic and identical**: a repeated astronomy paragraph and -a fixed question, with BOS included in the 32,768-token count. Context prefix -reuse is enabled; generation reuse is disabled, so each decode request receives -its own KV through disaggregation. Ratio-4 sparse layers have 8,192 compressed -history rows at this input length. Their GPU fetch cache holds 1,056 rows with -TopK 512, `gpu_cache_factor=2`, and 128-token pages. The complete history is -therefore larger than the GPU fetch cache and uses the CPU hot tier. Dense pools -and ratio-128 history remain on the GPU. - -The topology is context TP2/EP2 and generation TP2/EP2 with attention DP2. Both -runs use NVFP4 KV, NIXL Python transfer, CUDA graphs, overlap scheduling, and -2,048-token chunked prefill. Greedy sampling requests exactly 160 tokens and -ignores EOS. The kernel-specific determinism flags listed above remain enabled, -with global `FORCE_DETERMINISTIC=0` and autotuning disabled. - -Exact parity compares output-token **multisets**, allowing identical prompts to -exchange attention-DP rank assignments. A test-only receive-publication barrier -starts matching decode batches, and context `max_batch_size=1` fixes prefill -shapes. The barrier leaves transfers, cache operations, model kernels, and -sampling unchanged. Each passing trial must also sustain its full requested -decode concurrency for all 159 decode iterations; the first output token comes -from context execution. - -| Offered requests | Exact resident/offload matches | Active decode requests | -| --- | --- | --- | -| 1, cold prefix | 1/1 | 1 | -| 1, reused prefix | 1/1 | 1 | -| 512 | 512/512 | 512 | -| 1,024 | 1,024/1,024 | 1,024 | -| 1,134, first trial | 1,134/1,134 | 1,134 | -| 1,134, repeat | 1,134/1,134 | 1,134 | - -All six comparisons pass: **3,806 matched request pairs and 608,960 exactly -matched output tokens per mode**. Both maximum-size repeats sustain 567 decodes -per rank for all 159 decode iterations with zero paused requests. The repeated -output multisets also match within each mode. - -The GPU-resident reference needs different allocation settings to hold the same -batch: `free_gpu_memory_fraction=0.99`, `max_util_for_resume=1.0`, -`fp8_context_mla_kv_len_cap=33024`, and `avg_seq_len=32770`. These change memory -allocation, while model computation, input/output lengths, graph sizes, and -sampling remain matched. The last setting sizes windowed pools for decode -crossing a page boundary. The default typical sequence ends exactly on a page -boundary and models one window page where this workload can need two. Earlier -resident attempts paused requests despite spare full-history blocks, changing -the decode batch and failing the comparison. Those attempts are preserved. -[Pool-sizing evidence](../sparse-disagg-e2e-artifacts/long-context-20261017T231142Z/resident-pool-sizing-diagnostic.json). - -**Ordinary asynchronous high-concurrency parity remains unresolved.** Initial -runs without the barrier showed token differences as their batch schedules -changed. The controlled results establish resident/offload parity for matching -batches, not invariance to arbitrary request arrival or batching. Warm prefix -reuse also makes this a capacity test with repeated prompts rather than a test -of sustained throughput on distinct 32K documents. -[Asynchronous comparison](../sparse-disagg-e2e-artifacts/long-context-20261017T231142Z/initial-comparison.json). - -The larger test exposed one additional **general disaggregation fix**. In -generation-first serving, receive descriptors could expire after a fixed 120 -seconds while prefill remained queued, even with a 600-second transfer timeout. -`Sender` now retains them for `max(120 seconds, configured transfer timeout)`. -The regression checks the default lifetime, a shorter configured timeout, and -retention plus eventual expiry for a longer timeout. Validation passed -**90 tests, with one skipped**, across the unsubmitted-session and bounded-polling -suites. This fix concerns transfer metadata lifetime and is not specific to -sparse KV offload. -[Implementation](tensorrt_llm/_torch/disaggregation/native/transfer.py), -[regression tests](tests/unittest/disaggregated/test_unsubmitted_session.py), -[validation evidence](../sparse-disagg-e2e-artifacts/long-context-20261017T231142Z/regression-validation.json). - -Complete results, configuration files, source hashes, failed attempts, and -iteration metrics are linked from the -[32K final report](../sparse-disagg-e2e-artifacts/long-context-20261017T231142Z/final-report.json) -and [exact comparison](../sparse-disagg-e2e-artifacts/long-context-20261017T231142Z/verified-comparison.json). From 99df0cd8ba13ca85a5247693ad86cf59d9af4abf Mon Sep 17 00:00:00 2001 From: Guiju Zhang <7135567+cascade812@users.noreply.github.com> Date: Thu, 29 Oct 2026 22:36:52 -0700 Subject: [PATCH 04/11] [TRTLLM-16282][chore] Remove the benchmark-only transfer TTL and determinism changes These changes served only benchmark and validation runs, not production. - The disaggregated sender kept generation-first REQUEST_DATA for the whole transfer timeout instead of 120 s. Only the 32K capacity run needed it: it used generation-first scheduling, a 600 s transfer timeout and a test-only barrier that held decodes until every prefill finished. With the default context-first schedule and 60 s timeout it changed nothing. The sender is back to main's fixed 120 s TTL. - FORCE_DETERMINISTIC, and the new FORCE_TOPK_KERNEL_DETERMINISTIC and FORCE_MHC_KERNEL_DETERMINISTIC switches, made TopK, mHC and the trtllm-gen FMHA KV reduction repeatable for exact resident-versus- offload token comparisons. top_k.py, mhc_cuda.py and fmhaKernels.h are back to main. Their tests are removed, along with fmhaDeterministicTest and the test_top_k.py entry in l0_b200.yml. Signed-off-by: Guiju Zhang <7135567+cascade812@users.noreply.github.com> --- .../trtllmGenKernels/fmha/fmhaKernels.h | 12 +- cpp/tests/unit_tests/kernels/CMakeLists.txt | 12 - .../kernels/fmhaDeterministicTest.cpp | 246 ------------------ .../_torch/disaggregation/native/transfer.py | 44 +--- tensorrt_llm/_torch/modules/mhc/mhc_cuda.py | 34 +-- tensorrt_llm/_torch/modules/top_k.py | 12 - .../test_lists/test-db/l0_b200.yml | 1 - tests/unittest/_torch/modules/test_mhc.py | 53 +--- tests/unittest/_torch/modules/test_top_k.py | 41 --- .../disaggregated/test_unsubmitted_session.py | 56 +--- 10 files changed, 22 insertions(+), 489 deletions(-) delete mode 100644 cpp/tests/unit_tests/kernels/fmhaDeterministicTest.cpp diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h index 4b303a0be53f..deb8884ad99d 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h @@ -626,7 +626,7 @@ class TllmGenFmhaKernel TLLM_CHECK_WITH_INFO(findIter != functions.end(), "Trtllm-gen kernels not found: " + info); auto const& kernelMeta = mKernelMeta[findIter->second.mMetaInfoIndex]; - CUfunction const func = findIter->second.mDeviceFunction; + const CUfunction func = findIter->second.mDeviceFunction; // mGroupsHeadsQ, mGroupsTokensHeadsQ, and mNumInsts{Q,Kv} are part of the hashID, so the // matched cubin already agrees with the AutoTuner decision. Re-assign from kernelMeta @@ -1215,12 +1215,6 @@ class TllmGenFmhaKernel // Multi-CTA KV mode (bool -> enum conversion) options.mMultiCtasKvMode = params.mMultiCtasKvMode ? tensorrt_llm::kernels::MultiCtasKvMode::GmemReduction : tensorrt_llm::kernels::MultiCtasKvMode::Disabled; - if (tensorrt_llm::common::getEnvForceDeterministicAttention()) - { - // Keep one KV reduction partition even when the selector considers CGA reduction. - options.mMultiCtasKvMode = tensorrt_llm::kernels::MultiCtasKvMode::Disabled; - options.mMaxNumCtasPerSeqKv = 1; - } options.mStoresSoftmaxStats = true; // Attention features @@ -1583,7 +1577,7 @@ class TllmFmhaKernelFactory public: using KernelType = TllmGenFmhaKernel; - KernelType* getKernels(typename KernelType::KernelMeta const* pKernelList, unsigned int nbKernels, Data_type dtypeQ, + KernelType* getKernels(const typename KernelType::KernelMeta* pKernelList, unsigned int nbKernels, Data_type dtypeQ, Data_type dtypeK, Data_type dtypeV, Data_type dtypeOut, unsigned int sm, int numEltsPerSageAttnBlkQ = 0, int numEltsPerSageAttnBlkK = 0, int numEltsPerSageAttnBlkP = 0, int numEltsPerSageAttnBlkV = 0, bool fusesDsv4InvRopeFp8Quant = false) @@ -1656,7 +1650,7 @@ class TllmFmhaKernelFactory | (static_cast(fusesDsv4InvRopeFp8Quant) << 44); } - std::unordered_map const> mKernels; + std::unordered_map> mKernels; }; inline TllmGenFmhaKernel* getTllmFmhaKernels(Data_type dtypeQ, Data_type dtypeK, Data_type dtypeV, Data_type dtypeOut, diff --git a/cpp/tests/unit_tests/kernels/CMakeLists.txt b/cpp/tests/unit_tests/kernels/CMakeLists.txt index 21f67b7685d0..f75cb1de4299 100644 --- a/cpp/tests/unit_tests/kernels/CMakeLists.txt +++ b/cpp/tests/unit_tests/kernels/CMakeLists.txt @@ -228,18 +228,6 @@ add_tllm_gtest( LABELS "unit;kernels;attention-mask;gpu") add_gtest(fmhaDsv4EpilogueFusionTest fmhaDsv4EpilogueFusionTest.cpp) -# The trtllm-gen FMHA runner has no component target, so this links the full -# stack. -add_tllm_gtest( - fmhaDeterministicTest - OWNER - kernel - SOURCES - fmhaDeterministicTest.cpp - FULL_STACK - NO_GTEST_MAIN - LABELS - "unit;kernels;fmha;full-stack;gpu") add_tllm_gtest( nvfp4ColdPageKernelsTest diff --git a/cpp/tests/unit_tests/kernels/fmhaDeterministicTest.cpp b/cpp/tests/unit_tests/kernels/fmhaDeterministicTest.cpp deleted file mode 100644 index 3275b612a086..000000000000 --- a/cpp/tests/unit_tests/kernels/fmhaDeterministicTest.cpp +++ /dev/null @@ -1,246 +0,0 @@ -/* - * 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 - -#include -#include -#include -#include -#include -#include -#include - -#include "tensorrt_llm/common/cudaUtils.h" -#include "tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaRunner.h" - -namespace -{ - -using tensorrt_llm::kernels::DATA_TYPE_BF16; -using tensorrt_llm::kernels::DATA_TYPE_E4M3; -using tensorrt_llm::kernels::FmhaKernelType; -using tensorrt_llm::kernels::QkvLayout; -using tensorrt_llm::kernels::SparseType; -using tensorrt_llm::kernels::TileScheduler; -using tensorrt_llm::kernels::TllmGenFmhaRunner; -using tensorrt_llm::kernels::TllmGenFmhaRunnerParams; -using tensorrt_llm::kernels::TrtllmGenAttentionMaskType; - -class FmhaDeterministicTest : public ::testing::Test -{ -protected: - static constexpr int32_t kHeadDim = 512; - static constexpr int32_t kNumHeads = 32; - static constexpr int32_t kSparseTopK = 640; - static constexpr int32_t kCaptureKvLength = 9216; - - void SetUp() override - { - if (!tensorrt_llm::common::isSM100Family(tensorrt_llm::common::getSMVersion())) - { - GTEST_SKIP() << "TRTLLM-Gen sparse MLA requires Blackwell or newer."; - } - TLLM_CUDA_CHECK(cudaStreamCreate(&mStream)); - } - - void TearDown() override - { - if (mStream != nullptr) - { - cudaStreamSynchronize(mStream); - } - if (mGraphExec != nullptr) - { - cudaGraphExecDestroy(mGraphExec); - } - if (mGraph != nullptr) - { - cudaGraphDestroy(mGraph); - } - for (void* ptr : mAllocations) - { - cudaFree(ptr); - } - if (mStream != nullptr) - { - cudaStreamDestroy(mStream); - } - } - - template - T* allocate(size_t count) - { - T* ptr = nullptr; - TLLM_CUDA_CHECK(cudaMalloc(&ptr, count * sizeof(T))); - mAllocations.push_back(ptr); - return ptr; - } - - template - void copyToDevice(T* dst, std::vector const& src) - { - TLLM_CUDA_CHECK(cudaMemcpy(dst, src.data(), src.size() * sizeof(T), cudaMemcpyHostToDevice)); - } - - std::vector readOutput(uint16_t const* output) - { - std::vector result(kNumHeads * kHeadDim); - TLLM_CUDA_CHECK( - cudaMemcpyAsync(result.data(), output, result.size() * sizeof(uint16_t), cudaMemcpyDeviceToHost, mStream)); - TLLM_CUDA_CHECK(cudaStreamSynchronize(mStream)); - return result; - } - - TllmGenFmhaRunnerParams makeParams(int32_t numHeads = kNumHeads) - { - TllmGenFmhaRunnerParams params{}; - params.mQkvLayout = QkvLayout::PagedKv; - params.mMaskType = TrtllmGenAttentionMaskType::Dense; - params.mKernelType = FmhaKernelType::Generation; - params.mTileScheduler = TileScheduler::Persistent; - params.mMultiCtasKvMode = false; - params.mHeadDimQk = kHeadDim; - params.mHeadDimV = kHeadDim; - params.mNumHeadsQ = numHeads; - params.mNumHeadsKv = 1; - params.mNumHeadsQPerKv = numHeads; - params.mBatchSize = 1; - params.mMaxSeqLenCacheKv = kCaptureKvLength; - params.mMaxSeqLenQ = 1; - params.mMaxSeqLenKv = kCaptureKvLength; - // Negative reaches are unbounded, which is what a full-attention decode wants. - params.mLeftSlidingWindow = -1; - params.mRightSlidingWindow = -1; - params.mChunkedAttentionSize = std::numeric_limits::max(); - params.mSumOfSeqLensQ = 1; - params.mSumOfSeqLensKv = kCaptureKvLength; - params.mMaxNumPagesPerSeqKv = kSparseTopK; - params.mNumTokensPerPage = 1; - params.mNumPagesInMemPool = kSparseTopK; - params.mMultiProcessorCount = tensorrt_llm::common::getMultiProcessorCount(); - params.mScaleQ = 1.0F; - params.mSparseAttention = SparseType::DynamicTokenSparse; - params.mSparseTopK = kSparseTopK; - params.stream = mStream; - return params; - } - - cudaStream_t mStream{}; - cudaGraph_t mGraph{}; - cudaGraphExec_t mGraphExec{}; - std::vector mAllocations; -}; - -TEST_F(FmhaDeterministicTest, LongSparseKvKeepsOneReductionPartition) -{ - TllmGenFmhaRunner runner(DATA_TYPE_E4M3, DATA_TYPE_E4M3, DATA_TYPE_E4M3, DATA_TYPE_BF16, 0, 0, 0, 0, false); - for (int32_t const numHeads : {32, 64}) - { - for (int32_t const kvLength : {129, 640, kCaptureKvLength}) - { - auto params = makeParams(numHeads); - params.mMaxSeqLenKv = kvLength; - auto const [supported, info] = runner.isSupportedWithInfo(params); - ASSERT_TRUE(supported) << info; - EXPECT_NE(info.find("multiCtasKvMode=0"), std::string::npos) - << "heads=" << numHeads << ", kvLength=" << kvLength << ": " << info; - } - } -} - -TEST_F(FmhaDeterministicTest, CudaGraphMatchesEagerForChangingSparseLengths) -{ - auto params = makeParams(); - auto* q = allocate(kNumHeads * kHeadDim); - auto* kv = allocate(kSparseTopK * kHeadDim); - auto* output = allocate(kNumHeads * kHeadDim); - auto* sequenceLength = allocate(1); - auto* topkLength = allocate(1); - auto* indices = allocate(kSparseTopK); - auto* scale = allocate(1); - auto* softmaxScale = allocate(1); - std::vector qValues(kNumHeads * kHeadDim); - std::vector kvValues(kSparseTopK * kHeadDim); - for (size_t i = 0; i < qValues.size(); ++i) - { - qValues[i] = static_cast(24 + (i * 7 % 24) + ((i / 3 % 2) << 7)); - } - for (size_t i = 0; i < kvValues.size(); ++i) - { - kvValues[i] = static_cast(24 + (i * 11 % 24) + ((i / 5 % 2) << 7)); - } - std::vector indexValues(kSparseTopK); - for (int32_t i = 0; i < kSparseTopK; ++i) - { - indexValues[i] = i; - } - copyToDevice(q, qValues); - copyToDevice(kv, kvValues); - copyToDevice(indices, indexValues); - copyToDevice(scale, std::vector{1.0F}); - copyToDevice(softmaxScale, std::vector{static_cast(std::log2(std::exp(1.0)) / std::sqrt(kHeadDim))}); - copyToDevice(sequenceLength, std::vector{kCaptureKvLength}); - copyToDevice(topkLength, std::vector{kSparseTopK}); - params.qPtr = q; - params.kvPtr = kv; - params.slidingWindowKvPoolBasePtr = kv; - params.oPtr = output; - params.seqLensKvPtr = sequenceLength; - params.kvPageIdxPtr = indices; - params.ptrSparseMlaTopKLens = topkLength; - params.outputScalePtr = scale; - params.scaleSoftmaxLog2Ptr = softmaxScale; - params.kvSfScalePtr = scale; - params.oSfScalePtr = scale; - TllmGenFmhaRunner runner(DATA_TYPE_E4M3, DATA_TYPE_E4M3, DATA_TYPE_E4M3, DATA_TYPE_BF16, 0, 0, 0, 0, false); - runner.run(params); - TLLM_CUDA_CHECK(cudaStreamSynchronize(mStream)); - TLLM_CUDA_CHECK(cudaStreamBeginCapture(mStream, cudaStreamCaptureModeGlobal)); - runner.run(params); - TLLM_CUDA_CHECK(cudaStreamEndCapture(mStream, &mGraph)); - TLLM_CUDA_CHECK(cudaGraphInstantiate(&mGraphExec, mGraph, nullptr, nullptr, 0)); - - for (int32_t const kvLength : {127, 128, 129, 160, 513, 2049, 8193}) - { - SCOPED_TRACE(kvLength); - int32_t const selected = std::min(kvLength, 128) + std::min(kvLength / 4, 512); - copyToDevice(sequenceLength, std::vector{kvLength}); - copyToDevice(topkLength, std::vector{selected}); - params.mMaxSeqLenKv = kvLength; - params.mSumOfSeqLensKv = kvLength; - runner.run(params); - auto const eager = readOutput(output); - ASSERT_TRUE( - std::all_of(eager.begin(), eager.end(), [](uint16_t value) { return (value & 0x7f80U) != 0x7f80U; })); - TLLM_CUDA_CHECK(cudaGraphLaunch(mGraphExec, mStream)); - auto const graph = readOutput(output); - EXPECT_EQ(eager, graph); - } -} - -} // namespace - -int main(int argc, char** argv) -{ - if (setenv("FORCE_ATTENTION_KERNEL_DETERMINISTIC", "1", 1) != 0) - { - return EXIT_FAILURE; - } - ::testing::InitGoogleTest(&argc, argv); - return RUN_ALL_TESTS(); -} diff --git a/tensorrt_llm/_torch/disaggregation/native/transfer.py b/tensorrt_llm/_torch/disaggregation/native/transfer.py index d3bc158524c3..99f5c35d7a0b 100644 --- a/tensorrt_llm/_torch/disaggregation/native/transfer.py +++ b/tensorrt_llm/_torch/disaggregation/native/transfer.py @@ -589,7 +589,7 @@ class _SessionQuiescence: class Sender(SenderBase): - # Minimum time-to-live for orphaned RecvReqInfo entries (seconds). + # Time-to-live for orphaned RecvReqInfo entries (seconds). # In gen-first ADP broadcast, non-assigned DP ranks accumulate # RecvReqInfo that never gets consumed. Entries older than this # are evicted during periodic sweeps. @@ -602,7 +602,6 @@ def __init__( bounce=None, enforce_physical_ownership: bool = False, retirement_watchdog: Optional[RetirementWatchdog] = None, - request_info_timeout_s: Optional[float] = None, ) -> None: self._registrar = peer_registrar self._device_id = peer_registrar.self_rank_info.device_id @@ -610,9 +609,6 @@ def __init__( self._bounce = bounce self._enforce_physical_ownership = enforce_physical_ownership self._retirement_watchdog = retirement_watchdog - # Generation-first metadata can arrive while prefill is still queued, so it - # is kept for at least the configured transfer timeout (_request_info_ttl). - self._request_info_timeout_s = request_info_timeout_s or 0.0 self._peer_requests: dict = {} self._peer_requests_timestamps: dict[int, float] = {} # unique_rid -> insert time self._peer_requests_lock = threading.Lock() @@ -708,40 +704,27 @@ def sweep_stale_req_infos(self): """ now = time.monotonic() with self._peer_requests_lock: - stale_rids = {} - for rid, ts in self._peer_requests_timestamps.items(): - ttl = self._request_info_ttl(rid) - if now - ts > ttl: - stale_rids[rid] = ttl + stale_rids = [ + rid + for rid, ts in self._peer_requests_timestamps.items() + if now - ts > self._STALE_REQ_INFO_TTL_S + ] if not stale_rids: return - orphaned_infos: list[tuple[RecvReqInfo, float]] = [] - for rid, ttl in stale_rids.items(): + orphaned_infos: list[RecvReqInfo] = [] + for rid in stale_rids: with self._sessions_lock, self._peer_requests_lock: if rid not in self._sessions and rid in self._peer_requests: req_info_map = self._peer_requests.pop(rid) self._peer_requests_timestamps.pop(rid, None) orphaned_infos.extend( - (info, ttl) for info in req_info_map.values() if info.context_first + info for info in req_info_map.values() if info.context_first ) logger.debug(f"Swept stale RecvReqInfo for rid={rid}") - for info, ttl in orphaned_infos: - self._fail_orphaned_receiver(info, ttl) - - def _request_info_ttl(self, unique_rid: int) -> float: - """How long REQUEST_DATA for ``unique_rid`` waits for its TxSession. - - A generation-first request can arrive while its prefill is still queued, - so it waits as long as the transfer may take. A context-first one names a - session that should already exist, so it is failed after the base TTL. - Called with ``_peer_requests_lock`` held. - """ - infos = self._peer_requests.get(unique_rid) or {} - if any(not info.context_first for info in infos.values()): - return max(self._STALE_REQ_INFO_TTL_S, self._request_info_timeout_s) - return self._STALE_REQ_INFO_TTL_S + for info in orphaned_infos: + self._fail_orphaned_receiver(info) - def _fail_orphaned_receiver(self, info: RecvReqInfo, waited_s: float) -> None: + def _fail_orphaned_receiver(self, info: RecvReqInfo) -> None: """Report FAILED for a context-first REQUEST_DATA that no TxSession claimed. This resolves the receiver's task, so its RxSession fails through the @@ -751,7 +734,7 @@ def _fail_orphaned_receiver(self, info: RecvReqInfo, waited_s: float) -> None: logger.warning( f"KV transfer request rid={info.unique_rid} from receiver " f"{info.instance_name} rank {info.instance_rank} has no TxSession on this " - f"rank after {waited_s:.0f}s; failing it" + f"rank after {self._STALE_REQ_INFO_TTL_S:.0f}s; failing it" ) self._send_failed_result_to_receiver(info, defer_to_worker=True) @@ -4433,7 +4416,6 @@ def _setup_transfer_engine(self): bounce=self._bounce, enforce_physical_ownership=self._config.enforce_physical_ownership, retirement_watchdog=self._retirement_watchdog, - request_info_timeout_s=self._config.tx_overall_timeout_s, ) self._receiver = Receiver( self._peer_registrar, diff --git a/tensorrt_llm/_torch/modules/mhc/mhc_cuda.py b/tensorrt_llm/_torch/modules/mhc/mhc_cuda.py index a91567d8a1a8..a7ff8409fcc6 100644 --- a/tensorrt_llm/_torch/modules/mhc/mhc_cuda.py +++ b/tensorrt_llm/_torch/modules/mhc/mhc_cuda.py @@ -826,17 +826,11 @@ def _alloc_fused_hc_scratch( _FUSED_HC_FALLBACK_TACTIC_FMA = ("fused_half_fma", 2, 1, 256, 1) -def _get_fused_hc_fallback_tactic( - hidden_size: int | None = None, *, deterministic: bool = False -) -> tuple[str, int, int, int, int]: +def _get_fused_hc_fallback_tactic(hidden_size: int | None = None): mma_ok = _fused_hc_mma_supported() if hidden_size is not None: mma_ok = mma_ok and hidden_size in _FUSED_HC_MMA_SUPPORTED_HIDDEN_SIZES - if mma_ok: - if deterministic: - return ("fused_half_mma", 0, 1, 0, 1) - return _FUSED_HC_FALLBACK_TACTIC_MMA - return _FUSED_HC_FALLBACK_TACTIC_FMA + return _FUSED_HC_FALLBACK_TACTIC_MMA if mma_ok else _FUSED_HC_FALLBACK_TACTIC_FMA class MhcFusedHcRunner(TunableRunner): @@ -860,11 +854,6 @@ class MhcFusedHcRunner(TunableRunner): Fallback (-1): delegates to the C++ heuristic (fused_half_mma + auto ks + auto bs). - - ``FORCE_DETERMINISTIC=1`` or ``FORCE_MHC_KERNEL_DETERMINISTIC=1`` uses a - fixed half-fused tactic with one K split for tuned and fallback execution, - except that the FMA fallback tactic is kept: ``mhc_fused_hc`` passes it under - ``FORCE_DETERMINISTIC=1`` for its row-count-independent reduction order. """ tuning_config = TuningConfig( @@ -898,10 +887,6 @@ def __init__( ): self.n = n self.hidden_size = hidden_size - self._force_deterministic = ( - os.getenv("FORCE_DETERMINISTIC", "0") == "1" - or os.getenv("FORCE_MHC_KERNEL_DETERMINISTIC", "0") == "1" - ) self.rms_eps = rms_eps self.hc_pre_eps = hc_pre_eps self.hc_sinkhorn_eps = hc_sinkhorn_eps @@ -909,13 +894,9 @@ def __init__( self.sinkhorn_repeat = sinkhorn_repeat def unique_id(self): - if self._force_deterministic: - return (self.n, self.hidden_size, "deterministic") return (self.n, self.hidden_size) def get_valid_tactics(self, inputs, profile: OptimizationProfile, **kwargs): - if self._force_deterministic: - return [_get_fused_hc_fallback_tactic(self.hidden_size, deterministic=True)] M = inputs[0].shape[0] tactics = [] seen = set() @@ -1001,15 +982,8 @@ def forward(self, inputs, *, tactic=-1, **kwargs): hc_base_cur = hc_base_cur.to(torch.float32).contiguous() pre_mix_ext = kwargs.get("pre_mix_ext") - if tactic == -1 or (self._force_deterministic and tactic != _FUSED_HC_FALLBACK_TACTIC_FMA): - # One fixed tactic avoids both atomic split-K reductions and - # arithmetic changes from independently tuned rank/process choices. - # The FMA fallback stays: it has neither, and mhc_fused_hc passes it - # under FORCE_DETERMINISTIC because its reduction order does not - # depend on the number of rows. - tactic = _get_fused_hc_fallback_tactic( - self.hidden_size, deterministic=self._force_deterministic - ) + if tactic == -1: + tactic = _get_fused_hc_fallback_tactic(self.hidden_size) backend, tile_n, num_k_splits, bigfuse_bs, tile_m = tactic backend_code = _FUSED_HC_BACKEND_CODE[backend] diff --git a/tensorrt_llm/_torch/modules/top_k.py b/tensorrt_llm/_torch/modules/top_k.py index 4ab65437b64a..2a2fda1b87d3 100644 --- a/tensorrt_llm/_torch/modules/top_k.py +++ b/tensorrt_llm/_torch/modules/top_k.py @@ -4,7 +4,6 @@ from __future__ import annotations -import os from enum import Enum import torch @@ -37,10 +36,6 @@ class TopK(nn.Module): request metadata and retain a stable address across CUDA Graph replays. Native CUDA scratch is supplied by the caller so its lifetime follows the caller's eager or CUDA-graph metadata. - - ``FORCE_DETERMINISTIC=1`` or ``FORCE_TOPK_KERNEL_DETERMINISTIC=1`` selects - the PyTorch implementation for repeatable indices, including their ordering - and selection among tied scores. """ def __init__( @@ -56,13 +51,6 @@ def __init__( ) -> None: super().__init__() self.top_k = top_k - if ( - os.getenv("FORCE_DETERMINISTIC", "0") == "1" - or os.getenv("FORCE_TOPK_KERNEL_DETERMINISTIC", "0") == "1" - ): - # Native selection uses atomic appends, including for ties at the - # cutoff. Fix both the selected set and its order for repeatability. - prefill_implementation = decode_implementation = TopKImplementation.TORCH self.prefill_implementation = TopKImplementation( prefill_implementation or TopKImplementation.CUDA_RADIX ) diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 64eb2abf2c28..6ce9b590f6d2 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -138,7 +138,6 @@ l0_b200: - unittest/_torch/modeling/test_glm5_next_contracts.py - unittest/llmapi/test_deepseek_v4_tokenizer.py - unittest/_torch/modules/test_mhc.py - - unittest/_torch/modules/test_top_k.py - unittest/_torch/modules/test_engram.py - unittest/_torch/modules/test_gate_up_swiglu_quack.py - unittest/_torch/modules/test_multi_stream_utils.py diff --git a/tests/unittest/_torch/modules/test_mhc.py b/tests/unittest/_torch/modules/test_mhc.py index 154daaefff88..0a072bcfb4b1 100644 --- a/tests/unittest/_torch/modules/test_mhc.py +++ b/tests/unittest/_torch/modules/test_mhc.py @@ -22,46 +22,6 @@ timing_stats = defaultdict(dict) -@skip_pre_blackwell -@pytest.mark.parametrize("num_tokens", [1, 129]) -@pytest.mark.parametrize("force_fma", [False, True]) -@pytest.mark.parametrize( - "deterministic_env", ["FORCE_DETERMINISTIC", "FORCE_MHC_KERNEL_DETERMINISTIC"] -) -def test_mhc_deterministic_mode_repeats_with_cached_tactic( - monkeypatch: pytest.MonkeyPatch, num_tokens: int, force_fma: bool, deterministic_env: str -) -> None: - from tensorrt_llm._torch.modules.mhc import mhc_cuda - - monkeypatch.delenv("FORCE_DETERMINISTIC", raising=False) - monkeypatch.delenv("FORCE_MHC_KERNEL_DETERMINISTIC", raising=False) - monkeypatch.setenv(deterministic_env, "1") - if force_fma: - monkeypatch.setattr(mhc_cuda, "_fused_hc_mma_supported", lambda: False) - runner = mhc_cuda.MhcFusedHcRunner(4, 4096, 1e-6, 1e-6, 1e-6, 2.0, 20) - generator = torch.Generator(device="cuda").manual_seed(0) - - def random_tensor(*shape: int, dtype: torch.dtype = torch.float32) -> torch.Tensor: - return torch.randn(shape, generator=generator, dtype=dtype, device="cuda") - - inputs = [ - random_tensor(num_tokens, 4096, dtype=torch.bfloat16), - random_tensor(num_tokens, 4, 4096, dtype=torch.bfloat16), - random_tensor(num_tokens, 4).sigmoid(), - random_tensor(num_tokens, 4, 4).softmax(-1), - random_tensor(24, 16384) * 0.01, - torch.ones(3, device="cuda"), - torch.zeros(24, device="cuda"), - ] - assert len(runner.get_valid_tactics(inputs, None)) == 1 - reference = runner(inputs) - for _ in range(20): - actual = runner(inputs, tactic=("fused_half_mma", 0, 16, 0, 1)) - for expected, value in zip(reference, actual): - assert torch.isfinite(value).all() - assert torch.equal(value, expected) - - def _mhc_fused_hc_mma_available() -> bool: try: from tensorrt_llm._torch.modules.mhc.mhc_cuda import _fused_hc_mma_supported @@ -1238,9 +1198,7 @@ def _make_fused_hc_runner_case(n: int, hidden_size: int, hc_mult: int, seed: int return runner, inputs -# 300 rows take a different automatic BigFuse block size than one row, which the -# row-invariance check below must not see. -@pytest.mark.parametrize("num_tokens", [1, 129, 300]) +@pytest.mark.parametrize("num_tokens", [1, 129]) @pytest.mark.parametrize("hidden_size", [4096, 7168]) @pytest.mark.parametrize("fuse_norm", [False, True]) @pytest.mark.parametrize("lagged", [False, True]) @@ -1302,15 +1260,6 @@ def run( single_pre_mix = pre_mix_ext[row : row + 1].clone() if pre_mix_ext is not None else None for actual, reference in zip(run(single_inputs, single_pre_mix), expected, strict=True): torch.testing.assert_close(actual, reference[row : row + 1], rtol=0, atol=0) - # Nor may splitting the batch, which checks every row. - half = (num_tokens + 1) // 2 - for start, stop in ((0, half), (half, num_tokens)): - if start == stop: - continue - part_inputs = [v[start:stop].clone() for v in inputs[:4]] + inputs[4:] - part_pre_mix = pre_mix_ext[start:stop].clone() if pre_mix_ext is not None else None - for actual, reference in zip(run(part_inputs, part_pre_mix), expected, strict=True): - torch.testing.assert_close(actual, reference[start:stop], rtol=0, atol=0) stream = torch.cuda.Stream() stream.wait_stream(torch.cuda.current_stream()) diff --git a/tests/unittest/_torch/modules/test_top_k.py b/tests/unittest/_torch/modules/test_top_k.py index a2828160a525..0c9b7522e027 100644 --- a/tests/unittest/_torch/modules/test_top_k.py +++ b/tests/unittest/_torch/modules/test_top_k.py @@ -13,47 +13,6 @@ from tensorrt_llm._torch.modules.top_k import _CUTE_DSL_PREFILL_COPY_BITS, TopK, TopKImplementation -@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") -@pytest.mark.parametrize("is_prefill", [False, True]) -@pytest.mark.parametrize("implementation", list(TopKImplementation)) -@pytest.mark.parametrize( - "deterministic_env", ["FORCE_DETERMINISTIC", "FORCE_TOPK_KERNEL_DETERMINISTIC"] -) -def test_deterministic_topk_repeats_tied_selection( - monkeypatch: pytest.MonkeyPatch, - is_prefill: bool, - implementation: TopKImplementation, - deterministic_env: str, -) -> None: - monkeypatch.delenv("FORCE_DETERMINISTIC", raising=False) - monkeypatch.delenv("FORCE_TOPK_KERNEL_DETERMINISTIC", raising=False) - monkeypatch.setenv(deterministic_env, "1") - generator = torch.Generator(device="cuda").manual_seed(0) - scores = torch.randn(3, 2048, generator=generator, device="cuda").round() - output = torch.empty(3, 512, dtype=torch.int32, device="cuda") - lengths = torch.tensor([2048, 3, 0], dtype=torch.int32, device="cuda") - top_k = TopK( - 512, - prefill_implementation=implementation, - decode_implementation=implementation, - compress_ratio=4, - ) - kwargs = dict( - is_prefill=is_prefill, - row_starts=torch.zeros_like(lengths), - row_ends=lengths, - sequence_lengths=lengths * 4, - scan_lengths=lengths, - ) - reference = top_k(scores, output, **kwargs).clone() - assert (reference >= 0).sum(dim=1).tolist() == [512, 3, 0] - selected_scores = scores[0, reference[0].long()].sort().values - torch.testing.assert_close(selected_scores, scores[0].topk(512).values.sort().values) - for _ in range(20): - actual = top_k(scores, output, **kwargs) - assert torch.equal(actual, reference) - - def test_prefill_torch_masks_dirty_scores_and_pads_output() -> None: scores = torch.tensor( [ diff --git a/tests/unittest/disaggregated/test_unsubmitted_session.py b/tests/unittest/disaggregated/test_unsubmitted_session.py index 5b9460418858..fbe719015676 100644 --- a/tests/unittest/disaggregated/test_unsubmitted_session.py +++ b/tests/unittest/disaggregated/test_unsubmitted_session.py @@ -17,7 +17,7 @@ import queue import threading from types import SimpleNamespace -from unittest.mock import Mock, patch +from unittest.mock import Mock import numpy as np import pytest @@ -47,7 +47,6 @@ def _sender(*, ownership: bool = True) -> transfer_mod.Sender: sender._pre_cancelled_rids = {} sender._peer_requests_lock = threading.Lock() sender._peer_requests, sender._peer_requests_timestamps = {}, {} - sender._request_info_timeout_s = 0.0 sender._shutdown, sender._shutdown_requested = True, False sender._retirement_watchdog = None sender._instance_rank = 7 @@ -140,59 +139,6 @@ def test_metadata_expiry_does_not_prevent_cancel_acknowledgment() -> None: assert not sender._sessions -@pytest.mark.parametrize("transfer_timeout_s", [None, 60.0, 600.0]) -def test_queued_prefill_metadata_uses_transfer_timeout( - monkeypatch: pytest.MonkeyPatch, transfer_timeout_s: float | None -) -> None: - monkeypatch.setattr(transfer_mod, "ZMQMessenger", Mock()) - monkeypatch.setattr(transfer_mod, "KV_TRANSFER_NUM_THREADS", 0) - now = 0.0 - monkeypatch.setattr(transfer_mod.time, "monotonic", lambda: now) - sender = transfer_mod.Sender(Mock(), Mock(), request_info_timeout_s=transfer_timeout_s) - info = transfer_mod.RecvReqInfo.from_bytes(_request_data()[1]) - try: - sender._add_req_info(401, 0, info) - now = 121.0 - sender.sweep_stale_req_infos() - if transfer_timeout_s == 600.0: - assert sender._get_req_info(401) == {0: info} - else: - assert sender._get_req_info(401) is None - - now = 601.0 - sender.sweep_stale_req_infos() - assert sender._get_req_info(401) is None - finally: - sender.shutdown() - - -def test_orphaned_context_first_request_fails_after_the_base_ttl( - monkeypatch: pytest.MonkeyPatch, -) -> None: - """A context-first request's long transfer timeout does not keep its receiver waiting. - - Its REQUEST_DATA names a session that should already exist. - """ - monkeypatch.setattr(transfer_mod, "ZMQMessenger", Mock()) - monkeypatch.setattr(transfer_mod, "KV_TRANSFER_NUM_THREADS", 0) - now = 0.0 - monkeypatch.setattr(transfer_mod.time, "monotonic", lambda: now) - sender = transfer_mod.Sender(Mock(), Mock(), request_info_timeout_s=600.0) - info = transfer_mod.RecvReqInfo.from_bytes(_request_data()[1]) - info.context_first = True - try: - sender._add_req_info(401, 0, info) - with patch.object(sender, "_fail_orphaned_receiver") as fail: - now = 121.0 - sender.sweep_stale_req_infos() - - assert sender._get_req_info(401) is None - # The log names the TTL the entry actually waited. - fail.assert_called_once_with(info, transfer_mod.Sender._STALE_REQ_INFO_TTL_S) - finally: - sender.shutdown() - - @pytest.mark.parametrize("pending_settlement", [False, True]) @pytest.mark.parametrize("ack_fails_once", [False, True]) def test_worker_orders_session_ack_after_piece_reports( From bcf6df1a6b2f8aeb5ed3bab96ca83d4b1b34b9d7 Mon Sep 17 00:00:00 2001 From: Guiju Zhang <7135567+cascade812@users.noreply.github.com> Date: Fri, 30 Oct 2026 00:08:08 -0700 Subject: [PATCH 05/11] [TRTLLM-16282][test] Add disaggregated sparse KV offload E2E tests for DeepSeek-V4-Flash and GLM-5.2 Each test runs one context server and one generation server, turns on sparse KV cache offload on the generation server only, and checks GSM8K against the model's reference. - TestDeepSeekV4Flash::test_kv_cache_offload: TP2 + TP2 on 4 GPUs with an NVFP4 KV cache. The generation server uses attention DP and turns off block reuse, which offload does not support. - TestGLM52NVFP4::test_kv_cache_offload: TP4 + TP4 on 8 GPUs with the checkpoint's FP8 KV cache and attention DP. It drops the MTP of test_nvfp4_nixl, since offload takes no speculative decoding. The generation server's batch is 16: each row holds a 177.5 MiB GPU cache for its selected history. Both join the 8-GPU B200 pre-merge stage in l0_dgx_b200.yml. Signed-off-by: Guiju Zhang <7135567+cascade812@users.noreply.github.com> --- .../accuracy/test_disaggregated_serving.py | 142 ++++++++++++++++++ .../test_lists/test-db/l0_dgx_b200.yml | 2 + 2 files changed, 144 insertions(+) diff --git a/tests/integration/defs/accuracy/test_disaggregated_serving.py b/tests/integration/defs/accuracy/test_disaggregated_serving.py index 4bf85c4041ce..a4fd04f240b7 100644 --- a/tests/integration/defs/accuracy/test_disaggregated_serving.py +++ b/tests/integration/defs/accuracy/test_disaggregated_serving.py @@ -2275,6 +2275,79 @@ def test_nvfp4_nixl(self, use_kv_cache_manager_v2): max_draft_len=speculative_config["max_draft_len"]) run_accuracy_test(llm, self.MODEL_NAME, ["GSM8K"]) + @pytest.mark.skip_less_device(8) + def test_kv_cache_offload(self): + """The generation server keeps the MLA latent history in host memory. + + Each decode step fetches the indexer-selected tokens into a GPU cache. + Only a decode-only worker can offload, so the context server keeps its + history resident. Offload takes no speculative decoding and no block + reuse. + """ + kv_cache_config = { + "free_gpu_memory_fraction": 0.7, + "enable_block_reuse": False, + "use_kv_cache_manager_v2": True, + } + cache_transceiver_config = { + "backend": "NIXL", + "transceiver_runtime": "PYTHON", + } + moe_config = {"backend": "CUTEDSL"} + ctx_server_config = { + "tensor_parallel_size": 4, + "pipeline_parallel_size": 1, + "moe_expert_parallel_size": 4, + "enable_attention_dp": True, + "disable_overlap_scheduler": True, + "enable_chunked_prefill": True, + "cuda_graph_config": None, + "trust_remote_code": True, + "max_seq_len": 8192, + "kv_cache_config": kv_cache_config, + "moe_config": moe_config, + "cache_transceiver_config": cache_transceiver_config, + } + gen_server_config = { + "tensor_parallel_size": 4, + "pipeline_parallel_size": 1, + "moe_expert_parallel_size": 4, + "enable_attention_dp": True, + "disable_overlap_scheduler": False, + "enable_chunked_prefill": True, + "trust_remote_code": True, + # Each batch row holds a GPU cache for its selected history, about + # 180 MiB on GLM-5.2, so the batch stays small. + "max_batch_size": 16, + "max_seq_len": 8192, + "kv_cache_config": kv_cache_config, + "moe_config": moe_config, + "cache_transceiver_config": cache_transceiver_config, + "sparse_attention_config": { + "algorithm": "dsa", + "enable_kv_cache_offload": True, + }, + } + disaggregated_server_config = { + "hostname": "localhost", + "backend": "pytorch", + "context_servers": { + "num_instances": 1 + }, + "generation_servers": { + "num_instances": 1 + } + } + with launch_disaggregated_llm(disaggregated_server_config, + ctx_server_config, + gen_server_config, + self.MODEL_PATH, + max_workers=16) as llm: + # The DuckLLM's bare LlmArgs takes NVFP4 from the model path; the + # checkpoint's FP8 KV cache selects the reference entry. + llm.args.quant_config.kv_cache_quant_algo = "FP8" + run_accuracy_test(llm, self.MODEL_NAME, ["GSM8K"]) + @pytest.mark.timeout(DEFAULT_TEST_TIMEOUT) @skip_pre_blackwell @@ -2389,6 +2462,75 @@ def test_gen_first(self): task = MMLU(self.MODEL_NAME) task.evaluate(llm, is_integration_test=True) + # The servers take up to an hour to start, before a full GSM8K run. + @pytest.mark.timeout(5400) + @pytest.mark.skip_less_device(4) + def test_kv_cache_offload(self): + """The generation server keeps the ratio-4 compressed history in host memory. + + Each decode step fetches the indexer-selected units into a GPU cache. + Only a decode-only worker can offload, so the context server keeps its + history resident. Offload takes no block reuse. + """ + cache_transceiver_config = { + "backend": "NIXL", + "transceiver_runtime": "PYTHON", + } + # Both servers lay out the NVFP4 KV cache alike for the transfer. + kv_cache_config = { + "dtype": "nvfp4", + "free_gpu_memory_fraction": 0.5, + } + ctx_server_config = { + "tensor_parallel_size": 2, + "moe_expert_parallel_size": 2, + "disable_overlap_scheduler": True, + "max_batch_size": DEEPSEEKV4_TEST_MAX_BATCH_SIZE, + "max_seq_len": 4096, + "kv_cache_config": kv_cache_config, + "cache_transceiver_config": cache_transceiver_config, + } + gen_server_config = { + "tensor_parallel_size": 2, + "moe_expert_parallel_size": 2, + "enable_attention_dp": True, + "disable_overlap_scheduler": True, + "max_batch_size": DEEPSEEKV4_TEST_MAX_BATCH_SIZE, + "max_seq_len": 4096, + "moe_config": { + "backend": "TRTLLM", + }, + "kv_cache_config": { + **kv_cache_config, + "enable_block_reuse": False, + }, + "cache_transceiver_config": cache_transceiver_config, + "sparse_attention_config": { + "algorithm": "deepseek_v4", + "enable_kv_cache_offload": True, + }, + } + disaggregated_server_config = { + "hostname": "localhost", + "backend": "pytorch", + "context_servers": { + "num_instances": 1 + }, + "generation_servers": { + "num_instances": 1 + }, + } + with launch_disaggregated_llm(disaggregated_server_config, + ctx_server_config, + gen_server_config, + self.MODEL_PATH, + server_waiting_timeout=3600) as llm: + # The DuckLLM's bare LlmArgs knows no checkpoint quantization; these + # select the reference entry. + llm.args.quant_config.quant_algo = QuantAlgo.FP8_BLOCK_SCALES + llm.args.quant_config.kv_cache_quant_algo = QuantAlgo.NVFP4 + run_accuracy_test(llm, self.MODEL_NAME, ["GSM8K"]) + @pytest.mark.timeout(14400) @skip_pre_blackwell diff --git a/tests/integration/test_lists/test-db/l0_dgx_b200.yml b/tests/integration/test_lists/test-db/l0_dgx_b200.yml index 5dfdc6fbcdd4..b4b664cb7e37 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_b200.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_b200.yml @@ -229,6 +229,8 @@ l0_dgx_b200: - accuracy/test_disaggregated_serving.py::TestNemotron3Super120B::test_auto_dtype[mtp_nextn=0-block_reuse=False-use_py_transceiver=True] TIMEOUT (60) - accuracy/test_disaggregated_serving.py::TestNemotron3Super120B::test_auto_dtype[mtp_nextn=3-block_reuse=True-use_py_transceiver=True] TIMEOUT (60) - accuracy/test_disaggregated_serving.py::TestGLM52NVFP4::test_nvfp4_nixl[cache_mgr_v2] TIMEOUT (60) + - accuracy/test_disaggregated_serving.py::TestGLM52NVFP4::test_kv_cache_offload TIMEOUT (60) + - accuracy/test_disaggregated_serving.py::TestDeepSeekV4Flash::test_kv_cache_offload TIMEOUT (90) # DeepSeek-V4 EPLB pre-merge sanity (uncomment once DeepSeek-V4-Flash/Flash-Base # checkpoints are staged under llm_models_root()). # - accuracy/test_llm_api_pytorch.py::TestDeepSeekV4Flash::test_nvfp4_4gpus_static_eplb[moe_backend=TRTLLM] TIMEOUT (120) From 2b9147496b1a1de1c95591e0c1ee91da902c98c4 Mon Sep 17 00:00:00 2001 From: Guiju Zhang <7135567+cascade812@users.noreply.github.com> Date: Fri, 30 Oct 2026 22:42:39 -0700 Subject: [PATCH 06/11] [TRTLLM-16282][test] Give the Helix KV-transfer test fakes a VRAM pool view Sender._build_kv_write_meta groups writes by the memory type of each pool view's physical pool, so a page table it reads needs pool views and pool groups. The Helix tests' fake page tables had neither and failed with AttributeError on pool_views. Both helpers now build a one-layer-group page table whose only pool view is backed by GPU memory. Signed-off-by: Guiju Zhang <7135567+cascade812@users.noreply.github.com> --- .../disaggregated/test_kv_transfer.py | 24 ++++++++++++++----- 1 file changed, 18 insertions(+), 6 deletions(-) diff --git a/tests/unittest/disaggregated/test_kv_transfer.py b/tests/unittest/disaggregated/test_kv_transfer.py index 4fd2b17cf6d0..7723e719c102 100644 --- a/tests/unittest/disaggregated/test_kv_transfer.py +++ b/tests/unittest/disaggregated/test_kv_transfer.py @@ -2248,6 +2248,22 @@ def collect_on_worker() -> None: _shutdown_transfer_worker_setup(healthy) +def _single_vram_pool_page_table(tpb, kind): + """A one-layer-group page table whose only pool view is backed by GPU memory. + + The sender reads each pool view's physical pool to group writes by memory type. + """ + return SimpleNamespace( + tokens_per_block=tpb, + layer_groups={ + 0: SimpleNamespace( + kind=kind, pool_group_idx=0, pool_views=[SimpleNamespace(pool_idx=0)] + ) + }, + pool_groups=[SimpleNamespace(pools=[SimpleNamespace(memory_type="VRAM")])], + ) + + def _helix_kv_write_meta_block_pairs(src_table, dst_table, cp_rank, cp_size, prompt_len, tpb): """Run Sender._build_kv_write_meta for one Helix receiver; return (src, dst) block pairs. @@ -2267,9 +2283,7 @@ def _helix_kv_write_meta_block_pairs(src_table, dst_table, cp_rank, cp_size, pro class _BlockExtractor: def __init__(self, kind): - self.page_table = SimpleNamespace( - tokens_per_block=tpb, layer_groups={0: SimpleNamespace(kind=kind)} - ) + self.page_table = _single_vram_pool_page_table(tpb, kind) def extract(self, block_ids, layer_group, pool_idx): return np.asarray(block_ids, dtype=np.int64) @@ -2371,9 +2385,7 @@ def _helix_ring_write_meta( class _PtrExtractor: def __init__(self, base): self.base = base - self.page_table = SimpleNamespace( - tokens_per_block=tpb, layer_groups={0: SimpleNamespace(kind=CacheKind.PAGED)} - ) + self.page_table = _single_vram_pool_page_table(tpb, CacheKind.PAGED) def extract(self, block_ids, layer_group, pool_idx): ptrs = self.base + np.asarray(block_ids, dtype=np.int64) * 1_000_000 From 1f09b8ff17d3b9156a69a1a47ad67ef1673bd9a2 Mon Sep 17 00:00:00 2001 From: Guiju Zhang <7135567+cascade812@users.noreply.github.com> Date: Sat, 31 Oct 2026 22:25:36 -0700 Subject: [PATCH 07/11] [TRTLLM-16288][test] Restore real admission in disagg receive ordering tests Initialize the mocked manager without host-history reservations and bind _prepare_disagg_gen_init_cache to its real implementation. This keeps the receive ordering tests exercising cache allocation and readiness event recording through the extracted admission helper. Validation: reproduced the five CI failures with the original fixture; all 14 CPU cases pass with the updated fixture using the matching CI wheel in an isolated pytest run. Pre-commit checks pass. Signed-off-by: Guiju Zhang <7135567+cascade812@users.noreply.github.com> --- .../_torch/pyexecutor/test_disagg_receive_ordering.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/unittest/_torch/pyexecutor/test_disagg_receive_ordering.py b/tests/unittest/_torch/pyexecutor/test_disagg_receive_ordering.py index 89d9e727ff11..130359ac5d23 100644 --- a/tests/unittest/_torch/pyexecutor/test_disagg_receive_ordering.py +++ b/tests/unittest/_torch/pyexecutor/test_disagg_receive_ordering.py @@ -52,6 +52,7 @@ def _manager() -> Mock: """Run the real admission and resource hooks with mocked cache allocation.""" manager = Mock(spec=KVCacheManagerV2) manager._sharing = None # no lender: free_resources closes and removes the cache itself + manager._host_history_reservations = None manager.impl = Mock() manager._stream = Mock() manager._disagg_receive_ready = {} @@ -66,6 +67,9 @@ def _manager() -> Mock: manager.kv_cache_map = {} manager.prepare_context_cache.return_value = 0 manager.prepare_disagg_gen_init = MethodType(KVCacheManagerV2.prepare_disagg_gen_init, manager) + manager._prepare_disagg_gen_init_cache = MethodType( + KVCacheManagerV2._prepare_disagg_gen_init_cache, manager + ) manager.prepare_resources = MethodType(KVCacheManagerV2.prepare_resources, manager) return manager From b33e1f4e96a9cc510e6b8faf32286657ed9f882a Mon Sep 17 00:00:00 2001 From: Guiju Zhang <7135567+cascade812@users.noreply.github.com> Date: Sun, 1 Nov 2026 11:47:30 -0800 Subject: [PATCH 08/11] [None][fix] Restore auxiliary KV pool selection for DSA FMHA Prefer the sparse runtime auxiliary KV pool when translating DSA parameters for native FMHA, retaining the mapped primary pool as the resident fallback. Offload indices address the fetched GPU arena, so pairing them with the host-backed history pool produces incorrect attention and excessive latency. The phased-FMHA refactor dropped this override from the non-dynamic sparse path. Keep DeepSeek-V4 sliding-window and auxiliary-pool behavior unchanged. Add native host tests covering both pointer-table layouts, DSA offload, resident fallback, and dynamic sparse attention with and without an auxiliary pool. Fixes: 120ea6b8c7a15107aaafd7f696cc532308487ee7 ("[None][refactor] Split `thop.attention` into phased fmha (#17558)") Signed-off-by: Guiju Zhang <7135567+cascade812@users.noreply.github.com> --- cpp/tensorrt_llm/thop/attentionOp.h | 4 +- cpp/tests/unit_tests/thop/CMakeLists.txt | 20 +++++ cpp/tests/unit_tests/thop/fmhaParamsTest.cpp | 89 ++++++++++++++++++++ 3 files changed, 112 insertions(+), 1 deletion(-) create mode 100644 cpp/tests/unit_tests/thop/fmhaParamsTest.cpp diff --git a/cpp/tensorrt_llm/thop/attentionOp.h b/cpp/tensorrt_llm/thop/attentionOp.h index 964eb52eadd8..10c787fe62fd 100644 --- a/cpp/tensorrt_llm/thop/attentionOp.h +++ b/cpp/tensorrt_llm/thop/attentionOp.h @@ -514,7 +514,9 @@ struct FmhaParams } else { - sparse.sparse_kv_cache_pool = kvCachePool; + sparse.sparse_kv_cache_pool = fwd.sparse_runtime_params.aux_kv_cache_pool_ptr.has_value() + ? reinterpret_cast(fwd.sparse_runtime_params.aux_kv_cache_pool_ptr.value()) + : kvCachePool; } return sparse; } diff --git a/cpp/tests/unit_tests/thop/CMakeLists.txt b/cpp/tests/unit_tests/thop/CMakeLists.txt index 70733afed359..f01282ec9dae 100644 --- a/cpp/tests/unit_tests/thop/CMakeLists.txt +++ b/cpp/tests/unit_tests/thop/CMakeLists.txt @@ -14,6 +14,26 @@ # the License. if(${BUILD_PYT}) + add_tllm_gtest( + fmhaParamsTest + SOURCES + fmhaParamsTest.cpp + LINK_LIBRARIES + tllm::common_logging + trtllm_gen_fmha_interface + ${Python3_LIBRARIES} + ${TORCH_LIBRARIES} + OWNER + runtime + LABELS + "unit;runtime;host") + get_directory_property( + fmha_params_generated_include_dir + DIRECTORY "${PROJECT_SOURCE_DIR}/tensorrt_llm" DEFINITION + TRTLLM_FMHA_PARAMS_GENERATED_INCLUDE_DIR) + target_include_directories(fmhaParamsTest + PRIVATE ${fmha_params_generated_include_dir}) + add_tllm_gtest( thUtilsTest SOURCES diff --git a/cpp/tests/unit_tests/thop/fmhaParamsTest.cpp b/cpp/tests/unit_tests/thop/fmhaParamsTest.cpp new file mode 100644 index 000000000000..c19436b3650f --- /dev/null +++ b/cpp/tests/unit_tests/thop/fmhaParamsTest.cpp @@ -0,0 +1,89 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + * + * 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/thop/attentionOp.h" + +#include + +namespace +{ + +class FmhaSparseParamsTest : public ::testing::TestWithParam +{ +protected: + void SetUp() override + { + mPrimaryPool = torch::zeros({16}, torch::kInt32); + mAuxPool = torch::ones({16}, torch::kInt32); + auto poolPointers = torch::zeros({2, 2}, torch::kInt64); + poolPointers.index_put_({1, 0}, reinterpret_cast(mPrimaryPool.data_ptr())); + mParams.meta.host_kv_cache_pool_pointers = GetParam() ? poolPointers.unsqueeze(-1) : poolPointers; + mParams.meta.host_kv_cache_pool_mapping = torch::tensor({{0, 0}, {1, 0}}, torch::kInt32); + mParams.meta.kv_cache_block_offsets = torch::zeros({2, 1, 2, 1}, torch::kInt32); + mParams.local_layer_idx = 1; + mParams.fwd.sparse_runtime_params.sparse_attn_indices = torch::tensor({{0, 1}}, torch::kInt32); + mParams.meta.num_sparse_topk = 2; + } + + tensorrt_llm::torch_ext::FmhaParams mParams; + torch::Tensor mPrimaryPool; + torch::Tensor mAuxPool; +}; + +TEST_P(FmhaSparseParamsTest, DsaUsesAuxiliaryPool) +{ + mParams.fwd.sparse_runtime_params.aux_kv_cache_pool_ptr = reinterpret_cast(mAuxPool.data_ptr()); + + auto const sparse = mParams.getSparseParams(); + + EXPECT_EQ(sparse.sparse_kv_cache_pool, mAuxPool.data_ptr()); + EXPECT_EQ(sparse.sliding_window_kv_cache_pool, nullptr); + EXPECT_EQ(sparse.sparse_attn_indices, mParams.fwd.sparse_runtime_params.sparse_attn_indices->data_ptr()); + EXPECT_EQ(sparse.num_sparse_topk, 2); +} + +TEST_P(FmhaSparseParamsTest, ResidentDsaUsesMappedPool) +{ + auto const sparse = mParams.getSparseParams(); + + EXPECT_EQ(sparse.sparse_kv_cache_pool, mPrimaryPool.data_ptr()); + EXPECT_EQ(sparse.sliding_window_kv_cache_pool, nullptr); +} + +TEST_P(FmhaSparseParamsTest, DynamicSparseKeepsSlidingWindowAndAuxiliaryPoolsSeparate) +{ + mParams.fwd.sparse_runtime_params.sparse_attn_kv_lens = torch::tensor({2}, torch::kInt32); + mParams.fwd.sparse_runtime_params.aux_kv_cache_pool_ptr = reinterpret_cast(mAuxPool.data_ptr()); + + auto const sparse = mParams.getSparseParams(); + + EXPECT_EQ(sparse.sparse_kv_cache_pool, mAuxPool.data_ptr()); + EXPECT_EQ(sparse.sliding_window_kv_cache_pool, mPrimaryPool.data_ptr()); +} + +TEST_P(FmhaSparseParamsTest, DynamicSparseWithoutAuxiliaryPoolUsesOnlySlidingWindowPool) +{ + mParams.fwd.sparse_runtime_params.sparse_attn_kv_lens = torch::tensor({2}, torch::kInt32); + + auto const sparse = mParams.getSparseParams(); + + EXPECT_EQ(sparse.sparse_kv_cache_pool, nullptr); + EXPECT_EQ(sparse.sliding_window_kv_cache_pool, mPrimaryPool.data_ptr()); +} + +INSTANTIATE_TEST_SUITE_P(PoolPointerLayouts, FmhaSparseParamsTest, ::testing::Bool()); + +} // namespace From fd1c0ba8b6b24a76823e732672623500bdd763e8 Mon Sep 17 00:00:00 2001 From: Guiju Zhang Date: Mon, 2 Nov 2026 11:20:24 -0800 Subject: [PATCH 09/11] [TRTLLM-16282][perf] Enable THP by default for host KV pools Request huge-page backing for host KV allocations by default, while retaining TLLM_KV_CACHE_MANAGER_V2_THP=0 as an opt-out. Update the unit test for the unset-variable policy. Validation: 19 host-memory tests and 14 host-backing tests passed on the matching native build. Signed-off-by: Guiju Zhang --- .../batch_manager/kv_cache_manager_v2/utils/hostMem.cpp | 9 ++------- .../batch_manager/kvCacheManagerV2HostMemTest.cpp | 2 +- 2 files changed, 3 insertions(+), 8 deletions(-) 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 532b86d43a24..5c26d92273c0 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 @@ -95,14 +95,9 @@ void resizeFile(int fd, size_t newSize) bool hostUseThp() { - // Off unless asked for. MADV_HUGEPAGE opts a region's faults into direct - // compaction wherever /sys/kernel/mm/transparent_hugepage/defrag is - // "madvise" or "always". On a node whose free memory has fragmented below - // the 2 MiB order that costs far more than the huge pages save, and - // fragmentation accumulates with uptime, so the penalty grows while the - // benefit does not. + // Enable transparent huge pages by default. Set TLLM_KV_CACHE_MANAGER_V2_THP=0 to opt out. char const* value = std::getenv("TLLM_KV_CACHE_MANAGER_V2_THP"); - return value != nullptr && std::string_view(value) == "1"; + return value == nullptr || std::string_view(value) == "1"; } size_t hostTotalMemory() diff --git a/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2HostMemTest.cpp b/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2HostMemTest.cpp index e726bd5c4644..a14b805032cc 100644 --- a/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2HostMemTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2HostMemTest.cpp @@ -110,7 +110,7 @@ TEST(KvCacheManagerV2HostMemTest, SelectsConfiguredPageMode) EXPECT_EQ(gCapturedAdvice, MADV_NOHUGEPAGE); ScopedEnv defaultThp("TLLM_KV_CACHE_MANAGER_V2_THP", std::nullopt); - EXPECT_FALSE(hostUseThp()); + EXPECT_TRUE(hostUseThp()); { ScopedEnv enableThp("TLLM_KV_CACHE_MANAGER_V2_THP", "1"); EXPECT_TRUE(hostUseThp()); From b10d99c20198bc3e2854c86d71429617572df361 Mon Sep 17 00:00:00 2001 From: Guiju Zhang Date: Mon, 2 Nov 2026 11:21:46 -0800 Subject: [PATCH 10/11] [TRTLLM-16282][fix] Size sparse-offload pools at decode block boundaries Sample the first token of the last logical block when deriving default sparse-offload pool ratios. This retains full-history page counts while reserving the peak sliding-window footprint across decode block boundaries. Validation: 32 focused sizing tests passed. Six regression cases fail against the original code, including the DeepSeek V4 33024-token geometry. Signed-off-by: Guiju Zhang --- .../kv_cache/kv_cache_manager_v2.py | 6 ++ .../kv_cache/test_kvcm2_integration.py | 67 +++++++++++++++++-- 2 files changed, 66 insertions(+), 7 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index 616908575bed..888a09163cde 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 @@ -3913,6 +3913,12 @@ def _build_base_config( # Size its step for the length offload is sized for, which is # max_seq_len when avg_seq_len is not set. typical_seq_len = self._get_sparse_offload_seq_len(kv_cache_config) + # Use the first token in the last ledger block: full-history pools + # still need the same number of pages, while sliding windows span + # their maximum number of pages as decode crosses a block boundary. + typical_seq_len = ( + (typical_seq_len - 1) // self._ledger_tokens_per_block + ) * self._ledger_tokens_per_block + 1 if typical_seq_len is not None: # Model enough generation requests to fill the resident request diff --git a/tests/unittest/_torch/pyexecutor/kv_cache/test_kvcm2_integration.py b/tests/unittest/_torch/pyexecutor/kv_cache/test_kvcm2_integration.py index 764f4f73586a..abd0c6cc4317 100644 --- a/tests/unittest/_torch/pyexecutor/kv_cache/test_kvcm2_integration.py +++ b/tests/unittest/_torch/pyexecutor/kv_cache/test_kvcm2_integration.py @@ -1404,11 +1404,14 @@ def test_extend_swa_windows_for_reuse_preserves_non_attention_windows() -> None: ) == [None, 0, recurrent_states, 6, None] -def test_pool_ratio_overrides_constraints() -> None: +@pytest.mark.cpu_only +@pytest.mark.parametrize("offload", [False, True]) +def test_pool_ratio_overrides_constraints(offload: bool) -> None: config = _make_cache_config_for_test( KvCacheConfig(pool_ratio=[1.0], avg_seq_len=256, host_cache_size=0), max_batch_size=3, max_num_tokens=2048, + sparse_kv_offload=SparseKvCacheOffloadConfig() if offload else None, ) assert config.initial_pool_ratio == pytest.approx([1.0]) @@ -1446,8 +1449,9 @@ def test_avg_seq_len_builds_context_warmup_constraint() -> None: assert config.constraints == [BatchDesc([KVCacheDesc(capacity=2048, history_length=0)])] +@pytest.mark.cpu_only @pytest.mark.parametrize( - ("avg_seq_len", "capacity"), [(1024, 1024), (None, 2048)], ids=["avg_seq_len", "max_seq_len"] + ("avg_seq_len", "capacity"), [(1024, 897), (None, 1921)], ids=["avg_seq_len", "max_seq_len"] ) def test_sparse_kv_offload_typical_step_models_only_decoding( avg_seq_len: int | None, capacity: int @@ -1457,20 +1461,69 @@ def test_sparse_kv_offload_typical_step_models_only_decoding( max_batch_size=3, max_seq_len=2048, max_num_tokens=2048, - max_draft_len=2, sparse_kv_offload=SparseKvCacheOffloadConfig(), ) - # A worker that offloads sparse history only decodes requests of avg_seq_len, or - # else of max_seq_len, and its warmup runs no context batch to hold room for. The - # manager splits the quota from this step, host-backed pools included. + # Decode-only sizing samples a block boundary within avg_seq_len or max_seq_len, + # with no context warmup batch. The native manager still derives the pool ratio. assert config.typical_step == BatchDesc( - [KVCacheDesc(capacity=capacity, history_length=capacity - 3)] * 3 + [KVCacheDesc(capacity=capacity, history_length=capacity - 1)] * 3 ) assert config.constraints == [] assert config.initial_pool_ratio is None +@pytest.mark.cpu_only +@pytest.mark.parametrize( + ("max_seq_len", "full_history_pages"), + [ + (1, 1), + (127, 1), + (128, 1), + (129, 2), + (255, 2), + (256, 2), + (257, 3), + (32768, 256), + (32769, 257), + (33024, 258), + (33025, 259), + ], +) +@pytest.mark.parametrize("num_sink_tokens", [0, 129]) +def test_sparse_kv_offload_pool_sizing_covers_window_boundaries( + max_seq_len: int, full_history_pages: int, num_sink_tokens: int +) -> None: + config = _make_cache_config_for_test( + KvCacheConfig(host_cache_size=0), + max_seq_len=max_seq_len, + sparse_kv_offload=SparseKvCacheOffloadConfig(), + ) + tokens_per_block = config.tokens_per_block + generation = config.typical_step.kv_caches[0] + assert 0 < generation.capacity <= max_seq_len + assert generation.history_length == generation.capacity - 1 + assert (generation.capacity + tokens_per_block - 1) // tokens_per_block == full_history_pages + + def required_pages(life_cycle: AttnLifeCycle, capacity: int) -> int: + stale = life_cycle.get_stale_range(capacity - 1, tokens_per_block) + return (capacity + tokens_per_block - 1) // tokens_per_block - len(stale) + + for window in (None, 1, 8, 128, 129, 256, 65536): + life_cycle = AttnLifeCycle.make( + window, num_sink_tokens if window is not None else 0, tokens_per_block + ) + # Native lifecycle accounting must cover every decode phase, including + # windows shorter than a page, non-aligned windows and retained sinks. + peak_pages = max( + required_pages(life_cycle, capacity) + for capacity in range(max(1, max_seq_len - tokens_per_block + 1), max_seq_len + 1) + ) + assert required_pages(life_cycle, generation.capacity) == peak_pages, window + if max_seq_len == 33024 and num_sink_tokens == 0 and window in (8, 128): + assert peak_pages == 2 + + def test_avg_seq_len_updates_typical_step() -> None: config = _make_cache_config_for_test( KvCacheConfig(avg_seq_len=256), From 50e1e45cc92406a1b53087a4724dcb71d352b83d Mon Sep 17 00:00:00 2001 From: Guiju Zhang Date: Mon, 2 Nov 2026 21:05:14 -0800 Subject: [PATCH 11/11] [None][fix] Restore sparse indexer defaults and refine offload coverage Preserve the standard CuTe DSL indexer default with sparse KV offload, while respecting explicit backend choices. Regenerate serving schemas and retain focused coverage for defaults, overrides, and serialization. Use batched GovReport summaries with resident ROUGE-L references to check GLM-5.2 FP8/NVFP4 KV and DeepSeek-V4 Flash NVFP4 KV offload. Require sustained batched decode and histories larger than the GPU working set, and update the B200 CI cases. Validation: 155 CPU schema tests and 24 focused configuration/schema cases passed. The GovReport YAML references come from saved TP2 GB300 GPU calibration and response replay; TP4 B200 CI has not been rerun. Signed-off-by: Guiju Zhang --- scripts/generate_trtllm_serve_schemas.py | 23 -- tensorrt_llm/llmapi/llm_args.py | 8 +- .../schemas/trtllm-serve-config.schema.json | 62 +---- .../trtllm-serve-disagg-config.schema.json | 62 +---- .../defs/accuracy/accuracy_core.py | 22 ++ .../defs/accuracy/references/gov_report.yaml | 38 +++ .../accuracy/test_disaggregated_serving.py | 245 +++++++++++++++--- .../test_lists/test-db/l0_dgx_b200.yml | 5 +- tests/unittest/llmapi/test_llm_args.py | 49 +--- .../test_trtllm_serve_config_schemas.py | 8 +- 10 files changed, 280 insertions(+), 242 deletions(-) create mode 100644 tests/integration/defs/accuracy/references/gov_report.yaml diff --git a/scripts/generate_trtllm_serve_schemas.py b/scripts/generate_trtllm_serve_schemas.py index 198cfb4cb05e..07bd42e8c8f9 100644 --- a/scripts/generate_trtllm_serve_schemas.py +++ b/scripts/generate_trtllm_serve_schemas.py @@ -104,7 +104,6 @@ def _schema_for(annotation: Any, parent: dict | None = None) -> dict: from pydantic import TypeAdapter schema = TypeAdapter(annotation).json_schema(schema_generator=_schema_generator()) - _apply_sparse_indexer_defaults(schema) if parent is not None: definitions = parent.setdefault("$defs", {}) for name, definition in schema.pop("$defs", {}).items(): @@ -114,28 +113,6 @@ def _schema_for(annotation: Any, parent: dict | None = None) -> dict: return schema -def _apply_sparse_indexer_defaults(schema: dict) -> None: - """Describe the offload-dependent default resolved by the config validator.""" - for definition in [schema, *schema.get("$defs", {}).values()]: - if definition.get("title") not in { - "DeepSeekSparseAttentionConfig", - "DeepSeekV4SparseAttentionConfig", - }: - continue - field = "use_cute_dsl_paged_mqa_logits" - resident_default = definition["properties"][field].pop("default") - definition.setdefault("allOf", []).append( - { - "if": { - "properties": {"enable_kv_cache_offload": {"const": True}}, - "required": ["enable_kv_cache_offload"], - }, - "then": {"properties": {field: {"default": False}}}, - "else": {"properties": {field: {"default": resident_default}}}, - } - ) - - def _nullable(schema: dict) -> dict: return {"anyOf": [schema, {"type": "null"}]} diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index e6add33fe7df..2504e2bc86db 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -1383,9 +1383,7 @@ class DeepSeekSparseAttentionConfig(SeqLenAwareSparseAttentionConfig): default=True, description= "Whether to use the CuTe DSL paged MQA logits kernel on SM100-family " - "GPUs instead of C++ DeepGEMM. Defaults to False with KV cache " - "offload enabled and True otherwise; an explicit value overrides " - "this default. Falls back to DeepGEMM off the SM100 " + "GPUs instead of C++ DeepGEMM. Falls back to DeepGEMM off the SM100 " "family or when the cutlass-dsl package is missing.") q_split_threshold: int = Field( default=8192, @@ -1467,7 +1465,7 @@ class DeepSeekSparseAttentionConfig(SeqLenAwareSparseAttentionConfig): @model_validator(mode="after") def _validate_kv_cache_offload(self): - """Validate offload sizing and resolve the default indexer backend. + """Reject offload sizing without offload and a non-positive top-k. ``index_topk`` may still be None here; the checkpoint fills it later. """ @@ -1475,8 +1473,6 @@ def _validate_kv_cache_offload(self): if self.index_topk is not None and self.index_topk <= 0: raise ValueError( "KV cache offload requires a positive index_topk.") - if "use_cute_dsl_paged_mqa_logits" not in self.model_fields_set: - self.use_cute_dsl_paged_mqa_logits = False elif self.kv_cache_offload_config != SparseKvCacheOffloadConfig(): raise ValueError("kv_cache_offload_config requires " "enable_kv_cache_offload=True.") diff --git a/tensorrt_llm/schemas/trtllm-serve-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-config.schema.json index f70f0178a6e0..7769857d471a 100644 --- a/tensorrt_llm/schemas/trtllm-serve-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-config.schema.json @@ -1312,34 +1312,6 @@ }, "DeepSeekSparseAttentionConfig": { "additionalProperties": false, - "allOf": [ - { - "else": { - "properties": { - "use_cute_dsl_paged_mqa_logits": { - "default": true - } - } - }, - "if": { - "properties": { - "enable_kv_cache_offload": { - "const": true - } - }, - "required": [ - "enable_kv_cache_offload" - ] - }, - "then": { - "properties": { - "use_cute_dsl_paged_mqa_logits": { - "default": false - } - } - } - } - ], "description": "Configuration for DeepSeek Sparse Attention.", "properties": { "algorithm": { @@ -1480,7 +1452,8 @@ "type": "boolean" }, "use_cute_dsl_paged_mqa_logits": { - "description": "Whether to use the CuTe DSL paged MQA logits kernel on SM100-family GPUs instead of C++ DeepGEMM. Defaults to False with KV cache offload enabled and True otherwise; an explicit value overrides this default. Falls back to DeepGEMM off the SM100 family or when the cutlass-dsl package is missing.", + "default": true, + "description": "Whether to use the CuTe DSL paged MQA logits kernel on SM100-family GPUs instead of C++ DeepGEMM. Falls back to DeepGEMM off the SM100 family or when the cutlass-dsl package is missing.", "title": "Use Cute Dsl Paged Mqa Logits", "type": "boolean" }, @@ -1508,34 +1481,6 @@ }, "DeepSeekV4SparseAttentionConfig": { "additionalProperties": false, - "allOf": [ - { - "else": { - "properties": { - "use_cute_dsl_paged_mqa_logits": { - "default": true - } - } - }, - "if": { - "properties": { - "enable_kv_cache_offload": { - "const": true - } - }, - "required": [ - "enable_kv_cache_offload" - ] - }, - "then": { - "properties": { - "use_cute_dsl_paged_mqa_logits": { - "default": false - } - } - } - } - ], "description": "Configuration for DeepSeek-V4 Sparse Attention.", "properties": { "algorithm": { @@ -1684,7 +1629,8 @@ "type": "boolean" }, "use_cute_dsl_paged_mqa_logits": { - "description": "Whether to use the CuTe DSL paged MQA logits kernel on SM100-family GPUs instead of C++ DeepGEMM. Defaults to False with KV cache offload enabled and True otherwise; an explicit value overrides this default. Falls back to DeepGEMM off the SM100 family or when the cutlass-dsl package is missing.", + "default": true, + "description": "Whether to use the CuTe DSL paged MQA logits kernel on SM100-family GPUs instead of C++ DeepGEMM. Falls back to DeepGEMM off the SM100 family or when the cutlass-dsl package is missing.", "title": "Use Cute Dsl Paged Mqa Logits", "type": "boolean" }, diff --git a/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json index 95e43bb84588..33a0551922b7 100644 --- a/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json @@ -1312,34 +1312,6 @@ }, "DeepSeekSparseAttentionConfig": { "additionalProperties": false, - "allOf": [ - { - "else": { - "properties": { - "use_cute_dsl_paged_mqa_logits": { - "default": true - } - } - }, - "if": { - "properties": { - "enable_kv_cache_offload": { - "const": true - } - }, - "required": [ - "enable_kv_cache_offload" - ] - }, - "then": { - "properties": { - "use_cute_dsl_paged_mqa_logits": { - "default": false - } - } - } - } - ], "description": "Configuration for DeepSeek Sparse Attention.", "properties": { "algorithm": { @@ -1480,7 +1452,8 @@ "type": "boolean" }, "use_cute_dsl_paged_mqa_logits": { - "description": "Whether to use the CuTe DSL paged MQA logits kernel on SM100-family GPUs instead of C++ DeepGEMM. Defaults to False with KV cache offload enabled and True otherwise; an explicit value overrides this default. Falls back to DeepGEMM off the SM100 family or when the cutlass-dsl package is missing.", + "default": true, + "description": "Whether to use the CuTe DSL paged MQA logits kernel on SM100-family GPUs instead of C++ DeepGEMM. Falls back to DeepGEMM off the SM100 family or when the cutlass-dsl package is missing.", "title": "Use Cute Dsl Paged Mqa Logits", "type": "boolean" }, @@ -1508,34 +1481,6 @@ }, "DeepSeekV4SparseAttentionConfig": { "additionalProperties": false, - "allOf": [ - { - "else": { - "properties": { - "use_cute_dsl_paged_mqa_logits": { - "default": true - } - } - }, - "if": { - "properties": { - "enable_kv_cache_offload": { - "const": true - } - }, - "required": [ - "enable_kv_cache_offload" - ] - }, - "then": { - "properties": { - "use_cute_dsl_paged_mqa_logits": { - "default": false - } - } - } - } - ], "description": "Configuration for DeepSeek-V4 Sparse Attention.", "properties": { "algorithm": { @@ -1684,7 +1629,8 @@ "type": "boolean" }, "use_cute_dsl_paged_mqa_logits": { - "description": "Whether to use the CuTe DSL paged MQA logits kernel on SM100-family GPUs instead of C++ DeepGEMM. Defaults to False with KV cache offload enabled and True otherwise; an explicit value overrides this default. Falls back to DeepGEMM off the SM100 family or when the cutlass-dsl package is missing.", + "default": true, + "description": "Whether to use the CuTe DSL paged MQA logits kernel on SM100-family GPUs instead of C++ DeepGEMM. Falls back to DeepGEMM off the SM100 family or when the cutlass-dsl package is missing.", "title": "Use Cute Dsl Paged Mqa Logits", "type": "boolean" }, diff --git a/tests/integration/defs/accuracy/accuracy_core.py b/tests/integration/defs/accuracy/accuracy_core.py index 405edd7c2f23..b47c19a3a0a5 100644 --- a/tests/integration/defs/accuracy/accuracy_core.py +++ b/tests/integration/defs/accuracy/accuracy_core.py @@ -775,6 +775,28 @@ class LongBenchV1(AccuracyTask): apply_chat_template=True) +class GovReport(AccuracyTask): + """ROUGE-L references for complete reports in the batched serving test. + + The serving test selects reports using the model's chat tokenizer and + scores responses with LongBench's ``longbench_gov_report`` task. + """ + + DATASET = "gov_report" + DATASET_DIR = LongBenchV1.DATASET_DIR + METRIC_NAME = "ROUGE-L" + ALPHA = 0.05 + BETA = 0.2 + SIGMA = 50.0 + + NUM_SAMPLES = 32 + MIN_INPUT_LEN = 6144 + MAX_INPUT_LEN = 8192 + MAX_OUTPUT_LEN = 512 + MAX_BATCH_SIZE = 16 + RANDOM_SEED = 0 + + class LlmapiAccuracyTestHarness: # Model MODEL_NAME = None diff --git a/tests/integration/defs/accuracy/references/gov_report.yaml b/tests/integration/defs/accuracy/references/gov_report.yaml new file mode 100644 index 000000000000..12d0202c3ce6 --- /dev/null +++ b/tests/integration/defs/accuracy/references/gov_report.yaml @@ -0,0 +1,38 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# LongBench v1 GovReport ROUGE-L F1 on a 0-100 scale, with offload disabled. +# Select the first 32 complete reports with 6144-8192 chat-template tokens +# after shuffling with seed 0, using each model's tokenizer. No truncation. +# Greedy generation, seed 0, thinking disabled, up to 512 tokens, EOS allowed. +# Measured on 4 GB300 GPUs: one TP2/EP2 context worker and one TP2/EP2 +# generation worker, max_batch_size=16, concurrency=32, no speculative +# decoding, no block reuse, and autotuning disabled. GLM CI uses TP4/EP4. +# Explicit thresholds allow one absolute ROUGE-L point below the resident +# score for numerical and scheduling variation on this fixed sample set. +zai-org/GLM-5.2: + - quant_algo: NVFP4 + kv_cache_quant_algo: FP8 + extra_acc_spec: "input=6144-8192,output=512,seed=0,thinking=off" + index_topk: 2048 + unit_tokens: 1 + num_samples: 32 + rouge-l: 31.620 + threshold: 30.620 + - quant_algo: NVFP4 + kv_cache_quant_algo: NVFP4 + extra_acc_spec: "input=6144-8192,output=512,seed=0,thinking=off" + index_topk: 2048 + unit_tokens: 1 + num_samples: 32 + rouge-l: 31.831 + threshold: 30.831 +deepseek-ai/DeepSeek-V4-Flash: + # NVFP4 routed experts; the checkpoint declares MIXED_PRECISION. + - quant_algo: MIXED_PRECISION + kv_cache_quant_algo: NVFP4 + extra_acc_spec: "input=6144-8192,output=512,seed=0,thinking=off" + index_topk: 512 + unit_tokens: 4 + num_samples: 32 + rouge-l: 32.814 + threshold: 31.814 diff --git a/tests/integration/defs/accuracy/test_disaggregated_serving.py b/tests/integration/defs/accuracy/test_disaggregated_serving.py index b2f7561b0a9e..d91ed7e7052c 100644 --- a/tests/integration/defs/accuracy/test_disaggregated_serving.py +++ b/tests/integration/defs/accuracy/test_disaggregated_serving.py @@ -15,6 +15,7 @@ import concurrent import contextlib +import copy import functools import itertools import json @@ -37,19 +38,21 @@ from datasets import load_dataset from defs.common import wait_for_reported_addr +from tensorrt_llm.evaluate.lm_eval import (LmEvalEvaluator, + _replace_fuzzywuzzy_with_rapidfuzz) from tensorrt_llm.executor.result import GenerationResultBase from tensorrt_llm.llmapi import CompletionOutput, RequestOutput, SamplingParams from tensorrt_llm.llmapi.llm_args import (DSparkDecodingConfig, Eagle3DecodingConfig, LlmArgs, MTPDecodingConfig) -from tensorrt_llm.llmapi.tokenizer import load_hf_tokenizer +from tensorrt_llm.llmapi.tokenizer import DeepseekV4Tokenizer, load_hf_tokenizer from tensorrt_llm.quantization import QuantAlgo from ..conftest import (get_device_count, llm_models_root, parametrize_with_ids, skip_no_hopper, skip_no_rubin, skip_pre_blackwell, skip_pre_hopper) from ..trt_test_alternative import popen -from .accuracy_core import (GSM8K, MMLU, GSM8KInferenceX, +from .accuracy_core import (GSM8K, MMLU, GovReport, GSM8KInferenceX, LlmapiAccuracyTestHarness, acceptance_length_from_iteration_stats, assert_acceptance_length, get_accuracy_task) @@ -97,8 +100,8 @@ def pick_distinct_free_ports(count: int) -> List[int]: sock.close() -def compute_disagg_acceptance_length(serve_url: str) -> float: - """Acceptance length from the generation workers' iteration stats. +def _get_disagg_generation_stats(serve_url: str) -> list[dict[str, Any]]: + """Read the generation workers' iteration stats. The disaggregated server does not proxy ``/metrics``, so the workers are resolved through ``/cluster_info`` -- the only authoritative source for @@ -119,13 +122,19 @@ def compute_disagg_acceptance_length(serve_url: str) -> float: stats = [] for worker in gen_workers: - # /metrics drains the buffer, so call it exactly once per worker, - # after the workload has finished. + # /metrics drains the buffer. Read it before a workload to discard + # startup statistics, then after the workload to collect its iterations. response = requests.get( f"http://{worker['host']}:{worker['port']}/metrics", timeout=60) response.raise_for_status() stats.extend(response.json()) - return acceptance_length_from_iteration_stats(stats) + return stats + + +def compute_disagg_acceptance_length(serve_url: str) -> float: + """Acceptance length from the generation workers' iteration stats.""" + return acceptance_length_from_iteration_stats( + _get_disagg_generation_stats(serve_url)) # Timeout for the entire test @@ -676,6 +685,169 @@ def _show_kvcache_time(kv_cache_perf_dir, max_lines=100): "the required transfer path was not exercised") +def _test_sparse_offload_long_context( + disaggregated_server_config: dict[str, Any], + ctx_server_config: dict[str, Any], + gen_server_config: dict[str, Any], + model_path: str, + *, + model_name: str, + quant_algo: str, + index_topk: int, + unit_tokens: int, + server_waiting_timeout: int = 1200, +) -> None: + """Compare batched GovReport quality against the resident YAML reference.""" + min_input_len = GovReport.MIN_INPUT_LEN + max_input_len = GovReport.MAX_INPUT_LEN + max_batch_size = GovReport.MAX_BATCH_SIZE + concurrency = 32 + output_len = GovReport.MAX_OUTPUT_LEN + random_seed = GovReport.RANDOM_SEED + tokens_per_block = 128 + reference = GovReport(model_name).get_hypothesis_testing_params( + quant_algo=quant_algo, + kv_cache_quant_algo=gen_server_config["kv_cache_config"] + ["dtype"].upper(), + extra_acc_spec=(f"input={min_input_len}-{max_input_len}," + f"output={output_len},seed={random_seed},thinking=off"), + index_topk=index_topk, + unit_tokens=unit_tokens, + num_samples=GovReport.NUM_SAMPLES) + num_samples = reference.num_samples + assert min_input_len >= 3 * index_topk * unit_tokens + if gen_server_config["sparse_attention_config"][ + "algorithm"] == "deepseek_v4": + tokenizer = DeepseekV4Tokenizer.from_pretrained(model_path) + else: + tokenizer = load_hf_tokenizer(model_path) + # Use the same local dataset and task prompt/metric as the LongBench suite. + with _replace_fuzzywuzzy_with_rapidfuzz(): + evaluator = LmEvalEvaluator("longbench_gov_report", + dataset_path=GovReport.DATASET_DIR, + random_seed=random_seed) + task = evaluator.task_dict["longbench_gov_report"] + samples = [] + prompts = [] + for sample in task.test_docs(): + prompt = tokenizer.apply_chat_template( + [{ + "role": "user", + "content": task.doc_to_text(sample) + }], + tokenize=True, + return_dict=False, + add_generation_prompt=True, + enable_thinking=False) + # Keep complete documents: truncation could remove the facts scored by + # the reference summary. Tokenize with each model's actual tokenizer. + if not min_input_len <= len(prompt) <= max_input_len: + continue + samples.append(sample) + prompts.append(prompt) + if len(samples) == num_samples: + break + assert len(samples) == num_samples, ( + f"Need {num_samples} complete GovReport prompts between {min_input_len} " + f"and {max_input_len} tokens, found {len(samples)}") + assert len({sample["_id"] for sample in samples}) == num_samples + + for config in (ctx_server_config, gen_server_config): + config.update(max_batch_size=max_batch_size, + max_seq_len=max(map(len, prompts)) + output_len, + enable_autotuner=False) + config["kv_cache_config"] = { + **config["kv_cache_config"], + "enable_block_reuse": False, + "use_kv_cache_manager_v2": True, + "tokens_per_block": tokens_per_block, + } + ctx_server_config.update(enable_chunked_prefill=True, max_num_tokens=8192) + gen_server_config.update(enable_iter_perf_stats=True, + iter_stats_max_iterations=-1, + max_stats_len=-1) + sparse_config = { + **gen_server_config["sparse_attention_config"], + "index_topk": index_topk, + # Match the indexer used to measure the resident YAML reference. + "use_cute_dsl_paged_mqa_logits": False, + } + ctx_server_config["sparse_attention_config"] = { + **sparse_config, + "enable_kv_cache_offload": False, + } + # At factor 1, the cache holds top-k plus at most one page of slack. + # The history cannot fit. Check the actual decode lengths below, since + # benchmark summaries are allowed to stop at EOS. + units_per_page = tokens_per_block // unit_tokens + cache_units = (index_topk // units_per_page + 1) * units_per_page + assert all(len(prompt) // unit_tokens > cache_units for prompt in prompts) + + gen_server_config["sparse_attention_config"] = { + **sparse_config, + "enable_kv_cache_offload": True, + "kv_cache_offload_config": { + "gpu_cache_factor": 1.0 + }, + } + # The launcher adds discovery and authentication settings in place. + with launch_disaggregated_llm( + copy.deepcopy(disaggregated_server_config), + copy.deepcopy(ctx_server_config), + copy.deepcopy(gen_server_config), + model_path, + server_waiting_timeout=server_waiting_timeout, + max_workers=concurrency) as llm, openai.OpenAI( + api_key="1234567890", + base_url=f"{llm.serve_url}/v1", + timeout=300, + max_retries=0) as client: + _get_disagg_generation_stats(llm.serve_url) + with ThreadPoolExecutor(max_workers=concurrency) as pool: + futures = [ + pool.submit(client.completions.create, + model=model_path, + prompt=prompt, + max_tokens=output_len, + temperature=0, + seed=random_seed, + extra_body={"detokenize": False}) + for prompt in prompts + ] + responses = [future.result() for future in futures] + stats = _get_disagg_generation_stats(llm.serve_url) + batched_steps = sum( + stat.get("inflightBatchingStats", {}).get("numGenRequests", 0) > 1 + for stat in stats) + assert batched_steps >= tokens_per_block, ( + f"Only {batched_steps} decode iterations had multiple requests; " + "the long-context test must exercise sustained batched decode") + + scores = [] + long_decodes = 0 + for sample, prompt, response in zip(samples, prompts, responses): + assert response.usage.prompt_tokens == len(prompt) + choice = response.choices[0] + assert choice.finish_reason in ("length", "stop") + token_ids = choice.model_dump()["token_ids"] + assert 0 < len(token_ids) <= output_len + assert response.usage.completion_tokens == len(token_ids) + long_decodes += ((len(token_ids) - 1) // unit_tokens + > cache_units - index_topk) + prediction = tokenizer.decode(token_ids, skip_special_tokens=True) + scores.append( + task.process_results(sample, [prediction])["rouge_score"]) + assert long_decodes >= max_batch_size, ( + f"Only {long_decodes} summaries decoded beyond the cache's " + "free slots; need at least one full batch") + score = 100 * sum(scores) / len(scores) + print(f"GovReport offload=True: n={num_samples}, " + f"ROUGE-L={score:.3f}, batched_steps={batched_steps}, " + f"long_decodes={long_decodes}") + print(reference.report(score)) + reference.assert_passing(score) + + def run_parallel_test(model_name: str, model_path: str, *, @@ -2275,16 +2447,20 @@ def test_nvfp4_nixl(self, use_kv_cache_manager_v2): max_draft_len=speculative_config["max_draft_len"]) run_accuracy_test(llm, self.MODEL_NAME, ["GSM8K"]) + @pytest.mark.timeout(1800) @pytest.mark.skip_less_device(8) - def test_kv_cache_offload(self): - """The generation server keeps the MLA latent history in host memory. + @parametrize_with_ids("kv_dtype", ["fp8", "nvfp4"]) + def test_kv_cache_offload(self, kv_dtype: str) -> None: + """GovReport summary quality is preserved with MLA latent history offloaded. Each decode step fetches the indexer-selected tokens into a GPU cache. Only a decode-only worker can offload, so the context server keeps its history resident. Offload takes no speculative decoding and no block reuse. """ + # Both servers use the same KV dtype for transfer and decode. kv_cache_config = { + "dtype": kv_dtype, "free_gpu_memory_fraction": 0.7, "enable_block_reuse": False, "use_kv_cache_manager_v2": True, @@ -2303,7 +2479,6 @@ def test_kv_cache_offload(self): "enable_chunked_prefill": True, "cuda_graph_config": None, "trust_remote_code": True, - "max_seq_len": 8192, "kv_cache_config": kv_cache_config, "moe_config": moe_config, "cache_transceiver_config": cache_transceiver_config, @@ -2316,10 +2491,6 @@ def test_kv_cache_offload(self): "disable_overlap_scheduler": False, "enable_chunked_prefill": True, "trust_remote_code": True, - # Each batch row holds a GPU cache for its selected history, about - # 180 MiB on GLM-5.2, so the batch stays small. - "max_batch_size": 16, - "max_seq_len": 8192, "kv_cache_config": kv_cache_config, "moe_config": moe_config, "cache_transceiver_config": cache_transceiver_config, @@ -2338,15 +2509,14 @@ def test_kv_cache_offload(self): "num_instances": 1 } } - with launch_disaggregated_llm(disaggregated_server_config, - ctx_server_config, - gen_server_config, - self.MODEL_PATH, - max_workers=16) as llm: - # The DuckLLM's bare LlmArgs takes NVFP4 from the model path; the - # checkpoint's FP8 KV cache selects the reference entry. - llm.args.quant_config.kv_cache_quant_algo = "FP8" - run_accuracy_test(llm, self.MODEL_NAME, ["GSM8K"]) + _test_sparse_offload_long_context(disaggregated_server_config, + ctx_server_config, + gen_server_config, + self.MODEL_PATH, + model_name=self.MODEL_NAME, + quant_algo="NVFP4", + index_topk=2048, + unit_tokens=1) @pytest.mark.timeout(DEFAULT_TEST_TIMEOUT) @@ -2462,11 +2632,10 @@ def test_gen_first(self): task = MMLU(self.MODEL_NAME) task.evaluate(llm, is_integration_test=True) - # The servers take up to an hour to start, before a full GSM8K run. - @pytest.mark.timeout(5400) + @pytest.mark.timeout(1800) @pytest.mark.skip_less_device(4) - def test_kv_cache_offload(self): - """The generation server keeps the ratio-4 compressed history in host memory. + def test_kv_cache_offload(self) -> None: + """GovReport summary quality is preserved with ratio-4 history offloaded. Each decode step fetches the indexer-selected units into a GPU cache. Only a decode-only worker can offload, so the context server keeps its @@ -2485,8 +2654,6 @@ def test_kv_cache_offload(self): "tensor_parallel_size": 2, "moe_expert_parallel_size": 2, "disable_overlap_scheduler": True, - "max_batch_size": DEEPSEEKV4_TEST_MAX_BATCH_SIZE, - "max_seq_len": 4096, "kv_cache_config": kv_cache_config, "cache_transceiver_config": cache_transceiver_config, } @@ -2495,8 +2662,6 @@ def test_kv_cache_offload(self): "moe_expert_parallel_size": 2, "enable_attention_dp": True, "disable_overlap_scheduler": True, - "max_batch_size": DEEPSEEKV4_TEST_MAX_BATCH_SIZE, - "max_seq_len": 4096, "moe_config": { "backend": "TRTLLM", }, @@ -2520,16 +2685,14 @@ def test_kv_cache_offload(self): "num_instances": 1 }, } - with launch_disaggregated_llm(disaggregated_server_config, - ctx_server_config, - gen_server_config, - self.MODEL_PATH, - server_waiting_timeout=3600) as llm: - # The DuckLLM's bare LlmArgs knows no checkpoint quantization; these - # select the reference entry. - llm.args.quant_config.quant_algo = QuantAlgo.FP8_BLOCK_SCALES - llm.args.quant_config.kv_cache_quant_algo = QuantAlgo.NVFP4 - run_accuracy_test(llm, self.MODEL_NAME, ["GSM8K"]) + _test_sparse_offload_long_context(disaggregated_server_config, + ctx_server_config, + gen_server_config, + self.MODEL_PATH, + model_name=self.MODEL_NAME, + quant_algo="MIXED_PRECISION", + index_topk=512, + unit_tokens=4) @pytest.mark.timeout(14400) diff --git a/tests/integration/test_lists/test-db/l0_dgx_b200.yml b/tests/integration/test_lists/test-db/l0_dgx_b200.yml index 3ec7aee77a47..1dcbb0f1baf1 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_b200.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_b200.yml @@ -215,8 +215,9 @@ l0_dgx_b200: - accuracy/test_disaggregated_serving.py::TestNemotron3Super120B::test_auto_dtype[mtp_nextn=0-block_reuse=False-use_py_transceiver=True] TIMEOUT (60) - accuracy/test_disaggregated_serving.py::TestNemotron3Super120B::test_auto_dtype[mtp_nextn=3-block_reuse=True-use_py_transceiver=True] TIMEOUT (60) - accuracy/test_disaggregated_serving.py::TestGLM52NVFP4::test_nvfp4_nixl[cache_mgr_v2] TIMEOUT (60) - - accuracy/test_disaggregated_serving.py::TestGLM52NVFP4::test_kv_cache_offload TIMEOUT (60) - - accuracy/test_disaggregated_serving.py::TestDeepSeekV4Flash::test_kv_cache_offload TIMEOUT (90) + - accuracy/test_disaggregated_serving.py::TestGLM52NVFP4::test_kv_cache_offload[kv_dtype=fp8] TIMEOUT (30) + - accuracy/test_disaggregated_serving.py::TestGLM52NVFP4::test_kv_cache_offload[kv_dtype=nvfp4] TIMEOUT (30) + - accuracy/test_disaggregated_serving.py::TestDeepSeekV4Flash::test_kv_cache_offload TIMEOUT (30) # DeepSeek-V4 EPLB pre-merge sanity (uncomment once DeepSeek-V4-Flash/Flash-Base # checkpoints are staged under llm_models_root()). # - accuracy/test_llm_api_pytorch.py::TestDeepSeekV4Flash::test_nvfp4_4gpus_static_eplb[moe_backend=TRTLLM] TIMEOUT (120) diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index 68a09f00a1b1..69fd3b5ecabc 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -5062,7 +5062,7 @@ def test_default_override_and_round_trip( kwargs = {"enable_kv_cache_offload": enabled} if requested is not None: kwargs["use_cute_dsl_paged_mqa_logits"] = requested - expected = not enabled if requested is None else requested + expected = True if requested is None else requested config = config_cls(**kwargs) for resolved in ( @@ -5077,53 +5077,6 @@ def test_default_override_and_round_trip( assert resolved.to_sparse_metadata_params( ).use_cute_dsl_paged_mqa_logits is expected - @pytest.mark.parametrize("enabled", [False, True]) - def test_yaml_bool_coercion( - self, - config_cls: type[llm_args_mod.DeepSeekSparseAttentionConfig], - enabled: bool, - ) -> None: - config = config_cls(enable_kv_cache_offload=str(enabled).lower()) - assert config.use_cute_dsl_paged_mqa_logits is not enabled - - -@pytest.mark.cpu_only -@pytest.mark.parametrize("enabled", [False, True]) -@pytest.mark.parametrize("requested", [None, False, True]) -def test_glm_model_config_preserves_indexer_choice( - tmp_path: Path, - monkeypatch: pytest.MonkeyPatch, - enabled: bool, - requested: bool | None, -) -> None: - checkpoint = PretrainedConfig( - architectures=["GlmMoeDsaForCausalLM"], - num_hidden_layers=7, - index_n_heads=32, - index_head_dim=128, - index_topk=2048, - ) - monkeypatch.setattr( - "tensorrt_llm._torch.model_config.load_pretrained_config", - lambda *args, **kwargs: checkpoint) - kwargs = {"enable_kv_cache_offload": enabled} - if requested is not None: - kwargs["use_cute_dsl_paged_mqa_logits"] = requested - config = llm_args_mod.DeepSeekSparseAttentionConfig(**kwargs) - - resolved = ModelConfig.from_pretrained( - str(tmp_path), - sparse_attention_config=config, - attn_backend="TRTLLM", - moe_backend="TRTLLM", - ).sparse_attention_config - - expected = not enabled if requested is None else requested - assert resolved.use_cute_dsl_paged_mqa_logits is expected - assert resolved.to_sparse_params().use_cute_dsl_paged_mqa_logits is expected - assert resolved.to_sparse_metadata_params( - ).use_cute_dsl_paged_mqa_logits is expected - @pytest.mark.cpu_only class TestDeepSeekV4SparseAttentionConfig: diff --git a/tests/unittest/llmapi/test_trtllm_serve_config_schemas.py b/tests/unittest/llmapi/test_trtllm_serve_config_schemas.py index e1c223158498..509f5ab3f0ad 100644 --- a/tests/unittest/llmapi/test_trtllm_serve_config_schemas.py +++ b/tests/unittest/llmapi/test_trtllm_serve_config_schemas.py @@ -429,12 +429,8 @@ def test_sparse_offload_indexer_default( schema = schemas[filename] definition = schema["$defs"][config_name] field = "use_cute_dsl_paged_mqa_logits" - assert "default" not in definition["properties"][field] - condition = definition["allOf"][0] - branch = "then" if Draft202012Validator(condition["if"]).is_valid(sparse_config) else "else" - expected = enabled is not True - assert condition[branch]["properties"][field]["default"] is expected - assert config_cls.model_validate(sparse_config).use_cute_dsl_paged_mqa_logits is expected + assert definition["properties"][field]["default"] is True + assert config_cls.model_validate(sparse_config).use_cute_dsl_paged_mqa_logits is True validator = Draft202012Validator(schema) assert validator.is_valid({"sparse_attention_config": sparse_config}) for value in (False, True):