Support mtp run in full graph mode

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

Author: anon189Ty

vllm_ascend/ascend_forward_context.py

@@ -71,7 +71,8 @@ def set_ascend_forward_context(
         batch_descriptor: Optional[BatchDescriptor] = None,
         prefetch_stream: torch.npu.Stream = None,
         model_instance: torch.nn.Module = None,
-        weight_prefetch_method: Optional[WeightPrefetchMethod] = None):
+        weight_prefetch_method: Optional[WeightPrefetchMethod] = None,
+        is_mtp_model=False):
     """A context manager that stores the current forward context,
     can be attention metadata, etc.
     We add some additional param into forward_context.
@@ -153,6 +154,7 @@ def set_ascend_forward_context(
         forward_context.prefetch_mlp_enabled = prefetch_mlp_enabled
         forward_context.model_instance = model_instance
         forward_context.weight_prefetch_method = weight_prefetch_method
+        forward_context.is_mtp_model = is_mtp_model
 
         # TODO(rjg-lyh): The current implementation is somewhat brute force and not elegant.
         # It will be improved later by implementing operator fusion through the FX graph.

vllm_ascend/attention/mla_v1.py

@@ -26,6 +26,7 @@
                                          trans_rope_weight, transdata,
                                          wait_for_kv_layer_from_connector)
 from vllm_ascend.compilation.acl_graph import (get_graph_params,
+                                               get_mtp_graph_params,
                                                update_graph_params_workspaces)
 from vllm_ascend.multistream.base import MSAttentionMetadataSplitConfig
 from vllm_ascend.multistream.context import get_multistream_comm_context
@@ -1028,8 +1029,11 @@ def _forward_decode(
             "actual_seq_lengths": actual_seq_lengths,
             "actual_seq_lengths_kv": decode_meta.seq_lens_list,
         }
-        graph_params = get_graph_params()
         forward_context: ForwardContext = get_forward_context()
+        if forward_context.is_mtp_model:
+            graph_params = get_mtp_graph_params()
+        else:
+            graph_params = get_graph_params()
         if forward_context.capturing:
             stream = torch_npu.npu.current_stream()
 

vllm_ascend/compilation/acl_graph.py

@@ -249,7 +249,10 @@ def update_attn_params(update_stream, forward_context, runtime_shape):
 
 def update_mla_attn_params(update_stream, forward_context, runtime_shape,
                            speculative_config):
-    graph_params = get_graph_params()
+    if forward_context.is_mtp_model:
+        graph_params = get_mtp_graph_params()
+    else:
+        graph_params = get_graph_params()
     # FIXME: Behold! We are using a temporary hack here to update the args
     # for each layer's attention op in the graph.
     with torch.npu.stream(update_stream):
@@ -265,7 +268,8 @@ def update_mla_attn_params(update_stream, forward_context, runtime_shape,
              softmax_lse) = param
             seq_lens_list = forward_context.attn_metadata[
                 key].decode.seq_lens_list
-            if speculative_config and speculative_config.method == "deepseek_mtp":
+            if speculative_config and speculative_config.method == "deepseek_mtp" \
+                and not forward_context.is_mtp_model:
                 actual_seq_lengths = forward_context.attn_metadata[
                     key].decode.actual_seq_lengths_q
                 spec_multiple = speculative_config.num_speculative_tokens + 1
@@ -340,3 +344,40 @@ def update_graph_params_workspaces(num_tokens: int, workspace: Any):
 
 def get_graph_params():
     return _graph_params
+
+
+@dataclass
+class MTPGraphParams:
+    events: dict[int, list[torch.npu.ExternalEvent]]
+    workspaces: dict[int, torch.Tensor]
+    handles: dict[int, list[torch_npu._C._NPUTaskGroupHandle]]
+    attn_params: dict[int, list[tuple]]
+
+
+_mtp_graph_params: Optional[MTPGraphParams] = None
+
+
+def set_mtp_graph_params(aclgraph_capture_sizes: set[int]):
+    global _mtp_graph_params
+    if _mtp_graph_params is not None:
+        raise ValueError("MTPGraph parameters have already been set!")
+    _mtp_graph_params = MTPGraphParams(
+        {size: []
+         for size in aclgraph_capture_sizes},
+        {size: None
+         for size in aclgraph_capture_sizes},
+        {size: []
+         for size in aclgraph_capture_sizes},
+        {size: []
+         for size in aclgraph_capture_sizes},
+    )
+
+
+def update_mtp_graph_params_workspaces(num_tokens: int, workspace: Any):
+    global _mtp_graph_params
+    if _mtp_graph_params is not None:
+        _mtp_graph_params.workspaces[num_tokens] = workspace
+
+
+def get_mtp_graph_params():
+    return _mtp_graph_params

vllm_ascend/spec_decode/mtp_proposer.py

@@ -17,9 +17,13 @@
 
 from vllm_ascend.ascend_config import get_ascend_config
 from vllm_ascend.ascend_forward_context import set_ascend_forward_context
+from vllm_ascend.attention.attention_v1 import AscendAttentionState
 from vllm_ascend.attention.utils import AscendCommonAttentionMetadata
 from vllm_ascend.patch.worker.patch_deepseek_mtp import \
     AscendDeepSeekMTP as DeepSeekMTP
+from vllm_ascend.compilation.acl_graph import (ACLGraphWrapper,
+                                               set_mtp_graph_params,
+                                               update_mla_attn_params)
 from vllm_ascend.spec_decode.interface import Proposer, SpecDcodeType
 from vllm_ascend.torchair.models.torchair_deepseek_mtp import \
     TorchairDeepSeekMTP
@@ -71,6 +75,23 @@ def __init__(
         self.use_sparse = hasattr(vllm_config.model_config.hf_config,
                                   "index_topk")
 
+        self.actual_seq_lengths_q = list(
+            range(1, self.runner.max_num_tokens + 1, 1))
+        self.query_start_loc = torch.zeros(self.runner.max_num_reqs + 1,
+                                           dtype=torch.int32,
+                                           device=self.device)
+        self.query_start_loc_cpu = torch.zeros(self.runner.max_num_reqs + 1,
+                                               dtype=torch.int32,
+                                               device="cpu",
+                                               pin_memory=True)
+        self.slot_mapping = torch.zeros(self.runner.max_num_tokens,
+                                        dtype=torch.int32,
+                                        device=self.device)
+        self.seq_lens_cpu = torch.zeros(self.runner.max_num_reqs,
+                                        dtype=torch.int32,
+                                        device="cpu",
+                                        pin_memory=True)
+
     def load_model(self, model) -> None:
         loader = get_model_loader(self.vllm_config.load_config)
 
@@ -106,6 +127,15 @@ def load_model(self, model) -> None:
         process_weights_after_loading(self.model, draft_model_config,
                                       target_device)
 
+        if self.vllm_config.compilation_config.cudagraph_mode.has_full_cudagraphs(
+        ):
+            self.update_stream: torch.npu.Stream = torch.npu.Stream()
+            set_mtp_graph_params(
+                self.vllm_config.compilation_config.cudagraph_capture_sizes)
+            self.model = ACLGraphWrapper(self.model,
+                                         self.vllm_config,
+                                         runtime_mode=CUDAGraphMode.FULL)
+
     @torch.inference_mode()
     def dummy_run(self,
                   num_tokens: int,
@@ -131,7 +161,7 @@ def dummy_run(self,
             skip_attn = False
         if skip_attn:
             attn_metadata = None
-        else:
+        elif is_running_torchair:
             common_attn_metadata = TorchairCommonAttentionMetadata(
                 num_reqs=num_reqs,
                 num_actual_tokens=1,
@@ -142,6 +172,56 @@ def dummy_run(self,
             )
             attn_metadata = self.runner.attn_metadata_builder.build_torchair_graph_dummy(
                 common_attn_metadata)
+        elif aclgraph_runtime_mode == CUDAGraphMode.FULL:
+            assert with_prefill is False, \
+                "Full decode graph only supports uniform batch now."
+            num_reqs = num_tokens
+            max_seq_lens = self.runner.model_config.max_model_len
+            self.seq_lens_cpu[:num_reqs] = max_seq_lens
+            self.seq_lens_cpu[num_reqs:] = 0
+            if len(self.runner.attn_groups) > 0:
+                num_computed_tokens_cpu = (
+                    self.runner.input_batch.
+                    num_computed_tokens_cpu_tensor[:num_reqs])
+                query_start_loc = torch.tensor(
+                    [0] + self.actual_seq_lengths_q[:num_reqs],
+                    device=self.runner.device,
+                    dtype=torch.int32)
+                self.query_start_loc[:num_reqs + 1].copy_(query_start_loc)
+                common_attn_metadata = AscendCommonAttentionMetadata(
+                    query_start_loc=self.query_start_loc[:num_reqs + 1],
+                    query_start_loc_cpu=self.query_start_loc_cpu[:num_reqs +
+                                                                 1],
+                    seq_lens_cpu=self.seq_lens_cpu,
+                    seq_lens=self.seq_lens_cpu[:num_reqs],
+                    num_reqs=num_reqs,
+                    num_actual_tokens=num_tokens,
+                    max_query_len=self.num_speculative_tokens + 1,
+                    num_computed_tokens_cpu=num_computed_tokens_cpu,
+                    actual_seq_lengths_q=self.runner.actual_seq_lengths_q,
+                    block_table_tensor=self.runner.input_batch.block_table[0].
+                    get_device_tensor()[:num_reqs],
+                    slot_mapping=self.slot_mapping,
+                    positions=self.positions,
+                    attn_mask=self.runner.attn_mask,
+                    spec_attn_mask=self.runner.spec_attn_mask,
+                    attn_state=self.runner.attn_state,
+                    decode_token_per_req=self.runner.decode_token_per_req,
+                    cos=self.runner.cos,  # 考虑mrope，是否可以共用？
+                    sin=self.runner.sin,
+                )
+
+                builder = self.runner.attn_groups[0][0].get_metadata_builder()
+                attn_metadata_mtp = builder.build_for_graph_capture(
+                    common_attn_metadata, AscendAttentionState.SpecDecoding,
+                    self.runner.get_model())
+                attn_metadata = {}
+                for layer_name in self.attn_layer_name:
+                    attn_metadata[layer_name] = attn_metadata_mtp
+            else:
+                attn_metadata = None
+        else:
+            attn_metadata = None
 
         input_ids = self.input_ids[:num_tokens]
         positions = self.positions[:num_tokens]
@@ -158,7 +238,8 @@ def dummy_run(self,
                     in_profile_run=self.runner.in_profile_run,
                     num_actual_tokens=0,
                     aclgraph_runtime_mode=aclgraph_runtime_mode,
-                    batch_descriptor=batch_descriptor):
+                    batch_descriptor=batch_descriptor,
+                    is_mtp_model=True):
                 if is_running_torchair:
                     assert attn_metadata is not None
                     torch._dynamo.mark_static(input_ids)
@@ -188,6 +269,14 @@ def dummy_run(self,
                     self.model(input_ids=input_ids,
                                positions=positions,
                                hidden_states=previous_hidden_states)
+                    forward_context = get_forward_context()
+                    if forward_context.cudagraph_runtime_mode == CUDAGraphMode.FULL and \
+                        not forward_context.capturing:
+                        if self.vllm_config.model_config.use_mla:
+                            update_mla_attn_params(
+                                self.update_stream, forward_context,
+                                positions.shape[0],
+                                self.vllm_config.speculative_config)
             if with_prefill:
                 break
 
@@ -260,7 +349,8 @@ def generate_token_ids(self,
             cu_num_tokens=cu_num_tokens,
             block_table=attn_metadata.block_tables,
             sampling_metadata=sampling_metadata,
-            token_indices=accepted_token_indices)
+            token_indices=accepted_token_indices,
+            scheduler_output=scheduler_output)
         spec_token_ids = draft_token_ids.tolist()
         return spec_token_ids
 
@@ -322,6 +412,17 @@ def _prepare_inputs(
             target_positions = positions[token_indices]
             target_hidden_states = hidden_states[token_indices]
             target_slot_mapping = slot_mapping[token_indices]
+
+        batch_size = num_rejected_tokens.shape[0]
+        self.query_start_loc[:batch_size + 1].copy_(cu_num_tokens[:batch_size +
+                                                                  1])
+        self.query_start_loc_cpu[:batch_size + 1].copy_(
+            self.query_start_loc[:batch_size + 1], non_blocking=True)
+        target_positions_len = target_positions.shape[0]
+        self.positions[:target_positions_len].copy_(target_positions)
+        target_slot_mapping_len = target_slot_mapping.shape[0]
+        self.slot_mapping[:target_slot_mapping_len].copy_(target_slot_mapping)
+
         return cu_num_tokens, token_indices, target_token_ids, target_positions, target_hidden_states, target_slot_mapping
 
     def _propose(
@@ -341,7 +442,8 @@ def _propose(
             # [batch_size, max_num_blocks_per_req]
             block_table: torch.Tensor,
             sampling_metadata: SamplingMetadata,
-            token_indices=None) -> torch.Tensor:
+            token_indices=None,
+            scheduler_output: SchedulerOutput = None) -> torch.Tensor:
         num_tokens = target_token_ids.shape[0]
         batch_size = next_token_ids.shape[0]
         last_token_indices = cu_num_tokens[1:] - 1
@@ -385,18 +487,20 @@ def _propose(
 
         seq_lens = target_positions[last_token_indices] + 1
         seq_lens = seq_lens.int()
+        seq_lens_len = seq_lens.shape[0]
+        self.seq_lens_cpu[:seq_lens_len].copy_(seq_lens, non_blocking=True)
         common_attn_metadata = AscendCommonAttentionMetadata(
-            query_start_loc=cu_num_tokens[:batch_size + 1],
-            query_start_loc_cpu=cu_num_tokens[:batch_size + 1].cpu(),
-            seq_lens_cpu=seq_lens.cpu(),
+            query_start_loc=self.query_start_loc[:batch_size + 1],
+            query_start_loc_cpu=self.query_start_loc_cpu[:batch_size + 1],
+            seq_lens_cpu=self.seq_lens_cpu[:seq_lens_len],
             num_reqs=batch_size,
             num_actual_tokens=num_tokens,
             max_query_len=max_query_len,
             actual_seq_lengths_q=self.runner.actual_seq_lengths_q,
             block_table_tensor=self.runner.input_batch.block_table[0].
             get_device_tensor(),
-            slot_mapping=target_slot_mapping,
-            positions=target_positions,
+            slot_mapping=self.slot_mapping[:target_slot_mapping.shape[0]],
+            positions=self.positions[:target_positions.shape[0]],
             attn_mask=self.runner.attn_mask,
             spec_attn_mask=self.runner.spec_attn_mask,
             attn_state=self.runner.attn_state,
@@ -434,8 +538,18 @@ def _propose(
 
         moe_comm_type = self.runner._select_moe_comm_method(
             num_input_tokens, with_prefill)
-        batch_descriptor = BatchDescriptor(num_tokens=num_input_tokens,
-                                           uniform_decode=False)
+
+        if scheduler_output:
+            uniform_decode = (max_query_len in list(
+                range(1, self.num_speculative_tokens + 2))) and (
+                    scheduler_output.total_num_scheduled_tokens //
+                    (self.num_speculative_tokens + 2 - max_query_len)
+                    == self.runner.input_batch.num_reqs * max_query_len)
+            batch_descriptor = BatchDescriptor(num_tokens=num_input_tokens,
+                                               uniform_decode=uniform_decode)
+        else:
+            batch_descriptor = BatchDescriptor(num_tokens=num_input_tokens,
+                                               uniform_decode=False)
         aclgraph_runtime_mode, batch_descriptor = \
             self.runner.aclgraph_dispatcher.dispatch(batch_descriptor)
 
@@ -451,7 +565,8 @@ def _propose(
                     aclgraph_runtime_mode=aclgraph_runtime_mode,
                     batch_descriptor=batch_descriptor,
                     in_profile_run=self.runner.in_profile_run,
-                    num_actual_tokens=num_tokens):
+                    num_actual_tokens=num_tokens,
+                    is_mtp_model=True):
                 with ProfileExecuteDuration().capture_async('mtp_forward'):
                     model_kwargs = {}
                     model_kwargs["attn_metadata"] = attn_metadata
@@ -475,6 +590,13 @@ def _propose(
                             positions=self.positions[:num_input_tokens],
                             hidden_states=self.hidden_states[:num_input_tokens]
                         )
+                        forward_context = get_forward_context()
+                        if forward_context.cudagraph_runtime_mode == CUDAGraphMode.FULL:
+                            if self.vllm_config.model_config.use_mla:
+                                update_mla_attn_params(
+                                    self.update_stream, forward_context,
+                                    num_input_tokens,
+                                    self.vllm_config.speculative_config)
 
             num_indices = last_token_indices.shape[0]
             if lmhead_tp_enable():

vllm_ascend/worker/model_runner_v1.py

@@ -2501,7 +2501,6 @@ def dummy_compute_logits(hidden_states):
                 self.drafter.dummy_run(
                     num_tokens=num_tokens,
                     with_prefill=with_prefill,
-                    skip_attn=True,
                     num_reqs=num_reqs,
                     num_tokens_across_dp=num_tokens_across_dp,
                     aclgraph_runtime_mode=aclgraph_runtime_mode,