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 14, 05:30 UTC
[Bugfix] Fix dp_chunking enablement logic in FusedMoE layer (#27220)
Signed-off-by: Alexander Matveev <amatveev@redhat.com>
Alexander Matveev committed
Oct 23, 2025 at 12:03 UTC
9ef3d5b8759aa5c0731a44ab7c77fc0bc84b2c4e
1 file changed
+32
-26
vllm/model_executor/layers/fused_moe/layer.py
+32
-26
@@ -1102,6 +1102,7 @@ class FusedMoE(CustomOp):
1102
self.params_dtype = params_dtype
1103
1104
vllm_config = get_current_vllm_config()
1105
+ self.vllm_config = vllm_config
1106
1107
# FIXME (varun): We should have a better way of inferring the activation
1108
# datatype. This works for now as the tensor datatype entering the MoE
@@ -1342,26 +1343,6 @@ class FusedMoE(CustomOp):
1343
self.batched_hidden_states: torch.Tensor | None = None
1344
self.batched_router_logits: torch.Tensor | None = None
1345
1345
- if self.use_dp_chunking:
1346
- states_shape: tuple[int, ...]
1347
- logits_shape: tuple[int, ...]
1348
-
1349
- # Note here we use `num_experts` which is logical expert count
1350
- if vllm_config.parallel_config.enable_dbo:
1351
- states_shape = (2, moe.max_num_tokens, self.hidden_size)
1352
- logits_shape = (2, moe.max_num_tokens, num_experts)
1353
- else:
1354
- states_shape = (moe.max_num_tokens, self.hidden_size)
1355
- logits_shape = (moe.max_num_tokens, num_experts)
1356
-
1357
- self.batched_hidden_states = torch.zeros(
1358
- states_shape, dtype=moe.in_dtype, device=torch.cuda.current_device()
1359
- )
1360
-
1361
- self.batched_router_logits = torch.zeros(
1362
- logits_shape, dtype=moe.in_dtype, device=torch.cuda.current_device()
1363
- )
1364
-
1346
@property
1347
def shared_experts(self) -> torch.nn.Module | None:
1348
return None
@@ -1420,8 +1401,6 @@ class FusedMoE(CustomOp):
1401
1402
@property
1403
def use_dp_chunking(self) -> bool:
1423
- # Route to the chunked forward path using the FlashInfer Cutlass kernel
1424
- # only when data parallelism (DP) is enabled.
1404
return (
1405
self.moe_parallel_config.use_pplx_kernels
1406
or self.moe_parallel_config.use_deepep_ll_kernels
@@ -1988,12 +1967,40 @@ class FusedMoE(CustomOp):
1967
self.logical_to_physical_map = logical_to_physical_map[moe_layer_idx]
1968
self.logical_replica_count = logical_replica_count[moe_layer_idx]
1969
1991
- def ensure_moe_quant_config(self):
1970
+ def ensure_moe_quant_config_init(self):
1971
if self.quant_method.moe_quant_config is None:
1972
self.quant_method.moe_quant_config = (
1973
self.quant_method.get_fused_moe_quant_config(self)
1974
)
1975
1976
+ if self.moe_quant_config is None:
1977
+ self.moe_quant_config = self.quant_method.moe_quant_config
1978
+
1979
+ def ensure_dp_chunking_init(self):
1980
+ if not self.use_dp_chunking or self.batched_hidden_states is not None:
1981
+ return
1982
+
1983
+ states_shape: tuple[int, ...]
1984
+ logits_shape: tuple[int, ...]
1985
+
1986
+ moe = self.moe_config
1987
+
1988
+ # Note here we use `num_experts` which is logical expert count
1989
+ if self.vllm_config.parallel_config.enable_dbo:
1990
+ states_shape = (2, moe.max_num_tokens, self.hidden_size)
1991
+ logits_shape = (2, moe.max_num_tokens, moe.num_experts)
1992
+ else:
1993
+ states_shape = (moe.max_num_tokens, self.hidden_size)
1994
+ logits_shape = (moe.max_num_tokens, moe.num_experts)
1995
+
1996
+ self.batched_hidden_states = torch.zeros(
1997
+ states_shape, dtype=moe.in_dtype, device=torch.cuda.current_device()
1998
+ )
1999
+
2000
+ self.batched_router_logits = torch.zeros(
2001
+ logits_shape, dtype=moe.in_dtype, device=torch.cuda.current_device()
2002
+ )
2003
+
2004
@staticmethod
2005
def select_experts(
2006
hidden_states: torch.Tensor,
@@ -2224,8 +2231,6 @@ class FusedMoE(CustomOp):
2231
assert self.batched_hidden_states.size(-1) == full_hidden_states.size(-1)
2232
assert self.batched_router_logits.size(-1) == full_router_logits.size(-1)
2233
2227
- self.ensure_moe_quant_config()
2228
-
2234
full_fused_final_hidden_states = torch.empty_like(full_hidden_states)
2235
if self.shared_experts is not None:
2236
full_shared_final_hidden_states = torch.empty_like(full_hidden_states)
@@ -2383,7 +2388,8 @@ class FusedMoE(CustomOp):
2388
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
2389
assert self.quant_method is not None
2390
2386
- self.ensure_moe_quant_config()
2391
+ self.ensure_moe_quant_config_init()
2392
+ self.ensure_dp_chunking_init()
2393
2394
has_separate_shared_experts = (
2395
not isinstance(self.quant_method.fused_experts, FusedMoEModularKernel)