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 11, 01:15 UTC
[MoE Refactor][9/N] Use modular kernel for unquantized Triton MoE (#31052)
Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
Yongye Zhu committed
Dec 22, 2025 at 09:34 UTC
7b926e890195f026b60f36506503a80afc583b33
2 files changed
+22
-7
tests/kernels/moe/test_moe.py
+7
@@ -60,6 +60,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import quantize_w
60
from vllm.model_executor.models.mixtral import MixtralMoE
61
from vllm.platforms import current_platform
62
from vllm.scalar_type import ScalarType, scalar_types
63
+from vllm.v1.worker.workspace import init_workspace_manager
64
65
NUM_EXPERTS = [8, 64, 192]
66
EP_SIZE = [1, 4]
@@ -487,6 +488,7 @@ def test_mixtral_moe(
488
monkeypatch.setenv("MASTER_ADDR", "localhost")
489
monkeypatch.setenv("MASTER_PORT", "12345")
490
init_distributed_environment()
491
+ init_workspace_manager(torch.cuda.current_device())
492
493
# Instantiate our and huggingface's MoE blocks
494
vllm_config.compilation_config.static_forward_context = dict()
@@ -533,6 +535,11 @@ def test_mixtral_moe(
535
torch.cuda.synchronize()
536
torch.cuda.empty_cache()
537
538
+ # FIXME (zyongye) fix this after we move self.kernel
539
+ # assignment in FusedMoE.__init__
540
+
541
+ vllm_moe.experts.quant_method.process_weights_after_loading(vllm_moe.experts)
542
+
543
# Run forward passes for both MoE blocks
544
hf_states, _ = hf_moe.forward(hf_inputs)
545
vllm_states = vllm_moe.forward(vllm_inputs)
vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py
+15
-7
@@ -6,6 +6,7 @@ import torch
6
import torch.nn.functional as F
7
8
import vllm.envs as envs
9
+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
@@ -23,6 +24,9 @@ from vllm.model_executor.layers.fused_moe.modular_kernel import (
24
FusedMoEPermuteExpertsUnpermute,
25
FusedMoEPrepareAndFinalize,
26
)
27
+from vllm.model_executor.layers.fused_moe.prepare_finalize import (
28
+ MoEPrepareAndFinalizeNoEP,
29
+)
30
from vllm.model_executor.utils import set_weight_attrs
31
from vllm.platforms import current_platform
32
from vllm.platforms.interface import CpuArchEnum
@@ -30,9 +34,9 @@ from vllm.utils.flashinfer import has_flashinfer_cutlass_fused_moe
34
35
if current_platform.is_cuda_alike():
36
from .fused_batched_moe import BatchedTritonExperts
33
- from .fused_moe import TritonExperts, fused_experts
37
+ from .fused_moe import TritonExperts
38
else:
35
- fused_experts = None # type: ignore
39
+ TritonExperts = None # type: ignore
40
41
if current_platform.is_tpu():
42
from .moe_pallas import fused_moe as fused_moe_pallas
@@ -265,6 +269,13 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
269
layer.cpu_fused_moe = cpu_fused_moe.CPUFusedMOE(layer)
270
else:
271
layer.cpu_fused_moe = cpu_fused_moe.CPUFusedMOE(layer)
272
+ elif current_platform.is_cuda_alike():
273
+ 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
+ )
279
280
def apply(
281
self,
@@ -278,9 +289,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
289
router_logits=router_logits,
290
)
291
281
- def get_fused_moe_quant_config(
282
- self, layer: torch.nn.Module
283
- ) -> FusedMoEQuantConfig | None:
292
+ def get_fused_moe_quant_config(self, layer: torch.nn.Module) -> FusedMoEQuantConfig:
293
if self.moe.has_bias:
294
return biased_moe_quant_config(
295
layer.w13_bias,
@@ -322,7 +331,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
331
apply_router_weight_on_input=layer.apply_router_weight_on_input,
332
)
333
else:
325
- result = fused_experts(
334
+ result = self.kernel(
335
hidden_states=x,
336
w1=layer.w13_weight,
337
w2=layer.w2_weight,
@@ -330,7 +339,6 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
339
topk_ids=topk_ids,
340
inplace=True,
341
activation=layer.activation,
333
- quant_config=self.moe_quant_config,
342
apply_router_weight_on_input=layer.apply_router_weight_on_input,
343
global_num_experts=layer.global_num_experts,
344
expert_map=layer.expert_map,