Keep fused MoE weights intact through FSDP gather (1.5% perf improvement on DeepSeek-v3)

[abhinavgoel95]

Description

Preserve prefused MoE input weights through shard_map so FSDP performs one all-gather of the combined weight tensor before the grouped GEMM.

Previously, MaxText split the prefused [experts, embed, 2 * mlp] weight into separate w0 and w1 tensors before shard_map. This could cause XLA to all-gather each half independently and concatenate them again immediately before the grouped GEMM.

The updated flow is:

prefused weight
    -> one FSDP all-gather
    -> one grouped GEMM
    -> split output activations for SwiGLU

Only the GMM output is split into gate and up-projection activations. This split is required for SwiGLU and is considerably cheaper than splitting, gathering, and concatenating the large weight tensors.

The implementation:

• Passes the prefused weight directly into the sparse MoE shard_map.
• Uses one partition specification for the combined weight.
• Preserves the combined weight through one-stage and two-stage FSDP all-gather paths.
• Runs one grouped GEMM using the combined weight.
• Splits the resulting activations before applying the configured FFN activation.
• Preserves existing behavior for quantized serving, weight sparsity, dense MoE computation, and vLLM RPA inference.

The optimization currently applies to sparse MoE training when prefused weights can be consumed directly. Quantized serving and sparsified-weight paths continue using separate w0 and w1 tensors. Supporting fused quantized-weight retrieval could be considered as a future improvement.

Tests

The original experimental change was profiled with a DeepSeek MoE workload and removed the separate weight-slice operations before the grouped GEMM.

The version rebased onto current main has not been run end-to-end. The patch was checked with:

git diff --check

To reproduce the performance comparison, run the same sparse MoE workload with:

prefuse_moe_weights=true
sparse_matmul=true

Compare the generated HLO or GPU profile before and after this change. The updated version should contain one FSDP all-gather for the combined MoE input weight and should not split and concatenate the weight before the grouped GEMM.

Checklist

• 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.

[gobbleturk]

sync issue? I assume this file should be untouched in this PR, please note if this understanding is not correct

[gobbleturk]

can this be put into a separate PR? Keep PRs small for easier readability and reviewing

[gobbleturk]

can this be part of a separate PR? Is this quantization necessary?

[gobbleturk]

I see maybe this quantization technique has some dependency with fused vs non fused (which is surprising to me), maybe since fused just doesn't support quantization for now?

[gobbleturk]

would prefer to put this in a method (to keep train_loop as lean and readable as possible). Either a method defined somewhere outside of train_loop, or possibly in a new te_utils file