Skip to content

Commit 2201406

Browse files
Pr0Wh1teGiveeweijinqian0
authored andcommitted
[main] [refactor] refactor common_fused_moe.py (vllm-project#2706)
### What this PR does / why we need it? 1. Move prepare/finalize operation from moe_comm_method to /ops/moe/fused_moe_prepare_and_finalize 2. Adapt to token_dispatcher in moe_comm_method 3. Move moe_comm_method/experts_selector/token_dispatcher/fused_moe_prepare_and_finalize to /ops/moe ### Does this PR introduce _any_ user-facing change? no ### How was this patch tested? e2e & ut - vLLM version: v0.10.1.1 - vLLM main: vllm-project/vllm@f4962a6 Signed-off-by: weichen <calvin_zhu0210@outlook.com> Signed-off-by: Pr0Wh1teGivee <calvin_zhu0210@outlook.com> Co-authored-by: weijinqian0 <12153182+weijinqian0@users.noreply.github.com>
1 parent 1773ac8 commit 2201406

21 files changed

+1052
-932
lines changed

tests/e2e/singlecard/ops/test_fused_moe.py

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -28,9 +28,8 @@
2828
import torch_npu
2929
from vllm.model_executor.layers.activation import SiluAndMul
3030

31-
from vllm_ascend.ops.layers.experts_selector import select_experts
32-
from vllm_ascend.ops.moe_dispatcher.token_dispatcher import \
33-
TokenDispatcherWithAllGather
31+
from vllm_ascend.ops.moe.experts_selector import select_experts
32+
from vllm_ascend.ops.moe.token_dispatcher import TokenDispatcherWithAllGather
3433

3534
NUM_EXPERTS = [8, 64]
3635
EP_SIZE = [1]
@@ -209,7 +208,7 @@ def test_select_experts(
209208
dtype=torch.int32)
210209
custom_routing_function.return_value = (mock_weights, mock_ids)
211210

212-
with patch("vllm_ascend.ops.layers.experts_selector._native_grouped_topk"
211+
with patch("vllm_ascend.ops.moe.experts_selector._native_grouped_topk"
213212
) as mock_native_grouped_topk:
214213
mock_native_grouped_topk.side_effect = lambda x, num_groups, k: torch.randn_like(
215214
x)

tests/e2e/singlecard/ops/test_moe_comm.py

Lines changed: 0 additions & 175 deletions
This file was deleted.

0 commit comments

Comments
 (0)