From 69372f353da3659436ce02c8b633b59984e78a9b Mon Sep 17 00:00:00 2001 From: tianruih Date: Tue, 14 Jul 2026 21:40:05 -0700 Subject: [PATCH 1/5] [None][feat] Request-carried KV-cache compression lengths Signed-off-by: tianruih --- .../_torch/kv_cache_compression/__init__.py | 0 tensorrt_llm/_torch/model_config.py | 6 +- tensorrt_llm/_torch/pyexecutor/_util.py | 24 ++- .../_torch/pyexecutor/kv_cache_manager_v2.py | 9 +- tensorrt_llm/_torch/pyexecutor/llm_request.py | 2 + .../_torch/pyexecutor/model_engine.py | 36 ++-- .../_torch/pyexecutor/model_loader.py | 2 + .../_torch/pyexecutor/resource_manager.py | 80 ++++++--- .../integration/test_lists/test-db/l0_a10.yml | 1 + .../test_kv_cache_compression_manager.py | 161 ++++++++++++++++-- .../test_kv_cache_v2_capacity_only.py | 154 +++++++++++++++++ 11 files changed, 416 insertions(+), 59 deletions(-) create mode 100644 tensorrt_llm/_torch/kv_cache_compression/__init__.py create mode 100644 tests/unittest/_torch/executor/test_kv_cache_v2_capacity_only.py diff --git a/tensorrt_llm/_torch/kv_cache_compression/__init__.py b/tensorrt_llm/_torch/kv_cache_compression/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py index dd92337543c6..0bd6b64e1fcd 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/tensorrt_llm/_torch/model_config.py @@ -51,8 +51,9 @@ if TYPE_CHECKING: from tensorrt_llm.bindings import ModelConfig as ModelConfigCpp - from tensorrt_llm.llmapi.llm_args import (DecodingBaseConfig, LoraConfig, - SparseAttentionConfig, + from tensorrt_llm.llmapi.llm_args import (DecodingBaseConfig, + KvCacheCompressionConfig, + LoraConfig, SparseAttentionConfig, SpeculativeConfig) TConfig = TypeVar("TConfig", bound=transformers.PretrainedConfig) @@ -162,6 +163,7 @@ class ModelConfig(Generic[TConfig]): spec_config: Optional["DecodingBaseConfig"] = None lora_config: Optional["LoraConfig"] = None sparse_attention_config: Optional["SparseAttentionConfig"] = None + kv_cache_compression_config: Optional["KvCacheCompressionConfig"] = None is_generation: bool = True is_encoder_decoder: bool = False diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 48d0c0dccdc4..feacedf85151 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -2035,13 +2035,17 @@ def _create_kv_cache_manager( def create_kv_cache_compression_manager( config: KvCacheCompressionConfig, kv_cache_manager: KVCacheManagerV2, + draft_kv_cache_manager: Optional[KVCacheManagerV2] = None, + spec_config: Optional[SpeculativeConfig] = None, ) -> Optional[BaseKVCacheCompressionManager]: """Build the KV-cache compression manager for ``config.algorithm``, or return None if no algorithm matches. Called from ``create_py_executor`` and registered as a resource manager, like the KV cache manager itself. Concrete algorithms add a dispatch branch - here; the framework ships none. + here; the framework ships none. Speculative-decoding compatibility is also + decided here: a compression manager is only created when the speculative + mode supports it, otherwise the run stays uncompressed. """ logger.warning( "KV-cache compression algorithm '%s' is not registered; running without " @@ -2256,7 +2260,12 @@ def create_py_executor_instance( "kv_cache_compression_config", None) if kv_cache_compression_config is not None: compression_manager = create_kv_cache_compression_manager( - kv_cache_compression_config, kv_cache_manager) + kv_cache_compression_config, + kv_cache_manager, + draft_kv_cache_manager=resources.get( + ResourceManagerType.DRAFT_KV_CACHE_MANAGER), + spec_config=spec_config, + ) if compression_manager is not None: resources[ResourceManagerType.KV_CACHE_COMPRESSION_MANAGER] = ( compression_manager) @@ -2268,17 +2277,16 @@ def create_py_executor_instance( if kv_cache_manager is not None: resource_manager.resource_managers.move_to_end( ResourceManagerType.KV_CACHE_MANAGER, last=True) - # Compression manager runs after the cache manager: reconciles history once it's resized. - if (ResourceManagerType.KV_CACHE_COMPRESSION_MANAGER - in resource_manager.resource_managers): - resource_manager.resource_managers.move_to_end( - ResourceManagerType.KV_CACHE_COMPRESSION_MANAGER, last=True) - cross_kv_cache_manager = resources.get( ResourceManagerType.CROSS_KV_CACHE_MANAGER) if cross_kv_cache_manager is not None: resource_manager.resource_managers.move_to_end( ResourceManagerType.CROSS_KV_CACHE_MANAGER, last=True) + # Compression is the final reconciler after every native KV manager. + if (ResourceManagerType.KV_CACHE_COMPRESSION_MANAGER + in resource_manager.resource_managers): + resource_manager.resource_managers.move_to_end( + ResourceManagerType.KV_CACHE_COMPRESSION_MANAGER, last=True) # When scheduler_capacity == 1, attention dp dummy request will prevent the scheduling of DISAGG_GENERATION_INIT. # Enlarge scheduler capacity to avoid DISAGG_GENERATION_INIT stuck in the scheduler. diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py index 5d22a2d395f2..260d2526d649 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py @@ -647,6 +647,8 @@ def __init__( layer_mask=layer_mask, ) self.is_draft = is_draft + # Set True by a compression manager; generation-step resize then leaves history untouched. + self.kv_compression_manages_history: bool = False self.enable_swa_scratch_reuse = ( kv_cache_config.enable_swa_scratch_reuse and not self.is_draft ) @@ -3061,12 +3063,15 @@ def update_resources( if req.state in (LlmRequestState.GENERATION_COMPLETE, LlmRequestState.CONTEXT_INIT) else kv_cache.capacity - req.py_rewind_len ) - success = kv_cache.resize(new_capacity, req.max_beam_num_tokens - 1) + history_length = ( + None if self.kv_compression_manages_history else req.max_beam_num_tokens - 1 + ) + success = kv_cache.resize(new_capacity, history_length) if not success: raise ValueError( f"Failed to resize KV cache for request {req.py_request_id} " f"to capacity {new_capacity} and history length " - f"{req.max_beam_num_tokens - 1} tokens at generation update" + f"{history_length} tokens at generation update" ) def copy_batch_block_offsets( diff --git a/tensorrt_llm/_torch/pyexecutor/llm_request.py b/tensorrt_llm/_torch/pyexecutor/llm_request.py index 63bfb318e284..00d01f65df69 100644 --- a/tensorrt_llm/_torch/pyexecutor/llm_request.py +++ b/tensorrt_llm/_torch/pyexecutor/llm_request.py @@ -730,6 +730,8 @@ def __init__( self.py_batch_idx = None self.py_draft_pages_allocated = 0 self.py_rewind_len = 0 + # Tokens physically evicted by KV-cache compression; deducted in the engine. + self.py_num_compressed_tokens = 0 self.py_draft_tokens = [] if self.draft_tokens is None else self.draft_tokens self.py_last_context_chunk = (None, None) self.py_draft_logits = None diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index b1a4f1b49cd6..a5f1d079bc46 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -3087,6 +3087,10 @@ def _apply_incremental_update_target( else: prompt_lengths[idx] = request.py_prompt_len + # Physical KV length for the kernels: subtract the tokens a + # KV-cache compression manager evicted (tracked on the request, + # 0 without compression). Position ids and the cached_tokens stat + # keep the logical count. if request.is_dummy: num_cached_tokens_per_seq[idx] = base_past_seen request.cached_tokens = base_past_seen @@ -3097,9 +3101,11 @@ def _apply_incremental_update_target( num_previous_batch] = request.py_batch_idx num_previous_batch += 1 - num_cached_tokens_per_seq[ - idx] = base_past_seen + num_tokens_per_extend_request - request.cached_tokens = num_cached_tokens_per_seq[idx].item() + request.cached_tokens = (base_past_seen + + num_tokens_per_extend_request) + num_cached_tokens_per_seq[idx] = ( + base_past_seen + num_tokens_per_extend_request - + request.py_num_compressed_tokens) request.py_batch_idx = request.py_seq_slot @@ -3341,8 +3347,9 @@ def append_cross_attention_state(request: LlmRequest, py_request_id] = request.py_num_accepted_draft_tokens_indices prompt_lengths.append(len(prompt_tokens)) past_seen_token_num = begin_compute - num_cached_tokens_per_seq.append(past_seen_token_num) - request.cached_tokens = num_cached_tokens_per_seq[-1] + num_cached_tokens_per_seq.append(past_seen_token_num - + request.py_num_compressed_tokens) + request.cached_tokens = past_seen_token_num append_cross_attention_state( request, project_encoder_output=not request.py_skip_cross_kv_projection @@ -3527,8 +3534,9 @@ def append_cross_attention_state(request: LlmRequest, list( range(past_seen_token_num, past_seen_token_num + 1 + num_draft_tokens))) - num_cached_tokens_per_seq.append(past_seen_token_num) - request.cached_tokens = num_cached_tokens_per_seq[-1] + num_cached_tokens_per_seq.append( + past_seen_token_num - request.py_num_compressed_tokens) + request.cached_tokens = past_seen_token_num # update batch index request.py_batch_idx = request.py_seq_slot else: @@ -3555,9 +3563,11 @@ def append_cross_attention_state(request: LlmRequest, previous_pos_indices.extend([previous_batch_idx] * runtime_tokens_per_gen_step) - num_cached_tokens_per_seq.append(past_seen_token_num + - runtime_tokens_per_gen_step) - request.cached_tokens = num_cached_tokens_per_seq[-1] + num_cached_tokens_per_seq.append( + past_seen_token_num + runtime_tokens_per_gen_step - + request.py_num_compressed_tokens) + request.cached_tokens = (past_seen_token_num + + runtime_tokens_per_gen_step) if self.enable_spec_decode and spec_config.spec_dec_mode.extend_ctx( self.attn_backend) and spec_config.is_linear_tree: prompt_lengths.append(runtime_tokens_per_gen_step) @@ -3614,7 +3624,8 @@ def append_cross_attention_state(request: LlmRequest, py_request_id] = request.py_num_accepted_draft_tokens_indices prompt_lengths.append(request.py_prompt_len) past_seen_token_num = begin_compute - num_cached_tokens_per_seq.append(past_seen_token_num) + num_cached_tokens_per_seq.append(past_seen_token_num - + request.py_num_compressed_tokens) append_cross_attention_state(request, project_encoder_output=False) # update batch index @@ -3691,7 +3702,8 @@ def append_cross_attention_state(request: LlmRequest, request.cached_tokens = past_seen_token_num for beam in range(beam_width): position_ids.append(position_id) - num_cached_tokens_per_seq.append(past_seen_token_num) + num_cached_tokens_per_seq.append( + past_seen_token_num - request.py_num_compressed_tokens) prompt_lengths.append(request.py_prompt_len) gather_ids.append(len(position_ids) - 1) diff --git a/tensorrt_llm/_torch/pyexecutor/model_loader.py b/tensorrt_llm/_torch/pyexecutor/model_loader.py index 4aa8ce703d4c..b0c2ed44ef34 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_loader.py +++ b/tensorrt_llm/_torch/pyexecutor/model_loader.py @@ -1166,6 +1166,8 @@ def _load_and_validate_config( force_dynamic_quantization=self.llm_args.force_dynamic_quantization, spec_config=self.spec_config, sparse_attention_config=self.sparse_attention_config, + kv_cache_compression_config=( + self.llm_args.kv_cache_compression_config), max_num_tokens=self.max_num_tokens, max_seq_len=self.max_seq_len, moe_max_num_tokens=self.llm_args.moe_config.max_num_tokens, diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index 486c43e8e269..c5f3b1947069 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -19,8 +19,8 @@ from abc import ABC, abstractmethod from collections import OrderedDict, defaultdict, deque from dataclasses import dataclass -from typing import (TYPE_CHECKING, Dict, Iterable, List, Optional, Sequence, - Set, Tuple, Union) +from typing import (TYPE_CHECKING, ClassVar, Dict, Iterable, List, Optional, + Sequence, Set, Tuple, Union) import torch from mpi4py import MPI @@ -2341,14 +2341,37 @@ class BaseKVCacheCompressionManager(BaseResourceManager): base implementations below translate those callbacks into the lifecycle hooks. - Concrete compression methods subclass this directly. All 4 hooks default to + Concrete compression methods subclass this directly. The hooks default to no-op; subclasses override what they need. The manager never inherits from any cache manager because this layer decides *how* the physical KV is used, not *what* physical KV exists. Subclasses hold ``KVCacheManagerV2`` as a tool. + + A subclass compacts through the ``KVCacheManagerV2`` it holds and records + the evicted count on ``LlmRequest.py_num_compressed_tokens``; the model + engine subtracts that count when building ``num_cached_tokens_per_seq``. """ - def __init__(self, kv_cache_manager: "KVCacheManagerV2"): + adjusts_generation_kv_length: ClassVar[bool] = False + """Whether this manager can make target and logical KV lengths diverge.""" + + physically_evicts_cached_tokens: ClassVar[bool] = False + """True for evicting methods; attention modules then keep RoPE unfused.""" + + def __init__( + self, + kv_cache_manager: "KVCacheManagerV2", + draft_kv_cache_manager: Optional["KVCacheManagerV2"] = None, + ): + from .kv_cache_manager_v2 import KVCacheManagerV2 + + if not isinstance(kv_cache_manager, KVCacheManagerV2): + raise TypeError("KV-cache compression requires KVCacheManagerV2") + if draft_kv_cache_manager is not None and not isinstance( + draft_kv_cache_manager, KVCacheManagerV2): + raise TypeError( + "draft KV-cache compression requires KVCacheManagerV2") self.kv_cache_manager = kv_cache_manager + self.draft_kv_cache_manager = draft_kv_cache_manager # Compression evicts/rewrites stored keys and values, so a shared prefix # block is no longer safe to reuse (same constraint as RocketKVCacheManager). if kv_cache_manager.enable_block_reuse: @@ -2356,9 +2379,23 @@ def __init__(self, kv_cache_manager: "KVCacheManagerV2"): f"{type(self).__name__} changes stored keys and values and cannot " f"run with KV-cache block reuse. Set " f"KvCacheConfig.enable_block_reuse to False.") + kv_cache_manager.kv_compression_manages_history = self.adjusts_generation_kv_length + if draft_kv_cache_manager is not None: + # The draft cache is compacted together with the target. + draft_kv_cache_manager.kv_compression_manages_history = ( + self.adjusts_generation_kv_length) + + @classmethod + def is_eviction_method(cls) -> bool: + """Whether this method physically evicts cached tokens.""" + return cls.physically_evicts_cached_tokens + + @property + def has_independent_draft_kv_cache(self) -> bool: + return self.draft_kv_cache_manager is not None # ================================================================== # - # KV-cache lifecycle hooks (4, in temporal order). # + # KV-cache lifecycle hooks (5, in temporal order). # # Subclasses override what they need; all default to no-op. # # ================================================================== # @@ -2369,20 +2406,23 @@ def on_request_init(self, request: "LlmRequest", **kwargs) -> None: scoring buffers). """ - def on_context_step_end( + def on_context_step_end(self, requests: List["LlmRequest"], + **kwargs) -> None: + """Fired once per iteration with the requests whose prefill finished + (their final chunk) this step. Batched like the generation hook so a + one-shot prefill-end eviction can process the cohort in one launch. + """ + + def on_generation_step_begin( self, - request: "LlmRequest", - metadata: "AttentionMetadata", + scheduled_batch: "ScheduledRequests", **kwargs, ) -> None: - """Fired once per request, when its prefill finishes (its final - chunk). Override for a one-shot prefill-end eviction. - """ + """Fired once per generation step before this step's forward.""" def on_generation_step_end( self, scheduled_batch: "ScheduledRequests", - attn_metadata: "AttentionMetadata", **kwargs, ) -> None: """Fired once per generation step, after every layer's forward @@ -2422,6 +2462,7 @@ def prepare_resources(self, scheduled_batch: "ScheduledRequests") -> None: for req in scheduled_batch.context_requests: if req.is_first_context_chunk: self.on_request_init(req) + self.on_generation_step_begin(scheduled_batch) def update_resources( self, @@ -2429,8 +2470,8 @@ def update_resources( attn_metadata: Optional["AttentionMetadata"] = None, kv_cache_dtype_byte_size: Optional[float] = None, ) -> None: - """Fire :meth:`on_context_step_end` once per request, on the iteration its - final prefill chunk runs, then :meth:`on_generation_step_end` once. + """Fire :meth:`on_context_step_end` with the requests whose final + prefill chunk ran this iteration, then :meth:`on_generation_step_end`. Uses the scheduler's ``context_requests_last_chunk`` split (computed at schedule time from ``is_last_context_chunk``) rather than tracking @@ -2441,9 +2482,10 @@ def update_resources( managers so PyExecutor passes ``attn_metadata`` / ``kv_cache_dtype_byte_size`` through transparently. """ - for req in scheduled_batch.context_requests_last_chunk: - self.on_context_step_end(req, attn_metadata) - self.on_generation_step_end(scheduled_batch, attn_metadata) + if scheduled_batch.context_requests_last_chunk: + self.on_context_step_end( + scheduled_batch.context_requests_last_chunk) + self.on_generation_step_end(scheduled_batch) def free_resources(self, request: "LlmRequest") -> None: """Fire :meth:`on_request_finish`.""" @@ -2480,9 +2522,9 @@ def update_resources( attn_metadata: Optional["AttentionMetadata"] = None, kv_cache_dtype_byte_size: Optional[float] = None, ): - for _, resource_manager in self.resource_managers.items(): + for resource_type, resource_manager in self.resource_managers.items(): if hasattr(resource_manager, "update_resources"): - if isinstance(resource_manager, KVCacheManager): + if resource_type == ResourceManagerType.KV_CACHE_MANAGER: resource_manager.update_resources(scheduled_batch, attn_metadata, kv_cache_dtype_byte_size) diff --git a/tests/integration/test_lists/test-db/l0_a10.yml b/tests/integration/test_lists/test-db/l0_a10.yml index 1e6836fc7cd5..cdd02b720890 100644 --- a/tests/integration/test_lists/test-db/l0_a10.yml +++ b/tests/integration/test_lists/test-db/l0_a10.yml @@ -37,6 +37,7 @@ l0_a10: - unittest/_torch/executor/test_kv_pool_rebalance.py - unittest/_torch/executor/test_disagg_index_mapper_early_release.py - unittest/_torch/executor/test_kv_cache_compression_manager.py + - unittest/_torch/executor/test_kv_cache_v2_capacity_only.py - unittest/_torch/executor/test_error_classification.py - unittest/_torch/modules/dwdp/test_dwdp_fixup_moe_backends.py - unittest/_torch/modules/dwdp/test_dwdp_manager.py diff --git a/tests/unittest/_torch/executor/test_kv_cache_compression_manager.py b/tests/unittest/_torch/executor/test_kv_cache_compression_manager.py index 13cc7fd9d6d2..e7448b79ff33 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_compression_manager.py +++ b/tests/unittest/_torch/executor/test_kv_cache_compression_manager.py @@ -12,8 +12,8 @@ - The resource-manager API -> lifecycle-hook translation, gated on PyExecutor's own signals: ``prepare_resources`` fires ``on_request_init`` on each request's first prefill chunk (``is_first_context_chunk``); - ``update_resources`` fires ``on_context_step_end`` for each request in - ``context_requests_last_chunk`` + one ``on_generation_step_end`` per + ``update_resources`` fires ``on_context_step_end`` once with the + ``context_requests_last_chunk`` list + one ``on_generation_step_end`` per iteration; ``free_resources`` fires ``on_request_finish``. - :func:`create_kv_cache_compression_manager` factory. @@ -22,6 +22,8 @@ lives in ``_util.py`` next to ``_create_kv_cache_manager``. """ +from types import SimpleNamespace +from typing import ClassVar from unittest.mock import MagicMock, patch import pytest @@ -31,6 +33,8 @@ from tensorrt_llm._torch.pyexecutor.resource_manager import ( BaseKVCacheCompressionManager, BaseResourceManager, + ResourceManager, + ResourceManagerType, ) # ---------------------------------------------------------------------- # @@ -57,23 +61,35 @@ class _MockCompressionManager(_RecordingMixin, BaseKVCacheCompressionManager): def on_request_init(self, request): self._record("on_request_init") - def on_context_step_end(self, request, metadata): - self._record("on_context_step_end") + def on_context_step_end(self, requests): + self._record(f"on_context_step_end[{len(requests)}]") - def on_generation_step_end(self, scheduled_batch, attn_metadata): + def on_generation_step_end(self, scheduled_batch): self._record("on_generation_step_end") def on_request_finish(self, request): self._record("on_request_finish") +class _LengthAdjustingCompressionManager(BaseKVCacheCompressionManager): + adjusts_generation_kv_length: ClassVar[bool] = True + + +def _v2_manager(*, is_draft: bool): + from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2 + + manager = KVCacheManagerV2.__new__(KVCacheManagerV2) + manager.enable_block_reuse = False + manager.kv_compression_manages_history = False + manager.is_draft = is_draft + return manager + + @pytest.fixture def fake_kv_cache_manager(): """A stand-in KVCacheManagerV2. The framework reads enable_block_reuse off it in __init__; default it to False, like a normal run with reuse off.""" - m = MagicMock(name="fake_KVCacheManagerV2") - m.enable_block_reuse = False - return m + return _v2_manager(is_draft=False) def _req(rid, first_chunk=True): @@ -103,10 +119,10 @@ def test_inherits_base_resource_manager(self): def test_four_hooks_default_noop(self, fake_kv_cache_manager): m = BaseKVCacheCompressionManager(fake_kv_cache_manager) - meta = MagicMock() assert m.on_request_init(MagicMock()) is None - assert m.on_context_step_end(MagicMock(), meta) is None - assert m.on_generation_step_end(MagicMock(), meta) is None + assert m.on_context_step_end([MagicMock()]) is None + assert m.on_generation_step_begin(MagicMock()) is None + assert m.on_generation_step_end(MagicMock()) is None assert m.on_request_finish(MagicMock()) is None def test_hooks_accept_extra_kwargs(self, fake_kv_cache_manager): @@ -114,7 +130,7 @@ def test_hooks_accept_extra_kwargs(self, fake_kv_cache_manager): # existing overrides. m = BaseKVCacheCompressionManager(fake_kv_cache_manager) assert m.on_request_init(MagicMock(), future_arg=1) is None - assert m.on_generation_step_end(MagicMock(), MagicMock(), future_arg=1) is None + assert m.on_generation_step_end(MagicMock(), future_arg=1) is None def test_resource_counts_are_zero(self, fake_kv_cache_manager): m = BaseKVCacheCompressionManager(fake_kv_cache_manager) @@ -123,6 +139,42 @@ def test_resource_counts_are_zero(self, fake_kv_cache_manager): assert m.get_max_resource_count() == 0 assert m.get_needed_resource_to_completion(MagicMock()) == 0 + def test_length_adjustment_marks_target_and_draft_v2(self): + # The draft cache is compacted together with the target, so both + # managers diverge from the logical length in the same way. + target = _v2_manager(is_draft=False) + draft = _v2_manager(is_draft=True) + + manager = _LengthAdjustingCompressionManager(target, draft) + + assert manager.kv_cache_manager is target + assert manager.draft_kv_cache_manager is draft + assert manager.has_independent_draft_kv_cache + assert target.kv_compression_manages_history is True + assert draft.kv_compression_manages_history is True + + def test_rejects_non_v2_ownership(self): + with pytest.raises(TypeError, match="requires KVCacheManagerV2"): + BaseKVCacheCompressionManager(MagicMock()) + with pytest.raises(TypeError, match="requires KVCacheManagerV2"): + BaseKVCacheCompressionManager(_v2_manager(is_draft=False), MagicMock()) + + def test_request_field_defaults_to_zero(self): + """LlmRequest carries the compression count (the manager's only + channel to the runtime); a fresh request must default to 0 so runs + without a compression manager are unchanged.""" + from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest + from tensorrt_llm.bindings import SamplingConfig + + request = LlmRequest( + request_id=1, + max_new_tokens=8, + input_tokens=[1, 2, 3], + sampling_config=SamplingConfig(), + is_streaming=False, + ) + assert request.py_num_compressed_tokens == 0 + # ---------------------------------------------------------------------- # # 2. Resource-manager API -> lifecycle-hook translation # @@ -131,6 +183,46 @@ def test_resource_counts_are_zero(self, fake_kv_cache_manager): class TestResourceManagerAPI: + def test_target_update_receives_metadata_before_final_compression(self): + calls = [] + metadata = MagicMock(name="attention_metadata") + draft = MagicMock(name="draft_kv_cache_manager") + target = MagicMock(name="target_kv_cache_manager") + compression = MagicMock(name="compression_manager") + draft.update_resources.side_effect = lambda *args: calls.append(("draft", args)) + target.update_resources.side_effect = lambda *args: calls.append(("target", args)) + compression.update_resources.side_effect = lambda *args: calls.append(("compression", args)) + manager = ResourceManager( + { + ResourceManagerType.DRAFT_KV_CACHE_MANAGER: draft, + ResourceManagerType.KV_CACHE_MANAGER: target, + ResourceManagerType.KV_CACHE_COMPRESSION_MANAGER: compression, + } + ) + batch = _batch(generation=[_req(1)]) + + manager.update_resources(batch, metadata, 2.0) + + assert calls == [ + ("draft", (batch,)), + ("target", (batch, metadata, 2.0)), + ("compression", (batch,)), + ] + + def test_real_v2_target_receives_relocation_metadata(self): + from tensorrt_llm._torch.pyexecutor import kv_cache_manager_v2 as kv_cache_v2_module + + target = _v2_manager(is_draft=False) + target.kv_cache_map = {} + batch = _batch(generation=[_req(1)]) + metadata = MagicMock(name="attention_metadata") + manager = ResourceManager({ResourceManagerType.KV_CACHE_MANAGER: target}) + + with patch.object(kv_cache_v2_module, "_update_kv_cache_draft_token_location") as relocate: + manager.update_resources(batch, metadata, 2.0) + + relocate.assert_called_once_with(target, batch, metadata, 2.0) + def test_prepare_fires_init_on_first_chunk_only(self, fake_kv_cache_manager): rec = [] m = _MockCompressionManager(fake_kv_cache_manager, rec, "s") @@ -144,9 +236,10 @@ def test_update_fires_context_end_on_last_chunk(self, fake_kv_cache_manager): rec = [] m = _MockCompressionManager(fake_kv_cache_manager, rec, "s") req = _req(1) - # Request's final prefill chunk this iteration -> context_step_end fires. - m.update_resources(_batch(generation=[req], last_chunk=[req]), attn_metadata=MagicMock()) - assert "s:on_context_step_end" in rec + # Final prefill chunks this iteration -> one batched context_step_end. + req2 = _req(2) + m.update_resources(_batch(generation=[req], last_chunk=[req, req2])) + assert "s:on_context_step_end[2]" in rec assert rec[-1] == "s:on_generation_step_end" # Subsequent decode iteration (not in last_chunk) -> no context_step_end. rec.clear() @@ -185,6 +278,42 @@ def test_warns_for_unregistered_algorithm(self, fake_kv_cache_manager): create_kv_cache_compression_manager(cfg, fake_kv_cache_manager) mock_logger.warning.assert_called_once() + def test_factory_accepts_independent_draft_manager(self): + from tensorrt_llm._torch.speculative.interface import SpeculativeDecodingMode + + cfg = MagicMock() + cfg.algorithm = "made_up_method" + target = _v2_manager(is_draft=False) + draft = _v2_manager(is_draft=True) + spec_config = SimpleNamespace(spec_dec_mode=SpeculativeDecodingMode.EAGLE3_ONE_MODEL) + + assert ( + create_kv_cache_compression_manager( + cfg, + target, + draft_kv_cache_manager=draft, + spec_config=spec_config, + ) + is None + ) + + def test_eviction_method_predicate_defaults_false(self): + # The base is not an evicting method, so the factory's speculative + # mode gate never restricts it: methods that do not touch the draft + # KV (e.g. offloading) work with any speculative mode. + m = BaseKVCacheCompressionManager(_v2_manager(is_draft=False)) + assert m.is_eviction_method() is False + assert not hasattr(m, "spec_config") + + def test_eviction_method_predicate_follows_class_flag(self): + # The factory's speculative gate reads this predicate: an evicting + # method (physically_evicts_cached_tokens True) only accepts modes + # whose draft KV is a standard paged cache in the same forward. + class _EvictingManager(BaseKVCacheCompressionManager): + physically_evicts_cached_tokens = True + + assert _EvictingManager.is_eviction_method() is True + # ---------------------------------------------------------------------- # # 4. Canonical names live in resource_manager, not in the sparse module # @@ -219,7 +348,7 @@ class TestBlockReuseGuard: and values, the same check RocketKVCacheManager makes.""" def _mgr(self, enable_block_reuse): - m = MagicMock(name="KVCacheManagerV2") + m = _v2_manager(is_draft=False) m.enable_block_reuse = enable_block_reuse return m diff --git a/tests/unittest/_torch/executor/test_kv_cache_v2_capacity_only.py b/tests/unittest/_torch/executor/test_kv_cache_v2_capacity_only.py new file mode 100644 index 000000000000..df38d7bba8fc --- /dev/null +++ b/tests/unittest/_torch/executor/test_kv_cache_v2_capacity_only.py @@ -0,0 +1,154 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest + +import tensorrt_llm +import tensorrt_llm.bindings +from tensorrt_llm._torch.pyexecutor import kv_cache_manager_v2 as kv_cache_v2_module +from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2 +from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, LlmRequestState, SamplingConfig + +DataType = tensorrt_llm.bindings.DataType +CacheType = tensorrt_llm.bindings.internal.batch_manager.CacheType + + +def _manager(*, is_draft: bool, kv_compression_manages_history: bool = False) -> KVCacheManagerV2: + manager = KVCacheManagerV2.__new__(KVCacheManagerV2) + manager.is_draft = is_draft + manager.kv_compression_manages_history = kv_compression_manages_history + manager.kv_cache_map = {} + return manager + + +def _request(request_id: int, *, rewind: int = 0, complete: bool = False) -> SimpleNamespace: + return SimpleNamespace( + py_request_id=request_id, + py_rewind_len=rewind, + max_beam_num_tokens=201, + state=LlmRequestState.GENERATION_COMPLETE + if complete + else LlmRequestState.GENERATION_IN_PROGRESS, + ) + + +def _cache(*, capacity: int = 256, active: bool = True) -> MagicMock: + cache = MagicMock() + cache.capacity = capacity + cache.is_active = active + cache.resize.return_value = True + return cache + + +@pytest.fixture(autouse=True) +def _disable_draft_token_relocation(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(kv_cache_v2_module, "_update_kv_cache_draft_token_location", MagicMock()) + + +def test_manager_initializes_capacity_only_policy_to_false() -> None: + class StopInitialization(RuntimeError): + pass + + class StopAfterPolicyConfig: + @property + def enable_swa_scratch_reuse(self): + raise StopInitialization + + manager = KVCacheManagerV2.__new__(KVCacheManagerV2) + mapping = SimpleNamespace(cp_config={}) + + with ( + patch.object(kv_cache_v2_module, "get_pp_layers", return_value=([0], 1)), + pytest.raises(StopInitialization), + ): + manager.__init__( + StopAfterPolicyConfig(), + kv_cache_v2_module.CacheTypeCpp.SELF, + num_layers=1, + num_kv_heads=1, + head_dim=128, + tokens_per_block=64, + max_seq_len=256, + max_batch_size=1, + mapping=mapping, + ) + + assert manager.kv_compression_manages_history is False + + +def test_default_generation_resize_updates_capacity_and_history() -> None: + manager = _manager(is_draft=False) + request = _request(1, rewind=3) + cache = _cache() + manager.kv_cache_map[request.py_request_id] = cache + + manager.update_resources(SimpleNamespace(generation_requests=[request])) + + cache.resize.assert_called_once_with(253, 200) + + +def test_capacity_only_is_scoped_to_target_manager() -> None: + request = _request(1, rewind=3) + batch = SimpleNamespace(generation_requests=[request]) + target = _manager(is_draft=False, kv_compression_manages_history=True) + draft = _manager(is_draft=True) + target_cache = _cache() + draft_cache = _cache() + target.kv_cache_map[request.py_request_id] = target_cache + draft.kv_cache_map[request.py_request_id] = draft_cache + + draft.update_resources(batch) + target.update_resources(batch) + + draft_cache.resize.assert_called_once_with(253, 200) + target_cache.resize.assert_called_once_with(253, None) + + +def test_capacity_only_completion_preserves_history() -> None: + manager = _manager(is_draft=False, kv_compression_manages_history=True) + request = _request(1, complete=True) + cache = _cache() + manager.kv_cache_map[request.py_request_id] = cache + + manager.update_resources(SimpleNamespace(generation_requests=[request])) + + cache.resize.assert_called_once_with(None, None) + + +def test_capacity_only_skips_suspended_cache() -> None: + manager = _manager(is_draft=False, kv_compression_manages_history=True) + request = _request(1, rewind=3) + cache = _cache(active=False) + manager.kv_cache_map[request.py_request_id] = cache + + manager.update_resources(SimpleNamespace(generation_requests=[request])) + + cache.resize.assert_not_called() + + +def test_generation_update_has_no_request_compaction_marker() -> None: + manager = _manager(is_draft=False, kv_compression_manages_history=True) + request = _request(1, rewind=3) + cache = _cache() + manager.kv_cache_map[request.py_request_id] = cache + + manager.update_resources(SimpleNamespace(generation_requests=[request])) + + assert "py_kv_cache_kv_compression_manages_history" not in vars(request) + assert "py_kv_cache_compaction" not in vars(request) + + +def test_llm_request_has_no_compression_consumer_marker() -> None: + request = LlmRequest( + request_id=1, + max_new_tokens=1, + input_tokens=[1], + sampling_config=SamplingConfig(1), + is_streaming=False, + ) + + assert "py_kv_cache_kv_compression_manages_history" not in vars(request) + assert "py_kv_cache_compaction" not in vars(request) From ec747810cd32194daddad0c3db2270dd576ff494 Mon Sep 17 00:00:00 2001 From: tianruih Date: Wed, 15 Jul 2026 21:16:27 -0700 Subject: [PATCH 2/5] [None][refactor] Decide speculative compatibility at the executor call site Signed-off-by: tianruih --- tensorrt_llm/_torch/pyexecutor/_util.py | 54 ++++++++++++++----- .../_torch/pyexecutor/resource_manager.py | 8 --- tensorrt_llm/llmapi/llm_args.py | 4 ++ .../test_kv_cache_compression_manager.py | 29 +++++----- 4 files changed, 58 insertions(+), 37 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index feacedf85151..f9775b4ee6f6 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -2032,20 +2032,43 @@ def _create_kv_cache_manager( return kv_cache_manager +def kv_cache_compression_supported_with_spec( + config: KvCacheCompressionConfig, + spec_config: Optional[SpeculativeConfig], + draft_kv_cache_manager: Optional[KVCacheManagerV2], +) -> bool: + """Decide, before any manager is created, whether the configured + compression method supports the speculative setup. Logs the reason and + returns False when it does not; the run then stays uncompressed.""" + if spec_config is None or not config.is_eviction_method(): + return True + # Evicting methods co-compact the draft KV, so they only support spec + # modes whose draft KV is a standard paged cache in the same forward + # (one-model speculation). + mode = spec_config.spec_dec_mode + if not (mode.is_mtp_one_model() or mode.is_eagle3_one_model()): + logger.warning( + "KV-cache compression algorithm '%s' evicts cached tokens and " + "does not support speculative decoding mode %s (the draft KV must " + "be a standard paged cache compacted together with the target); " + "running without a compression manager.", config.algorithm, + mode.name) + return False + return True + + def create_kv_cache_compression_manager( config: KvCacheCompressionConfig, kv_cache_manager: KVCacheManagerV2, draft_kv_cache_manager: Optional[KVCacheManagerV2] = None, - spec_config: Optional[SpeculativeConfig] = None, ) -> Optional[BaseKVCacheCompressionManager]: """Build the KV-cache compression manager for ``config.algorithm``, or return None if no algorithm matches. Called from ``create_py_executor`` and registered as a resource manager, like the KV cache manager itself. Concrete algorithms add a dispatch branch - here; the framework ships none. Speculative-decoding compatibility is also - decided here: a compression manager is only created when the speculative - mode supports it, otherwise the run stays uncompressed. + here; the framework ships none. Speculative-decoding compatibility is + decided by the caller via ``kv_cache_compression_supported_with_spec``. """ logger.warning( "KV-cache compression algorithm '%s' is not registered; running without " @@ -2259,16 +2282,19 @@ def create_py_executor_instance( kv_cache_compression_config = getattr(llm_args, "kv_cache_compression_config", None) if kv_cache_compression_config is not None: - compression_manager = create_kv_cache_compression_manager( - kv_cache_compression_config, - kv_cache_manager, - draft_kv_cache_manager=resources.get( - ResourceManagerType.DRAFT_KV_CACHE_MANAGER), - spec_config=spec_config, - ) - if compression_manager is not None: - resources[ResourceManagerType.KV_CACHE_COMPRESSION_MANAGER] = ( - compression_manager) + draft_kv_cache_manager = resources.get( + ResourceManagerType.DRAFT_KV_CACHE_MANAGER) + if kv_cache_compression_supported_with_spec(kv_cache_compression_config, + spec_config, + draft_kv_cache_manager): + compression_manager = create_kv_cache_compression_manager( + kv_cache_compression_config, + kv_cache_manager, + draft_kv_cache_manager=draft_kv_cache_manager, + ) + if compression_manager is not None: + resources[ResourceManagerType.KV_CACHE_COMPRESSION_MANAGER] = ( + compression_manager) resource_manager = ResourceManager(resources) diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index c5f3b1947069..58ec6ab41485 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -2354,9 +2354,6 @@ class BaseKVCacheCompressionManager(BaseResourceManager): adjusts_generation_kv_length: ClassVar[bool] = False """Whether this manager can make target and logical KV lengths diverge.""" - physically_evicts_cached_tokens: ClassVar[bool] = False - """True for evicting methods; attention modules then keep RoPE unfused.""" - def __init__( self, kv_cache_manager: "KVCacheManagerV2", @@ -2385,11 +2382,6 @@ def __init__( draft_kv_cache_manager.kv_compression_manages_history = ( self.adjusts_generation_kv_length) - @classmethod - def is_eviction_method(cls) -> bool: - """Whether this method physically evicts cached tokens.""" - return cls.physically_evicts_cached_tokens - @property def has_independent_draft_kv_cache(self) -> bool: return self.draft_kv_cache_manager is not None diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index fd3276163831..5fa1c4012264 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -3268,6 +3268,10 @@ class KvCacheCompressionConfig(StrictBaseModel): "compression manager is built. Concrete algorithm configs subclass this " "and set the value.") + def is_eviction_method(self) -> bool: + """Whether this method physically evicts cached tokens.""" + return False + @PybindMirror.mirror_pybind_fields(_AgentTreeConfig) class AgentTreeConfig(StrictBaseModel, PybindMirror): diff --git a/tests/unittest/_torch/executor/test_kv_cache_compression_manager.py b/tests/unittest/_torch/executor/test_kv_cache_compression_manager.py index e7448b79ff33..b95810fe503c 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_compression_manager.py +++ b/tests/unittest/_torch/executor/test_kv_cache_compression_manager.py @@ -279,40 +279,39 @@ def test_warns_for_unregistered_algorithm(self, fake_kv_cache_manager): mock_logger.warning.assert_called_once() def test_factory_accepts_independent_draft_manager(self): - from tensorrt_llm._torch.speculative.interface import SpeculativeDecodingMode - cfg = MagicMock() cfg.algorithm = "made_up_method" target = _v2_manager(is_draft=False) draft = _v2_manager(is_draft=True) - spec_config = SimpleNamespace(spec_dec_mode=SpeculativeDecodingMode.EAGLE3_ONE_MODEL) assert ( create_kv_cache_compression_manager( cfg, target, draft_kv_cache_manager=draft, - spec_config=spec_config, ) is None ) def test_eviction_method_predicate_defaults_false(self): - # The base is not an evicting method, so the factory's speculative - # mode gate never restricts it: methods that do not touch the draft - # KV (e.g. offloading) work with any speculative mode. + # Non-evicting methods (e.g. offloading) are never restricted by the + # speculative mode: the call-site gate reads this config predicate. + from tensorrt_llm.llmapi.llm_args import KvCacheCompressionConfig + + config = KvCacheCompressionConfig(algorithm="offload") + assert config.is_eviction_method() is False m = BaseKVCacheCompressionManager(_v2_manager(is_draft=False)) - assert m.is_eviction_method() is False assert not hasattr(m, "spec_config") - def test_eviction_method_predicate_follows_class_flag(self): - # The factory's speculative gate reads this predicate: an evicting - # method (physically_evicts_cached_tokens True) only accepts modes - # whose draft KV is a standard paged cache in the same forward. - class _EvictingManager(BaseKVCacheCompressionManager): - physically_evicts_cached_tokens = True + def test_spec_gate_only_restricts_eviction_methods(self): + from tensorrt_llm._torch.pyexecutor._util import kv_cache_compression_supported_with_spec + from tensorrt_llm._torch.speculative.interface import SpeculativeDecodingMode + from tensorrt_llm.llmapi.llm_args import KvCacheCompressionConfig - assert _EvictingManager.is_eviction_method() is True + config = KvCacheCompressionConfig(algorithm="offload") + spec_config = SimpleNamespace(spec_dec_mode=SpeculativeDecodingMode.DFLASH) + assert kv_cache_compression_supported_with_spec(config, spec_config, None) is True + assert kv_cache_compression_supported_with_spec(config, None, None) is True # ---------------------------------------------------------------------- # From 9eec866a828adee6fab2542ca21476ccf1fc1ab1 Mon Sep 17 00:00:00 2001 From: tianruih Date: Wed, 15 Jul 2026 21:20:01 -0700 Subject: [PATCH 3/5] [None][refactor] Carry algorithm traits on a KV-cache compression mode Signed-off-by: tianruih --- .../_torch/kv_cache_compression/interface.py | 30 +++++++++++++++++++ tensorrt_llm/_torch/pyexecutor/_util.py | 3 +- tensorrt_llm/llmapi/llm_args.py | 10 +++++-- .../test_kv_cache_compression_manager.py | 2 +- 4 files changed, 40 insertions(+), 5 deletions(-) create mode 100644 tensorrt_llm/_torch/kv_cache_compression/interface.py diff --git a/tensorrt_llm/_torch/kv_cache_compression/interface.py b/tensorrt_llm/_torch/kv_cache_compression/interface.py new file mode 100644 index 000000000000..8ac47c98a640 --- /dev/null +++ b/tensorrt_llm/_torch/kv_cache_compression/interface.py @@ -0,0 +1,30 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from enum import IntEnum, auto +from typing import Optional + + +class KvCacheCompressionMode(IntEnum): + """Algorithm-level traits of a KV-cache compression method. + + Mirrors ``SpeculativeDecodingMode``: configs map their ``algorithm`` + string to a member here, and callers read ``is_*`` predicates instead of + comparing strings. + """ + + NONE = auto() + + def is_eviction_method(self): + """Whether this method physically evicts cached tokens. Evicting + algorithms add their member and extend this predicate.""" + return False + + @staticmethod + def from_string(name: Optional[str]) -> "KvCacheCompressionMode": + if name is None: + return KvCacheCompressionMode.NONE + try: + return KvCacheCompressionMode[name.upper()] + except KeyError: + return KvCacheCompressionMode.NONE diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index f9775b4ee6f6..8ce2d2564f85 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -2040,7 +2040,8 @@ def kv_cache_compression_supported_with_spec( """Decide, before any manager is created, whether the configured compression method supports the speculative setup. Logs the reason and returns False when it does not; the run then stays uncompressed.""" - if spec_config is None or not config.is_eviction_method(): + if (spec_config is None + or not config.kv_cache_compression_mode.is_eviction_method()): return True # Evicting methods co-compact the draft KV, so they only support spec # modes whose draft KV is a standard paged cache in the same forward diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 5fa1c4012264..63002d47e7b3 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -3268,9 +3268,13 @@ class KvCacheCompressionConfig(StrictBaseModel): "compression manager is built. Concrete algorithm configs subclass this " "and set the value.") - def is_eviction_method(self) -> bool: - """Whether this method physically evicts cached tokens.""" - return False + @property + def kv_cache_compression_mode(self): + # The mode carries algorithm-level traits (``is_*`` predicates) the + # raw algorithm string does not. + from tensorrt_llm._torch.kv_cache_compression.interface import \ + KvCacheCompressionMode + return KvCacheCompressionMode.from_string(self.algorithm) @PybindMirror.mirror_pybind_fields(_AgentTreeConfig) diff --git a/tests/unittest/_torch/executor/test_kv_cache_compression_manager.py b/tests/unittest/_torch/executor/test_kv_cache_compression_manager.py index b95810fe503c..8a9d16974146 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_compression_manager.py +++ b/tests/unittest/_torch/executor/test_kv_cache_compression_manager.py @@ -299,7 +299,7 @@ def test_eviction_method_predicate_defaults_false(self): from tensorrt_llm.llmapi.llm_args import KvCacheCompressionConfig config = KvCacheCompressionConfig(algorithm="offload") - assert config.is_eviction_method() is False + assert config.kv_cache_compression_mode.is_eviction_method() is False m = BaseKVCacheCompressionManager(_v2_manager(is_draft=False)) assert not hasattr(m, "spec_config") From 6cedd04dc11357383c438cec887b0de2e90fe4f5 Mon Sep 17 00:00:00 2001 From: tianruih Date: Wed, 15 Jul 2026 21:28:42 -0700 Subject: [PATCH 4/5] [None][chore] Trim the compression-mode docstring Signed-off-by: tianruih --- tensorrt_llm/_torch/kv_cache_compression/interface.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/interface.py b/tensorrt_llm/_torch/kv_cache_compression/interface.py index 8ac47c98a640..cd4e7bf32068 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/interface.py +++ b/tensorrt_llm/_torch/kv_cache_compression/interface.py @@ -8,9 +8,8 @@ class KvCacheCompressionMode(IntEnum): """Algorithm-level traits of a KV-cache compression method. - Mirrors ``SpeculativeDecodingMode``: configs map their ``algorithm`` - string to a member here, and callers read ``is_*`` predicates instead of - comparing strings. + Configs map their ``algorithm`` string to a member here; callers read the + ``is_*`` predicates instead of comparing strings. """ NONE = auto() From c4d80d3523c029e7d396786ab31bfbb15506d4b2 Mon Sep 17 00:00:00 2001 From: tianruih Date: Wed, 15 Jul 2026 22:04:37 -0700 Subject: [PATCH 5/5] [None][fix] Fail fast on speculative modes compression cannot support Signed-off-by: tianruih --- tensorrt_llm/_torch/pyexecutor/_util.py | 52 ++++++++----------- .../test_kv_cache_compression_manager.py | 7 +-- 2 files changed, 27 insertions(+), 32 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 8ce2d2564f85..673ad835a67e 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -2032,30 +2032,24 @@ def _create_kv_cache_manager( return kv_cache_manager -def kv_cache_compression_supported_with_spec( +def validate_kv_cache_compression_with_spec( config: KvCacheCompressionConfig, spec_config: Optional[SpeculativeConfig], draft_kv_cache_manager: Optional[KVCacheManagerV2], -) -> bool: - """Decide, before any manager is created, whether the configured - compression method supports the speculative setup. Logs the reason and - returns False when it does not; the run then stays uncompressed.""" +) -> None: + """Reject speculative setups the compression method cannot run with.""" if (spec_config is None or not config.kv_cache_compression_mode.is_eviction_method()): - return True - # Evicting methods co-compact the draft KV, so they only support spec - # modes whose draft KV is a standard paged cache in the same forward - # (one-model speculation). + return + # Evicting methods co-compact the draft KV, so the draft must be a + # standard paged cache in the same forward (one-model speculation). mode = spec_config.spec_dec_mode if not (mode.is_mtp_one_model() or mode.is_eagle3_one_model()): - logger.warning( - "KV-cache compression algorithm '%s' evicts cached tokens and " - "does not support speculative decoding mode %s (the draft KV must " - "be a standard paged cache compacted together with the target); " - "running without a compression manager.", config.algorithm, - mode.name) - return False - return True + raise ValueError( + f"KV-cache compression algorithm {config.algorithm!r} does not " + f"support speculative decoding mode {mode.name}: the draft KV " + "must be a standard paged cache compacted together with the " + "target (one-model MTP/EAGLE3).") def create_kv_cache_compression_manager( @@ -2069,7 +2063,7 @@ def create_kv_cache_compression_manager( Called from ``create_py_executor`` and registered as a resource manager, like the KV cache manager itself. Concrete algorithms add a dispatch branch here; the framework ships none. Speculative-decoding compatibility is - decided by the caller via ``kv_cache_compression_supported_with_spec``. + checked by the caller via ``validate_kv_cache_compression_with_spec``. """ logger.warning( "KV-cache compression algorithm '%s' is not registered; running without " @@ -2285,17 +2279,17 @@ def create_py_executor_instance( if kv_cache_compression_config is not None: draft_kv_cache_manager = resources.get( ResourceManagerType.DRAFT_KV_CACHE_MANAGER) - if kv_cache_compression_supported_with_spec(kv_cache_compression_config, - spec_config, - draft_kv_cache_manager): - compression_manager = create_kv_cache_compression_manager( - kv_cache_compression_config, - kv_cache_manager, - draft_kv_cache_manager=draft_kv_cache_manager, - ) - if compression_manager is not None: - resources[ResourceManagerType.KV_CACHE_COMPRESSION_MANAGER] = ( - compression_manager) + validate_kv_cache_compression_with_spec(kv_cache_compression_config, + spec_config, + draft_kv_cache_manager) + compression_manager = create_kv_cache_compression_manager( + kv_cache_compression_config, + kv_cache_manager, + draft_kv_cache_manager=draft_kv_cache_manager, + ) + if compression_manager is not None: + resources[ResourceManagerType.KV_CACHE_COMPRESSION_MANAGER] = ( + compression_manager) resource_manager = ResourceManager(resources) diff --git a/tests/unittest/_torch/executor/test_kv_cache_compression_manager.py b/tests/unittest/_torch/executor/test_kv_cache_compression_manager.py index 8a9d16974146..5c6bed5af1b6 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_compression_manager.py +++ b/tests/unittest/_torch/executor/test_kv_cache_compression_manager.py @@ -304,14 +304,15 @@ def test_eviction_method_predicate_defaults_false(self): assert not hasattr(m, "spec_config") def test_spec_gate_only_restricts_eviction_methods(self): - from tensorrt_llm._torch.pyexecutor._util import kv_cache_compression_supported_with_spec + from tensorrt_llm._torch.pyexecutor._util import validate_kv_cache_compression_with_spec from tensorrt_llm._torch.speculative.interface import SpeculativeDecodingMode from tensorrt_llm.llmapi.llm_args import KvCacheCompressionConfig + # Non-evicting methods pass with any speculative mode; no exception. config = KvCacheCompressionConfig(algorithm="offload") spec_config = SimpleNamespace(spec_dec_mode=SpeculativeDecodingMode.DFLASH) - assert kv_cache_compression_supported_with_spec(config, spec_config, None) is True - assert kv_cache_compression_supported_with_spec(config, None, None) is True + validate_kv_cache_compression_with_spec(config, spec_config, None) + validate_kv_cache_compression_with_spec(config, None, None) # ---------------------------------------------------------------------- #