[BUGFIX] [v0.9.1] Fix mtp with disaggregated-prefill (#1694)

### What this PR does / why we need it?

[BUGFIX] [v0.9.1] Fix mtp with disaggregated-prefill
### Does this PR introduce _any_ user-facing change?


### How was this patch tested?

Signed-off-by: xuyexiong <xuyexiong@huawei.com>

Author: JC-ut0

vllm_ascend/attention/mla_v1.py

@@ -538,6 +538,11 @@ def build(
                 actual_seq_q_lens = query_start_loc[1:].tolist(
                 ) + self.runner.actual_seq_q_lens[num_reqs:num_reqs +
                                                   num_reqs_pad_size]
+                # mtp torchair + PD scenario, last element of actual_seq_q_lens must equal to num_reqs_pad_size
+                num_padded_token_size = slot_mapping.size(0)
+                if actual_seq_q_lens[-1] != num_padded_token_size:
+                    actual_seq_q_lens.append(num_padded_token_size)
+                    seq_lens_list.append(0)
             else:
                 seq_lens_list = seq_lens.tolist()
 

vllm_ascend/worker/model_runner_v1.py

@@ -1691,6 +1691,9 @@ def _dummy_run(
                         torch._dynamo.mark_static(
                             get_forward_context().mc2_mask)
                         torch._dynamo.mark_static(attn_metadata.slot_mapping)
+                        if self.speculative_config:
+                            torch._dynamo.mark_static(
+                                attn_metadata.decode.attn_mask)
                         for kv in self.kv_caches:
                             assert isinstance(
                                 kv, tuple), "kv_cache must be a tuple"
@@ -1720,7 +1723,7 @@ def _dummy_run(
                         **model_kwargs)
             if self.speculative_config and self.speculative_config.method == "deepseek_mtp":
                 assert isinstance(self.drafter, MtpProposer)
-                self.drafter.dummy_run(num_reqs)
+                self.drafter.dummy_run(num_reqs, with_prefill=with_prefill)
             return hidden_states
 
     @contextmanager

vllm_ascend/worker/mtp_proposer_v1.py

@@ -211,7 +211,8 @@ def propose(
 
         with set_ascend_forward_context(attn_metadata,
                                         self.vllm_config,
-                                        num_tokens=num_input_tokens):
+                                        num_tokens=num_input_tokens,
+                                        with_prefill=self.runner.with_prefill):
             with ProfileExecuteDuration().capture_async('mtp_forward'):
                 model_kwargs = {}
                 model_kwargs["attn_metadata"] = attn_metadata
@@ -305,15 +306,13 @@ def load_model(self) -> None:
                     ge_cache=False)
 
     @torch.inference_mode()
-    def dummy_run(
-        self,
-        num_tokens: int,
-    ) -> None:
+    def dummy_run(self, num_tokens: int, with_prefill: bool = False) -> None:
         attn_metadata = self.runner.attn_metadata_builder.build_torchair_graph_dummy(
             num_reqs=num_tokens, num_actual_tokens=1, is_mtp_model=True)
         with set_ascend_forward_context(None,
                                         self.vllm_config,
-                                        num_tokens=num_tokens):
+                                        num_tokens=num_tokens,
+                                        with_prefill=with_prefill):
             self.model(input_ids=self.input_ids[:num_tokens],
                        positions=self.positions[:num_tokens],
                        previous_hidden_states=self.hidden_states[:num_tokens],