From fc05acc2c7b7aab93c8631eb0a60ca2d41fe731b Mon Sep 17 00:00:00 2001 From: Yinghai Lu Date: Sat, 20 Dec 2025 21:59:36 -0800 Subject: [PATCH] [FusedMoE] Fix fused w13 tp sharded weight loading (#15432) --- .../srt/layers/moe/fused_moe_triton/layer.py | 29 +++++++++++++++++-- 1 file changed, 26 insertions(+), 3 deletions(-) diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py index 71b8e661a..73125e229 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py +++ b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py @@ -404,9 +404,32 @@ class FusedMoE(torch.nn.Module): if not is_bias and self.use_triton_kernels: # do not transpose for bias loaded_weight = loaded_weight.transpose(-2, -1) - loaded_weight = loaded_weight.narrow( - shard_dim, shard_size * tp_rank, shard_size - ) + + if shard_id == "w13": + # For fused w13 weights [w1; w3], we need to shard w1 and w3 + # independently, then concatenate them back to maintain the + # [w1_shard; w3_shard] structure per TP rank. + single_shard_size = shard_size // 2 # Size of w1 or w3 per TP rank + full_single_size = ( + loaded_weight.shape[shard_dim] // 2 + ) # Full size of w1 or w3 + + # Shard w1 (from position 0) + w1_shard = loaded_weight.narrow( + shard_dim, single_shard_size * tp_rank, single_shard_size + ) + # Shard w3 (from position full_single_size) + w3_shard = loaded_weight.narrow( + shard_dim, + full_single_size + single_shard_size * tp_rank, + single_shard_size, + ) + + loaded_weight = torch.cat([w1_shard, w3_shard], dim=shard_dim) + else: + loaded_weight = loaded_weight.narrow( + shard_dim, shard_size * tp_rank, shard_size + ) expert_data = expert_data.narrow(shard_dim, start, shard_size) expert_data.copy_(loaded_weight)