chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Integration mori backend for EP a2a data communication (#17012)

Browse Source 
Co-authored-by: Duyi-Wang <duyi.wang@amd.com>
Co-authored-by: billishyahao <bill.he@amd.com>
Co-authored-by: HaiShaw <hixiao@gmail.com>
This commit is contained in:
kk
2026-01-28 19:07:34 -08:00
committed by GitHub
co-authored by Duyi-Wang billishyahao HaiShaw
parent 673dc09d9b
commit f1384f5293
15 changed files with 934 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/models/deepseek_v2.py
+22 -4
View File
@@ -467,6 +467,7 @@ class DeepseekV2MoE(nn.Module):
dict(tp_rank=0, tp_size=1)
if get_moe_a2a_backend().is_deepep()
or get_moe_a2a_backend().is_mooncake()
or get_moe_a2a_backend().is_mori()
or get_moe_a2a_backend().is_ascend_fuseep()
or get_moe_a2a_backend().is_flashinfer()
or should_use_flashinfer_cutlass_moe_fp4_allgather()
@@ -510,6 +511,7 @@ class DeepseekV2MoE(nn.Module):
if (
get_moe_a2a_backend().is_deepep()
or get_moe_a2a_backend().is_mooncake()
or get_moe_a2a_backend().is_mori()
or get_moe_a2a_backend().is_ascend_fuseep()
):
# TODO: we will support tp < ep in the future
@@ -530,6 +532,7 @@ class DeepseekV2MoE(nn.Module):
self._enable_a2a_moe = (
get_moe_a2a_backend().is_deepep()
or get_moe_a2a_backend().is_mooncake()
or get_moe_a2a_backend().is_mori()
or get_moe_a2a_backend().is_ascend_fuseep()
or get_moe_a2a_backend().is_flashinfer()
)
@@ -945,13 +948,17 @@ class DeepseekV2MoE(nn.Module):
torch.cuda.current_stream().wait_event(shared_event)
if shared_output is not None:
x = shared_output
if self.experts.should_fuse_routed_scaling_factor_in_topk:
# aiter moe call will handle routed_scaling_factor in the function
# so add _use_aiter condition to eliminate to use self.routed_scaling_factor in add_ call
if self.experts.should_fuse_routed_scaling_factor_in_topk or _use_aiter:
x.add_(final_hidden_states)
else:
x.add_(final_hidden_states, alpha=self.routed_scaling_factor)
final_hidden_states = x
else:
if not self.experts.should_fuse_routed_scaling_factor_in_topk:
if not (
self.experts.should_fuse_routed_scaling_factor_in_topk or _use_aiter
):
final_hidden_states *= self.routed_scaling_factor
return final_hidden_states
@@ -2576,7 +2583,16 @@ class DeepseekV2Model(nn.Module):
allocate_size = 0
for i in range(len(self.layers)):
if isinstance(self.layers[i].mlp, DeepseekV2MoE):
tp_size = get_tensor_model_parallel_world_size()
# tp_size = get_tensor_model_parallel_world_size()
a2a_backend = get_moe_a2a_backend()
is_a2a_moe = (
a2a_backend.is_deepep()
or a2a_backend.is_mori()
or a2a_backend.is_mooncake()
)
tp_size = (
1 if is_a2a_moe else get_tensor_model_parallel_world_size()
)
intermediate_size = (
config.moe_intermediate_size * config.n_shared_experts
)
@@ -2842,7 +2858,9 @@ class DeepseekV2ForCausalLM(nn.Module, DeepseekV2WeightLoaderMixin):
not _is_hip or torch.cuda.get_device_capability("cuda") < (9, 4)
):
disable_reason = "Only Deepseek V3/R1 on AMD-platform with capability >= gfx942(MI30x) can use shared experts fusion optimization under expert parallelism."
elif disable_reason is None and get_moe_a2a_backend().is_deepep():
elif disable_reason is None and (
get_moe_a2a_backend().is_deepep() or get_moe_a2a_backend().is_mori()
):
disable_reason = "Deepseek V3/R1 can not use shared experts fusion optimization under deepep expert parallelism."
elif self.quant_config and self.quant_config.get_name() == "w4afp8":
disable_reason = "Deepseek V3/R1 W4AFP8 model uses different quant method for routed experts and shared experts."