[0.9.1][MTP V1]MTP model adapt torchair

Signed-off-by: xuyexiong <[email protected]>

fix mtp torchair

fix

Author: JC-ut0

vllm_ascend/attention/mla_v1.py

@@ -94,7 +94,7 @@ class AscendMLADecodeMetadata:
     seq_lens: torch.Tensor
     max_seq_lens: int
     seq_lens_list: list[int]
-    actual_seq_q_lens: Optional[list[int]] = None
+    actual_seq_lengths_q: Optional[list[int]] = None
     attn_mask: Optional[torch.Tensor] = None
     sin: torch.Tensor = None
     cos: torch.Tensor = None
@@ -131,7 +131,6 @@ 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,7 +312,6 @@ def build_torchair_graph_dummy(
         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
@@ -337,7 +335,7 @@ def build_torchair_graph_dummy(
                                      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:
+            self.runner.speculative_config.method == 'deepseek_mtp':
             attn_state = AscendAttentionState.SpecDecoding
             num_decode_tokens = 2
         else:
@@ -362,7 +360,7 @@ def build_torchair_graph_dummy(
             seq_lens_list=seq_lens_list,
             max_seq_lens=1,
             attn_mask=self.runner.spec_attn_mask,
-            actual_seq_q_lens=self.runner.actual_seq_q_lens[:num_reqs],
+            actual_seq_lengths_q=self.runner.actual_seq_lengths_q[:num_reqs],
             sin=sin,
             cos=cos)
         return self.metadata_cls(  # type: ignore
@@ -380,7 +378,6 @@ def build_torchair_graph_dummy(
             query_start_loc=query_start_loc,
             seq_lens=seq_lens,
             block_tables=block_table,
-            is_mtp_model=is_mtp_model,
         )
 
     def build(
@@ -393,7 +390,6 @@ def build(
         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
 
@@ -498,7 +494,7 @@ def build(
         decode_metadata = None
         use_torchair_graph = num_token_pad_size != -1
         if self._num_decodes > 0:
-            actual_seq_q_lens = query_start_loc[1:].tolist()
+            actual_seq_lengths_q = query_start_loc[1:].tolist()
             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]
@@ -534,16 +530,21 @@ def build(
                                         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 +
+                actual_seq_lengths_q = query_start_loc[1:].tolist(
+                ) + self.runner.actual_seq_lengths_q[num_reqs:num_reqs +
                                                   num_reqs_pad_size]
-                # mtp torchair + PD scenario, last element of actual_seq_q_lens must equal to num_padded_token_size
+                # mtp torchair + PD scenario, last element of actual_seq_lengths_q must equal to num_padded_token_size
                 num_padded_token_size = slot_mapping.size(0)
-                if actual_seq_q_lens[-1] != num_padded_token_size \
+                if actual_seq_lengths_q[-1] != num_padded_token_size \
                     and self.runner.attn_state == AscendAttentionState.SpecDecoding:
-                    actual_seq_q_lens[-1] = num_padded_token_size
+                    actual_seq_lengths_q[-1] = num_padded_token_size
             else:
                 seq_lens_list = seq_lens.tolist()
+            # mtp torchair + PD scenario, last element of actual_seq_lengths_q must equal to num_padded_token_size
+            num_padded_token_size = slot_mapping.size(0)
+            if actual_seq_lengths_q[-1] != num_padded_token_size \
+                and self.runner.attn_state == AscendAttentionState.SpecDecoding:
+                actual_seq_lengths_q[-1] = num_padded_token_size
 
             cos = self.cos_cache[input_positions].unsqueeze(  # type: ignore
                 1).unsqueeze(2)
@@ -557,7 +558,7 @@ def build(
                 seq_lens_list=seq_lens_list,
                 max_seq_lens=max_seq_lens,
                 attn_mask=self.runner.spec_attn_mask,
-                actual_seq_q_lens=actual_seq_q_lens,
+                actual_seq_lengths_q=actual_seq_lengths_q,
                 sin=sin,
                 cos=cos)
 
@@ -577,7 +578,6 @@ def build(
             block_tables=block_table,
             seq_lens=seq_lens,
             enable_dbo_across_dp=enable_dbo_across_dp,
-            is_mtp_model=is_mtp_model,
         )
 
 
@@ -1017,16 +1017,13 @@ def _forward_decode(
 
             if attn_metadata.attn_state == AscendAttentionState.SpecDecoding:
                 assert num_tokens % self.spec_token_num == 0
-                if self.enable_kv_nz:
-                    input_layout = "TND_NTD"
-                else:
-                    input_layout = "TND"
+                input_layout = "TND"
                 # [bs * q_seq_len, num_heads_per_rank, dim]
                 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
+                actual_seq_lengths = decode_meta.actual_seq_lengths_q
             else:
                 if self.enable_kv_nz:
                     q_nope = q_nope.view(num_tokens, 1, self.num_heads, -1)
@@ -1110,8 +1107,6 @@ 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_mtp.py

@@ -24,6 +24,7 @@
 from transformers import PretrainedConfig
 from vllm.attention.backends.abstract import AttentionMetadata
 from vllm.config import CacheConfig, ModelConfig, VllmConfig
+from vllm.forward_context import get_forward_context
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
 from vllm.model_executor.layers.quantization import QuantizationConfig
@@ -98,9 +99,11 @@ def forward(
             inputs_embeds = self.embed_tokens(input_ids)
         assert inputs_embeds is not None
         # masking inputs at position 0, as not needed by MTP
-        inputs_embeds = torch.where((positions == 0).unsqueeze(-1),
-                                    torch.zeros_like(inputs_embeds),
-                                    inputs_embeds)
+        forward_context = get_forward_context()
+        if forward_context.with_prefill:
+            inputs_embeds = torch.where((positions == 0).unsqueeze(-1),
+                                        torch.zeros_like(inputs_embeds),
+                                        inputs_embeds)
         inputs_embeds = self.enorm(inputs_embeds)
         previous_hidden_states = self.hnorm(previous_hidden_states)
 

vllm_ascend/models/deepseek_v2.py

@@ -485,8 +485,7 @@ def forward(
             hidden_states_or_q_c = self.q_a_layernorm(ckq)
         else:
             hidden_states_or_q_c = hidden_states
-        is_mtp_model = attn_metadata is not None and attn_metadata.is_mtp_model
-        if self.torchair_graph_enabled and not is_mtp_model:
+        if self.torchair_graph_enabled:
             if envs.VLLM_USE_V1:
                 output_shape = hidden_states.shape
                 output = torch.empty(output_shape,

vllm_ascend/ops/vocab_parallel_embedding.py

@@ -30,15 +30,22 @@ def get_masked_input_and_mask(
         added_vocab_end_index: int) -> Tuple[torch.Tensor, torch.Tensor]:
     # torch.compile will fuse all of the pointwise ops below
     # into a single kernel, making it very fast
-    org_vocab_mask = (input_ >= org_vocab_start_index) & (input_ <
-                                                          org_vocab_end_index)
-    added_vocab_mask = (input_ >= added_vocab_start_index) & (
-        input_ < added_vocab_end_index)
-    added_offset = added_vocab_start_index - (
-        org_vocab_end_index - org_vocab_start_index) - num_org_vocab_padding
-    valid_offset = (org_vocab_start_index *
-                    org_vocab_mask) + (added_offset * added_vocab_mask)
-    vocab_mask = org_vocab_mask | added_vocab_mask
+    org_vocab_mask = (input_ >= org_vocab_start_index) & (
+        input_ < org_vocab_end_index)
+    # Adapt: avoid create added_vocab_mask when added_vocab_start_index == added_vocab_end_index.
+    if added_vocab_start_index == added_vocab_end_index:
+        valid_offset = (org_vocab_start_index * org_vocab_mask)
+        vocab_mask = org_vocab_mask
+    else:
+        added_vocab_mask = (input_ >= added_vocab_start_index) & (
+            input_ < added_vocab_end_index)
+        added_offset = added_vocab_start_index - (
+            org_vocab_end_index -
+            org_vocab_start_index) - num_org_vocab_padding
+        valid_offset = (org_vocab_start_index *
+                        org_vocab_mask) + (added_offset * added_vocab_mask)
+        vocab_mask = org_vocab_mask | added_vocab_mask
+    # Adapt end.
     input_ = vocab_mask * (input_ - valid_offset)
     return input_, ~vocab_mask
 

vllm_ascend/worker/model_runner_v1.py

@@ -212,7 +212,7 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device):
         # Set up speculative decoding.
         self.use_spec_decode = False
         self.spec_attn_mask = None
-        self.actual_seq_q_lens = []
+        self.actual_seq_lengths_q = []
         self.spec_token_num = 0
         self.decode_token_per_req = 1
         if self.speculative_config:
@@ -232,7 +232,7 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device):
                 elif self.speculative_config.method == 'deepseek_mtp':
                     self.drafter = MtpProposer(self.vllm_config, self)
                     self.decode_token_per_req = 1 + self.spec_token_num
-                    self.actual_seq_q_lens = [
+                    self.actual_seq_lengths_q = [
                         len for len in
                         range(self.decode_token_per_req, self.max_num_tokens +
                               1, self.decode_token_per_req)
@@ -1009,6 +1009,7 @@ def _process_reqs(
         common_attn_metadata = CommonAttentionMetadata(
             query_start_loc=query_start_loc,
             seq_lens=self.seq_lens_cpu[:num_reqs])
+        self.common_attn_metadata = common_attn_metadata
         self.seq_lens_list = self.seq_lens_np.tolist()[:num_input_tokens]
         with_prefill = attn_state not in [
             AscendAttentionState.DecodeOnly, AscendAttentionState.SpecDecoding
@@ -1040,7 +1041,9 @@ def _process_reqs(
             extra_builder_kwargs['num_token_pad_size'] = num_token_pad_size
             extra_builder_kwargs['num_reqs_pad_size'] = num_reqs_pad_size
             self.num_reqs_pad_size = num_reqs_pad_size
+            self.num_token_pad_size = num_token_pad_size
         self.extra_builder_kwargs = extra_builder_kwargs
+        self.num_tokens_across_dp = num_tokens_across_dp
 
         if self.vllm_config.model_config.use_mla:
             attn_metadata = self.attn_metadata_builder.build(  # type: ignore
@@ -1740,7 +1743,12 @@ 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, with_prefill=with_prefill)
+                self.drafter.dummy_run(
+                    num_tokens=num_tokens,
+                    with_prefill=with_prefill,
+                    skip_attn=skip_attn,
+                    num_reqs=num_reqs,
+                    num_tokens_across_dp=num_tokens_across_dp)
             return hidden_states
 
     @contextmanager
@@ -2123,7 +2131,7 @@ def _generate_mtp_token_ids(
         next_token_ids = torch.tensor(next_token_ids,
                                       dtype=torch.int32,
                                       device=self.device)
-
+        token_indices = None
         if spec_decode_metadata is None:
             # input_ids can be None for multimodal models.
             target_token_ids = self.input_ids[:num_scheduled_tokens]
@@ -2146,12 +2154,20 @@ def _generate_mtp_token_ids(
             cu_num_tokens, token_indices = self.drafter.prepare_inputs(
                 attn_metadata.query_start_loc,
                 num_rejected_tokens,
-                force_one_token=True,
-            )
-            target_token_ids = self.input_ids[token_indices]
-            target_positions = positions[token_indices]
-            target_hidden_states = hidden_states[token_indices]
-            target_slot_mapping = attn_metadata.slot_mapping[token_indices]
+                force_one_token=False,
+                is_torchair_graph=self.torchair_graph_enabled)
+            if self.torchair_graph_enabled:
+                # the seq len of each bath is padded to 2, thus input is same as the main model
+                target_token_ids = self.input_ids[:num_scheduled_tokens]
+                target_positions = positions[:num_scheduled_tokens]
+                target_hidden_states = hidden_states[:num_scheduled_tokens]
+                target_slot_mapping = attn_metadata.slot_mapping[:
+                                                                 num_scheduled_tokens]
+            else:
+                target_token_ids = self.input_ids[token_indices]
+                target_positions = positions[token_indices]
+                target_hidden_states = hidden_states[token_indices]
+                target_slot_mapping = attn_metadata.slot_mapping[token_indices]
 
         draft_token_ids = self.drafter.propose(
             target_token_ids=target_token_ids,
@@ -2162,7 +2178,7 @@ def _generate_mtp_token_ids(
             cu_num_tokens=cu_num_tokens,
             block_table=attn_metadata.block_tables,
             sampling_metadata=sampling_metadata,
-        )
+            token_indices=token_indices)
         spec_token_ids = draft_token_ids.tolist()
         return spec_token_ids