Merge pull request #3736 from abhinavgoel95:abgoel/fused-moe-mlp

PiperOrigin-RevId: 939895106

Author: Google-ML-Automation

docs/reference/core_concepts/moe_configuration.md

@@ -99,6 +99,8 @@ Dropping:
 
 `mlp_bias`: If enabled, add learnable bias terms for MLP matmul. Originally implemented to support the GPT-OSS model architecture.
 
+`prefuse_moe_weights`: If enabled alongside `sparse_matmul=True`, fuses the two FFN1 grouped GEMMs (wi\_0 and wi\_1) into a single grouped GEMM call. Expert weights are stored in a concatenated `(num_experts, embed_dim, 2 * mlp_dim)` shape, so input activations are loaded from HBM once per forward pass instead of twice. Backend-agnostic (works with Megablox, JAX Ragged Dot, and Tokamax). When used with `attention=vllm_rpa`, the fused weight tensor is passed directly to the vLLM-TPU serving kernel without splitting.
+
 `use_batch_split_schedule` (experimental): If enabled, split batch into micro-batches to hide communications that yields performance benefits.
 
 ## 2. Sharding

src/maxtext/configs/types.py

@@ -845,7 +845,7 @@ class MoEGeneral(BaseModel):
   prefuse_moe_weights: bool = Field(
       False,
       description="Whether to pre-fuse MoE weights (w0 and w1) during initialization. "
-      "This is useful for inference performance in vllm_rpa mode.",
+      "This enables a single FFN1 grouped GEMM in sparse MoE paths and passes fused weights directly in vllm_rpa mode.",
   )
   fuse_expert_scales: bool = Field(
       False,

src/maxtext/layers/moe.py

@@ -499,6 +499,12 @@ def __init__(
 
     kernel_in_axis = np.arange(1)
     kernel_out_axis = np.arange(1, 2)
+    # Pad the MoE input weight kernels for GMM_v2 execution when configured.
+    moe_intermediate_dim = (
+        self.config.padded_base_moe_mlp_dim
+        if self.config.padded_base_moe_mlp_dim is not None
+        else self.intermediate_dim
+    )
 
     if quantizations.in_serve_mode(self.quant):
       # During aqt convert state we delete kernel weight from params to save
@@ -507,13 +513,7 @@ def __init__(
       self.wi_0 = jnp.zeros((num_experts, self.moe_expert_input_dim, intermediate_dim))
       self.wi_1 = jnp.zeros((num_experts, self.moe_expert_input_dim, intermediate_dim))
       self.wo = jnp.zeros((num_experts, intermediate_dim, self.moe_expert_input_dim))
-    elif self.config.prefuse_moe_weights and self.config.attention == "vllm_rpa":
-      # Pad model dimension in Fused MoE weight kernels for GMM_v2 execution.
-      moe_intermediate_dim = (
-          self.config.padded_base_moe_mlp_dim
-          if self.config.padded_base_moe_mlp_dim is not None
-          else self.intermediate_dim
-      )
+    elif self.config.prefuse_moe_weights:
       self.wi = nnx.Param(
           self.kernel_init(
               self.rngs.params(),
@@ -535,12 +535,6 @@ def __init__(
           out_sharding=self.wo_kernel_axes,
       )
     else:
-      # Pad model dimension in Unfused MoE weight kernels for GMM_v2 execution.
-      moe_intermediate_dim = (
-          self.config.padded_base_moe_mlp_dim
-          if self.config.padded_base_moe_mlp_dim is not None
-          else self.intermediate_dim
-      )
       self.wi_0 = nnx.Param(
           self.kernel_init(
               self.rngs.params(),
@@ -1669,19 +1663,44 @@ def get_wo_gmm_params():
     def gmm_up(x, w0, w1, w0_bias, w1_bias, gmm_fn, weight_gather):
       """Run the two up-projections (gate + up) and apply the FFN activation."""
       wi_gather_axes, wi_tile_size = get_wi_gmm_params()
-      layer_w0 = gmm_fn(x, w0, tiling=wi_tile_size, weight_gather_axes=wi_gather_axes)
-      if self.get_tensor_transpose_parallelism_size() > 1:
-        layer_w0 = jax.lax.psum(layer_w0, "tensor_transpose")
-      if self.config.mlp_bias:
-        layer_w0 = layer_w0 + w0_bias
-      layer_w0 = adc.checkpoint_name(adc.checkpoint_name(layer_w0, "mlpwi_0"), "moe_mlpwi_0")
+      if self.config.prefuse_moe_weights:
+        # Weights are stored as (G,K,2N); w0/w1 are adjacent slices so XLA elides this concat.
+        w_fused = jnp.concatenate([w0, w1], axis=-1)
+        out = gmm_fn(x, w_fused, tiling=wi_tile_size, weight_gather_axes=wi_gather_axes)
+        n = out.shape[-1] // 2
+        layer_w0, layer_w1 = out[:, :n], out[:, n:]
+        if self.get_tensor_transpose_parallelism_size() > 1:
+          layer_w0 = jax.lax.psum(layer_w0, "tensor_transpose")
+          layer_w1 = jax.lax.psum(layer_w1, "tensor_transpose")
+        if self.config.mlp_bias:
+          layer_w0 = layer_w0 + w0_bias
+          layer_w1 = layer_w1 + w1_bias
+        layer_w0 = adc.checkpoint_name(adc.checkpoint_name(layer_w0, "mlpwi_0"), "moe_mlpwi_0")
+        layer_w1 = adc.checkpoint_name(layer_w1, "moe_mlpwi_1")
+      else:
+        layer_w0 = gmm_fn(
+            x,
+            w0,
+            tiling=wi_tile_size,
+            weight_gather_axes=wi_gather_axes,
+        )
+        if self.get_tensor_transpose_parallelism_size() > 1:
+          layer_w0 = jax.lax.psum(layer_w0, "tensor_transpose")
+        if self.config.mlp_bias:
+          layer_w0 = layer_w0 + w0_bias
+        layer_w0 = adc.checkpoint_name(adc.checkpoint_name(layer_w0, "mlpwi_0"), "moe_mlpwi_0")
 
-      layer_w1 = gmm_fn(x, w1, tiling=wi_tile_size, weight_gather_axes=wi_gather_axes)
-      if self.get_tensor_transpose_parallelism_size() > 1:
-        layer_w1 = jax.lax.psum(layer_w1, "tensor_transpose")
-      if self.config.mlp_bias:
-        layer_w1 = layer_w1 + w1_bias
-      layer_w1 = adc.checkpoint_name(layer_w1, "moe_mlpwi_1")
+        layer_w1 = gmm_fn(
+            x,
+            w1,
+            tiling=wi_tile_size,
+            weight_gather_axes=wi_gather_axes,
+        )
+        if self.get_tensor_transpose_parallelism_size() > 1:
+          layer_w1 = jax.lax.psum(layer_w1, "tensor_transpose")
+        if self.config.mlp_bias:
+          layer_w1 = layer_w1 + w1_bias
+        layer_w1 = adc.checkpoint_name(layer_w1, "moe_mlpwi_1")
       return self.apply_ffn_activation(layer_w0, layer_w1)
 
     def get_gmm_for_local_experts(x, routing, route_metadata):
@@ -2592,6 +2611,11 @@ def __call__(
     w1_kernel = None
     if cfg.prefuse_moe_weights and cfg.attention == "vllm_rpa" and not self.is_hash_routing:
       fused_kernel = jnp.asarray(self.wi[...], self.dtype)
+    elif cfg.prefuse_moe_weights:
+      wi = jnp.asarray(self.wi[...], self.dtype)
+      n = wi.shape[-1] // 2
+      w0_kernel = wi[..., :n]
+      w1_kernel = wi[..., n:]
     else:
       w0_kernel = jnp.asarray(self.wi_0[...], self.dtype)
       w1_kernel = jnp.asarray(self.wi_1[...], self.dtype)

tests/unit/moe_test.py

@@ -1643,5 +1643,69 @@ def test_moe_dispatch_no_expert_sharding_dense_forward():
   assert out.shape == inputs.shape
 
 
+@pytest.mark.tpu_only
+class FusedMlpMoETest(unittest.TestCase):
+  """Tests that prefuse_moe_weights=True and prefuse_moe_weights=False produce identical outputs for MoE."""
+
+  def __init__(self, *args, **kwargs):
+    super().__init__(*args, **kwargs)
+    self._B = 1
+    self._S = 16
+
+  def setUp(self):
+    super().setUp()
+    self.rng = jax.random.PRNGKey(0)
+    self.ref_cfg = pyconfig.initialize(
+        [None, get_test_config_path()],
+        run_name="fused_mlp_moe_ref",
+        enable_checkpointing=False,
+        model_name="mixtral-8x7b",
+        dtype="bfloat16",
+        sparse_matmul=True,
+        megablox=True,
+        prefuse_moe_weights=False,
+        ici_expert_parallelism=jax.device_count(),
+        max_target_length=self._S,
+        per_device_batch_size=self._B,
+    )
+    ref_devices = maxtext_utils.create_device_mesh(self.ref_cfg)
+    self.ref_mesh = Mesh(ref_devices, self.ref_cfg.mesh_axes)
+    self.ref_model = make_moe(self.ref_cfg, self.ref_mesh)
+
+  def _inputs(self):
+    return jax.random.normal(self.rng, (self._B, self._S, self.ref_cfg.base_emb_dim), dtype=jnp.bfloat16)
+
+  def test_prefuse_moe_weights_matches_unfused(self):
+    """prefuse_moe_weights=True output matches prefuse_moe_weights=False with sparse_matmul (Megablox)."""
+    fused_cfg = pyconfig.initialize(
+        [None, get_test_config_path()],
+        run_name="fused_mlp_moe_fused",
+        enable_checkpointing=False,
+        model_name="mixtral-8x7b",
+        dtype="bfloat16",
+        sparse_matmul=True,
+        megablox=True,
+        prefuse_moe_weights=True,
+        ici_expert_parallelism=jax.device_count(),
+        max_target_length=self._S,
+        per_device_batch_size=self._B,
+    )
+    fused_devices = maxtext_utils.create_device_mesh(fused_cfg)
+    fused_mesh = Mesh(fused_devices, fused_cfg.mesh_axes)
+    fused_model = make_moe(fused_cfg, fused_mesh)
+    copy_weights_prefused(self.ref_model, fused_model)
+
+    inputs = self._inputs()
+    ref_out, _, _ = self.ref_model(inputs)
+    fused_out, _, _ = fused_model(inputs)
+
+    np.testing.assert_allclose(
+        np.array(ref_out, dtype=np.float32),
+        np.array(fused_out, dtype=np.float32),
+        rtol=1e-2,
+        atol=1e-2,
+    )
+
+
 if __name__ == "__main__":
   unittest.main()