Keep fused MoE weights intact through FSDP gather

[abhinavgoel95]

Description

Keep prefused MoE input weights intact through the FSDP all-gather, avoiding an unnecessary weight slice/concatenate sequence and improving DeepSeek V3 performance by approximately 1.5% on GB300.

Previously, prefuse_moe_weights stored the two MoE input projections as one [experts, embed, 2 * mlp] tensor, but the sparse MoE path split this tensor into w0 and w1 before entering shard_map. This caused FSDP to all-gather the two halves separately, after which they were concatenated again immediately before the grouped GEMM.

This change passes the prefused weight through shard_map intact. The sparse MoE path now performs:

1. One FSDP all-gather of the [experts, embed, 2 * mlp] weight.
2. One grouped GEMM producing [tokens, 2 * mlp].
3. A split of the activation output immediately before the SwiGLU activation.

This removes the pre-GEMM weight slicing and avoids the second independent FSDP all-gather. The implementation also preserves the fused weight through two-stage FSDP all-gather and applies the appropriate partition specifications to the combined tensor.

The optimization is restricted to sparse matmul with prefused weights, outside quantized serving and weight-sparsity paths. Existing split-weight paths remain unchanged as fallbacks. There is no change to the checkpoint format or default behavior because the optimization is controlled by the existing prefuse_moe_weights option.

Tests

Tested with DeepSeek V3 training on NVIDIA GB300 using:

prefuse_moe_weights: true
sparse_matmul: true

Compared against the same workload, configuration, and container without this change:

• HLO/profile inspection confirmed that the prefused weight remains intact through FSDP all-gather.
• The unnecessary pre-GEMM weight slice operations were removed.
• The two independent input-weight all-gathers were replaced by one combined all-gather.
• End-to-end training performance improved by approximately 1.5%.
• git diff --check passes.

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. No documentation changes are required for this internal optimization.