From 123b8efb8c1903a1f1e799e80d27896842b50640 Mon Sep 17 00:00:00 2001 From: jcarin-sys Date: Mon, 17 Aug 2026 13:53:59 +0000 Subject: [PATCH] MoE: remove the replicated-activation repeat in RoutedMoE.permute Gather directly from the original activations with divided indices instead of materializing top_k copies before the sorted gather. Forward, VJP, and JVP outputs are unchanged. Co-authored-by: Sudarsanan Co-authored-by: Armin Co-authored-by: utlz --- src/maxtext/layers/moe.py | 52 +++++++++++++++++++++++++++++++++++---- tests/unit/moe_test.py | 50 +++++++++++++++++++++++++++++++++++++ 2 files changed, 97 insertions(+), 5 deletions(-) diff --git a/src/maxtext/layers/moe.py b/src/maxtext/layers/moe.py index eccaf295c6..c98d75fde8 100644 --- a/src/maxtext/layers/moe.py +++ b/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 diff --git a/tests/unit/moe_test.py b/tests/unit/moe_test.py index 6d85874993..1273fca3a0 100644 --- a/tests/unit/moe_test.py +++ b/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()