[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

@@ -446,6 +446,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'):