Move expert_mask_gpu from FusedMoE layer to StandardDispatcher (#23585)

Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

python/sglang/srt/layers/moe/fused_moe_triton/layer.py

@@ -213,8 +213,6 @@ class FusedMoE(torch.nn.Module):
         self.num_local_experts = self._num_local_routed + num_fused_shared_experts
         self._has_fused_shared = num_fused_shared_experts > 0
 
-        self.expert_mask_gpu = None
-
         assert intermediate_size % self.moe_tp_size == 0
         self.intermediate_size_per_partition = intermediate_size // self.moe_tp_size
         self.reduce_results = reduce_results
@@ -1010,16 +1008,6 @@ class FusedMoE(torch.nn.Module):
             hidden_states=hidden_states, topk_output=topk_output
         )
 
-        if _use_aiter and self.dispatcher.local_expert_mapping is not None:
-            self.expert_mask_gpu = (
-                (
-                    (self.dispatcher.local_expert_mapping >= 0)
-                    & (self.dispatcher.local_expert_mapping < self.num_local_experts)
-                )
-                .to(torch.int32)
-                .to(device="cuda")
-            )
-
         combine_input = self.run_moe_core(
             dispatch_output=dispatch_output,
         )

python/sglang/srt/layers/moe/token_dispatcher/standard.py

@@ -99,12 +99,14 @@ class StandardDispatcher(BaseDispatcher):
             get_moe_runner_backend().is_flashinfer_trtllm_routed()
         )
         self.num_experts = moe_runner_config.num_experts
+        self.num_local_experts = moe_runner_config.num_local_experts
         self.num_local_shared_experts = moe_runner_config.num_fused_shared_experts
         self.num_local_routed_experts = (
-            moe_runner_config.num_local_experts - self.num_local_shared_experts
+            self.num_local_experts - self.num_local_shared_experts
         )
         self.moe_ep_rank = get_moe_expert_parallel_rank()
         self.local_expert_mapping = None
+        self.expert_mask_gpu = None
 
     def dispatch(
         self, hidden_states: torch.Tensor, topk_output: TopKOutput
@@ -187,7 +189,17 @@ class StandardDispatcher(BaseDispatcher):
                         )
                     )
 
-        if self.local_expert_mapping is not None and not _use_aiter:
+        if self.local_expert_mapping is not None:
+            if _use_aiter:
+                self.expert_mask_gpu = (
+                    (
+                        (self.local_expert_mapping >= 0)
+                        & (self.local_expert_mapping < self.num_local_experts)
+                    )
+                    .to(torch.int32)
+                    .to(device="cuda")
+                )
+            else:
                 if TopKOutputChecker.format_is_standard(topk_output):
                     topk_output = topk_output._replace(
                         topk_ids=self.local_expert_mapping[topk_output.topk_ids]

python/sglang/srt/layers/quantization/fp8.py

@@ -1821,7 +1821,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
                         if activation == "silu"
                         else ActivationType.Gelu
                     ),
-                    expert_mask=layer.expert_mask_gpu,
+                    expert_mask=layer.dispatcher.expert_mask_gpu,
                 )
             else:
                 return fused_moe(
@@ -1838,7 +1838,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
                         if activation == "silu"
                         else ActivationType.Gelu
                     ),
-                    expert_mask=layer.expert_mask_gpu,
+                    expert_mask=layer.dispatcher.expert_mask_gpu,
                 )
         return None
 

python/sglang/srt/layers/quantization/mxfp4.py

@@ -856,7 +856,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
                 w2_weight,
                 topk_weights,
                 topk_ids,
-                expert_mask=layer.expert_mask_gpu,
+                expert_mask=layer.dispatcher.expert_mask_gpu,
                 activation=ActivationType.Swiglu,
                 quant_type=QuantType.per_1x32,
                 w1_scale=layer.w13_weight_scale,
@@ -1045,6 +1045,6 @@ class Mxfp4DynamicQuantMoEMethod(FusedMoEMethodBase):
                 else ActivationType.Gelu
             ),
             doweight_stage1=False,
-            expert_mask=layer.expert_mask_gpu,
+            expert_mask=layer.dispatcher.expert_mask_gpu,
         )
         return StandardCombineInput(hidden_states=output)

python/sglang/srt/layers/quantization/quark/schemes/quark_w4a4_mxfp4_moe.py

@@ -227,6 +227,6 @@ class QuarkW4A4MXFp4MoE(QuarkMoEScheme):
                 else ActivationType.Gelu
             ),
             doweight_stage1=False,
-            expert_mask=layer.expert_mask_gpu,
+            expert_mask=layer.dispatcher.expert_mask_gpu,
         )
         return StandardCombineInput(hidden_states=output)

python/sglang/srt/layers/quantization/unquant.py

@@ -484,7 +484,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, MultiPlatformOp):
                             if moe_runner_config.activation == "silu"
                             else ActivationType.Gelu
                         ),
-                        expert_mask=layer.expert_mask_gpu,
+                        expert_mask=layer.dispatcher.expert_mask_gpu,
                     )
                     return StandardCombineInput(hidden_states=output)
                 except RuntimeError as e: