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

[systalyze-ai]

Description

Changes the non-ragged RoutedMoE.permute path to gather directly from the original activations instead of materializing top_k copies before the sorted gather. This removes an activation buffer that is 6x the dispatch input for DeepSeek-V4-284B.

Performance

Step time changed by ~1% on v6e-128 with DeepSeek-V4-Flash 284B LoRA fine-tuning at seq 16384; the benefit of the change is the removed activation buffer. Setup: ici_expert_parallelism=16, ici_tensor_parallelism=8, ici_fsdp_parallelism=1, max_target_length=16384, per_device_batch_size=0.125 (global batch 16), LIBTPU_INIT_ARGS=--xla_tpu_scoped_vmem_limit_kib=98304, dynamic-splash attention path enabled.

Tests

• JAX_PLATFORMS=cpu pytest -q tests/unit/moe_test.py -k GatherReplicatedActivationsTest: 2 passed.

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[gemini-code-assist]

Code Review

This pull request introduces '_gather_replicated_activations' along with a custom VJP implementation to optimize the gathering of sorted rows without materializing repeated inputs in the MoE layer, and adds corresponding unit tests. Feedback is provided to optimize the backward pass of the custom VJP by replacing the inefficient 'jnp.argsort' with a scatter-based permutation inversion to improve TPU performance.

[gemini-code-assist]

Using jnp.argsort to invert a permutation on TPU is highly inefficient because it compiles to an $O(N \log N)$ sorting kernel (like RadixSort or BitonicSort), which has poor hardware utilization on TPU.

Since sort_indices is guaranteed to be a permutation of 0..N-1, we can invert it in $O(N)$ time using a scatter operation: jnp.empty_like(sort_indices).at[sort_indices].set(jnp.arange(sort_indices.shape[0])). This compiles to a highly optimized scatter/transpose-like operation on TPU, significantly speeding up the backward pass.

Note: The same optimization can be applied to _sort_activations_custom_bwd and unpermute in a separate PR to further improve performance.

Suggested change

	sort_indices = residuals
	unsorted_grads = grads[jnp.argsort(sort_indices), ...]
	sort_indices = residuals
	inverse_indices = jnp.empty_like(sort_indices).at[sort_indices].set(jnp.arange(sort_indices.shape[0]))
	unsorted_grads = grads[inverse_indices, ...]