Merge pull request vllm-project#2 from luccafong/mtp_config_enablement

fix mtp config and padding

Author: heheda12345

vllm/config/speculative.py

@@ -142,7 +142,7 @@ def compute_hash(self) -> str:
 
     @staticmethod
     def hf_config_override(hf_config: PretrainedConfig) -> PretrainedConfig:
-        if hf_config.model_type == "deepseek_v3":
+        if hf_config.model_type in ("deepseek_v3", "deepseek_v32"):
             hf_config.model_type = "deepseek_mtp"
         if hf_config.model_type == "deepseek_mtp":
             n_predict = getattr(hf_config, "num_nextn_predict_layers", None)
@@ -204,9 +204,8 @@ def __post_init__(self):
             # mtp acceleration for more models besides deepseek_v3
             if self.target_model_config and \
                 (self.target_model_config.hf_text_config.model_type \
-                        == "deepseek_v3" or
-                    self.target_model_config.hf_text_config.model_type in
-                        ("mimo","ernie4_5_moe", "qwen3_next")):
+                    in ("deepseek_v3", "deepseek_v32",
+                        "mimo","ernie4_5_moe", "qwen3_next")):
                 # use the draft model from the same model:
                 self.model = self.target_model_config.model
                 # Align the quantization of draft model for cases such as

vllm/model_executor/models/deepseek_mtp.py

@@ -59,7 +59,7 @@ def __init__(self, vllm_config: VllmConfig, prefix: str) -> None:
             config, "index_topk"
         )
         if self.is_v32:
-            topk_tokens = config.attn_module_list_cfg[0]["topk_tokens"]
+            topk_tokens = config.index_topk
             topk_indices_buffer = torch.empty(vllm_config.scheduler_config.max_num_batched_tokens,
                                               topk_tokens,
                                               dtype=torch.int32,

vllm/model_executor/models/deepseek_v2.py

@@ -690,18 +690,18 @@ def sparse_attn_indexer(
             padded_q_fp8_decode_tokens = q_fp8[:num_decode_tokens].reshape(decode_lens.shape[0], -1, *q_fp8.shape[1:])
         # TODO: move and optimize below logic with triton kernels
         batch_size = padded_q_fp8_decode_tokens.shape[0]
+        next_n = padded_q_fp8_decode_tokens.shape[1]
         assert batch_size == decode_metadata.seq_lens.shape[0]
+        num_padded_tokens = batch_size * next_n
         logits = fp8_paged_mqa_logits(
             padded_q_fp8_decode_tokens,
             kv_cache,
-            weights[:num_decode_tokens],
+            weights[:num_padded_tokens],
             decode_metadata.seq_lens,
             decode_metadata.block_table,
             decode_metadata.schedule_metadata,
             max_model_len=max_model_len,
         )
-        # [B, N, L]
-        next_n = padded_q_fp8_decode_tokens.shape[1]
         # padded query len
         current_device = padded_q_fp8_decode_tokens.device
         padded_num_tokens = batch_size * next_n
@@ -721,7 +721,7 @@ def sparse_attn_indexer(
         topk_indices[topk_indices > index_end_pos] = -1
         if decode_metadata.requires_padding:
             # if padded, we need to unpack the topk indices removing padded tokens
-            topk_indices = unpack_seq_triton(topk_indices.reshape(batch_size, -1, logits.shape[-1]), decode_lens)
+            topk_indices = unpack_seq_triton(topk_indices.reshape(batch_size, -1, topk_indices.shape[-1]), decode_lens)
         topk_indices_buffer[:num_decode_tokens, :topk_indices.
                             shape[-1]] = topk_indices.to(
                                 dtype=torch.int32)

vllm/v1/attention/backends/mla/flashmla_sparse.py

@@ -273,8 +273,6 @@ class FlashMLASparseMetadataBuilder(
     cudagraph_support: ClassVar[AttentionCGSupport] = \
         AttentionCGSupport.UNIFORM_BATCH
 
-    reorder_batch_threshold: ClassVar[int] = 128  # TODO(lucas): tune this
-
     reorder_batch_threshold: ClassVar[int] = 1
 
     def __init__(self, kv_cache_spec: AttentionSpec, layer_names: list[str],
@@ -309,7 +307,8 @@ def __init__(self, kv_cache_spec: AttentionSpec, layer_names: list[str],
             vllm_config.speculative_config.num_speculative_tokens 
             if vllm_config.speculative_config else 0
         )
-        self.reorder_batch_threshold += self.num_speculative_tokens
+        # Now deepgemm fp8_paged_mqa_logits does not support next_n > 2
+        self.reorder_batch_threshold += min(self.num_speculative_tokens, 1) 
 
         # Equation taken from FlashMLA/csrc/pybind.cpp
         h_q, h_k = self.num_heads, 1

vllm/v1/attention/backends/mla/indexer.py

@@ -175,7 +175,8 @@ def __init__(self, *args, **kwargs):
             self.vllm_config.speculative_config.num_speculative_tokens 
             if self.vllm_config.speculative_config else 0
         )
-        self.reorder_batch_threshold += self.num_speculative_tokens
+        # Now deepgemm fp8_paged_mqa_logits does not support next_n > 2
+        self.reorder_batch_threshold += min(self.num_speculative_tokens, 1) 
 
         props = torch.cuda.get_device_properties(self.device)
         sm_count = props.multi_processor_count