[None][feat] Add TensorRT-LLM runtime integration for KV cache compression

tensorrt_llm/_torch/kv_cache_compression/interface.py

@@ -0,0 +1,29 @@
+# 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.
+
+    Configs map their ``algorithm`` string to a member here; callers read the
+    ``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

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)
@@ -148,6 +149,7 @@ class ModelConfig(Generic[TConfig]):
     lm_head_gather_output: bool = True
     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

tensorrt_llm/_torch/pyexecutor/_util.py

@@ -2067,16 +2067,38 @@ def _create_kv_cache_manager(
     return kv_cache_manager
 
 
+def validate_kv_cache_compression_with_spec(
+    config: KvCacheCompressionConfig,
+    spec_config: Optional[SpeculativeConfig],
+    draft_kv_cache_manager: Optional[KVCacheManagerV2],
+) -> 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
+    # 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()):
+        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(
     config: KvCacheCompressionConfig,
     kv_cache_manager: KVCacheManagerV2,
+    draft_kv_cache_manager: Optional[KVCacheManagerV2] = 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
+    checked by the caller via ``validate_kv_cache_compression_with_spec``.
     """
     logger.warning(
         "KV-cache compression algorithm '%s' is not registered; running without "
@@ -2330,8 +2352,16 @@ 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:
+        draft_kv_cache_manager = resources.get(
+            ResourceManagerType.DRAFT_KV_CACHE_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)
+            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)
@@ -2343,17 +2373,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.

tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py

@@ -657,6 +657,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
         )
@@ -3179,12 +3181,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(

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

tensorrt_llm/_torch/pyexecutor/model_engine.py

@@ -3169,6 +3169,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
@@ -3179,9 +3183,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
 
@@ -3427,8 +3433,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
@@ -3619,8 +3626,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:
@@ -3647,9 +3655,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)
@@ -3706,7 +3716,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
@@ -3783,7 +3794,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)
 

tensorrt_llm/_torch/pyexecutor/model_loader.py

@@ -1272,6 +1272,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,