add mtp3 support.

Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>

Author: lfr-0531

tensorrt_llm/_torch/attention_backend/sparse/dsa.py

@@ -431,6 +431,18 @@ def __post_init__(self):
             dtype=torch.int32,
             capture_graph=capture_graph,
         )
+        self.kv_lens_expanded_cuda = self.get_empty(
+            self.cuda_graph_buffers,
+            (self.max_num_sequences * (1 + self.max_draft_tokens), ),
+            cache_name="kv_lens_expanded_cuda",
+            dtype=torch.int32,
+            capture_graph=capture_graph,
+        )
+        self.kv_lens_expanded_host = torch.zeros_like(
+            self.kv_lens_expanded_cuda,
+            device='cpu',
+            pin_memory=True,
+        )
 
     def prepare(self):
         super().prepare()
@@ -534,6 +546,10 @@ def prepare(self):
         else:
             self.max_gen_seq_len = 0
 
+        # Expand kv_lens_cuda for draft tokens (only generation)
+        gen_kv_lens = kv_lens[self.num_contexts:self.num_seqs]
+        self.kv_lens_expanded_host = torch.cat([gen_kv_lens] * (1+self.max_draft_tokens), dim=0)
+
         # Prepare metadata for indexer
         Indexer.prepare(metadata=self)
 
@@ -1047,9 +1063,15 @@ def sparse_attn_indexer(
             # Reshape q for decode phase: [num_gen_tokens, ...] -> [batch_size, next_n, ...]
             q_decode = q_fp8[num_ctx_tokens:num_ctx_tokens + num_gen_tokens,
                              ...]
-            q_decode = q_decode.view(num_generations, -1, *q_fp8.shape[1:])
-            batch_size = q_decode.shape[0]
-            next_n = q_decode.shape[1]
+            batch_size = num_generations
+            next_n = num_gen_tokens // num_generations
+            if next_n <= 2:
+                q_decode = q_decode.view(num_generations, -1, *q_fp8.shape[1:])
+                context_lens = metadata.kv_lens_cuda_runtime[num_contexts:num_contexts +
+                                                             num_generations]
+            else:
+                q_decode = q_decode.view(-1, 1, *q_fp8.shape[1:])
+
             assert num_gen_tokens == batch_size * next_n
             weights_decode = weights[num_ctx_tokens:num_ctx_tokens +
                                      num_gen_tokens, ...]

tensorrt_llm/_torch/attention_backend/trtllm.py

@@ -597,6 +597,8 @@ class TrtllmAttentionMetadata(AttentionMetadata):
     is_spec_decoding_enabled: bool = False
     # use_spec_decoding determines if the attention layer should be run in spec-dec mode at the specific step / layer.
     use_spec_decoding: bool = False
+    # max number of draft tokens
+    max_draft_tokens: int = 0
 
     # if spec-dec tree is a tree or a chain (linear tree)
     is_spec_dec_tree: bool = False
@@ -1067,7 +1069,7 @@ def update_spec_dec_param(
         max_draft_tokens,
         spec_decoding_tensor: Optional['SpecDecodingTensor'] = None,
     ):
-
+        self.max_draft_tokens = max_draft_tokens
         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