Skip to content

Commit effe6fb

Browse files
jiaganclowsfer
authored and
lfr-0531
committed
[TRTLLM-12229][fix] Fix MTP by scratch reuse rewind (NVIDIA#14403)
Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com> Signed-off-by: Jiagan Cheng <jiaganc@nvidia.com> Co-authored-by: Yao Yao <lowsfer@users.noreply.github.com>
1 parent 293a7a4 commit effe6fb

8 files changed

Lines changed: 156 additions & 48 deletions

File tree

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

Lines changed: 12 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,7 @@
4545
LayerId,
4646
PageIndexMode,
4747
ScratchDesc,
48+
SwaScratchReuseConfig,
4849
)
4950
from tensorrt_llm.runtime.kv_cache_manager_v2 import KVCacheManagerConfig as KVCacheManagerConfigPy
5051
from tensorrt_llm.runtime.kv_cache_manager_v2._common import BAD_PAGE_INDEX
@@ -60,8 +61,8 @@
6061
is_overlap_compressor,
6162
)
6263

63-
# Keep DSV4 scratch reuse opt-in so per-layer block tables can be tested
64-
# without scratch-page remapping in the default path.
64+
# Keep the env override so per-layer block tables can still be tested
65+
# without scratch-page remapping when needed.
6566
DSV4_ENABLE_SWA_SCRATCH_REUSE_ENV = "TRTLLM_DSV4_ENABLE_SWA_SCRATCH_REUSE"
6667

6768

@@ -119,7 +120,7 @@ def _get_index_mode(attn_type: DeepseekV4AttentionType) -> PageIndexMode:
119120

120121

121122
def _enable_swa_scratch_reuse_from_env() -> bool:
122-
value = os.environ.get(DSV4_ENABLE_SWA_SCRATCH_REUSE_ENV, "0")
123+
value = os.environ.get(DSV4_ENABLE_SWA_SCRATCH_REUSE_ENV, "1")
123124
return value.strip() == "1"
124125

125126

@@ -798,13 +799,20 @@ def _add_layer(
798799
if max_num_tokens is not None:
799800
constraints.append(BatchDesc([KVCacheDesc(capacity=max_num_tokens, history_length=0)]))
800801

802+
scratch_reuse_config = None
803+
if self.enable_swa_scratch_reuse:
804+
# Context requests will allocate num_extra_kv_tokens tokens for spec decoding.
805+
# Cache manager should not take them into account when calculating scratch range.
806+
# Therefore set max_rewind_len to num_extra_kv_tokens.
807+
scratch_reuse_config = SwaScratchReuseConfig(max_rewind_len=self.num_extra_kv_tokens)
808+
801809
return KVCacheManagerConfigPy(
802810
tokens_per_block=tokens_per_block,
803811
vocab_size=vocab_size,
804812
cache_tiers=cache_tiers,
805813
max_util_for_resume=kv_cache_config.max_util_for_resume,
806814
enable_stats=self.enable_stats,
807-
enable_swa_scratch_reuse=self.enable_swa_scratch_reuse,
815+
swa_scratch_reuse=scratch_reuse_config,
808816
layers=layers,
809817
typical_step=typical_step,
810818
constraints=constraints,

‎tensorrt_llm/_torch/pyexecutor/resource_manager.py‎

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -49,7 +49,8 @@
4949
# isort: off
5050
from tensorrt_llm.runtime.kv_cache_manager_v2 import (
5151
DEFAULT_BEAM_INDEX, AttentionLayerConfig, BufferConfig, CacheTierConfig,
52-
GpuCacheTierConfig, HostCacheTierConfig, PageIndexMode, ReuseScope)
52+
GpuCacheTierConfig, HostCacheTierConfig, PageIndexMode, ReuseScope,
53+
SwaScratchReuseConfig)
5354
# isort: on
5455
from tensorrt_llm.runtime.kv_cache_manager_v2 import KVCacheIterationStatsDelta
5556
from tensorrt_llm.runtime.kv_cache_manager_v2 import \
@@ -1941,6 +1942,7 @@ def __init__(
19411942
self.kv_factor = 1 if kv_cache_type == CacheTypeCpp.SELFKONLY else 2
19421943
from ..speculative import get_num_extra_kv_tokens
19431944
self.num_extra_kv_tokens = get_num_extra_kv_tokens(spec_config)
1945+
self.max_draft_len = spec_config.max_draft_len if spec_config is not None else 0
19441946
self.max_total_draft_tokens = spec_config.max_total_draft_tokens if spec_config is not None else 0
19451947

19461948
self.event_buffer_max_size = kv_cache_config.event_buffer_max_size
@@ -2460,13 +2462,21 @@ def _build_cache_config(
24602462
if self.kv_cache_type != CacheTypeCpp.SELFKONLY:
24612463
buffer_type.append(Role.VALUE_BLOCK_SCALE)
24622464

2465+
scratch_reuse_config = None
2466+
if self.enable_swa_scratch_reuse:
2467+
# Context requests will allocate num_extra_kv_tokens tokens for spec decoding.
2468+
# Cache manager should not take them into account when calculating scratch range.
2469+
# Therefore set max_rewind_len to num_extra_kv_tokens.
2470+
scratch_reuse_config = SwaScratchReuseConfig(
2471+
max_rewind_len=self.num_extra_kv_tokens)
2472+
24632473
return KVCacheManagerConfigPy(
24642474
tokens_per_block=tokens_per_block,
24652475
vocab_size=vocab_size,
24662476
cache_tiers=cache_tiers,
24672477
max_util_for_resume=kv_cache_config.max_util_for_resume,
24682478
enable_stats=self.enable_stats,
2469-
enable_swa_scratch_reuse=self.enable_swa_scratch_reuse,
2479+
swa_scratch_reuse=scratch_reuse_config,
24702480
layers=[
24712481
AttentionLayerConfig(
24722482
layer_id=layer_id,

‎tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,7 @@
4343
KVCacheDesc,
4444
KVCacheManagerConfig,
4545
SsmLayerConfig,
46+
SwaScratchReuseConfig,
4647
)
4748
from ._core import (
4849
DEFAULT_BEAM_INDEX,
@@ -95,6 +96,7 @@
9596
"MemAddress",
9697
"NDEBUG",
9798
"KVCacheManagerConfig",
99+
"SwaScratchReuseConfig",
98100
"AttentionLayerConfig",
99101
"SsmLayerConfig",
100102
"BufferConfig",

‎tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -157,6 +157,10 @@ class BatchDesc:
157157
kv_caches: list[KVCacheDesc]
158158
system_prompt_length: int = 0
159159

160+
@dataclass(slots=True)
161+
class SwaScratchReuseConfig:
162+
max_rewind_len: int = 0
163+
160164
@dataclass(slots=True)
161165
class KVCacheManagerConfig:
162166
tokens_per_block: int
@@ -168,9 +172,11 @@ class KVCacheManagerConfig:
168172
constraints: list[BatchDesc] = ...
169173
typical_step: BatchDesc | None = None
170174
ssm_reuse_interval: int = 512
171-
enable_swa_scratch_reuse: bool = False
175+
swa_scratch_reuse: SwaScratchReuseConfig | None = None
172176
enable_stats: bool = True
173177
helix_config: HelixConfig | None = None
178+
@property
179+
def enable_swa_scratch_reuse(self) -> bool: ...
174180

175181
# From _event_manager.py
176182
EventBlockHash: TypeAlias = int | str

‎tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py‎

Lines changed: 33 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -650,11 +650,16 @@ def resize(self, capacity: int | None, history_length: int | None = None) -> boo
650650
raise ValueError("History length cannot be decreased")
651651
if capacity < history_length:
652652
raise ValueError("History length cannot be greater than capacity")
653+
manager = self.manager
653654
# Scratch reuse: compute scratch ranges and slot delta
654655
enable_scratch = self.enable_swa_scratch_reuse
655656
if enable_scratch and capacity != self._capacity:
656-
assert history_length == self._capacity, (
657-
f"SWA scratch requires history_length ({history_length}) == "
657+
max_rewind_len = self._swa_scratch_max_rewind_len()
658+
min_history_length = max(0, self._capacity - max_rewind_len)
659+
assert min_history_length <= history_length <= self._capacity, (
660+
"SWA scratch requires "
661+
f"old_capacity - max_rewind_len ({min_history_length}) <= "
662+
f"history_length ({history_length}) <= "
658663
f"old_capacity ({self._capacity})"
659664
)
660665
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
665670
):
666671
self._refresh_generation_alloc_ready()
667672
return True
668-
ssm_lc_id = self.manager._life_cycles.ssm_life_cycle_id
673+
ssm_lc_id = manager._life_cycles.ssm_life_cycle_id
669674
beam_width = self.beam_width
670675
backup_holders = self._unlock_stale_blocks(history_length)
671676
old_num_blocks = BlockOrdinal(div_up(self._capacity, tokens_per_block))
672677
new_num_blocks = BlockOrdinal(div_up(capacity, tokens_per_block))
673-
num_life_cycles = self.manager._life_cycles.size
678+
num_life_cycles = manager._life_cycles.size
674679
if new_num_blocks < old_num_blocks:
675680
assert not self.has_scratch_slots, "Cannot shrink while scratch slots exist"
676681
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
694699
num_new_slots = filled_list(0, num_life_cycles)
695700
stale_ranges = [
696701
_KVCache._get_stale_range(tokens_per_block, history_length, lc)
697-
for _, lc in self.manager._life_cycles.items()
702+
for _, lc in manager._life_cycles.items()
698703
]
699704
for lc in typed_range(num_life_cycles):
700705
if lc == ssm_lc_id:
@@ -1567,23 +1572,43 @@ def _get_scratch_range(
15671572
Range of blocks that should use scratch (shared) slots during SWA prefill.
15681573
15691574
Scratch = stale_at_capacity ∩ input_blocks, where:
1570-
- stale_at_capacity: blocks out-of-window when all capacity tokens become history.
1575+
- stale_at_capacity: blocks out-of-window when all non-rewindable capacity tokens
1576+
become history.
15711577
- input_blocks: [div_up(history_length, tpb), div_up(capacity, tpb)) — new blocks
15721578
for the current chunk. Blocks before this range already contain real KV data
15731579
from previous chunks and must not be overwritten.
1580+
1581+
The configured max_rewind_len excludes a speculative tail from scratch reuse.
15741582
"""
15751583
if not self.enable_swa_scratch_reuse:
15761584
return HalfOpenRange(BlockOrdinal(0), BlockOrdinal(0))
15771585
history_length = value_or(history_length_override, self.history_length)
15781586
capacity = value_or(capacity_override, self.capacity)
1579-
return compute_scratch_range(life_cycle, history_length, capacity, self.tokens_per_block)
1587+
max_rewind_len = self._swa_scratch_max_rewind_len()
1588+
return compute_scratch_range(
1589+
life_cycle,
1590+
history_length,
1591+
capacity,
1592+
self.tokens_per_block,
1593+
max_rewind_len,
1594+
)
15801595

15811596
def _would_use_swa_scratch_blocks(self) -> bool:
1597+
max_rewind_len = self._swa_scratch_max_rewind_len()
15821598
return any(
1583-
compute_scratch_range(lc, self.history_length, self.capacity, self.tokens_per_block)
1599+
compute_scratch_range(
1600+
lc,
1601+
self.history_length,
1602+
self.capacity,
1603+
self.tokens_per_block,
1604+
max_rewind_len,
1605+
)
15841606
for lc in self.manager._life_cycles
15851607
)
15861608

1609+
def _swa_scratch_max_rewind_len(self) -> int:
1610+
return unwrap_optional(self.manager.init_config.swa_scratch_reuse).max_rewind_len
1611+
15871612
@staticmethod
15881613
def _get_stale_range(
15891614
tokens_per_block: int,

‎tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache_manager.py‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -257,7 +257,7 @@ def __init__(
257257
self._life_cycles,
258258
storage_config,
259259
config.tokens_per_block,
260-
config.enable_swa_scratch_reuse,
260+
config.swa_scratch_reuse,
261261
typical_batch=config.typical_step,
262262
constraints=config.constraints,
263263
event_manager=event_manager,
@@ -819,3 +819,7 @@ def is_enough(num_blocks: int) -> bool:
819819
else:
820820
ub = mid
821821
return min(lb * tokens_per_block, token_num_upper_bound)
822+
823+
@property
824+
def init_config(self) -> KVCacheManagerConfig:
825+
return self._init_config

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

Lines changed: 40 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -77,7 +77,7 @@ def scratch_reuse_enabled(request, monkeypatch) -> bool:
7777
if request.param:
7878
monkeypatch.setenv(DSV4_ENABLE_SWA_SCRATCH_REUSE_ENV, "1")
7979
else:
80-
monkeypatch.delenv(DSV4_ENABLE_SWA_SCRATCH_REUSE_ENV, raising=False)
80+
monkeypatch.setenv(DSV4_ENABLE_SWA_SCRATCH_REUSE_ENV, "0")
8181
return request.param
8282

8383

@@ -171,6 +171,7 @@ def _create_deepseek_v4_cache_manager(
171171
is_draft: bool = False,
172172
tp_size: int = 1,
173173
enable_attention_dp: bool = False,
174+
spec_config: object | None = None,
174175
) -> Tuple[DeepseekV4CacheManager, DeepSeekV4SparseAttentionConfig]:
175176
"""Helper to create a DeepseekV4CacheManager for testing."""
176177

@@ -217,6 +218,7 @@ def _create_deepseek_v4_cache_manager(
217218
max_num_tokens=max_batch_size * (max_input_len + 1),
218219
sparse_attn_config=sparse_attn_config,
219220
is_draft=is_draft,
221+
spec_config=spec_config,
220222
)
221223

222224
return cache_manager, sparse_attn_config
@@ -1423,7 +1425,7 @@ def test_copy_batch_indexer_compress_block_tables_matches_python_converter(
14231425
cache_manager.free_resources(req)
14241426
cache_manager.shutdown()
14251427

1426-
def test_swa_scratch_reuse_disabled_by_default_for_main_manager(self, monkeypatch):
1428+
def test_swa_scratch_reuse_enabled_by_default_for_main_manager(self, monkeypatch):
14271429
monkeypatch.delenv(DSV4_ENABLE_SWA_SCRATCH_REUSE_ENV, raising=False)
14281430
cache_manager, _ = self._create_deepseek_v4_cache_manager(
14291431
tokens_per_block=self.tokens_per_block,
@@ -1434,28 +1436,59 @@ def test_swa_scratch_reuse_disabled_by_default_for_main_manager(self, monkeypatc
14341436
compressor_dtype=DataType.FLOAT,
14351437
)
14361438

1439+
try:
1440+
assert cache_manager.enable_swa_scratch_reuse
1441+
assert cache_manager.kv_cache_manager_py_config.enable_swa_scratch_reuse
1442+
assert cache_manager.kv_cache_manager_py_config.swa_scratch_reuse is not None
1443+
assert cache_manager.kv_cache_manager_py_config.swa_scratch_reuse.max_rewind_len == 0
1444+
assert cache_manager.num_attention_op_pools == cache_manager.num_local_layers
1445+
finally:
1446+
cache_manager.shutdown()
1447+
1448+
def test_swa_scratch_reuse_disabled_by_env_for_main_manager(self, monkeypatch):
1449+
monkeypatch.setenv(DSV4_ENABLE_SWA_SCRATCH_REUSE_ENV, "0")
1450+
cache_manager, _ = self._create_deepseek_v4_cache_manager(
1451+
tokens_per_block=self.tokens_per_block,
1452+
max_batch_size=1,
1453+
max_seq_len=1024,
1454+
compress_ratios=[1],
1455+
dtype=DataType.BF16,
1456+
compressor_dtype=DataType.FLOAT,
1457+
)
1458+
14371459
try:
14381460
assert not cache_manager.enable_swa_scratch_reuse
14391461
assert not cache_manager.kv_cache_manager_py_config.enable_swa_scratch_reuse
1462+
assert cache_manager.kv_cache_manager_py_config.swa_scratch_reuse is None
14401463
assert cache_manager.num_attention_op_pools == cache_manager.num_local_layers
14411464
finally:
14421465
cache_manager.shutdown()
14431466

1444-
def test_swa_scratch_reuse_enabled_by_env_for_main_manager(self, monkeypatch):
1467+
def test_swa_scratch_reuse_uses_extra_kv_tokens_for_rewind(self, monkeypatch):
14451468
monkeypatch.setenv(DSV4_ENABLE_SWA_SCRATCH_REUSE_ENV, "1")
1469+
spec_config = SimpleNamespace(
1470+
max_draft_len=7,
1471+
max_total_draft_tokens=7,
1472+
spec_dec_mode=SimpleNamespace(
1473+
is_eagle3_one_model=lambda: False,
1474+
is_mtp_one_model=lambda: False,
1475+
use_one_engine=lambda: True,
1476+
),
1477+
)
14461478
cache_manager, _ = self._create_deepseek_v4_cache_manager(
14471479
tokens_per_block=self.tokens_per_block,
14481480
max_batch_size=1,
14491481
max_seq_len=1024,
14501482
compress_ratios=[1],
14511483
dtype=DataType.BF16,
14521484
compressor_dtype=DataType.FLOAT,
1485+
spec_config=spec_config,
14531486
)
14541487

14551488
try:
1456-
assert cache_manager.enable_swa_scratch_reuse
1457-
assert cache_manager.kv_cache_manager_py_config.enable_swa_scratch_reuse
1458-
assert cache_manager.num_attention_op_pools == cache_manager.num_local_layers
1489+
scratch_reuse = cache_manager.kv_cache_manager_py_config.swa_scratch_reuse
1490+
assert scratch_reuse is not None
1491+
assert scratch_reuse.max_rewind_len == spec_config.max_draft_len - 1
14591492
finally:
14601493
cache_manager.shutdown()
14611494

@@ -1474,6 +1507,7 @@ def test_draft_cache_manager_disables_swa_scratch_reuse(self, monkeypatch):
14741507
try:
14751508
assert not cache_manager.enable_swa_scratch_reuse
14761509
assert not cache_manager.kv_cache_manager_py_config.enable_swa_scratch_reuse
1510+
assert cache_manager.kv_cache_manager_py_config.swa_scratch_reuse is None
14771511
assert cache_manager.num_attention_op_pools == cache_manager.num_local_layers
14781512
finally:
14791513
cache_manager.shutdown()

0 commit comments

Comments
 (0)