From c7816b4d0e8b9959613f9dc841d94aac0d77ac0a Mon Sep 17 00:00:00 2001 From: Yao Yao Date: Wed, 20 May 2026 11:08:14 +0000 Subject: [PATCH 1/5] [None][fix] support SWA scratch reuse rewind Signed-off-by: Yao Yao --- .../sparse/deepseek_v4/cache_manager.py | 7 +- .../_torch/pyexecutor/resource_manager.py | 8 +- .../runtime/kv_cache_manager_v2/__init__.py | 2 + .../runtime/kv_cache_manager_v2/__init__.pyi | 8 +- .../runtime/kv_cache_manager_v2/_config.py | 28 ++++++- .../kv_cache_manager_v2/_core/_kv_cache.py | 41 +++++++++-- .../_core/_kv_cache_manager.py | 6 +- .../_life_cycle_registry.py | 11 ++- .../kv_cache_manager_v2/_storage_manager.py | 41 ++++++++--- .../test_deepseek_v4_cache_manager.py | 30 ++++++++ .../test_kv_cache_manager_v2.py | 73 ++++++++++++------- 11 files changed, 201 insertions(+), 54 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py index f16da05fa965..375d9253b9ea 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py +++ b/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 @@ -804,7 +805,11 @@ def _add_layer( 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=( + SwaScratchReuseConfig(max_rewind_len=self._max_draft_len) + if self.enable_swa_scratch_reuse + else None + ), layers=layers, typical_step=typical_step, constraints=constraints, diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index f29073a10561..d5315db8d8eb 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/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 @@ -2453,7 +2455,9 @@ def _build_cache_config( 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=(SwaScratchReuseConfig( + max_rewind_len=self.max_draft_len) + if self.enable_swa_scratch_reuse else None), layers=[ AttentionLayerConfig( layer_id=layer_id, diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py index 26c223cc12fc..f13558ab4eec 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py +++ b/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", diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi index beec4baae37b..4c9f9447d201 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi +++ b/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 diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_config.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_config.py index b03b0e611444..9341141a139a 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_config.py +++ b/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), ( 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 869da913b66a..beddc339d339 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 @@ -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, diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache_manager.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache_manager.py index 99234ba8b5d6..cc914a768e42 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache_manager.py +++ b/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 diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_life_cycle_registry.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_life_cycle_registry.py index 4200e8b11822..551f5cec0217 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_life_cycle_registry.py +++ b/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)), diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_storage_manager.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_storage_manager.py index 70a78c64b207..b0cd3abfaab4 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_storage_manager.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_storage_manager.py @@ -34,7 +34,14 @@ MemAddress, PageStatus, ) -from ._config import BatchDesc, CacheTierConfig, DataRole, DiskCacheTierConfig, KVCacheDesc +from ._config import ( + BatchDesc, + CacheTierConfig, + DataRole, + DiskCacheTierConfig, + KVCacheDesc, + SwaScratchReuseConfig, +) from ._copy_engine import CopyTask, batched_copy from ._event_manager import KVCacheEventDiff from ._eviction_controller import EvictablePage, PerLevelEvictionController @@ -201,7 +208,7 @@ def __init__( life_cycles: LifeCycleRegistry, config: StorageConfig, tokens_per_block: int, - enable_swa_scratch_reuse: bool, + swa_scratch_reuse: SwaScratchReuseConfig | None, typical_batch: BatchDesc | None = None, constraints: list[BatchDesc] | None = None, event_manager: "KVCacheEventManager | None" = None, @@ -231,13 +238,13 @@ def __init__( gpu_granularity = CacheLevelManager.cache_tier_granularity(CacheTier.GPU_MEM, gpu_quota) self._min_slots = self._compute_min_slots_from_constraints( - constraints or [], tokens_per_block, enable_swa_scratch_reuse + constraints or [], tokens_per_block, swa_scratch_reuse ) # Compute init_ratio from typical_batch, constraints, or fallback. if typical_batch is not None: init_ratio = self.ratio_from_batch( - typical_batch, tokens_per_block, enable_swa_scratch_reuse, gpu_granularity + typical_batch, tokens_per_block, swa_scratch_reuse, gpu_granularity ) elif constraints: # Use the constraint slot counts as the ratio basis. @@ -248,7 +255,7 @@ def __init__( init_ratio = self.ratio_from_batch( BatchDesc([KVCacheDesc(capacity=2049, history_length=2048)]), tokens_per_block, - enable_swa_scratch_reuse, + swa_scratch_reuse, gpu_granularity, ) @@ -847,18 +854,21 @@ def ratio_from_batch( self, batch: BatchDesc, tokens_per_block: int, - enable_swa_scratch_reuse: bool, + swa_scratch_reuse: SwaScratchReuseConfig | None, granularity: int, ) -> TypedIndexList[PoolGroupIndex, float]: """Compute the ratio of bytes needed per pool group for a batch described by a BatchDesc.""" - num_slots = self._compute_slots_for_batch(batch, tokens_per_block, enable_swa_scratch_reuse) + num_slots = self._compute_slots_for_batch(batch, tokens_per_block, swa_scratch_reuse) num_bytes = self._slots_to_bytes(num_slots, granularity) total = sum(num_bytes) assert total > 0 return typed_map(num_bytes, lambda x: x / total) def _compute_min_slots_from_constraints( - self, constraints: list[BatchDesc], tokens_per_block: int, enable_swa_scratch_reuse: bool + self, + constraints: list[BatchDesc], + tokens_per_block: int, + swa_scratch_reuse: SwaScratchReuseConfig | None, ) -> TypedIndexList[PoolGroupIndex, int]: """Compute the minimum slots per pool group across all constraints (element-wise max). @@ -869,13 +879,16 @@ def _compute_min_slots_from_constraints( for pg_idx in self._life_cycle_grouping: max_slots[pg_idx] += 1 for batch in constraints: - slots = self._compute_slots_for_batch(batch, tokens_per_block, enable_swa_scratch_reuse) + slots = self._compute_slots_for_batch(batch, tokens_per_block, swa_scratch_reuse) for pg_idx in typed_range(self.num_pool_groups): max_slots[pg_idx] = max(max_slots[pg_idx], slots[pg_idx]) return max_slots def _compute_slots_for_batch( - self, batch: BatchDesc, tokens_per_block: int, enable_swa_scratch_reuse: bool + self, + batch: BatchDesc, + tokens_per_block: int, + swa_scratch_reuse: SwaScratchReuseConfig | None, ) -> TypedIndexList[PoolGroupIndex, int]: """Compute the minimum number of slots per pool group to support a BatchDesc.""" num_slots = filled_list(0, self.num_pool_groups) @@ -905,9 +918,13 @@ def _compute_slots_for_batch( # Non-stale sys blocks for this request. non_stale_sys = sys_blocks - len(intersect(stale, sys_range)) unique_non_stale = max(0, non_stale - non_stale_sys) - if enable_swa_scratch_reuse: + if swa_scratch_reuse is not None: scratch = compute_scratch_range( - lc, kv.history_length, kv.capacity, tokens_per_block + lc, + kv.history_length, + kv.capacity, + tokens_per_block, + swa_scratch_reuse.max_rewind_len, ) # Scratch blocks are always input blocks, so they never # overlap with shared sys blocks (which are history). diff --git a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py index 18f83227c09b..473bdccb91f6 100644 --- a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py +++ b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py @@ -171,6 +171,7 @@ def _create_deepseek_v4_cache_manager( is_draft: bool = False, tp_size: int = 1, enable_attention_dp: bool = False, + spec_config: object | None = None, ) -> Tuple[DeepseekV4CacheManager, DeepSeekV4SparseAttentionConfig]: """Helper to create a DeepseekV4CacheManager for testing.""" @@ -217,6 +218,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 @@ -1437,6 +1439,7 @@ def test_swa_scratch_reuse_disabled_by_default_for_main_manager(self, monkeypatc 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() @@ -1455,10 +1458,36 @@ def test_swa_scratch_reuse_enabled_by_env_for_main_manager(self, monkeypatch): 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_spec_draft_len_for_rewind(self, monkeypatch): + monkeypatch.setenv(DSV4_ENABLE_SWA_SCRATCH_REUSE_ENV, "1") + spec_config = SimpleNamespace( + max_draft_len=7, + max_total_draft_tokens=7, + spec_dec_mode=SimpleNamespace(use_one_engine=lambda: False), + ) + 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, + ) + + 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 + finally: + cache_manager.shutdown() + def test_draft_cache_manager_disables_swa_scratch_reuse(self, monkeypatch): monkeypatch.setenv(DSV4_ENABLE_SWA_SCRATCH_REUSE_ENV, "1") cache_manager, _ = self._create_deepseek_v4_cache_manager( @@ -1474,6 +1503,7 @@ def test_draft_cache_manager_disables_swa_scratch_reuse(self, monkeypatch): 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() 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 87c344d46460..4a88cd2f49cf 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 @@ -47,6 +47,7 @@ LayerId, ReuseScope, SsmLayerConfig, + SwaScratchReuseConfig, TokenId, TokenIdExt, _KVCache, @@ -98,6 +99,7 @@ LayerId, ReuseScope, SsmLayerConfig, + SwaScratchReuseConfig, TokenId, TokenIdExt, _KVCache, @@ -1629,7 +1631,7 @@ def _make_config( is non-trivial and constraint clamping is exercised. With num_windowed_layers / num_full_layers > 1 and - enable_swa_scratch_reuse=True, multiple layers per lifecycle give + scratch reuse enabled, multiple layers per lifecycle give frac_max < 1, making scratch savings visible in capacity planning. """ cache_tiers: list = [GpuCacheTierConfig(quota=gpu_quota)] @@ -1662,7 +1664,7 @@ def _make_config( layers=layers, typical_step=typical_step, constraints=constraints or [], - enable_swa_scratch_reuse=enable_swa_scratch_reuse, + swa_scratch_reuse=(SwaScratchReuseConfig() if enable_swa_scratch_reuse else None), ) def test_default_init_ratio(self): @@ -2050,6 +2052,7 @@ def _prepare_scratch( tokens_per_block: int = 32, gpu_quota: int = 64 << 20, sink_tokens: int = 0, + max_rewind_len: int = 0, ): """Prepare a manager with scratch reuse enabled.""" kv_buf_size = 8192 @@ -2069,7 +2072,7 @@ def _prepare_scratch( ) for i in range(num_layers) ], - enable_swa_scratch_reuse=True, + swa_scratch_reuse=SwaScratchReuseConfig(max_rewind_len=max_rewind_len), ) self.engine = FakeEngine(self.cfg) self.manager = KVCacheManager(self.cfg) @@ -2244,13 +2247,20 @@ def test_scratch_slot_count(self): kv3.close() self.manager.clear_reusable_blocks() - def test_scratch_shared_slot_ids(self): + @parameterized.expand([(0, 7), (64, 5)]) + def test_scratch_shared_slot_ids(self, rewind_len: int, expected_scratch_blocks: int): """Verify that scratch blocks share coalesced slot IDs via ScratchDesc.""" # 8 layers, window=32, tokens_per_block=32, prompt=256 # num_sub_pages = 8 (all layers in one group) - # blocks 0-6 are scratch (7 blocks), block 7 is in-window (normal) - # 7 scratch blocks / 8 sub_pages = 1 scratch slot - self._prepare_scratch(num_layers=8, window_size=32, tokens_per_block=32, gpu_quota=16 << 20) + # rewind_len=0: blocks 0-6 are scratch, block 7 is in-window. + # rewind_len=64: blocks 5-7 are protected from scratch by the rewind tail. + self._prepare_scratch( + num_layers=8, + window_size=32, + tokens_per_block=32, + gpu_quota=16 << 20, + max_rewind_len=rewind_len, + ) prompt = [self.next_token() for _ in range(256)] kv = self.manager.create_kv_cache(None, prompt) @@ -2269,10 +2279,10 @@ def test_scratch_shared_slot_ids(self): self.assertIsNotNone(scratch_desc) num_scratch_blocks = scratch_desc.range.end - scratch_desc.range.beg - self.assertEqual(num_scratch_blocks, 7) # blocks 0-6 + self.assertEqual(num_scratch_blocks, expected_scratch_blocks) - # 7 blocks / 8 sub_pages = ceil = 1 scratch slot - self.assertEqual(len(scratch_desc.slot_ids), 1) + expected_scratch_slots = div_up(expected_scratch_blocks, 8) + self.assertEqual(len(scratch_desc.slot_ids), expected_scratch_slots) # Verify scratch blocks have BAD_PAGE_INDEX in base_page_indices indices = kv.get_base_page_indices(lg_id) @@ -2307,22 +2317,28 @@ def test_scratch_shared_slot_ids(self): kv.close() self.manager.clear_reusable_blocks() - def test_scratch_chunk_size_variation(self): + @parameterized.expand([(0, (0, 6), (8, 9)), (32, (0, 5), None)]) + def test_scratch_chunk_size_variation( + self, + rewind_len: int, + chunk1_scratch_range: tuple[int, int], + chunk2_scratch_range: tuple[int, int] | None, + ): """Verify scratch block allocation with changing chunk sizes and multiple window sizes. This ensures both positive and negative net_alloc_counts code paths are tested - simultaneously across different layers. + simultaneously across different layers. The rewind_len parameter verifies + that the protected rewind tail is kept out of scratch ranges. Layer 0: window_size = 64 (2 blocks) Layer 1: window_size = 256 (8 blocks) Chunk 1: resize(256) -> 8 blocks. - - Layer 0 (stale 0-6): needs 6 scratch blocks (net_alloc_counts = 6 > 0) + - Layer 0 needs 6 scratch blocks without rewind, or 5 with rewind_len=32. - Layer 1 (stale 0-0): needs 0 scratch blocks (net_alloc_counts = 8 > 0) Chunk 2: resize(352, 256) -> 11 blocks. - - Layer 0 (stale 0-9): needs 1 scratch block [8, 9). delta_scratch = -5. New normal = 2. - net_alloc_counts = -3 < 0 + - Layer 0 needs 1 scratch block without rewind, or 0 with rewind_len=32. - Layer 1 (stale 0-3): needs 0 scratch blocks. delta_scratch = 0. New normal = 3. net_alloc_counts = 3 > 0 """ @@ -2352,7 +2368,7 @@ def test_scratch_chunk_size_variation(self): sliding_window_size=256, ), ], - enable_swa_scratch_reuse=True, + swa_scratch_reuse=SwaScratchReuseConfig(max_rewind_len=rewind_len), ) self.engine = FakeEngine(self.cfg) self.manager = KVCacheManager(self.cfg) @@ -2373,12 +2389,14 @@ def test_scratch_chunk_size_variation(self): lg_id_0 = LayerGroupId(0) lg_id_1 = LayerGroupId(1) - # Layer 0 should have 6 scratch blocks: range [0, 6) scratch_desc_0 = kv.get_scratch_desc(lg_id_0) self.assertIsNotNone(scratch_desc_0) - self.assertEqual(scratch_desc_0.range.beg, 0) - self.assertEqual(scratch_desc_0.range.end, 6) - self.assertEqual(len(scratch_desc_0.slot_ids), 6) + self.assertEqual(scratch_desc_0.range.beg, chunk1_scratch_range[0]) + self.assertEqual(scratch_desc_0.range.end, chunk1_scratch_range[1]) + self.assertEqual( + len(scratch_desc_0.slot_ids), + chunk1_scratch_range[1] - chunk1_scratch_range[0], + ) # Layer 1 should have 0 scratch blocks scratch_desc_1 = kv.get_scratch_desc(lg_id_1) @@ -2394,15 +2412,20 @@ def test_scratch_chunk_size_variation(self): self.assertTrue(success) self.assertFalse(kv.has_scratch_slots) - # Chunk 2: resize to 352 with history_length=256 + # Chunk 2: resize to 352 with history_length=256. + if rewind_len > 0: + with self.assertRaisesRegex(AssertionError, "old_capacity - max_rewind_len"): + kv.resize(352, 223) success = kv.resize(352, 256) self.assertTrue(success) - # Layer 0 should have 1 scratch block: range [8, 9) scratch_desc_0 = kv.get_scratch_desc(lg_id_0) - self.assertIsNotNone(scratch_desc_0) - self.assertEqual(scratch_desc_0.range.beg, 8) - self.assertEqual(scratch_desc_0.range.end, 9) + if chunk2_scratch_range is None: + self.assertIsNone(scratch_desc_0) + else: + self.assertIsNotNone(scratch_desc_0) + self.assertEqual(scratch_desc_0.range.beg, chunk2_scratch_range[0]) + self.assertEqual(scratch_desc_0.range.end, chunk2_scratch_range[1]) # Layer 1 should still have 0 scratch blocks scratch_desc_1 = kv.get_scratch_desc(lg_id_1) From d8f5b9b086ba1ca8edf9516174487281e81de789 Mon Sep 17 00:00:00 2001 From: Jiagan Cheng Date: Wed, 20 May 2026 23:19:52 -0700 Subject: [PATCH 2/5] [None][fix] use extra KV tokens for scratch rewind Signed-off-by: Jiagan Cheng --- .../sparse/deepseek_v4/cache_manager.py | 13 ++++++++----- tensorrt_llm/_torch/pyexecutor/resource_manager.py | 12 +++++++++--- 2 files changed, 17 insertions(+), 8 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py index 375d9253b9ea..31575a446c9f 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py @@ -799,17 +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, - swa_scratch_reuse=( - SwaScratchReuseConfig(max_rewind_len=self._max_draft_len) - if self.enable_swa_scratch_reuse - else None - ), + swa_scratch_reuse=scratch_reuse_config, layers=layers, typical_step=typical_step, constraints=constraints, diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index d5315db8d8eb..7752a77b2194 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -2449,15 +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, - swa_scratch_reuse=(SwaScratchReuseConfig( - max_rewind_len=self.max_draft_len) - if self.enable_swa_scratch_reuse else None), + swa_scratch_reuse=scratch_reuse_config, layers=[ AttentionLayerConfig( layer_id=layer_id, From eb67fa40eb50e485575910f0bf56e246f2220c41 Mon Sep 17 00:00:00 2001 From: Jiagan Cheng Date: Thu, 21 May 2026 03:13:20 -0700 Subject: [PATCH 3/5] [None][fix] enable DSV4 scratch reuse by default Signed-off-by: Jiagan Cheng --- .../sparse/deepseek_v4/cache_manager.py | 6 ++-- .../test_deepseek_v4_cache_manager.py | 32 +++++++++++-------- 2 files changed, 21 insertions(+), 17 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py index 31575a446c9f..21fd4cec24a2 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py @@ -61,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" @@ -120,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" diff --git a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py index 473bdccb91f6..02db4593cc15 100644 --- a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py +++ b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py @@ -77,7 +77,7 @@ def scratch_reuse_enabled(request, monkeypatch) -> bool: if request.param: monkeypatch.setenv(DSV4_ENABLE_SWA_SCRATCH_REUSE_ENV, "1") else: - monkeypatch.delenv(DSV4_ENABLE_SWA_SCRATCH_REUSE_ENV, raising=False) + monkeypatch.setenv(DSV4_ENABLE_SWA_SCRATCH_REUSE_ENV, "0") return request.param @@ -1425,7 +1425,7 @@ def test_copy_batch_indexer_compress_block_tables_matches_python_converter( cache_manager.free_resources(req) cache_manager.shutdown() - def test_swa_scratch_reuse_disabled_by_default_for_main_manager(self, monkeypatch): + def test_swa_scratch_reuse_enabled_by_default_for_main_manager(self, monkeypatch): monkeypatch.delenv(DSV4_ENABLE_SWA_SCRATCH_REUSE_ENV, raising=False) cache_manager, _ = self._create_deepseek_v4_cache_manager( tokens_per_block=self.tokens_per_block, @@ -1437,15 +1437,16 @@ def test_swa_scratch_reuse_disabled_by_default_for_main_manager(self, monkeypatc ) 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.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_enabled_by_env_for_main_manager(self, monkeypatch): - monkeypatch.setenv(DSV4_ENABLE_SWA_SCRATCH_REUSE_ENV, "1") + def test_swa_scratch_reuse_disabled_by_env_for_main_manager(self, monkeypatch): + monkeypatch.setenv(DSV4_ENABLE_SWA_SCRATCH_REUSE_ENV, "0") cache_manager, _ = self._create_deepseek_v4_cache_manager( tokens_per_block=self.tokens_per_block, max_batch_size=1, @@ -1456,20 +1457,23 @@ def test_swa_scratch_reuse_enabled_by_env_for_main_manager(self, monkeypatch): ) 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 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() - def test_swa_scratch_reuse_uses_spec_draft_len_for_rewind(self, monkeypatch): + def test_swa_scratch_reuse_uses_extra_kv_tokens_for_rewind(self, monkeypatch): monkeypatch.setenv(DSV4_ENABLE_SWA_SCRATCH_REUSE_ENV, "1") spec_config = SimpleNamespace( max_draft_len=7, max_total_draft_tokens=7, - spec_dec_mode=SimpleNamespace(use_one_engine=lambda: False), + 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, @@ -1484,7 +1488,7 @@ def test_swa_scratch_reuse_uses_spec_draft_len_for_rewind(self, monkeypatch): 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 + assert scratch_reuse.max_rewind_len == spec_config.max_draft_len - 1 finally: cache_manager.shutdown() From d05ac24b73680b328e4ad4872a5518ba95f5a2a6 Mon Sep 17 00:00:00 2001 From: Jiagan Cheng Date: Thu, 21 May 2026 19:54:03 -0700 Subject: [PATCH 4/5] test: update KV cache scratch resize coverage Signed-off-by: Jiagan Cheng --- .../kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py | 4 ---- 1 file changed, 4 deletions(-) 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 4a88cd2f49cf..1d6a26b9c147 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 @@ -2411,11 +2411,7 @@ def test_scratch_chunk_size_variation( success = kv.resume(stream) self.assertTrue(success) self.assertFalse(kv.has_scratch_slots) - # Chunk 2: resize to 352 with history_length=256. - if rewind_len > 0: - with self.assertRaisesRegex(AssertionError, "old_capacity - max_rewind_len"): - kv.resize(352, 223) success = kv.resize(352, 256) self.assertTrue(success) From b0209a404a1862d9b640b963d1ad3fb18efbdfaa Mon Sep 17 00:00:00 2001 From: Jiagan Cheng Date: Thu, 21 May 2026 23:52:44 -0700 Subject: [PATCH 5/5] style: format resource manager import Signed-off-by: Jiagan Cheng --- tensorrt_llm/_torch/pyexecutor/resource_manager.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index 7752a77b2194..3992e204aee6 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -48,8 +48,8 @@ # isort: off from tensorrt_llm.runtime.kv_cache_manager_v2 import ( DEFAULT_BEAM_INDEX, AttentionLayerConfig, BufferConfig, CacheTierConfig, - GpuCacheTierConfig, HostCacheTierConfig, PageIndexMode, - ReuseScope, SwaScratchReuseConfig) + 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 \