adapt mtp with graph mode in v1

Signed-off-by: whx-sjtu <2952154980@qq.com>

Author: whx-sjtu

vllm_ascend/attention/mla_v1.py

@@ -85,6 +85,7 @@ class AscendMLADecodeMetadata:
     seq_lens: torch.Tensor
     max_seq_lens: int
     seq_lens_list: list[int]
+    attn_mask: torch.Tensor
 
 
 @dataclass
@@ -170,11 +171,12 @@ def reorder_batch(self, input_batch: "InputBatch",
 
         for i, req_id in enumerate(input_batch.req_ids):
             num_tokens = scheduler_output.num_scheduled_tokens[req_id]
+            num_spec_tokens = len(scheduler_output.scheduled_spec_decode_tokens.get(req_id, []))
             # for now treat 1 scheduled token as "decode" even if its not,
             # we should update this to something like < 8 in the future but
             # currently the TritonMLA._forward_decode only supports
             # num_tokens = 1
-            if num_tokens == 1:
+            if num_tokens - num_spec_tokens == 1:
                 decodes.append(i)
                 num_decode_tokens += num_tokens
             else:
@@ -335,7 +337,8 @@ def build(self,
                 block_table=block_table,
                 seq_lens=seq_lens,
                 seq_lens_list=seq_lens.tolist(),
-                max_seq_lens=max_seq_lens)
+                max_seq_lens=max_seq_lens,
+                attn_mask=self.runner.spec_attn_mask)
 
         return self.metadata_cls(  # type: ignore
             num_actual_tokens=num_actual_tokens,
@@ -424,6 +427,17 @@ def __init__(
 
         self.enable_graph_mode = False
         additional_config = get_current_vllm_config().additional_config
+        speculative_config = get_current_vllm_config().speculative_config
+        self.fia_sparse_mode = 0
+        self.use_spec_decode = False
+        # We need to set the sparse_mode of fused_infer_attention op to 3
+        # in spec decoding scenario in order to pass in attention mask.
+        if speculative_config is not None:
+            self.fia_sparse_mode = 3
+            self.use_spec_decode = True
+            self.spec_token_num = speculative_config.num_speculative_tokens
+            assert self.spec_token_num > 0
+
         if additional_config:
             self.enable_graph_mode = additional_config.get(
                 "enable_graph_mode", False)
@@ -628,9 +642,32 @@ def _forward_decode(
             dtype=q.dtype,
             device=q.device)
         if self.running_in_graph:
-            # TorchAir's shape is [bs, num_heads_per_rank, seq_len, dim]
-            q_nope = q_nope.view(num_tokens, self.num_heads, 1, -1)
-            q_pe = q_pe.view(num_tokens, self.num_heads, 1, -1)
+            # TorchAir's shape is [bs, num_heads_per_rank, q_seq_len, dim]
+            if self.use_spec_decode:
+                assert num_tokens % self.spec_token_num == 0
+                q_nope = (
+                    q_nope.view(
+                        num_tokens // self.spec_token_num,
+                        self.spec_token_num,
+                        self.num_heads,
+                        -1,
+                    )
+                    .transpose(1, 2)
+                    .contiguous()
+                )
+                q_pe = (
+                    q_pe.view(
+                        num_tokens // self.spec_token_num,
+                        self.spec_token_num,
+                        self.num_heads,
+                        -1,
+                    )
+                    .transpose(1, 2)
+                    .contiguous()
+                )
+            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)
             # 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]
@@ -648,7 +685,8 @@ def _forward_decode(
                 num_heads=self.num_heads,
                 num_key_value_heads=self.num_kv_heads,
                 input_layout="BNSD",
-                atten_mask=attn_metadata.attn_mask,
+                atten_mask=attn_metadata.decode.attn_mask, # type:ignore
+                sparse_mode=self.fia_sparse_mode,
                 scale=self.scale,
                 antiquant_mode=0,
                 antiquant_scale=None,

vllm_ascend/worker/model_runner_v1.py

@@ -195,8 +195,12 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device):
 
         # Set up speculative decoding.
         self.use_spec_decode = False
+        self.spec_attn_mask = None
         if self.speculative_config:
             self.use_spec_decode = True
+            # TODO: Need to find out the right value of spec_attn_mask to make sure
+            # that accuracy is right.
+            self.spec_attn_mask = torch.zeros(2048, 2048, dtype=torch.bool).to("npu")
             if get_pp_group().is_last_rank:
                 if self.speculative_config.method == "ngram":
                     self.drafter = NgramProposer(self.vllm_config)
@@ -534,10 +538,13 @@ def _process_reqs(
         # Get the number of scheduled tokens for each request.
         # TODO: The Python loop can be slow. Optimize.
         num_scheduled_tokens = np.empty(num_reqs, dtype=np.int32)
+        num_valid_tokens = np.empty(num_reqs, dtype=np.int32)
         max_num_scheduled_tokens = 0
         for i, req_id in enumerate(self.input_batch.req_ids):
             num_tokens = scheduler_output.num_scheduled_tokens[req_id]
             num_scheduled_tokens[i] = num_tokens
+            num_valid_tokens[i] = num_tokens - \
+                len(scheduler_output.scheduled_spec_decode_tokens.get(req_id, []))
             max_num_scheduled_tokens = max(max_num_scheduled_tokens,
                                            num_tokens)
 
@@ -584,7 +591,7 @@ def _process_reqs(
         if np.array_equal(self.seq_lens_np[:num_reqs], num_scheduled_tokens):
             attn_state = AscendAttentionState.PrefillNoCache
         # We assume it is the decode stage, where prefill occurs but only one token is not hit in cache.
-        elif np.all(num_scheduled_tokens == 1):
+        elif np.all(num_valid_tokens == 1):
             attn_state = AscendAttentionState.DecodeOnly
         # splitfuse
         elif not self.use_v0_scheduler or self.chunked_prefill_enabled:
@@ -618,7 +625,7 @@ def _process_reqs(
             query_start_loc=query_start_loc, seq_lens=seq_lens)
         # Add graph_pad_size here
         if self.enable_torchair_graph_mode:
-            graph_pad_size = self.scheduler_config.max_num_seqs - len(seq_lens)
+            graph_pad_size = self.scheduler_config.max_num_seqs - sum(num_scheduled_tokens)
             extra_builder_kwargs['graph_pad_size'] = graph_pad_size
 
         if self.vllm_config.model_config.use_mla:

vllm_ascend/worker/mtp_proposer_v1.py

@@ -4,7 +4,7 @@
                          set_current_vllm_config)
 from vllm.forward_context import set_forward_context
 from vllm.model_executor.model_loader import get_model_loader
-from vllm.model_executor.model_loader.utils import set_default_torch_dtype
+from vllm.model_executor.model_loader.utils import set_default_torch_dtype, process_weights_after_loading
 from vllm.v1.sample.metadata import SamplingMetadata
 
 from vllm_ascend.attention.mla_v1 import CommonAttentionMetadata
@@ -199,6 +199,7 @@ def load_model(self) -> None:
             loader.get_all_weights(
                 self.vllm_config.speculative_config.draft_model_config,
                 self.model))
+        process_weights_after_loading(self.model, draft_model_config, target_device)
 
 
 # TODO Using torch instead of triton may result in poor performance