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, 02:45 UTC
[MoE Refactor][17/N] Apply Refactor to Bf16 (#31827)
Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com>
Yongye Zhu committed
Jan 15, 2026 at 12:53 UTC
31c29257c852c429a655dd35a1c373ccd46f97da
12 files changed
+257
-87
tests/evals/gsm8k/configs/moe-refactor-dp-ep/Qwen3-30B-A3B-BF16-triton.yaml
new
+5
@@ -0,0 +1,5 @@
1
+model_name: "Qwen/Qwen3-30B-A3B"
2
+accuracy_threshold: 0.88
3
+num_questions: 1319
4
+num_fewshot: 5
5
+server_args: "--enforce-eager --max-model-len 8192 --data-parallel-size 2 --enable-expert-parallel"
tests/evals/gsm8k/configs/moe-refactor-dp-ep/config-b200.txt
+1
@@ -8,3 +8,4 @@ Qwen3-30B-A3B-Fp8-AutoFp8-deepgemm.yaml
8
Qwen3-30B-A3B-Fp8-CT-Block-deepgemm-deepep-ht.yaml
9
Qwen3-30B-A3B-Fp8-CT-Block-deepgemm-deepep-ll.yaml
10
Qwen3-30B-A3B-Fp8-CT-Block-deepgemm.yaml
11
+Qwen3-30B-A3B-BF16-triton.yaml
tests/evals/gsm8k/configs/moe-refactor/Llama-4-Scout-BF16-fi-cutlass.yaml
new
+7
@@ -0,0 +1,7 @@
1
+model_name: "meta-llama/Llama-4-Scout-17B-16E-Instruct"
2
+accuracy_threshold: 0.92
3
+num_questions: 1319
4
+num_fewshot: 5
5
+server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2 --enable-expert-parallel"
6
+env:
7
+ VLLM_USE_FLASHINFER_MOE_FP16: "1"
tests/evals/gsm8k/configs/moe-refactor/Llama-4-Scout-BF16-triton.yaml
new
+6
@@ -0,0 +1,6 @@
1
+model_name: "meta-llama/Llama-4-Scout-17B-16E-Instruct"
2
+accuracy_threshold: 0.92
3
+num_questions: 1319
4
+num_fewshot: 5
5
+server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2"
6
+
tests/evals/gsm8k/configs/moe-refactor/Mixtral-8x7B-BF16-fi-cutlass.yaml
new
+7
@@ -0,0 +1,7 @@
1
+model_name: "mistralai/Mixtral-8x7B-v0.1"
2
+accuracy_threshold: 0.58
3
+num_questions: 1319
4
+num_fewshot: 5
5
+server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2 --enable-expert-parallel"
6
+env:
7
+ VLLM_USE_FLASHINFER_MOE_FP16: "1"
tests/evals/gsm8k/configs/moe-refactor/Mixtral-8x7B-BF16-triton.yaml
new
+5
@@ -0,0 +1,5 @@
1
+model_name: "mistralai/Mixtral-8x7B-v0.1"
2
+accuracy_threshold: 0.58
3
+num_questions: 1319
4
+num_fewshot: 5
5
+server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2"
tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-BF16-fi-cutlass.yaml
new
+7
@@ -0,0 +1,7 @@
1
+model_name: "Qwen/Qwen3-30B-A3B"
2
+accuracy_threshold: 0.88
3
+num_questions: 1319
4
+num_fewshot: 5
5
+server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2 --enable-expert-parallel"
6
+env:
7
+ VLLM_USE_FLASHINFER_MOE_FP16: "1"
tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-BF16-triton.yaml
new
+5
@@ -0,0 +1,5 @@
1
+model_name: "Qwen/Qwen3-30B-A3B"
2
+accuracy_threshold: 0.88
3
+num_questions: 1319
4
+num_fewshot: 5
5
+server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2"
tests/evals/gsm8k/configs/moe-refactor/config-b200.txt
+4
@@ -11,3 +11,7 @@ Qwen3-30B-A3B-NvFp4-ModelOpt-marlin.yaml
11
Qwen3-30B-A3B-NvFp4-ModelOpt-fi-trtllm.yaml
12
Qwen3-30B-A3B-NvFp4-ModelOpt-fi-cutlass.yaml
13
Qwen3-30B-A3B-NvFp4-ModelOpt-fi-cutlass-dp-ep.yaml
14
+Llama-4-Scout-BF16-fi-cutlass.yaml
15
+Llama-4-Scout-BF16-triton.yaml
16
+Mixtral-8x7B-BF16-fi-cutlass.yaml
17
+Mixtral-8x7B-BF16-triton.yaml
tests/evals/gsm8k/configs/moe-refactor/config-h100.txt
+2
@@ -11,3 +11,5 @@ Qwen3-30B-A3B-Fp8-CT-Channel-vllm-cutlass.yaml
11
Llama-4-Scout-Fp8-ModelOpt-fi-cutlass.yaml
12
Llama-4-Scout-Fp8-ModelOpt-marlin.yaml
13
Llama-4-Scout-Fp8-ModelOpt-triton.yaml
14
+Qwen3-30B-A3B-BF16-fi-cutlass.yaml
15
+Qwen3-30B-A3B-BF16-triton.yaml
\ No newline at end of file
vllm/model_executor/layers/fused_moe/oracle/unquantized.py
new
+161
@@ -0,0 +1,161 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from enum import Enum
5
+
6
+import torch
7
+from torch.nn import Module
8
+
9
+import vllm.envs as envs
10
+import vllm.model_executor.layers.fused_moe.modular_kernel as mk
11
+from vllm._aiter_ops import rocm_aiter_ops
12
+from vllm.logger import init_logger
13
+from vllm.model_executor.layers.fused_moe.config import (
14
+ FusedMoEConfig,
15
+ FusedMoEQuantConfig,
16
+)
17
+from vllm.model_executor.layers.fused_moe.prepare_finalize import (
18
+ MoEPrepareAndFinalizeNoEP,
19
+)
20
+from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
21
+ swap_w13_to_w31,
22
+)
23
+from vllm.platforms import current_platform
24
+from vllm.utils.flashinfer import has_flashinfer_cutlass_fused_moe
25
+
26
+logger = init_logger(__name__)
27
+
28
+
29
+class UnquantizedMoeBackend(Enum):
30
+ FLASHINFER_CUTLASS = "FlashInfer CUTLASS"
31
+ AITER = "ROCm AITER"
32
+ TRITON = "TRITON"
33
+ CPU = "CPU"
34
+ XPU = "XPU"
35
+
36
+
37
+# NOTE(zyongye): Unsupported backend means backend
38
+# that is not conform with Modular kernel format.
39
+# We will directly call the kernel for those backend
40
+UNSUPPORTED_BACKEND = [
41
+ UnquantizedMoeBackend.CPU,
42
+ UnquantizedMoeBackend.XPU,
43
+]
44
+
45
+
46
+def select_unquantized_moe_backend(
47
+ use_ep: bool,
48
+ use_dp: bool,
49
+) -> UnquantizedMoeBackend:
50
+ """
51
+ Select the primary FP8 MoE backend
52
+ Note: Shape-specific fallbacks may still occur at runtime.
53
+ """
54
+
55
+ def _make_log_backend(backend: UnquantizedMoeBackend):
56
+ return f"Using {backend.value} backend for Unquantized MoE"
57
+
58
+ rocm_aiter_moe_enabled = rocm_aiter_ops.is_fused_moe_enabled()
59
+
60
+ # FlashInfer CUTLASS MoE is only supported on Hopper and later GPUS
61
+ flashinfer_cutlass_moe_enabled = (
62
+ has_flashinfer_cutlass_fused_moe()
63
+ and envs.VLLM_USE_FLASHINFER_MOE_FP16
64
+ and use_ep
65
+ and (not use_dp)
66
+ and current_platform.get_device_capability()[0] >= 9
67
+ )
68
+ if current_platform.is_rocm():
69
+ if rocm_aiter_moe_enabled:
70
+ backend = UnquantizedMoeBackend.AITER
71
+ else:
72
+ backend = UnquantizedMoeBackend.TRITON
73
+ if current_platform.is_cuda():
74
+ if flashinfer_cutlass_moe_enabled:
75
+ backend = UnquantizedMoeBackend.FLASHINFER_CUTLASS
76
+ else:
77
+ if use_ep and (not use_dp):
78
+ logger.info_once(
79
+ "FlashInfer CUTLASS MoE is available for EP"
80
+ " but not enabled, consider setting"
81
+ " VLLM_USE_FLASHINFER_MOE_FP16=1 to enable it.",
82
+ scope="local",
83
+ )
84
+ elif use_dp:
85
+ logger.info_once(
86
+ "FlashInfer CUTLASS MoE is currently not available for DP.",
87
+ scope="local",
88
+ )
89
+ backend = UnquantizedMoeBackend.TRITON
90
+ if current_platform.is_xpu():
91
+ backend = UnquantizedMoeBackend.XPU
92
+ if current_platform.is_cpu():
93
+ backend = UnquantizedMoeBackend.CPU
94
+
95
+ logger.info_once(_make_log_backend(backend), scope="local")
96
+ return backend
97
+
98
+
99
+def convert_to_unquantized_kernel_format(
100
+ unquantized_backend: UnquantizedMoeBackend,
101
+ layer: Module,
102
+ w13_weight: torch.Tensor | None = None,
103
+ w2_weight: torch.Tensor | None = None,
104
+) -> tuple[torch.Tensor, torch.Tensor]:
105
+ if unquantized_backend == UnquantizedMoeBackend.AITER:
106
+ w13_weight, w2_weight = rocm_aiter_ops.shuffle_weights(
107
+ layer.w13_weight.data, layer.w2_weight.data
108
+ )
109
+
110
+ elif unquantized_backend == UnquantizedMoeBackend.FLASHINFER_CUTLASS:
111
+ # Swap halves to arrange as [w3; w1] (kernel expectation)
112
+ w13_weight = swap_w13_to_w31(layer.w13_weight.data)
113
+
114
+ return w13_weight, w2_weight
115
+
116
+
117
+def make_unquantized_moe_kernel(
118
+ layer: torch.nn.Module,
119
+ backend: UnquantizedMoeBackend,
120
+ quant_config: FusedMoEQuantConfig,
121
+ moe_config: FusedMoEConfig,
122
+) -> tuple[mk.FusedMoEModularKernel | None, bool]:
123
+ use_inplace = True
124
+
125
+ if backend in UNSUPPORTED_BACKEND:
126
+ return None, use_inplace
127
+
128
+ if backend == UnquantizedMoeBackend.FLASHINFER_CUTLASS:
129
+ from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import (
130
+ FlashInferExperts,
131
+ )
132
+
133
+ kernel = mk.FusedMoEModularKernel(
134
+ MoEPrepareAndFinalizeNoEP(),
135
+ FlashInferExperts(
136
+ out_dtype=layer.params_dtype,
137
+ quant_config=quant_config,
138
+ tp_rank=moe_config.moe_parallel_config.tp_rank,
139
+ tp_size=moe_config.moe_parallel_config.tp_size,
140
+ ep_rank=moe_config.moe_parallel_config.ep_rank,
141
+ ep_size=moe_config.moe_parallel_config.ep_size,
142
+ ),
143
+ )
144
+ use_inplace = False
145
+ elif backend == UnquantizedMoeBackend.AITER:
146
+ from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
147
+ AiterExperts,
148
+ )
149
+
150
+ kernel = mk.FusedMoEModularKernel(
151
+ MoEPrepareAndFinalizeNoEP(),
152
+ AiterExperts(quant_config),
153
+ )
154
+ elif backend == UnquantizedMoeBackend.TRITON:
155
+ from vllm.model_executor.layers.fused_moe import TritonExperts
156
+
157
+ kernel = mk.FusedMoEModularKernel(
158
+ MoEPrepareAndFinalizeNoEP(),
159
+ TritonExperts(quant_config),
160
+ )
161
+ return kernel, use_inplace
vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py
+47
-87
@@ -4,10 +4,10 @@
4
5
import torch
6
import torch.nn.functional as F
7
+from torch.nn import Module
8
9
import vllm.envs as envs
10
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
10
-from vllm._aiter_ops import rocm_aiter_ops
11
from vllm.logger import init_logger
12
from vllm.model_executor.custom_op import CustomOp
13
from vllm.model_executor.layers.fused_moe.config import (
@@ -16,9 +16,6 @@ from vllm.model_executor.layers.fused_moe.config import (
16
FusedMoEQuantConfig,
17
biased_moe_quant_config,
18
)
19
-from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import (
20
- FlashInferExperts,
21
-)
19
from vllm.model_executor.layers.fused_moe.fused_moe_method_base import (
20
FusedMoEMethodBase,
21
)
@@ -28,19 +25,15 @@ from vllm.model_executor.layers.fused_moe.modular_kernel import (
25
FusedMoEPermuteExpertsUnpermute,
26
FusedMoEPrepareAndFinalize,
27
)
31
-from vllm.model_executor.layers.fused_moe.prepare_finalize import (
32
- MoEPrepareAndFinalizeNoEP,
33
-)
34
-from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
35
- AiterExperts,
36
-)
37
-from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
38
- swap_w13_to_w31,
28
+from vllm.model_executor.layers.fused_moe.oracle.unquantized import (
29
+ UnquantizedMoeBackend,
30
+ convert_to_unquantized_kernel_format,
31
+ make_unquantized_moe_kernel,
32
+ select_unquantized_moe_backend,
33
)
34
from vllm.model_executor.utils import replace_parameter, set_weight_attrs
35
from vllm.platforms import current_platform
36
from vllm.platforms.interface import CpuArchEnum
43
-from vllm.utils.flashinfer import has_flashinfer_cutlass_fused_moe
37
38
if current_platform.is_cuda_alike():
39
from .fused_batched_moe import BatchedTritonExperts
@@ -57,41 +50,13 @@ logger = init_logger(__name__)
50
class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
51
"""MoE method without quantization."""
52
60
- # --8<-- [end:unquantized_fused_moe]
61
-
53
def __init__(self, moe: FusedMoEConfig):
54
super().__init__(moe)
64
-
65
- self.rocm_aiter_moe_enabled = rocm_aiter_ops.is_fused_moe_enabled()
66
-
67
- # FlashInfer CUTLASS MoE is only supported on Hopper and later GPUS
68
- self.flashinfer_cutlass_moe_enabled = (
69
- has_flashinfer_cutlass_fused_moe()
70
- and envs.VLLM_USE_FLASHINFER_MOE_FP16
71
- and self.moe.moe_parallel_config.use_ep
72
- and self.moe.moe_parallel_config.dp_size == 1
73
- and current_platform.get_device_capability()[0] >= 9
55
+ self.unquantized_backend = select_unquantized_moe_backend(
56
+ use_ep=self.moe.moe_parallel_config.use_ep,
57
+ use_dp=self.moe.moe_parallel_config.dp_size > 1,
58
)
75
- if self.flashinfer_cutlass_moe_enabled:
76
- logger.info_once(
77
- "Enabling FlashInfer CUTLASS MoE for UnquantizedFusedMoEMethod"
78
- )
79
- else:
80
- if (
81
- self.moe.moe_parallel_config.use_ep
82
- and self.moe.moe_parallel_config.dp_size == 1
83
- ):
84
- logger.info_once(
85
- "FlashInfer CUTLASS MoE is available for EP"
86
- " but not enabled, consider setting"
87
- " VLLM_USE_FLASHINFER_MOE_FP16=1 to enable it.",
88
- scope="local",
89
- )
90
- elif self.moe.moe_parallel_config.dp_size > 1:
91
- logger.info_once(
92
- "FlashInfer CUTLASS MoE is currently not available for DP.",
93
- scope="local",
94
- )
59
+ self.kernel: mk.FusedMoEModularKernel | None = None
60
61
@property
62
def supports_eplb(self) -> bool:
@@ -105,7 +70,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
70
self,
71
routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
72
) -> FusedMoEPrepareAndFinalize | None:
108
- if self.rocm_aiter_moe_enabled:
73
+ if self.unquantized_backend == UnquantizedMoeBackend.AITER:
74
return None
75
else:
76
return super().maybe_make_prepare_finalize(routing_tables)
@@ -197,6 +162,33 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
162
163
return weight
164
165
+ def _setup_kernel(
166
+ self,
167
+ layer: Module,
168
+ w13: torch.Tensor,
169
+ w2: torch.Tensor,
170
+ ) -> None:
171
+ # Shuffle weights to runtime format.
172
+ w13, w2 = convert_to_unquantized_kernel_format(
173
+ self.unquantized_backend,
174
+ layer=layer,
175
+ w13_weight=w13,
176
+ w2_weight=w2,
177
+ )
178
+ replace_parameter(layer, "w13_weight", w13)
179
+ replace_parameter(layer, "w2_weight", w2)
180
+
181
+ # Setup Modular Kernel for TP Case
182
+ self.moe_quant_config = self.get_fused_moe_quant_config(layer)
183
+ assert self.moe_quant_config is not None
184
+
185
+ self.kernel, self.use_inplace = make_unquantized_moe_kernel(
186
+ layer=layer,
187
+ backend=self.unquantized_backend,
188
+ quant_config=self.moe_quant_config,
189
+ moe_config=self.moe,
190
+ )
191
+
192
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
193
super().process_weights_after_loading(layer)
194
@@ -204,7 +196,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
196
layer.w13_weight.data = self._maybe_pad_weight(layer.w13_weight.data)
197
layer.w2_weight.data = self._maybe_pad_weight(layer.w2_weight.data)
198
207
- if current_platform.is_xpu():
199
+ if self.unquantized_backend == UnquantizedMoeBackend.XPU:
200
import intel_extension_for_pytorch as ipex
201
202
ep_rank_start = self.moe.ep_rank * self.moe.num_local_experts
@@ -214,7 +206,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
206
use_prepack=True,
207
experts_start_id=ep_rank_start,
208
)
217
- elif current_platform.is_cpu():
209
+ elif self.unquantized_backend == UnquantizedMoeBackend.CPU:
210
from vllm.model_executor.layers.fused_moe import cpu_fused_moe
211
212
if current_platform.get_cpu_architecture() == CpuArchEnum.X86:
@@ -246,45 +238,11 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
238
else:
239
layer.cpu_fused_moe = cpu_fused_moe.CPUFusedMOE(layer)
240
elif current_platform.is_cuda_alike():
249
- self.moe_quant_config = self.get_fused_moe_quant_config(layer)
250
- if self.rocm_aiter_moe_enabled:
251
- shuffled_w13, shuffled_w2 = rocm_aiter_ops.shuffle_weights(
252
- layer.w13_weight.data, layer.w2_weight.data
253
- )
254
- replace_parameter(layer, "w13_weight", shuffled_w13)
255
- replace_parameter(layer, "w2_weight", shuffled_w2)
256
-
257
- self.use_inplace = True
258
- self.kernel = mk.FusedMoEModularKernel(
259
- MoEPrepareAndFinalizeNoEP(),
260
- AiterExperts(self.moe_quant_config),
261
- shared_experts=None,
262
- )
263
-
264
- elif self.flashinfer_cutlass_moe_enabled:
265
- self.use_inplace = False
266
- # Swap halves to arrange as [w3; w1] (kernel expectation)
267
- w13_weight = swap_w13_to_w31(layer.w13_weight.data)
268
- replace_parameter(layer, "w13_weight", w13_weight)
269
-
270
- self.kernel = mk.FusedMoEModularKernel(
271
- MoEPrepareAndFinalizeNoEP(),
272
- FlashInferExperts(
273
- out_dtype=layer.params_dtype,
274
- quant_config=self.moe_quant_config,
275
- tp_rank=self.moe.moe_parallel_config.tp_rank,
276
- tp_size=self.moe.moe_parallel_config.tp_size,
277
- ep_rank=self.moe.moe_parallel_config.ep_rank,
278
- ep_size=self.moe.moe_parallel_config.ep_size,
279
- ),
280
- )
281
- else:
282
- self.use_inplace = True
283
- self.kernel = mk.FusedMoEModularKernel(
284
- MoEPrepareAndFinalizeNoEP(),
285
- TritonExperts(self.moe_quant_config),
286
- shared_experts=None,
287
- )
241
+ self._setup_kernel(
242
+ layer=layer,
243
+ w13=layer.w13_weight,
244
+ w2=layer.w2_weight,
245
+ )
246
247
def apply(
248
self,
@@ -316,6 +274,8 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
274
x: torch.Tensor,
275
router_logits: torch.Tensor,
276
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
277
+ assert self.kernel
278
+
279
topk_weights, topk_ids = router.select_experts(
280
hidden_states=x,
281
router_logits=router_logits,