[TRTLLM-12229][fix] Fix MTP by scratch reuse rewind

tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py

@@ -45,6 +45,7 @@
     LayerId,
     PageIndexMode,
     ScratchDesc,
+    SwaScratchReuseConfig,
 )
 from tensorrt_llm.runtime.kv_cache_manager_v2 import KVCacheManagerConfig as KVCacheManagerConfigPy
 from tensorrt_llm.runtime.kv_cache_manager_v2._common import BAD_PAGE_INDEX
@@ -60,8 +61,8 @@
     is_overlap_compressor,
 )
 
-# Keep DSV4 scratch reuse opt-in so per-layer block tables can be tested
-# without scratch-page remapping in the default path.
+# Keep the env override so per-layer block tables can still be tested
+# without scratch-page remapping when needed.
 DSV4_ENABLE_SWA_SCRATCH_REUSE_ENV = "TRTLLM_DSV4_ENABLE_SWA_SCRATCH_REUSE"
 
 
@@ -119,7 +120,7 @@ def _get_index_mode(attn_type: DeepseekV4AttentionType) -> PageIndexMode:
 
 
 def _enable_swa_scratch_reuse_from_env() -> bool:
-    value = os.environ.get(DSV4_ENABLE_SWA_SCRATCH_REUSE_ENV, "0")
+    value = os.environ.get(DSV4_ENABLE_SWA_SCRATCH_REUSE_ENV, "1")
     return value.strip() == "1"
 
 
@@ -798,13 +799,20 @@ def _add_layer(
         if max_num_tokens is not None:
             constraints.append(BatchDesc([KVCacheDesc(capacity=max_num_tokens, history_length=0)]))
 
+        scratch_reuse_config = None
+        if self.enable_swa_scratch_reuse:
+            # Context requests will allocate num_extra_kv_tokens tokens for spec decoding.
+            # Cache manager should not take them into account when calculating scratch range.
+            # Therefore set max_rewind_len to num_extra_kv_tokens.
+            scratch_reuse_config = SwaScratchReuseConfig(max_rewind_len=self.num_extra_kv_tokens)
+
         return KVCacheManagerConfigPy(
             tokens_per_block=tokens_per_block,
             vocab_size=vocab_size,
             cache_tiers=cache_tiers,
             max_util_for_resume=kv_cache_config.max_util_for_resume,
             enable_stats=self.enable_stats,
-            enable_swa_scratch_reuse=self.enable_swa_scratch_reuse,
+            swa_scratch_reuse=scratch_reuse_config,
             layers=layers,
             typical_step=typical_step,
             constraints=constraints,

tensorrt_llm/_torch/pyexecutor/resource_manager.py

@@ -48,7 +48,8 @@
 # isort: off
 from tensorrt_llm.runtime.kv_cache_manager_v2 import (
     DEFAULT_BEAM_INDEX, AttentionLayerConfig, BufferConfig, CacheTierConfig,
-    GpuCacheTierConfig, HostCacheTierConfig, PageIndexMode, ReuseScope)
+    GpuCacheTierConfig, HostCacheTierConfig, PageIndexMode,
+    ReuseScope, SwaScratchReuseConfig)
 # isort: on
 from tensorrt_llm.runtime.kv_cache_manager_v2 import KVCacheIterationStatsDelta
 from tensorrt_llm.runtime.kv_cache_manager_v2 import \
@@ -1928,6 +1929,7 @@ def __init__(
         self.kv_factor = 1 if kv_cache_type == CacheTypeCpp.SELFKONLY else 2
         from ..speculative import get_num_extra_kv_tokens
         self.num_extra_kv_tokens = get_num_extra_kv_tokens(spec_config)
+        self.max_draft_len = spec_config.max_draft_len if spec_config is not None else 0
         self.max_total_draft_tokens = spec_config.max_total_draft_tokens if spec_config is not None else 0
 
         self.event_buffer_max_size = kv_cache_config.event_buffer_max_size
@@ -2447,13 +2449,21 @@ def _build_cache_config(
             if self.kv_cache_type != CacheTypeCpp.SELFKONLY:
                 buffer_type.append(Role.VALUE_BLOCK_SCALE)
 
+        scratch_reuse_config = None
+        if self.enable_swa_scratch_reuse:
+            # Context requests will allocate num_extra_kv_tokens tokens for spec decoding.
+            # Cache manager should not take them into account when calculating scratch range.
+            # Therefore set max_rewind_len to num_extra_kv_tokens.
+            scratch_reuse_config = SwaScratchReuseConfig(
+                max_rewind_len=self.num_extra_kv_tokens)
+
         return KVCacheManagerConfigPy(
             tokens_per_block=tokens_per_block,
             vocab_size=vocab_size,
             cache_tiers=cache_tiers,
             max_util_for_resume=kv_cache_config.max_util_for_resume,
             enable_stats=self.enable_stats,
-            enable_swa_scratch_reuse=self.enable_swa_scratch_reuse,
+            swa_scratch_reuse=scratch_reuse_config,
             layers=[
                 AttentionLayerConfig(
                     layer_id=layer_id,

tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py

@@ -39,6 +39,7 @@
     KVCacheDesc,
     KVCacheManagerConfig,
     SsmLayerConfig,
+    SwaScratchReuseConfig,
 )
 from ._core import (
     DEFAULT_BEAM_INDEX,
@@ -91,6 +92,7 @@
     "MemAddress",
     "NDEBUG",
     "KVCacheManagerConfig",
+    "SwaScratchReuseConfig",
     "AttentionLayerConfig",
     "SsmLayerConfig",
     "BufferConfig",

tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi

@@ -157,6 +157,10 @@ class BatchDesc:
     kv_caches: list[KVCacheDesc]
     system_prompt_length: int = 0
 
+@dataclass(slots=True)
+class SwaScratchReuseConfig:
+    max_rewind_len: int = 0
+
 @dataclass(slots=True)
 class KVCacheManagerConfig:
     tokens_per_block: int
@@ -168,9 +172,11 @@ class KVCacheManagerConfig:
     constraints: list[BatchDesc] = ...
     typical_step: BatchDesc | None = None
     ssm_reuse_interval: int = 512
-    enable_swa_scratch_reuse: bool = False
+    swa_scratch_reuse: SwaScratchReuseConfig | None = None
     enable_stats: bool = True
     helix_config: HelixConfig | None = None
+    @property
+    def enable_swa_scratch_reuse(self) -> bool: ...
 
 # From _event_manager.py
 EventBlockHash: TypeAlias = int | str

tensorrt_llm/runtime/kv_cache_manager_v2/_config.py

@@ -171,6 +171,23 @@ class HelixConfig:
     shared_comm_port: int
 
 
+@dataclass(slots=True)
+class SwaScratchReuseConfig:
+    """
+    Configuration for SWA scratch reuse.
+
+    Args:
+        max_rewind_len: Maximum number of tail tokens that can be rewound after
+            scratch-enabled allocation. Scratch reuse will not cover blocks that
+            may be needed to preserve those tokens.
+    """
+
+    max_rewind_len: int = 0
+
+    def __post_init__(self) -> None:
+        assert self.max_rewind_len >= 0, "max_rewind_len must be non-negative"
+
+
 @dataclass(slots=True)
 class KVCacheManagerConfig:
     """
@@ -212,13 +229,16 @@ class KVCacheManagerConfig:
     Must be a positive multiple of tokens_per_block. Only takes effect when SSM layers are present.
     """
 
-    enable_swa_scratch_reuse: bool = False
+    swa_scratch_reuse: SwaScratchReuseConfig | None = None
     """
-    When True, SWA layers reuse physical pages for out-of-window blocks during prefill.
+    When set, SWA layers reuse physical pages for out-of-window blocks during prefill.
     Scratch blocks share coalesced slot sub-pages across blocks for the currently executing
     layer, reducing peak memory. Trade-off: KV cache reuse is degraded because scratch blocks
     have no preserved data after the step.
 
+    If max_rewind_len is non-zero, the rewindable tail is excluded from scratch reuse so
+    draft/target shared KV cache can preserve tokens that may survive speculative rewind.
+
     Most useful for disaggregated prefill servers handling long prompts or long prompt chunks,
     where the number of out-of-window blocks dominates memory usage.
     """
@@ -231,6 +251,10 @@ class KVCacheManagerConfig:
     # unsupported yet
     helix_config: HelixConfig | None = None
 
+    @property
+    def enable_swa_scratch_reuse(self) -> bool:
+        return self.swa_scratch_reuse is not None
+
     def __post_init__(self) -> None:
         assert self.cache_tiers and self.cache_tiers[0].tier == CacheTier.GPU_MEM
         assert len(set(layer.layer_id for layer in self.layers)) == len(self.layers), (

tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py

@@ -650,11 +650,16 @@ def resize(self, capacity: int | None, history_length: int | None = None) -> boo
             raise ValueError("History length cannot be decreased")
         if capacity < history_length:
             raise ValueError("History length cannot be greater than capacity")
+        manager = self.manager
         # Scratch reuse: compute scratch ranges and slot delta
         enable_scratch = self.enable_swa_scratch_reuse
         if enable_scratch and capacity != self._capacity:
-            assert history_length == self._capacity, (
-                f"SWA scratch requires history_length ({history_length}) == "
+            max_rewind_len = self._swa_scratch_max_rewind_len()
+            min_history_length = max(0, self._capacity - max_rewind_len)
+            assert min_history_length <= history_length <= self._capacity, (
+                "SWA scratch requires "
+                f"old_capacity - max_rewind_len ({min_history_length}) <= "
+                f"history_length ({history_length}) <= "
                 f"old_capacity ({self._capacity})"
             )
         record_generation_alloc_stats = self._should_record_generation_alloc_stats(capacity)
@@ -665,12 +670,12 @@ def resize(self, capacity: int | None, history_length: int | None = None) -> boo
         ):
             self._refresh_generation_alloc_ready()
             return True
-        ssm_lc_id = self.manager._life_cycles.ssm_life_cycle_id
+        ssm_lc_id = manager._life_cycles.ssm_life_cycle_id
         beam_width = self.beam_width
         backup_holders = self._unlock_stale_blocks(history_length)
         old_num_blocks = BlockOrdinal(div_up(self._capacity, tokens_per_block))
         new_num_blocks = BlockOrdinal(div_up(capacity, tokens_per_block))
-        num_life_cycles = self.manager._life_cycles.size
+        num_life_cycles = manager._life_cycles.size
         if new_num_blocks < old_num_blocks:
             assert not self.has_scratch_slots, "Cannot shrink while scratch slots exist"
             self._subtract_pending_allocation_range(new_num_blocks, old_num_blocks)
@@ -694,7 +699,7 @@ def resize(self, capacity: int | None, history_length: int | None = None) -> boo
             num_new_slots = filled_list(0, num_life_cycles)
             stale_ranges = [
                 _KVCache._get_stale_range(tokens_per_block, history_length, lc)
-                for _, lc in self.manager._life_cycles.items()
+                for _, lc in manager._life_cycles.items()
             ]
             for lc in typed_range(num_life_cycles):
                 if lc == ssm_lc_id:
@@ -1567,23 +1572,43 @@ def _get_scratch_range(
         Range of blocks that should use scratch (shared) slots during SWA prefill.
 
         Scratch = stale_at_capacity ∩ input_blocks, where:
-        - stale_at_capacity: blocks out-of-window when all capacity tokens become history.
+        - stale_at_capacity: blocks out-of-window when all non-rewindable capacity tokens
+          become history.
         - input_blocks: [div_up(history_length, tpb), div_up(capacity, tpb)) — new blocks
           for the current chunk. Blocks before this range already contain real KV data
           from previous chunks and must not be overwritten.
+
+        The configured max_rewind_len excludes a speculative tail from scratch reuse.
         """
         if not self.enable_swa_scratch_reuse:
             return HalfOpenRange(BlockOrdinal(0), BlockOrdinal(0))
         history_length = value_or(history_length_override, self.history_length)
         capacity = value_or(capacity_override, self.capacity)
-        return compute_scratch_range(life_cycle, history_length, capacity, self.tokens_per_block)
+        max_rewind_len = self._swa_scratch_max_rewind_len()
+        return compute_scratch_range(
+            life_cycle,
+            history_length,
+            capacity,
+            self.tokens_per_block,
+            max_rewind_len,
+        )
 
     def _would_use_swa_scratch_blocks(self) -> bool:
+        max_rewind_len = self._swa_scratch_max_rewind_len()
         return any(
-            compute_scratch_range(lc, self.history_length, self.capacity, self.tokens_per_block)
+            compute_scratch_range(
+                lc,
+                self.history_length,
+                self.capacity,
+                self.tokens_per_block,
+                max_rewind_len,
+            )
             for lc in self.manager._life_cycles
         )
 
+    def _swa_scratch_max_rewind_len(self) -> int:
+        return unwrap_optional(self.manager.init_config.swa_scratch_reuse).max_rewind_len
+
     @staticmethod
     def _get_stale_range(
         tokens_per_block: int,

tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache_manager.py

@@ -257,7 +257,7 @@ def __init__(
             self._life_cycles,
             storage_config,
             config.tokens_per_block,
-            config.enable_swa_scratch_reuse,
+            config.swa_scratch_reuse,
             typical_batch=config.typical_step,
             constraints=config.constraints,
             event_manager=event_manager,
@@ -819,3 +819,7 @@ def is_enough(num_blocks: int) -> bool:
             else:
                 ub = mid
         return min(lb * tokens_per_block, token_num_upper_bound)
+
+    @property
+    def init_config(self) -> KVCacheManagerConfig:
+        return self._init_config

tensorrt_llm/runtime/kv_cache_manager_v2/_life_cycle_registry.py

@@ -137,19 +137,26 @@ def compute_scratch_range(
     history_length: int,
     capacity: int,
     tokens_per_block: int,
+    max_rewind_len: int,
 ) -> HalfOpenRange[BlockOrdinal]:
     """
     Range of blocks that should use scratch (shared) slots during SWA prefill.
 
     Scratch = stale_at_capacity ∩ input_blocks, where:
-    - stale_at_capacity: blocks out-of-window when all capacity tokens become history.
+    - stale_at_capacity: blocks out-of-window when all non-rewindable capacity tokens
+      become history.
     - input_blocks: [div_up(history_length, tpb), div_up(capacity, tpb)) — new blocks
       for the current chunk. Blocks before this range already contain real KV data
       from previous chunks and must not be overwritten.
+
+    max_rewind_len protects the speculative tail from scratch reuse. Those tokens may
+    survive after rejected draft tokens are rewound, so their KV data must remain in
+    normal per-block pages.
     """
     if not isinstance(life_cycle, AttnLifeCycle) or life_cycle.window_size is None:
         return HalfOpenRange(BlockOrdinal(0), BlockOrdinal(0))
-    cap_stale = life_cycle.get_stale_range(capacity, tokens_per_block)
+    non_rewindable_capacity = max(0, capacity - max_rewind_len)
+    cap_stale = life_cycle.get_stale_range(non_rewindable_capacity, tokens_per_block)
     input_range = HalfOpenRange(
         BlockOrdinal(div_up(history_length, tokens_per_block)),
         BlockOrdinal(div_up(capacity, tokens_per_block)),