[0.11.0][MTP][Aclgraph] Fix the support aclgraph with MTP (#3912)
### What this PR does / why we need it? Fix 2 breaks of aclgraph with MTP: 1. deepseekmtp in vllm 0.11.0 does not support aclgraph and lack the `support_torch_compile` decorator 2. There is a d2h synchornization in the original forward of mtp predictor. The fix pr in vllm https://github.com/vllm-project/vllm/pull/27643 As we'll fix it in vllm main, this fix pr is only needed in branch v0.11.0-dev The profling shows that MTP replays in aclgraph now: <img width="1612" height="1866" alt="a7d7f04155df4ed454b7eb20a92b2e2a" src="https://github.com/user-attachments/assets/eaa4b9ff-aeb0-416d-964f-5a06e497f155" /> ### How was this patch tested? <!-- CI passed with new added/existing test. If it was tested in a way different from regular unit tests, please clarify how you tested step by step, ideally copy and paste-able, so that other reviewers can test and check, and descendants can verify in the future. If tests were not added, please describe why they were not added and/or why it was difficult to add. --> --------- Signed-off-by: MengqingCao <cmq0113@163.com>
This commit is contained in:
@@ -11,7 +11,6 @@ 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 (
|
||||
process_weights_after_loading, set_default_torch_dtype)
|
||||
from vllm.model_executor.models.deepseek_mtp import DeepSeekMTP
|
||||
from vllm.v1.core.sched.output import SchedulerOutput
|
||||
from vllm.v1.sample.metadata import SamplingMetadata
|
||||
from vllm.v1.spec_decode.metadata import SpecDecodeMetadata
|
||||
@@ -19,6 +18,8 @@ from vllm.v1.spec_decode.metadata import SpecDecodeMetadata
|
||||
from vllm_ascend.ascend_config import get_ascend_config
|
||||
from vllm_ascend.ascend_forward_context import set_ascend_forward_context
|
||||
from vllm_ascend.attention.utils import AscendCommonAttentionMetadata
|
||||
from vllm_ascend.patch.worker.patch_deepseek_mtp import \
|
||||
AscendDeepSeekMTP as DeepSeekMTP
|
||||
from vllm_ascend.spec_decode.interface import Proposer, SpecDcodeType
|
||||
from vllm_ascend.torchair.models.torchair_deepseek_mtp import \
|
||||
TorchairDeepSeekMTP
|
||||
|
||||
Reference in New Issue
Block a user