[TRTLLM-15944][feat] BREAKING: Enable SWA scratch reuse by default

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

cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp

@@ -717,7 +717,8 @@ void KvCache::_recordDroppedPages(std::vector<SharedPtr<Page>> const& pages, Cac
 }
 
 void KvCache::_recordResizePendingAllocations(BlockOrdinal blockBegin, BlockOrdinal blockEnd,
-    TypedVec<LifeCycleId, HalfOpenRange<BlockOrdinal>> const& excludedRanges, bool countAsGeneration)
+    TypedVec<LifeCycleId, HalfOpenRange<BlockOrdinal>> const& excludedRanges, bool countAsGeneration,
+    bool excludedIsScratch)
 {
     bool const recordManagerStats = _shouldRecordManagerStats();
     bool const recordRequestStats = _shouldRecordRequestStats();
@@ -743,13 +744,32 @@ 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<BlockOrdinal>{blockBegin, blockEnd}).length());
+        }
     }
     if (changed)
     {
         mManager->markStatsDirty(id);
     }
 }
 
+void KvCache::_recordScratchIterationStats(LifeCycleId lifeCycle, int64_t numBlocks)
+{
+    if (numBlocks <= 0)
+    {
+        return;
+    }
+    KVCacheIterationStatsDelta iterationStats;
+    iterationStats.iterScratchBlocks = numBlocks;
+    iterationStats.iterScratchSlotsInUse = static_cast<int64_t>(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<int> capacity, std::optional<int> 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(),

cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.h

@@ -492,7 +492,9 @@ class KvCache : public std::enable_shared_from_this<KvCache>
     void _recordDroppedPages(std::vector<SharedPtr<Page>> const& pages, CacheLevel cacheLevel);
     void _refreshGenerationAllocReady();
     void _recordResizePendingAllocations(BlockOrdinal blockBegin, BlockOrdinal blockEnd,
-        TypedVec<LifeCycleId, HalfOpenRange<BlockOrdinal>> const& excludedRanges, bool countAsGeneration);
+        TypedVec<LifeCycleId, HalfOpenRange<BlockOrdinal>> 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);

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;
     }
 };
 

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_<tbk::BlockKey>(m, "BlockKey")
         .def(nb::init<>())

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,14 +1199,17 @@ 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,
             nb::arg("iter_partial_reused_blocks") = 0, nb::arg("iter_missed_blocks") = 0,
             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_<kv::SsmSnapshotIterationStatsDelta>(m, "SsmSnapshotIterationStatsDelta")
         .def(

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` | `<class 'bool'>` | `value` |  |  |
 | `kv_cache_config.enable_kv_pool_rebalance` | `<class 'bool'>` | `value` |  |  |
 | `kv_cache_config.enable_partial_reuse` | `<class 'bool'>` | `value` |  |  |
-| `kv_cache_config.enable_swa_scratch_reuse` | `<class 'bool'>` | `value` |  |  |
+| `kv_cache_config.enable_swa_scratch_reuse` | `Union[bool, Literal['auto']]` | `value` |  | `auto` |
 | `kv_cache_config.event_buffer_max_size` | `<class 'int'>` | `value` |  |  |
 | `kv_cache_config.fp8_context_mla_kv_len_cap` | `Optional[int]` | `value` |  |  |
 | `kv_cache_config.free_gpu_memory_fraction` | `Optional[float]` | `value` |  |  |

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)
 

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