[https://nvbugs/6463967][fix] DeepSeek-V4 one-model MTP separate draft kv cache (TEP)

tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py

@@ -1283,9 +1283,9 @@ def copy_batch_block_offsets(
         assert beam_width == 1, "DSV4 only supports beam width 1 now"
         assert dst_tensor.is_cuda, "copy_batch_block_offsets expects a CUDA destination"
         dst_tensor.fill_(BAD_PAGE_INDEX)
-        dst_tensor[:, : self._num_tables, 0, :].copy_(
+        dst_tensor[:, :num_seqs, 0, :].copy_(
             self._precomputed_sliding_block_tables[
-                :, DeepseekV4AttentionType.SWA.value, : self._num_tables, :
+                :, DeepseekV4AttentionType.SWA.value, :num_seqs, :
             ],
             non_blocking=True,
         )
@@ -1303,8 +1303,8 @@ def copy_batch_sliding_block_tables(
         """
         assert dst_tensor.is_cuda, "copy_batch_sliding_block_tables expects a CUDA destination"
         dst_tensor.fill_(BAD_PAGE_INDEX)
-        dst_tensor[:, :, : self._num_tables, :].copy_(
-            self._precomputed_sliding_block_tables[:, :, : self._num_tables, :],
+        dst_tensor[:, :, :num_seqs, :].copy_(
+            self._precomputed_sliding_block_tables[:, :, :num_seqs, :],
             non_blocking=True,
         )
 

tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/metadata.py

@@ -5,7 +5,7 @@
 from __future__ import annotations
 
 import math
-from typing import Dict, Optional, Set, Tuple
+from typing import TYPE_CHECKING, Dict, Optional, Set, Tuple
 
 import torch
 
@@ -23,6 +23,9 @@
     is_compress_layer,
 )
 
+if TYPE_CHECKING:
+    from .cache_manager import DeepseekV4CacheManager
+
 
 class DeepseekV4TrtllmAttentionMetadata(DSAtrtllmAttentionMetadata):
     # The set of compress ratios for the layers
@@ -282,6 +285,9 @@ def __post_init__(self):
         # so compute them once during initialization instead of every prepare().
         self._init_cache_buffer_data_pointers()
 
+        # Draft-sized sparse buffers for one-model MTP separate draft KV cache.
+        self._init_draft_sparse_buffers()
+
     def prepare_for_indexer_k_cache(self):
         """Prepare the shared indexer K-cache decode table for DSA kernels."""
         # INDEXER_COMPRESS uses shared page indices, so the generic DSA
@@ -413,31 +419,95 @@ def _prepare_deepseek_v4_indices_compiled(
                 raise ValueError(f"Unsupported compress_ratio: {compress_ratio}")
             sparse_mla_topk_lens_bufs[compress_ratio][:num_tokens] = total_count.to(torch.int32)
 
+    def _build_cache_buffer_data_pointers(
+        self, manager: "DeepseekV4CacheManager", compress_ratios_by_layer: list[int]
+    ) -> tuple[dict[int, int], dict[int, int], dict[int, int]]:
+        """Build sparse cache pointers for a target or draft manager."""
+        sparse_mla_base_ptrs = {1: manager.swa_pool_ptr}
+        for ratio, compress_pool_ptr in manager.compress_pool_ptrs.items():
+            sparse_mla_base_ptrs[ratio] = compress_pool_ptr
+
+        swa_buffer_ptrs = {layer_idx: manager.swa_pool_ptr for layer_idx in manager.pp_layers}
+        compressed_buffer_ptrs = {
+            layer_idx: manager.get_buffers(layer_idx, DeepseekV4AttentionType.COMPRESS).data_ptr()
+            for layer_idx in manager.pp_layers
+            if is_compress_layer(compress_ratios_by_layer[layer_idx])
+        }
+        return sparse_mla_base_ptrs, swa_buffer_ptrs, compressed_buffer_ptrs
+
     def _init_cache_buffer_data_pointers(self):
-        # If MTP is enabled, enlarge the compress ratios by max_draft_tokens - 1
+        # If MTP is enabled, enlarge the compress ratios by max_draft_tokens - 1.
         extend_compress_ratios = self.compress_ratios + [self.compress_ratios[-1]] * (
             self.max_draft_tokens - 1
         )
-        # SWA uses PER_LAYER indices; COMPRESS uses SHARED indices. The sparse
-        # MLA conversion kernel receives a representative base pointer per pool
-        # and a per-layer buffer pointer so it can account for any layer offset.
-        self.sparse_mla_base_ptrs = {
-            1: self.kv_cache_manager.swa_pool_ptr,
-        }
-        for ratio, compress_pool_ptr in self.kv_cache_manager.compress_pool_ptrs.items():
-            self.sparse_mla_base_ptrs[ratio] = compress_pool_ptr
+        (
+            self.sparse_mla_base_ptrs,
+            self.swa_buffer_ptrs,
+            self.compressed_buffer_ptrs,
+        ) = self._build_cache_buffer_data_pointers(self.kv_cache_manager, extend_compress_ratios)
+
+    def _init_draft_sparse_buffers(self):
+        """Initialize sparse buffers for a separate one-model MTP draft cache."""
+        self.draft_sliding_block_tables = None
+        self.draft_sparse_mla_base_ptrs = None
+        self.draft_swa_buffer_ptrs = None
+
+        draft_mgr = self.draft_kv_cache_manager
+        from .cache_manager import DeepseekV4CacheManager
+
+        if not isinstance(draft_mgr, DeepseekV4CacheManager):
+            return
+
+        # Current DSv4 MTP layers are SWA-only. Fail fast if a future model
+        # introduces a compressed or indexer MTP layer.
+        draft_ratio = self.compress_ratios[-1]
+        if draft_ratio != 1:
+            raise NotImplementedError(
+                "Separate DeepSeek-V4 draft KV cache supports only SWA-only "
+                f"(compress_ratio 1) MTP draft layers; got ratio {draft_ratio}."
+            )
 
-        self.swa_buffer_ptrs = {
-            layer_idx: self.kv_cache_manager.swa_pool_ptr
-            for layer_idx in self.kv_cache_manager.pp_layers
-        }
-        self.compressed_buffer_ptrs = {
-            layer_idx: self.kv_cache_manager.get_buffers(
-                layer_idx, DeepseekV4AttentionType.COMPRESS
-            ).data_ptr()
-            for layer_idx in self.kv_cache_manager.pp_layers
-            if is_compress_layer(extend_compress_ratios[layer_idx])
-        }
+        draft_block_table_shape = (
+            draft_mgr.num_local_layers,
+            len(DEEPSEEK_V4_SLIDING_ATTENTION),
+            self.max_num_sequences,
+            draft_mgr.max_blocks_per_seq,
+        )
+        self.draft_sliding_block_tables = self.get_empty(
+            self.cuda_graph_buffers,
+            draft_block_table_shape,
+            cache_name="draft_sliding_block_tables",
+            dtype=torch.int32,
+            capture_graph=self.is_cuda_graph,
+        )
+        extend_compress_ratios = self.compress_ratios + [draft_ratio] * (self.max_draft_tokens - 1)
+        (
+            self.draft_sparse_mla_base_ptrs,
+            self.draft_swa_buffer_ptrs,
+            _,
+        ) = self._build_cache_buffer_data_pointers(draft_mgr, extend_compress_ratios)
+
+    _DRAFT_SPARSE_FIELDS = (
+        "sliding_block_tables",
+        "sparse_mla_base_ptrs",
+        "swa_buffer_ptrs",
+    )
+
+    def prepare_for_draft_forward(self) -> dict | None:
+        """Repoint sparse fields to the draft buffers for a draft forward."""
+        if self.draft_sliding_block_tables is None:
+            return None
+        saved_state = {field: getattr(self, field) for field in self._DRAFT_SPARSE_FIELDS}
+        for field in self._DRAFT_SPARSE_FIELDS:
+            setattr(self, field, getattr(self, f"draft_{field}"))
+        return saved_state
+
+    def restore_after_draft_forward(self, saved_state: dict | None) -> None:
+        """Restore the target sparse fields after a draft forward."""
+        if saved_state is None:
+            return
+        for field in self._DRAFT_SPARSE_FIELDS:
+            setattr(self, field, saved_state[field])
 
     def prepare(self):
         assert self.kv_cache_manager is not None
@@ -448,6 +518,22 @@ def prepare(self):
             self.num_contexts,
         )
 
+        # Prepare the draft manager's tables before the generic prepare copies
+        # its block offsets, and populate the dedicated DSv4 draft sparse table.
+        draft_mgr = self.draft_kv_cache_manager
+        if draft_mgr is not None and hasattr(draft_mgr, "compute_sliding_block_tables"):
+            draft_mgr.compute_sliding_block_tables(
+                self.request_ids,
+                self.num_contexts,
+            )
+            if self.draft_sliding_block_tables is not None:
+                draft_mgr.copy_batch_sliding_block_tables(
+                    self.draft_sliding_block_tables,
+                    self.request_ids,
+                    self.num_contexts,
+                    self.num_seqs,
+                )
+
         TrtllmAttentionMetadata.prepare(self)
 
         num_requests = self.num_contexts + self.num_generations

tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py

@@ -218,6 +218,90 @@ def prepare(self):
         # Prepare metadata for indexer
         Indexer.prepare(metadata=self)
 
+    def prepare_for_draft_forward(self) -> dict | None:
+        """Select native DSA indexer metadata for a draft forward."""
+        # DeepSeek-V4 metadata inherits DSA metadata, but its cache manager uses a
+        # different dual-pool layout. Only native DSA cache managers use the DSA
+        # draft-replay buffers below.
+        if not is_dsa_cache_manager(self.kv_cache_manager):
+            return None
+
+        saved_state = {
+            "host_indexer_k_cache_block_offsets": self.host_indexer_k_cache_block_offsets,
+            "indexer_k_cache_block_offsets": self.indexer_k_cache_block_offsets,
+            "host_slot_mapping_fp8": self.host_slot_mapping_fp8,
+            "host_slot_mapping_scale": self.host_slot_mapping_scale,
+            "slot_mapping_fp8": self.slot_mapping_fp8,
+            "slot_mapping_scale": self.slot_mapping_scale,
+            "block_table": self.block_table,
+            "block_table_expanded": self.block_table_expanded,
+            "host_block_table_expanded": self.host_block_table_expanded,
+        }
+        # The cached-KV feature owns these references even when an optimized
+        # path aliases them to slot_mapping_*. With the feature disabled, the
+        # aliases are lazy and may not exist on the first generation replay.
+        if self.enable_context_mla_with_cached_kv:
+            saved_state.update(
+                {
+                    "slot_mapping_fp8_fullkv": self.slot_mapping_fp8_fullkv,
+                    "slot_mapping_scale_fullkv": self.slot_mapping_scale_fullkv,
+                }
+            )
+
+        # Rebind to the draft manager's dedicated buffers instead of
+        # overwriting the target tensors in place. Rebinding is invisible to
+        # CUDA graph capture, so the target and draft segments of the graph
+        # bake distinct addresses (like draft_kv_cache_block_offsets) and no
+        # graph-recorded copy from a transient host buffer is needed.
+        self.host_indexer_k_cache_block_offsets = self.host_draft_indexer_k_cache_block_offsets
+        self.indexer_k_cache_block_offsets = self.draft_indexer_k_cache_block_offsets
+        self.host_slot_mapping_fp8 = self.host_draft_slot_mapping_fp8
+        self.slot_mapping_fp8 = self.draft_slot_mapping_fp8
+        self.host_slot_mapping_scale = self.host_draft_slot_mapping_scale
+        self.slot_mapping_scale = self.draft_slot_mapping_scale
+        self.block_table = self.draft_block_table
+        self.block_table_expanded = self.draft_block_table_expanded
+        self.host_block_table_expanded = self.host_draft_block_table_expanded
+        self._invalidate_pool_view_cache()
+
+        # Recording a capture executes no kernels, so the draft mappings only
+        # need refreshing when the transfers actually run: eager forwards
+        # (warmup) and the pre-replay call from model_engine. The per-step
+        # advance inside the captured graph re-derives slot mappings on
+        # device from the rebound block-offset buffer.
+        # kv_cache_manager was already swapped to the draft manager above.
+        if not torch.cuda.is_current_stream_capturing():
+            self.prepare_for_indexer_k_cache()
+            self._refresh_expanded_block_table()
+            Indexer.recompute_slot_mappings(self)
+        Indexer.recompute_context_kv_gather_mappings(self)
+
+        return saved_state
+
+    def restore_after_draft_forward(self, saved_state: dict | None) -> None:
+        """Restore native DSA indexer metadata after a draft forward."""
+        if saved_state is None:
+            return
+
+        self.host_indexer_k_cache_block_offsets = saved_state["host_indexer_k_cache_block_offsets"]
+        self.indexer_k_cache_block_offsets = saved_state["indexer_k_cache_block_offsets"]
+        self.host_slot_mapping_fp8 = saved_state["host_slot_mapping_fp8"]
+        self.host_slot_mapping_scale = saved_state["host_slot_mapping_scale"]
+        self.slot_mapping_fp8 = saved_state["slot_mapping_fp8"]
+        self.slot_mapping_scale = saved_state["slot_mapping_scale"]
+        self.block_table = saved_state["block_table"]
+        self.block_table_expanded = saved_state["block_table_expanded"]
+        self.host_block_table_expanded = saved_state["host_block_table_expanded"]
+        self._invalidate_pool_view_cache()
+        if "slot_mapping_fp8_fullkv" in saved_state:
+            self.slot_mapping_fp8_fullkv = saved_state["slot_mapping_fp8_fullkv"]
+            self.slot_mapping_scale_fullkv = saved_state["slot_mapping_scale_fullkv"]
+        else:
+            # The draft recomputation rebound the aliases to the draft tensors;
+            # point them back at the restored target tensors.
+            self.slot_mapping_fp8_fullkv = self.slot_mapping_fp8
+            self.slot_mapping_scale_fullkv = self.slot_mapping_scale
+
     def get_indexer_kv_lens(self, kv_lens: torch.Tensor) -> torch.Tensor:
         if self._indexer_compress_ratio <= 1:
             return kv_lens

tensorrt_llm/_torch/attention_backend/trtllm.py

@@ -652,6 +652,14 @@ def mla_prepare_ctx_cu_seqlens(self) -> Optional[torch.Tensor]:
             self._mla_ctx_cu_seqlens_valid = True
         return self.mla_ctx_cu_q_seqlens[:num_ctx + 1]
 
+    def prepare_for_draft_forward(self) -> dict | None:
+        """Prepare backend state shared by draft-forward execution paths."""
+        return None
+
+    def restore_after_draft_forward(self, saved_state: dict | None) -> None:
+        """Restore backend state modified for draft-forward execution."""
+        return None
+
     def prepare(self) -> None:
         super().prepare()
         # Recomputed on first use this iteration; see mla_prepare_scheduler_buffers.

tensorrt_llm/_torch/pyexecutor/_util.py

@@ -1391,6 +1391,16 @@ def _should_create_separate_draft_kv_cache(self) -> bool:
                 "Attention DP is enabled, separate draft KV cache is not supported."
             )
             return False
+
+        sparse_cfg = self._sparse_attention_config
+        if (sparse_cfg is not None
+                and getattr(sparse_cfg, "algorithm", None) == "deepseek_v4"
+                and self._mapping.pp_size > 1):
+            logger.info(
+                "DeepSeek-V4 separate draft KV cache is only supported for PP=1; "
+                "folding draft layers into the unified manager for pp_size=%d.",
+                self._mapping.pp_size)
+            return False
         return should_use_separate_draft_kv_cache(self._speculative_config)
 
     def _get_effective_draft_config(self) -> ModelConfig: