Co-authored-by: Byron Hsu <byronhsu@noreply.github.com>
This commit is contained in:
@@ -50,6 +50,7 @@ from sglang.srt.layers.linear import (
|
|||||||
from sglang.srt.layers.logits_processor import LogitsProcessor
|
from sglang.srt.layers.logits_processor import LogitsProcessor
|
||||||
from sglang.srt.layers.moe import (
|
from sglang.srt.layers.moe import (
|
||||||
get_moe_a2a_backend,
|
get_moe_a2a_backend,
|
||||||
|
should_use_dp_reduce_scatterv,
|
||||||
should_use_flashinfer_cutlass_moe_fp4_allgather,
|
should_use_flashinfer_cutlass_moe_fp4_allgather,
|
||||||
)
|
)
|
||||||
from sglang.srt.layers.moe.ep_moe.layer import get_moe_impl_class
|
from sglang.srt.layers.moe.ep_moe.layer import get_moe_impl_class
|
||||||
@@ -331,7 +332,11 @@ class Qwen3MoeSparseMoeBlock(nn.Module):
|
|||||||
topk_output = self.topk(hidden_states, router_logits)
|
topk_output = self.topk(hidden_states, router_logits)
|
||||||
final_hidden_states = self.experts(hidden_states, topk_output)
|
final_hidden_states = self.experts(hidden_states, topk_output)
|
||||||
|
|
||||||
if self.ep_size > 1 and not should_allreduce_fusion:
|
if (
|
||||||
|
self.ep_size > 1
|
||||||
|
and not should_allreduce_fusion
|
||||||
|
and not should_use_dp_reduce_scatterv()
|
||||||
|
):
|
||||||
final_hidden_states = moe_expert_parallel_all_reduce(final_hidden_states)
|
final_hidden_states = moe_expert_parallel_all_reduce(final_hidden_states)
|
||||||
|
|
||||||
if (
|
if (
|
||||||
@@ -339,6 +344,7 @@ class Qwen3MoeSparseMoeBlock(nn.Module):
|
|||||||
and not should_allreduce_fusion
|
and not should_allreduce_fusion
|
||||||
and not use_reduce_scatter
|
and not use_reduce_scatter
|
||||||
and not should_use_flashinfer_cutlass_moe_fp4_allgather()
|
and not should_use_flashinfer_cutlass_moe_fp4_allgather()
|
||||||
|
and not should_use_dp_reduce_scatterv()
|
||||||
):
|
):
|
||||||
final_hidden_states = moe_tensor_model_parallel_all_reduce(
|
final_hidden_states = moe_tensor_model_parallel_all_reduce(
|
||||||
final_hidden_states
|
final_hidden_states
|
||||||
|
|||||||
Reference in New Issue
Block a user