@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 995bbf38f1

Fork 0

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 8, 16:30 UTC

[Bugfix] Fix shared expert input for latent MoE in EP+DP (Nemotron-H) (#34087)

Signed-off-by: Tomer Natan <tbarnatan@nvidia.com> Co-authored-by: Cursor <cursoragent@cursor.com>

TomerBN-Nvidia committed Feb 9, 2026 at 18:44 UTC 995bbf38f114a0e1bd7e34d6fd92d255ac2efca7

6 files changed +30 -3

vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py

+1 -1

@@ -139,7 +139,7 @@ class FlashInferExperts(mk.FusedMoEPermuteExpertsUnpermute):

139 # work with SP. This will be removed in follow up after we get

140 # rid of the FlashInfer specific P/F function.

141 # TODO: the per-tensor fp8 kernels don't work with MNNVL FI A2As.

142 - return not moe_parallel_config.is_sequence_parallel

142 + return True

143

144 @staticmethod

145 def activation_format() -> mk.FusedMoEActivationFormat:

vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py

@@ -101,4 +101,5 @@ class FusedMoEModularMethod(FusedMoEMethodBase, CustomOp):

101 global_num_experts=layer.global_num_experts,

102 apply_router_weight_on_input=layer.apply_router_weight_on_input,

103 expert_map=None if self.disable_expert_map else layer.expert_map,

104 + shared_experts_input=layer._get_shared_experts_input(x),

105 )

vllm/model_executor/layers/fused_moe/modular_kernel.py

+22 -2

@@ -1228,13 +1228,28 @@ class FusedMoEModularKernel(torch.nn.Module):

1228 topk_weights: torch.Tensor,

1229 topk_ids: torch.Tensor,

1230 apply_router_weight_on_input: bool,

1231 + shared_experts_input: torch.Tensor | None = None,

1232 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:

1233 """

1234 The _finalize method is a wrapper around self.prepare_finalize.finalize

1235 that handles DBO, async and shared expert overlap.

1236 +

1237 + Args:

1238 + shared_experts_input: Optional separate input for shared experts.

1239 + When latent MoE is used, hidden_states is the latent-projected

1240 + tensor (smaller dimension) used by routed experts, while

1241 + shared_experts_input is the original hidden_states (full

1242 + dimension) needed by the shared expert MLP.

1243 """

1244 shared_output: torch.Tensor | None = None

1245

1246 + # For latent MoE: shared experts need the original hidden_states

1247 + # (full hidden_size), not the latent-projected version used by

1248 + # routed experts.

1249 + se_hidden_states = (

1250 + shared_experts_input if shared_experts_input is not None else hidden_states

1251 + )

1252 +

1253 if not self.prepare_finalize.supports_async():

1254 assert not dbo_enabled()

1255

@@ -1247,7 +1262,7 @@ class FusedMoEModularKernel(torch.nn.Module):

1262 self.fused_experts.finalize_weight_and_reduce_impl(),

1263 )

1264 if self.shared_experts is not None:

1250 - shared_output = self.shared_experts(hidden_states)

1265 + shared_output = self.shared_experts(se_hidden_states)

1266 else:

1267 finalize_ret = self.prepare_finalize.finalize_async(

1268 output,

@@ -1258,7 +1273,7 @@ class FusedMoEModularKernel(torch.nn.Module):

1273 self.fused_experts.finalize_weight_and_reduce_impl(),

1274 )

1275 if self.shared_experts is not None:

1261 - shared_output = self.shared_experts(hidden_states)

1276 + shared_output = self.shared_experts(se_hidden_states)

1277

1278 # TODO(lucas): refactor this in the alternative schedules followup

1279 # currently unpack if we have hook + receiver pair or just

@@ -1298,6 +1313,7 @@ class FusedMoEModularKernel(torch.nn.Module):

1313 global_num_experts: int = -1,

1314 expert_map: torch.Tensor | None = None,

1315 apply_router_weight_on_input: bool = False,

1316 + shared_experts_input: torch.Tensor | None = None,

1317 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:

1318 """

1319 This function computes a Mixture of Experts (MoE) layer using two sets

@@ -1320,6 +1336,9 @@ class FusedMoEModularKernel(torch.nn.Module):

1336 - apply_router_weight_on_input (bool): When true, the topk weights are

1337 applied directly on the inputs. This is only applicable when topk is

1338 1.

1339 + - shared_experts_input (Optional[torch.Tensor]): Optional separate

1340 + input for shared experts. For latent MoE, this is the original

1341 + hidden_states before latent projection.

1342

1343 Returns:

1344 - torch.Tensor: The output tensor after applying the MoE layer.

@@ -1368,4 +1387,5 @@ class FusedMoEModularKernel(torch.nn.Module):

1387 topk_weights,

1388 topk_ids,

1389 apply_router_weight_on_input,

1390 + shared_experts_input=shared_experts_input,

1391 )

vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py

@@ -361,6 +361,7 @@ class CompressedTensorsW4A4Mxfp4MoEMethod(CompressedTensorsMoEMethod):

361 global_num_experts=layer.global_num_experts,

362 expert_map=layer.expert_map,

363 apply_router_weight_on_input=layer.apply_router_weight_on_input,

364 + shared_experts_input=layer._get_shared_experts_input(x),

365 )

366

367

@@ -672,6 +673,7 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):

673 global_num_experts=layer.global_num_experts,

674 expert_map=layer.expert_map,

675 apply_router_weight_on_input=layer.apply_router_weight_on_input,

676 + shared_experts_input=layer._get_shared_experts_input(x),

677 )

678

679

@@ -1077,6 +1079,7 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):

1079 # https://github.com/vllm-project/vllm/commit/84166fee9770e6fba71a96978b3e7d149392fb28 # noqa: E501

1080 expert_map=layer.expert_map,

1081 apply_router_weight_on_input=layer.apply_router_weight_on_input,

1082 + shared_experts_input=layer._get_shared_experts_input(x),

1083 )

1084

1085 @property

vllm/model_executor/layers/quantization/fp8.py

@@ -1023,6 +1023,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):

1023 global_num_experts=layer.global_num_experts,

1024 expert_map=layer.expert_map,

1025 apply_router_weight_on_input=layer.apply_router_weight_on_input,

1026 + shared_experts_input=layer._get_shared_experts_input(x),

1027 )

1028

1029

vllm/model_executor/layers/quantization/modelopt.py

@@ -980,6 +980,7 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):

980 global_num_experts=layer.global_num_experts,

981 expert_map=layer.expert_map,

982 apply_router_weight_on_input=layer.apply_router_weight_on_input,

983 + shared_experts_input=layer._get_shared_experts_input(x),

984 )

985

986

@@ -1550,6 +1551,7 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):

1551 global_num_experts=layer.global_num_experts,

1552 expert_map=layer.expert_map,

1553 apply_router_weight_on_input=layer.apply_router_weight_on_input,

1554 + shared_experts_input=layer._get_shared_experts_input(x),

1555 )

1556

1557