diff --git a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h index 4ce5fa94b131..21aa57704ff3 100644 --- a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h +++ b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h @@ -379,6 +379,17 @@ struct KvCacheIterationStats // to GPU during their stay at that tier (i.e. fully dropped from the hierarchy). SizeType32 iterHostDroppedBlocks{0}; std::size_t iterHostDroppedBytes{0}; + + // KVCacheManagerV2 SWA scratch reuse. Scratch blocks take no per-request KV page and are + // therefore excluded from iterAlloc*; these are the attribution for that exclusion. + // Always 0 for the V1 manager, which has no scratch reuse. + // + // iterScratchBlocks is a COUNT (blocks served from scratch this iteration) and is additive + // across iterations. iterScratchSlotsInUse is a GAUGE (slots concurrently occupied) and is + // NOT additive across iterations -- see KVCacheIterationStatsDelta in + // kv_cache_manager_v2/stats.h for the full contract. + SizeType32 iterScratchBlocks{0}; + SizeType32 iterScratchSlotsInUse{0}; }; // Basic building block of a paged KV cache - a single diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp index 29eacc617d11..0a8c27d21d9c 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp @@ -717,7 +717,8 @@ void KvCache::_recordDroppedPages(std::vector> const& pages, Cac } void KvCache::_recordResizePendingAllocations(BlockOrdinal blockBegin, BlockOrdinal blockEnd, - TypedVec> const& excludedRanges, bool countAsGeneration) + TypedVec> const& excludedRanges, bool countAsGeneration, + bool excludedIsScratch) { bool const recordManagerStats = _shouldRecordManagerStats(); bool const recordRequestStats = _shouldRecordRequestStats(); @@ -743,6 +744,13 @@ void KvCache::_recordResizePendingAllocations(BlockOrdinal blockBegin, BlockOrdi changed |= mPendingStats.recordAllocationRange(lifeCycle, secondBegin, blockEnd, mBeamWidth.value(), !countAsGeneration, countAsGeneration, recordManagerStats, recordRequestStats); } + if (excludedIsScratch) + { + // The piece dropped between the two recorded outer ranges is exactly the scratch + // block count for this lifecycle. + _recordScratchIterationStats( + lifeCycle, intersect(excluded, HalfOpenRange{blockBegin, blockEnd}).length()); + } } if (changed) { @@ -750,6 +758,18 @@ void KvCache::_recordResizePendingAllocations(BlockOrdinal blockBegin, BlockOrdi } } +void KvCache::_recordScratchIterationStats(LifeCycleId lifeCycle, int64_t numBlocks) +{ + if (numBlocks <= 0) + { + return; + } + KVCacheIterationStatsDelta iterationStats; + iterationStats.iterScratchBlocks = numBlocks; + iterationStats.iterScratchSlotsInUse = static_cast(mScratchSlots[lifeCycle].size()); + _recordDirectIterationStats(lifeCycle, iterationStats); +} + void KvCache::_subtractPendingAllocationRange(BlockOrdinal blockBegin, BlockOrdinal blockEnd) { if (mPendingStats.subtractAllocationRange(blockBegin, blockEnd)) @@ -1207,7 +1227,8 @@ bool KvCache::resize(std::optional capacity, std::optional historyLeng // Scratch and stale blocks do not consume per-request KV pages, so exclude them from allocation stats. auto const& excludedRanges = enableScratch ? scratchRanges : staleRanges; - _recordResizePendingAllocations(oldNumBlocks, newNumBlocks, excludedRanges, recordGenerationAllocStats); + _recordResizePendingAllocations( + oldNumBlocks, newNumBlocks, excludedRanges, recordGenerationAllocStats, enableScratch); // Resize page index buffers. TLLM_CHECK_DEBUG(std::all_of(mBasePageIndices.begin(), mBasePageIndices.end(), diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.h index 4192e7f3d035..23660aad2ff5 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.h @@ -492,7 +492,9 @@ class KvCache : public std::enable_shared_from_this void _recordDroppedPages(std::vector> const& pages, CacheLevel cacheLevel); void _refreshGenerationAllocReady(); void _recordResizePendingAllocations(BlockOrdinal blockBegin, BlockOrdinal blockEnd, - TypedVec> const& excludedRanges, bool countAsGeneration); + TypedVec> const& excludedRanges, bool countAsGeneration, + bool excludedIsScratch = false); + void _recordScratchIterationStats(LifeCycleId lifeCycle, int64_t numBlocks); void _subtractPendingAllocationRange(BlockOrdinal blockBegin, BlockOrdinal blockEnd); static bool _hasReuseSource(BlockPage const& page); void _increaseCapacity(BlockOrdinal newNumBlocks, int newHistoryLength); diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/stats.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/stats.h index e81e936f6e30..8da97c1d2789 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/stats.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/stats.h @@ -89,6 +89,24 @@ struct KVCacheIterationStatsDelta int64_t iterIntraDeviceCopyBytes = 0; int64_t iterHostDroppedBlocks = 0; int64_t iterHostDroppedBytes = 0; + // SWA scratch reuse attribution. Scratch blocks are excluded from iterAlloc* by design, so + // these are the only accounting for the saving: iterAllocNewBlocks drops by exactly + // iterScratchBlocks. + // + // The two fields have DIFFERENT semantics, which the names carry deliberately: + // iterScratchBlocks - a COUNT. Blocks served from shared scratch sub-pages during + // this iteration. Additive: summing over iterations gives the + // total number of blocks ever served from scratch. + // iterScratchSlotsInUse - a GAUGE, sampled per lifecycle as slots are recorded and + // therefore summed across lifecycles WITHIN one iteration to + // give the slots concurrently in use. It is reset every + // iteration along with the rest of the delta. Accumulating it + // ACROSS iterations is meaningless - it is an occupancy, not a + // flow. Consumers that sum every int field indiscriminately + // will produce a nonsense number here; the "InUse" suffix is + // the contract. + int64_t iterScratchBlocks = 0; + int64_t iterScratchSlotsInUse = 0; void add(KVCacheIterationStatsDelta const& other) noexcept { @@ -107,6 +125,8 @@ struct KVCacheIterationStatsDelta iterIntraDeviceCopyBytes += other.iterIntraDeviceCopyBytes; iterHostDroppedBlocks += other.iterHostDroppedBlocks; iterHostDroppedBytes += other.iterHostDroppedBytes; + iterScratchBlocks += other.iterScratchBlocks; + iterScratchSlotsInUse += other.iterScratchSlotsInUse; } void subtract(KVCacheIterationStatsDelta const& other) noexcept @@ -126,6 +146,8 @@ struct KVCacheIterationStatsDelta iterIntraDeviceCopyBytes -= other.iterIntraDeviceCopyBytes; iterHostDroppedBlocks -= other.iterHostDroppedBlocks; iterHostDroppedBytes -= other.iterHostDroppedBytes; + iterScratchBlocks -= other.iterScratchBlocks; + iterScratchSlotsInUse -= other.iterScratchSlotsInUse; } void clear() noexcept @@ -144,7 +166,8 @@ struct KVCacheIterationStatsDelta && iterFullReusedBlocks == 0 && iterPartialReusedBlocks == 0 && iterMissedBlocks == 0 && iterGenAllocBlocks == 0 && iterOnboardBlocks == 0 && iterOnboardBytes == 0 && iterOffloadBlocks == 0 && iterOffloadBytes == 0 && iterIntraDeviceCopyBlocks == 0 && iterIntraDeviceCopyBytes == 0 - && iterHostDroppedBlocks == 0 && iterHostDroppedBytes == 0; + && iterHostDroppedBlocks == 0 && iterHostDroppedBytes == 0 && iterScratchBlocks == 0 + && iterScratchSlotsInUse == 0; } [[nodiscard]] double iterCacheHitRate() const noexcept @@ -168,7 +191,8 @@ struct KVCacheIterationStatsDelta && iterIntraDeviceCopyBlocks == other.iterIntraDeviceCopyBlocks && iterIntraDeviceCopyBytes == other.iterIntraDeviceCopyBytes && iterHostDroppedBlocks == other.iterHostDroppedBlocks - && iterHostDroppedBytes == other.iterHostDroppedBytes; + && iterHostDroppedBytes == other.iterHostDroppedBytes && iterScratchBlocks == other.iterScratchBlocks + && iterScratchSlotsInUse == other.iterScratchSlotsInUse; } }; diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp index 222f2e528cc3..87acdbd9cbe5 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp @@ -406,7 +406,9 @@ void tb::kv_cache_manager::KVCacheManagerBindings::initBindings(nb::module_& m) .def_rw("iter_intra_device_copy_blocks", &tbk::KvCacheIterationStats::iterIntraDeviceCopyBlocks) .def_rw("iter_intra_device_copy_bytes", &tbk::KvCacheIterationStats::iterIntraDeviceCopyBytes) .def_rw("iter_host_dropped_blocks", &tbk::KvCacheIterationStats::iterHostDroppedBlocks) - .def_rw("iter_host_dropped_bytes", &tbk::KvCacheIterationStats::iterHostDroppedBytes); + .def_rw("iter_host_dropped_bytes", &tbk::KvCacheIterationStats::iterHostDroppedBytes) + .def_rw("iter_scratch_blocks", &tbk::KvCacheIterationStats::iterScratchBlocks) + .def_rw("iter_scratch_slots_in_use", &tbk::KvCacheIterationStats::iterScratchSlotsInUse); nb::class_(m, "BlockKey") .def(nb::init<>()) diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp index 91aad225b38b..636504b6e939 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp @@ -629,7 +629,9 @@ static std::string iterationStatsDeltaRepr(kv::KVCacheIterationStatsDelta const& << ", iter_intra_device_copy_blocks=" << stats.iterIntraDeviceCopyBlocks << ", iter_intra_device_copy_bytes=" << stats.iterIntraDeviceCopyBytes << ", iter_host_dropped_blocks=" << stats.iterHostDroppedBlocks - << ", iter_host_dropped_bytes=" << stats.iterHostDroppedBytes << ')'; + << ", iter_host_dropped_bytes=" << stats.iterHostDroppedBytes + << ", iter_scratch_blocks=" << stats.iterScratchBlocks + << ", iter_scratch_slots_in_use=" << stats.iterScratchSlotsInUse << ')'; return stream.str(); } @@ -1179,7 +1181,7 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) int64_t iterMissedBlocks, int64_t iterGenAllocBlocks, int64_t iterOnboardBlocks, int64_t iterOnboardBytes, int64_t iterOffloadBlocks, int64_t iterOffloadBytes, int64_t iterIntraDeviceCopyBlocks, int64_t iterIntraDeviceCopyBytes, int64_t iterHostDroppedBlocks, - int64_t iterHostDroppedBytes) + int64_t iterHostDroppedBytes, int64_t iterScratchBlocks, int64_t iterScratchSlotsInUse) { new (self) kv::KVCacheIterationStatsDelta{}; self->iterAllocTotalBlocks = iterAllocTotalBlocks; @@ -1197,6 +1199,8 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) self->iterIntraDeviceCopyBytes = iterIntraDeviceCopyBytes; self->iterHostDroppedBlocks = iterHostDroppedBlocks; self->iterHostDroppedBytes = iterHostDroppedBytes; + self->iterScratchBlocks = iterScratchBlocks; + self->iterScratchSlotsInUse = iterScratchSlotsInUse; }, nb::arg("iter_alloc_total_blocks") = 0, nb::arg("iter_alloc_new_blocks") = 0, nb::arg("iter_reused_blocks") = 0, nb::arg("iter_full_reused_blocks") = 0, @@ -1204,7 +1208,8 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) nb::arg("iter_gen_alloc_blocks") = 0, nb::arg("iter_onboard_blocks") = 0, nb::arg("iter_onboard_bytes") = 0, nb::arg("iter_offload_blocks") = 0, nb::arg("iter_offload_bytes") = 0, nb::arg("iter_intra_device_copy_blocks") = 0, nb::arg("iter_intra_device_copy_bytes") = 0, - nb::arg("iter_host_dropped_blocks") = 0, nb::arg("iter_host_dropped_bytes") = 0) + nb::arg("iter_host_dropped_blocks") = 0, nb::arg("iter_host_dropped_bytes") = 0, + nb::arg("iter_scratch_blocks") = 0, nb::arg("iter_scratch_slots_in_use") = 0) .def_rw("iter_alloc_total_blocks", &kv::KVCacheIterationStatsDelta::iterAllocTotalBlocks) .def_rw("iter_alloc_new_blocks", &kv::KVCacheIterationStatsDelta::iterAllocNewBlocks) .def_rw("iter_reused_blocks", &kv::KVCacheIterationStatsDelta::iterReusedBlocks) @@ -1220,6 +1225,8 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) .def_rw("iter_intra_device_copy_bytes", &kv::KVCacheIterationStatsDelta::iterIntraDeviceCopyBytes) .def_rw("iter_host_dropped_blocks", &kv::KVCacheIterationStatsDelta::iterHostDroppedBlocks) .def_rw("iter_host_dropped_bytes", &kv::KVCacheIterationStatsDelta::iterHostDroppedBytes) + .def_rw("iter_scratch_blocks", &kv::KVCacheIterationStatsDelta::iterScratchBlocks) + .def_rw("iter_scratch_slots_in_use", &kv::KVCacheIterationStatsDelta::iterScratchSlotsInUse) .def("add", &kv::KVCacheIterationStatsDelta::add, nb::arg("other")) .def("subtract", &kv::KVCacheIterationStatsDelta::subtract, nb::arg("other")) .def("clear", &kv::KVCacheIterationStatsDelta::clear) @@ -1233,7 +1240,7 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) "iter_alloc_new_blocks", "iter_reused_blocks", "iter_full_reused_blocks", "iter_partial_reused_blocks", "iter_missed_blocks", "iter_gen_alloc_blocks", "iter_onboard_blocks", "iter_onboard_bytes", "iter_offload_blocks", "iter_offload_bytes", "iter_intra_device_copy_blocks", "iter_intra_device_copy_bytes", - "iter_host_dropped_blocks", "iter_host_dropped_bytes"); + "iter_host_dropped_blocks", "iter_host_dropped_bytes", "iter_scratch_blocks", "iter_scratch_slots_in_use"); nb::class_(m, "SsmSnapshotIterationStatsDelta") .def( diff --git a/docs/source/developer-guide/telemetry.md b/docs/source/developer-guide/telemetry.md index 072cbd825c8a..a9a1eda22e7c 100644 --- a/docs/source/developer-guide/telemetry.md +++ b/docs/source/developer-guide/telemetry.md @@ -129,7 +129,7 @@ unset or when the safety sanitizer rejects the runtime value. | `kv_cache_config.enable_block_reuse` | `` | `value` | | | | `kv_cache_config.enable_kv_pool_rebalance` | `` | `value` | | | | `kv_cache_config.enable_partial_reuse` | `` | `value` | | | -| `kv_cache_config.enable_swa_scratch_reuse` | `` | `value` | | | +| `kv_cache_config.enable_swa_scratch_reuse` | `Union[bool, Literal['auto']]` | `value` | | `auto` | | `kv_cache_config.event_buffer_max_size` | `` | `value` | | | | `kv_cache_config.fp8_context_mla_kv_len_cap` | `Optional[int]` | `value` | | | | `kv_cache_config.free_gpu_memory_fraction` | `Optional[float]` | `value` | | | diff --git a/tensorrt_llm/_torch/attention_backend/flashinfer.py b/tensorrt_llm/_torch/attention_backend/flashinfer.py index a1fb669369a0..79abceb67786 100644 --- a/tensorrt_llm/_torch/attention_backend/flashinfer.py +++ b/tensorrt_llm/_torch/attention_backend/flashinfer.py @@ -564,6 +564,17 @@ def paged_kv_indices(self) -> torch.Tensor: return self._paged_kv_indices[:self.num_generation_blocks + self.num_context_blocks] + def _primary_space_rep_layer(self): + """Representative layer of the primary page-index space, or None. + + None keeps the historical layer-independent lookup for managers without + per-layer spaces. + """ + if not self._per_layer_index_spaces or self._vswa_layer_to_pool is None: + return None + primary_pool_id = self._vswa_layer_to_pool.get(0, 0) + return self._vswa_pool_to_rep_layer.get(primary_pool_id) + def get_paged_kv_indices_for_layer(self, layer_idx: int) -> torch.Tensor: """Return page indices for the pool that *layer_idx* belongs to. @@ -1019,6 +1030,7 @@ def _post_init_with_buffers(self, buffers) -> None: # use the indices that match its pool's buffer. self._vswa_layer_to_pool: Optional[Dict[int, int]] = None self._vswa_pool_indices_cache: Optional[Dict[int, torch.Tensor]] = None + self._per_layer_index_spaces: bool = False if self.kv_cache_manager is not None: blocks_in_primary_pool = self.kv_cache_manager.blocks_in_primary_pool @@ -1051,16 +1063,28 @@ def _post_init_with_buffers(self, buffers) -> None: mgr = self.kv_cache_manager get_scale = getattr(mgr, 'get_layer_page_index_scale', None) layer_space: Dict[int, int] = {} + # With SWA scratch reuse the page index is genuinely per-layer: + # a scratch block's sub-page rotates with block position, so it + # cannot be folded into a per-layer base pointer the way a fixed + # layer offset can. Give every layer its own index space; the + # per-space buffers and the per-layer swap below then work + # unchanged, just with more spaces. Bound outside the guard: it is + # read again below, where only `layer_space` being empty currently + # short-circuits the read. + per_layer_spaces = getattr(mgr, 'enable_swa_scratch_reuse', False) if hasattr(mgr, 'layer_to_pool_mapping_dict') and get_scale: space_ids = {} for layer_idx in getattr(mgr, 'layer_offsets', {}): layer_offset = mgr.layer_offsets[layer_idx] - key = (mgr.layer_to_pool_mapping_dict[layer_offset], - get_scale(layer_idx)) + key = layer_idx if per_layer_spaces else ( + mgr.layer_to_pool_mapping_dict[layer_offset], + get_scale(layer_idx)) space_ids.setdefault(key, len(space_ids)) layer_space[layer_idx] = space_ids[key] if layer_space and (getattr(mgr, 'is_vswa', False) + or per_layer_spaces or len(set(layer_space.values())) > 1): + self._per_layer_index_spaces = per_layer_spaces self._vswa_layer_to_pool = {} self._vswa_pool_to_rep_layer: Dict[int, int] = {} for layer_idx, pool_id in layer_space.items(): @@ -1541,9 +1565,16 @@ def _to_int32_tensor(arr: np.ndarray) -> torch.Tensor: self.num_context_blocks = sum(self.num_blocks[:self.num_contexts]) self.num_generation_blocks = sum(self.num_blocks[self.num_contexts:]) - # indices of used cache blocks for each sequence + # indices of used cache blocks for each sequence. + # Under per-layer index spaces (SWA scratch reuse) there is no + # layer-independent index: a scratch block has no base page, so a + # layer_idx-less call would yield BAD_PAGE_INDEX for it. Build the base + # table from the primary space's representative layer so the primary + # buffer below is populated on the same footing as every other space. paged_kv_indices = self.kv_cache_manager.get_batch_cache_indices_flat( - self.request_ids, self.num_blocks) + self.request_ids, + self.num_blocks, + layer_idx=self._primary_space_rep_layer()) self._paged_kv_indices[:paged_kv_indices.size(0)].copy_( paged_kv_indices, non_blocking=True) @@ -2438,7 +2469,11 @@ def forward_impl( ) return - # Key and Value + # Key and Value. The manager decides the addressing mode: with SWA + # scratch reuse the page indices are PER_LAYER, so the buffer must be + # based at the pool group rather than at this layer's sub-page. See + # KVCacheManagerV2.page_index_mode -- indices and buffer are derived + # from one source so they cannot disagree. kv_cache = metadata.kv_cache_manager.get_buffers( self.layer_idx, kv_layout=metadata.kv_layout) diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv4.py b/tensorrt_llm/_torch/models/modeling_deepseekv4.py index d3183c0856d3..a1dc16079878 100644 --- a/tensorrt_llm/_torch/models/modeling_deepseekv4.py +++ b/tensorrt_llm/_torch/models/modeling_deepseekv4.py @@ -2546,7 +2546,6 @@ class DeepseekV4ForCausalLM(SpecDecOneEngineForCausalLM[DeepseekV4Model, Pretrai def get_model_defaults(cls, llm_args: "TorchLlmArgs") -> dict: kv_cache_defaults = { "tokens_per_block": 128, - "enable_swa_scratch_reuse": True, } if llm_args is not None and llm_args.kv_cache_config.dtype == "fp8_ds_mla": kv_cache_defaults["tokens_per_block"] = 256 diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py index b1a965d20e32..0dcd9e8bcf16 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py @@ -18,7 +18,18 @@ import sys from collections import OrderedDict, defaultdict, deque from dataclasses import dataclass, field, replace -from typing import TYPE_CHECKING, Dict, Iterable, List, NamedTuple, Optional, Sequence, Tuple, Union +from typing import ( + TYPE_CHECKING, + Dict, + Iterable, + List, + NamedTuple, + Optional, + Sequence, + Set, + Tuple, + Union, +) import numpy as np import torch @@ -142,6 +153,120 @@ class BlockReusePolicy(StrEnum): PER_CONVERSATION = "per_conversation" +# Number of per-request SWA scratch-range lines emitted at debug level before +# the manager goes quiet. Enough to eyeball compute_scratch_range on real +# prompts without flooding a long run. +_SWA_SCRATCH_DEBUG_LOG_LIMIT = 32 + + +def _summarize_ints(values: Sequence[int]) -> str: + """Compact "5" / "5 x8" / "5,4096 x7" rendering for a BatchDesc column.""" + if not values: + return "-" + runs: List[Tuple[int, int]] = [] + for v in values: + if runs and runs[-1][0] == v: + runs[-1] = (v, runs[-1][1] + 1) + else: + runs.append((v, 1)) + return ",".join(f"{v}" if n == 1 else f"{v} x{n}" for v, n in runs) + + +def compute_scratch_flat_page_indices( + first: int, + last: int, + scratch_pages_per_block: int, + scale: int, + layer_offset: int, + slot_ids: Sequence[int], + div_factor: int, +) -> "np.ndarray": + """Flat PER_LAYER page indices for scratch block positions ``[first, last)``. + + ``first``/``last`` are positions *within* the scratch range, i.e. relative to + ``scratch_range.beg``. This is the host mirror of the arithmetic + ``_copy_swa_block_offsets_with_scratch_compiled`` performs on device for the + AttentionOp path:: + + total = position * scratch_pages_per_block + slot = slot_ids[total // scale] + sub = (total % scale + layer_offset) % scale + index = (slot * scale + sub) // div_factor + + The rotation of ``sub`` with block position is the whole reason scratch + cannot use SHARED addressing: a non-scratch layer sits at a fixed sub-page + that can be folded into a base pointer, a scratch block's does not. + + Kept a module-level pure function on purpose. It is the piece that silently + corrupts KV when wrong -- an off-by-one here reads another layer's cache + rather than failing -- so it must be testable without a GPU, a model, or a + manager instance. + """ + count = last - first + if count <= 0: + return np.empty(0, dtype=np.int32) + total = (np.arange(first, last, dtype=np.int64)) * scratch_pages_per_block + slots = np.asarray(slot_ids, dtype=np.int64)[total // scale] + sub = (total % scale + layer_offset) % scale + return ((slots * scale + sub) // div_factor).astype(np.int32) + + +def apply_scratch_to_block_segment( + seg: "np.ndarray", + beg: int, + end: int, + scratch_pages_per_block: int, + scale: int, + layer_offset: int, + slot_ids: Sequence[int], + div_factor: int, +) -> None: + """Rewrite one request's flat page indices in place for PER_LAYER addressing. + + ``seg`` is that request's slice of the page table, ``[beg, end)`` its scratch + range in block ordinals. Blocks inside the range are replaced by the rotation; + blocks outside it keep their base index and only gain ``layer_offset``, + because the caller now addresses the pool-group base rather than this layer's + sub-page. + + Separated from the batch loop so the clamping is testable: ``[beg, end)`` may + fall partly or entirely outside ``seg``, and getting that wrong silently + shifts the wrong blocks rather than raising. + """ + n = seg.shape[0] + lo, hi = max(beg, 0), min(end, n) + if hi <= lo: + # No overlap: every block in this segment is non-scratch. Collapse the + # window to the end so the two "outside" slices cover the whole segment. + # (A negative `beg` here would otherwise make `seg[hi:]` index from the + # end of the array and shift only a suffix.) + lo = hi = n + + if hi > lo: + slots_required = ((hi - beg - 1) * scratch_pages_per_block) // scale + 1 + if slots_required > len(slot_ids): + raise ValueError( + f"blocks [{lo}, {hi}) of scratch range [{beg}, {end}) at " + f"{scratch_pages_per_block} page(s) per block need {slots_required} scratch " + f"slot(s) at scale {scale}, but the descriptor holds {len(slot_ids)}" + ) + + shift = layer_offset // div_factor + for part in (seg[:lo], seg[hi:]): + if part.size: + np.add(part, shift, out=part, where=part != BAD_PAGE_INDEX) + if hi > lo: + seg[lo:hi] = compute_scratch_flat_page_indices( + lo - beg, + hi - beg, + scratch_pages_per_block, + scale, + layer_offset, + slot_ids, + div_factor, + ) + + def _request_conversation_id(request: LlmRequest) -> Optional[str]: if request.is_dummy_request: return None @@ -827,9 +952,25 @@ def __init__( # Set True by a compression manager; generation-step resize then leaves history untouched. self.kv_compression_manages_history: bool = False + # "auto" is normally resolved against the attention backend and the + # resolved manager version at config load time + # (ModelLoader.load_config_and_apply_defaults); paths that skip that + # step leave the sentinel in place and get the conservative answer. + # Narrowed again below, once the per-layer window vector exists: a + # model with no sliding-window layer has nothing to share. self.enable_swa_scratch_reuse = ( - kv_cache_config.enable_swa_scratch_reuse and not self.is_draft + kv_cache_config.enable_swa_scratch_reuse is True and not self.is_draft ) + # Budget for the per-request scratch-range debug lines emitted by + # _log_scratch_desc. Bounded so a debug-level run does not turn into one + # log line per context request forever. `_scratch_debug_seen` keys the + # budget by (request, pool) rather than by call site: the FlashInfer + # PER_LAYER path visits every scratch descriptor once *per layer* per + # iteration, so an unkeyed counter would spend the whole budget on the + # first request and emit num_layers duplicate lines while doing it. + self._scratch_debug_budget = _SWA_SCRATCH_DEBUG_LOG_LIMIT + self._scratch_debug_seen: Set[Tuple[int, int]] = set() + self._per_layer_flat_validated = False block_reuse_config = kv_cache_config.block_reuse_config self.block_reuse_policy = BlockReusePolicy(block_reuse_config.policy) self.num_local_layers = len(self.pp_layers) @@ -1156,6 +1297,22 @@ def append_to_kv_heads_per_layer( self._pool_layer_ids_by_role.setdefault( (pool_id, buffer_id.role), buffer_id.layer_id ) + # Scratch reuse shares the out-of-window part of a prefill block between + # the layers of one windowed lifecycle. With no windowed lifecycle there + # is nothing to share, and switching to PER_LAYER addressing would cost + # one virtual pool per layer for nothing. Since the feature is on by + # default, that has to be a silent no-op rather than a warning. + # + # The lifecycles are the authoritative source: a subclass may declare + # windows directly in its layer configs (DeepseekV4CacheManager) instead + # of through ``max_attention_window_vec``. + if self.enable_swa_scratch_reuse and not _introspection.swa_life_cycle_ids(self.impl): + logger.debug( + f"{type(self).__name__}: SWA scratch reuse is inactive; no attention " + "lifecycle uses a sliding window." + ) + self.enable_swa_scratch_reuse = False + # num_pools is the logical layer-group count. With SWA scratch reuse, # scratch slot IDs are only valid with per-layer page indices, so the # attention op sees one virtual pool per local layer while the @@ -1255,6 +1412,7 @@ def append_to_kv_heads_per_layer( self._prepare_page_table_tensor(index_mapper_capacity) self._log_kv_cache_pool_lifecycle_mapping() + self._log_swa_scratch_summary() def _get_pool_roles(self, pool_id: int) -> Tuple[DataRole, Optional[DataRole]]: """Return the roles represented by the two page-table index lanes. @@ -1628,6 +1786,148 @@ def _log_kv_cache_pool_lifecycle_mapping(self) -> None: for entry in entries: logger.info(entry) + def _swa_scratch_counterfactual_config(self) -> SwaScratchReuseConfig: + """The scratch config that is (or would be) in force. Pure sizing input.""" + return SwaScratchReuseConfig(max_rewind_len=self.num_extra_kv_tokens) + + def _log_swa_scratch_summary(self) -> None: + """Report what SWA scratch reuse saves (or would save) on this engine. + + ``_compute_slots_for_batch`` is a pure function of the registered batch + shapes, so both sides of the comparison are computable at startup with + no GPU work and no workload knowledge. Emitting both makes the feature + observable: today enabling scratch only makes ``iter_alloc_new_blocks`` + drop, with no attribution. + """ + swa_lc_ids = _introspection.swa_life_cycle_ids(self.impl) + enabled = self.enable_swa_scratch_reuse + if not swa_lc_ids: + return + + layers_per_lc: Dict[int, int] = defaultdict(int) + window_per_lc: Dict[int, Optional[int]] = {} + for layer in self.kv_cache_manager_py_config.layers: + # A hybrid model's layer list also holds SsmLayerConfig, which has + # no sliding_window_size to read. An SSM layer never joins an + # attention lifecycle, so skipping it leaves every count reported + # below unchanged -- same guard as _stats_life_cycle_metadata. + if not isinstance(layer, AttentionLayerConfig): + continue + lc_id = int(self.impl.get_layer_group_id(layer.layer_id)) + layers_per_lc[lc_id] += 1 + window_per_lc[lc_id] = layer.sliding_window_size + + constraints = self.kv_cache_manager_py_config.constraints or [] + tpb = self.tokens_per_block + scratch_config = self._swa_scratch_counterfactual_config() + + logger.info( + f"{type(self).__name__} SWA scratch reuse: {'ENABLED' if enabled else 'DISABLED'}" + ) + for lc_id in swa_lc_ids: + num_layers = layers_per_lc.get(lc_id, 0) + frac = f"1/{num_layers}" if num_layers else "?" + logger.info( + f" pool_group={_introspection.pool_group_index(self.impl, lc_id)} " + f"window={window_per_lc.get(lc_id)} swa_layers={num_layers} " + f"frac_max~={frac}" + ) + + any_saving = False + best_saving_pct = 0 + for idx, batch in enumerate(constraints): + without = _introspection.compute_slots_for_batch(self.impl, batch, tpb, None) + with_scratch = _introspection.compute_slots_for_batch( + self.impl, batch, tpb, scratch_config + ) + capacities = [kv.capacity for kv in batch.kv_caches] + histories = [kv.history_length for kv in batch.kv_caches] + shape = ( + f"{len(batch.kv_caches)} req(s), capacity={_summarize_ints(capacities)}, " + f"history={_summarize_ints(histories)}" + ) + logger.info(f" constraint[{idx}] ({shape}):") + for pg_idx, (no_s, yes_s) in enumerate(zip(without, with_scratch)): + delta = "" + if no_s > 0 and yes_s != no_s: + pct = (yes_s - no_s) * 100 // no_s + delta = f" ({pct}%)" + # Only a drop in slots is a saving. Counting a rise would + # leave best_saving_pct at 0, so the warning below would + # report "up to 0%", and it would also suppress the + # inert-configuration branch that should fire instead. + if yes_s < no_s: + any_saving = True + best_saving_pct = min(best_saving_pct, pct) + logger.info( + f" pool_group={pg_idx} slots without scratch: {no_s} " + f"with scratch: {yes_s}{delta}" + ) + + if constraints and not any_saving: + message = ( + f"{type(self).__name__}: SWA scratch reuse changes no slot count for any " + "registered batch shape. Every constraint is decode-shaped or its scratch " + "range is empty, so scratch reuse is inert on this configuration " + f"(max_num_tokens={self.max_num_tokens}, tokens_per_block={tpb}, " + f"windows={[window_per_lc.get(lc) for lc in swa_lc_ids]})." + ) + if enabled: + logger.warning(message) + else: + logger.debug(message) + elif any_saving and not enabled: + # The prefill constraint is registered for every model with + # max_num_tokens, scratch or no scratch, because the engine really + # must be able to run that shape. Its effect on capacity is + # quota-dependent: at a binding quota it raises allocatable tokens, + # but at a non-binding quota it can *reduce* them, because the SWA + # pool is now sized for the true prefill requirement instead of the + # old window-sized floor. Scratch reuse is what pays that back. A + # windowed model that registers the constraint and then declines + # the scratch saving takes the sizing cost with none of the benefit, + # and would otherwise do so silently. + logger.warning( + f"{type(self).__name__}: this model has sliding-window layers and SWA scratch " + f"reuse would cut windowed slots by up to {abs(best_saving_pct)}% on a " + "registered prefill shape, but it is disabled. The prefill constraint is " + "registered either way, so this configuration pays the sizing cost without the " + "saving. Set kv_cache_config.enable_swa_scratch_reuse=True (requires " + "attn_backend TRTLLM or FLASHINFER) to recover it." + ) + + def _log_scratch_desc(self, request_id: int, pool_id: int, desc, path: str) -> None: + """Emit one bounded debug line per (request, pool) that holds scratch. + + Both scratch-consuming paths call this, which is the point: the + AttentionOp path goes through ``_copy_scratch_metadata_to_device`` and + the FlashInfer PER_LAYER path through ``_apply_scratch_to_flat_indices``. + Logging from only the former left every FlashInfer model -- i.e. exactly + the models this feature was extended to cover -- with no per-request + scratch diagnostics at all, so ``compute_scratch_range`` could not be + eyeballed on the one backend where it was newest. + + ``path`` is recorded so a reader can tell which addressing mode produced + the descriptor without cross-referencing the attention backend. + """ + if self._scratch_debug_budget <= 0: + return + key = (request_id, pool_id) + if key in self._scratch_debug_seen: + return + self._scratch_debug_seen.add(key) + self._scratch_debug_budget -= 1 + kv_cache = self.kv_cache_map.get(request_id) + if kv_cache is None: + return + logger.debug( + f"SWA scratch: request={request_id} pool={pool_id} path={path} " + f"history_length={kv_cache.history_length} " + f"capacity={kv_cache.capacity} " + f"scratch_range=[{int(desc.range.beg)}, {int(desc.range.end)}) " + f"num_scratch_slots={len(desc.slot_ids)}" + ) + def _prepare_swa_scratch_copy_tensors(self, index_mapper_capacity: int) -> None: pool_ids = torch.empty( self.num_attention_op_pools, @@ -1765,6 +2065,7 @@ def _copy_scratch_metadata_to_device( if desc is None: continue slot_ids = desc.slot_ids + self._log_scratch_desc(request_id, pool_id, desc, path="attention_op") if len(slot_ids) > self._max_scratch_slots: raise RuntimeError( f"Scratch slot count {len(slot_ids)} exceeds staging capacity " @@ -1896,19 +2197,16 @@ def _build_base_config( ) ) - # 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 runs one iteration at the + # per-iteration token budget. Every input here is engine + # configuration -- no workload knowledge -- and warmup exercises + # exactly this shape, so the constraint is registered regardless of + # whether avg_seq_len was supplied. Without it, a model that does + # not set avg_seq_len gets no prefill BatchDesc at all: the SWA pool + # is then sized from swa_floor_blocks alone (only ~window-sized, + # silently below the real prefill requirement) and SWA scratch reuse + # has no prefill shape to shrink. + constraints.extend(self._prefill_constraints()) buffer_type = [Role.KEY] if self.kv_cache_type != CacheTypeCpp.SELFKONLY: @@ -1983,6 +2281,29 @@ def _build_cache_config(self, config: KVCacheManagerConfigPy) -> KVCacheManagerC """Customize the general cache config for a specialized cache manager.""" return config + def _prefill_constraints(self) -> List[BatchDesc]: + """Batch shapes the engine must be able to run during prefill. + + One fresh context request at the per-iteration token budget. Note this + models the budget as a single request; when ``max_num_tokens`` exceeds + ``max_seq_len`` the real iteration is several shorter requests, which + needs *more* windowed slots than modelled here. That stays a + conservative under-count rather than an over-provision, and refining it + is a separate sizing change. + """ + if self.max_num_tokens is None: + return [] + return [ + BatchDesc( + [ + KVCacheDesc( + capacity=self.max_num_tokens + self.num_extra_kv_tokens, + history_length=0, + ) + ] + ) + ] + def _get_typical_seq_len(self, kv_cache_config: KvCacheConfig) -> int | None: """Return the configured typical sequence length, if any.""" return kv_cache_config.avg_seq_len @@ -2043,17 +2364,135 @@ def blocks_in_primary_pool(self) -> int: """ return self.impl.get_page_index_upper_bound(0, Role.KEY) - def get_buffers(self, layer_idx: int, kv_layout: str = "NHD") -> Optional[torch.Tensor]: + def _page_index_upper_bound(self, local_layer_idx: int, index_mode: PageIndexMode) -> int: + """Page-index extent for a layer under the given addressing mode. + + ``get_page_index_upper_bound`` is SHARED-shaped: it subtracts the + layer's sub-page offset because the base pointer already skipped it. + PER_LAYER bases at the pool group, so the layer offset is addressable + and must be added back. + """ + upper = self.impl.get_page_index_upper_bound(local_layer_idx, Role.KEY) + if index_mode == PageIndexMode.SHARED: + return upper + return upper + int( + self.impl.get_page_index_converter(local_layer_idx, Role.KEY).layer_offset + ) + + def _validate_per_layer_kv_adjacency(self) -> None: + """Check the layout assumptions PER_LAYER page indices rely on. + + A FlashInfer-style page table carries *one* index per block plus a + ``kv_factor`` axis, so it can only address K and V if V sits exactly one + sub-page after K and K starts on a ``kv_factor`` boundary. The + AttentionOp path needs neither property (it carries separate K and V + indices), so this is checked only where it is relied upon. Under scratch + rotation the index becomes ``(offset + layer_offset) % scale``, which + preserves both properties only when ``scale`` and the per-block scratch + stride are multiples of ``kv_factor``. + """ + if self._per_layer_flat_validated: + return + if self.kv_cache_type != CacheTypeCpp.SELFKONLY: + # kv_factor == 1 under SELFKONLY: one sub-page per block, so there + # is no K/V pairing to preserve and every index divides through + # unchanged. Every other layout has to be checked. + self._check_per_layer_kv_adjacency() + # Latched only on success: memoizing a raise as "validated" would let a + # later call return silently and build a flat page table on a layout + # already proven unsupported. + self._per_layer_flat_validated = True + + def _check_per_layer_kv_adjacency(self) -> None: + """Raise if any layer's converter breaks a flat page table's assumptions.""" + for local_layer_idx in range(self.num_local_layers): + layer_id = LayerId(local_layer_idx) + conv_k = self.impl.get_page_index_converter(layer_id, Role.KEY) + conv_v = self.impl.get_page_index_converter(layer_id, Role.VALUE) + problems = [] + if int(conv_v.layer_offset) != int(conv_k.layer_offset) + 1: + problems.append( + f"V sub-page {int(conv_v.layer_offset)} is not adjacent to K " + f"{int(conv_k.layer_offset)}" + ) + if int(conv_k.layer_offset) % self.kv_factor != 0: + problems.append( + f"K sub-page {int(conv_k.layer_offset)} is not {self.kv_factor}-aligned" + ) + if int(conv_k.scale) % self.kv_factor != 0: + problems.append(f"slot scale {int(conv_k.scale)} is not {self.kv_factor}-aligned") + if int(conv_k.scratch_pages_per_block) % self.kv_factor != 0: + problems.append( + f"scratch stride {int(conv_k.scratch_pages_per_block)} is not " + f"{self.kv_factor}-aligned" + ) + if conv_k.expansion != 1 or conv_v.expansion != 1: + problems.append("expanded page indices are not supported") + if problems: + raise NotImplementedError( + f"SWA scratch reuse cannot produce a flat per-layer page table for local " + f"layer {local_layer_idx}: {'; '.join(problems)}. Set " + "kv_cache_config.enable_swa_scratch_reuse=False, or use an attention " + "backend that consumes copy_batch_block_offsets (which carries separate " + "K and V indices)." + ) + + @property + def page_index_mode(self) -> PageIndexMode: + """The addressing contract this manager's page indices are built with. + + This is a property of the manager, not a choice the caller makes: the + indices produced by ``get_batch_cache_indices_flat`` and the buffer + returned by ``get_buffers`` must agree, and only the manager knows + whether scratch is active. Handing a PER_LAYER index table to a SHARED + buffer reads the wrong memory -- that pairing is what produced an + illegal memory access during Gemma4 bring-up -- so the two are derived + from one source rather than threaded through backends independently. + """ + return PageIndexMode.PER_LAYER if self.enable_swa_scratch_reuse else PageIndexMode.SHARED + + def get_buffers( + self, + layer_idx: int, + kv_layout: str = "NHD", + index_mode: Optional[PageIndexMode] = None, + ) -> Optional[torch.Tensor]: + """Wrap this layer's KV pages as a tensor addressable by page index. + + The addressing contract must match the mode used to build the page + indices handed to the kernel: + + - ``SHARED``: base pointer is this layer's own sub-page, so the index + carries no layer offset. This is the historical behavior. + - ``PER_LAYER``: base pointer is the pool-group base and the index + carries the layer offset. Required whenever SWA scratch reuse is + active, because a scratch block's sub-page position rotates with the + block position and therefore cannot be folded into a fixed pointer. + + ``index_mode`` defaults to :attr:`page_index_mode`, so callers get the + mode that matches the indices this manager produces without having to + know the rule. It stays overridable for tests that need to construct a + specific pairing on purpose. + """ + if index_mode is None: + index_mode = self.page_index_mode layer_offset = self.layer_offsets[layer_idx] - addr_key = self.impl.get_mem_pool_base_address(layer_offset, Role.KEY, PageIndexMode.SHARED) + addr_key = self.impl.get_mem_pool_base_address(layer_offset, Role.KEY, index_mode) if self.kv_cache_type != CacheTypeCpp.SELFKONLY: - addr_value = self.impl.get_mem_pool_base_address( - layer_offset, Role.VALUE, PageIndexMode.SHARED - ) page_size_key = self.impl.get_page_stride(layer_offset, Role.KEY) page_size_value = self.impl.get_page_stride(layer_offset, Role.VALUE) - - assert addr_key + page_size_value == addr_value and page_size_key == page_size_value + if index_mode == PageIndexMode.SHARED: + addr_value = self.impl.get_mem_pool_base_address( + layer_offset, Role.VALUE, index_mode + ) + assert addr_key + page_size_value == addr_value and page_size_key == page_size_value + else: + # PER_LAYER shares one base pointer for both roles; K/V are + # distinguished by the layer offset carried in the index. The + # kv_factor axis of this tensor still assumes V sits one + # sub-page after K, which _validate_per_layer_kv_adjacency + # checks up front. + assert page_size_key == page_size_value assert kv_layout in ["NHD", "HND"], f"Unsupported kv_layout: {kv_layout}" @@ -2066,7 +2505,7 @@ def get_buffers(self, layer_idx: int, kv_layout: str = "NHD") -> Optional[torch. layer_head_dim = self.head_dim_per_layer[layer_offset] if kv_layout == "NHD": shape = [ - self.impl.get_page_index_upper_bound(layer_offset, Role.KEY) // self.kv_factor, + self._page_index_upper_bound(layer_offset, index_mode) // self.kv_factor, self.kv_factor, self.tokens_per_block, self.num_kv_heads_per_layer[layer_offset], @@ -2074,7 +2513,7 @@ def get_buffers(self, layer_idx: int, kv_layout: str = "NHD") -> Optional[torch. ] else: shape = [ - self.impl.get_page_index_upper_bound(layer_offset, Role.KEY) // self.kv_factor, + self._page_index_upper_bound(layer_offset, index_mode) // self.kv_factor, self.kv_factor, self.num_kv_heads_per_layer[layer_offset], self.tokens_per_block, @@ -3321,6 +3760,25 @@ def get_iteration_stats(self): ) def get_block_ids_per_seq(self, request_ids: List[int]) -> torch.Tensor: + # This flat view maps BAD_PAGE_INDEX to block 0. A scratch block holds no + # per-block page, so its base page index *is* BAD_PAGE_INDEX and would be + # silently rewritten to block 0 -- the consumer (flash_mla) would then read + # the wrong KV. Per-layer scratch slot ids only travel through + # copy_batch_block_offsets, so fail loudly instead. + if self.enable_swa_scratch_reuse: + for request_id in request_ids: + # A request may legitimately be absent (already released, or a + # dummy); that is not this guard's concern, and a bare + # kv_cache_map[...] would turn it into a KeyError that masks the + # scratch diagnostic below. + kv_cache = self.kv_cache_map.get(request_id) + if kv_cache is not None and kv_cache.has_scratch_slots: + raise RuntimeError( + f"get_block_ids_per_seq() is not compatible with SWA scratch " + f"reuse: request {request_id} holds scratch slots whose per-layer " + "page ids are not representable in a flat block-id table. Set " + "kv_cache_config.enable_swa_scratch_reuse=False for this backend." + ) block_ids_per_seq = self.get_batch_cache_indices(request_ids) block_ids_per_seq_tensors = [ torch.tensor( @@ -3636,6 +4094,16 @@ def get_batch_cache_indices_flat( Returns a CPU int32 tensor (pinned when supported) ready for an async H2D copy. """ + if layer_idx is None and self.enable_swa_scratch_reuse: + # Under scratch reuse the page index is genuinely per-layer: a + # scratch block holds no base page, so a layer-independent table + # would carry BAD_PAGE_INDEX for it and the caller would index its + # PER_LAYER buffer out of bounds. Fail here rather than let that + # reach a kernel as an illegal access. + raise ValueError( + "get_batch_cache_indices_flat() requires layer_idx when SWA scratch " + "reuse is enabled: scratch blocks have no layer-independent page index." + ) if layer_idx is None: pool_id = 0 scale = self._index_scale_ints[pool_id] @@ -3664,8 +4132,78 @@ def get_batch_cache_indices_flat( # get_batch_cache_indices. valid = out != BAD_PAGE_INDEX np.copyto(out, out * scale // div_factor, where=valid) + + if self.enable_swa_scratch_reuse and layer_idx is not None: + self._apply_scratch_to_flat_indices( + out, request_ids, num_blocks, layer_idx, pool_id, scale, div_factor + ) return out_tensor + def _apply_scratch_to_flat_indices( + self, + out: "np.ndarray", + request_ids: List[int], + num_blocks: List[int], + layer_idx: int, + pool_id: int, + scale: int, + div_factor: int, + ) -> None: + """Overwrite scratch blocks in a flat page table with PER_LAYER indices. + + A scratch block holds no per-block page, so the base-index transform + above left it at BAD_PAGE_INDEX. Its real address is a rotating + sub-page inside a shared slot -- the same arithmetic + ``_copy_swa_block_offsets_with_scratch_compiled`` performs on device for + the AttentionOp path, applied here on the host for one layer. + + Non-scratch blocks also gain ``layer_offset``, because the caller is + addressing the pool-group base (PER_LAYER), not this layer's sub-page. + """ + # Checked lazily: only a flat page table needs K/V adjacency, so + # backends that consume copy_batch_block_offsets are never subjected + # to this restriction. + self._validate_per_layer_kv_adjacency() + converter = self.impl.get_page_index_converter(self.layer_offsets[layer_idx], Role.KEY) + layer_offset = int(converter.layer_offset) + scratch_pages = int(converter.scratch_pages_per_block) + offset = 0 + for req_id, n in zip(request_ids, num_blocks): + kv_cache = self.kv_cache_map.get(req_id) + desc = None if kv_cache is None else kv_cache.get_scratch_desc(pool_id) + if desc is None: + # Non-scratch request: only the layer offset is missing. + out[offset : offset + n] = np.where( + out[offset : offset + n] != BAD_PAGE_INDEX, + out[offset : offset + n] + layer_offset // div_factor, + BAD_PAGE_INDEX, + ) + offset += n + continue + self._log_scratch_desc(req_id, pool_id, desc, path="flashinfer_per_layer") + beg, end = int(desc.range.beg), int(desc.range.end) + slot_ids = desc.slot_ids + try: + apply_scratch_to_block_segment( + out[offset : offset + n], + beg, + end, + scratch_pages, + scale, + layer_offset, + slot_ids, + div_factor, + ) + except ValueError as exc: + # Re-raise with the identity the helper cannot know, so a bad + # descriptor names the request and pool rather than surfacing as + # a bare arithmetic error from inside the hot path. + raise RuntimeError( + f"SWA scratch descriptor for request {req_id} pool {pool_id} is " + f"inconsistent: {exc}." + ) from exc + offset += n + def get_cache_bytes_per_token(self) -> int: data_roles = [Role.KEY] if self.kv_cache_type != CacheTypeCpp.SELFKONLY: diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_stats.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_stats.py index 580bd5a857cc..5bfe49b2e851 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_stats.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_stats.py @@ -46,6 +46,27 @@ "iterIntraDeviceCopyBytes", "iterHostDroppedBlocks", "iterHostDroppedBytes", + "iterScratchBlocks", + "iterScratchSlotsInUse", +) + +# Iteration fields deliberately absent from the per-pool-group view, with the +# reason. `test_pool_group_keys_cover_every_iteration_field` asserts that every +# field the serializer emits is either in the allowlist above or named here, so +# adding a field without deciding which side it belongs on fails CI rather than +# silently disappearing from the pool-group view. +KV_CACHE_ITERATION_STATS_NOT_PER_POOL_GROUP = frozenset( + { + # Reuse accounting is tracked per window size, not per pool group: a + # single pool group can back several windows, so these would double + # count. + "iterReusedBlocks", + "iterFullReusedBlocks", + "iterPartialReusedBlocks", + "iterMissedBlocks", + # Derived from the four above; meaningless once they are absent. + "iterCacheHitRate", + } ) # Subset of KV_CACHE_ITERATION_STATS_POOL_GROUP_KEYS reported per cold pool group. The primary_* keys @@ -148,6 +169,8 @@ def serialize_kv_cache_iteration_stats(stats, keys: tuple[str, ...] | None = Non "iterIntraDeviceCopyBytes": stats.iter_intra_device_copy_bytes, "iterHostDroppedBlocks": stats.iter_host_dropped_blocks, "iterHostDroppedBytes": stats.iter_host_dropped_bytes, + "iterScratchBlocks": stats.iter_scratch_blocks, + "iterScratchSlotsInUse": stats.iter_scratch_slots_in_use, } if keys is None: return fields diff --git a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py index dbdb4c582e67..f61f024a1a92 100644 --- a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py @@ -3533,13 +3533,21 @@ def blocks_in_primary_pool(self) -> int: LayerId(local_layer_idx), Role.KEY) return 0 - def get_buffers(self, - layer_idx: int, - kv_layout: str = "NHD") -> Optional[torch.Tensor]: + def get_buffers( + self, + layer_idx: int, + kv_layout: str = "NHD", + index_mode: Optional[PageIndexMode] = None + ) -> Optional[torch.Tensor]: + # `None` (not SHARED) mirrors KVCacheManagerV2.get_buffers: the base + # resolves it to `page_index_mode`, so a caller that omits the mode -- + # FlashInfer does -- still gets the buffer that matches the indices + # this manager produces. Defaulting to SHARED here would hand a SHARED + # buffer to PER_LAYER indices once scratch reuse is on. local_layer_idx = self.layer_offsets[layer_idx] if self._is_local_mamba_layer(local_layer_idx): return None - return super().get_buffers(layer_idx, kv_layout) + return super().get_buffers(layer_idx, kv_layout, index_mode) def _iter_cache_buffers_for_invalid_check(self) -> Iterable[torch.Tensor]: for global_layer_id, local_layer_id in self.layer_offsets.items(): diff --git a/tensorrt_llm/_torch/pyexecutor/model_loader.py b/tensorrt_llm/_torch/pyexecutor/model_loader.py index 9899ad7426ac..e223b63efb07 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_loader.py +++ b/tensorrt_llm/_torch/pyexecutor/model_loader.py @@ -28,6 +28,7 @@ ModelExpressConfig, SparseAttentionConfig, TorchLlmArgs) from tensorrt_llm.llmapi.llm_utils import (_resolve_kv_cache_manager_v2_auto, + _resolve_swa_scratch_reuse_auto, _resolve_transceiver_runtime_auto, apply_model_defaults_to_llm_args) from tensorrt_llm.logger import logger @@ -476,6 +477,7 @@ def load_config_and_apply_defaults( # "auto" sentinel so it never leaks past config loading. _resolve_transceiver_runtime_auto(llm_args) _resolve_kv_cache_manager_v2_auto(llm_args) + _resolve_swa_scratch_reuse_auto(llm_args) return llm_args config_kwargs = { @@ -540,13 +542,18 @@ def load_config_and_apply_defaults( config.pretrained_config) _resolve_kv_cache_manager_v2_auto(llm_args, preference_cls, config.pretrained_config) + # Depends on the resolved manager version, the model's own manager + # preference, and the attention backend, which model defaults may have + # just set. + _resolve_swa_scratch_reuse_auto(llm_args, preference_cls, + config.pretrained_config) _validate_and_adjust_mamba_snapshot_config(config, llm_args) if original_kv_cache_manager_setting == "auto": - logger.info( - "Resolved use_kv_cache_manager_v2='auto' to %s for %s", - llm_args.kv_cache_config.use_kv_cache_manager_v2, - model_cls.__name__ - if model_cls is not None else "unknown model") + model_name = (model_cls.__name__ + if model_cls is not None else "unknown model") + logger.info("Resolved use_kv_cache_manager_v2='auto' to " + f"{llm_args.kv_cache_config.use_kv_cache_manager_v2} " + f"for {model_name}") return llm_args diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index 7e4fed2a0eb2..9eaf3dc8e65f 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -468,6 +468,19 @@ def append_to_kv_heads_per_layer(num_kv_heads_per_layer: List[int], # The `w > 0` check excludes LinearCacheType.RECURRENT_STATES sentinel # values (negative) used by hybrid linear attention models. self.is_vswa = uses_vswa_kv_cache_layout(self.max_attention_window_vec) + + # Only an explicit request is worth a warning: the default resolves to + # False whenever the V1 manager is selected. Read through getattr: + # this constructor is also handed the bindings `executor.KvCacheConfig`, + # which mirrors the C++ fields only and carries none of the Python-only + # ones, and a legacy config that cannot express the request is a + # request that was never made. + if getattr(kv_cache_config, "enable_swa_scratch_reuse", False) is True: + logger.warning( + "kv_cache_config.enable_swa_scratch_reuse is set but the V1 KV cache " + "manager is in use; SWA scratch reuse is a V2-only feature and is " + "silently ignored here. Set kv_cache_config.use_kv_cache_manager_v2=True " + "to use it.") self.is_linear_attention = linear_attention_metadata is not None # Calculate kv cache blocks for each window size diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 6a1b0ec24ce3..4f99ac831aab 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -3869,6 +3869,16 @@ def _reject_renamed_block_reuse_policy(cls, data: Any) -> Any: "`policy` is 'per_conversation'.") +# Attention backends that can address a scratch block. TRTLLM consumes +# copy_batch_block_offsets, which carries separate K and V page indices; +# FlashInfer builds a flat page table and uses the PER_LAYER addressing path. +# Every other backend reads raw base page indices, where a scratch block -- +# which owns no per-block page -- appears as an invalid page. Defined here so +# the 'auto' resolver in llm_utils and validate_swa_scratch_reuse below cannot +# drift apart when a backend gains scratch-page addressing. +SWA_SCRATCH_CAPABLE_ATTN_BACKENDS = ("TRTLLM", "FLASHINFER") + + @PybindMirror.mirror_pybind_fields(_KvCacheConfig) class KvCacheConfig(StrictBaseModel, PybindMirror): """Configuration for the KV cache.""" @@ -4026,11 +4036,15 @@ class KvCacheConfig(StrictBaseModel, PybindMirror): "does not declare one.") # This is a pure python field, not a pybind field. It is only for the Pytorch backend. - enable_swa_scratch_reuse: bool = Field( - default=False, + enable_swa_scratch_reuse: bool | Literal["auto"] = Field( + default="auto", status="prototype", description= - "Whether KV cache manager v2 uses SWA scratch reuse during prefill.") + "Whether KV cache manager v2 uses SWA scratch reuse during prefill. " + "'auto' enables it wherever it is supported and can save memory: KV " + "cache manager v2, an attention backend that can address scratch " + "pages (TRTLLM or FlashInfer), and a model with sliding-window " + "layers. It is a no-op everywhere else.") kv_cache_event_hash_algo: Literal[ "auto", "v1_block_key", "v2_sha256", "v2_sha256_64"] = Field( @@ -6305,6 +6319,49 @@ def warn_on_unstable_feature_usage(self) -> 'TorchLlmArgs': return self + @model_validator(mode='after') + def validate_swa_scratch_reuse(self) -> 'TorchLlmArgs': + """Reject or warn on configurations SWA scratch reuse cannot serve. + + Scratch blocks get no per-block KV page; their address is a rotating + sub-page inside a shared slot. Backends that consume + ``copy_batch_block_offsets`` get that for free. FlashInfer builds its + own flat page table and is supported through the PER_LAYER addressing + path, which additionally requires K and V to stay adjacent under the + rotation -- checked at first use, since only a flat table needs it. + Backends that read raw base page indices without either path are + rejected here. + + Only an explicit ``True`` is rejected. The default ``'auto'`` resolves + to False on those backends instead, so enabling scratch reuse by + default never turns a working configuration into an error. + """ + kv_cache_config = self.kv_cache_config + if kv_cache_config.enable_swa_scratch_reuse is not True: + return self + + if self.attn_backend not in SWA_SCRATCH_CAPABLE_ATTN_BACKENDS: + raise ValueError( + "kv_cache_config.enable_swa_scratch_reuse is not supported with " + f"attn_backend={self.attn_backend!r}. Only 'TRTLLM' (via " + "copy_batch_block_offsets) and 'FLASHINFER' (via PER_LAYER page " + "indices) can address scratch blocks; other backends read raw base " + "page indices, where scratch blocks appear as invalid pages. Set " + "enable_swa_scratch_reuse=False for this backend.") + + if (kv_cache_config.enable_block_reuse and + kv_cache_config.block_reuse_config.policy == "all_reusable"): + logger.warning( + "kv_cache_config.enable_swa_scratch_reuse is enabled together with " + "block_reuse_config.policy='all_reusable'. Under 'all_reusable' the " + "manager would otherwise commit out-of-window blocks for prefix " + "reuse, but scratch blocks use shared storage that is never " + "preserved, so only non-scratch blocks stay reusable. Use " + "block_reuse_config.policy='per_request' to avoid the reuse loss." + ) + + return self + @model_validator(mode='after') def validate_ray_worker_extension_cls(self) -> 'TorchLlmArgs': if self.ray_worker_extension_cls is not None and self.orchestrator_type != "ray": diff --git a/tensorrt_llm/llmapi/llm_utils.py b/tensorrt_llm/llmapi/llm_utils.py index bc3b8fc659ae..20e8bb1fa058 100644 --- a/tensorrt_llm/llmapi/llm_utils.py +++ b/tensorrt_llm/llmapi/llm_utils.py @@ -25,7 +25,8 @@ read_modelopt_quant_config, warn_if_inline_diverges) # yapf: disable -from .llm_args import (CalibConfig, CudaGraphConfig, DecodeCudaGraphConfig, +from .llm_args import (SWA_SCRATCH_CAPABLE_ATTN_BACKENDS, CalibConfig, + CudaGraphConfig, DecodeCudaGraphConfig, DraftTargetDecodingConfig, Eagle3DecodingConfig, EagleDecodingConfig, EncodeCudaGraphConfig, KvCacheConfig, LlmArgs, LookaheadDecodingConfig, @@ -675,6 +676,47 @@ def _resolve_kv_cache_manager_v2_auto(llm_args: 'TorchLlmArgs', return use_v2 +def _resolve_swa_scratch_reuse_auto(llm_args: 'TorchLlmArgs', + model_cls: Optional[type] = None, + pretrained_config: Any = None) -> bool: + """Resolve the 'auto' sentinel in kv_cache_config.enable_swa_scratch_reuse. + + 'auto' turns SWA scratch reuse on wherever the engine can actually run it: + a V2 KV cache manager and a scratch-capable attention backend. Whether the + model has sliding-window layers at all is decided by the manager, which is + the first place the per-layer window vector exists; without them the + feature is inert rather than wrong. + + ``use_kv_cache_manager_v2`` alone is not the right question. Some models are + routed to a V2 manager structurally: the sparse-attention registry picks + ``DeepseekV4CacheManager`` / ``MiniMaxM3KVCacheManagerV2`` from the + algorithm, ignoring the flag, and those models keep a V2 manager even when + the flag is demoted to False (disaggregated serving on a non-Python + transceiver, two-model speculative decoding). A declared ``"V2"`` preference + is the signal that survives that demotion. + + Must run after model defaults are applied (a model may select the attention + backend) and after ``_resolve_kv_cache_manager_v2_auto``, so both inputs are + concrete. An explicit user value is returned untouched. + """ + kv_cache_config = llm_args.kv_cache_config + setting = kv_cache_config.enable_swa_scratch_reuse + if setting != "auto": + return setting + + uses_v2 = kv_cache_config.use_kv_cache_manager_v2 is True + if not uses_v2 and model_cls is not None: + get_preferred = getattr(model_cls, + 'get_preferred_kv_cache_manager_version', None) + uses_v2 = (get_preferred is not None + and get_preferred(pretrained_config) == "V2") + + enable = (uses_v2 + and llm_args.attn_backend in SWA_SCRATCH_CAPABLE_ATTN_BACKENDS) + kv_cache_config.enable_swa_scratch_reuse = enable + return enable + + def _resolve_transceiver_runtime_auto(llm_args: 'TorchLlmArgs', model_cls: Optional[type] = None, pretrained_config: Any = None) -> None: diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi index 2af6bc537252..8048625b820e 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi @@ -92,6 +92,8 @@ class KVCacheIterationStatsDelta: iter_intra_device_copy_bytes: int = 0 iter_host_dropped_blocks: int = 0 iter_host_dropped_bytes: int = 0 + iter_scratch_blocks: int = 0 + iter_scratch_slots_in_use: int = 0 @dataclass(slots=True) class SsmSnapshotIterationStatsDelta: diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py index 3bd6601b54b8..992a873b1fac 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py @@ -500,6 +500,7 @@ def _record_resize_pending_allocations( beam_width: BeamIndex, excluded_ranges: TypedIndexList[LifeCycleId, HalfOpenRange[BlockOrdinal]], count_as_generation: bool, + excluded_is_scratch: bool = False, ) -> None: record_manager_stats = self._should_record_manager_stats() record_request_stats = self._should_record_request_stats() @@ -523,9 +524,34 @@ def _record_resize_pending_allocations( record_manager_stats=record_manager_stats, record_request_stats=record_request_stats, ) + if excluded_is_scratch: + # The piece dropped between the two recorded outer ranges is + # exactly the scratch block count for this lifecycle. + self._record_scratch_iteration_stats( + lc_idx, + len(intersect(excluded_ranges[lc_idx], HalfOpenRange(block_begin, block_end))), + ) if changed: self.manager.mark_stats_dirty(self.id) + def _record_scratch_iteration_stats(self, life_cycle: LifeCycleId, num_blocks: int) -> None: + """Attribute scratch-served blocks and the slot occupancy behind them. + + ``iter_scratch_blocks`` is a count and accumulates; ``iter_scratch_slots_in_use`` + is the current occupancy of this lifecycle's scratch slots and is a gauge. + See ``KVCacheIterationStatsDelta`` for the contract -- the two must not be + aggregated the same way. + """ + if num_blocks <= 0: + return + self._record_direct_iteration_stats( + life_cycle, + KVCacheIterationStatsDelta( + iter_scratch_blocks=num_blocks, + iter_scratch_slots_in_use=len(self._scratch_slots[life_cycle]), + ), + ) + @staticmethod def _has_reuse_source(page: BlockPage) -> bool: if page is None or not isinstance(page.page, CommittedPage): @@ -947,6 +973,7 @@ def resize(self, capacity: int | None, history_length: int | None = None) -> boo beam_width, excluded_ranges, record_generation_alloc_stats, + excluded_is_scratch=enable_scratch, ) for ordinal in typed_range(old_num_blocks, new_num_blocks): block = make_typed( diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_stats.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_stats.py index 03d95d745df0..d55981eb361e 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_stats.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_stats.py @@ -69,6 +69,21 @@ class KVCacheIterationStatsDelta(_StatsDeltaMixin): # _storage_manager._prepare_free_slots when is_last_level(lvl). iter_host_dropped_blocks: int = 0 iter_host_dropped_bytes: int = 0 + # SWA scratch reuse attribution. Scratch blocks are excluded from + # iter_alloc_* by design (they take no per-request KV page), so without + # these enabling scratch just makes iter_alloc_new_blocks drop with no + # attribution: iter_alloc_new_blocks drops by exactly iter_scratch_blocks. + # + # Different semantics, carried in the names: + # iter_scratch_blocks COUNT. Additive; summing over iterations gives + # the total blocks ever served from scratch. + # iter_scratch_slots_in_use GAUGE. Occupancy, not flow. Summed across + # lifecycles within one iteration, reset with + # the rest of the delta each iteration. + # Accumulating it across iterations is + # meaningless. + iter_scratch_blocks: int = 0 + iter_scratch_slots_in_use: int = 0 @property def iter_cache_hit_rate(self) -> float: diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index d5e4ae74a35c..ab390ab91a2e 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -751,8 +751,10 @@ "path": "kv_cache_config.enable_partial_reuse" }, { - "allowed_values": [], - "annotation": "", + "allowed_values": [ + "auto" + ], + "annotation": "Union[bool, Literal['auto']]", "converter": "", "kind": "value", "path": "kv_cache_config.enable_swa_scratch_reuse" diff --git a/tests/integration/defs/llmapi/test_llm_api_connector.py b/tests/integration/defs/llmapi/test_llm_api_connector.py index 6a0fbd26c7ac..160f78832d0e 100644 --- a/tests/integration/defs/llmapi/test_llm_api_connector.py +++ b/tests/integration/defs/llmapi/test_llm_api_connector.py @@ -576,6 +576,59 @@ def test_connector_rejects_unsupported_config(enforce_single_worker, model_fn(**llm_kwargs) +@pytest.mark.threadleak(enabled=False) +@pytest.mark.parametrize("use_kv_cache_manager_v2", [True, "auto"]) +def test_connector_with_kv_cache_manager_v2(enforce_single_worker, + model_with_connector, + use_kv_cache_manager_v2): + """The connector and KV cache manager v2 must produce a working engine. + + v2 cannot drive the connector's per-block load/save hooks, so + `_validate_or_fallback_kv_cache_manager_v2` selects v1 instead. That is a + fallback, not an error: the run must still come up and generate, and the + connector must be driven exactly as it is without the v2 request. + """ + from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import \ + KVCacheManagerV2 + + NUM_TOKENS = 8 + + model_fn, scheduler, worker = model_with_connector + + model = model_fn(kv_cache_config=KvCacheConfig( + free_gpu_memory_fraction=0.1, + use_kv_cache_manager_v2=use_kv_cache_manager_v2)) + + # The connector is registered against the selected manager's pool, so this + # asserts the fallback actually happened rather than that it was requested. + assert worker.register_kv_caches.call_count == 1 + + scheduler.get_num_new_matched_tokens.return_value = 0, False + worker.get_finished.return_value = [], [] + + sampling_params = SamplingParams(max_tokens=NUM_TOKENS, ignore_eos=True) + outputs = generate_and_sleep(model, ["Hello, world"], sampling_params) + + assert len(outputs[0].outputs[0].token_ids) == NUM_TOKENS + + # The connector still sees the whole request lifecycle. + assert scheduler.update_state_after_alloc.call_count == 1 + assert scheduler.build_connector_meta.call_count == NUM_TOKENS + assert scheduler.request_finished.call_count == 1 + + kv_cache_manager = _get_kv_cache_manager(model) + assert not isinstance(kv_cache_manager, KVCacheManagerV2), ( + "KVCacheManagerV2 cannot serve the KV connector: it exposes no " + "single primary pool to register, no per-request block hashes, and no " + "block priorities. Selecting it here would hand the connector worker " + "block ids it cannot address.") + + +def _get_kv_cache_manager(model): + """The KV cache manager actually built for this engine.""" + return model._executor.engine.kv_cache_manager + + @pytest.mark.threadleak(enabled=False) def test_connector_e2e_persistent_cache(enforce_single_worker): """Test e2e KV cache connector using PersistentKvCacheConnector from examples. diff --git a/tests/integration/test_lists/test-db/l0_a10.yml b/tests/integration/test_lists/test-db/l0_a10.yml index 5b39eea23efa..89eea09cfc5e 100644 --- a/tests/integration/test_lists/test-db/l0_a10.yml +++ b/tests/integration/test_lists/test-db/l0_a10.yml @@ -136,6 +136,8 @@ l0_a10: - llmapi/test_llm_api_connector.py::test_connector_rejects_unsupported_config[host_offloading] - llmapi/test_llm_api_connector.py::test_connector_rejects_unsupported_config[beam_search] - llmapi/test_llm_api_connector.py::test_connector_rejects_unsupported_config[attention_dp] + - llmapi/test_llm_api_connector.py::test_connector_with_kv_cache_manager_v2[True] + - llmapi/test_llm_api_connector.py::test_connector_with_kv_cache_manager_v2[auto] - llmapi/test_llm_api_connector.py::test_connector_e2e_persistent_cache # third-party policy checks CPU-only - thirdparty/test_cmake_third_party.py::test_cmake_listfiles diff --git a/tests/unittest/_torch/executor/test_kv_cache_manager_v2.py b/tests/unittest/_torch/executor/test_kv_cache_manager_v2.py index ae79c8ef55a7..c7d50c27fa6c 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_manager_v2.py +++ b/tests/unittest/_torch/executor/test_kv_cache_manager_v2.py @@ -21,7 +21,11 @@ import pytest import torch -from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import BlockReusePolicy, KVCacheManagerV2 +from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import ( + BlockReusePolicy, + KVCacheManagerV2, + Role, +) from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests from tensorrt_llm.bindings import DataType from tensorrt_llm.bindings.internal.batch_manager import CacheType @@ -30,12 +34,16 @@ from tensorrt_llm.mapping import Mapping from tensorrt_llm.runtime.kv_cache_manager_v2 import ( DEFAULT_BEAM_INDEX, + AttentionLayerConfig, BatchDesc, + BufferConfig, DiskCacheTierConfig, GpuCacheTierConfig, HostCacheTierConfig, KVCacheDesc, KVCacheManagerConfig, + LayerId, + SsmLayerConfig, ) from tensorrt_llm.runtime.kv_cache_manager_v2._utils import init_cuda_once @@ -153,6 +161,7 @@ def build_cache_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, "_log_swa_scratch_summary"), ): manager = KVCacheManagerV2( kv_cache_config, @@ -218,7 +227,14 @@ def test_pool_ratio_overrides_constraints() -> None: assert config.constraints == [] -def test_default_uses_allocator_fallback() -> None: +def test_prefill_constraint_registered_without_avg_seq_len() -> None: + """The chunked-prefill constraint must not be gated behind avg_seq_len. + + Regression lock: when this constraint is missing, StorageManager falls back + to a DECODE-shaped BatchDesc whose scratch range is provably empty, so SWA + scratch reuse is inert for every model that does not set avg_seq_len, and + the SWA pool is sized from swa_floor_blocks alone. + """ config = _make_cache_config_for_test( KvCacheConfig(host_cache_size=0), max_batch_size=3, @@ -228,7 +244,31 @@ def test_default_uses_allocator_fallback() -> None: ) assert config.initial_pool_ratio is None + # typical_step stays opt-in: it needs avg_seq_len, which is workload knowledge. assert config.typical_step is None + assert config.constraints == [BatchDesc([KVCacheDesc(capacity=2048, history_length=0)])] + + +def test_prefill_constraint_includes_extra_kv_tokens() -> None: + config = _make_cache_config_for_test( + KvCacheConfig(host_cache_size=0), + max_batch_size=3, + max_seq_len=1024, + max_num_tokens=2048, + num_extra_kv_tokens=4, + ) + + assert config.constraints == [BatchDesc([KVCacheDesc(capacity=2052, history_length=0)])] + + +def test_no_prefill_constraint_without_max_num_tokens() -> None: + config = _make_cache_config_for_test( + KvCacheConfig(host_cache_size=0), + max_batch_size=3, + max_seq_len=1024, + max_num_tokens=None, + ) + assert config.constraints == [] @@ -368,6 +408,106 @@ def test_host_init_fallback_drops_only_host_tier(tmp_path) -> None: ] +def _attention_layer(layer_id: int, window: int | None) -> AttentionLayerConfig: + return AttentionLayerConfig( + layer_id=LayerId(layer_id), + buffers=[BufferConfig(role=Role.KEY, size=256)], + sliding_window_size=window, + ) + + +def _ssm_layer(layer_id: int) -> SsmLayerConfig: + return SsmLayerConfig( + layer_id=LayerId(layer_id), + buffers=[BufferConfig(role="ssm_state", size=64)], + ) + + +def _run_swa_scratch_summary( + layers: list[object], + *, + slots_without: list[int], + slots_with: list[int], + enabled: bool, +) -> tuple[list[str], list[str]]: + """Drive ``_log_swa_scratch_summary`` and return its warning/debug lines. + + The method is diagnostics-only, but it runs unconditionally at the end of + ``__init__``, so anything it cannot read aborts engine construction. Calling + it unbound keeps the check on the layer walk and the saving arithmetic + themselves, with no GPU and no real manager needed. + """ + manager = Mock() + manager.impl.get_layer_group_id.side_effect = lambda _: 0 + manager.kv_cache_manager_py_config = SimpleNamespace( + layers=layers, + constraints=[BatchDesc([KVCacheDesc(capacity=TOKENS_PER_BLOCK, history_length=0)])], + ) + manager.enable_swa_scratch_reuse = enabled + manager.tokens_per_block = TOKENS_PER_BLOCK + manager.max_num_tokens = TOKENS_PER_BLOCK + + module = "tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2" + with ( + patch(f"{module}._introspection") as introspection, + patch(f"{module}.logger") as logger, + ): + introspection.swa_life_cycle_ids.return_value = [0] + introspection.pool_group_index.return_value = 0 + introspection.compute_slots_for_batch.side_effect = [slots_without, slots_with] + KVCacheManagerV2._log_swa_scratch_summary(manager) + + return ( + [str(call.args[0]) for call in logger.warning.call_args_list], + [str(call.args[0]) for call in logger.debug.call_args_list], + ) + + +def test_swa_scratch_summary_skips_ssm_layers() -> None: + """A hybrid SSM + SWA engine must survive the startup summary. + + ``SsmLayerConfig`` carries no ``sliding_window_size``. Reading it unguarded + raised ``AttributeError`` out of ``__init__`` for any Mamba-hybrid model + that also had a real attention window. + """ + warnings, _ = _run_swa_scratch_summary( + [_attention_layer(0, window=TOKENS_PER_BLOCK), _ssm_layer(1)], + slots_without=[8], + slots_with=[4], + enabled=True, + ) + + assert warnings == [] + + +def test_swa_scratch_summary_warns_when_a_real_saving_is_declined() -> None: + warnings, _ = _run_swa_scratch_summary( + [_attention_layer(0, window=TOKENS_PER_BLOCK)], + slots_without=[8], + slots_with=[4], + enabled=False, + ) + + assert any("would cut windowed slots by up to 50%" in message for message in warnings) + + +def test_swa_scratch_summary_treats_a_slot_increase_as_no_saving() -> None: + """A rise in slot count is not a saving and must not be reported as one. + + Counting it left ``best_saving_pct`` at 0, so the declined-saving warning + advertised "up to 0%" while the inert-configuration branch never fired. + """ + warnings, debugs = _run_swa_scratch_summary( + [_attention_layer(0, window=TOKENS_PER_BLOCK)], + slots_without=[4], + slots_with=[8], + enabled=False, + ) + + assert not any("would cut windowed slots" in message for message in warnings) + assert any("scratch reuse is inert" in message for message in debugs) + + def test_extra_tokens_are_in_context_capacity() -> None: config = _make_cache_config_for_test( KvCacheConfig(avg_seq_len=264), @@ -832,3 +972,265 @@ def test_disagg_role_mapper_kinds_default_to_indexed(): Role.ALL: MapperKind.INDEXED, Role.INDEX_KEY: MapperKind.REPLICATED, } + + +# --------------------------------------------------------------------------- +# SWA scratch reuse: PER_LAYER flat page-index rotation. +# +# This is the arithmetic that addresses a scratch block on the FlashInfer path. +# It is the highest-risk code in the feature because it fails *silently*: a +# wrong index reads another layer's KV rather than raising, so an end-to-end run +# still exits 0 with plausible-looking output. The bug actually hit during +# Gemma4 bring-up (a layer_idx-less lookup yielding BAD_PAGE_INDEX) was found +# only by an illegal memory access on a B200, which is far too late and far too +# expensive a feedback loop for integer arithmetic. +# +# These tests pin the invariants the flat page table depends on, on a real +# Gemma4-12B-shaped configuration, with no GPU and no model. +# --------------------------------------------------------------------------- + +# Gemma4-12B: 48 layers, 40 sliding (W=1024) / 8 full, K and V per layer. +GEMMA4_NUM_SWA_LAYERS = 40 +GEMMA4_KV_FACTOR = 2 +# One slot holds `scale` sub-pages: kv_factor per layer across the shared group. +GEMMA4_SCALE = GEMMA4_NUM_SWA_LAYERS * GEMMA4_KV_FACTOR +# Each scratch block advances by one K/V pair. +GEMMA4_SCRATCH_PAGES_PER_BLOCK = GEMMA4_KV_FACTOR + + +def _reference_flat_index(position, scratch_pages, scale, layer_offset, slot_ids, div_factor): + """Independent restatement of the device kernel's arithmetic. + + Deliberately written as a scalar loop from the formula rather than by + calling the implementation, so agreement is evidence rather than tautology. + """ + total = position * scratch_pages + slot = int(slot_ids[total // scale]) + sub = (total % scale + layer_offset) % scale + return (slot * scale + sub) // div_factor + + +def _k_layer_offset(layer_idx): + return layer_idx * GEMMA4_KV_FACTOR + + +class TestSwaScratchFlatIndexRotation: + """Correctness of compute_scratch_flat_page_indices on a Gemma4 shape.""" + + NUM_BLOCKS = 41 # a realistic prefill scratch range (~2340 tokens, W=1024, tpb=32) + SLOT_IDS = tuple(range(7, 7 + 8)) # arbitrary non-contiguous-looking slot ids + + def _indices(self, layer_idx, div_factor=1, count=None): + from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import ( + compute_scratch_flat_page_indices, + ) + + return compute_scratch_flat_page_indices( + 0, + self.NUM_BLOCKS if count is None else count, + GEMMA4_SCRATCH_PAGES_PER_BLOCK, + GEMMA4_SCALE, + _k_layer_offset(layer_idx), + self.SLOT_IDS, + div_factor, + ) + + def test_matches_device_kernel_formula_for_every_swa_layer(self): + """Host rotation must equal the device kernel's, for all 40 SWA layers.""" + for layer_idx in range(GEMMA4_NUM_SWA_LAYERS): + got = self._indices(layer_idx) + expected = [ + _reference_flat_index( + pos, + GEMMA4_SCRATCH_PAGES_PER_BLOCK, + GEMMA4_SCALE, + _k_layer_offset(layer_idx), + self.SLOT_IDS, + 1, + ) + for pos in range(self.NUM_BLOCKS) + ] + assert got.tolist() == expected, f"layer {layer_idx} diverges from the kernel formula" + + def test_v_stays_exactly_one_subpage_after_k(self): + """The precondition a flat page table cannot express any other way. + + A flat table carries one index per block plus a kv_factor axis, so it can + only address V if V remains K+1 *after* the rotation. If this breaks, + attention reads K as V and produces silently wrong output rather than an + error. _validate_per_layer_kv_adjacency promises this; here it is checked + against the arithmetic that has to honour it. + """ + for layer_idx in range(GEMMA4_NUM_SWA_LAYERS): + k_off = _k_layer_offset(layer_idx) + k = self._indices(layer_idx) + v = [ + _reference_flat_index( + pos, GEMMA4_SCRATCH_PAGES_PER_BLOCK, GEMMA4_SCALE, k_off + 1, self.SLOT_IDS, 1 + ) + for pos in range(self.NUM_BLOCKS) + ] + assert [b - a for a, b in zip(k.tolist(), v)] == [1] * self.NUM_BLOCKS, ( + f"layer {layer_idx}: V is not adjacent to K under the scratch rotation" + ) + + def test_no_two_swa_layers_alias_the_same_subpage(self): + """Distinct layers must never resolve to the same page for a block. + + Aliasing is the failure mode that corrupts KV without any crash: two + layers would read and write each other's cache. With scale == 40 layers + x kv_factor, all 40 layers must land on 40 distinct K sub-pages. + """ + per_layer = [ + self._indices(layer_idx).tolist() for layer_idx in range(GEMMA4_NUM_SWA_LAYERS) + ] + for position in range(self.NUM_BLOCKS): + seen = {indices[position] for indices in per_layer} + assert len(seen) == GEMMA4_NUM_SWA_LAYERS, ( + f"block position {position}: only {len(seen)} distinct pages for " + f"{GEMMA4_NUM_SWA_LAYERS} layers -- layers alias each other's KV" + ) + + def test_indices_stay_inside_the_addressed_slots(self): + """Every index must fall inside a slot the descriptor actually holds.""" + valid = {slot * GEMMA4_SCALE + sub for slot in self.SLOT_IDS for sub in range(GEMMA4_SCALE)} + for layer_idx in range(GEMMA4_NUM_SWA_LAYERS): + assert set(self._indices(layer_idx).tolist()) <= valid, ( + f"layer {layer_idx} produced an index outside the descriptor's slots" + ) + + def test_rotation_advances_with_block_position(self): + """The rotation is the reason SHARED addressing cannot work. + + If a layer's sub-page were fixed across block positions it could be + folded into a base pointer and none of the PER_LAYER machinery would be + needed. Assert it genuinely moves, so this test fails if someone + "simplifies" the rotation away. + """ + idx = self._indices(layer_idx=3) + sub_pages = [int(i) % GEMMA4_SCALE for i in idx.tolist()] + assert len(set(sub_pages)) > 1, "sub-page did not rotate with block position" + + def test_kv_factor_division_preserves_pairing(self): + """div_factor halves the index space; K must stay kv_factor-aligned. + + The flat table indexes block-granular entries, so the caller divides by + kv_factor. That is only sound when K is kv_factor-aligned -- one of the + conditions _validate_per_layer_kv_adjacency enforces. + """ + for layer_idx in range(GEMMA4_NUM_SWA_LAYERS): + raw = self._indices(layer_idx, div_factor=1).tolist() + halved = self._indices(layer_idx, div_factor=GEMMA4_KV_FACTOR).tolist() + assert all(r % GEMMA4_KV_FACTOR == 0 for r in raw), ( + f"layer {layer_idx}: K index is not kv_factor-aligned, so dividing by " + "kv_factor would collapse K and V onto the same entry" + ) + assert halved == [r // GEMMA4_KV_FACTOR for r in raw] + + def test_empty_range_is_empty(self): + assert self._indices(layer_idx=0, count=0).tolist() == [] + + +class TestSwaScratchSegmentClamping: + """Range/segment clamping in apply_scratch_to_block_segment. + + The scratch range and a request's block count are computed independently, so + they can fail to overlap. Getting the clamp wrong does not raise -- it shifts + the wrong blocks by layer_offset and leaves them pointing at another layer's + pages, which reads as plausible output. These cases are cheap to pin and + impossible to notice at runtime. + """ + + SCALE = GEMMA4_SCALE + SPB = GEMMA4_SCRATCH_PAGES_PER_BLOCK + SLOT_IDS = tuple(range(7, 15)) + LAYER_OFFSET = 6 # layer 3, K + + def _apply(self, values, beg, end, div_factor=1): + import numpy as np + + from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import ( + apply_scratch_to_block_segment, + ) + + seg = np.asarray(values, dtype=np.int32).copy() + apply_scratch_to_block_segment( + seg, + beg, + end, + self.SPB, + self.SCALE, + self.LAYER_OFFSET, + self.SLOT_IDS, + div_factor, + ) + return seg.tolist() + + def test_range_entirely_before_segment_shifts_every_block(self): + """beg < end <= 0: nothing is scratch, so every block just gains the offset. + + Regression: a naive ``seg[hi:]`` with a negative ``hi`` indexes from the + end of the array and shifts only a suffix, silently leaving the leading + blocks addressed as if the buffer were still SHARED-based. + """ + values = [10, 11, 12, 13] + got = self._apply(values, beg=-3, end=-1) + assert got == [v + self.LAYER_OFFSET for v in values] + + def test_empty_range_shifts_every_block(self): + values = [10, 11, 12, 13] + assert self._apply(values, beg=2, end=2) == [v + self.LAYER_OFFSET for v in values] + + def test_range_entirely_after_segment_shifts_every_block(self): + values = [10, 11, 12, 13] + assert self._apply(values, beg=9, end=12) == [v + self.LAYER_OFFSET for v in values] + + def test_bad_page_index_is_never_shifted(self): + """BAD_PAGE_INDEX must stay the sentinel; shifting it makes it a real page.""" + from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import BAD_PAGE_INDEX + + got = self._apply([BAD_PAGE_INDEX, 11, BAD_PAGE_INDEX], beg=5, end=5) + assert got[0] == BAD_PAGE_INDEX and got[2] == BAD_PAGE_INDEX + assert got[1] == 11 + self.LAYER_OFFSET + + def test_partial_overlap_splits_scratch_and_non_scratch(self): + """Only blocks inside the range rotate; the rest are shifted.""" + from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import ( + compute_scratch_flat_page_indices, + ) + + values = [10, 11, 12, 13, 14] + got = self._apply(values, beg=1, end=3) + rotated = compute_scratch_flat_page_indices( + 0, 2, self.SPB, self.SCALE, self.LAYER_OFFSET, self.SLOT_IDS, 1 + ).tolist() + assert got[0] == 10 + self.LAYER_OFFSET + assert got[1:3] == rotated + assert got[3:] == [13 + self.LAYER_OFFSET, 14 + self.LAYER_OFFSET] + + def test_range_clipped_to_segment_does_not_false_trip_slot_guard(self): + """A range extending past the request's blocks is clipped, not rejected. + + Only the blocks actually addressed consume slots, so the bound must be + checked on the clipped range. Checking [beg, end) instead would reject + descriptors that are perfectly serviceable. + """ + from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import ( + compute_scratch_flat_page_indices, + ) + + # 2 blocks addressed needs 1 slot; the unclipped range would demand 13. + got = self._apply([10, 11], beg=0, end=500) + assert ( + got + == compute_scratch_flat_page_indices( + 0, 2, self.SPB, self.SCALE, self.LAYER_OFFSET, self.SLOT_IDS, 1 + ).tolist() + ) + + def test_insufficient_slots_raises_with_numbers(self): + import pytest as _pytest + + with _pytest.raises(ValueError, match="scratch slot"): + # 400 blocks x 2 pages / scale 80 needs 10 slots; only 8 provided. + self._apply(list(range(400)), beg=0, end=400) diff --git a/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py b/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py index 3abeb9471426..cf6e356e7d79 100644 --- a/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py +++ b/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py @@ -159,10 +159,11 @@ def test_deepseek_v4_fused_hc_default_enabled(monkeypatch): def test_deepseek_v4_kv_cache_defaults_and_v2_preference(): defaults = DeepseekV4ForCausalLM.get_model_defaults(None) + # SWA scratch reuse is a global default resolved by + # llm_utils._resolve_swa_scratch_reuse_auto, not a per-model opt-in. assert defaults == { "kv_cache_config": { "tokens_per_block": 128, - "enable_swa_scratch_reuse": True, } } assert DeepseekV4ForCausalLM.get_preferred_kv_cache_manager_version() == "V2" @@ -177,7 +178,6 @@ class LlmArgs: assert defaults == { "kv_cache_config": { "tokens_per_block": 256, - "enable_swa_scratch_reuse": True, } } diff --git a/tests/unittest/executor/test_stats_serializer.py b/tests/unittest/executor/test_stats_serializer.py index 244f30c91153..dbbf49d8480b 100644 --- a/tests/unittest/executor/test_stats_serializer.py +++ b/tests/unittest/executor/test_stats_serializer.py @@ -117,6 +117,11 @@ def _make_mock_kv_iter_stats( iter_intra_device_copy_bytes=8192, iter_host_dropped_blocks=0, iter_host_dropped_bytes=0, + # Non-zero on purpose: with 0 here every assertion below would pass + # even if the field were dropped from a view entirely, which is exactly + # how the pool-group allowlist omission went unnoticed. + iter_scratch_blocks=129, + iter_scratch_slots_in_use=4, ) return {window_size: s} @@ -437,6 +442,14 @@ def test_serializer_with_v2_pool_group_stats(self): assert pool_group["primaryPeakEvictableNumBlocks"] == 4 assert pool_group["iterGenAllocBlocks"] == 2 assert "iterReusedBlocks" not in pool_group + # SWA scratch reuse is a per-pool-group phenomenon -- a pool group is + # exactly the unit a scratch slot is shared within -- so the counters + # must survive the pool-group key filter. They were originally missing + # from KV_CACHE_ITERATION_STATS_POOL_GROUP_KEYS, which made the feature + # read as permanently inert to anything consuming this view. + assert pool_group["iterScratchBlocks"] == 129 + assert pool_group["iterScratchSlotsInUse"] == 4 + assert d["kvCacheIterationStats"]["16"]["iterScratchBlocks"] == 129 assert "iterMissedBlocks" not in pool_group assert "iterCacheHitRate" not in pool_group # A hot pool group cannot index a cold level, so it reports no cold blocks either. @@ -522,3 +535,42 @@ def test_v2_peak_block_stats_reset_tracks_interval_peak(self): assert [stats.available for stats in secondary_peak] == [4, 5] assert [stats.unavailable for stats in secondary_peak] == [1, 0] assert [stats.evictable for stats in secondary_peak] == [1, 0] + + +def test_pool_group_keys_cover_every_iteration_field(): + """Every serialized iteration field must be classified, not forgotten. + + ``KV_CACHE_ITERATION_STATS_POOL_GROUP_KEYS`` is a hand-maintained allowlist, + so a field added to ``serialize_kv_cache_iteration_stats`` without a + corresponding allowlist entry silently vanishes from + ``kvCacheIterationStatsByPoolGroup`` -- the view most consumers reach for, + since a pool group is the unit a scratch slot is shared within. That is what + happened to ``iterScratchBlocks``/``iterScratchSlotsInUse``: the counters were + emitted correctly by the manager and read as a permanent 0 by every + consumer of the pool-group view, which made SWA scratch reuse look inert. + + This test forces a decision: a new field either belongs in the pool-group + view (add it to the allowlist) or it does not (name it, with a reason, in + KV_CACHE_ITERATION_STATS_NOT_PER_POOL_GROUP). Forgetting both fails here. + """ + from tensorrt_llm._torch.pyexecutor.kv_cache_stats import ( + KV_CACHE_ITERATION_STATS_NOT_PER_POOL_GROUP, + KV_CACHE_ITERATION_STATS_POOL_GROUP_KEYS, + serialize_kv_cache_iteration_stats, + ) + + emitted = set(serialize_kv_cache_iteration_stats(_make_mock_kv_iter_stats()[16])) + allowed = set(KV_CACHE_ITERATION_STATS_POOL_GROUP_KEYS) + excluded = set(KV_CACHE_ITERATION_STATS_NOT_PER_POOL_GROUP) + + unclassified = emitted - allowed - excluded + assert not unclassified, ( + f"iteration stat field(s) {sorted(unclassified)} are neither in the " + "pool-group allowlist nor explicitly excluded. Add them to " + "KV_CACHE_ITERATION_STATS_POOL_GROUP_KEYS, or to " + "KV_CACHE_ITERATION_STATS_NOT_PER_POOL_GROUP with a reason." + ) + # The allowlist must not name fields the serializer does not emit, or the + # pool-group projection raises KeyError at runtime. + assert not (allowed - emitted), sorted(allowed - emitted) + assert not (allowed & excluded), sorted(allowed & excluded) 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 77a9c1cffaae..c87fa51c63ea 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 @@ -3805,6 +3805,107 @@ def test_constraint_with_scratch_accounts_for_scratch(self): kv.close() manager.shutdown() + # ----- scratch savings on real model-shaped lifecycles ----- + + def _slots_for_batch(self, cfg: KVCacheManagerConfig, batch: BatchDesc, scratch: bool): + """Per-pool-group slot counts for ``batch``, with and without scratch. + + Uses the same pure ``compute_slots_for_batch`` entry point that the + startup counterfactual log reports, so the numbers asserted here are + exactly the numbers a user sees at engine init. + """ + manager = KVCacheManager(cfg) + try: + return _introspection.compute_slots_for_batch( + manager, + batch, + self.TOKENS_PER_BLOCK, + SwaScratchReuseConfig() if scratch else None, + ) + finally: + manager.shutdown() + + def _expected_prefill_slots(self, capacity: int, num_windowed_layers: int): + """(without_scratch, with_scratch) slots for one fresh context request. + + Mirrors compute_scratch_range: scratch = stale_at_capacity INTERSECT + input_blocks, and input_blocks is the whole request when history is 0. + """ + tpb = self.TOKENS_PER_BLOCK + total_blocks = div_up(capacity, tpb) + num_sink_blocks = div_up(self.SINK_TOKENS, tpb) + stale_beg = min(total_blocks, num_sink_blocks) + stale_end = max(stale_beg, (capacity + 1 - self.WINDOW_SIZE) // tpb) + num_scratch = min(stale_end, total_blocks) - stale_beg + with_scratch = (total_blocks - num_scratch) + div_up(num_scratch, num_windowed_layers) + return total_blocks, with_scratch, num_scratch + + def test_vswa_prefill_scratch_reduces_windowed_slots(self): + """GPT-OSS / Gemma3 shape: N windowed layers alternating with N full layers. + + The windowed pool group shrinks by ~frac_max on the prefill batch; the + full-attention pool group is untouched. + """ + num_windowed = 12 + capacity = 8192 + cfg = self._make_config( + gpu_quota=8 << 30, + num_windowed_layers=num_windowed, + num_full_layers=12, + ) + batch = BatchDesc(kv_caches=[KVCacheDesc(capacity=capacity, history_length=0)]) + without = self._slots_for_batch(cfg, batch, scratch=False) + with_scratch = self._slots_for_batch(cfg, batch, scratch=True) + + exp_without, exp_with, num_scratch = self._expected_prefill_slots(capacity, num_windowed) + self.assertGreater(num_scratch, 0, "prefill batch must produce a non-empty scratch range") + self.assertEqual(without[0], exp_without) + self.assertEqual(with_scratch[0], exp_with) + # Full-attention pool group has no window, so scratch is a no-op there. + self.assertEqual(with_scratch[1], without[1]) + + def test_pure_swa_prefill_scratch_reduces_slots(self): + """Mistral shape: every layer windowed, so one lifecycle holds all N. + + frac_max = 1/N is maximal here, so the saving is larger than the VSWA + case even though ``is_vswa`` would be False for this model. + """ + num_windowed = 24 + capacity = 8192 + cfg = self._make_config( + gpu_quota=8 << 30, + num_windowed_layers=num_windowed, + num_full_layers=0, + ) + batch = BatchDesc(kv_caches=[KVCacheDesc(capacity=capacity, history_length=0)]) + without = self._slots_for_batch(cfg, batch, scratch=False) + with_scratch = self._slots_for_batch(cfg, batch, scratch=True) + + exp_without, exp_with, num_scratch = self._expected_prefill_slots(capacity, num_windowed) + self.assertGreater(num_scratch, 0) + self.assertEqual(len(without), 1, "pure SWA model has a single pool group") + self.assertEqual(without[0], exp_without) + self.assertEqual(with_scratch[0], exp_with) + # More layers per lifecycle -> smaller frac_max -> bigger saving. + self.assertLess(with_scratch[0] / without[0], 0.25) + + def test_decode_shaped_batch_yields_no_scratch(self): + """The fallback DECODE BatchDesc has an empty scratch range. + + This is why the chunked-prefill constraint has to be registered + unconditionally: without it, StorageManager falls back to this shape + and scratch reuse is provably inert. + """ + cfg = self._make_config( + gpu_quota=8 << 30, + num_windowed_layers=12, + num_full_layers=12, + ) + batch = BatchDesc(kv_caches=[KVCacheDesc(capacity=2049, history_length=2048)]) + without = self._slots_for_batch(cfg, batch, scratch=False) + with_scratch = self._slots_for_batch(cfg, batch, scratch=True) + self.assertEqual(list(with_scratch), list(without)) + class TestScratchReuse(TestKVCacheManagerV2): """Tests for SWA prefill memory reuse (scratch slots).""" diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index 529851947843..249023f1f3fd 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -62,6 +62,7 @@ update_llm_args_with_extra_options) # fmt: on from tensorrt_llm.llmapi.llm_utils import (_resolve_kv_cache_manager_v2_auto, + _resolve_swa_scratch_reuse_auto, _resolve_transceiver_runtime_auto, apply_model_defaults_to_llm_args) from tensorrt_llm.llmapi.mm_encoder import MultimodalEncoder @@ -606,6 +607,166 @@ def get_model_defaults(cls, llm_args): assert "enable_block_reuse" in error_str or "max_tokens" in error_str +@pytest.mark.cpu_only +class TestSwaScratchReuseAutoResolution: + """SWA scratch reuse is on by default wherever the engine can run it.""" + + @staticmethod + def _args(**kwargs): + kv_cache_config = KvCacheConfig( + **{ + k: kwargs.pop(k) + for k in ("enable_swa_scratch_reuse", "use_kv_cache_manager_v2") + if k in kwargs + }) + return TorchLlmArgs(model="/tmp/dummy_model", + kv_cache_config=kv_cache_config, + **kwargs) + + @pytest.mark.parametrize("attn_backend", ["TRTLLM", "FLASHINFER"]) + def test_auto_enables_on_v2_with_capable_backend(self, attn_backend): + llm_args = self._args(use_kv_cache_manager_v2=True, + attn_backend=attn_backend) + + assert _resolve_swa_scratch_reuse_auto(llm_args) is True + assert llm_args.kv_cache_config.enable_swa_scratch_reuse is True + + @pytest.mark.parametrize("attn_backend", + ["VANILLA", "FLASHINFER_STAR_ATTENTION"]) + def test_auto_disables_on_backend_that_cannot_address_scratch( + self, attn_backend): + # A backend reading raw base page indices sees a scratch block as an + # invalid page. Turning the feature off beats failing the run. + llm_args = self._args(use_kv_cache_manager_v2=True, + attn_backend=attn_backend) + + assert _resolve_swa_scratch_reuse_auto(llm_args) is False + assert llm_args.kv_cache_config.enable_swa_scratch_reuse is False + + def test_auto_disables_on_v1(self): + # Scratch reuse is a V2-only feature. + llm_args = self._args(use_kv_cache_manager_v2=False, + attn_backend="TRTLLM") + + assert _resolve_swa_scratch_reuse_auto(llm_args) is False + assert llm_args.kv_cache_config.enable_swa_scratch_reuse is False + + def test_auto_disables_when_manager_version_unresolved(self): + # 'auto' must be resolved to a bool first; anything else is treated as + # "not V2" rather than silently enabling the feature. + llm_args = self._args(attn_backend="TRTLLM") + assert llm_args.kv_cache_config.use_kv_cache_manager_v2 == "auto" + + assert _resolve_swa_scratch_reuse_auto(llm_args) is False + + def test_auto_follows_a_declared_v2_preference(self): + """Sparse models keep a v2 manager even when the flag says v1. + + get_sparse_attn_kv_cache_manager picks DeepseekV4CacheManager / + MiniMaxM3KVCacheManagerV2 from the algorithm alone, so a v2 manager can + outlive a demoted use_kv_cache_manager_v2 (disagg on a non-Python + transceiver, two-model speculative decoding). The declared preference is + what survives that. + """ + + class _PrefersV2: + + @classmethod + def get_preferred_kv_cache_manager_version(cls, + pretrained_config=None): + return "V2" + + llm_args = self._args(use_kv_cache_manager_v2=False, + attn_backend="TRTLLM") + + assert _resolve_swa_scratch_reuse_auto(llm_args, _PrefersV2) is True + + def test_auto_ignores_a_v1_preference(self): + + class _PrefersV1: + + @classmethod + def get_preferred_kv_cache_manager_version(cls, + pretrained_config=None): + return "V1" + + llm_args = self._args(use_kv_cache_manager_v2=False, + attn_backend="TRTLLM") + + assert _resolve_swa_scratch_reuse_auto(llm_args, _PrefersV1) is False + + @pytest.mark.parametrize("architecture", [ + "DeepseekV4ForCausalLM", + "MiniMaxM3SparseForCausalLM", + ]) + def test_structurally_v2_models_keep_scratch_reuse(self, architecture): + """Regression guard for the models the sparse registry pins to v2.""" + from tensorrt_llm._torch.models.modeling_utils import \ + get_registered_model_class + + model_cls = get_registered_model_class(architecture) + assert model_cls is not None, f"{architecture} is not registered" + + llm_args = self._args(use_kv_cache_manager_v2=False, + attn_backend="TRTLLM") + + assert _resolve_swa_scratch_reuse_auto(llm_args, model_cls) is True + + @pytest.mark.parametrize("user_setting", [True, False]) + def test_explicit_value_is_untouched(self, user_setting): + llm_args = self._args(enable_swa_scratch_reuse=user_setting, + use_kv_cache_manager_v2=True, + attn_backend="TRTLLM") + + assert _resolve_swa_scratch_reuse_auto(llm_args) is user_setting + assert (llm_args.kv_cache_config.enable_swa_scratch_reuse + is user_setting) + + def test_resolution_is_idempotent(self): + llm_args = self._args(use_kv_cache_manager_v2=True, + attn_backend="TRTLLM") + + assert _resolve_swa_scratch_reuse_auto(llm_args) is True + assert _resolve_swa_scratch_reuse_auto(llm_args) is True + + def test_explicit_true_rejects_incapable_backend(self): + # The default degrades quietly; an explicit request must not. + with pytest.raises(ValidationError, match="enable_swa_scratch_reuse"): + self._args(enable_swa_scratch_reuse=True, attn_backend="VANILLA") + + def test_auto_does_not_reject_incapable_backend(self): + # Same config, default value: constructing it must still work. + self._args(attn_backend="VANILLA") + + @pytest.mark.parametrize("architecture", [ + "Gemma3ForCausalLM", + "Gemma4ForCausalLM", + "Gemma4ForConditionalGeneration", + "GptOssForCausalLM", + "DeepseekV4ForCausalLM", + ]) + def test_no_model_declares_a_per_model_opt_in(self, architecture): + """Enablement lives in one place: the 'auto' resolution above. + + A model that repeats ``enable_swa_scratch_reuse`` in + ``get_model_defaults`` would reintroduce a second, silently diverging + source of truth -- and, being an explicit value, would turn an + unsupported attention backend into a hard error for that model. + """ + from tensorrt_llm._torch.models.modeling_utils import \ + get_registered_model_class + + model_cls = get_registered_model_class(architecture) + assert model_cls is not None, f"{architecture} is not registered" + + defaults = getattr(model_cls, "get_model_defaults", + lambda _: {})(None) or {} + assert "enable_swa_scratch_reuse" not in ( + defaults.get("kv_cache_config") or {}), ( + f"{architecture} pins enable_swa_scratch_reuse in " + "get_model_defaults; remove it and rely on the global default") + + @pytest.mark.cpu_only class TestKvCacheManagerV2AutoResolution: """Test model preferences for the KV cache manager version.""" @@ -774,7 +935,13 @@ def test_KvCacheConfig_declaration(): assert KvCacheConfig().mamba_state_config.periodic_snapshot_interval == 0 assert KvCacheConfig().kv_cache_event_hash_algo == "auto" assert KvCacheConfig().block_reuse_config == BlockReuseConfig() - assert KvCacheConfig().enable_swa_scratch_reuse is False + assert KvCacheConfig().enable_swa_scratch_reuse == "auto" + assert KvCacheConfig( + enable_swa_scratch_reuse=True).enable_swa_scratch_reuse is True + assert KvCacheConfig( + enable_swa_scratch_reuse=False).enable_swa_scratch_reuse is False + with pytest.raises(ValidationError, match="enable_swa_scratch_reuse"): + KvCacheConfig(enable_swa_scratch_reuse="invalid") assert KvCacheConfig().use_kv_cache_manager_v2 == "auto" assert KvCacheConfig( use_kv_cache_manager_v2=True).use_kv_cache_manager_v2 is True @@ -819,7 +986,7 @@ def test_KvCacheConfig_declaration(): assert config.disk_cache_size == 2048 assert config.disk_cache_path == "/tmp" assert config.enable_swa_scratch_reuse is True - assert KvCacheConfig().enable_swa_scratch_reuse is False + assert KvCacheConfig().enable_swa_scratch_reuse == "auto" assert pybind_config.cross_kv_cache_fraction == 0.5 assert pybind_config.secondary_offload_min_priority == 1 assert pybind_config.event_buffer_max_size == 0