minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Revert "[FusedMoE] Fix fused w13 tp sharded weight loading" (#15579)

Browse Source
This commit is contained in:
Liangsheng Yin
2025-12-22 04:15:24 +08:00
committed by GitHub
parent 8766a1dd24
commit 393e2f9b62
1 changed files with 3 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/moe/fused_moe_triton/layer.py
+3 -26
View File
@@ -410,32 +410,9 @@ class FusedMoE(torch.nn.Module):
if not is_bias and self.use_triton_kernels: if not is_bias and self.use_triton_kernels:
# do not transpose for bias # do not transpose for bias
loaded_weight = loaded_weight.transpose(-2, -1) loaded_weight = loaded_weight.transpose(-2, -1)
loaded_weight = loaded_weight.narrow(
if shard_id == "w13": shard_dim, shard_size * tp_rank, shard_size
# 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 = expert_data.narrow(shard_dim, start, shard_size)
expert_data.copy_(loaded_weight) expert_data.copy_(loaded_weight)