MoE: add moe_late_tp_combine so routed experts take the tensor axis as expert parallelism

[systalyze-ai]

Description

Adds moe_late_tp_combine so routed experts use the tensor axis as additional expert parallelism instead of sharding the MoE hidden dimension. This removes tensor-axis collectives from expert-expanded activations.

Performance

• On v6e-128 with DeepSeek-V4-Flash 284B LoRA at seq 16384, step time measured 6.811 s/step with the flag off and 4.632 s/step with it on (1.47x).

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 XLA_FLAGS=--xla_force_host_platform_device_count=8 pytest tests/unit/moe_late_tp_combine_test.py passed 3 tests.

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 the moe_late_tp_combine configuration flag to use the tensor axis as expert parallelism in the default sparse MoE path, moving tensor-axis communication outside of expert-expanded activations. It updates the configurations, MoE layer sharding logic, and adds a new unit test suite. The reviewer recommended ensuring that "tensor" is not already present in the sequence before appending it in token_split_pspec to avoid duplicate mesh axes in the PartitionSpec under custom mesh configurations.

[gemini-code-assist]

To prevent potential JAX sharding errors under custom mesh configurations, it is safer to ensure that "tensor" is not already present in seq before appending it. If "tensor" is already in seq, appending it again would result in a duplicate mesh axis in the PartitionSpec, which is invalid.

Suggested change

	def token_split_pspec(logical_axes):
	spec = list(self._logical_to_mesh_axes(logical_axes))
	seq = spec[1]
	seq = () if seq is None else ((seq,) if isinstance(seq, str) else tuple(seq))
	spec[1] = seq + ("tensor",)
	return P(*spec)
	def token_split_pspec(logical_axes):
	spec = list(self._logical_to_mesh_axes(logical_axes))
	seq = spec[1]
	seq = () if seq is None else ((seq,) if isinstance(seq, str) else tuple(seq))
	if "tensor" not in seq:
	seq = seq + ("tensor",)
	spec[1] = seq
	return P(*spec)