Support EPLB in FusedMoE (#8448)

2025-07-29 16:02:41 -07:00
parent 1992ef9ba7
commit 9effeb5bdd
15 changed files with 107 additions and 11 deletions
@@ -69,6 +69,7 @@ class MixtralMoE(nn.Module):
        top_k: int,
        hidden_size: int,
        intermediate_size: int,
+        layer_id: int,
        params_dtype: Optional[torch.dtype] = None,
        quant_config: Optional[QuantizationConfig] = None,
        tp_size: Optional[int] = None,
@@ -97,6 +98,7 @@ class MixtralMoE(nn.Module):
        self.experts = MoEImpl(
            num_experts=num_experts,
            top_k=top_k,
+            layer_id=layer_id,
            hidden_size=hidden_size,
            intermediate_size=intermediate_size,
            params_dtype=params_dtype,
@@ -226,6 +228,7 @@ class MixtralDecoderLayer(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("block_sparse_moe", prefix),
        )