[MTP] follow custom deepseek modeling changes to support graph mode

Signed-off-by: mengwei805 <mengwei25@huawei.com>

Author: mengwei805

vllm_ascend/models/deepseek_mtp.py

@@ -17,11 +17,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Optional
+from typing import List, Optional
 
 import torch
 import torch.nn as nn
 from transformers import PretrainedConfig
+from vllm.attention.backends.abstract import AttentionMetadata
 from vllm.config import CacheConfig, ModelConfig, VllmConfig
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
@@ -34,6 +35,7 @@
     SharedHead)
 from vllm.model_executor.models.utils import maybe_prefix
 from vllm.model_executor.sampling_metadata import SamplingMetadata
+from vllm.sequence import IntermediateTensors
 
 from .deepseek_v2 import CustomDeepseekV2DecoderLayer
 
@@ -69,6 +71,8 @@ def forward(
         self,
         input_ids: torch.Tensor,
         positions: torch.Tensor,
+        kv_cache: torch.Tensor,
+        attn_metadata: AttentionMetadata,
         previous_hidden_states: torch.Tensor,
         inputs_embeds: Optional[torch.Tensor] = None,
         spec_step_index: int = 0,
@@ -88,6 +92,8 @@ def forward(
 
         hidden_states, residual = self.mtp_block(positions=positions,
                                                  hidden_states=hidden_states,
+                                                 kv_cache=kv_cache,
+                                                 attn_metadata=attn_metadata,
                                                  residual=None)
         hidden_states = residual + hidden_states
         return hidden_states
@@ -125,14 +131,20 @@ def forward(
         self,
         input_ids: torch.Tensor,
         positions: torch.Tensor,
+        kv_caches: torch.Tensor,
+        attn_metadata: AttentionMetadata,
         previous_hidden_states: torch.Tensor,
         inputs_embeds: Optional[torch.Tensor] = None,
         spec_step_idx: int = 0,
     ) -> torch.Tensor:
         current_step_idx = (spec_step_idx % self.num_mtp_layers)
+        step_kv_cache = kv_caches[
+            current_step_idx] if kv_caches is not None else None
         return self.layers_list[current_step_idx](
             input_ids,
             positions,
+            step_kv_cache,
+            attn_metadata,
             previous_hidden_states,
             inputs_embeds,
             current_step_idx,
@@ -170,3 +182,19 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
                                                            prefix, "model"))
 
         self.sampler = get_sampler()
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        kv_caches: Optional[List[torch.Tensor]] = None,
+        attn_metadata: Optional[AttentionMetadata] = None,
+        previous_hidden_states: Optional[torch.Tensor] = None,
+        intermediate_tensors: Optional[IntermediateTensors] = None,
+        inputs_embeds: Optional[torch.Tensor] = None,
+        spec_step_idx: int = 0,
+    ) -> torch.Tensor:
+        hidden_states = self.model(input_ids, positions, kv_caches,
+                                   attn_metadata, previous_hidden_states,
+                                   inputs_embeds, spec_step_idx)
+        return hidden_states

vllm_ascend/patch/worker/patch_common/patch_multi_step_worker.py

@@ -61,15 +61,19 @@ def sampler_output(
     else:
         # Here we run multi-step directly, with every step prepared
         # on the CPU.
-        # TODO: Remove this branch once DraftModelRunner supports TP>1
+        # TODO Remove this branch once DraftModelRunner supports TP>1
         # and other restrictions that are part of DraftModelRunner's
         # supports_gpu_multi_step(..)
+        if expanded_request.previous_hidden_states is not None:
+            self.worker.model_runner.return_hidden_states = True
         for _ in range(sample_len):
             model_output: List[SamplerOutput] = self.worker.execute_model(
                 execute_model_req=expanded_request)
             assert (len(model_output) == 1
                     ), "composing multistep workers not supported"
             model_output = model_output[0]
+            self._maybe_update_previous_hidden_states(model_output,
+                                                      expanded_request)
 
             self._append_new_tokens(model_output,
                                     expanded_request.seq_group_metadata_list,

vllm_ascend/patch/worker/patch_common/patch_spec_decode_worker.py

@@ -93,7 +93,7 @@ def create_worker(
 
             proposer_worker = MultiStepWorker(**draft_worker_kwargs)
             if draft_model_config.hf_config.model_type == "deepseek_mtp":
-                num_spec_prefill_steps = num_speculative_tokens
+                num_spec_prefill_steps = draft_model_config.hf_config.n_predict
 
         proposer_worker = SmallerTpProposerWorker.maybe_wrap_worker(
             proposer_worker, draft_tp, target_tp)