[OPEN] [MTP V1] MTP adapt torchair graph mode (#1294)

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

1. Add MTP dummy_run, and adapt main model dummy_run when mtp is enabled
2. Adapt main model torchair graph mode, when mtp is enabled
3. mtp model torchair graph mode will be supported in the future

### Does this PR introduce _any_ user-facing change?


### How was this patch tested?

This patch is tested by
`vllm-ascend/tests/long_term/spec_decode_v1/test_v1_mtp_correctness.py`


### Usage
online example
```shell
export VLLM_USE_V1=1
export VLLM_ENABLE_MC2=1
export VLLM_VERSION=0.9.1
export ASCEND_LAUNCH_BLOCKING=0

python -m vllm.entrypoints.openai.api_server \
 --model="/model_weight_path" \
 --trust-remote-code \
 --max-model-len 40000 \
 --tensor-parallel-size 4 \
 --data_parallel_size 4 \
 --enable_expert_parallel \
 --served-model-name deepseekr1 \
 --quantization ascend \
 --host 0.0.0.0 \
 --port 1234 \
 --speculative-config '{"num_speculative_tokens": 1, "method":"deepseek_mtp"}' \
 --additional-config '{"ascend_scheduler_config":{"enabled":true,"enable_chunked_prefill":false},"torchair_graph_config":{"enabled":true}}' \
 --gpu_memory_utilization 0.95
```

offline example
```
lm = LLM(
    model="/model_weight_path",
    tensor_parallel_size=16,
    max_num_seqs=128,
    gpu_memory_utilization=0.95,
    distributed_executor_backend="mp",
    enable_expert_parallel=True,
    speculative_config={
        "method": "deepseek_mtp",
        "num_speculative_tokens": 1,
    },
    trust_remote_code=True,
    enforce_eager=False,
    additional_config = {
       'torchair_graph_config': {
            'enabled': True,
            'enable_multistream_shared_expert': False
        },
       "ascend_scheduler_config": {
            "enabled": True
        },
    }
)
```

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

Author: JC-ut0

tests/long_term/spec_decode_v1/test_v1_mtp_correctness.py

@@ -94,3 +94,62 @@ def test_mtp_correctness(
         # Upon failure, inspect the outputs to check for inaccuracy.
         assert matches > int(0.66 * len(ref_outputs))
         del spec_llm
+
+
+def test_mtp_torchair_correctness(
+    monkeypatch: pytest.MonkeyPatch,
+    test_prompts: list[list[dict[str, Any]]],
+    sampling_config: SamplingParams,
+    model_name: str,
+):
+    '''
+    Compare the outputs of a original LLM and a speculative LLM
+    should be the same when using mtp speculative decoding.
+    '''
+    with monkeypatch.context() as m:
+        m.setenv("VLLM_USE_V1", "1")
+
+        ref_llm = LLM(model=model_name,
+                      max_model_len=256,
+                      enforce_eager=False,
+                      additional_config={
+                          "torchair_graph_config": {
+                              "enabled": True
+                          },
+                          "ascend_scheduler_config": {
+                              "enabled": True
+                          },
+                      })
+        ref_outputs = ref_llm.chat(test_prompts, sampling_config)
+        del ref_llm
+
+        spec_llm = LLM(model=model_name,
+                       trust_remote_code=True,
+                       enforce_eager=False,
+                       speculative_config={
+                           "method": "deepseek_mtp",
+                           "num_speculative_tokens": 1,
+                       },
+                       additional_config={
+                           "torchair_graph_config": {
+                               "enabled": True
+                           },
+                           "ascend_scheduler_config": {
+                               "enabled": True
+                           },
+                       })
+        spec_outputs = spec_llm.chat(test_prompts, sampling_config)
+        matches = 0
+        misses = 0
+        for ref_output, spec_output in zip(ref_outputs, spec_outputs):
+            if ref_output.outputs[0].text == spec_output.outputs[0].text:
+                matches += 1
+            else:
+                misses += 1
+                print(f"ref_output: {ref_output.outputs[0].text}")
+                print(f"spec_output: {spec_output.outputs[0].text}")
+
+        # Heuristic: expect at least 66% of the prompts to match exactly
+        # Upon failure, inspect the outputs to check for inaccuracy.
+        assert matches > int(0.66 * len(ref_outputs))
+        del spec_llm

vllm_ascend/attention/mla_v1.py

@@ -104,6 +104,7 @@ class AscendMLADecodeMetadata:
     seq_lens: torch.Tensor
     max_seq_lens: int
     seq_lens_list: list[int]
+    actual_seq_q_lens: Optional[list[int]] = None
     attn_mask: Optional[torch.Tensor] = None
 
 
@@ -138,6 +139,7 @@ class AscendMLAMetadata:
     num_input_tokens: int = 0  # Number of tokens including padding.
 
     enable_dbo_across_dp: bool = False
+    is_mtp_model: bool = False
 
     query_lens: Optional[list[int]] = None
     # The dimension of the attention heads
@@ -313,48 +315,64 @@ def _get_graph_runner_block_tables(
         return graph_block_tables[:num_seqs, :max_blocks]
 
     def build_torchair_graph_dummy(
-            self, num_reqs: int, num_actual_tokens: int) -> AscendMLAMetadata:
+        self,
+        num_reqs: int,
+        num_actual_tokens: int,
+        is_mtp_model: bool = False,
+    ) -> AscendMLAMetadata:
         device = self.runner.device
         _, max_blocks = self.runner.graph_block_tables.shape
         block_table = torch.zeros((num_reqs, max_blocks),
                                   dtype=torch.int32,
                                   device=device)
         block_table = self._get_graph_runner_block_tables(
             num_reqs, block_table)
-        seq_lens = torch.ones(num_reqs, dtype=torch.int32, device=device)
-        input_positions = torch.zeros(num_reqs,
+        num_tokens = num_reqs * self.runner.decode_token_per_req
+        seq_lens = torch.zeros(num_reqs, dtype=torch.int32, device=device)
+        seq_lens_list = seq_lens.tolist()
+        input_positions = torch.zeros(num_tokens,
                                       dtype=torch.int32,
                                       device=device).long()
-        slot_mapping = torch.full((num_reqs, ),
+        slot_mapping = torch.full((num_tokens, ),
                                   PAD_SLOT_ID,
                                   dtype=torch.int32,
                                   device=device)
         query_start_loc = torch.full((num_reqs, ),
                                      -1,
                                      dtype=torch.int32,
                                      device=device)
+        if self.runner.speculative_config is not None and\
+            self.runner.speculative_config.method == 'deepseek_mtp' and not is_mtp_model:
+            attn_state = AscendAttentionState.SpecDecoding
+            num_decode_tokens = 2
+        else:
+            attn_state = AscendAttentionState.DecodeOnly
+            num_decode_tokens = 1
         decode_metadata = AscendMLADecodeMetadata(
             input_positions=input_positions,
             block_table=block_table,
             seq_lens=seq_lens,
-            seq_lens_list=seq_lens.tolist(),
+            seq_lens_list=seq_lens_list,
             max_seq_lens=1,
-            attn_mask=self.runner.spec_attn_mask)
+            attn_mask=self.runner.spec_attn_mask,
+            actual_seq_q_lens=self.runner.actual_seq_q_lens[:num_reqs],
+        )
         return self.metadata_cls(  # type: ignore
             num_input_tokens=num_actual_tokens,
             num_actual_tokens=num_actual_tokens,
             slot_mapping=slot_mapping,
             head_dim=self.runner.model_config.get_head_size(),
             num_decodes=1,
-            num_decode_tokens=1,
+            num_decode_tokens=num_decode_tokens,
             num_prefills=0,
             attn_mask=self.runner.attn_mask,
-            attn_state=AscendAttentionState.DecodeOnly,
+            attn_state=attn_state,
             prefill=None,
             decode=decode_metadata,
             query_start_loc=query_start_loc,
             seq_lens=seq_lens,
             block_tables=block_table,
+            is_mtp_model=is_mtp_model,
         )
 
     def build(
@@ -364,8 +382,10 @@ def build(
         max_query_len: int,
         common_attn_metadata: CommonAttentionMetadata,
         common_prefix_len: Optional[int] = None,
-        graph_pad_size: int = -1,
+        num_token_pad_size: int = -1,
+        num_reqs_pad_size: int = 0,
         enable_dbo_across_dp: bool = False,
+        is_mtp_model: bool = False,
     ) -> AscendMLAMetadata:
         assert self._num_decodes + self._num_prefills == num_reqs
 
@@ -449,8 +469,9 @@ def build(
             )
 
         decode_metadata = None
-        use_torchair_graph = graph_pad_size != -1
+        use_torchair_graph = num_token_pad_size != -1
         if self._num_decodes > 0:
+            actual_seq_q_lens = None
             max_seq_lens = seq_lens[:self._num_decodes].max().item()
             seq_lens = seq_lens[:self._num_decode_tokens]
             input_positions = input_positions[:self._num_decode_tokens]
@@ -459,41 +480,48 @@ def build(
                     AscendAttentionState.DecodeOnly,
                     AscendAttentionState.SpecDecoding
             ]:
-                num_seqs = len(seq_lens)
-                if graph_pad_size != 0:
-                    pad_value = 1
-                    padded_seq_lens = seq_lens.tolist() + [pad_value
-                                                           ] * graph_pad_size
+                if num_token_pad_size != 0:
+                    pad_value = 0
+                    padded_seq_lens = seq_lens.tolist(
+                    ) + [pad_value] * num_reqs_pad_size
                 else:
                     padded_seq_lens = seq_lens.tolist()
 
                 seq_lens = torch.from_numpy(
                     np.array(padded_seq_lens).astype(np.int32))
-                padding = torch.full((graph_pad_size, ),
+                seq_lens_list = padded_seq_lens
+                padding = torch.full((num_token_pad_size, ),
                                      PAD_SLOT_ID,
                                      dtype=slot_mapping.dtype,
                                      device=slot_mapping.device)
                 slot_mapping = torch.cat([slot_mapping, padding])
                 block_table_padding = torch.zeros(
-                    (graph_pad_size, ) + block_table.shape[1:],
+                    (num_reqs_pad_size, ) + block_table.shape[1:],
                     dtype=block_table.dtype,
                     device=block_table.device)
                 block_table = torch.cat([block_table, block_table_padding],
                                         dim=0)
                 block_table = self._get_graph_runner_block_tables(
-                    num_seqs + graph_pad_size, block_table)
-                padding_0 = torch.zeros(graph_pad_size,
+                    num_reqs + num_reqs_pad_size, block_table)
+                padding_0 = torch.zeros(num_token_pad_size,
                                         dtype=input_positions.dtype,
                                         device=input_positions.device)
                 input_positions = torch.cat([input_positions, padding_0])
+                actual_seq_q_lens = query_start_loc[1:].tolist(
+                ) + self.runner.actual_seq_q_lens[num_reqs:num_reqs +
+                                                  num_reqs_pad_size]
+            else:
+                seq_lens_list = seq_lens.tolist()
 
             decode_metadata = AscendMLADecodeMetadata(
                 input_positions=input_positions,
                 block_table=block_table,
                 seq_lens=seq_lens,
-                seq_lens_list=seq_lens.tolist(),
+                seq_lens_list=seq_lens_list,
                 max_seq_lens=max_seq_lens,
-                attn_mask=self.runner.spec_attn_mask)
+                attn_mask=self.runner.spec_attn_mask,
+                actual_seq_q_lens=actual_seq_q_lens,
+            )
 
         return self.metadata_cls(  # type: ignore
             num_actual_tokens=num_actual_tokens,
@@ -510,7 +538,9 @@ def build(
             query_start_loc=query_start_loc,
             block_tables=block_table,
             seq_lens=seq_lens,
-            enable_dbo_across_dp=enable_dbo_across_dp)
+            enable_dbo_across_dp=enable_dbo_across_dp,
+            is_mtp_model=is_mtp_model,
+        )
 
 
 class AscendMLAImpl(MLAAttentionImpl):
@@ -933,31 +963,10 @@ def _forward_decode(
         assert decode_meta is not None
         num_tokens = q_nope.size(0)
         if self.running_in_graph:
-            # TorchAir's shape is [bs, num_heads_per_rank, q_seq_len, dim]
-            if attn_metadata.attn_state == AscendAttentionState.SpecDecoding:
-                assert num_tokens % self.spec_token_num == 0
-                q_nope = q_nope.view(num_tokens // (self.spec_token_num + 1),
-                                     self.spec_token_num + 1, self.num_heads,
-                                     -1)
-                q_pe = q_pe.view(num_tokens // (self.spec_token_num + 1),
-                                 self.spec_token_num + 1, self.num_heads, -1)
-                if not self.enable_kv_nz:
-                    q_nope = q_nope.transpose(1, 2).contiguous()
-                    q_pe = q_pe.transpose(1, 2).contiguous()
-                sparse_mode = 3
-                spec_attn_mask = attn_metadata.decode.attn_mask  # type:ignore
-            else:
-                if self.enable_kv_nz:
-                    q_nope = q_nope.view(num_tokens, 1, self.num_heads, -1)
-                    q_pe = q_pe.view(num_tokens, 1, self.num_heads, -1)
-                else:
-                    q_nope = q_nope.view(num_tokens, self.num_heads, 1, -1)
-                    q_pe = q_pe.view(num_tokens, self.num_heads, 1, -1)
-                sparse_mode = 0
-                spec_attn_mask = None
             # shape of knope/k_pe for npu graph mode should be:
             # [num_blocks, num_kv_heads, block_size, self.kv_lora_rank/self.qk_rope_head_dim]
             block_size = kv_c_and_k_pe_cache[0].shape[1]
+            actual_seq_lengths = None
             if self.enable_kv_nz:
                 k_nope = k_nope.view(-1, self.num_kv_heads,
                                      self.kv_lora_rank // 16, block_size, 16)
@@ -971,6 +980,26 @@ def _forward_decode(
                                  self.qk_rope_head_dim)
                 input_layout = "BNSD"
 
+            # TorchAir's shape is [bs, num_heads_per_rank, q_seq_len, dim]
+            if attn_metadata.attn_state == AscendAttentionState.SpecDecoding:
+                assert num_tokens % self.spec_token_num == 0
+                # [bs * q_seq_len, num_heads_per_rank, dim]
+                input_layout = "TND"
+                q_nope = q_nope.view(num_tokens, self.num_heads, -1)
+                q_pe = q_pe.view(num_tokens, self.num_heads, -1)
+                sparse_mode = 3
+                spec_attn_mask = attn_metadata.decode.attn_mask  # type:ignore
+                actual_seq_lengths = decode_meta.actual_seq_q_lens
+            else:
+                if self.enable_kv_nz:
+                    q_nope = q_nope.view(num_tokens, 1, self.num_heads, -1)
+                    q_pe = q_pe.view(num_tokens, 1, self.num_heads, -1)
+                else:
+                    q_nope = q_nope.view(num_tokens, self.num_heads, 1, -1)
+                    q_pe = q_pe.view(num_tokens, self.num_heads, 1, -1)
+                sparse_mode = 0
+                spec_attn_mask = None
+
             attn_output, _ = torch_npu.npu_fused_infer_attention_score(
                 q_nope,
                 k_nope,
@@ -988,7 +1017,7 @@ def _forward_decode(
                 block_table=decode_meta.block_table,
                 block_size=block_size,
                 actual_seq_lengths_kv=decode_meta.seq_lens_list,
-            )
+                actual_seq_lengths=actual_seq_lengths)
         else:
             # The MLA_PA path will be used as default path in the future, `_npu_paged_attention_mla` will
             # be removed after the torch_npu contains `torch_npu.atb.npu_multi_head_latent_attention` become
@@ -1042,6 +1071,8 @@ def forward(
         if attn_metadata is None:
             # Profiling run.
             return output
+        # mtp model is not support for graph mode yet
+        self.torchair_graph_enabled = self.torchair_graph_enabled and not attn_metadata.is_mtp_model
         self.running_in_graph = self.torchair_graph_enabled and attn_metadata.attn_state in [
             AscendAttentionState.DecodeOnly, AscendAttentionState.SpecDecoding
         ]

vllm_ascend/models/deepseek_v2.py

@@ -482,7 +482,8 @@ def forward(
                 hidden_states_or_q_c = self.q_a_layernorm(ckq)
         else:
             hidden_states_or_q_c = hidden_states
-        if self.torchair_graph_enabled:
+        is_mtp_model = attn_metadata is not None and attn_metadata.is_mtp_model
+        if self.torchair_graph_enabled and not is_mtp_model:
             forward_kwargs = {}
             if envs.VLLM_USE_V1:
                 output_shape = hidden_states.shape