Skip to content

Commit 986ed47

Browse files
committed
refatcor torchair fused_moe 1/N
Signed-off-by: hust17yixuan <303660421@qq.com> Signed-off-by: hust17yixuan <303660421@qq.com> Signed-off-by: hust17yixuan <303660421@qq.com>
1 parent 9554116 commit 986ed47

File tree

3 files changed

+1569
-3
lines changed

3 files changed

+1569
-3
lines changed

vllm_ascend/torchair/models/torchair_deepseek_v2.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -70,7 +70,7 @@
7070
from vllm.sequence import IntermediateTensors
7171

7272
from vllm_ascend.ascend_config import get_ascend_config
73-
from vllm_ascend.ops.fused_moe import AscendFusedMoE
73+
from vllm_ascend.torchair.ops.torchair_fused_moe import TorchairAscendFusedMoE
7474
from vllm_ascend.quantization.quant_config import AscendLinearMethod
7575
from vllm_ascend.quantization.w8a8_dynamic import AscendW8A8DynamicLinearMethod
7676
from vllm_ascend.utils import dispose_tensor, npu_prefetch
@@ -335,7 +335,7 @@ def __init__(
335335
else:
336336
self.gate.e_score_correction_bias = None
337337

338-
self.experts = AscendFusedMoE(
338+
self.experts = TorchairAscendFusedMoE(
339339
num_experts=config.n_routed_experts,
340340
top_k=config.num_experts_per_tok,
341341
hidden_size=config.hidden_size,
@@ -951,7 +951,7 @@ def load_weights(self, weights: Iterable[tuple[str,
951951

952952
# Params for weights, fp8 weight scales, fp8 activation scales
953953
# (param_name, weight_name, expert_id, shard_id)
954-
expert_params_mapping = AscendFusedMoE.make_expert_params_mapping(
954+
expert_params_mapping = TorchairAscendFusedMoE.make_expert_params_mapping(
955955
ckpt_gate_proj_name="gate_proj",
956956
ckpt_down_proj_name="down_proj",
957957
ckpt_up_proj_name="up_proj",

vllm_ascend/torchair/ops/__init__.py

Whitespace-only changes.

0 commit comments

Comments
 (0)