[WIP]mtp aclgraph support

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

Author: anon189Ty

vllm_ascend/models/deepseek_mtp.py

@@ -23,6 +23,7 @@
 import torch.nn as nn
 from transformers import PretrainedConfig
 from vllm.attention.backends.abstract import AttentionMetadata
+from vllm.compilation.decorators import support_torch_compile
 from vllm.config import (CacheConfig, ModelConfig, VllmConfig,
                          get_current_vllm_config)
 from vllm.model_executor.layers.layernorm import RMSNorm
@@ -177,6 +178,7 @@ def compute_logits(
         return logits
 
 
+@support_torch_compile
 class CustomDeepSeekMTP(DeepSeekMTP):
 
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

vllm_ascend/spec_decode/mtp_proposer.py

@@ -5,8 +5,8 @@
 import torchair
 from torchair import patch_for_hcom
 from vllm.attention.layer import Attention
-from vllm.config import (VllmConfig, get_layers_from_vllm_config,
-                         set_current_vllm_config)
+from vllm.config import (CUDAGraphMode, VllmConfig,
+                         get_layers_from_vllm_config, set_current_vllm_config)
 from vllm.forward_context import BatchDescriptor, get_forward_context
 from vllm.model_executor.model_loader import get_model_loader
 from vllm.model_executor.model_loader.utils import (
@@ -110,7 +110,9 @@ def dummy_run(self,
                   with_prefill: bool = False,
                   skip_attn: bool = False,
                   num_reqs: int = 0,
-                  num_tokens_across_dp=None) -> None:
+                  num_tokens_across_dp=None,
+                  aclgraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE,
+                  batch_descriptor=None) -> None:
         if not self.torchair_graph_enabled:
             # TODO: adapt enable_dbo later
             (num_tokens, num_tokens_across_dp, with_prefill,
@@ -152,7 +154,9 @@ def dummy_run(self,
                     reserved_mc2_mask=self.runner.reserved_mc2_mask,
                     moe_comm_type=moe_comm_type,
                     in_profile_run=self.runner.in_profile_run,
-                    num_actual_tokens=0):
+                    num_actual_tokens=0,
+                    aclgraph_runtime_mode=aclgraph_runtime_mode,
+                    batch_descriptor=batch_descriptor):
                 if is_running_torchair:
                     assert attn_metadata is not None
                     torch._dynamo.mark_static(input_ids)
@@ -446,6 +450,7 @@ def _propose(
                     reserved_mc2_mask=self.runner.reserved_mc2_mask,
                     moe_comm_type=moe_comm_type,
                     aclgraph_runtime_mode=aclgraph_runtime_mode,
+                    batch_descriptor=batch_descriptor,
                     in_profile_run=self.runner.in_profile_run,
                     num_actual_tokens=num_tokens):
                 with ProfileExecuteDuration().capture_async('mtp_forward'):

vllm_ascend/worker/model_runner_v1.py

@@ -2484,7 +2484,9 @@ def dummy_compute_logits(hidden_states):
                     with_prefill=with_prefill,
                     skip_attn=True,
                     num_reqs=num_reqs,
-                    num_tokens_across_dp=num_tokens_across_dp)
+                    num_tokens_across_dp=num_tokens_across_dp,
+                    aclgraph_runtime_mode=aclgraph_runtime_mode,
+                    batch_descriptor=batch_descriptor)
                 if need_dummy_logits:
                     dummy_compute_logits(hidden_states)
             if self.in_profile_run and self.dynamic_eplb: