0
vllm_ascend/_310p/fused_moe/__init__.py
Normal file
0
vllm_ascend/_310p/fused_moe/__init__.py
Normal file
89
vllm_ascend/_310p/fused_moe/experts_selector.py
Normal file
89
vllm_ascend/_310p/fused_moe/experts_selector.py
Normal file
@@ -0,0 +1,89 @@
|
||||
#
|
||||
# Copyright (c) 2026 Huawei Technologies Co., Ltd. All Rights Reserved.
|
||||
# This file is a part of the vllm-ascend project.
|
||||
#
|
||||
# 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.
|
||||
#
|
||||
from collections.abc import Callable
|
||||
|
||||
import torch
|
||||
import torch_npu
|
||||
|
||||
from vllm_ascend.ops.fused_moe.experts_selector import _native_select_experts, _renormalize_topk_weights
|
||||
|
||||
|
||||
def select_experts(
|
||||
hidden_states: torch.Tensor,
|
||||
router_logits: torch.Tensor,
|
||||
top_k: int,
|
||||
use_grouped_topk: bool,
|
||||
renormalize: bool,
|
||||
topk_group: int | None = None,
|
||||
num_expert_group: int | None = None,
|
||||
custom_routing_function: Callable | None = None,
|
||||
scoring_func: str = "softmax",
|
||||
routed_scaling_factor: float = 1.0,
|
||||
e_score_correction_bias: torch.Tensor | None = None,
|
||||
global_num_experts: int = -1,
|
||||
):
|
||||
"""
|
||||
Fused experts with select experts.
|
||||
|
||||
Args:
|
||||
router_logits: router logits of shape (num_tokens, hidden_size).
|
||||
hidden_states: Hidden states of shape (num_tokens, hidden_size).
|
||||
top_k: number of top k experts.
|
||||
use_grouped_topk: Whether to group experts before selecting top-k.
|
||||
renormalize: Whether to renormalize the routing weights.
|
||||
topk_group: Number of expert groups to select from.
|
||||
num_expert_group: Number of experts in each group.
|
||||
custom_routing_function: Custom routing function.
|
||||
scoring_func: Scoring function to use.
|
||||
e_score_correction_bias: Correction bias to apply to expert scores.
|
||||
routed_scaling_factor: Scaling factor applied to routing weights.
|
||||
global_num_experts: Global number of experts.
|
||||
|
||||
Returns:
|
||||
topk_weights: router weights of shape (num_tokens, top_k).
|
||||
topk_ids: selected expert IDs of shape (num_tokens, top_k).
|
||||
"""
|
||||
if scoring_func == "softmax" and not use_grouped_topk and custom_routing_function is None:
|
||||
# 310P returns invalid routing results when this op receives more than 1024 tokens.
|
||||
if router_logits.shape[0] > 1024:
|
||||
topk_results = [
|
||||
torch_npu.npu_moe_gating_top_k_softmax(router_logits_chunk, k=top_k)
|
||||
for router_logits_chunk in router_logits.split(1024, dim=0)
|
||||
]
|
||||
topk_weights = torch.cat([result[0] for result in topk_results], dim=0)
|
||||
topk_ids = torch.cat([result[1] for result in topk_results], dim=0)
|
||||
else:
|
||||
topk_weights, topk_ids, _ = torch_npu.npu_moe_gating_top_k_softmax(router_logits, k=top_k)
|
||||
topk_weights = _renormalize_topk_weights(topk_weights, renormalize)
|
||||
else:
|
||||
topk_weights, topk_ids = _native_select_experts(
|
||||
hidden_states=hidden_states,
|
||||
router_logits=router_logits,
|
||||
top_k=top_k,
|
||||
use_grouped_topk=use_grouped_topk,
|
||||
renormalize=renormalize,
|
||||
topk_group=topk_group,
|
||||
num_expert_group=num_expert_group,
|
||||
custom_routing_function=custom_routing_function,
|
||||
scoring_func=scoring_func,
|
||||
e_score_correction_bias=e_score_correction_bias,
|
||||
)
|
||||
# Apply routed scaling factor to weights
|
||||
if routed_scaling_factor != 1.0:
|
||||
topk_weights = topk_weights * routed_scaling_factor
|
||||
|
||||
return topk_weights, topk_ids
|
||||
367
vllm_ascend/_310p/fused_moe/fused_moe.py
Normal file
367
vllm_ascend/_310p/fused_moe/fused_moe.py
Normal file
@@ -0,0 +1,367 @@
|
||||
#
|
||||
# Copyright (c) 2026 Huawei Technologies Co., Ltd. All Rights Reserved.
|
||||
# This file is a part of the vllm-ascend project.
|
||||
#
|
||||
# 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.
|
||||
#
|
||||
from collections.abc import Callable
|
||||
|
||||
import torch
|
||||
from vllm.distributed import get_dp_group, get_ep_group, get_tp_group
|
||||
from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig
|
||||
|
||||
from vllm_ascend.ascend_forward_context import _EXTRA_CTX, MoECommType
|
||||
from vllm_ascend.ops.fused_moe.experts_selector import zero_experts_compute
|
||||
from vllm_ascend.ops.fused_moe.fused_moe import AscendMoERunner
|
||||
from vllm_ascend.ops.fused_moe.moe_comm_method import (
|
||||
AllGatherCommImpl,
|
||||
FusedExpertsResult,
|
||||
_MoECommMethods,
|
||||
)
|
||||
from vllm_ascend.ops.fused_moe.moe_runtime_args import build_fused_experts_input
|
||||
from vllm_ascend.quantization.quant_type import QuantType
|
||||
from vllm_ascend.utils import maybe_trans_nz, vllm_version_is
|
||||
|
||||
from .experts_selector import select_experts
|
||||
from .moe_comm_method import AllGatherCommImpl310
|
||||
|
||||
if vllm_version_is("0.23.0"):
|
||||
from vllm.model_executor.layers.fused_moe.layer import FusedMoE as _LegacyFusedMoEBase
|
||||
from vllm.model_executor.layers.fused_moe.layer import UnquantizedFusedMoEMethod
|
||||
else:
|
||||
from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import UnquantizedFusedMoEMethod
|
||||
|
||||
try:
|
||||
from vllm.model_executor.layers.fused_moe.layer import FusedMoE as _LegacyFusedMoEBase
|
||||
except ImportError:
|
||||
_LegacyFusedMoEBase = torch.nn.Module
|
||||
|
||||
if not isinstance(_LegacyFusedMoEBase, type):
|
||||
_LegacyFusedMoEBase = torch.nn.Module
|
||||
|
||||
|
||||
class AscendUnquantizedFusedMoEMethod310(UnquantizedFusedMoEMethod):
|
||||
def __init__(self, moe: FusedMoEConfig = None):
|
||||
super().__init__(moe=moe)
|
||||
|
||||
@property
|
||||
def is_monolithic(self) -> bool:
|
||||
return False
|
||||
|
||||
def maybe_make_prepare_finalize(self, routing_tables=None):
|
||||
# Ascend 310P uses its own MoE communication and forward_impl path.
|
||||
# Do not let upstream modular-kernel initialization replace it.
|
||||
return None
|
||||
|
||||
def process_weights_after_loading(self, layer):
|
||||
super().process_weights_after_loading(layer)
|
||||
|
||||
# Fused gate_up_proj (column parallel)
|
||||
w13_data = self._maybe_pad_weight(layer.w13_weight.data).transpose(1, 2).contiguous()
|
||||
w13_data = maybe_trans_nz(w13_data)
|
||||
layer.w13_weight = torch.nn.Parameter(w13_data, requires_grad=False)
|
||||
# down_proj (row parallel)
|
||||
w2_data = self._maybe_pad_weight(layer.w2_weight.data).transpose(1, 2).contiguous()
|
||||
w2_data = maybe_trans_nz(w2_data)
|
||||
layer.w2_weight = torch.nn.Parameter(w2_data, requires_grad=False)
|
||||
|
||||
def apply(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
x: torch.Tensor,
|
||||
use_grouped_topk: bool,
|
||||
top_k: int,
|
||||
router_logits: torch.Tensor,
|
||||
renormalize: bool,
|
||||
topk_group: int | None = None,
|
||||
num_expert_group: int | None = None,
|
||||
custom_routing_function: Callable | None = None,
|
||||
scoring_func: str = "softmax",
|
||||
e_score_correction_bias: torch.Tensor | None = None,
|
||||
num_experts: int = -1,
|
||||
expert_map: torch.Tensor | None = None,
|
||||
apply_router_weight_on_input: bool = False,
|
||||
**kwargs,
|
||||
) -> torch.Tensor:
|
||||
zero_expert_num = getattr(layer, "zero_expert_num", 0)
|
||||
zero_expert_type = getattr(layer, "zero_expert_type", None)
|
||||
|
||||
topk_weights, topk_ids = select_experts(
|
||||
hidden_states=x,
|
||||
router_logits=router_logits,
|
||||
top_k=top_k,
|
||||
use_grouped_topk=use_grouped_topk,
|
||||
renormalize=renormalize,
|
||||
topk_group=topk_group,
|
||||
num_expert_group=num_expert_group,
|
||||
custom_routing_function=custom_routing_function,
|
||||
scoring_func=scoring_func,
|
||||
e_score_correction_bias=e_score_correction_bias,
|
||||
global_num_experts=num_experts,
|
||||
)
|
||||
|
||||
if zero_expert_num > 0 and zero_expert_type is not None:
|
||||
topk_ids, topk_weights, zero_expert_result = zero_experts_compute(
|
||||
expert_indices=topk_ids,
|
||||
expert_scales=topk_weights,
|
||||
num_experts=num_experts,
|
||||
zero_expert_type=zero_expert_type,
|
||||
hidden_states=x,
|
||||
)
|
||||
|
||||
topk_weights = topk_weights.to(x.dtype)
|
||||
|
||||
moe_comm_method = _EXTRA_CTX.moe_comm_method
|
||||
final_hidden_states = moe_comm_method.fused_experts(
|
||||
fused_experts_input=build_fused_experts_input(
|
||||
hidden_states=x,
|
||||
topk_weights=topk_weights,
|
||||
topk_ids=topk_ids,
|
||||
w1=layer.w13_weight,
|
||||
w2=layer.w2_weight,
|
||||
quant_type=QuantType.NONE,
|
||||
dynamic_eplb=False,
|
||||
expert_map=expert_map,
|
||||
apply_router_weight_on_input=apply_router_weight_on_input,
|
||||
),
|
||||
)
|
||||
if zero_expert_num > 0 and zero_expert_type is not None:
|
||||
final_hidden_states += zero_expert_result
|
||||
return final_hidden_states
|
||||
|
||||
|
||||
if not vllm_version_is("0.23.0"):
|
||||
|
||||
class AscendMoERunner310(AscendMoERunner):
|
||||
def __init__(
|
||||
self,
|
||||
layer_name,
|
||||
moe_config,
|
||||
router,
|
||||
routed_experts,
|
||||
enable_dbo=False,
|
||||
gate=None,
|
||||
shared_experts=None,
|
||||
shared_expert_gate=None,
|
||||
routed_input_transform=None,
|
||||
routed_output_transform=None,
|
||||
routed_scaling_factor=1,
|
||||
tid2eid=None,
|
||||
n_shared_experts: int = 0,
|
||||
):
|
||||
super().__init__(
|
||||
layer_name,
|
||||
moe_config,
|
||||
router,
|
||||
routed_experts,
|
||||
enable_dbo,
|
||||
gate,
|
||||
shared_experts,
|
||||
shared_expert_gate,
|
||||
routed_input_transform,
|
||||
routed_output_transform,
|
||||
routed_scaling_factor,
|
||||
tid2eid,
|
||||
n_shared_experts,
|
||||
)
|
||||
|
||||
if routed_experts.quant_config is None:
|
||||
routed_experts.quant_method = AscendUnquantizedFusedMoEMethod310(self.moe_config)
|
||||
self.quant_type = self._get_quant_type()
|
||||
|
||||
self.multistream_overlap_gate = False
|
||||
self.shared_multistream_overlap_gate = False
|
||||
self.multistream_overlap_shared_expert = False
|
||||
_MoECommMethods[MoECommType.ALLGATHER] = AllGatherCommImpl310(self.moe_config)
|
||||
|
||||
|
||||
class AscendFusedMoE310(_LegacyFusedMoEBase):
|
||||
def __init__(self, *args, **kwargs):
|
||||
if _LegacyFusedMoEBase is torch.nn.Module:
|
||||
raise RuntimeError("AscendFusedMoE310 is only kept for the legacy FusedMoE class API.")
|
||||
super().__init__(*args, **kwargs)
|
||||
|
||||
self._routed_input_transform = kwargs.get("routed_input_transform")
|
||||
self._shared_experts = kwargs.get("shared_experts")
|
||||
self.global_num_experts = kwargs["num_experts"]
|
||||
|
||||
if self.quant_config is None:
|
||||
self.quant_method = AscendUnquantizedFusedMoEMethod310(self.moe_config)
|
||||
else:
|
||||
self.quant_method = self.quant_config.get_quant_method(self, self.layer_name)
|
||||
|
||||
assert self.quant_method is not None
|
||||
# Keep base_quant_method aligned with the Ascend-replaced quant_method
|
||||
# so FusedMoE.maybe_init_modular_kernel doesn't dispatch into the
|
||||
# upstream UnquantizedFusedMoEMethod.maybe_make_prepare_finalize.
|
||||
self.base_quant_method = self.quant_method
|
||||
|
||||
self.moe_config.tp_group = get_tp_group()
|
||||
self.moe_config.dp_group = get_dp_group()
|
||||
self.moe_config.ep_group = get_ep_group()
|
||||
self.moe_config.supports_eplb = False
|
||||
|
||||
# init moe
|
||||
self.global_expert_map = None
|
||||
self.local_expert_map = None
|
||||
if self.moe_config.ep_size > 1:
|
||||
raise RuntimeError("Expert Parallel is not supported on 310P. Please remove --enable-expert-parallel.")
|
||||
self.local_num_experts = self.global_num_experts
|
||||
|
||||
self.moe_config.num_experts = self.global_num_experts
|
||||
self.moe_config.num_local_experts = self.local_num_experts
|
||||
self.moe_config.global_redundant_expert_num = 0
|
||||
|
||||
moe_quant_params = {
|
||||
"num_experts": self.local_num_experts,
|
||||
"hidden_size": self.hidden_size,
|
||||
"intermediate_size_per_partition": self.intermediate_size_per_partition,
|
||||
"params_dtype": self.params_dtype,
|
||||
"weight_loader": self.weight_loader,
|
||||
}
|
||||
|
||||
self.quant_method.create_weights(layer=self, **moe_quant_params)
|
||||
self.quant_type = self.get_quant_type()
|
||||
|
||||
_MoECommMethods[MoECommType.ALLGATHER] = AllGatherCommImpl310(self.moe_config)
|
||||
|
||||
if vllm_version_is("0.23.0"):
|
||||
self.runner = AscendMoERunner(
|
||||
self.layer_name,
|
||||
self.moe_config,
|
||||
self.router,
|
||||
self._routed_input_transform,
|
||||
kwargs.pop("gate", None),
|
||||
kwargs.pop("shared_experts", None),
|
||||
self.quant_method,
|
||||
self.vllm_config.parallel_config.enable_dbo,
|
||||
)
|
||||
else:
|
||||
self.runner = AscendMoERunner310(
|
||||
self.layer_name,
|
||||
self.moe_config,
|
||||
self.router,
|
||||
self._routed_input_transform,
|
||||
kwargs.pop("gate", None),
|
||||
kwargs.pop("shared_experts", None),
|
||||
self.quant_method,
|
||||
self.vllm_config.parallel_config.enable_dbo,
|
||||
)
|
||||
|
||||
@property
|
||||
def is_internal_router(self) -> bool:
|
||||
# 310P Ascend path expects router logits from the model forward path.
|
||||
return False
|
||||
|
||||
def init_experts_map(self, moe_config):
|
||||
"""
|
||||
Initialize expert mapping for MoE (Mixture of Experts) model.
|
||||
|
||||
This function creates mappings between global expert indices and local expert indices
|
||||
for each rank in the expert parallel group. It divides the total experts among
|
||||
different ranks and creates both global and local expert maps that are used
|
||||
during MoE computation to determine which experts are handled by which rank.
|
||||
|
||||
Args:
|
||||
moe_config: Configuration object containing MoE parameters including
|
||||
number of experts, expert parallel size, and expert parallel rank.
|
||||
|
||||
Returns:
|
||||
tuple: A tuple containing:
|
||||
- global_expert_map: Stack of expert maps for all ranks
|
||||
- local_expert_map: Expert map for the current rank (transferred to NPU)
|
||||
"""
|
||||
n_experts = moe_config.num_experts
|
||||
ep_size = moe_config.ep_size
|
||||
all_experts = torch.arange(n_experts, dtype=torch.int32)
|
||||
experts_groups = all_experts.chunk(ep_size)
|
||||
global_expert_map = []
|
||||
local_expert_map = None
|
||||
for rankid in range(ep_size):
|
||||
expert_map = torch.full((n_experts,), -1, dtype=torch.int32)
|
||||
local_experts = experts_groups[rankid]
|
||||
expert_map[local_experts] = torch.arange(local_experts.shape[0], dtype=torch.int32)
|
||||
global_expert_map.append(expert_map)
|
||||
if rankid == moe_config.ep_rank:
|
||||
local_expert_map = expert_map.npu()
|
||||
return torch.stack(global_expert_map), local_expert_map
|
||||
|
||||
def get_quant_type(self) -> QuantType:
|
||||
quant_method = self.quant_method
|
||||
if not hasattr(quant_method, "quant_method") or quant_method.quant_method is None:
|
||||
return QuantType.NONE
|
||||
|
||||
method = quant_method.quant_method
|
||||
quant_type = getattr(method, "quant_type", QuantType.NONE)
|
||||
if quant_type not in [QuantType.NONE, QuantType.W8A8]:
|
||||
raise RuntimeError("Only Unquant and W8A8 is supported.")
|
||||
return quant_type
|
||||
|
||||
def forward_impl( # type: ignore[override]
|
||||
self, hidden_states: torch.Tensor, router_logits: torch.Tensor
|
||||
) -> torch.Tensor:
|
||||
assert self.quant_method is not None
|
||||
assert self.routed_scaling_factor == 1.0, "routed_scaling_factor != 1.0 is not supported."
|
||||
|
||||
prepare_output = _EXTRA_CTX.moe_comm_method.prepare(
|
||||
hidden_states=hidden_states, router_logits=router_logits, quant_type=self.quant_type
|
||||
)
|
||||
hidden_states = prepare_output.hidden_states
|
||||
router_logits = prepare_output.router_logits
|
||||
pertoken_scale = prepare_output.pertoken_scale
|
||||
padded_hidden_states_shape = prepare_output.padded_hidden_states_shape
|
||||
|
||||
# Matrix multiply.
|
||||
fused_experts_results: FusedExpertsResult = self.quant_method.apply(
|
||||
layer=self,
|
||||
x=hidden_states,
|
||||
use_grouped_topk=self.use_grouped_topk,
|
||||
top_k=self.top_k,
|
||||
router_logits=router_logits,
|
||||
renormalize=self.renormalize,
|
||||
topk_group=self.topk_group,
|
||||
num_expert_group=self.num_expert_group,
|
||||
custom_routing_function=self.custom_routing_function,
|
||||
scoring_func=self.scoring_func,
|
||||
e_score_correction_bias=self.e_score_correction_bias,
|
||||
num_experts=self.global_num_experts,
|
||||
expert_map=self.local_expert_map,
|
||||
apply_router_weight_on_input=self.apply_router_weight_on_input,
|
||||
pertoken_scale=pertoken_scale,
|
||||
)
|
||||
|
||||
routed_out = _EXTRA_CTX.moe_comm_method.finalize(
|
||||
hidden_states=fused_experts_results.routed_out,
|
||||
reduce_results=isinstance(_EXTRA_CTX.moe_comm_method, AllGatherCommImpl),
|
||||
padded_hidden_states_shape=padded_hidden_states_shape,
|
||||
)
|
||||
|
||||
return routed_out
|
||||
|
||||
def _forward_shared_experts(self, hidden_states: torch.Tensor):
|
||||
if self._shared_experts is None:
|
||||
return None
|
||||
return self._shared_experts(hidden_states)
|
||||
|
||||
def shared_forward_impl( # type: ignore[override]
|
||||
self, hidden_states: torch.Tensor, router_logits: torch.Tensor
|
||||
):
|
||||
routed_out = AscendFusedMoE310.forward_impl(
|
||||
self,
|
||||
hidden_states=hidden_states,
|
||||
router_logits=router_logits,
|
||||
)
|
||||
if self._shared_experts is None:
|
||||
return routed_out
|
||||
shared_out = self._forward_shared_experts(hidden_states)
|
||||
return shared_out, routed_out
|
||||
50
vllm_ascend/_310p/fused_moe/moe_comm_method.py
Normal file
50
vllm_ascend/_310p/fused_moe/moe_comm_method.py
Normal file
@@ -0,0 +1,50 @@
|
||||
# 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
|
||||
|
||||
from typing import Any
|
||||
|
||||
from vllm_ascend.ops.fused_moe.moe_comm_method import AllGatherCommImpl
|
||||
from vllm_ascend.ops.fused_moe.moe_runtime_args import MoEMlpComputeInput
|
||||
|
||||
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 __init__(self, moe_config):
|
||||
super().__init__(moe_config)
|
||||
self.use_fusion_ops = False
|
||||
|
||||
def _apply_mlp(self, mlp_compute_input: MoEMlpComputeInput) -> Any:
|
||||
return unified_apply_mlp(mlp_compute_input=mlp_compute_input), None
|
||||
|
||||
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,
|
||||
)
|
||||
102
vllm_ascend/_310p/fused_moe/moe_mlp.py
Normal file
102
vllm_ascend/_310p/fused_moe/moe_mlp.py
Normal file
@@ -0,0 +1,102 @@
|
||||
# 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.
|
||||
|
||||
|
||||
import torch
|
||||
import torch_npu
|
||||
|
||||
from vllm_ascend.ops.fused_moe.moe_runtime_args import MoEMlpComputeInput
|
||||
|
||||
|
||||
def quant_apply_mlp(
|
||||
hidden_states: torch.Tensor,
|
||||
w1: torch.Tensor,
|
||||
w1_scale: torch.Tensor,
|
||||
w2: torch.Tensor,
|
||||
w2_scale: torch.Tensor,
|
||||
group_list: torch.Tensor,
|
||||
group_list_type: int = 1,
|
||||
) -> torch.Tensor:
|
||||
if group_list_type == 1:
|
||||
# Convert group_list to cumulative sum format if group_list is count format
|
||||
group_list = torch.cumsum(group_list, dim=0)
|
||||
|
||||
hidden_states = torch_npu.npu_quant_grouped_matmul_dequant(
|
||||
x=hidden_states, quantized_weight=w1, weight_scale=w1_scale, group_list=group_list, quant_mode="pertoken"
|
||||
)
|
||||
hidden_states = torch_npu.npu_swiglu(hidden_states)
|
||||
hidden_states = torch_npu.npu_quant_grouped_matmul_dequant(
|
||||
x=hidden_states, quantized_weight=w2, weight_scale=w2_scale, group_list=group_list, quant_mode="pertoken"
|
||||
)
|
||||
return hidden_states
|
||||
|
||||
|
||||
def unquant_apply_mlp(
|
||||
hidden_states: torch.Tensor, w1: torch.Tensor, w2: torch.Tensor, group_list: torch.Tensor, group_list_type: int = 1
|
||||
) -> torch.Tensor:
|
||||
gate_up_out = torch_npu.npu_grouped_matmul(
|
||||
x=[hidden_states],
|
||||
weight=[w1],
|
||||
split_item=2,
|
||||
group_list_type=group_list_type,
|
||||
group_type=0,
|
||||
group_list=group_list,
|
||||
)[0]
|
||||
act_out = torch_npu.npu_swiglu(gate_up_out)
|
||||
|
||||
hidden_states = torch_npu.npu_grouped_matmul(
|
||||
x=[act_out],
|
||||
weight=[w2],
|
||||
split_item=2,
|
||||
group_list_type=group_list_type,
|
||||
group_type=0,
|
||||
group_list=group_list,
|
||||
)[0]
|
||||
return hidden_states
|
||||
|
||||
|
||||
def unified_apply_mlp(*, mlp_compute_input: MoEMlpComputeInput) -> torch.Tensor:
|
||||
hidden_states = mlp_compute_input.hidden_states
|
||||
w1 = mlp_compute_input.weights.w1
|
||||
w2 = mlp_compute_input.weights.w2
|
||||
w1_scale = mlp_compute_input.weights.w1_scale
|
||||
w2_scale = mlp_compute_input.weights.w2_scale
|
||||
group_list = mlp_compute_input.group_list
|
||||
group_list_type = mlp_compute_input.group_list_type
|
||||
assert isinstance(w1, torch.Tensor)
|
||||
assert isinstance(w2, torch.Tensor)
|
||||
|
||||
if mlp_compute_input.quant.is_quant:
|
||||
assert isinstance(w1_scale, torch.Tensor)
|
||||
assert isinstance(w2_scale, torch.Tensor)
|
||||
assert w1_scale is not None and w2_scale is not None
|
||||
return quant_apply_mlp(
|
||||
hidden_states=hidden_states,
|
||||
w1=w1,
|
||||
w1_scale=w1_scale,
|
||||
w2=w2,
|
||||
w2_scale=w2_scale,
|
||||
group_list=group_list,
|
||||
group_list_type=group_list_type,
|
||||
)
|
||||
|
||||
return unquant_apply_mlp(
|
||||
hidden_states=hidden_states,
|
||||
w1=w1,
|
||||
w2=w2,
|
||||
group_list=group_list,
|
||||
group_list_type=group_list_type,
|
||||
)
|
||||
88
vllm_ascend/_310p/fused_moe/token_dispatcher.py
Normal file
88
vllm_ascend/_310p/fused_moe/token_dispatcher.py
Normal file
@@ -0,0 +1,88 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# Copyright (c) 2024; NVIDIA CORPORATION. All rights reserved.
|
||||
# Copyright (c) 2026 Huawei Technologies Co., Ltd. All Rights Reserved.
|
||||
# Copyright 2023 The vLLM team.
|
||||
# Copyright 2023 DeepSeek-AI and the HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
|
||||
# and OPT implementations in this library. It has been modified from its
|
||||
# original forms to accommodate minor architectural differences compared
|
||||
# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
|
||||
#
|
||||
# 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.
|
||||
|
||||
|
||||
import torch
|
||||
import torch_npu
|
||||
from vllm.distributed.parallel_state import get_ep_group
|
||||
|
||||
from vllm_ascend.ops.fused_moe.moe_runtime_args import MoEAllGatherCombineMetadata, MoETokenDispatchInput
|
||||
from vllm_ascend.ops.fused_moe.token_dispatcher import MoETokenDispatchOutput, TokenDispatcherWithAllGather
|
||||
|
||||
|
||||
class TokenDispatcherWithAllGather310(TokenDispatcherWithAllGather):
|
||||
def __init__(self, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
|
||||
def token_dispatch(
|
||||
self,
|
||||
token_dispatch_input: MoETokenDispatchInput,
|
||||
):
|
||||
hidden_states = token_dispatch_input.hidden_states
|
||||
topk_weights = token_dispatch_input.topk_weights
|
||||
topk_ids = token_dispatch_input.topk_ids
|
||||
expert_map = token_dispatch_input.routing.expert_map
|
||||
apply_router_weight_on_input = token_dispatch_input.routing.apply_router_weight_on_input
|
||||
restore_shape = hidden_states.shape
|
||||
|
||||
num_tokens = hidden_states.shape[:-1].numel()
|
||||
if apply_router_weight_on_input:
|
||||
assert topk_weights.dim() == 2, "`topk_weights` should be in shape (num_tokens, topk)"
|
||||
_, topk = topk_weights.shape
|
||||
assert topk == 1, "Only support topk=1 when `apply_router_weight_on_input` is True"
|
||||
hidden_states = hidden_states * topk_weights.to(hidden_states.dtype)
|
||||
if expert_map is not None:
|
||||
mask = expert_map[topk_ids] != -1
|
||||
topk_weights = topk_weights * mask
|
||||
first_expert_idx = get_ep_group().rank_in_group * self.num_experts_local
|
||||
last_expert_idx = first_expert_idx + self.num_experts_local
|
||||
else:
|
||||
first_expert_idx = 0
|
||||
last_expert_idx = self.num_experts_local
|
||||
|
||||
assert hidden_states.shape[-1] % 16 == 0, (
|
||||
f"The last dim of hidden_states {hidden_states.shape[-1]} should be aligned with 16."
|
||||
)
|
||||
sorted_hidden_states, expanded_row_idx, expert_tokens, _ = torch_npu.npu_moe_init_routing_v2(
|
||||
hidden_states,
|
||||
topk_ids,
|
||||
active_num=num_tokens * self.top_k,
|
||||
expert_num=self.num_experts_local,
|
||||
drop_pad_mode=0,
|
||||
active_expert_range=[first_expert_idx, last_expert_idx],
|
||||
quant_mode=-1,
|
||||
row_idx_type=0,
|
||||
)
|
||||
expert_tokens = expert_tokens.to(torch.int64)
|
||||
group_list_type = 0 # `cumsum` mode
|
||||
|
||||
return MoETokenDispatchOutput(
|
||||
hidden_states=sorted_hidden_states,
|
||||
group_list=expert_tokens,
|
||||
group_list_type=group_list_type,
|
||||
combine_metadata=MoEAllGatherCombineMetadata(
|
||||
topk_weights=topk_weights,
|
||||
expanded_row_idx=expanded_row_idx,
|
||||
restore_shape=restore_shape,
|
||||
),
|
||||
)
|
||||
Reference in New Issue
Block a user