[TRTLLM-12229][fix] Fix MTP by scratch reuse rewind (NVIDIA#14403)

Source-Commit: 7a4b113
Signed-off-by: Jiagan Cheng <jiaganc@nvidia.com>

Author: jiaganc

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

@@ -800,6 +800,9 @@ def _add_layer(
 
         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(

tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py

@@ -563,6 +563,7 @@ def __init__(
         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
         )

tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py

@@ -177,6 +177,7 @@ def _create_deepseek_v4_cache_manager(
         tp_size: int = 1,
         enable_attention_dp: bool = False,
         enable_swa_scratch_reuse: bool = False,
+        spec_config: object | None = None,
     ) -> Tuple[DeepseekV4CacheManager, DeepSeekV4SparseAttentionConfig]:
         """Helper to create a DeepseekV4CacheManager for testing."""
 
@@ -224,6 +225,7 @@ def _create_deepseek_v4_cache_manager(
             max_num_tokens=max_batch_size * (max_input_len + 1),
             sparse_attn_config=sparse_attn_config,
             is_draft=is_draft,
+            spec_config=spec_config,
         )
 
         return cache_manager, sparse_attn_config
@@ -1450,6 +1452,7 @@ def test_swa_scratch_reuse_disabled_by_default_for_main_manager(self):
         try:
             assert not cache_manager.enable_swa_scratch_reuse
             assert not cache_manager.kv_cache_manager_py_config.enable_swa_scratch_reuse
+            assert cache_manager.kv_cache_manager_py_config.swa_scratch_reuse is None
             assert cache_manager.num_attention_op_pools == cache_manager.num_local_layers
         finally:
             cache_manager.shutdown()
@@ -1468,10 +1471,40 @@ def test_swa_scratch_reuse_enabled_by_config_for_main_manager(self):
         try:
             assert cache_manager.enable_swa_scratch_reuse
             assert cache_manager.kv_cache_manager_py_config.enable_swa_scratch_reuse
+            assert cache_manager.kv_cache_manager_py_config.swa_scratch_reuse is not None
+            assert cache_manager.kv_cache_manager_py_config.swa_scratch_reuse.max_rewind_len == 0
             assert cache_manager.num_attention_op_pools == cache_manager.num_local_layers
         finally:
             cache_manager.shutdown()
 
+    def test_swa_scratch_reuse_uses_extra_kv_tokens_for_rewind(self):
+        spec_config = SimpleNamespace(
+            max_draft_len=7,
+            max_total_draft_tokens=7,
+            spec_dec_mode=SimpleNamespace(
+                is_eagle3_one_model=lambda: False,
+                is_mtp_one_model=lambda: False,
+                use_one_engine=lambda: True,
+            ),
+        )
+        cache_manager, _ = self._create_deepseek_v4_cache_manager(
+            tokens_per_block=self.tokens_per_block,
+            max_batch_size=1,
+            max_seq_len=1024,
+            compress_ratios=[1],
+            dtype=DataType.BF16,
+            compressor_dtype=DataType.FLOAT,
+            spec_config=spec_config,
+            enable_swa_scratch_reuse=True,
+        )
+
+        try:
+            scratch_reuse = cache_manager.kv_cache_manager_py_config.swa_scratch_reuse
+            assert scratch_reuse is not None
+            assert scratch_reuse.max_rewind_len == spec_config.max_draft_len - 1
+        finally:
+            cache_manager.shutdown()
+
     def test_draft_cache_manager_disables_swa_scratch_reuse(self):
         cache_manager, _ = self._create_deepseek_v4_cache_manager(
             tokens_per_block=self.tokens_per_block,
@@ -1487,6 +1520,7 @@ def test_draft_cache_manager_disables_swa_scratch_reuse(self):
         try:
             assert not cache_manager.enable_swa_scratch_reuse
             assert not cache_manager.kv_cache_manager_py_config.enable_swa_scratch_reuse
+            assert cache_manager.kv_cache_manager_py_config.swa_scratch_reuse is None
             assert cache_manager.num_attention_op_pools == cache_manager.num_local_layers
         finally:
             cache_manager.shutdown()