A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 10, 01:30 UTC
[MoE Refactor] Move Test Impl into Test Dirs (#32129)
Signed-off-by: Robert Shaw <rshaw@neuralmagic.com> Co-authored-by: Robert Shaw <rshaw@neuralmagic.com>
Robert Shaw committed
Jan 17, 2026 at 23:16 UTC
4a6af8813fb53ebde21ffadc2bfce94c965e98b1
3 files changed
+58
-64
docs/design/moe_kernel_features.md
-1
@@ -92,7 +92,6 @@ To be used with a particular `FusedMoEPrepareAndFinalize` subclass, MoE kernels
92
| gpt oss triton | standard | N/A | N/A | <sup>5</sup> | Y | Y | [`triton_kernel_fused_experts`][vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe.triton_kernel_fused_experts],</br>[`OAITritonExperts`][vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe.OAITritonExperts] |
93
| marlin | standard,</br>batched | <sup>3</sup> / N/A | <sup>3</sup> / N/A | silu,</br>swigluoai | Y | Y | [`fused_marlin_moe`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.fused_marlin_moe],</br>[`MarlinExperts`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.MarlinExperts],</br>[`BatchedMarlinExperts`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.BatchedMarlinExperts] |
94
| trtllm | standard | mxfp4,</br>nvfp4 | G(16),G(32) | <sup>5</sup> | N | Y | [`TrtLlmGenExperts`][vllm.model_executor.layers.fused_moe.trtllm_moe.TrtLlmGenExperts] |
95
-| iterative | standard | N/A | N/A | silu | N | N | [`fused_moe`][vllm.model_executor.layers.fused_moe.moe_torch_iterative.fused_moe] |
95
| rocm aiter moe | standard | fp8 | G(128),A,T | silu, gelu | Y | N | [`rocm_aiter_fused_experts`][vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe.rocm_aiter_fused_experts] |
96
| cpu_fused_moe | standard | N/A | N/A | silu | N | N | [`CPUFusedMOE`][vllm.model_executor.layers.fused_moe.cpu_fused_moe.CPUFusedMOE] |
97
| naive batched<sup>4</sup> | batched | int8,</br>fp8 | G,A,T | silu, gelu | <sup>6</sup> | Y | [`NaiveBatchedExperts`][vllm.model_executor.layers.fused_moe.fused_batched_moe.NaiveBatchedExperts] |
tests/kernels/moe/test_moe.py
+58
-3
@@ -37,9 +37,6 @@ from vllm.model_executor.layers.fused_moe.fused_moe import (
37
fused_topk,
38
modular_triton_fused_moe,
39
)
40
-from vllm.model_executor.layers.fused_moe.moe_torch_iterative import (
41
- fused_moe as iterative_moe,
42
-)
40
from vllm.model_executor.layers.quantization.utils.marlin_utils import (
41
marlin_permute_bias,
42
)
@@ -61,6 +58,64 @@ from vllm.scalar_type import ScalarType, scalar_types
58
from vllm.utils.torch_utils import set_random_seed
59
from vllm.v1.worker.workspace import init_workspace_manager
60
61
+
62
+def iterative_moe(
63
+ hidden_states: torch.Tensor,
64
+ w1: torch.Tensor,
65
+ w2: torch.Tensor,
66
+ gating_output: torch.Tensor,
67
+ topk: int,
68
+ global_num_experts: int,
69
+ expert_map: torch.Tensor = None,
70
+ renormalize: bool = False,
71
+) -> torch.Tensor:
72
+ """
73
+ Baseline implementation of fused moe.
74
+
75
+ Args:
76
+ hidden_states: [*, hidden_size]
77
+ w1: [num_experts, intermediate_size * 2, hidden_size]
78
+ w2: [num_experts, hidden_size, intermediate_size]
79
+ gating_output: [*, num_experts]
80
+ expert_map: [num_experts]
81
+ """
82
+ orig_shape = hidden_states.shape
83
+ hidden_size = hidden_states.shape[-1]
84
+ num_tokens = hidden_states.shape[:-1].numel()
85
+ num_experts = w1.shape[0]
86
+ intermediate_size = w2.shape[-1]
87
+ dtype = hidden_states.dtype
88
+
89
+ hidden_states = hidden_states.view(num_tokens, hidden_size)
90
+ gating_output = gating_output.view(num_tokens, global_num_experts)
91
+ topk_weights = gating_output.softmax(dim=-1, dtype=torch.float)
92
+ topk_weights, selected_experts = topk_weights.topk(topk, dim=-1)
93
+ if renormalize:
94
+ topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
95
+ topk_weights = topk_weights.to(dtype)
96
+
97
+ if expert_map is not None:
98
+ selected_experts = expert_map[selected_experts]
99
+
100
+ final_hidden_states = None
101
+ for expert_idx in range(num_experts):
102
+ expert_w1 = w1[expert_idx]
103
+ expert_w2 = w2[expert_idx]
104
+ expert_mask = selected_experts == expert_idx
105
+ expert_weights = (topk_weights * expert_mask).sum(dim=-1, keepdim=True)
106
+ x = F.linear(hidden_states, expert_w1)
107
+ gate = F.silu(x[:, :intermediate_size])
108
+ x = x[:, intermediate_size:] * gate
109
+ x = F.linear(x, expert_w2)
110
+ current_hidden_states = x * expert_weights
111
+ if final_hidden_states is None:
112
+ final_hidden_states = current_hidden_states
113
+ else:
114
+ final_hidden_states = final_hidden_states + current_hidden_states
115
+
116
+ return final_hidden_states.view(orig_shape) # type: ignore
117
+
118
+
119
NUM_EXPERTS = [8, 64, 192]
120
NUM_EXPERTS_LARGE = [128, 256]
121
EP_SIZE = [1, 4]
vllm/model_executor/layers/fused_moe/moe_torch_iterative.py
deleted
-60
@@ -1,60 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-import torch
5
-import torch.nn.functional as F
6
-
7
-
8
-def fused_moe(
9
- hidden_states: torch.Tensor,
10
- w1: torch.Tensor,
11
- w2: torch.Tensor,
12
- gating_output: torch.Tensor,
13
- topk: int,
14
- global_num_experts: int,
15
- expert_map: torch.Tensor = None,
16
- renormalize: bool = False,
17
-) -> torch.Tensor:
18
- """
19
- Args:
20
- hidden_states: [*, hidden_size]
21
- w1: [num_experts, intermediate_size * 2, hidden_size]
22
- w2: [num_experts, hidden_size, intermediate_size]
23
- gating_output: [*, num_experts]
24
- expert_map: [num_experts]
25
- """
26
- orig_shape = hidden_states.shape
27
- hidden_size = hidden_states.shape[-1]
28
- num_tokens = hidden_states.shape[:-1].numel()
29
- num_experts = w1.shape[0]
30
- intermediate_size = w2.shape[-1]
31
- dtype = hidden_states.dtype
32
-
33
- hidden_states = hidden_states.view(num_tokens, hidden_size)
34
- gating_output = gating_output.view(num_tokens, global_num_experts)
35
- topk_weights = gating_output.softmax(dim=-1, dtype=torch.float)
36
- topk_weights, selected_experts = topk_weights.topk(topk, dim=-1)
37
- if renormalize:
38
- topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
39
- topk_weights = topk_weights.to(dtype)
40
-
41
- if expert_map is not None:
42
- selected_experts = expert_map[selected_experts]
43
-
44
- final_hidden_states = None
45
- for expert_idx in range(num_experts):
46
- expert_w1 = w1[expert_idx]
47
- expert_w2 = w2[expert_idx]
48
- expert_mask = selected_experts == expert_idx
49
- expert_weights = (topk_weights * expert_mask).sum(dim=-1, keepdim=True)
50
- x = F.linear(hidden_states, expert_w1)
51
- gate = F.silu(x[:, :intermediate_size])
52
- x = x[:, intermediate_size:] * gate
53
- x = F.linear(x, expert_w2)
54
- current_hidden_states = x * expert_weights
55
- if final_hidden_states is None:
56
- final_hidden_states = current_hidden_states
57
- else:
58
- final_hidden_states = final_hidden_states + current_hidden_states
59
-
60
- return final_hidden_states.view(orig_shape) # type: ignore