[TRTLLM-7954][feat] Target model KV cache rellocation (#8421)

Signed-off-by: qgai <qgai@nvidia.com>

Author: sunnyqgg

tensorrt_llm/_torch/attention_backend/interface.py

@@ -3,12 +3,15 @@
 from collections import namedtuple
 from dataclasses import dataclass, field
 from enum import Enum, IntEnum
-from typing import (Dict, Generic, List, Optional, Protocol, Tuple, Type,
-                    TypeVar, Union)
+from typing import (TYPE_CHECKING, Dict, Generic, List, Optional, Protocol,
+                    Tuple, Type, TypeVar, Union)
 
 import torch
 from typing_extensions import Self
 
+if TYPE_CHECKING:
+    from ..speculative.utils import SpecDecodingTensor
+
 from tensorrt_llm.functional import (PositionEmbeddingType, RopeEmbeddingUtils,
                                      RotaryScalingType)
 from tensorrt_llm.mapping import Mapping
@@ -330,8 +333,13 @@ def restore_from_spec_dec(self) -> None:
             setattr(self, f, v)
         self._saved_tensors.clear()
 
-    def update_spec_dec_param(self, is_spec_decoding_enabled, is_spec_dec_tree,
-                              is_spec_dec_dynamic_tree, max_draft_tokens):
+    def update_spec_dec_param(
+            self,
+            is_spec_decoding_enabled,
+            is_spec_dec_tree,
+            is_spec_dec_dynamic_tree,
+            max_draft_tokens,
+            spec_decoding_tensor: Optional['SpecDecodingTensor'] = None):
         """
         Hook to be called when using TRTLLM attention backend in spec-dec mode.
         """

tensorrt_llm/_torch/attention_backend/sparse/rocket.py

@@ -8,6 +8,7 @@
 
 import tensorrt_llm
 import tensorrt_llm.bindings
+from tensorrt_llm._torch.attention_backend.interface import AttentionMetadata
 from tensorrt_llm._torch.attention_backend.trtllm import (
     TrtllmAttention, TrtllmAttentionMetadata)
 from tensorrt_llm._torch.attention_backend.vanilla import (
@@ -949,7 +950,10 @@ def prepare_resources(self, scheduled_batch):
             kt_token_num = math.ceil(req.max_beam_num_tokens / self.page_size)
             self.add_kt_tokens(request_id, kt_token_num)
 
-    def update_resources(self, scheduled_batch):
+    def update_resources(self,
+                         scheduled_batch,
+                         attn_metadata: AttentionMetadata = None,
+                         kv_cache_dtype_byte_size: float = None):
         for request in scheduled_batch.context_requests:
             if request.state != LlmRequestState.GENERATION_COMPLETE:
                 seq_len = request.get_num_tokens(0)

tensorrt_llm/_torch/attention_backend/trtllm.py

@@ -2,10 +2,13 @@
 import os
 import weakref
 from dataclasses import dataclass, field
-from typing import Optional, Tuple, Union
+from typing import TYPE_CHECKING, Optional, Tuple, Union
 
 import torch
 
+if TYPE_CHECKING:
+    from ..speculative.utils import SpecDecodingTensor
+
 from tensorrt_llm._utils import get_sm_version
 from tensorrt_llm.bindings.internal import thop
 from tensorrt_llm.functional import AttentionMaskType
@@ -1045,12 +1048,32 @@ def prepare_context_mla_with_cached_kv(self,
         self.ctx_kv_indptr[:self.num_contexts + 1].copy_(
             self.host_ctx_kv_indptr[:self.num_contexts + 1], non_blocking=True)
 
-    def update_spec_dec_param(self, is_spec_decoding_enabled, is_spec_dec_tree,
-                              is_spec_dec_dynamic_tree, max_draft_tokens):
+    def update_spec_dec_param(
+        self,
+        is_spec_decoding_enabled,
+        is_spec_dec_tree,
+        is_spec_dec_dynamic_tree,
+        max_draft_tokens,
+        spec_decoding_tensor: Optional['SpecDecodingTensor'] = None,
+    ):
+
+        if spec_decoding_tensor is not None:
+            spec_decoding_position_offsets = spec_decoding_tensor.position_offsets
+            spec_decoding_packed_mask = spec_decoding_tensor.packed_mask
+            spec_decoding_generation_lengths = spec_decoding_tensor.generation_lengths
+        else:
+            spec_decoding_position_offsets = None
+            spec_decoding_packed_mask = None
+            spec_decoding_generation_lengths = None
         # spec_dec mode should only be enabled for pre-Blackwell machines and when there's a spec-dec tree.
         self.is_spec_decoding_enabled = is_spec_decoding_enabled and get_sm_version(
         ) < 100
 
+        if get_sm_version() >= 100:
+            if is_spec_dec_tree or is_spec_dec_dynamic_tree:
+                assert not is_spec_dec_tree, "Spec-dec tree is not supported on this machine. Please use a pre-Blackwell machine for a spec-dec tree."
+                assert not is_spec_dec_dynamic_tree, "Spec-dec dynamic tree is not supported on this machine. Please use a pre-Blackwell machine for a spec-dec dynamic tree."
+
         # use_spec_decoding is default to true by default, change in runtime by layers / requests
         self.use_spec_decoding = self.is_spec_decoding_enabled
 
@@ -1068,7 +1091,7 @@ def update_spec_dec_param(self, is_spec_decoding_enabled, is_spec_dec_tree,
             self.spec_decoding_packed_mask = torch.empty(
                 [
                     self.max_num_requests, max_draft_tokens + 1,
-                    math.ceil(max_draft_tokens / 32)
+                    math.ceil((max_draft_tokens + 1) / 32)
                 ],
                 dtype=torch.int,
                 device='cuda',
@@ -1081,7 +1104,18 @@ def update_spec_dec_param(self, is_spec_decoding_enabled, is_spec_dec_tree,
             )
 
             if self.is_spec_dec_dynamic_tree:
-                assert False, "currently dynamic tree is not supported"
+                assert spec_decoding_position_offsets is not None, "spec_decoding_position_offsets is required for dynamic tree"
+                assert spec_decoding_packed_mask is not None, "spec_decoding_packed_mask is required for dynamic tree"
+                self.spec_decoding_position_offsets.copy_(
+                    spec_decoding_position_offsets, non_blocking=True)
+                self.spec_decoding_packed_mask.copy_(spec_decoding_packed_mask,
+                                                     non_blocking=True)
+                if spec_decoding_generation_lengths is not None:
+                    self.spec_decoding_generation_lengths.copy_(
+                        spec_decoding_generation_lengths, non_blocking=True)
+                else:
+                    self.generate_spec_decoding_generation_length(
+                        max_draft_tokens=max_draft_tokens)
             else:
                 # Populate the mask that won't change during inference phase.
                 self.generate_spec_decoding_position_offsets(
@@ -1092,7 +1126,6 @@ def update_spec_dec_param(self, is_spec_decoding_enabled, is_spec_dec_tree,
                     max_draft_tokens=max_draft_tokens)
 
     def generate_spec_decoding_position_offsets(self, max_draft_tokens):
-        assert not self.is_spec_dec_tree, "only chained/linear tree is supported now"
         position_offset = torch.arange(max_draft_tokens + 1,
                                        dtype=torch.int,
                                        device='cpu',
@@ -1103,7 +1136,6 @@ def generate_spec_decoding_position_offsets(self, max_draft_tokens):
                                                   non_blocking=True)
 
     def generate_spec_decoding_packed_mask(self, max_draft_tokens):
-        assert not self.is_spec_dec_tree, "only chained/linear tree is supported now"
         dummy_idx = torch.arange(max_draft_tokens + 1)
         spec_decoding_packed_mask = torch.pow(2, dummy_idx + 1) - 1
         self.spec_decoding_packed_mask[:, :, 0].copy_(spec_decoding_packed_mask,

tensorrt_llm/_torch/pyexecutor/llm_request.py

@@ -482,6 +482,8 @@ def __init__(
         self.py_target_probs = None
         self.py_last_draft_tokens = None
         self.py_num_accepted_draft_tokens = 0
+        self.py_num_accepted_draft_tokens_indices = []
+        self.py_rewind_draft_token_separate_adjustment = 0
         self.py_decoding_iter = 0
         self.is_attention_dp_dummy = False
         self.is_cuda_graph_dummy = False

tensorrt_llm/_torch/pyexecutor/model_engine.py

@@ -49,6 +49,7 @@
 from ..speculative.drafting_loops import ChainDrafter
 from ..speculative.eagle3 import Eagle3ResourceManager
 from ..speculative.mtp import SampleStateTensorsMTP
+from ..speculative.utils import SpecDecodingTensor
 from ..utils import (get_model_extra_attrs,
                      set_per_request_piecewise_cuda_graph_flag,
                      set_torch_compiling, with_model_extra_attrs)
@@ -370,9 +371,24 @@ def __init__(
         else:
             self.cache_indirection_attention = None
 
+        self.kv_cache_dtype_byte_size = self.get_kv_cache_dtype_byte_size()
+
     def register_forward_pass_callable(self, callable: Callable):
         self.forward_pass_callable = callable
 
+    def get_kv_cache_dtype_byte_size(self) -> float:
+        """
+        Returns the size (in bytes) occupied by kv cache type.
+        """
+        layer_quant_mode = self.model.model_config.quant_config.layer_quant_mode
+        if layer_quant_mode.has_fp4_kv_cache():
+            return 1 / 2
+        elif layer_quant_mode.has_fp8_kv_cache(
+        ) or layer_quant_mode.has_int8_kv_cache():
+            return 1
+        else:
+            return 2
+
     @property
     def runtime_draft_len(self):
         return self.max_total_draft_tokens if self.enable_spec_decode else 0
@@ -2249,6 +2265,7 @@ def forward(
         new_tensors_device: Optional[SampleStateTensors] = None,
         gather_context_logits: bool = False,
         cache_indirection_buffer: Optional[torch.Tensor] = None,
+        spec_decoding_tensor: Optional[SpecDecodingTensor] = None,
     ):
         kv_cache_manager = resource_manager.get_resource_manager(
             self.kv_cache_manager_key)
@@ -2267,7 +2284,7 @@ def forward(
             attn_metadata.update_spec_dec_param(
                 is_spec_dec_mode, spec_metadata.is_spec_dec_tree,
                 spec_metadata.is_spec_dec_dynamic_tree,
-                self.original_max_draft_len)
+                self.original_max_draft_len, spec_decoding_tensor)
         else:
             spec_resource_manager = None
             spec_metadata = None

tensorrt_llm/_torch/pyexecutor/py_executor.py

@@ -989,8 +989,13 @@ def _executor_loop_pp(self):
 
                         finished_requests = self._handle_responses()
                         previous_scheduled_batch = previous_batch.sample_state.scheduled_requests
+                        attn_metadata = getattr(self.model_engine,
+                                                'attn_metadata', None)
+                        kv_cache_dtype_byte_size = getattr(
+                            self.model_engine, 'kv_cache_dtype_byte_size', None)
                         self.resource_manager.update_resources(
-                            previous_scheduled_batch)
+                            previous_scheduled_batch, attn_metadata,
+                            kv_cache_dtype_byte_size)
                         self._remove_inflight_ids(previous_scheduled_batch)
 
                     self.wait_on_pp_send_handles(prev_microbatch_id)
@@ -1200,7 +1205,13 @@ def _executor_loop(self):
 
                     self._handle_canceled_requests()
                     finished_requests = self._handle_responses()
-                    self.resource_manager.update_resources(scheduled_batch)
+                    attn_metadata = getattr(self.model_engine, 'attn_metadata',
+                                            None)
+                    kv_cache_dtype_byte_size = getattr(
+                        self.model_engine, 'kv_cache_dtype_byte_size', None)
+                    self.resource_manager.update_resources(
+                        scheduled_batch, attn_metadata,
+                        kv_cache_dtype_byte_size)
                     if self.enable_kv_cache_events:
                         self._add_kv_cache_events()
 
@@ -1403,7 +1414,12 @@ def _process_previous_batch(self):
         self._handle_canceled_requests()
         finished_requests = self._handle_responses()
         scheduled_requests = self.previous_batch.sample_state.scheduled_requests
-        self.resource_manager.update_resources(scheduled_requests)
+        attn_metadata = getattr(self.model_engine, 'attn_metadata', None)
+        kv_cache_dtype_byte_size = getattr(self.model_engine,
+                                           'kv_cache_dtype_byte_size', None)
+        self.resource_manager.update_resources(scheduled_requests,
+                                               attn_metadata,
+                                               kv_cache_dtype_byte_size)
         if self.enable_kv_cache_events:
             self._add_kv_cache_events()