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, 19:00 UTC
[MoE Refactor] Explicit construct mk for flashinfer bf16 kernel (#31504)
Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com>
Yongye Zhu committed
Jan 2, 2026 at 13:54 UTC
a3f2f40947cfcf9446ab967cb90d9ebc92833f30
1 file changed
+32
-36
vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py
+32
-36
@@ -16,6 +16,9 @@ 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
+)
22
from vllm.model_executor.layers.fused_moe.fused_moe_method_base import (
23
FusedMoEMethodBase,
24
)
@@ -27,7 +30,10 @@ from vllm.model_executor.layers.fused_moe.modular_kernel import (
30
from vllm.model_executor.layers.fused_moe.prepare_finalize import (
31
MoEPrepareAndFinalizeNoEP,
32
)
30
-from vllm.model_executor.utils import set_weight_attrs
33
+from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
34
+ swap_w13_to_w31,
35
+)
36
+from vllm.model_executor.utils import replace_parameter, set_weight_attrs
37
from vllm.platforms import current_platform
38
from vllm.platforms.interface import CpuArchEnum
39
from vllm.utils.flashinfer import has_flashinfer_cutlass_fused_moe
@@ -73,18 +79,6 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
79
logger.info_once(
80
"Enabling FlashInfer CUTLASS MoE for UnquantizedFusedMoEMethod"
81
)
76
- from functools import partial
77
-
78
- from .flashinfer_cutlass_moe import flashinfer_cutlass_moe
79
-
80
- self.flashinfer_cutlass_moe = partial(
81
- flashinfer_cutlass_moe,
82
- quant_config=FUSED_MOE_UNQUANTIZED_CONFIG,
83
- tp_rank=self.moe.moe_parallel_config.tp_rank,
84
- tp_size=self.moe.moe_parallel_config.tp_size,
85
- ep_rank=self.moe.moe_parallel_config.ep_rank,
86
- ep_size=self.moe.moe_parallel_config.ep_size,
87
- )
82
else:
83
if (
84
self.moe.moe_parallel_config.use_ep
@@ -101,7 +95,6 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
95
"FlashInfer CUTLASS MoE is currently not available for DP.",
96
scope="local",
97
)
104
- self.flashinfer_cutlass_moe = None # type: ignore
98
99
@property
100
def supports_eplb(self) -> bool:
@@ -222,12 +215,6 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
215
layer.w13_weight.data = shuffled_w13
216
layer.w2_weight.data = shuffled_w2
217
225
- if self.flashinfer_cutlass_moe_enabled:
226
- # Swap halves to arrange as [w3; w1] (kernel expectation)
227
- w1_w, w3_w = torch.chunk(layer.w13_weight.data, 2, dim=1)
228
- w13_weight_swapped = torch.cat([w3_w, w1_w], dim=1)
229
- layer.w13_weight.data = w13_weight_swapped.contiguous()
230
-
218
if current_platform.is_xpu():
219
import intel_extension_for_pytorch as ipex
220
@@ -271,11 +258,30 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
258
layer.cpu_fused_moe = cpu_fused_moe.CPUFusedMOE(layer)
259
elif current_platform.is_cuda_alike():
260
self.moe_quant_config = self.get_fused_moe_quant_config(layer)
274
- self.kernel = mk.FusedMoEModularKernel(
275
- MoEPrepareAndFinalizeNoEP(),
276
- TritonExperts(self.moe_quant_config),
277
- shared_experts=None,
278
- )
261
+ if self.flashinfer_cutlass_moe_enabled:
262
+ self.use_inplace = False
263
+ # Swap halves to arrange as [w3; w1] (kernel expectation)
264
+ w13_weight = swap_w13_to_w31(layer.w13_weight.data)
265
+ replace_parameter(layer, "w13_weight", w13_weight)
266
+
267
+ self.kernel = mk.FusedMoEModularKernel(
268
+ MoEPrepareAndFinalizeNoEP(),
269
+ FlashInferExperts(
270
+ out_dtype=layer.params_dtype,
271
+ quant_config=self.moe_quant_config,
272
+ tp_rank=self.moe.moe_parallel_config.tp_rank,
273
+ tp_size=self.moe.moe_parallel_config.tp_size,
274
+ ep_rank=self.moe.moe_parallel_config.ep_rank,
275
+ ep_size=self.moe.moe_parallel_config.ep_size,
276
+ ),
277
+ )
278
+ else:
279
+ self.use_inplace = True
280
+ self.kernel = mk.FusedMoEModularKernel(
281
+ MoEPrepareAndFinalizeNoEP(),
282
+ TritonExperts(self.moe_quant_config),
283
+ shared_experts=None,
284
+ )
285
286
def apply(
287
self,
@@ -320,16 +326,6 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
326
activation=layer.activation,
327
apply_router_weight_on_input=layer.apply_router_weight_on_input,
328
)
323
- elif self.flashinfer_cutlass_moe_enabled:
324
- return self.flashinfer_cutlass_moe(
325
- hidden_states=x,
326
- w1=layer.w13_weight,
327
- w2=layer.w2_weight,
328
- topk_weights=topk_weights,
329
- topk_ids=topk_ids,
330
- activation=layer.activation,
331
- apply_router_weight_on_input=layer.apply_router_weight_on_input,
332
- )
329
else:
330
result = self.kernel(
331
hidden_states=x,
@@ -337,7 +333,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
333
w2=layer.w2_weight,
334
topk_weights=topk_weights,
335
topk_ids=topk_ids,
340
- inplace=True,
336
+ inplace=self.use_inplace,
337
activation=layer.activation,
338
apply_router_weight_on_input=layer.apply_router_weight_on_input,
339
global_num_experts=layer.global_num_experts,