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, 11:15 UTC
[Feat] Refactor for `parallel_config` in `FusedMoEModularKernel` (#30282)
Signed-off-by: yewentao256 <zhyanwentao@126.com> Signed-off-by: Robert Shaw <robshaw@redhat.com> Co-authored-by: Robert Shaw <robshaw@redhat.com> Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com>
Wentao Ye committed
Dec 14, 2025 at 23:21 UTC
3778673ea81bf5241f40e9c5e90f989bde377acf
8 files changed
+32
-27
tests/kernels/moe/modular_kernel_tools/common.py
+2
-1
@@ -594,7 +594,8 @@ def make_modular_kernel(
594
)
595
596
modular_kernel = mk.FusedMoEModularKernel(
597
- prepare_finalize=prepare_finalize, fused_experts=fused_experts
597
+ prepare_finalize=prepare_finalize,
598
+ fused_experts=fused_experts,
599
)
600
601
return modular_kernel
tests/kernels/moe/test_flashinfer.py
+14
@@ -5,6 +5,7 @@ from dataclasses import dataclass
5
import pytest
6
import torch
7
8
+import vllm.model_executor.layers.fused_moe.modular_kernel as mk
9
from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config
10
from vllm.model_executor.layers.fused_moe.config import (
11
FusedMoEQuantConfig,
@@ -107,6 +108,19 @@ class TestData:
108
layer.w2_input_scale = a2_scale
109
layer.w13_weight_scale = w13_weight_scale
110
layer.w2_weight_scale = w2_weight_scale
111
+ # Setup dummy config.
112
+ layer.moe_parallel_config = mk.FusedMoEParallelConfig(
113
+ tp_size=1,
114
+ pcp_size=1,
115
+ dp_size=1,
116
+ ep_size=1,
117
+ tp_rank=1,
118
+ pcp_rank=1,
119
+ dp_rank=1,
120
+ ep_rank=1,
121
+ use_ep=False,
122
+ all2all_backend="naive",
123
+ )
124
125
register_moe_scaling_factors(layer)
126
vllm/model_executor/layers/fused_moe/cutlass_moe.py
-2
@@ -460,7 +460,6 @@ def cutlass_moe_fp8(
460
expert_map: torch.Tensor | None = None,
461
apply_router_weight_on_input: bool = False,
462
global_num_experts: int = -1,
463
- parallel_config=None,
463
) -> torch.Tensor:
464
"""
465
This function computes a a8w8-quantized Mixture of Experts (MoE) layer
@@ -538,7 +537,6 @@ def cutlass_moe_fp8(
537
c_strides2=c_strides2,
538
quant_config=quant_config,
539
),
541
- parallel_config=parallel_config,
540
)
541
542
return fn(
vllm/model_executor/layers/fused_moe/deep_gemm_moe.py
+1
-1
@@ -293,7 +293,7 @@ def deep_gemm_moe_fp8(
293
expert_map: torch.Tensor | None = None,
294
a1_scale: torch.Tensor | None = None,
295
a2_scale: torch.Tensor | None = None,
296
- apply_router_weight_on_input=False,
296
+ apply_router_weight_on_input: bool = False,
297
) -> torch.Tensor:
298
"""
299
This function computes a a8w8-quantized Mixture of Experts (MoE) layer
vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py
+1
-6
@@ -43,11 +43,6 @@ class FusedMoEModularMethod(FusedMoEMethodBase, CustomOp):
43
prepare_finalize: FusedMoEPrepareAndFinalize,
44
shared_experts: torch.nn.Module | None,
45
) -> "FusedMoEModularMethod":
46
- parallel_config = getattr(
47
- getattr(moe_layer, "vllm_config", None),
48
- "parallel_config",
49
- None,
50
- )
46
return FusedMoEModularMethod(
47
old_quant_method,
48
FusedMoEModularKernel(
@@ -55,7 +50,7 @@ class FusedMoEModularMethod(FusedMoEMethodBase, CustomOp):
50
old_quant_method.select_gemm_impl(prepare_finalize, moe_layer),
51
shared_experts,
52
getattr(moe_layer, "shared_experts_stream", None),
58
- parallel_config=parallel_config,
53
+ moe_parallel_config=moe_layer.moe_parallel_config,
54
),
55
)
56
vllm/model_executor/layers/fused_moe/modular_kernel.py
+13
-8
@@ -10,10 +10,12 @@ from typing import final
10
import torch
11
12
import vllm.envs as envs
13
-from vllm.config import ParallelConfig, get_current_vllm_config
13
from vllm.forward_context import get_forward_context, is_forward_context_available
14
from vllm.logger import init_logger
16
-from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
15
+from vllm.model_executor.layers.fused_moe.config import (
16
+ FusedMoEParallelConfig,
17
+ FusedMoEQuantConfig,
18
+)
19
from vllm.model_executor.layers.fused_moe.utils import (
20
_resize_cache,
21
count_expert_num_tokens,
@@ -681,7 +683,7 @@ class FusedMoEModularKernel(torch.nn.Module):
683
fused_experts: FusedMoEPermuteExpertsUnpermute,
684
shared_experts: torch.nn.Module | None = None,
685
shared_experts_stream: torch.cuda.Stream | None = None,
684
- parallel_config: ParallelConfig | None = None,
686
+ moe_parallel_config: FusedMoEParallelConfig | None = None,
687
):
688
super().__init__()
689
self.prepare_finalize = prepare_finalize
@@ -689,12 +691,15 @@ class FusedMoEModularKernel(torch.nn.Module):
691
self.shared_experts = shared_experts
692
self.shared_experts_stream = shared_experts_stream
693
692
- # cache whether this worker is using DP+EP
693
- if parallel_config is None:
694
- parallel_config = get_current_vllm_config().parallel_config
694
+ # prefer an explicit FusedMoEParallelConfig when available (from
695
+ # FusedMoE layers / tests).
696
+ # if not provided, assume this kernel is
697
+ # running in a non-DP+EP context
698
+ self.moe_parallel_config: FusedMoEParallelConfig | None = moe_parallel_config
699
self.is_dp_ep = (
696
- parallel_config.data_parallel_size > 1
697
- and parallel_config.enable_expert_parallel
700
+ moe_parallel_config is not None
701
+ and moe_parallel_config.dp_size > 1
702
+ and moe_parallel_config.use_ep
703
)
704
705
self._post_init_setup()
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py
-3
@@ -1266,9 +1266,6 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
1266
ab_strides2=self.ab_strides2,
1267
c_strides1=self.c_strides1,
1268
c_strides2=self.ab_strides1_c_strides2,
1269
- parallel_config=getattr(
1270
- getattr(layer, "vllm_config", None), "parallel_config", None
1271
- ),
1269
)
1270
1271
else:
vllm/model_executor/layers/quantization/utils/flashinfer_utils.py
+1
-6
@@ -247,11 +247,6 @@ def flashinfer_cutlass_moe_fp8(
247
assert quant_config is not None
248
249
# Construct modular kernel with block-scale support when requested.
250
- parallel_config = getattr(
251
- getattr(layer, "vllm_config", None),
252
- "parallel_config",
253
- None,
254
- )
250
fused_experts = mk.FusedMoEModularKernel(
251
build_flashinfer_fp8_cutlass_moe_prepare_finalize(
252
moe=moe, use_deepseek_fp8_block_scale=use_deepseek_fp8_block_scale
@@ -262,7 +257,7 @@ def flashinfer_cutlass_moe_fp8(
257
out_dtype=hidden_states.dtype,
258
use_deepseek_fp8_block_scale=use_deepseek_fp8_block_scale,
259
),
265
- parallel_config=parallel_config,
260
+ moe_parallel_config=layer.moe_parallel_config,
261
)
262
263
return fused_experts(