MoE: remove the replicated-activation repeat in RoutedMoE.permute

src/maxtext/layers/moe.py

@@ -171,6 +171,50 @@ def _sort_activations_custom_bwd(residuals: jax.Array, grads: jax.Array) -> tupl
 _sort_activations_custom.defvjp(_sort_activations_custom_fwd, _sort_activations_custom_bwd)
 
 
+def _gather_replicated_activations(
+    inputs: jax.Array,
+    sort_indices: jax.Array,
+    num_repeats: int,
+    use_custom_vjp: bool,
+) -> jax.Array:
+  """Gathers sorted rows without materializing repeated inputs.
+
+  `sort_indices` must permute `inputs.shape[0] * num_repeats` rows.
+  """
+  assert sort_indices.shape[0] == inputs.shape[0] * num_repeats
+
+  if not use_custom_vjp:
+    return _sort_activations(jnp.repeat(inputs, num_repeats, axis=0), sort_indices, use_custom_vjp)
+
+  with jax.named_scope("sort_activations"):
+    return _gather_replicated_custom(inputs, sort_indices, num_repeats)
+
+
+@functools.partial(jax.custom_vjp, nondiff_argnums=(2,))
+def _gather_replicated_custom(inputs: jax.Array, sort_indices: jax.Array, num_repeats: int) -> jax.Array:
+  """Replicated-row gather with custom vjp."""
+  return inputs[sort_indices // num_repeats, ...]
+
+
+def _gather_replicated_custom_fwd(
+    inputs: jax.Array, sort_indices: jax.Array, num_repeats: int
+) -> tuple[jax.Array, jax.Array]:
+  """Forward pass of the custom vjp for `_gather_replicated_activations()`."""
+  return _gather_replicated_custom(inputs, sort_indices, num_repeats), sort_indices
+
+
+def _gather_replicated_custom_bwd(num_repeats: int, residuals: jax.Array, grads: jax.Array) -> tuple[jax.Array, None]:
+  """Backward pass of the custom vjp for `_gather_replicated_activations()`."""
+  sort_indices = residuals
+  unsorted_grads = grads[jnp.argsort(sort_indices), ...]
+  reshaped_grads = unsorted_grads.reshape(-1, num_repeats, *unsorted_grads.shape[1:])
+  # Match the cotangent-dtype accumulation used by the transpose of `jnp.repeat`.
+  return reshaped_grads.sum(axis=1, dtype=grads.dtype), None
+
+
+_gather_replicated_custom.defvjp(_gather_replicated_custom_fwd, _gather_replicated_custom_bwd)
+
+
 def get_batchsplit_init_kernel_axes():
   return (
       ("expert_only", "embed_moe", None),
@@ -952,11 +996,9 @@ def permute(
       if roll_to_expert_id is not None:
         flatten_selected_experts = (flatten_selected_experts - roll_to_expert_id) % self.num_experts
       sorted_selected_experts = jnp.argsort(flatten_selected_experts)
-      # sort inputs for number of selected experts
-      replicated_inputs_2d = jnp.repeat(inputs_2d, self.num_experts_per_tok, axis=0)
-      sorted_inputs = _sort_activations(replicated_inputs_2d, sorted_selected_experts, use_custom_sort_vjp).astype(
-          self.dtype
-      )
+      sorted_inputs = _gather_replicated_activations(
+          inputs_2d, sorted_selected_experts, self.num_experts_per_tok, use_custom_sort_vjp
+      ).astype(self.dtype)
       group_size = jnp.bincount(flatten_selected_experts, length=self.num_experts)
 
     num_tokens = bsz_times_seq_len * self.num_experts_per_tok

tests/unit/moe_test.py

@@ -2119,5 +2119,55 @@ def test_prefuse_moe_weights_matches_unfused(self):
     )
 
 
+class GatherReplicatedActivationsTest(unittest.TestCase):
+  """Tests the replicated activation gather."""
+
+  # pylint: disable=protected-access
+
+  @staticmethod
+  def _reference_chain(inputs, sort_indices, num_repeats, use_custom_vjp):
+    """Pre-optimization dispatch gather: materialized repeat, then sorted gather."""
+    return moe._sort_activations(jnp.repeat(inputs, num_repeats, axis=0), sort_indices, use_custom_vjp)
+
+  @staticmethod
+  def _random_case(seed, num_tokens, hidden, num_experts, num_experts_per_tok, dtype):
+    """Random activations plus a permute()-style argsort of flattened top-k expert ids."""
+    rng = np.random.default_rng(seed)
+    inputs = jnp.asarray(rng.standard_normal((num_tokens, hidden)), dtype=dtype)
+    selected_experts = rng.integers(0, num_experts, size=(num_tokens, num_experts_per_tok))
+    sort_indices = jnp.argsort(jnp.ravel(jnp.asarray(selected_experts, dtype=jnp.int32)))
+    return inputs, sort_indices
+
+  def test_forward_and_vjp_match_reference(self):
+    """Tests bitwise forward and vjp equality on the custom-vjp path."""
+    for num_tokens, hidden, num_experts, top_k in ((24, 16, 8, 6), (7, 5, 3, 4)):
+      for seed, dtype in enumerate((jnp.float32, jnp.bfloat16)):
+        with self.subTest(num_tokens=num_tokens, dtype=dtype):
+          inputs, sort_indices = self._random_case(seed, num_tokens, hidden, num_experts, top_k, dtype)
+
+          def ref_fn(x):
+            return self._reference_chain(x, sort_indices, top_k, True)
+
+          def new_fn(x):
+            return moe._gather_replicated_activations(x, sort_indices, top_k, True)
+
+          ref_out, ref_vjp = jax.vjp(ref_fn, inputs)
+          new_out, new_vjp = jax.vjp(new_fn, inputs)
+
+          self.assertEqual(ref_out.dtype, new_out.dtype)
+          np.testing.assert_array_equal(np.asarray(ref_out), np.asarray(new_out))
+          ct = jnp.asarray(np.random.default_rng(seed + 1000).standard_normal(ref_out.shape), dtype=dtype)
+          np.testing.assert_array_equal(np.asarray(ref_vjp(ct)[0]), np.asarray(new_vjp(ct)[0]))
+
+  def test_jvp_matches_reference(self):
+    """custom_vjp forbids jvp, so forward-mode has to go through the use_custom_vjp=False fallback."""
+    inputs, sort_indices = self._random_case(11, 24, 16, 8, 6, jnp.float32)
+    tangent = jnp.asarray(np.random.default_rng(12).standard_normal(inputs.shape), dtype=jnp.float32)
+    ref = jax.jvp(lambda x: self._reference_chain(x, sort_indices, 6, False), (inputs,), (tangent,))
+    new = jax.jvp(lambda x: moe._gather_replicated_activations(x, sort_indices, 6, False), (inputs,), (tangent,))
+    for reference, actual in zip(ref, new):
+      np.testing.assert_array_equal(np.asarray(reference), np.asarray(actual))
+
+
 if __name__ == "__main__":
   unittest.main()