c980e68d40
### What this PR does / why we need it?
This PR aims to support aclgraph for model runner v2, please see RFC
#5208. The PR contains these modifications:
- adapt to newest commit of vllm main branch.
- supply a unified interface of extra forward context for both model
runner v1 and model runner v2.
- implement graph mode for main model.
### Does this PR introduce _any_ user-facing change?
no
### How was this patch tested?
- vLLM version: v0.16.0
- vLLM main:
4034c3d32e
---------
Signed-off-by: Ronald1995 <ronaldautomobile@163.com>
91 lines
3.5 KiB
Python
91 lines
3.5 KiB
Python
# Copyright (c) 2026 Huawei Technologies Co., Ltd. All Rights Reserved.
|
|
# Copyright 2023 The vLLM team.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
# This file is a part of the vllm-ascend project.
|
|
from __future__ import annotations
|
|
|
|
import torch
|
|
|
|
from vllm_ascend.ascend_forward_context import _EXTRA_CTX
|
|
from vllm_ascend.ops.fused_moe.moe_comm_method import AllGatherCommImpl, FusedExpertsResult
|
|
|
|
from .moe_mlp import unified_apply_mlp
|
|
from .token_dispatcher import TokenDispatcherWithAllGather310
|
|
|
|
|
|
class AllGatherCommImpl310(AllGatherCommImpl):
|
|
"""This implementation is the same as NativeAllGatherCommImpl,
|
|
but uses NPU-specific ops for better performance.
|
|
|
|
This implementation should be compatible with all scenarios, and
|
|
thus it is the default implementation for MoE communication methods.
|
|
It uses `torch_npu.npu_moe_init_routing_v2` for pre-processing
|
|
and `torch_npu.npu_moe_token_unpermute` for post-processing
|
|
to handle the token-to-expert mapping and communication efficiently.
|
|
"""
|
|
|
|
def fused_experts( # type: ignore[override]
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
w1: torch.Tensor,
|
|
w2: torch.Tensor,
|
|
topk_weights: torch.Tensor,
|
|
topk_ids: torch.Tensor,
|
|
expert_map: torch.Tensor | None = None,
|
|
use_int8_w8a8: bool = False,
|
|
w1_scale: torch.Tensor | None = None,
|
|
w2_scale: torch.Tensor | None = None,
|
|
apply_router_weight_on_input: bool = False,
|
|
) -> FusedExpertsResult:
|
|
# This method is overridden to use the 310p-specific unified_apply_mlp
|
|
# which provides optimized MLP computation for the 310p platform
|
|
moe_comm_method = _EXTRA_CTX.moe_comm_method
|
|
assert moe_comm_method is not None, "Missing communication context"
|
|
|
|
dispatch_results = self.token_dispatcher.token_dispatch(
|
|
hidden_states=hidden_states,
|
|
topk_weights=topk_weights,
|
|
topk_ids=topk_ids,
|
|
expert_map=expert_map,
|
|
apply_router_weight_on_input=apply_router_weight_on_input,
|
|
)
|
|
|
|
mlp_output = unified_apply_mlp(
|
|
hidden_states=dispatch_results.hidden_states,
|
|
w1=w1,
|
|
w2=w2,
|
|
w1_scale=w1_scale,
|
|
w2_scale=w2_scale,
|
|
group_list=dispatch_results.group_list,
|
|
group_list_type=dispatch_results.group_list_type,
|
|
with_quant=use_int8_w8a8,
|
|
)
|
|
|
|
combine_results = self.token_dispatcher.token_combine(
|
|
hidden_states=mlp_output, context_metadata=dispatch_results.context_metadata
|
|
)
|
|
|
|
return FusedExpertsResult(
|
|
routed_out=combine_results.routed_out,
|
|
group_list_type=dispatch_results.group_list_type,
|
|
expert_tokens=dispatch_results.group_list,
|
|
)
|
|
|
|
def _get_token_dispatcher(self):
|
|
return TokenDispatcherWithAllGather310(
|
|
top_k=self.moe_config.experts_per_token,
|
|
num_experts=self.moe_config.num_experts,
|
|
num_local_experts=self.moe_config.num_local_experts,
|
|
)
|