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, 22:30 UTC
[Bug] Fix vLLM config is not set error (#29999)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Dec 5, 2025 at 16:42 UTC
7b5575fa7dcf76ac86ab8d18501b9cc04f74f6bb
5 files changed
+47
-27
vllm/model_executor/layers/fused_moe/cutlass_moe.py
+2
@@ -460,6 +460,7 @@ 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,
464
) -> torch.Tensor:
465
"""
466
This function computes a a8w8-quantized Mixture of Experts (MoE) layer
@@ -537,6 +538,7 @@ def cutlass_moe_fp8(
538
c_strides2=c_strides2,
539
quant_config=quant_config,
540
),
541
+ parallel_config=parallel_config,
542
)
543
544
return fn(
vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py
+6
@@ -44,6 +44,11 @@ class FusedMoEModularMethod(FusedMoEMethodBase, CustomOp):
44
prepare_finalize: FusedMoEPrepareAndFinalize,
45
shared_experts: torch.nn.Module | None,
46
) -> "FusedMoEModularMethod":
47
+ parallel_config = getattr(
48
+ getattr(moe_layer, "vllm_config", None),
49
+ "parallel_config",
50
+ None,
51
+ )
52
return FusedMoEModularMethod(
53
old_quant_method,
54
FusedMoEModularKernel(
@@ -51,6 +56,7 @@ class FusedMoEModularMethod(FusedMoEMethodBase, CustomOp):
56
old_quant_method.select_gemm_impl(prepare_finalize, moe_layer),
57
shared_experts,
58
getattr(moe_layer, "shared_experts_stream", None),
59
+ parallel_config=parallel_config,
60
),
61
)
62
vllm/model_executor/layers/fused_moe/modular_kernel.py
+30
-27
@@ -10,7 +10,7 @@ from typing import final
10
import torch
11
12
import vllm.envs as envs
13
-from vllm.config import get_current_vllm_config
13
+from vllm.config import ParallelConfig, get_current_vllm_config
14
from vllm.forward_context import get_forward_context, is_forward_context_available
15
from vllm.logger import init_logger
16
from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
@@ -716,6 +716,7 @@ class FusedMoEModularKernel(torch.nn.Module):
716
fused_experts: FusedMoEPermuteExpertsUnpermute,
717
shared_experts: torch.nn.Module | None = None,
718
shared_experts_stream: torch.cuda.Stream | None = None,
719
+ parallel_config: ParallelConfig | None = None,
720
):
721
super().__init__()
722
self.prepare_finalize = prepare_finalize
@@ -723,6 +724,14 @@ class FusedMoEModularKernel(torch.nn.Module):
724
self.shared_experts = shared_experts
725
self.shared_experts_stream = shared_experts_stream
726
727
+ # cache whether this worker is using DP+EP
728
+ if parallel_config is None:
729
+ parallel_config = get_current_vllm_config().parallel_config
730
+ self.is_dp_ep = (
731
+ parallel_config.data_parallel_size > 1
732
+ and parallel_config.enable_expert_parallel
733
+ )
734
+
735
self._post_init_setup()
736
assert (
737
prepare_finalize.activation_format == fused_experts.activation_formats[0]
@@ -811,33 +820,27 @@ class FusedMoEModularKernel(torch.nn.Module):
820
is_forward_context_available()
821
and get_forward_context().attn_metadata is None
822
)
814
- if is_profile_run and self.fused_experts.supports_chunking():
815
- parallel_config = get_current_vllm_config().parallel_config
816
- is_dp_ep = (
817
- parallel_config.data_parallel_size > 1
818
- and parallel_config.enable_expert_parallel
819
- )
820
- if is_dp_ep:
821
- max_workspace_13, max_workspace_2, max_fused_out_shape = (
822
- self.fused_experts.workspace_shapes(
823
- envs.VLLM_FUSED_MOE_CHUNK_SIZE,
824
- N,
825
- K,
826
- top_k,
827
- global_num_experts,
828
- local_num_experts,
829
- expert_tokens_meta,
830
- )
831
- )
832
- buffers.workspace13.get(
833
- max_workspace_13, device=device, dtype=workspace_dtype
834
- )
835
- buffers.workspace2.get(
836
- max_workspace_2, device=device, dtype=workspace_dtype
837
- )
838
- buffers.fused_out.get(
839
- max_fused_out_shape, device=device, dtype=workspace_dtype
823
+ if is_profile_run and self.fused_experts.supports_chunking() and self.is_dp_ep:
824
+ max_workspace_13, max_workspace_2, max_fused_out_shape = (
825
+ self.fused_experts.workspace_shapes(
826
+ envs.VLLM_FUSED_MOE_CHUNK_SIZE,
827
+ N,
828
+ K,
829
+ top_k,
830
+ global_num_experts,
831
+ local_num_experts,
832
+ expert_tokens_meta,
833
)
834
+ )
835
+ buffers.workspace13.get(
836
+ max_workspace_13, device=device, dtype=workspace_dtype
837
+ )
838
+ buffers.workspace2.get(
839
+ max_workspace_2, device=device, dtype=workspace_dtype
840
+ )
841
+ buffers.fused_out.get(
842
+ max_fused_out_shape, device=device, dtype=workspace_dtype
843
+ )
844
845
# Get intermediate workspace shapes based off the chunked M size.
846
workspace13_shape, workspace2_shape, _ = self.fused_experts.workspace_shapes(
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py
+3
@@ -1287,6 +1287,9 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
1287
ab_strides2=self.ab_strides2,
1288
c_strides1=self.c_strides1,
1289
c_strides2=self.ab_strides1_c_strides2,
1290
+ parallel_config=getattr(
1291
+ getattr(layer, "vllm_config", None), "parallel_config", None
1292
+ ),
1293
)
1294
1295
else:
vllm/model_executor/layers/quantization/utils/flashinfer_utils.py
+6
@@ -247,6 +247,11 @@ 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
+ )
255
fused_experts = mk.FusedMoEModularKernel(
256
build_flashinfer_fp8_cutlass_moe_prepare_finalize(
257
moe=moe, use_deepseek_fp8_block_scale=use_deepseek_fp8_block_scale
@@ -257,6 +262,7 @@ def flashinfer_cutlass_moe_fp8(
262
out_dtype=hidden_states.dtype,
263
use_deepseek_fp8_block_scale=use_deepseek_fp8_block_scale,
264
),
265
+ parallel_config=parallel_config,
266
)
267
268
return fused_experts(