Support EPLB in FusedMoE (#8448)

python/sglang/srt/models/olmoe.py

@@ -63,6 +63,7 @@ class OlmoeMoE(nn.Module):
         params_dtype: Optional[torch.dtype] = None,
         quant_config: Optional[QuantizationConfig] = None,
         tp_size: Optional[int] = None,
+        layer_id: int = 0,
         prefix: str = "",
     ):
         super().__init__()
@@ -89,6 +90,7 @@ class OlmoeMoE(nn.Module):
             reduce_results=True,
             quant_config=quant_config,
             tp_size=tp_size,
+            layer_id=layer_id,
             prefix=add_prefix("experts", prefix),
         )
 
@@ -224,6 +226,7 @@ class OlmoeDecoderLayer(nn.Module):
             top_k=config.num_experts_per_tok,
             hidden_size=config.hidden_size,
             intermediate_size=config.intermediate_size,
+            layer_id=layer_id,
             quant_config=quant_config,
             prefix=add_prefix("mlp", prefix),
         )