adding support for fused_moe_gmm

Author: NicoGrande

src/maxtext/configs/inference/vllm.yml

@@ -57,19 +57,18 @@ logical_axis_rules: [
                       ['activation_norm_length', []],
                       ['activation_norm_length_moe', []],
                       ['activation_exp', ['expert', 'attn_dp_expert']],
-                      ['decode_batch', ['expert', 'attn_dp_expert']],
+                      ['decode_batch', []],
                       ['decode_length', []],
                       ['mlp', ['model', 'attn_dp']],
                       ['mlp_no_fsdp', ['model', 'attn_dp']],
                       ['moe_mlp', ['model', 'attn_dp']],
                       ['vocab', ['model', 'attn_dp']],
-                      ['heads', ['model']],
+                      ['heads', ['model', 'expert']],
                       ['q_heads', ['model', 'expert']],
                       ['kv_heads', ['model', 'expert']],
                       ['kv_head_dim', []],
                       ['kv', []],
-                      ['embed', ['expert', 'attn_dp_expert']],
-                      ['embed', ['attn_dp_expert']],
+                      ['embed', []],
                       ['embed_moe', ['expert', 'attn_dp_expert']],
                       ['embed_moe', ['attn_dp_expert']],
                       ['embed_tensor_transpose', ['attn_dp', 'model']],

src/maxtext/configs/types.py

@@ -686,6 +686,11 @@ class MoEGeneral(BaseModel):
       False,
       description="Whether to cast inputs to fp32 to compute MoE gate logits for numerical stability.",
   )
+  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.",
+  )
 
 
 class MoEKernels(BaseModel):

src/maxtext/layers/moe.py

@@ -384,7 +384,7 @@ def __init__(
         kernel_init=self.kernel_init,
         kernel_axes=self.kernel_axes,
         use_bias=self.config.routed_bias,
-        score_func=self.config.routed_score_func,
+        score_func="" if self.config.attention == "vllm_rpa" else self.config.routed_score_func,
         matmul_precision=self.config.matmul_precision,
         shard_mode=config.shard_mode,
         rngs=self.rngs,
@@ -403,6 +403,27 @@ def __init__(
       self.wi_0 = jnp.zeros((num_experts, self.config.emb_dim, intermediate_dim))
       self.wi_1 = jnp.zeros((num_experts, self.config.emb_dim, intermediate_dim))
       self.wo = jnp.zeros((num_experts, intermediate_dim, self.config.emb_dim))
+    elif self.config.prefuse_moe_weights and self.config.attention == "vllm_rpa":
+      self.wi = nnx.Param(
+          self.kernel_init(
+              self.rngs.params(),
+              (num_experts, self.config.emb_dim, intermediate_dim * 2),
+              weight_dtype,
+              kernel_in_axis,
+              kernel_out_axis,
+          ),
+          sharding=self.wi_kernel_axes,
+      )
+      self.wo = nnx.Param(
+          self.kernel_init(
+              self.rngs.params(),
+              (self.num_experts, self.intermediate_dim, self.config.emb_dim),
+              self.weight_dtype,
+              kernel_in_axis,
+              kernel_out_axis,
+          ),
+          sharding=self.wo_kernel_axes,
+      )
     else:
       self.wi_0 = nnx.Param(
           self.kernel_init(
@@ -1988,6 +2009,68 @@ def dense_matmul(
         ).astype(self.dtype)
       return output, lb_loss, bias_updates
 
+  def fused_moe_matmul(
+      self,
+      inputs,
+      gate_logits,
+      wo_kernel,
+      w0_kernel=None,
+      w1_kernel=None,
+      fused_kernel=None,
+  ) -> tuple[jax.Array, None, None]:
+    """Fused MoE via tpu_inference fused_moe_func (vllm_rpa path only).
+
+    fused_moe_func handles routing, GMM, and weighted combination internally.
+    It does not compute lb_loss or bias_updates (inference-only).
+    """
+    try:
+      # pylint: disable=import-outside-toplevel
+      # pytype: disable=import-error
+      from tpu_inference.layers.common.fused_moe_gmm import fused_moe_func
+    except ImportError as e:
+      raise ImportError("fused_moe_matmul requires the tpu-inference package.") from e
+
+    # Reshape 3D [B, S, D] -> 2D [T, D] (fused_moe_func expects 2D input)
+    batch_size, seq_len, emb_dim = inputs.shape
+    hidden_states = jnp.reshape(inputs, (batch_size * seq_len, emb_dim))
+    gating_output = jnp.reshape(gate_logits, (batch_size * seq_len, self.num_experts))
+
+    # Concatenate gate and up projections: [E, D, H] + [E, D, H] -> [E, D, 2H]
+    # fused_moe_func splits this internally: gate=w1[..., :H], up=w1[..., H:]
+    if fused_kernel is None:
+      fused_kernel = jnp.concatenate([w0_kernel, w1_kernel], axis=-1)
+
+    # Use expert parallelism if the expert axis has size > 1
+    use_ep = self.get_expert_parallelism_size() > 1
+
+    # Map MaxText config fields to fused_moe_func args
+    activation = self.config.mlp_activations[0]  # e.g. "silu"
+    scoring_fn = self.config.routed_score_func if self.config.routed_score_func else "softmax"
+    renormalize = self.config.norm_topk_prob
+
+    output_2d = fused_moe_func(
+        hidden_states=hidden_states,
+        w1=fused_kernel,
+        w2=wo_kernel,
+        w1_scale=None,
+        w2_scale=None,
+        w1_bias=None,
+        w2_bias=None,
+        gating_output=gating_output,
+        topk=self.num_experts_per_tok,
+        renormalize=renormalize,
+        mesh=self.mesh,
+        use_ep=use_ep,
+        activation=activation,
+        scoring_fn=scoring_fn,
+        sc_kernel_threshold=16777216,
+        sc_kernel_col_chunk_size=1024,
+    )
+
+    # Reshape output 2D [T, D] -> 3D [B, S, D]
+    output = jnp.reshape(output_2d, (batch_size, seq_len, emb_dim))
+    return output, None, None
+
   def retrieve_quantized_weight(
       self,
       inputs,
@@ -2026,10 +2109,17 @@ def __call__(
     routing_inputs = inputs if gate_inputs is None else gate_inputs.astype(gate_dtype)
     gate_logits, pre_bias_logits = self.gate(routing_inputs)
 
-    w0_kernel = jnp.asarray(self.wi_0[...], self.dtype)
-    w1_kernel = jnp.asarray(self.wi_1[...], self.dtype)
     wo_kernel = jnp.asarray(self.wo[...], self.dtype)
 
+    fused_kernel = None
+    w0_kernel = None
+    w1_kernel = None
+    if cfg.prefuse_moe_weights and cfg.attention == "vllm_rpa":
+      fused_kernel = jnp.asarray(self.wi[...], self.dtype)
+    else:
+      w0_kernel = jnp.asarray(self.wi_0[...], self.dtype)
+      w1_kernel = jnp.asarray(self.wi_1[...], self.dtype)
+
     if self.per_expert_scale is not None:
       wo_kernel = wo_kernel * jnp.asarray(self.per_expert_scale[...], self.dtype)[:, None, None]
 
@@ -2040,7 +2130,12 @@ def __call__(
     else:
       w0_bias, w1_bias, wo_bias = None, None, None
 
-    if cfg.sparse_matmul:
+    # vllm_rpa uses fused_moe_func from tpu_inference (highest priority)
+    if cfg.attention == "vllm_rpa":
+      output, lb_loss, bias_updates = self.fused_moe_matmul(
+          inputs, gate_logits, wo_kernel, w0_kernel=w0_kernel, w1_kernel=w1_kernel, fused_kernel=fused_kernel
+      )
+    elif cfg.sparse_matmul:
       if quantizations.in_serve_mode(self.quant):
         w0_kernel, w1_kernel, wo_kernel = self.retrieve_quantized_weight(
             inputs,

src/maxtext/utils/model_creation_utils.py

@@ -25,6 +25,7 @@
 import flax.linen as nn
 import jax
 import jax.numpy as jnp
+import numpy as np
 from jax.sharding import AxisType, Mesh
 from maxtext.configs import pyconfig
 from maxtext.common.common_types import MODEL_MODE_TRAIN, ShardMode
@@ -298,6 +299,35 @@ def create_sharded_state():
         # Get the structure of checkpoint in `config.load_parameters_path`
         metadata = ckptr.metadata(config.load_parameters_path)
 
+        def _adjust_target_for_moe_fusion(target, meta_tree, is_nnx):
+          if not hasattr(target, "items") or not hasattr(meta_tree, "items"):
+            return target
+          new_target = {}
+          for k, v in target.items():
+            if k == "wi" and "wi" not in meta_tree and "wi_0" in meta_tree and "wi_1" in meta_tree:
+              if not is_nnx:
+                arr = v
+                half_dim = arr.shape[-1] // 2
+                new_target["wi_0"] = jax.ShapeDtypeStruct(
+                    shape=arr.shape[:-1] + (half_dim,), dtype=arr.dtype, sharding=arr.sharding
+                )
+                new_target["wi_1"] = jax.ShapeDtypeStruct(
+                    shape=arr.shape[:-1] + (half_dim,), dtype=arr.dtype, sharding=arr.sharding
+                )
+              else:
+                arr = v["value"]
+                half_dim = arr.shape[-1] // 2
+                new_target["wi_0"] = {
+                    "value": jax.ShapeDtypeStruct(shape=arr.shape[:-1] + (half_dim,), dtype=arr.dtype, sharding=arr.sharding)
+                }
+                new_target["wi_1"] = {
+                    "value": jax.ShapeDtypeStruct(shape=arr.shape[:-1] + (half_dim,), dtype=arr.dtype, sharding=arr.sharding)
+                }
+            else:
+              new_target[k] = _adjust_target_for_moe_fusion(v, meta_tree.get(k, {}), is_nnx)
+
+          return new_target
+
         is_nnx_checkpoint = True
         if (
             "params" in metadata.item_metadata.tree.keys()
@@ -311,6 +341,10 @@ def create_sharded_state():
               is_leaf=lambda n: hasattr(n, "value"),
           )
 
+          target_for_restore = _adjust_target_for_moe_fusion(
+              target_for_restore, metadata.item_metadata.tree["params"]["params"], False
+          )
+
           item_to_restore = {"params": {"params": target_for_restore}}
           base_restore_args = ocp.checkpoint_utils.construct_restore_args(target_for_restore)
           restore_args = {
@@ -329,6 +363,7 @@ def create_sharded_state():
               sharded_state,
               is_leaf=lambda n: isinstance(n, nnx.Variable),
           )
+          target_for_restore = _adjust_target_for_moe_fusion(target_for_restore, metadata.item_metadata.tree, True)
           item_to_restore = target_for_restore
           base_restore_args = ocp.checkpoint_utils.construct_restore_args(target_for_restore)
           restore_args = _fix_restore_args_for_shape_mismatch(
@@ -368,6 +403,36 @@ def create_sharded_state():
               sharded_state,
               is_leaf=lambda n: isinstance(n, nnx.Variable),
           )
+
+          def to_dict(tree):
+            if hasattr(tree, "items"):
+              return {k: to_dict(v) for k, v in tree.items()}
+            return tree
+
+          model_arrays = to_dict(model_arrays)
+          checkpoint = to_dict(checkpoint)
+
+          def _fuse_moe_weights(ckpt_tree, model_arrays_tree):
+            if not hasattr(ckpt_tree, "items") or not hasattr(model_arrays_tree, "items"):
+              return ckpt_tree
+            new_ckpt = {}
+            for k, v in ckpt_tree.items():
+              if k in ("wi_0", "wi_1") and "wi" in model_arrays_tree:
+                continue
+              new_ckpt[k] = _fuse_moe_weights(v, model_arrays_tree.get(k, {}))
+
+            if "wi" in model_arrays_tree and "wi_0" in ckpt_tree and "wi_1" in ckpt_tree:
+              wi_0 = ckpt_tree["wi_0"]
+              wi_1 = ckpt_tree["wi_1"]
+              new_ckpt["wi"] = np.concatenate([wi_0, wi_1], axis=-1)
+
+            return new_ckpt
+
+          checkpoint = _fuse_moe_weights(checkpoint, model_arrays)
+          # Release the raw restored buffers now that wi_0/wi_1 have been fused (if needed).
+          # This prevents the replicated intermediate copies from persisting until function return.
+          del restored
+
           checkpoint = jax.tree.map(_expand_checkpoint_to_model_shapes, checkpoint, model_arrays)
           nnx.update(model, checkpoint)