[TRTLLM-10279][feat] Support different KV cache layout for one-model spec dec

Signed-off-by: ziyixiong-nv <219238287+ziyixiong-nv@users.noreply.github.com>

Author: ziyixiong-nv

tensorrt_llm/_torch/attention_backend/flashinfer.py

@@ -10,6 +10,7 @@
 from typing_extensions import Self
 
 from tensorrt_llm.functional import AttentionMaskType
+from tensorrt_llm.logger import logger
 from tensorrt_llm.models.modeling_utils import QuantConfig
 
 from ..utils import get_global_attrs, get_model_extra_attrs
@@ -61,6 +62,9 @@ class FlashInferAttentionMetadata(AttentionMetadata):
     # so set kv_layout as "HND" here
     kv_layout: Literal["NHD", "HND"] = "HND"
 
+    # Draft KV cache manager for one-model speculative decoding.
+    draft_kv_cache_manager: Optional[object] = None
+
     paged_kv_indptr_decode: torch.Tensor = field(init=False)
     paged_kv_indptr_prefill: torch.Tensor = field(init=False)
     _paged_kv_indices: torch.Tensor = field(init=False, repr=False)
@@ -127,6 +131,11 @@ def positions(self) -> torch.Tensor:
 
     def __post_init__(self) -> None:
         super().__post_init__()
+        if self.draft_kv_cache_manager is not None:
+            logger.warning(
+                "draft_kv_cache_manager is not supported in FlashInfer backend. "
+                "One-model speculative decoding with separate KV cache layouts "
+                "may not work correctly.")
         self._post_init_with_buffers(self.cuda_graph_buffers)
 
     def _post_init_with_buffers(self, buffers) -> None:

tensorrt_llm/_torch/attention_backend/interface.py

@@ -56,6 +56,9 @@ class AttentionMetadata:
     max_num_sequences: Optional[int] = None
     # The KV cache manager.
     kv_cache_manager: KVCacheManager
+    # Draft KV cache manager for one-model speculative decoding.
+    # Used when draft and target models have different KV cache layouts.
+    draft_kv_cache_manager: Optional[KVCacheManager] = None
     mapping: Optional[Mapping] = None
 
     enable_flash_mla: bool = False

tensorrt_llm/_torch/attention_backend/sparse/rocket.py

@@ -974,6 +974,7 @@ def add_dummy_requests(
         use_mrope: bool = False,
         max_beam_width: int = 1,
         num_extra_decoding_steps: int = 0,
+        draft_kv_cache_manager=None,
     ):
         requests = super().add_dummy_requests(
             request_ids=request_ids,
@@ -984,6 +985,7 @@ def add_dummy_requests(
             use_mrope=use_mrope,
             max_beam_width=max_beam_width,
             num_extra_decoding_steps=num_extra_decoding_steps,
+            draft_kv_cache_manager=draft_kv_cache_manager,
         )
         if prepare_resource:
             for req in requests:

tensorrt_llm/_torch/attention_backend/trtllm.py

@@ -10,6 +10,7 @@
     from ..speculative.utils import SpecDecodingTensor
     from ..speculative.interface import SpecMetadata
     from ..speculative.spec_tree_manager import SpecTreeManager
+    from ..pyexecutor.resource_manager import KVCacheManager
 
 from tensorrt_llm._utils import get_sm_version
 from tensorrt_llm.bindings.internal import thop
@@ -647,6 +648,9 @@ class TrtllmAttentionMetadata(AttentionMetadata):
                                                 init=True,
                                                 repr=False)
 
+    # Draft KV cache manager for one-model speculative decoding with separate KV cache layouts
+    draft_kv_cache_manager: Optional["KVCacheManager"] = None
+
     # Flags to enable spec-dec mode (multi-query mode) in TRTLLM XQA Kernels
     # spec decoding mode can be enabled for non-TRTLLM-gen kernels (pre-Blackwell XQA kernels)
     # is_spec_decoding_enabled specifies if spec-dec mode is supported for the entire runtime.
@@ -796,6 +800,29 @@ def _post_init_with_buffers(self, buffers) -> None:
             )
             self.block_ids_per_seq = None
             self.kv_block_ids_per_seq = None
+
+            # Allocate separate block offset tensors for draft KV cache manager
+            # Used in one-model speculative decoding with different KV cache layouts
+            if self.draft_kv_cache_manager is not None:
+                self.draft_kv_cache_block_offsets = self.get_empty(
+                    buffers,
+                    [
+                        self.draft_kv_cache_manager.num_pools,
+                        self.max_num_sequences, 2,
+                        self.draft_kv_cache_manager.max_blocks_per_seq
+                    ],
+                    cache_name="draft_kv_cache_block_offsets",
+                    dtype=torch.int32,
+                    capture_graph=capture_graph,
+                )
+                self.draft_host_kv_cache_block_offsets = torch.empty_like(
+                    self.draft_kv_cache_block_offsets,
+                    device='cpu',
+                    pin_memory=True,
+                )
+            else:
+                self.draft_kv_cache_block_offsets = None
+                self.draft_host_kv_cache_block_offsets = None
             if self.enable_flash_mla:
                 self.block_ids_per_seq = self.get_empty(
                     buffers,
@@ -1007,6 +1034,25 @@ def prepare(self) -> None:
             assert self.kv_lens[:self.num_seqs].max(
             ) <= self.kv_cache_manager.max_seq_len, error_message
 
+            # Also prepare draft KV cache block offsets if draft_kv_cache_manager exists
+            if self.draft_kv_cache_manager is not None:
+                # Copy blocks for all context requests
+                self.draft_kv_cache_manager.impl.copy_batch_block_offsets(
+                    self.draft_host_kv_cache_block_offsets,
+                    self.request_ids[:self.num_contexts], 1, 0)
+                # Copy blocks for all generation requests
+                self.draft_kv_cache_manager.impl.copy_batch_block_offsets(
+                    self.draft_host_kv_cache_block_offsets,
+                    self.request_ids[self.num_contexts:], self.beam_width,
+                    self.num_contexts)
+                for pool_idx in range(
+                        self.draft_host_kv_cache_block_offsets.shape[0]):
+                    self.draft_kv_cache_block_offsets[
+                        pool_idx, :self.num_seqs].copy_(
+                            self.draft_host_kv_cache_block_offsets[
+                                pool_idx, :self.num_seqs],
+                            non_blocking=True)
+
         self.kv_lens_cuda_runtime = self.kv_lens_cuda[:self.num_seqs]
         # Don't use self.kv_lens here because it includes extra tokens.
         # Use actual KV length (without extra tokens) for kv_lens_runtime,

tensorrt_llm/_torch/models/modeling_speculative.py

@@ -17,7 +17,8 @@
                               WeightsLoadingConfig)
 from ..modules.rms_norm import RMSNorm
 from ..pyexecutor.guided_decoder import CapturableGuidedDecoder
-from ..speculative import SpecMetadata, get_spec_worker
+from ..speculative import (SpecMetadata, get_spec_worker,
+                           should_use_separate_draft_kv_cache)
 from ..utils import AuxStreamType
 from .checkpoints.base_weight_mapper import BaseWeightMapper
 from .modeling_utils import (DecoderModel, DecoderModelForCausalLM, TModel,
@@ -880,6 +881,7 @@ def __init__(self, model: TModel, model_config: ModelConfig[TConfig]):
                          vocab_size=model_config.pretrained_config.vocab_size)
         self.draft_model = None
         self.draft_config = None
+        self.use_separate_draft_kv_cache = False
         spec_config = getattr(model_config, 'spec_config', None)
         if spec_config and spec_config.spec_dec_mode.use_one_engine():
             if spec_config.spec_dec_mode.is_eagle3_one_model():
@@ -913,11 +915,16 @@ def __init__(self, model: TModel, model_config: ModelConfig[TConfig]):
                 self.draft_config.quant_config.kv_cache_quant_algo = \
                 model_config.quant_config.kv_cache_quant_algo
 
+            self.use_separate_draft_kv_cache = should_use_separate_draft_kv_cache(
+                spec_config)
+
             self.draft_model = get_draft_model(model_config, self.draft_config,
                                                self.lm_head, self.model)
-            self.spec_worker = get_spec_worker(model_config.spec_config,
-                                               model_config,
-                                               model_config.mapping)
+            self.spec_worker = get_spec_worker(
+                model_config.spec_config,
+                model_config,
+                model_config.mapping,
+                use_separate_draft_kv_cache=self.use_separate_draft_kv_cache)
 
             if self.draft_config is not None and model_config.spec_config.eagle3_model_arch == "llama3":
                 for key, value in self.draft_config.extra_attrs.items():
@@ -934,6 +941,7 @@ def forward(
         inputs_embeds: Optional[torch.FloatTensor] = None,
         return_context_logits: bool = False,
         spec_metadata: Optional[SpecMetadata] = None,
+        resource_manager=None,
         **kwargs,
     ) -> torch.Tensor:
         hidden_states = self.model(
@@ -978,7 +986,8 @@ def forward(
                                     logits=logits,
                                     attn_metadata=attn_metadata,
                                     spec_metadata=spec_metadata,
-                                    draft_model=self.draft_model)
+                                    draft_model=self.draft_model,
+                                    resource_manager=resource_manager)
         else:
             logits = self.logits_processor.forward(
                 hidden_states,