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

[refactor] Move MLP collective flags onto ForwardFlags (#30802)

Browse Source
This commit is contained in:
Cheng Wan
2026-07-10 17:43:39 -07:00
committed by GitHub
parent 7de33ce806
commit fc2ef35308
37 changed files with 469 additions and 546 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py
-2
View File
@@ -742,7 +742,6 @@ class Mamba2AttnBackend(MambaAttnBackendBase):
forward_batch: ForwardBatch, forward_batch: ForwardBatch,
mup_vector: Optional[torch.Tensor] = None, mup_vector: Optional[torch.Tensor] = None,
use_triton_causal_conv: bool = False, use_triton_causal_conv: bool = False,
should_allreduce_fusion: bool = False,
): ):
assert isinstance(self.forward_metadata, Mamba2Metadata) assert isinstance(self.forward_metadata, Mamba2Metadata)
# Page-major stores state strided; only the stride-aware Triton causal-conv # Page-major stores state strided; only the stride-aware Triton causal-conv
@@ -759,7 +758,6 @@ class Mamba2AttnBackend(MambaAttnBackendBase):
forward_batch=forward_batch, forward_batch=forward_batch,
mup_vector=mup_vector, mup_vector=mup_vector,
use_triton_causal_conv=use_triton_causal_conv, use_triton_causal_conv=use_triton_causal_conv,
should_allreduce_fusion=should_allreduce_fusion,
) )
if forward_batch.mamba_track_mask is not None: if forward_batch.mamba_track_mask is not None:
python/sglang/srt/layers/attention/mamba/mamba.py
+1 -4
View File
@@ -448,7 +448,6 @@ class MambaMixer2(torch.nn.Module):
forward_batch: ForwardBatch, forward_batch: ForwardBatch,
mup_vector: Optional[torch.Tensor] = None, mup_vector: Optional[torch.Tensor] = None,
use_triton_causal_conv: bool = False, use_triton_causal_conv: bool = False,
should_allreduce_fusion: bool = False,
): ):
# Returns the projected result. When `output` is given it is also # Returns the projected result. When `output` is given it is also
# written into that buffer (required by the cuda-graph split ops, which # written into that buffer (required by the cuda-graph split ops, which
@@ -761,9 +760,7 @@ class MambaMixer2(torch.nn.Module):
# norm usage # norm usage
hidden_states = self.norm(preallocated_ssm_out, gate) hidden_states = self.norm(preallocated_ssm_out, gate)
mixer_out, _ = self.out_proj( mixer_out, _ = self.out_proj(hidden_states)
hidden_states, skip_all_reduce=should_allreduce_fusion
)
if output is not None: if output is not None:
output[:padded_num_tokens].copy_(mixer_out) output[:padded_num_tokens].copy_(mixer_out)
python/sglang/srt/layers/communicator.py
+1 -1
View File
@@ -1383,7 +1383,7 @@ class CommunicateSummableTensorPairFn:
"""Scatter MoE output back to TP_ATTN_FULL after MOE_FULL computation. """Scatter MoE output back to TP_ATTN_FULL after MOE_FULL computation.
After moe_tensor_model_parallel_all_reduce (which runs unconditionally since After moe_tensor_model_parallel_all_reduce (which runs unconditionally since
use_reduce_scatter=False for this path), all ranks in the moe_cp group hold the mlp_reduce_scatter=False for this path), all ranks in the moe_cp group hold the
full MoE result for all cp_per_moe token chunks. We simply slice out this rank's full MoE result for all cp_per_moe token chunks. We simply slice out this rank's
CP-local portion. CP-local portion.
python/sglang/srt/layers/linear.py
+10 -1
View File
@@ -27,6 +27,7 @@ from sglang.srt.distributed.device_communicators.pynccl_allocator import (
from sglang.srt.layers.dp_attention import ( from sglang.srt.layers.dp_attention import (
is_allocation_symmetric, is_allocation_symmetric,
) )
from sglang.srt.layers.moe.utils import should_skip_mlp_all_reduce
from sglang.srt.layers.parameter import ( from sglang.srt.layers.parameter import (
BasevLLMParameter, BasevLLMParameter,
BlockQuantScaleParameter, BlockQuantScaleParameter,
@@ -1537,7 +1538,15 @@ class RowParallelLinear(LinearBase):
with symm_ctx: with symm_ctx:
output_parallel = self.quant_method.apply(self, input_parallel, bias=bias_) output_parallel = self.quant_method.apply(self, input_parallel, bias=bias_)
if self.reduce_results and self.tp_size > 1 and not skip_all_reduce: # skip_all_reduce: explicit call-site override. Also honor
# ForwardFlags (fuse_mlp_allreduce / mlp_reduce_scatter) published by
# the decoder — callers should not thread those flags into modules.
if (
self.reduce_results
and self.tp_size > 1
and not skip_all_reduce
and not should_skip_mlp_all_reduce()
):
if self.use_dp_attention_reduce: if self.use_dp_attention_reduce:
output = get_parallel().attn_tp_group.all_reduce(output_parallel) output = get_parallel().attn_tp_group.all_reduce(output_parallel)
else: else:
python/sglang/srt/layers/moe/__init__.py
+2
View File
@@ -10,6 +10,7 @@ from sglang.srt.layers.moe.utils import (
get_tbo_token_distribution_threshold, get_tbo_token_distribution_threshold,
initialize_moe_config, initialize_moe_config,
is_tbo_enabled, is_tbo_enabled,
should_skip_mlp_all_reduce,
should_skip_post_experts_all_reduce, should_skip_post_experts_all_reduce,
should_use_dp_reduce_scatterv, should_use_dp_reduce_scatterv,
should_use_flashinfer_cutlass_moe_fp4_allgather, should_use_flashinfer_cutlass_moe_fp4_allgather,
@@ -25,6 +26,7 @@ __all__ = [
"get_moe_a2a_backend", "get_moe_a2a_backend",
"get_moe_runner_backend", "get_moe_runner_backend",
"get_deepep_mode", "get_deepep_mode",
"should_skip_mlp_all_reduce",
"should_skip_post_experts_all_reduce", "should_skip_post_experts_all_reduce",
"should_use_dp_reduce_scatterv", "should_use_dp_reduce_scatterv",
"should_use_flashinfer_cutlass_moe_fp4_allgather", "should_use_flashinfer_cutlass_moe_fp4_allgather",
python/sglang/srt/layers/moe/utils.py
+21 -14
View File
@@ -12,7 +12,7 @@ from sglang.srt.environ import envs
from sglang.srt.layers.dp_attention import ( from sglang.srt.layers.dp_attention import (
is_dp_attention_enabled, is_dp_attention_enabled,
) )
from sglang.srt.runtime_context import get_flags, get_parallel from sglang.srt.runtime_context import get_flags, get_forward, get_parallel
from sglang.srt.utils import is_cuda, is_npu from sglang.srt.utils import is_cuda, is_npu
_is_npu = is_npu() _is_npu = is_npu()
@@ -411,20 +411,27 @@ def should_use_dp_reduce_scatterv():
) )
def should_skip_post_experts_all_reduce( def should_skip_mlp_all_reduce() -> bool:
*, """Whether dense MLP / row-parallel projections should skip their all-reduce.
is_tp_path: bool,
use_reduce_scatter: bool = False, True when the decoder published ``fuse_mlp_allreduce`` (next residual+LN
should_allreduce_fusion: bool = False, absorbs the AR) or ``mlp_reduce_scatter`` (postprocess will reduce-scatter)
) -> bool: on ``get_forward()``.
"""
f = get_forward()
return f.fuse_mlp_allreduce or f.mlp_reduce_scatter
def should_skip_post_experts_all_reduce(*, is_tp_path: bool) -> bool:
"""Whether to skip the post-experts all-reduce (EP or TP) because a """Whether to skip the post-experts all-reduce (EP or TP) because a
downstream component will fuse, replace, or absorb it. downstream component will fuse, replace, or absorb it.
Skip reasons, in order: Skip reasons, in order:
- ``should_allreduce_fusion``: LayerCommunicator will fuse the all-reduce - ``get_forward().fuse_mlp_allreduce``: LayerCommunicator will fuse the
with the next layer's residual all-reduce. all-reduce with the next layer's residual all-reduce.
- ``use_reduce_scatter``: LayerCommunicator's post-attention scatter will - ``get_forward().mlp_reduce_scatter``: LayerCommunicator's post-attention
do reduce-scatter, which would double-reduce on top of an all-reduce. scatter will do reduce-scatter, which would double-reduce on top of
an all-reduce.
- ``should_use_dp_reduce_scatterv()``: the standard dispatcher's combine - ``should_use_dp_reduce_scatterv()``: the standard dispatcher's combine
path replaces the all-reduce with a reduce-scatterv. path replaces the all-reduce with a reduce-scatterv.
- ``should_use_flashinfer_cutlass_moe_fp4_allgather()`` (TP path only): - ``should_use_flashinfer_cutlass_moe_fp4_allgather()`` (TP path only):
@@ -437,11 +444,11 @@ def should_skip_post_experts_all_reduce(
``not enable_alltoall`` gate ``not enable_alltoall`` gate
(``tensorrt_llm/_torch/modules/fused_moe/interface.py:879``). (``tensorrt_llm/_torch/modules/fused_moe/interface.py:879``).
The first two args are layer-context flags from ``LayerCommunicator`` and The first two reasons come from per-layer ``ForwardFlags`` published by
default to ``False`` for models that don't use it. Pass ``is_tp_path=True`` the decoder via ``get_forward().scoped(...)``. Pass ``is_tp_path=True``
for the post-experts TP all-reduce, ``False`` for the EP all-reduce. for the post-experts TP all-reduce, ``False`` for the EP all-reduce.
""" """
if should_allreduce_fusion or use_reduce_scatter: if should_skip_mlp_all_reduce():
return True return True
if should_use_dp_reduce_scatterv(): if should_use_dp_reduce_scatterv():
return True return True
python/sglang/srt/lora/layers.py
+2
View File
@@ -19,6 +19,7 @@ from sglang.srt.layers.linear import (
) )
from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE
from sglang.srt.layers.moe.topk import TopKOutput from sglang.srt.layers.moe.topk import TopKOutput
from sglang.srt.layers.moe.utils import should_skip_mlp_all_reduce
from sglang.srt.layers.vocab_parallel_embedding import ( from sglang.srt.layers.vocab_parallel_embedding import (
ParallelLMHead, ParallelLMHead,
VocabParallelEmbedding, VocabParallelEmbedding,
@@ -732,6 +733,7 @@ class RowParallelLinearWithLoRA(BaseLayerWithLoRA):
self.base_layer.reduce_results self.base_layer.reduce_results
and self.base_layer.tp_size > 1 and self.base_layer.tp_size > 1
and not skip_all_reduce and not skip_all_reduce
and not should_skip_mlp_all_reduce()
) )
if self.set_lora and should_reduce: if self.set_lora and should_reduce:
python/sglang/srt/lora/trtllm_lora_temp/attention.py
+2
View File
@@ -14,6 +14,7 @@ from sglang.srt.distributed import (
tensor_model_parallel_all_gather, tensor_model_parallel_all_gather,
tensor_model_parallel_all_reduce, tensor_model_parallel_all_reduce,
) )
from sglang.srt.layers.moe.utils import should_skip_mlp_all_reduce
from sglang.srt.lora.trtllm_lora_temp import ( from sglang.srt.lora.trtllm_lora_temp import (
get_lora_side_stream, get_lora_side_stream,
get_original_column_forward, get_original_column_forward,
@@ -125,6 +126,7 @@ def row_parallel_lora_forward(
self.base_layer.reduce_results self.base_layer.reduce_results
and self.base_layer.tp_size > 1 and self.base_layer.tp_size > 1
and not skip_all_reduce and not skip_all_reduce
and not should_skip_mlp_all_reduce()
) )
if should_reduce: if should_reduce:
python/sglang/srt/models/apertus.py
+1 -5
View File
@@ -96,15 +96,11 @@ class ApertusMLP(nn.Module):
self, self,
x, x,
forward_batch=None, forward_batch=None,
use_reduce_scatter: bool = False,
): ):
# note: with xielu, there's no gate_proj # note: with xielu, there's no gate_proj
x, _ = self.up_proj(x) x, _ = self.up_proj(x)
x = self.act_fn(x) x = self.act_fn(x)
x, _ = self.down_proj( x, _ = self.down_proj(x)
x,
skip_all_reduce=use_reduce_scatter,
)
return x return x
python/sglang/srt/models/bailing_moe.py
+16 -23
View File
@@ -77,7 +77,12 @@ from sglang.srt.models.utils import (
create_fused_set_kv_buffer_arg, create_fused_set_kv_buffer_arg,
enable_fused_set_kv_buffer, enable_fused_set_kv_buffer,
) )
from sglang.srt.runtime_context import get_parallel, get_server_args, get_stream from sglang.srt.runtime_context import (
get_forward,
get_parallel,
get_server_args,
get_stream,
)
from sglang.srt.utils import add_prefix, is_cuda, is_non_idle_and_non_empty, make_layers from sglang.srt.utils import add_prefix, is_cuda, is_non_idle_and_non_empty, make_layers
LoraConfig = None LoraConfig = None
@@ -127,17 +132,13 @@ class BailingMoEMLP(nn.Module):
self, self,
hidden_states: torch.Tensor, hidden_states: torch.Tensor,
forward_batch: Optional[ForwardBatch] = None, forward_batch: Optional[ForwardBatch] = None,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
) -> torch.Tensor: ) -> torch.Tensor:
if (self.tp_size == 1) and hidden_states.shape[0] == 0: if (self.tp_size == 1) and hidden_states.shape[0] == 0:
return hidden_states return hidden_states
gate_up, _ = self.gate_up_proj(hidden_states) gate_up, _ = self.gate_up_proj(hidden_states)
hidden_states = self.act_fn(gate_up) hidden_states = self.act_fn(gate_up)
hidden_states, _ = self.down_proj( hidden_states, _ = self.down_proj(hidden_states)
hidden_states, skip_all_reduce=should_allreduce_fusion or use_reduce_scatter
)
return hidden_states return hidden_states
@@ -305,15 +306,9 @@ class BailingMoESparseMoeBlock(nn.Module):
self, self,
hidden_states: torch.Tensor, hidden_states: torch.Tensor,
forward_batch: Optional[ForwardBatch] = None, forward_batch: Optional[ForwardBatch] = None,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
) -> torch.Tensor: ) -> torch.Tensor:
if not get_moe_a2a_backend().is_deepep(): if not get_moe_a2a_backend().is_deepep():
return self.forward_normal( return self.forward_normal(hidden_states)
hidden_states,
should_allreduce_fusion,
use_reduce_scatter,
)
else: else:
return self.forward_deepep(hidden_states, forward_batch) return self.forward_deepep(hidden_states, forward_batch)
@@ -356,8 +351,6 @@ class BailingMoESparseMoeBlock(nn.Module):
def forward_normal( def forward_normal(
self, self,
hidden_states: torch.Tensor, hidden_states: torch.Tensor,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
) -> torch.Tensor: ) -> torch.Tensor:
num_tokens, hidden_size = hidden_states.shape num_tokens, hidden_size = hidden_states.shape
hidden_states = hidden_states.view(-1, hidden_size) hidden_states = hidden_states.view(-1, hidden_size)
@@ -379,8 +372,6 @@ class BailingMoESparseMoeBlock(nn.Module):
if self.tp_size > 1 and not should_skip_post_experts_all_reduce( if self.tp_size > 1 and not should_skip_post_experts_all_reduce(
is_tp_path=True, is_tp_path=True,
use_reduce_scatter=use_reduce_scatter,
should_allreduce_fusion=should_allreduce_fusion,
): ):
final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states)
return final_hidden_states.view(num_tokens, hidden_size) return final_hidden_states.view(num_tokens, hidden_size)
@@ -672,22 +663,24 @@ class BailingMoEBlock(nn.Module):
forward_batch=forward_batch, forward_batch=forward_batch,
) )
should_allreduce_fusion = ( fuse_mlp_allreduce = (
self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer( self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer(
forward_batch forward_batch
) )
) )
# For DP with padding, reduce scatter can be used instead of all-reduce. # For DP with padding, reduce scatter can be used instead of all-reduce.
use_reduce_scatter = self.layer_communicator.should_use_reduce_scatter( mlp_reduce_scatter = self.layer_communicator.should_use_reduce_scatter(
forward_batch forward_batch
) )
hidden_states = self.mlp( with get_forward().scoped(
hidden_states, forward_batch, should_allreduce_fusion, use_reduce_scatter fuse_mlp_allreduce=fuse_mlp_allreduce,
) mlp_reduce_scatter=mlp_reduce_scatter,
):
hidden_states = self.mlp(hidden_states, forward_batch)
if should_allreduce_fusion: if fuse_mlp_allreduce:
hidden_states._sglang_needs_allreduce_fusion = True hidden_states._sglang_needs_allreduce_fusion = True
else: else:
hidden_states, residual = self.layer_communicator.postprocess_layer( hidden_states, residual = self.layer_communicator.postprocess_layer(
python/sglang/srt/models/bailing_moe_linear.py
+20 -19
View File
@@ -58,7 +58,12 @@ from sglang.srt.model_executor.runner import get_is_capture_mode
from sglang.srt.model_loader.weight_utils import default_weight_loader from sglang.srt.model_loader.weight_utils import default_weight_loader
from sglang.srt.models.deepseek_v2 import DeepseekV2AttentionMLA, DeepseekV2MLP, _is_hip from sglang.srt.models.deepseek_v2 import DeepseekV2AttentionMLA, DeepseekV2MLP, _is_hip
from sglang.srt.models.utils import WeightsMapper from sglang.srt.models.utils import WeightsMapper
from sglang.srt.runtime_context import get_parallel, get_server_args, get_stream from sglang.srt.runtime_context import (
get_forward,
get_parallel,
get_server_args,
get_stream,
)
from sglang.srt.utils import ( from sglang.srt.utils import (
BumpAllocator, BumpAllocator,
add_prefix, add_prefix,
@@ -189,15 +194,10 @@ class BailingMLP(nn.Module):
def forward( def forward(
self, self,
x, x,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
): ):
x, _ = self.gate_up_proj(x) x, _ = self.gate_up_proj(x)
x = self.act_fn(x) x = self.act_fn(x)
x, _ = self.down_proj( x, _ = self.down_proj(x)
x,
skip_all_reduce=use_reduce_scatter or should_allreduce_fusion,
)
return x return x
@@ -332,8 +332,6 @@ class BailingMoE(nn.Module):
def forward( def forward(
self, self,
hidden_states: torch.Tensor, hidden_states: torch.Tensor,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
) -> torch.Tensor: ) -> torch.Tensor:
num_tokens, hidden_size = hidden_states.shape num_tokens, hidden_size = hidden_states.shape
hidden_states = hidden_states.view(-1, hidden_size) hidden_states = hidden_states.view(-1, hidden_size)
@@ -369,8 +367,6 @@ class BailingMoE(nn.Module):
if self.tp_size > 1 and not should_skip_post_experts_all_reduce( if self.tp_size > 1 and not should_skip_post_experts_all_reduce(
is_tp_path=True, is_tp_path=True,
use_reduce_scatter=use_reduce_scatter,
should_allreduce_fusion=should_allreduce_fusion,
): ):
final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states)
return final_hidden_states return final_hidden_states
@@ -890,20 +886,25 @@ class BailingMoELinearDecoderLayer(nn.Module):
# logger.warning( # logger.warning(
# f"===={self.layer_id=}, 3 shape= {hidden_states.shape}, {residual.shape}" # f"===={self.layer_id=}, 3 shape= {hidden_states.shape}, {residual.shape}"
# ) # )
should_allreduce_fusion = ( fuse_mlp_allreduce = (
self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer( self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer(
forward_batch forward_batch
) )
) )
use_reduce_scatter = self.layer_communicator.should_use_reduce_scatter( mlp_reduce_scatter = self.layer_communicator.should_use_reduce_scatter(
forward_batch forward_batch
) )
hidden_states = self.mlp( with get_forward().scoped(
hidden_states, should_allreduce_fusion, use_reduce_scatter fuse_mlp_allreduce=fuse_mlp_allreduce,
) mlp_reduce_scatter=mlp_reduce_scatter,
hidden_states, residual = self.layer_communicator.postprocess_layer( ):
hidden_states, residual, forward_batch hidden_states = self.mlp(hidden_states)
) if fuse_mlp_allreduce:
hidden_states._sglang_needs_allreduce_fusion = True
else:
hidden_states, residual = self.layer_communicator.postprocess_layer(
hidden_states, residual, forward_batch
)
return hidden_states, residual return hidden_states, residual
@staticmethod @staticmethod
python/sglang/srt/models/deepseek_v2.py
+27 -52
View File
@@ -182,7 +182,12 @@ from sglang.srt.models.deepseek_common.utils import (
_use_aiter_bpreshuffle_gfx95, _use_aiter_bpreshuffle_gfx95,
_use_aiter_gfx95, _use_aiter_gfx95,
) )
from sglang.srt.runtime_context import get_flags, get_parallel, get_server_args from sglang.srt.runtime_context import (
get_flags,
get_forward,
get_parallel,
get_server_args,
)
from sglang.srt.speculative.spec_info import SpeculativeAlgorithm from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
from sglang.srt.utils import ( from sglang.srt.utils import (
BumpAllocator, BumpAllocator,
@@ -292,8 +297,6 @@ class DeepseekV2MLP(nn.Module):
self, self,
x, x,
forward_batch=None, forward_batch=None,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
gemm_output_zero_allocator: BumpAllocator = None, gemm_output_zero_allocator: BumpAllocator = None,
): ):
if (self.tp_size == 1) and x.shape[0] == 0: if (self.tp_size == 1) and x.shape[0] == 0:
@@ -321,10 +324,7 @@ class DeepseekV2MLP(nn.Module):
self.down_proj.input_scale_inv, self.down_proj.input_scale_inv,
enable_pdl=True, enable_pdl=True,
) )
out, _ = self.down_proj( out, _ = self.down_proj((out_fp4, out_scale))
(out_fp4, out_scale),
skip_all_reduce=should_allreduce_fusion or use_reduce_scatter,
)
return out return out
if ( if (
@@ -426,10 +426,7 @@ class DeepseekV2MLP(nn.Module):
silu_and_mul_clamp(gate_up, x, float(self.swiglu_limit)) silu_and_mul_clamp(gate_up, x, float(self.swiglu_limit))
else: else:
x = self.act_fn(gate_up) x = self.act_fn(gate_up)
x, _ = self.down_proj( x, _ = self.down_proj(x)
x,
skip_all_reduce=should_allreduce_fusion or use_reduce_scatter,
)
return x return x
@@ -857,8 +854,6 @@ class DeepseekV2MoE(nn.Module):
self, self,
hidden_states: torch.Tensor, hidden_states: torch.Tensor,
forward_batch: Optional[ForwardBatch] = None, forward_batch: Optional[ForwardBatch] = None,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
gemm_output_zero_allocator: BumpAllocator = None, gemm_output_zero_allocator: BumpAllocator = None,
input_ids: Optional[torch.Tensor] = None, input_ids: Optional[torch.Tensor] = None,
input_ids_global: Optional[torch.Tensor] = None, input_ids_global: Optional[torch.Tensor] = None,
@@ -880,8 +875,6 @@ class DeepseekV2MoE(nn.Module):
return dsv2_flashinfer_moe_dual_stream_graph( return dsv2_flashinfer_moe_dual_stream_graph(
hidden_states, hidden_states,
self.layer_id, self.layer_id,
should_allreduce_fusion,
use_reduce_scatter,
) )
elif ( elif (
self.alt_stream is not None self.alt_stream is not None
@@ -897,8 +890,6 @@ class DeepseekV2MoE(nn.Module):
): ):
return self.forward_normal_dual_stream( return self.forward_normal_dual_stream(
hidden_states, hidden_states,
should_allreduce_fusion,
use_reduce_scatter,
gemm_output_zero_allocator, gemm_output_zero_allocator,
input_ids, input_ids,
input_ids_global=input_ids_global, input_ids_global=input_ids_global,
@@ -906,8 +897,6 @@ class DeepseekV2MoE(nn.Module):
else: else:
return self.forward_normal( return self.forward_normal(
hidden_states, hidden_states,
should_allreduce_fusion,
use_reduce_scatter,
gemm_output_zero_allocator, gemm_output_zero_allocator,
input_ids, input_ids,
input_ids_global=input_ids_global, input_ids_global=input_ids_global,
@@ -921,19 +910,16 @@ class DeepseekV2MoE(nn.Module):
def forward_normal_dual_stream( def forward_normal_dual_stream(
self, self,
hidden_states: torch.Tensor, hidden_states: torch.Tensor,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
gemm_output_zero_allocator: BumpAllocator = None, gemm_output_zero_allocator: BumpAllocator = None,
input_ids: Optional[torch.Tensor] = None, input_ids: Optional[torch.Tensor] = None,
input_ids_global: Optional[torch.Tensor] = None, input_ids_global: Optional[torch.Tensor] = None,
*,
use_flashinfer_trtllm_bypass: bool = False,
) -> torch.Tensor: ) -> torch.Tensor:
# Note(kpham-sgl): issue order satisfies 3 constraints: # Note(kpham-sgl): issue order satisfies 3 constraints:
# - no stream explosion: main (routed) issued before alt block -> capture reuses 1 alt stream; # - no stream explosion: main (routed) issued before alt block -> capture reuses 1 alt stream;
# - PDL overlap: routed is the last main-stream kernel (fuses w/ residual add); # - PDL overlap: routed is the last main-stream kernel (fuses w/ residual add);
# - dispose_tensor: disabled during capture (CaptureFlags.disable_dispose_tensor) so the routed # - dispose_tensor: disabled during capture (CaptureFlags.disable_dispose_tensor) so the routed
# deep_gemm does not free hidden_states, which the shared expert reads on the alt stream. # deep_gemm does not free hidden_states, which the shared expert reads on the alt stream.
use_flashinfer_trtllm_bypass = get_forward().flashinfer_trtllm_bypass
current_stream = torch.cuda.current_stream() current_stream = torch.cuda.current_stream()
self.alt_stream.wait_stream(current_stream) self.alt_stream.wait_stream(current_stream)
has_shared_output = ( has_shared_output = (
@@ -1014,8 +1000,6 @@ class DeepseekV2MoE(nn.Module):
if self.tp_size > 1 and not should_skip_post_experts_all_reduce( if self.tp_size > 1 and not should_skip_post_experts_all_reduce(
is_tp_path=True, is_tp_path=True,
use_reduce_scatter=use_reduce_scatter,
should_allreduce_fusion=should_allreduce_fusion,
): ):
final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states)
# TP1 shared experts are replicated, so add them after all-reduce to # TP1 shared experts are replicated, so add them after all-reduce to
@@ -1027,8 +1011,6 @@ class DeepseekV2MoE(nn.Module):
def forward_normal( def forward_normal(
self, self,
hidden_states: torch.Tensor, hidden_states: torch.Tensor,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
gemm_output_zero_allocator: BumpAllocator = None, gemm_output_zero_allocator: BumpAllocator = None,
input_ids: Optional[torch.Tensor] = None, input_ids: Optional[torch.Tensor] = None,
input_ids_global: Optional[torch.Tensor] = None, input_ids_global: Optional[torch.Tensor] = None,
@@ -1037,7 +1019,7 @@ class DeepseekV2MoE(nn.Module):
if hasattr(self, "shared_experts") and use_intel_amx_backend( if hasattr(self, "shared_experts") and use_intel_amx_backend(
self.shared_experts.gate_up_proj self.shared_experts.gate_up_proj
): ):
return self.forward_cpu(hidden_states, should_allreduce_fusion) return self.forward_cpu(hidden_states)
server_args = get_server_args() server_args = get_server_args()
dispatch_info = ( dispatch_info = (
ExpertLocationDispatchInfo.init_new(layer_id=self.layer_id) ExpertLocationDispatchInfo.init_new(layer_id=self.layer_id)
@@ -1140,8 +1122,6 @@ class DeepseekV2MoE(nn.Module):
if self.tp_size > 1 and not should_skip_post_experts_all_reduce( if self.tp_size > 1 and not should_skip_post_experts_all_reduce(
is_tp_path=True, is_tp_path=True,
use_reduce_scatter=use_reduce_scatter,
should_allreduce_fusion=should_allreduce_fusion,
): ):
final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states)
# TP1 shared experts are replicated, so add them after all-reduce to # TP1 shared experts are replicated, so add them after all-reduce to
@@ -1153,7 +1133,6 @@ class DeepseekV2MoE(nn.Module):
def forward_cpu( def forward_cpu(
self, self,
hidden_states: torch.Tensor, hidden_states: torch.Tensor,
should_allreduce_fusion: bool = False,
) -> torch.Tensor: ) -> torch.Tensor:
# router_logits: (num_tokens, n_experts) # router_logits: (num_tokens, n_experts)
router_logits = self.gate(hidden_states) router_logits = self.gate(hidden_states)
@@ -1202,7 +1181,7 @@ class DeepseekV2MoE(nn.Module):
), # block_size ), # block_size
True, # is_vnni True, # is_vnni
) )
if self.tp_size > 1 and not should_allreduce_fusion: if self.tp_size > 1 and not get_forward().fuse_mlp_allreduce:
final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states)
return final_hidden_states return final_hidden_states
@@ -2259,14 +2238,14 @@ class DeepseekV2DecoderLayer(nn.Module):
hidden_states, residual, forward_batch hidden_states, residual, forward_batch
) )
should_allreduce_fusion = ( fuse_mlp_allreduce = (
self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer( self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer(
forward_batch forward_batch
) )
) )
# For DP with padding, reduce scatter can be used instead of all-reduce. # For DP with padding, reduce scatter can be used instead of all-reduce.
use_reduce_scatter = self.layer_communicator.should_use_reduce_scatter( mlp_reduce_scatter = self.layer_communicator.should_use_reduce_scatter(
forward_batch forward_batch
) )
@@ -2284,22 +2263,24 @@ class DeepseekV2DecoderLayer(nn.Module):
else: else:
_mlp_ctx = nullcontext() _mlp_ctx = nullcontext()
with _mlp_ctx: with get_forward().scoped(
hidden_states = self.mlp( fuse_mlp_allreduce=fuse_mlp_allreduce,
hidden_states, mlp_reduce_scatter=mlp_reduce_scatter,
forward_batch, ):
should_allreduce_fusion, with _mlp_ctx:
use_reduce_scatter, hidden_states = self.mlp(
gemm_output_zero_allocator, hidden_states,
) forward_batch,
gemm_output_zero_allocator,
)
if ( if (
not (self.dsa_enable_prefill_cp or self.mla_enable_prefill_cp) not (self.dsa_enable_prefill_cp or self.mla_enable_prefill_cp)
and should_allreduce_fusion and fuse_mlp_allreduce
): ):
hidden_states._sglang_needs_allreduce_fusion = True hidden_states._sglang_needs_allreduce_fusion = True
if not should_allreduce_fusion: if not fuse_mlp_allreduce:
hidden_states, residual = self.layer_communicator.postprocess_layer( hidden_states, residual = self.layer_communicator.postprocess_layer(
hidden_states, residual, forward_batch hidden_states, residual, forward_batch
) )
@@ -2996,8 +2977,6 @@ class DeepseekV32ForCausalLM(DeepseekV2ForCausalLM):
def dsv2_flashinfer_moe_dual_stream_graph( def dsv2_flashinfer_moe_dual_stream_graph(
hidden_states: torch.Tensor, hidden_states: torch.Tensor,
layer_id: int, layer_id: int,
should_allreduce_fusion: bool,
use_reduce_scatter: bool,
) -> torch.Tensor: ) -> torch.Tensor:
forward_context = get_tc_piecewise_forward_context() forward_context = get_tc_piecewise_forward_context()
assert forward_context is not None assert forward_context is not None
@@ -3005,12 +2984,8 @@ def dsv2_flashinfer_moe_dual_stream_graph(
moe_fusion = forward_context.moe_fusions[layer_id] moe_fusion = forward_context.moe_fusions[layer_id]
assert moe_fusion is not None assert moe_fusion is not None
return moe_fusion.forward_normal_dual_stream( with get_forward().scoped(flashinfer_trtllm_bypass=True):
hidden_states, return moe_fusion.forward_normal_dual_stream(hidden_states)
should_allreduce_fusion=should_allreduce_fusion,
use_reduce_scatter=use_reduce_scatter,
use_flashinfer_trtllm_bypass=True,
)
EntryClass = [DeepseekV2ForCausalLM, DeepseekV3ForCausalLM, DeepseekV32ForCausalLM] EntryClass = [DeepseekV2ForCausalLM, DeepseekV3ForCausalLM, DeepseekV32ForCausalLM]
python/sglang/srt/models/deepseek_v4.py
+25 -25
View File
@@ -122,7 +122,7 @@ from sglang.srt.models.deepseek_v2 import (
_is_npu, _is_npu,
_is_xpu, _is_xpu,
) )
from sglang.srt.runtime_context import get_parallel, get_server_args from sglang.srt.runtime_context import get_forward, get_parallel, get_server_args
if not _is_hip: if not _is_hip:
from sglang.srt.layers.utils.cp_utils import ( from sglang.srt.layers.utils.cp_utils import (
@@ -134,7 +134,6 @@ if _is_xpu:
else: else:
from sglang.srt.layers.mhc import hc_split_sinkhorn, mhc_fused_post_pre, npu_hc_pre from sglang.srt.layers.mhc import hc_split_sinkhorn, mhc_fused_post_pre, npu_hc_pre
from sglang.srt.runtime_context import get_server_args
from sglang.srt.utils import ( from sglang.srt.utils import (
LazyValue, LazyValue,
add_prefix, add_prefix,
@@ -1639,7 +1638,7 @@ class DeepseekV4DecoderLayer(nn.Module):
# The experts ARE TP-sharded by intermediate (moe_tp_size==tp_size), so # The experts ARE TP-sharded by intermediate (moe_tp_size==tp_size), so
# the post-experts reduce is a SUM. reduce_scatterv does that sum+scatter # the post-experts reduce is a SUM. reduce_scatterv does that sum+scatter
# in ONE op, REPLACING the MoE-internal post-experts all_reduce — so we # in ONE op, REPLACING the MoE-internal post-experts all_reduce — so we
# MUST tell the MoE to skip it (use_reduce_scatter=True) or it # MUST tell the MoE to skip it (mlp_reduce_scatter=True) or it
# double-reduces. Env-gated via SGLANG_DP_USE_GATHERV, default OFF. # double-reduces. Env-gated via SGLANG_DP_USE_GATHERV, default OFF.
_use_reduce_scatterv = ( _use_reduce_scatterv = (
_use_tp_moe_gather _use_tp_moe_gather
@@ -1663,6 +1662,7 @@ class DeepseekV4DecoderLayer(nn.Module):
and forward_batch.dp_padding_mode.is_max_len() and forward_batch.dp_padding_mode.is_max_len()
and get_parallel().tp_size == get_parallel().attn_dp_size and get_parallel().tp_size == get_parallel().attn_dp_size
) )
mlp_reduce_scatter = _use_cp or _use_reduce_scatterv or _use_reduce_scatter
# PoC (SGLANG_DP_SHARED_EXPERT_LOCAL): compute the replicated shared expert # PoC (SGLANG_DP_SHARED_EXPERT_LOCAL): compute the replicated shared expert
# on LOCAL hidden before the gather and add it back after the combine # on LOCAL hidden before the gather and add it back after the combine
# (reduce_scatterv OR dp_scatter), instead of on the gathered global buffer. # (reduce_scatterv OR dp_scatter), instead of on the gathered global buffer.
@@ -1702,17 +1702,17 @@ class DeepseekV4DecoderLayer(nn.Module):
hidden_states = _a2a_scatter_chunks[r].contiguous() hidden_states = _a2a_scatter_chunks[r].contiguous()
input_ids = input_ids.tensor_split(s)[r].contiguous() input_ids = input_ids.tensor_split(s)[r].contiguous()
input_ids_global = input_ids_global.tensor_split(s)[r].contiguous() input_ids_global = input_ids_global.tensor_split(s)[r].contiguous()
hidden_states = self.mlp( # Skip the MoE-internal post-experts all_reduce when we will do the
hidden_states, # reduce via reduce_scatterv/reduce_scatter at the combine below
forward_batch, # (else double-reduce).
input_ids=input_ids, with get_forward().scoped(mlp_reduce_scatter=mlp_reduce_scatter):
input_ids_global=input_ids_global, hidden_states = self.mlp(
# Skip the MoE-internal post-experts all_reduce when we will do the hidden_states,
# reduce via reduce_scatterv/reduce_scatter at the combine below forward_batch,
# (else double-reduce). input_ids=input_ids,
use_reduce_scatter=_use_cp or _use_reduce_scatterv or _use_reduce_scatter, input_ids_global=input_ids_global,
skip_shared_experts=_do_shared_local, skip_shared_experts=_do_shared_local,
) )
if _use_cp and get_moe_a2a_backend().is_none(): if _use_cp and get_moe_a2a_backend().is_none():
hidden_states = dsa_cp_reduce_scatter_hidden_states(hidden_states) hidden_states = dsa_cp_reduce_scatter_hidden_states(hidden_states)
elif _use_tp_moe_gather: elif _use_tp_moe_gather:
@@ -1723,7 +1723,7 @@ class DeepseekV4DecoderLayer(nn.Module):
if should_use_dp_reduce_scatterv() or _use_reduce_scatterv: if should_use_dp_reduce_scatterv() or _use_reduce_scatterv:
# SUM the TP-sharded per-rank partial expert outputs AND scatter # SUM the TP-sharded per-rank partial expert outputs AND scatter
# each rank its own token slice, in one op. Correct because the # each rank its own token slice, in one op. Correct because the
# MoE-internal all_reduce was skipped (use_reduce_scatter above). # MoE-internal all_reduce was skipped (mlp_reduce_scatter above).
# This is the symmetric inverse of the all_gatherv gather. # This is the symmetric inverse of the all_gatherv gather.
get_tp_group().reduce_scatterv( get_tp_group().reduce_scatterv(
global_hidden_states, global_hidden_states,
@@ -1735,7 +1735,7 @@ class DeepseekV4DecoderLayer(nn.Module):
# expert outputs AND scatter each rank its own (MAX_LEN-padded) # expert outputs AND scatter each rank its own (MAX_LEN-padded)
# token chunk in one op (symmetric inverse of the MAX_LEN # token chunk in one op (symmetric inverse of the MAX_LEN
# all_gather). Correct because the MoE-internal all_reduce was # all_gather). Correct because the MoE-internal all_reduce was
# skipped (use_reduce_scatter above). dp_reduce_scatter_tensor # skipped (mlp_reduce_scatter above). dp_reduce_scatter_tensor
# routes to the equal-chunk reduce_scatter_tensor here (its # routes to the equal-chunk reduce_scatter_tensor here (its
# variable-length reduce_scatterv branch is gated by # variable-length reduce_scatterv branch is gated by
# is_dp_gatherv_active(), which is False under MAX_LEN), which in # is_dp_gatherv_active(), which is False under MAX_LEN), which in
@@ -1920,19 +1920,19 @@ class DeepseekV4DecoderLayer(nn.Module):
state.pop("gather_keepalive") state.pop("gather_keepalive")
def op_moe(self, state): def op_moe(self, state):
# MoE (gate/topk/experts) on the GLOBAL gathered buffer. use_reduce_scatter # MoE (gate/topk/experts) on the GLOBAL gathered buffer. mlp_reduce_scatter
# skips the MoE-internal all_reduce (we reduce_scatterv in op_combine). # skips the MoE-internal all_reduce (we reduce_scatterv in op_combine).
fb = state.forward_batch fb = state.forward_batch
global_hidden = state.pop("global_hidden") global_hidden = state.pop("global_hidden")
global_ids = fb._tbo_global_input_ids global_ids = fb._tbo_global_input_ids
state.global_expert_out = self.mlp( with get_forward().scoped(mlp_reduce_scatter=True):
global_hidden, state.global_expert_out = self.mlp(
fb, global_hidden,
use_reduce_scatter=True, fb,
input_ids=global_ids, input_ids=global_ids,
input_ids_global=global_ids, input_ids_global=global_ids,
skip_shared_experts=state.do_shared_local, skip_shared_experts=state.do_shared_local,
) )
def op_combine_a(self, state): def op_combine_a(self, state):
# Launch reduce_scatterv (global partial expert sums -> per-rank local) on # Launch reduce_scatterv (global partial expert sums -> per-rank local) on
python/sglang/srt/models/exaone_moe.py
+1 -8
View File
@@ -121,15 +121,10 @@ class ExaoneMoEMLP(nn.Module):
self, self,
x, x,
forward_batch=None, forward_batch=None,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
): ):
gate_up, _ = self.gate_up_proj(x) gate_up, _ = self.gate_up_proj(x)
x = self.act_fn(gate_up) x = self.act_fn(gate_up)
x, _ = self.down_proj( x, _ = self.down_proj(x)
x,
skip_all_reduce=should_allreduce_fusion or use_reduce_scatter,
)
return x return x
@@ -276,7 +271,6 @@ class ExaoneMoESparseMoEBlock(nn.Module):
self, self,
hidden_states: torch.Tensor, hidden_states: torch.Tensor,
forward_batch: Optional[ForwardBatch] = None, forward_batch: Optional[ForwardBatch] = None,
use_reduce_scatter: bool = False,
) -> torch.Tensor: ) -> torch.Tensor:
num_tokens, hidden_dim = hidden_states.shape num_tokens, hidden_dim = hidden_states.shape
hidden_states = hidden_states.view(-1, hidden_dim) hidden_states = hidden_states.view(-1, hidden_dim)
@@ -300,7 +294,6 @@ class ExaoneMoESparseMoEBlock(nn.Module):
final_hidden_states = final_hidden_states + shared_output final_hidden_states = final_hidden_states + shared_output
if self.tp_size > 1 and not should_skip_post_experts_all_reduce( if self.tp_size > 1 and not should_skip_post_experts_all_reduce(
is_tp_path=True, is_tp_path=True,
use_reduce_scatter=use_reduce_scatter,
): ):
final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states)
python/sglang/srt/models/falcon_h1.py
+10 -10
View File
@@ -33,7 +33,12 @@ from sglang.srt.layers.vocab_parallel_embedding import (
from sglang.srt.model_executor.forward_batch_info import ForwardBatch from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.model_executor.forward_context import get_attn_backend from sglang.srt.model_executor.forward_context import get_attn_backend
from sglang.srt.model_loader.weight_utils import default_weight_loader from sglang.srt.model_loader.weight_utils import default_weight_loader
from sglang.srt.runtime_context import get_parallel, get_server_args, get_stream from sglang.srt.runtime_context import (
get_forward,
get_parallel,
get_server_args,
get_stream,
)
from sglang.srt.utils import add_prefix, is_cuda, make_layers from sglang.srt.utils import add_prefix, is_cuda, make_layers
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@@ -85,16 +90,12 @@ class FalconH1MLP(nn.Module):
self, self,
x, x,
forward_batch=None, forward_batch=None,
use_reduce_scatter: bool = False,
): ):
gate_up, _ = self.gate_up_proj(x) gate_up, _ = self.gate_up_proj(x)
gate_up[:, : self.intermediate_size // self.tp_size] *= self.gate_multiplier gate_up[:, : self.intermediate_size // self.tp_size] *= self.gate_multiplier
x = self.act_fn(gate_up) x = self.act_fn(gate_up)
x, _ = self.down_proj( x, _ = self.down_proj(x)
x,
skip_all_reduce=use_reduce_scatter,
)
x = x * self.down_multiplier x = x * self.down_multiplier
return x return x
@@ -358,12 +359,11 @@ class FalconH1HybridAttentionDecoderLayer(nn.Module):
hidden_states, residual = self.layer_communicator.prepare_mlp( hidden_states, residual = self.layer_communicator.prepare_mlp(
hidden_states, residual, forward_batch hidden_states, residual, forward_batch
) )
use_reduce_scatter = self.layer_communicator.should_use_reduce_scatter( mlp_reduce_scatter = self.layer_communicator.should_use_reduce_scatter(
forward_batch forward_batch
) )
hidden_states = self.feed_forward( with get_forward().scoped(mlp_reduce_scatter=mlp_reduce_scatter):
hidden_states, forward_batch, use_reduce_scatter hidden_states = self.feed_forward(hidden_states, forward_batch)
)
hidden_states, residual = self.layer_communicator.postprocess_layer( hidden_states, residual = self.layer_communicator.postprocess_layer(
hidden_states, residual, forward_batch hidden_states, residual, forward_batch
python/sglang/srt/models/glm4.py
+1 -5
View File
@@ -94,14 +94,10 @@ class Glm4MLP(nn.Module):
self, self,
x, x,
forward_batch=None, forward_batch=None,
use_reduce_scatter: bool = False,
): ):
gate_up, _ = self.gate_up_proj(x) gate_up, _ = self.gate_up_proj(x)
x = self.act_fn(gate_up) x = self.act_fn(gate_up)
x, _ = self.down_proj( x, _ = self.down_proj(x)
x,
skip_all_reduce=use_reduce_scatter,
)
return x return x
python/sglang/srt/models/glm4_moe.py
+12 -31
View File
@@ -84,6 +84,7 @@ from sglang.srt.models.deepseek_nextn import DeepseekV3ForCausalLMNextN
from sglang.srt.models.deepseek_v2 import DeepseekV2ForCausalLM from sglang.srt.models.deepseek_v2 import DeepseekV2ForCausalLM
from sglang.srt.models.utils import WeightsMapper, apply_qk_norm from sglang.srt.models.utils import WeightsMapper, apply_qk_norm
from sglang.srt.runtime_context import ( from sglang.srt.runtime_context import (
get_forward,
get_parallel, get_parallel,
get_server_args, get_server_args,
get_stream, get_stream,
@@ -167,17 +168,13 @@ class Glm4MoeMLP(nn.Module):
self, self,
x, x,
forward_batch=None, forward_batch=None,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
): ):
if (self.tp_size == 1) and x.shape[0] == 0: if (self.tp_size == 1) and x.shape[0] == 0:
return x return x
gate_up, _ = self.gate_up_proj(x) gate_up, _ = self.gate_up_proj(x)
x = self.act_fn(gate_up) x = self.act_fn(gate_up)
x, _ = self.down_proj( x, _ = self.down_proj(x)
x, skip_all_reduce=should_allreduce_fusion or use_reduce_scatter
)
return x return x
@@ -564,8 +561,6 @@ class Glm4MoeSparseMoeBlock(nn.Module):
self, self,
hidden_states: torch.Tensor, hidden_states: torch.Tensor,
forward_batch: Optional[ForwardBatch] = None, forward_batch: Optional[ForwardBatch] = None,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
) -> torch.Tensor: ) -> torch.Tensor:
if not self._enable_a2a_moe: if not self._enable_a2a_moe:
if ( if (
@@ -574,25 +569,15 @@ class Glm4MoeSparseMoeBlock(nn.Module):
and hidden_states.shape[0] > 0 and hidden_states.shape[0] > 0
and get_is_capture_mode() and get_is_capture_mode()
): ):
return self.forward_normal_dual_stream( return self.forward_normal_dual_stream(hidden_states)
hidden_states,
should_allreduce_fusion,
use_reduce_scatter,
)
else: else:
return self.forward_normal( return self.forward_normal(hidden_states)
hidden_states,
should_allreduce_fusion,
use_reduce_scatter,
)
else: else:
return self.forward_deepep(hidden_states, forward_batch) return self.forward_deepep(hidden_states, forward_batch)
def forward_normal_dual_stream( def forward_normal_dual_stream(
self, self,
hidden_states: torch.Tensor, hidden_states: torch.Tensor,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
) -> torch.Tensor: ) -> torch.Tensor:
current_stream = torch.cuda.current_stream() current_stream = torch.cuda.current_stream()
self.alt_stream.wait_stream(current_stream) self.alt_stream.wait_stream(current_stream)
@@ -610,8 +595,6 @@ class Glm4MoeSparseMoeBlock(nn.Module):
final_hidden_states += shared_output final_hidden_states += shared_output
if self.tp_size > 1 and not should_skip_post_experts_all_reduce( if self.tp_size > 1 and not should_skip_post_experts_all_reduce(
is_tp_path=True, is_tp_path=True,
use_reduce_scatter=use_reduce_scatter,
should_allreduce_fusion=should_allreduce_fusion,
): ):
final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states)
return final_hidden_states return final_hidden_states
@@ -619,8 +602,6 @@ class Glm4MoeSparseMoeBlock(nn.Module):
def forward_normal( def forward_normal(
self, self,
hidden_states: torch.Tensor, hidden_states: torch.Tensor,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
) -> torch.Tensor: ) -> torch.Tensor:
if hidden_states.shape[0] > 0: if hidden_states.shape[0] > 0:
shared_output = self._forward_shared_experts(hidden_states) shared_output = self._forward_shared_experts(hidden_states)
@@ -643,8 +624,6 @@ class Glm4MoeSparseMoeBlock(nn.Module):
final_hidden_states = final_hidden_states_out final_hidden_states = final_hidden_states_out
if self.tp_size > 1 and not should_skip_post_experts_all_reduce( if self.tp_size > 1 and not should_skip_post_experts_all_reduce(
is_tp_path=True, is_tp_path=True,
use_reduce_scatter=use_reduce_scatter,
should_allreduce_fusion=should_allreduce_fusion,
): ):
final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states)
return final_hidden_states return final_hidden_states
@@ -958,22 +937,24 @@ class Glm4MoeDecoderLayer(nn.Module):
hidden_states, residual, forward_batch hidden_states, residual, forward_batch
) )
should_allreduce_fusion = ( fuse_mlp_allreduce = (
self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer( self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer(
forward_batch forward_batch
) )
) )
# For DP with padding, reduce scatter can be used instead of all-reduce. # For DP with padding, reduce scatter can be used instead of all-reduce.
use_reduce_scatter = self.layer_communicator.should_use_reduce_scatter( mlp_reduce_scatter = self.layer_communicator.should_use_reduce_scatter(
forward_batch forward_batch
) )
hidden_states = self.mlp( with get_forward().scoped(
hidden_states, forward_batch, should_allreduce_fusion, use_reduce_scatter fuse_mlp_allreduce=fuse_mlp_allreduce,
) mlp_reduce_scatter=mlp_reduce_scatter,
):
hidden_states = self.mlp(hidden_states, forward_batch)
if should_allreduce_fusion: if fuse_mlp_allreduce:
hidden_states._sglang_needs_allreduce_fusion = True hidden_states._sglang_needs_allreduce_fusion = True
else: else:
hidden_states, residual = self.layer_communicator.postprocess_layer( hidden_states, residual = self.layer_communicator.postprocess_layer(
python/sglang/srt/models/glm4_moe_lite.py
+17 -28
View File
@@ -74,7 +74,12 @@ from sglang.srt.models.deepseek_common.deepseek_weight_loader import (
) )
from sglang.srt.models.deepseek_common.utils import _is_cuda, _use_aiter from sglang.srt.models.deepseek_common.utils import _is_cuda, _use_aiter
from sglang.srt.models.deepseek_v2 import DeepseekV2AttentionMLA from sglang.srt.models.deepseek_v2 import DeepseekV2AttentionMLA
from sglang.srt.runtime_context import get_parallel, get_server_args, get_stream from sglang.srt.runtime_context import (
get_forward,
get_parallel,
get_server_args,
get_stream,
)
from sglang.srt.utils import ( from sglang.srt.utils import (
BumpAllocator, BumpAllocator,
LazyValue, LazyValue,
@@ -132,17 +137,13 @@ class Glm4MoeLiteMLP(nn.Module):
self, self,
x, x,
forward_batch=None, forward_batch=None,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
): ):
if (self.tp_size == 1) and x.shape[0] == 0: if (self.tp_size == 1) and x.shape[0] == 0:
return x return x
gate_up, _ = self.gate_up_proj(x) gate_up, _ = self.gate_up_proj(x)
x = self.act_fn(gate_up) x = self.act_fn(gate_up)
x, _ = self.down_proj( x, _ = self.down_proj(x)
x, skip_all_reduce=should_allreduce_fusion or use_reduce_scatter
)
return x return x
@@ -313,8 +314,6 @@ class Glm4MoeLiteSparseMoeBlock(nn.Module):
self, self,
hidden_states: torch.Tensor, hidden_states: torch.Tensor,
forward_batch: Optional[ForwardBatch] = None, forward_batch: Optional[ForwardBatch] = None,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
) -> torch.Tensor: ) -> torch.Tensor:
if not self._enable_a2a_moe: if not self._enable_a2a_moe:
if ( if (
@@ -323,21 +322,15 @@ class Glm4MoeLiteSparseMoeBlock(nn.Module):
and hidden_states.shape[0] > 0 and hidden_states.shape[0] > 0
and get_is_capture_mode() and get_is_capture_mode()
): ):
return self.forward_normal_dual_stream( return self.forward_normal_dual_stream(hidden_states)
hidden_states, should_allreduce_fusion, use_reduce_scatter
)
else: else:
return self.forward_normal( return self.forward_normal(hidden_states)
hidden_states, should_allreduce_fusion, use_reduce_scatter
)
else: else:
return self.forward_deepep(hidden_states, forward_batch) return self.forward_deepep(hidden_states, forward_batch)
def forward_normal_dual_stream( def forward_normal_dual_stream(
self, self,
hidden_states: torch.Tensor, hidden_states: torch.Tensor,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
) -> torch.Tensor: ) -> torch.Tensor:
current_stream = torch.cuda.current_stream() current_stream = torch.cuda.current_stream()
self.alt_stream.wait_stream(current_stream) self.alt_stream.wait_stream(current_stream)
@@ -355,8 +348,6 @@ class Glm4MoeLiteSparseMoeBlock(nn.Module):
final_hidden_states += shared_output final_hidden_states += shared_output
if self.tp_size > 1 and not should_skip_post_experts_all_reduce( if self.tp_size > 1 and not should_skip_post_experts_all_reduce(
is_tp_path=True, is_tp_path=True,
use_reduce_scatter=use_reduce_scatter,
should_allreduce_fusion=should_allreduce_fusion,
): ):
final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states)
return final_hidden_states return final_hidden_states
@@ -364,8 +355,6 @@ class Glm4MoeLiteSparseMoeBlock(nn.Module):
def forward_normal( def forward_normal(
self, self,
hidden_states: torch.Tensor, hidden_states: torch.Tensor,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
) -> torch.Tensor: ) -> torch.Tensor:
if hidden_states.shape[0] > 0: if hidden_states.shape[0] > 0:
shared_output = self._forward_shared_experts(hidden_states) shared_output = self._forward_shared_experts(hidden_states)
@@ -388,8 +377,6 @@ class Glm4MoeLiteSparseMoeBlock(nn.Module):
final_hidden_states = final_hidden_states_out final_hidden_states = final_hidden_states_out
if self.tp_size > 1 and not should_skip_post_experts_all_reduce( if self.tp_size > 1 and not should_skip_post_experts_all_reduce(
is_tp_path=True, is_tp_path=True,
use_reduce_scatter=use_reduce_scatter,
should_allreduce_fusion=should_allreduce_fusion,
): ):
final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states)
return final_hidden_states return final_hidden_states
@@ -679,22 +666,24 @@ class Glm4MoeLiteDecoderLayer(nn.Module):
hidden_states, residual, forward_batch hidden_states, residual, forward_batch
) )
should_allreduce_fusion = ( fuse_mlp_allreduce = (
self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer( self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer(
forward_batch forward_batch
) )
) )
# For DP with padding, reduce scatter can be used instead of all-reduce. # For DP with padding, reduce scatter can be used instead of all-reduce.
use_reduce_scatter = self.layer_communicator.should_use_reduce_scatter( mlp_reduce_scatter = self.layer_communicator.should_use_reduce_scatter(
forward_batch forward_batch
) )
hidden_states = self.mlp( with get_forward().scoped(
hidden_states, forward_batch, should_allreduce_fusion, use_reduce_scatter fuse_mlp_allreduce=fuse_mlp_allreduce,
) mlp_reduce_scatter=mlp_reduce_scatter,
):
hidden_states = self.mlp(hidden_states, forward_batch)
if should_allreduce_fusion: if fuse_mlp_allreduce:
hidden_states._sglang_needs_allreduce_fusion = True hidden_states._sglang_needs_allreduce_fusion = True
else: else:
hidden_states, residual = self.layer_communicator.postprocess_layer( hidden_states, residual = self.layer_communicator.postprocess_layer(
python/sglang/srt/models/gpt_oss.py
+8 -9
View File
@@ -68,7 +68,7 @@ from sglang.srt.models.utils import (
create_fused_set_kv_buffer_arg, create_fused_set_kv_buffer_arg,
enable_fused_set_kv_buffer, enable_fused_set_kv_buffer,
) )
from sglang.srt.runtime_context import get_parallel, get_server_args from sglang.srt.runtime_context import get_forward, get_parallel, get_server_args
from sglang.srt.utils import ( from sglang.srt.utils import (
LazyValue, LazyValue,
add_prefix, add_prefix,
@@ -254,10 +254,9 @@ class GptOssSparseMoeBlock(nn.Module):
self, self,
hidden_states: torch.Tensor, hidden_states: torch.Tensor,
forward_batch: Optional[ForwardBatch] = None, forward_batch: Optional[ForwardBatch] = None,
should_allreduce_fusion: bool = False,
) -> torch.Tensor: ) -> torch.Tensor:
if not get_moe_a2a_backend().is_deepep(): if not get_moe_a2a_backend().is_deepep():
return self.forward_normal(hidden_states, should_allreduce_fusion) return self.forward_normal(hidden_states)
else: else:
raise Exception("forward_deepep branch not implemented yet") raise Exception("forward_deepep branch not implemented yet")
@@ -274,7 +273,6 @@ class GptOssSparseMoeBlock(nn.Module):
def forward_normal( def forward_normal(
self, self,
hidden_states: torch.Tensor, hidden_states: torch.Tensor,
should_allreduce_fusion: bool = False,
) -> torch.Tensor: ) -> torch.Tensor:
# `hidden_states` may arrive pre-padded along the last dim when the # `hidden_states` may arrive pre-padded along the last dim when the
# preceding RMSNorm fused the MoE input pad (gated by # preceding RMSNorm fused the MoE input pad (gated by
@@ -299,7 +297,7 @@ class GptOssSparseMoeBlock(nn.Module):
topk_output = self.topk(router_input, router_logits) topk_output = self.topk(router_input, router_logits)
final_hidden_states = self.experts(hidden_states, topk_output) final_hidden_states = self.experts(hidden_states, topk_output)
if self.tp_size > 1 and not should_allreduce_fusion: if self.tp_size > 1 and not get_forward().fuse_mlp_allreduce:
final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states)
# When input was pre-padded, FusedMoE.forward_impl captured the # When input was pre-padded, FusedMoE.forward_impl captured the
@@ -603,18 +601,19 @@ class GptOssDecoderLayer(nn.Module):
hidden_states, residual, forward_batch hidden_states, residual, forward_batch
) )
should_allreduce_fusion = ( fuse_mlp_allreduce = (
self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer( self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer(
forward_batch forward_batch
) )
) )
hidden_states = self.mlp(hidden_states, forward_batch, should_allreduce_fusion) with get_forward().scoped(fuse_mlp_allreduce=fuse_mlp_allreduce):
hidden_states = self.mlp(hidden_states, forward_batch)
if should_allreduce_fusion: if fuse_mlp_allreduce:
hidden_states._sglang_needs_allreduce_fusion = True hidden_states._sglang_needs_allreduce_fusion = True
if not should_allreduce_fusion: if not fuse_mlp_allreduce:
hidden_states, residual = self.layer_communicator.postprocess_layer( hidden_states, residual = self.layer_communicator.postprocess_layer(
hidden_states, residual, forward_batch hidden_states, residual, forward_batch
) )
python/sglang/srt/models/laguna.py
+15 -22
View File
@@ -53,7 +53,7 @@ from sglang.srt.layers.vocab_parallel_embedding import (
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
from sglang.srt.model_loader.weight_utils import default_weight_loader from sglang.srt.model_loader.weight_utils import default_weight_loader
from sglang.srt.models.utils import apply_qk_norm from sglang.srt.models.utils import apply_qk_norm
from sglang.srt.runtime_context import get_parallel, get_server_args from sglang.srt.runtime_context import get_forward, get_parallel, get_server_args
from sglang.srt.utils import LazyValue, add_prefix, make_layers from sglang.srt.utils import LazyValue, add_prefix, make_layers
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@@ -101,17 +101,12 @@ class LagunaMLP(nn.Module):
self, self,
x: torch.Tensor, x: torch.Tensor,
forward_batch: Optional[ForwardBatch] = None, forward_batch: Optional[ForwardBatch] = None,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
) -> torch.Tensor: ) -> torch.Tensor:
gate_up, _ = self.gate_up_proj(x) gate_up, _ = self.gate_up_proj(x)
x = self.act_fn(gate_up) x = self.act_fn(gate_up)
# Skip the in-block reduce when LayerCommunicator will fuse it or when # RowParallelLinear honors ForwardFlags (fuse_mlp_allreduce /
# the next layer expects reduce-scatter — otherwise we'd double-reduce. # mlp_reduce_scatter) published by the decoder via scoped().
x, _ = self.down_proj( x, _ = self.down_proj(x)
x,
skip_all_reduce=should_allreduce_fusion or use_reduce_scatter,
)
return x return x
@@ -204,8 +199,6 @@ class LagunaMoE(nn.Module):
self, self,
hidden_states: torch.Tensor, hidden_states: torch.Tensor,
forward_batch: Optional[ForwardBatch] = None, forward_batch: Optional[ForwardBatch] = None,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
) -> torch.Tensor: ) -> torch.Tensor:
if hidden_states.shape[0] == 0: if hidden_states.shape[0] == 0:
return hidden_states return hidden_states
@@ -233,8 +226,6 @@ class LagunaMoE(nn.Module):
if self.tp_size > 1 and not should_skip_post_experts_all_reduce( if self.tp_size > 1 and not should_skip_post_experts_all_reduce(
is_tp_path=True, is_tp_path=True,
use_reduce_scatter=use_reduce_scatter,
should_allreduce_fusion=should_allreduce_fusion,
): ):
final = tensor_model_parallel_all_reduce(final) final = tensor_model_parallel_all_reduce(final)
if self._shared_expert_tp1: if self._shared_expert_tp1:
@@ -502,23 +493,25 @@ class LagunaDecoderLayer(nn.Module):
hidden_states, residual, forward_batch hidden_states, residual, forward_batch
) )
should_allreduce_fusion = ( fuse_mlp_allreduce = (
self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer( self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer(
forward_batch forward_batch
) )
) )
use_reduce_scatter = self.layer_communicator.should_use_reduce_scatter( mlp_reduce_scatter = self.layer_communicator.should_use_reduce_scatter(
forward_batch forward_batch
) )
hidden_states = self.mlp( with get_forward().scoped(
hidden_states, fuse_mlp_allreduce=fuse_mlp_allreduce,
forward_batch=forward_batch, mlp_reduce_scatter=mlp_reduce_scatter,
should_allreduce_fusion=should_allreduce_fusion, ):
use_reduce_scatter=use_reduce_scatter, hidden_states = self.mlp(
) hidden_states,
forward_batch=forward_batch,
)
if should_allreduce_fusion: if fuse_mlp_allreduce:
hidden_states._sglang_needs_allreduce_fusion = True hidden_states._sglang_needs_allreduce_fusion = True
else: else:
hidden_states, residual = self.layer_communicator.postprocess_layer( hidden_states, residual = self.layer_communicator.postprocess_layer(
python/sglang/srt/models/llada2.py
+11 -11
View File
@@ -76,7 +76,12 @@ from sglang.srt.models.utils import (
create_fused_set_kv_buffer_arg, create_fused_set_kv_buffer_arg,
enable_fused_set_kv_buffer, enable_fused_set_kv_buffer,
) )
from sglang.srt.runtime_context import get_parallel, get_server_args, get_stream from sglang.srt.runtime_context import (
get_forward,
get_parallel,
get_server_args,
get_stream,
)
from sglang.srt.utils import ( from sglang.srt.utils import (
add_prefix, add_prefix,
is_cuda, is_cuda,
@@ -134,16 +139,13 @@ class LLaDA2MoeMLP(nn.Module):
self, self,
hidden_states: torch.Tensor, hidden_states: torch.Tensor,
forward_batch: Optional[ForwardBatch] = None, forward_batch: Optional[ForwardBatch] = None,
use_reduce_scatter: bool = False,
) -> torch.Tensor: ) -> torch.Tensor:
if (self.tp_size == 1) and hidden_states.shape[0] == 0: if (self.tp_size == 1) and hidden_states.shape[0] == 0:
return hidden_states return hidden_states
gate_up, _ = self.gate_up_proj(hidden_states) gate_up, _ = self.gate_up_proj(hidden_states)
hidden_states = self.act_fn(gate_up) hidden_states = self.act_fn(gate_up)
hidden_states, _ = self.down_proj( hidden_states, _ = self.down_proj(hidden_states)
hidden_states, skip_all_reduce=use_reduce_scatter
)
return hidden_states return hidden_states
@@ -317,10 +319,9 @@ class LLaDA2MoeSparseMoeBlock(nn.Module):
self, self,
hidden_states: torch.Tensor, hidden_states: torch.Tensor,
forward_batch: Optional[ForwardBatch] = None, forward_batch: Optional[ForwardBatch] = None,
use_reduce_scatter: bool = False,
) -> torch.Tensor: ) -> torch.Tensor:
if not get_moe_a2a_backend().is_deepep(): if not get_moe_a2a_backend().is_deepep():
return self.forward_normal(hidden_states, use_reduce_scatter) return self.forward_normal(hidden_states)
else: else:
return self.forward_deepep(hidden_states, forward_batch) return self.forward_deepep(hidden_states, forward_batch)
@@ -360,7 +361,6 @@ class LLaDA2MoeSparseMoeBlock(nn.Module):
def forward_normal( def forward_normal(
self, self,
hidden_states: torch.Tensor, hidden_states: torch.Tensor,
use_reduce_scatter: bool = False,
) -> torch.Tensor: ) -> torch.Tensor:
num_tokens, hidden_size = hidden_states.shape num_tokens, hidden_size = hidden_states.shape
hidden_states = hidden_states.view(-1, hidden_size) hidden_states = hidden_states.view(-1, hidden_size)
@@ -382,7 +382,6 @@ class LLaDA2MoeSparseMoeBlock(nn.Module):
if self.tp_size > 1 and not should_skip_post_experts_all_reduce( if self.tp_size > 1 and not should_skip_post_experts_all_reduce(
is_tp_path=True, is_tp_path=True,
use_reduce_scatter=use_reduce_scatter,
): ):
final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states)
return final_hidden_states.view(num_tokens, hidden_size) return final_hidden_states.view(num_tokens, hidden_size)
@@ -662,11 +661,12 @@ class LLaDA2MoeBlock(nn.Module):
) )
# For DP with padding, reduce scatter can be used instead of all-reduce. # For DP with padding, reduce scatter can be used instead of all-reduce.
use_reduce_scatter = self.layer_communicator.should_use_reduce_scatter( mlp_reduce_scatter = self.layer_communicator.should_use_reduce_scatter(
forward_batch forward_batch
) )
hidden_states = self.mlp(hidden_states, forward_batch, use_reduce_scatter) with get_forward().scoped(mlp_reduce_scatter=mlp_reduce_scatter):
hidden_states = self.mlp(hidden_states, forward_batch)
hidden_states, residual = self.layer_communicator.postprocess_layer( hidden_states, residual = self.layer_communicator.postprocess_layer(
hidden_states=hidden_states, hidden_states=hidden_states,
python/sglang/srt/models/llama.py
+1 -5
View File
@@ -112,14 +112,10 @@ class LlamaMLP(nn.Module):
self, self,
x, x,
forward_batch=None, forward_batch=None,
use_reduce_scatter: bool = False,
): ):
gate_up, _ = self.gate_up_proj(x) gate_up, _ = self.gate_up_proj(x)
x = self.act_fn(gate_up) x = self.act_fn(gate_up)
x, _ = self.down_proj( x, _ = self.down_proj(x)
x,
skip_all_reduce=use_reduce_scatter,
)
return x return x
python/sglang/srt/models/llama4.py
+4 -7
View File
@@ -52,7 +52,7 @@ from sglang.srt.model_executor.forward_batch_info import (
) )
from sglang.srt.models.llama import LlamaForCausalLM, LlamaMLP from sglang.srt.models.llama import LlamaForCausalLM, LlamaMLP
from sglang.srt.models.utils import apply_qk_norm from sglang.srt.models.utils import apply_qk_norm
from sglang.srt.runtime_context import get_parallel from sglang.srt.runtime_context import get_forward, get_parallel
from sglang.srt.utils import ( from sglang.srt.utils import (
add_prefix, add_prefix,
fast_topk, fast_topk,
@@ -139,7 +139,6 @@ class Llama4MoE(nn.Module):
self, self,
hidden_states, hidden_states,
forward_batch: ForwardBatch, forward_batch: ForwardBatch,
use_reduce_scatter: bool = False,
): ):
shared_out, routed_out = self._forward_core( shared_out, routed_out = self._forward_core(
hidden_states, forward_batch.forward_mode hidden_states, forward_batch.forward_mode
@@ -149,7 +148,6 @@ class Llama4MoE(nn.Module):
if self.tp_size > 1 and not should_skip_post_experts_all_reduce( if self.tp_size > 1 and not should_skip_post_experts_all_reduce(
is_tp_path=True, is_tp_path=True,
use_reduce_scatter=use_reduce_scatter,
): ):
out_aD = tensor_model_parallel_all_reduce(out_aD) out_aD = tensor_model_parallel_all_reduce(out_aD)
@@ -479,14 +477,13 @@ class Llama4DecoderLayer(nn.Module):
) )
# For DP with padding, reduce scatter can be used instead of all-reduce. # For DP with padding, reduce scatter can be used instead of all-reduce.
use_reduce_scatter = self.layer_communicator.should_use_reduce_scatter( mlp_reduce_scatter = self.layer_communicator.should_use_reduce_scatter(
forward_batch forward_batch
) )
# Fully Connected # Fully Connected
hidden_states = self.feed_forward( with get_forward().scoped(mlp_reduce_scatter=mlp_reduce_scatter):
hidden_states, forward_batch, use_reduce_scatter hidden_states = self.feed_forward(hidden_states, forward_batch)
)
hidden_states, residual = self.layer_communicator.postprocess_layer( hidden_states, residual = self.layer_communicator.postprocess_layer(
hidden_states, residual, forward_batch hidden_states, residual, forward_batch
) )
python/sglang/srt/models/mimo_v2.py
+11 -23
View File
@@ -76,7 +76,7 @@ from sglang.srt.model_loader.weight_utils import (
) )
from sglang.srt.models.mimo_audio import AudioEncoderMixin, MiMoAudioEncoderConfig from sglang.srt.models.mimo_audio import AudioEncoderMixin, MiMoAudioEncoderConfig
from sglang.srt.models.mimo_vl import MiMoVisionTransformer, MiMoVLVisionConfig from sglang.srt.models.mimo_vl import MiMoVisionTransformer, MiMoVLVisionConfig
from sglang.srt.runtime_context import get_parallel, get_server_args from sglang.srt.runtime_context import get_forward, get_parallel, get_server_args
from sglang.srt.utils import ( from sglang.srt.utils import (
LazyValue, LazyValue,
add_prefix, add_prefix,
@@ -166,17 +166,13 @@ class MiMoV2MLP(nn.Module):
self, self,
x, x,
forward_batch: ForwardBatch = None, forward_batch: ForwardBatch = None,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
): ):
if (self.tp_size == 1) and x.shape[0] == 0: if (self.tp_size == 1) and x.shape[0] == 0:
return x return x
gate_up, _ = self.gate_up_proj(x) gate_up, _ = self.gate_up_proj(x)
x = self.act_fn(gate_up) x = self.act_fn(gate_up)
x, _ = self.down_proj( x, _ = self.down_proj(x)
x, skip_all_reduce=should_allreduce_fusion or use_reduce_scatter
)
return x return x
@@ -315,23 +311,15 @@ class MiMoV2MoE(nn.Module):
self, self,
hidden_states: torch.Tensor, hidden_states: torch.Tensor,
forward_batch: Optional[ForwardBatch] = None, forward_batch: Optional[ForwardBatch] = None,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
) -> torch.Tensor: ) -> torch.Tensor:
if not self._enable_a2a_moe: if not self._enable_a2a_moe:
return self.forward_normal( return self.forward_normal(hidden_states)
hidden_states,
should_allreduce_fusion,
use_reduce_scatter,
)
else: else:
return self.forward_deepep(hidden_states, forward_batch) return self.forward_deepep(hidden_states, forward_batch)
def forward_normal( def forward_normal(
self, self,
hidden_states: torch.Tensor, hidden_states: torch.Tensor,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
) -> torch.Tensor: ) -> torch.Tensor:
if hidden_states.shape[0] > 0: if hidden_states.shape[0] > 0:
@@ -345,8 +333,6 @@ class MiMoV2MoE(nn.Module):
if self.tp_size > 1 and not should_skip_post_experts_all_reduce( if self.tp_size > 1 and not should_skip_post_experts_all_reduce(
is_tp_path=True, is_tp_path=True,
use_reduce_scatter=use_reduce_scatter,
should_allreduce_fusion=should_allreduce_fusion,
): ):
final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states)
@@ -746,22 +732,24 @@ class MiMoV2DecoderLayer(nn.Module):
hidden_states, residual, forward_batch hidden_states, residual, forward_batch
) )
should_allreduce_fusion = ( fuse_mlp_allreduce = (
self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer( self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer(
forward_batch forward_batch
) )
) )
# For DP with padding, reduce scatter can be used instead of all-reduce. # For DP with padding, reduce scatter can be used instead of all-reduce.
use_reduce_scatter = self.layer_communicator.should_use_reduce_scatter( mlp_reduce_scatter = self.layer_communicator.should_use_reduce_scatter(
forward_batch forward_batch
) )
hidden_states = self.mlp( with get_forward().scoped(
hidden_states, forward_batch, should_allreduce_fusion, use_reduce_scatter fuse_mlp_allreduce=fuse_mlp_allreduce,
) mlp_reduce_scatter=mlp_reduce_scatter,
):
hidden_states = self.mlp(hidden_states, forward_batch)
if should_allreduce_fusion: if fuse_mlp_allreduce:
hidden_states._sglang_needs_allreduce_fusion = True hidden_states._sglang_needs_allreduce_fusion = True
else: else:
hidden_states, residual = self.layer_communicator.postprocess_layer( hidden_states, residual = self.layer_communicator.postprocess_layer(
python/sglang/srt/models/minimax_m2.py
+11 -17
View File
@@ -80,7 +80,7 @@ from sglang.srt.model_loader.weight_utils import (
maybe_remap_kv_scale_name, maybe_remap_kv_scale_name,
narrow_padded_param_and_loaded_weight, narrow_padded_param_and_loaded_weight,
) )
from sglang.srt.runtime_context import get_parallel, get_server_args from sglang.srt.runtime_context import get_forward, get_parallel, get_server_args
# get_bool_env_var is defined in sglang.srt.utils.common, not sglang.srt.distributed. # get_bool_env_var is defined in sglang.srt.utils.common, not sglang.srt.distributed.
# Importing from the wrong module causes this file to fail import, which prevents the # Importing from the wrong module causes this file to fail import, which prevents the
@@ -553,24 +553,18 @@ class MiniMaxM2MoE(nn.Module):
self, self,
hidden_states: torch.Tensor, hidden_states: torch.Tensor,
forward_batch: Optional[ForwardBatch] = None, forward_batch: Optional[ForwardBatch] = None,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
) -> torch.Tensor: ) -> torch.Tensor:
if ( if (
not get_moe_a2a_backend().is_deepep() not get_moe_a2a_backend().is_deepep()
and not get_moe_a2a_backend().is_ascend_fuseep() and not get_moe_a2a_backend().is_ascend_fuseep()
): ):
return self.forward_normal( return self.forward_normal(hidden_states)
hidden_states, should_allreduce_fusion, use_reduce_scatter
)
else: else:
return self.forward_deepep(hidden_states, forward_batch) return self.forward_deepep(hidden_states, forward_batch)
def forward_normal( def forward_normal(
self, self,
hidden_states: torch.Tensor, hidden_states: torch.Tensor,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
) -> torch.Tensor: ) -> torch.Tensor:
num_tokens, hidden_dim = hidden_states.shape num_tokens, hidden_dim = hidden_states.shape
hidden_states = hidden_states.view(-1, hidden_dim) hidden_states = hidden_states.view(-1, hidden_dim)
@@ -584,9 +578,7 @@ class MiniMaxM2MoE(nn.Module):
final_hidden_states = self.experts(hidden_states, topk_output) final_hidden_states = self.experts(hidden_states, topk_output)
if self.tp_size > 1 and not should_skip_post_experts_all_reduce( if self.tp_size > 1 and not should_skip_post_experts_all_reduce(
is_tp_path=True, is_tp_path=True
use_reduce_scatter=use_reduce_scatter,
should_allreduce_fusion=should_allreduce_fusion,
): ):
final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states)
@@ -1012,21 +1004,23 @@ class MiniMaxM2DecoderLayer(nn.Module):
hidden_states, residual, forward_batch hidden_states, residual, forward_batch
) )
should_allreduce_fusion = ( fuse_mlp_allreduce = (
self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer( self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer(
forward_batch forward_batch
) )
) )
use_reduce_scatter = self.layer_communicator.should_use_reduce_scatter( mlp_reduce_scatter = self.layer_communicator.should_use_reduce_scatter(
forward_batch forward_batch
) )
hidden_states = self.block_sparse_moe( with get_forward().scoped(
hidden_states, forward_batch, should_allreduce_fusion, use_reduce_scatter fuse_mlp_allreduce=fuse_mlp_allreduce,
) mlp_reduce_scatter=mlp_reduce_scatter,
):
hidden_states = self.block_sparse_moe(hidden_states, forward_batch)
if should_allreduce_fusion: if fuse_mlp_allreduce:
hidden_states._sglang_needs_allreduce_fusion = True hidden_states._sglang_needs_allreduce_fusion = True
else: else:
hidden_states, residual = self.layer_communicator.postprocess_layer( hidden_states, residual = self.layer_communicator.postprocess_layer(
python/sglang/srt/models/nemotron_h.py
+36 -56
View File
@@ -87,7 +87,7 @@ from sglang.srt.models.nemotron_h_utils import (
pad_to_original_num_tokens, pad_to_original_num_tokens,
) )
from sglang.srt.models.utils import WeightsMapper from sglang.srt.models.utils import WeightsMapper
from sglang.srt.runtime_context import get_parallel, get_server_args from sglang.srt.runtime_context import get_forward, get_parallel, get_server_args
from sglang.srt.utils import ( from sglang.srt.utils import (
add_prefix, add_prefix,
get_current_device_stream_fast, get_current_device_stream_fast,
@@ -132,14 +132,10 @@ class NemotronHMLP(nn.Module):
def forward( def forward(
self, self,
x: torch.Tensor, x: torch.Tensor,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
): ):
x, _ = self.up_proj(x) x, _ = self.up_proj(x)
x = self.act_fn(x) x = self.act_fn(x)
x, _ = self.down_proj( x, _ = self.down_proj(x)
x, skip_all_reduce=should_allreduce_fusion or use_reduce_scatter
)
return x return x
@@ -306,8 +302,6 @@ class NemotronHMoE(nn.Module):
def forward( def forward(
self, self,
hidden_states: torch.Tensor, hidden_states: torch.Tensor,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
) -> torch.Tensor: ) -> torch.Tensor:
num_tokens, hidden_dim = hidden_states.shape num_tokens, hidden_dim = hidden_states.shape
# routed_scaling_factor is fused into the experts call (applied by the # routed_scaling_factor is fused into the experts call (applied by the
@@ -322,8 +316,6 @@ class NemotronHMoE(nn.Module):
if self.tp_size > 1 and not should_skip_post_experts_all_reduce( if self.tp_size > 1 and not should_skip_post_experts_all_reduce(
is_tp_path=True, is_tp_path=True,
use_reduce_scatter=use_reduce_scatter,
should_allreduce_fusion=should_allreduce_fusion,
): ):
final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states)
@@ -344,20 +336,20 @@ class NemotronHMLPLikeDecoderLayer(nn.Module):
hidden_states, residual = self.layer_communicator.prepare_mlp( hidden_states, residual = self.layer_communicator.prepare_mlp(
hidden_states, residual, forward_batch hidden_states, residual, forward_batch
) )
use_reduce_scatter = self.layer_communicator.should_use_reduce_scatter( mlp_reduce_scatter = self.layer_communicator.should_use_reduce_scatter(
forward_batch forward_batch
) )
should_allreduce_fusion = ( fuse_mlp_allreduce = (
self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer( self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer(
forward_batch forward_batch
) )
) )
hidden_states = self.mixer.forward( with get_forward().scoped(
hidden_states, fuse_mlp_allreduce=fuse_mlp_allreduce,
should_allreduce_fusion=should_allreduce_fusion, mlp_reduce_scatter=mlp_reduce_scatter,
use_reduce_scatter=use_reduce_scatter, ):
) hidden_states = self.mixer.forward(hidden_states)
if should_allreduce_fusion: if fuse_mlp_allreduce:
hidden_states._sglang_needs_allreduce_fusion = True hidden_states._sglang_needs_allreduce_fusion = True
else: else:
hidden_states, residual = self.layer_communicator.postprocess_layer( hidden_states, residual = self.layer_communicator.postprocess_layer(
@@ -369,15 +361,14 @@ class NemotronHMLPLikeDecoderLayer(nn.Module):
self.norm, hidden_states, residual self.norm, hidden_states, residual
) )
should_allreduce_fusion = ( fuse_mlp_allreduce = (
self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer( self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer(
forward_batch forward_batch
) )
) )
hidden_states = self.mixer.forward( with get_forward().scoped(fuse_mlp_allreduce=fuse_mlp_allreduce):
hidden_states, should_allreduce_fusion=should_allreduce_fusion hidden_states = self.mixer.forward(hidden_states)
) if fuse_mlp_allreduce:
if should_allreduce_fusion:
hidden_states._sglang_needs_allreduce_fusion = True hidden_states._sglang_needs_allreduce_fusion = True
return hidden_states, residual return hidden_states, residual
@@ -505,7 +496,6 @@ class NemotronHMambaDecoderLayer(NemotronHAttnLikeDecoderLayer):
self, self,
hidden_states: torch.Tensor, hidden_states: torch.Tensor,
forward_batch: ForwardBatch, forward_batch: ForwardBatch,
should_allreduce_fusion: bool = False,
) -> torch.Tensor: ) -> torch.Tensor:
"""Core Mamba forward logic, called directly or via split op.""" """Core Mamba forward logic, called directly or via split op."""
original_num_tokens = hidden_states.shape[0] original_num_tokens = hidden_states.shape[0]
@@ -523,7 +513,6 @@ class NemotronHMambaDecoderLayer(NemotronHAttnLikeDecoderLayer):
output=None, output=None,
forward_batch=forward_batch, forward_batch=forward_batch,
use_triton_causal_conv=True, use_triton_causal_conv=True,
should_allreduce_fusion=should_allreduce_fusion,
) )
return pad_to_original_num_tokens(output, original_num_tokens) return pad_to_original_num_tokens(output, original_num_tokens)
@@ -551,28 +540,25 @@ class NemotronHMambaDecoderLayer(NemotronHAttnLikeDecoderLayer):
self.norm, hidden_states, residual self.norm, hidden_states, residual
) )
should_allreduce_fusion = ( fuse_mlp_allreduce = (
self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer( self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer(
forward_batch forward_batch
) )
) )
if is_in_breakable_cuda_graph(): with get_forward().scoped(fuse_mlp_allreduce=fuse_mlp_allreduce):
output = torch.empty_like(hidden_states) if is_in_breakable_cuda_graph():
breakable_nemotron_mamba2_with_output( output = torch.empty_like(hidden_states)
hidden_states, output, self.layer_id, should_allreduce_fusion breakable_nemotron_mamba2_with_output(
) hidden_states, output, self.layer_id
elif is_in_tc_piecewise_cuda_graph(): )
output = torch.empty_like(hidden_states) elif is_in_tc_piecewise_cuda_graph():
nemotron_mamba2_with_output( output = torch.empty_like(hidden_states)
hidden_states, output, self.layer_id, should_allreduce_fusion nemotron_mamba2_with_output(hidden_states, output, self.layer_id)
) else:
else: output = self._forward_mamba(hidden_states, forward_batch)
output = self._forward_mamba(
hidden_states, forward_batch, should_allreduce_fusion
)
if should_allreduce_fusion: if fuse_mlp_allreduce:
output._sglang_needs_allreduce_fusion = True output._sglang_needs_allreduce_fusion = True
return output, residual return output, residual
@@ -647,15 +633,12 @@ class NemotronHAttention(nn.Module):
self, self,
hidden_states: torch.Tensor, hidden_states: torch.Tensor,
forward_batch: ForwardBatch, forward_batch: ForwardBatch,
should_allreduce_fusion: bool = False,
) -> torch.Tensor: ) -> torch.Tensor:
if not is_dp_attention_enabled(): if not is_dp_attention_enabled():
qkv, _ = self.qkv_proj(hidden_states) qkv, _ = self.qkv_proj(hidden_states)
q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
attn_output = self.attn.forward(q, k, v, forward_batch) attn_output = self.attn.forward(q, k, v, forward_batch)
output, _ = self.o_proj( output, _ = self.o_proj(attn_output)
attn_output, skip_all_reduce=should_allreduce_fusion
)
return output return output
padded_shape = hidden_states.shape[0] padded_shape = hidden_states.shape[0]
@@ -733,18 +716,18 @@ class NemotronHAttentionDecoderLayer(NemotronHAttnLikeDecoderLayer):
self.norm, hidden_states, residual self.norm, hidden_states, residual
) )
should_allreduce_fusion = ( fuse_mlp_allreduce = (
self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer( self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer(
forward_batch forward_batch
) )
) )
hidden_states = self.mixer.forward( with get_forward().scoped(fuse_mlp_allreduce=fuse_mlp_allreduce):
hidden_states=hidden_states, hidden_states = self.mixer.forward(
forward_batch=forward_batch, hidden_states=hidden_states,
should_allreduce_fusion=should_allreduce_fusion, forward_batch=forward_batch,
) )
if should_allreduce_fusion: if fuse_mlp_allreduce:
hidden_states._sglang_needs_allreduce_fusion = True hidden_states._sglang_needs_allreduce_fusion = True
return hidden_states, residual return hidden_states, residual
@@ -1206,7 +1189,6 @@ def nemotron_mamba2_with_output(
hidden_states: torch.Tensor, hidden_states: torch.Tensor,
output: torch.Tensor, output: torch.Tensor,
layer_id: int, layer_id: int,
should_allreduce_fusion: bool = False,
) -> None: ) -> None:
"""Split op for Mamba2 forward in piecewise CUDA graph mode.""" """Split op for Mamba2 forward in piecewise CUDA graph mode."""
context = get_tc_piecewise_forward_context() context = get_tc_piecewise_forward_context()
@@ -1226,9 +1208,7 @@ def nemotron_mamba2_with_output(
if hidden_states.shape[0] != num_actual_tokens: if hidden_states.shape[0] != num_actual_tokens:
hidden_states = hidden_states[:num_actual_tokens] hidden_states = hidden_states[:num_actual_tokens]
ret = mamba_layer._forward_mamba( ret = mamba_layer._forward_mamba(hidden_states, forward_batch)
hidden_states, forward_batch, should_allreduce_fusion
)
# Copy result back; output may be larger (padded) so only fill actual tokens # Copy result back; output may be larger (padded) so only fill actual tokens
output[:num_actual_tokens].view(ret.shape).copy_(ret) output[:num_actual_tokens].view(ret.shape).copy_(ret)
python/sglang/srt/models/qwen2_moe.py
+5 -12
View File
@@ -91,7 +91,7 @@ from sglang.srt.model_executor.cuda_graph_config import (
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
from sglang.srt.model_executor.runner import get_is_capture_mode from sglang.srt.model_executor.runner import get_is_capture_mode
from sglang.srt.model_loader.weight_utils import default_weight_loader from sglang.srt.model_loader.weight_utils import default_weight_loader
from sglang.srt.runtime_context import get_parallel, get_server_args from sglang.srt.runtime_context import get_forward, get_parallel, get_server_args
from sglang.srt.utils import ( from sglang.srt.utils import (
add_prefix, add_prefix,
cpu_has_amx_support, cpu_has_amx_support,
@@ -209,14 +209,10 @@ class Qwen2MoeMLP(nn.Module):
def forward( def forward(
self, self,
x, x,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
): ):
gate_up, _ = self.gate_up_proj(x) gate_up, _ = self.gate_up_proj(x)
x = self.act_fn(gate_up) x = self.act_fn(gate_up)
x, _ = self.down_proj( x, _ = self.down_proj(x)
x, skip_all_reduce=should_allreduce_fusion or use_reduce_scatter
)
return x return x
@@ -545,8 +541,6 @@ class Qwen2MoeSparseMoeBlock(nn.Module):
self, self,
hidden_states: torch.Tensor, hidden_states: torch.Tensor,
forward_batch: Optional[ForwardBatch] = None, forward_batch: Optional[ForwardBatch] = None,
use_reduce_scatter: bool = False,
should_allreduce_fusion: bool = False,
) -> torch.Tensor: ) -> torch.Tensor:
num_tokens, hidden_dim = hidden_states.shape num_tokens, hidden_dim = hidden_states.shape
hidden_states = hidden_states.view(-1, hidden_dim) hidden_states = hidden_states.view(-1, hidden_dim)
@@ -591,8 +585,6 @@ class Qwen2MoeSparseMoeBlock(nn.Module):
self.tp_size > 1 self.tp_size > 1
and not should_skip_post_experts_all_reduce( and not should_skip_post_experts_all_reduce(
is_tp_path=True, is_tp_path=True,
use_reduce_scatter=use_reduce_scatter,
should_allreduce_fusion=should_allreduce_fusion,
) )
and not get_moe_a2a_backend().is_flashinfer() and not get_moe_a2a_backend().is_flashinfer()
): ):
@@ -810,11 +802,12 @@ class Qwen2MoeDecoderLayer(nn.Module):
) )
# For DP with padding, reduce scatter can be used instead of all-reduce. # For DP with padding, reduce scatter can be used instead of all-reduce.
use_reduce_scatter = self.layer_communicator.should_use_reduce_scatter( mlp_reduce_scatter = self.layer_communicator.should_use_reduce_scatter(
forward_batch forward_batch
) )
hidden_states = self.mlp(hidden_states, forward_batch, use_reduce_scatter) with get_forward().scoped(mlp_reduce_scatter=mlp_reduce_scatter):
hidden_states = self.mlp(hidden_states, forward_batch)
hidden_states, residual = self.layer_communicator.postprocess_layer( hidden_states, residual = self.layer_communicator.postprocess_layer(
hidden_states, residual, forward_batch hidden_states, residual, forward_batch
python/sglang/srt/models/qwen3_5.py
+34 -29
View File
@@ -91,7 +91,12 @@ from sglang.srt.models.utils import (
fused_qk_gemma_rmsnorm, fused_qk_gemma_rmsnorm,
fused_qk_gemma_rmsnorm_with_gate, fused_qk_gemma_rmsnorm_with_gate,
) )
from sglang.srt.runtime_context import get_parallel, get_server_args, get_stream from sglang.srt.runtime_context import (
get_forward,
get_parallel,
get_server_args,
get_stream,
)
# Utils # Utils
from sglang.srt.utils import ( from sglang.srt.utils import (
@@ -682,27 +687,27 @@ class Qwen3_5LinearDecoderLayer(nn.Module):
hidden_states, residual, forward_batch hidden_states, residual, forward_batch
) )
use_reduce_scatter = self.layer_communicator.should_use_reduce_scatter( mlp_reduce_scatter = self.layer_communicator.should_use_reduce_scatter(
forward_batch forward_batch
) )
should_allreduce_fusion = ( fuse_mlp_allreduce = (
self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer( self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer(
forward_batch forward_batch
) )
) )
if isinstance(self.mlp, Qwen2MoeSparseMoeBlock): with get_forward().scoped(
hidden_states = self.mlp( fuse_mlp_allreduce=fuse_mlp_allreduce,
hidden_states, mlp_reduce_scatter=mlp_reduce_scatter,
forward_batch, ):
use_reduce_scatter, if isinstance(self.mlp, Qwen2MoeSparseMoeBlock):
should_allreduce_fusion, hidden_states = self.mlp(
) hidden_states,
else: forward_batch,
hidden_states = self.mlp( )
hidden_states, should_allreduce_fusion, use_reduce_scatter else:
) hidden_states = self.mlp(hidden_states)
if should_allreduce_fusion: if fuse_mlp_allreduce:
hidden_states._sglang_needs_allreduce_fusion = True hidden_states._sglang_needs_allreduce_fusion = True
else: else:
hidden_states, residual = self.layer_communicator.postprocess_layer( hidden_states, residual = self.layer_communicator.postprocess_layer(
@@ -1067,27 +1072,27 @@ class Qwen3_5AttentionDecoderLayer(nn.Module):
hidden_states, residual = self.layer_communicator.prepare_mlp( hidden_states, residual = self.layer_communicator.prepare_mlp(
hidden_states, residual, forward_batch hidden_states, residual, forward_batch
) )
use_reduce_scatter = self.layer_communicator.should_use_reduce_scatter( mlp_reduce_scatter = self.layer_communicator.should_use_reduce_scatter(
forward_batch forward_batch
) )
should_allreduce_fusion = ( fuse_mlp_allreduce = (
self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer( self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer(
forward_batch forward_batch
) )
) )
if isinstance(self.mlp, Qwen2MoeSparseMoeBlock): with get_forward().scoped(
hidden_states = self.mlp( fuse_mlp_allreduce=fuse_mlp_allreduce,
hidden_states, mlp_reduce_scatter=mlp_reduce_scatter,
forward_batch, ):
use_reduce_scatter, if isinstance(self.mlp, Qwen2MoeSparseMoeBlock):
should_allreduce_fusion, hidden_states = self.mlp(
) hidden_states,
else: forward_batch,
hidden_states = self.mlp( )
hidden_states, should_allreduce_fusion, use_reduce_scatter else:
) hidden_states = self.mlp(hidden_states)
if should_allreduce_fusion: if fuse_mlp_allreduce:
hidden_states._sglang_needs_allreduce_fusion = True hidden_states._sglang_needs_allreduce_fusion = True
else: else:
hidden_states, residual = self.layer_communicator.postprocess_layer( hidden_states, residual = self.layer_communicator.postprocess_layer(
python/sglang/srt/models/qwen3_moe.py
+17 -20
View File
@@ -72,7 +72,12 @@ from sglang.srt.models.utils import (
create_fused_set_kv_buffer_arg, create_fused_set_kv_buffer_arg,
enable_fused_set_kv_buffer, enable_fused_set_kv_buffer,
) )
from sglang.srt.runtime_context import get_parallel, get_server_args, get_stream from sglang.srt.runtime_context import (
get_forward,
get_parallel,
get_server_args,
get_stream,
)
from sglang.srt.utils import ( from sglang.srt.utils import (
LazyValue, LazyValue,
add_prefix, add_prefix,
@@ -286,17 +291,13 @@ class Qwen3MoeSparseMoeBlock(nn.Module):
self, self,
hidden_states: torch.Tensor, hidden_states: torch.Tensor,
forward_batch: Optional[ForwardBatch] = None, forward_batch: Optional[ForwardBatch] = None,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
) -> torch.Tensor: ) -> torch.Tensor:
if ( if (
not get_moe_a2a_backend().is_deepep() not get_moe_a2a_backend().is_deepep()
and not get_moe_a2a_backend().is_ascend_fuseep() and not get_moe_a2a_backend().is_ascend_fuseep()
): ):
return self.forward_normal( return self.forward_normal(hidden_states)
hidden_states, should_allreduce_fusion, use_reduce_scatter
)
else: else:
return self.forward_deepep(hidden_states, forward_batch) return self.forward_deepep(hidden_states, forward_batch)
@@ -313,8 +314,6 @@ class Qwen3MoeSparseMoeBlock(nn.Module):
def forward_normal( def forward_normal(
self, self,
hidden_states: torch.Tensor, hidden_states: torch.Tensor,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
) -> torch.Tensor: ) -> torch.Tensor:
num_tokens, hidden_dim = hidden_states.shape num_tokens, hidden_dim = hidden_states.shape
hidden_states = hidden_states.view(-1, hidden_dim) hidden_states = hidden_states.view(-1, hidden_dim)
@@ -325,16 +324,12 @@ class Qwen3MoeSparseMoeBlock(nn.Module):
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_skip_post_experts_all_reduce( if self.ep_size > 1 and not should_skip_post_experts_all_reduce(
is_tp_path=False, is_tp_path=False
use_reduce_scatter=use_reduce_scatter,
should_allreduce_fusion=should_allreduce_fusion,
): ):
final_hidden_states = moe_expert_parallel_all_reduce(final_hidden_states) final_hidden_states = moe_expert_parallel_all_reduce(final_hidden_states)
if self.tp_size > 1 and not should_skip_post_experts_all_reduce( if self.tp_size > 1 and not should_skip_post_experts_all_reduce(
is_tp_path=True, is_tp_path=True
use_reduce_scatter=use_reduce_scatter,
should_allreduce_fusion=should_allreduce_fusion,
): ):
final_hidden_states = moe_tensor_model_parallel_all_reduce( final_hidden_states = moe_tensor_model_parallel_all_reduce(
final_hidden_states final_hidden_states
@@ -828,22 +823,24 @@ class Qwen3MoeDecoderLayer(nn.Module):
hidden_states, residual, forward_batch hidden_states, residual, forward_batch
) )
should_allreduce_fusion = ( fuse_mlp_allreduce = (
self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer( self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer(
forward_batch forward_batch
) )
) )
# For DP with padding, reduce scatter can be used instead of all-reduce. # For DP with padding, reduce scatter can be used instead of all-reduce.
use_reduce_scatter = self.layer_communicator.should_use_reduce_scatter( mlp_reduce_scatter = self.layer_communicator.should_use_reduce_scatter(
forward_batch forward_batch
) )
hidden_states = self.mlp( with get_forward().scoped(
hidden_states, forward_batch, should_allreduce_fusion, use_reduce_scatter fuse_mlp_allreduce=fuse_mlp_allreduce,
) mlp_reduce_scatter=mlp_reduce_scatter,
):
hidden_states = self.mlp(hidden_states, forward_batch)
if should_allreduce_fusion: if fuse_mlp_allreduce:
hidden_states._sglang_needs_allreduce_fusion = True hidden_states._sglang_needs_allreduce_fusion = True
else: else:
hidden_states, residual = self.layer_communicator.postprocess_layer( hidden_states, residual = self.layer_communicator.postprocess_layer(
python/sglang/srt/models/qwen3_next.py
+20 -17
View File
@@ -47,7 +47,12 @@ from sglang.srt.model_loader.weight_utils import (
sharded_weight_loader, sharded_weight_loader,
) )
from sglang.srt.models.qwen2_moe import Qwen2MoeMLP, Qwen2MoeSparseMoeBlock from sglang.srt.models.qwen2_moe import Qwen2MoeMLP, Qwen2MoeSparseMoeBlock
from sglang.srt.runtime_context import get_parallel, get_server_args, get_stream from sglang.srt.runtime_context import (
get_forward,
get_parallel,
get_server_args,
get_stream,
)
from sglang.srt.utils import ( from sglang.srt.utils import (
LazyValue, LazyValue,
add_prefix, add_prefix,
@@ -471,30 +476,28 @@ def _apply_qwen3_next_mlp(
hidden_states, residual = layer.layer_communicator.prepare_mlp( hidden_states, residual = layer.layer_communicator.prepare_mlp(
hidden_states, residual, forward_batch hidden_states, residual, forward_batch
) )
use_reduce_scatter = layer.layer_communicator.should_use_reduce_scatter( mlp_reduce_scatter = layer.layer_communicator.should_use_reduce_scatter(
forward_batch forward_batch
) )
should_allreduce_fusion = ( fuse_mlp_allreduce = (
layer.layer_communicator.should_fuse_mlp_allreduce_with_next_layer( layer.layer_communicator.should_fuse_mlp_allreduce_with_next_layer(
forward_batch forward_batch
) )
) )
if isinstance(layer.mlp, Qwen2MoeSparseMoeBlock): with get_forward().scoped(
hidden_states = layer.mlp( fuse_mlp_allreduce=fuse_mlp_allreduce,
hidden_states, mlp_reduce_scatter=mlp_reduce_scatter,
forward_batch=forward_batch, ):
use_reduce_scatter=use_reduce_scatter, if isinstance(layer.mlp, Qwen2MoeSparseMoeBlock):
should_allreduce_fusion=should_allreduce_fusion, hidden_states = layer.mlp(
) hidden_states,
else: forward_batch=forward_batch,
hidden_states = layer.mlp( )
hidden_states, else:
should_allreduce_fusion=should_allreduce_fusion, hidden_states = layer.mlp(hidden_states)
use_reduce_scatter=use_reduce_scatter,
)
if should_allreduce_fusion: if fuse_mlp_allreduce:
hidden_states._sglang_needs_allreduce_fusion = True hidden_states._sglang_needs_allreduce_fusion = True
else: else:
hidden_states, residual = layer.layer_communicator.postprocess_layer( hidden_states, residual = layer.layer_communicator.postprocess_layer(
python/sglang/srt/models/sarvam_moe.py
+19 -30
View File
@@ -60,7 +60,12 @@ from sglang.srt.models.bailing_moe import BailingMoEForCausalLM
from sglang.srt.models.deepseek_common.attention_forward_methods.forward_mha import ( from sglang.srt.models.deepseek_common.attention_forward_methods.forward_mha import (
DeepseekMHAForwardMixin, DeepseekMHAForwardMixin,
) )
from sglang.srt.runtime_context import get_parallel, get_server_args, get_stream from sglang.srt.runtime_context import (
get_forward,
get_parallel,
get_server_args,
get_stream,
)
from sglang.srt.utils import ( from sglang.srt.utils import (
BumpAllocator, BumpAllocator,
add_prefix, add_prefix,
@@ -205,16 +210,12 @@ class SarvamMoEMLP(nn.Module):
self, self,
x, x,
forward_batch: ForwardBatch = None, forward_batch: ForwardBatch = None,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
): ):
if x.shape[0] == 0: if x.shape[0] == 0:
return x return x
gate_up, _ = self.gate_up_proj(x) gate_up, _ = self.gate_up_proj(x)
x = self.act_fn(gate_up) x = self.act_fn(gate_up)
x, _ = self.down_proj( x, _ = self.down_proj(x)
x, skip_all_reduce=should_allreduce_fusion or use_reduce_scatter
)
return x return x
@@ -315,8 +316,6 @@ class SarvamMoESparseMoeBlock(nn.Module):
self, self,
hidden_states: torch.Tensor, hidden_states: torch.Tensor,
forward_batch: Optional[ForwardBatch] = None, forward_batch: Optional[ForwardBatch] = None,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
gemm_output_zero_allocator: Optional[BumpAllocator] = None, gemm_output_zero_allocator: Optional[BumpAllocator] = None,
) -> torch.Tensor: ) -> torch.Tensor:
del gemm_output_zero_allocator del gemm_output_zero_allocator
@@ -327,13 +326,9 @@ class SarvamMoESparseMoeBlock(nn.Module):
and hidden_states.shape[0] > 0 and hidden_states.shape[0] > 0
and get_is_capture_mode() and get_is_capture_mode()
): ):
return self.forward_normal_dual_stream( return self.forward_normal_dual_stream(hidden_states)
hidden_states, should_allreduce_fusion, use_reduce_scatter
)
else: else:
return self.forward_normal( return self.forward_normal(hidden_states)
hidden_states, should_allreduce_fusion, use_reduce_scatter
)
def get_moe_weights(self): def get_moe_weights(self):
return [ return [
@@ -359,8 +354,6 @@ class SarvamMoESparseMoeBlock(nn.Module):
def forward_normal_dual_stream( def forward_normal_dual_stream(
self, self,
hidden_states: torch.Tensor, hidden_states: torch.Tensor,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
) -> torch.Tensor: ) -> torch.Tensor:
num_tokens, hidden_dim = hidden_states.shape num_tokens, hidden_dim = hidden_states.shape
current_stream = torch.cuda.current_stream() current_stream = torch.cuda.current_stream()
@@ -374,8 +367,6 @@ class SarvamMoESparseMoeBlock(nn.Module):
final_hidden_states = final_hidden_states + shared_out final_hidden_states = final_hidden_states + shared_out
if self.tp_size > 1 and not should_skip_post_experts_all_reduce( if self.tp_size > 1 and not should_skip_post_experts_all_reduce(
is_tp_path=True, is_tp_path=True,
use_reduce_scatter=use_reduce_scatter,
should_allreduce_fusion=should_allreduce_fusion,
): ):
final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states)
return final_hidden_states.view(num_tokens, hidden_dim) return final_hidden_states.view(num_tokens, hidden_dim)
@@ -383,8 +374,6 @@ class SarvamMoESparseMoeBlock(nn.Module):
def forward_normal( def forward_normal(
self, self,
hidden_states: torch.Tensor, hidden_states: torch.Tensor,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
) -> torch.Tensor: ) -> torch.Tensor:
if hidden_states.shape[0] == 0: if hidden_states.shape[0] == 0:
return hidden_states return hidden_states
@@ -416,8 +405,6 @@ class SarvamMoESparseMoeBlock(nn.Module):
if self.tp_size > 1 and not should_skip_post_experts_all_reduce( if self.tp_size > 1 and not should_skip_post_experts_all_reduce(
is_tp_path=True, is_tp_path=True,
use_reduce_scatter=use_reduce_scatter,
should_allreduce_fusion=should_allreduce_fusion,
): ):
final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states)
@@ -1111,25 +1098,27 @@ class SarvamMoEMLADecoderLayer(nn.Module):
hidden_states, residual = self.layer_communicator.prepare_mlp( hidden_states, residual = self.layer_communicator.prepare_mlp(
hidden_states, residual, forward_batch hidden_states, residual, forward_batch
) )
should_allreduce_fusion = ( fuse_mlp_allreduce = (
self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer( self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer(
forward_batch forward_batch
) )
) )
use_reduce_scatter = self.layer_communicator.should_use_reduce_scatter( mlp_reduce_scatter = self.layer_communicator.should_use_reduce_scatter(
forward_batch forward_batch
) )
hidden_states = self.mlp( with get_forward().scoped(
hidden_states, forward_batch, should_allreduce_fusion, use_reduce_scatter fuse_mlp_allreduce=fuse_mlp_allreduce,
) mlp_reduce_scatter=mlp_reduce_scatter,
):
hidden_states = self.mlp(hidden_states, forward_batch)
if ( if (
not self.is_layer_sparse not self.is_layer_sparse
and self.attn_tp_size > 1 and self.attn_tp_size > 1
and not use_reduce_scatter and not mlp_reduce_scatter
and not should_allreduce_fusion and not fuse_mlp_allreduce
): ):
hidden_states = tensor_model_parallel_all_reduce(hidden_states) hidden_states = tensor_model_parallel_all_reduce(hidden_states)
if should_allreduce_fusion: if fuse_mlp_allreduce:
hidden_states._sglang_needs_allreduce_fusion = True hidden_states._sglang_needs_allreduce_fusion = True
else: else:
hidden_states, residual = self.layer_communicator.postprocess_layer( hidden_states, residual = self.layer_communicator.postprocess_layer(
python/sglang/srt/models/sdar.py
+11 -7
View File
@@ -41,7 +41,12 @@ from sglang.srt.models.utils import (
create_fused_set_kv_buffer_arg, create_fused_set_kv_buffer_arg,
enable_fused_set_kv_buffer, enable_fused_set_kv_buffer,
) )
from sglang.srt.runtime_context import get_parallel, get_server_args, get_stream from sglang.srt.runtime_context import (
get_forward,
get_parallel,
get_server_args,
get_stream,
)
from sglang.srt.utils import add_prefix, is_cuda, make_layers from sglang.srt.utils import add_prefix, is_cuda, make_layers
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@@ -74,12 +79,10 @@ class SDARMLP(nn.Module):
) )
self.act_fn = SiluAndMul() self.act_fn = SiluAndMul()
def forward(self, hidden_states: torch.Tensor, use_reduce_scatter: bool = False): def forward(self, hidden_states: torch.Tensor):
gate_up, _ = self.gate_up_proj(hidden_states) gate_up, _ = self.gate_up_proj(hidden_states)
hidden_states = self.act_fn(gate_up) hidden_states = self.act_fn(gate_up)
hidden_states, _ = self.down_proj( hidden_states, _ = self.down_proj(hidden_states)
hidden_states, skip_all_reduce=use_reduce_scatter
)
return hidden_states return hidden_states
@@ -333,10 +336,11 @@ class SDARBlock(nn.Module):
forward_batch, forward_batch,
) )
use_reduce_scatter = self.layer_communicator.should_use_reduce_scatter( mlp_reduce_scatter = self.layer_communicator.should_use_reduce_scatter(
forward_batch forward_batch
) )
hidden_states = self.mlp(hidden_states, use_reduce_scatter=use_reduce_scatter) with get_forward().scoped(mlp_reduce_scatter=mlp_reduce_scatter):
hidden_states = self.mlp(hidden_states)
hidden_states, residual = self.layer_communicator.postprocess_layer( hidden_states, residual = self.layer_communicator.postprocess_layer(
hidden_states, residual, forward_batch hidden_states, residual, forward_batch
python/sglang/srt/models/sdar_moe.py
+18 -21
View File
@@ -57,7 +57,12 @@ from sglang.srt.models.utils import (
create_fused_set_kv_buffer_arg, create_fused_set_kv_buffer_arg,
enable_fused_set_kv_buffer, enable_fused_set_kv_buffer,
) )
from sglang.srt.runtime_context import get_parallel, get_server_args, get_stream from sglang.srt.runtime_context import (
get_forward,
get_parallel,
get_server_args,
get_stream,
)
from sglang.srt.utils import LazyValue, add_prefix, is_cuda, make_layers from sglang.srt.utils import LazyValue, add_prefix, is_cuda, make_layers
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@@ -126,18 +131,12 @@ class SDARMoeSparseMoeBlock(nn.Module):
self, self,
hidden_states: torch.Tensor, hidden_states: torch.Tensor,
forward_batch: Optional[ForwardBatch] = None, forward_batch: Optional[ForwardBatch] = None,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
) -> torch.Tensor: ) -> torch.Tensor:
if ( if (
not get_moe_a2a_backend().is_deepep() not get_moe_a2a_backend().is_deepep()
and not get_moe_a2a_backend().is_ascend_fuseep() and not get_moe_a2a_backend().is_ascend_fuseep()
): ):
return self.forward_normal( return self.forward_normal(hidden_states)
hidden_states,
should_allreduce_fusion=should_allreduce_fusion,
use_reduce_scatter=use_reduce_scatter,
)
else: else:
assert forward_batch is not None, "deepep/fuseep MoE needs forward_batch" assert forward_batch is not None, "deepep/fuseep MoE needs forward_batch"
return self.forward_deepep(hidden_states, forward_batch) return self.forward_deepep(hidden_states, forward_batch)
@@ -145,8 +144,6 @@ class SDARMoeSparseMoeBlock(nn.Module):
def forward_normal( def forward_normal(
self, self,
hidden_states: torch.Tensor, hidden_states: torch.Tensor,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
) -> torch.Tensor: ) -> torch.Tensor:
num_tokens, hidden_dim = hidden_states.shape num_tokens, hidden_dim = hidden_states.shape
hidden_states = hidden_states.view(-1, hidden_dim) hidden_states = hidden_states.view(-1, hidden_dim)
@@ -157,8 +154,6 @@ class SDARMoeSparseMoeBlock(nn.Module):
if self.tp_size > 1 and not should_skip_post_experts_all_reduce( if self.tp_size > 1 and not should_skip_post_experts_all_reduce(
is_tp_path=True, is_tp_path=True,
use_reduce_scatter=use_reduce_scatter,
should_allreduce_fusion=should_allreduce_fusion,
): ):
out = tensor_model_parallel_all_reduce(out) out = tensor_model_parallel_all_reduce(out)
@@ -408,23 +403,25 @@ class SDARMoeBlock(nn.Module):
hidden_states, residual, forward_batch hidden_states, residual, forward_batch
) )
should_allreduce_fusion = ( fuse_mlp_allreduce = (
self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer( self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer(
forward_batch forward_batch
) )
) )
use_reduce_scatter = self.layer_communicator.should_use_reduce_scatter( mlp_reduce_scatter = self.layer_communicator.should_use_reduce_scatter(
forward_batch forward_batch
) )
hidden_states = self.mlp( with get_forward().scoped(
hidden_states, fuse_mlp_allreduce=fuse_mlp_allreduce,
forward_batch=forward_batch, mlp_reduce_scatter=mlp_reduce_scatter,
should_allreduce_fusion=should_allreduce_fusion, ):
use_reduce_scatter=use_reduce_scatter, hidden_states = self.mlp(
) hidden_states,
forward_batch=forward_batch,
)
if should_allreduce_fusion: if fuse_mlp_allreduce:
hidden_states._sglang_needs_allreduce_fusion = True hidden_states._sglang_needs_allreduce_fusion = True
else: else:
hidden_states, residual = self.layer_communicator.postprocess_layer( hidden_states, residual = self.layer_communicator.postprocess_layer(
python/sglang/srt/models/step3p5.py
+18 -21
View File
@@ -46,7 +46,12 @@ from sglang.srt.layers.vocab_parallel_embedding import (
) )
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
from sglang.srt.model_loader.weight_utils import default_weight_loader from sglang.srt.model_loader.weight_utils import default_weight_loader
from sglang.srt.runtime_context import get_parallel, get_server_args, get_stream from sglang.srt.runtime_context import (
get_forward,
get_parallel,
get_server_args,
get_stream,
)
from sglang.srt.utils import add_prefix, is_cuda, is_non_idle_and_non_empty, make_layers from sglang.srt.utils import add_prefix, is_cuda, is_non_idle_and_non_empty, make_layers
Step3p5Config = None Step3p5Config = None
@@ -179,17 +184,13 @@ class Step3p5MoEMLP(nn.Module):
self, self,
hidden_states: torch.Tensor, hidden_states: torch.Tensor,
forward_batch: Optional[ForwardBatch] = None, forward_batch: Optional[ForwardBatch] = None,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
) -> torch.Tensor: ) -> torch.Tensor:
if ( if (
not get_moe_a2a_backend().is_deepep() not get_moe_a2a_backend().is_deepep()
and not get_moe_a2a_backend().is_ascend_fuseep() and not get_moe_a2a_backend().is_ascend_fuseep()
): ):
return self.forward_normal( return self.forward_normal(hidden_states)
hidden_states, should_allreduce_fusion, use_reduce_scatter
)
else: else:
return self.forward_deepep(hidden_states, forward_batch) return self.forward_deepep(hidden_states, forward_batch)
@@ -206,8 +207,6 @@ class Step3p5MoEMLP(nn.Module):
def forward_normal( def forward_normal(
self, self,
hidden_states: torch.Tensor, hidden_states: torch.Tensor,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
) -> torch.Tensor: ) -> torch.Tensor:
num_tokens, hidden_dim = hidden_states.shape num_tokens, hidden_dim = hidden_states.shape
hidden_states = hidden_states.view(-1, hidden_dim) hidden_states = hidden_states.view(-1, hidden_dim)
@@ -231,8 +230,6 @@ class Step3p5MoEMLP(nn.Module):
final_hidden_states = self.experts(hidden_states, topk_output) final_hidden_states = self.experts(hidden_states, topk_output)
if self.tp_size > 1 and not should_skip_post_experts_all_reduce( if self.tp_size > 1 and not should_skip_post_experts_all_reduce(
is_tp_path=True, is_tp_path=True,
use_reduce_scatter=use_reduce_scatter,
should_allreduce_fusion=should_allreduce_fusion,
): ):
final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states)
@@ -617,12 +614,12 @@ class Step3p5DecoderLayer(nn.Module):
forward_batch, forward_batch,
) )
should_allreduce_fusion = ( fuse_mlp_allreduce = (
self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer( self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer(
forward_batch forward_batch
) )
) )
use_reduce_scatter = self.layer_communicator.should_use_reduce_scatter( mlp_reduce_scatter = self.layer_communicator.should_use_reduce_scatter(
forward_batch forward_batch
) )
@@ -630,24 +627,24 @@ class Step3p5DecoderLayer(nn.Module):
# Both share_expert and MoE return unreduced (TP-partial) outputs. # Both share_expert and MoE return unreduced (TP-partial) outputs.
# Combine them first, then do a single all-reduce — saving one # Combine them first, then do a single all-reduce — saving one
# full-TP all-reduce per layer. # full-TP all-reduce per layer.
# Force fuse_mlp_allreduce=True so MoE skips its internal AR.
share_output = self.share_expert(hidden_states) share_output = self.share_expert(hidden_states)
moe_output = self.moe( with get_forward().scoped(
hidden_states, fuse_mlp_allreduce=True,
forward_batch, mlp_reduce_scatter=mlp_reduce_scatter,
should_allreduce_fusion=True, ):
use_reduce_scatter=use_reduce_scatter, moe_output = self.moe(hidden_states, forward_batch)
)
hidden_states = moe_output + share_output hidden_states = moe_output + share_output
if not should_allreduce_fusion and not use_reduce_scatter: if not fuse_mlp_allreduce and not mlp_reduce_scatter:
hidden_states = tensor_model_parallel_all_reduce(hidden_states) hidden_states = tensor_model_parallel_all_reduce(hidden_states)
else: else:
hidden_states = self.mlp(hidden_states) hidden_states = self.mlp(hidden_states)
# Dense MLP uses reduce_results=True, so the output is already # Dense MLP uses reduce_results=True, so the output is already
# all-reduced. Do NOT set the fusion flag — otherwise the next # all-reduced. Do NOT set the fusion flag — otherwise the next
# layer would all-reduce again, multiplying values by world_size. # layer would all-reduce again, multiplying values by world_size.
should_allreduce_fusion = False fuse_mlp_allreduce = False
if should_allreduce_fusion: if fuse_mlp_allreduce:
hidden_states._sglang_needs_allreduce_fusion = True hidden_states._sglang_needs_allreduce_fusion = True
else: else:
hidden_states, residual = self.layer_communicator.postprocess_layer( hidden_states, residual = self.layer_communicator.postprocess_layer(
python/sglang/srt/runtime_context.py
+14 -3
View File
@@ -407,17 +407,28 @@ class ForwardFlags:
# Sticky across forwards: every ForwardBatch construction writes it; # Sticky across forwards: every ForwardBatch construction writes it;
# graph runners force False around capture. # graph runners force False around capture.
"is_extend_in_batch": False, "is_extend_in_batch": False,
# Per-layer MLP collective control (set by decoder via scoped()
# around the MLP / MoE / hybrid mixer call).
# fuse_mlp_allreduce: next residual+LN absorbs the post-MLP all-reduce.
# mlp_reduce_scatter: postprocess will reduce-scatter (skip MLP AR).
# flashinfer_trtllm_bypass: deepseek dual-stream graph topk bypass.
"fuse_mlp_allreduce": False,
"mlp_reduce_scatter": False,
"flashinfer_trtllm_bypass": False,
} }
# Read/written inside compiled graphs (vocab embedding, communicator, # Read/written inside compiled graphs (vocab embedding, communicator,
# EP dispatch, DP gather/scatter): plain-slot backed. Before moving a # EP dispatch, DP gather/scatter, MLP/MoE skip-AR): plain-slot backed.
# flag out of this set, prove no read/write site sits under # Before moving a flag out of this set, prove no read/write site sits
# torch.compile. # under torch.compile.
_GRAPH_VISIBLE = frozenset( _GRAPH_VISIBLE = frozenset(
{ {
"attn_input_scattered", "attn_input_scattered",
"attn_inputs", "attn_inputs",
"is_extend_in_batch", "is_extend_in_batch",
"fuse_mlp_allreduce",
"mlp_reduce_scatter",
"flashinfer_trtllm_bypass",
} }
) )
test/registered/debug_utils/test_engine_dumper_comparator_e2e.py
+2 -8
View File
@@ -90,10 +90,7 @@ patches:
hidden_states, residual, forward_batch hidden_states, residual, forward_batch
) )
append: "dumper.dump('pre_mlp_residual', hidden_states, dims='t h # tp:replicated')" append: "dumper.dump('pre_mlp_residual', hidden_states, dims='t h # tp:replicated')"
- match: | - match: "hidden_states = self.mlp(hidden_states, forward_batch)"
hidden_states = self.mlp(
hidden_states, forward_batch, should_allreduce_fusion, use_reduce_scatter
)
append: "dumper.dump('mlp_output', hidden_states, dims='t h[moe_tp:partial] # tp:replicated')" append: "dumper.dump('mlp_output', hidden_states, dims='t h[moe_tp:partial] # tp:replicated')"
# --- attention internals --- # --- attention internals ---
@@ -150,10 +147,7 @@ patches:
hidden_states, residual, forward_batch hidden_states, residual, forward_batch
) )
append: "dumper.dump('pre_mlp_residual', hidden_states, dims='t h # tp:replicated')" append: "dumper.dump('pre_mlp_residual', hidden_states, dims='t h # tp:replicated')"
- match: | - match: "hidden_states = self.mlp(hidden_states, forward_batch)"
hidden_states = self.mlp(
hidden_states, forward_batch, should_allreduce_fusion, use_reduce_scatter
)
append: "dumper.dump('mlp_output', hidden_states, dims='t h[moe_tp:partial] # tp:replicated')" append: "dumper.dump('mlp_output', hidden_states, dims='t h[moe_tp:partial] # tp:replicated')"
# --- attention internals --- # --- attention internals ---
test/registered/unit/test_runtime_context.py
+45
View File
@@ -696,6 +696,12 @@ class TestForwardFlags(_IsolatedServerArgs):
x = x + 1 x = x + 1
if fwd.is_extend_in_batch: if fwd.is_extend_in_batch:
x = x + 2 x = x + 2
if fwd.fuse_mlp_allreduce:
x = x + 4
if fwd.mlp_reduce_scatter:
x = x + 8
if fwd.flashinfer_trtllm_bypass:
x = x + 16
return x return x
self.assertEqual(probe(torch.zeros(())).item(), 0) self.assertEqual(probe(torch.zeros(())).item(), 0)
@@ -704,6 +710,13 @@ class TestForwardFlags(_IsolatedServerArgs):
get_forward().set("is_extend_in_batch", True) get_forward().set("is_extend_in_batch", True)
self.assertEqual(probe(torch.zeros(())).item(), 2) self.assertEqual(probe(torch.zeros(())).item(), 2)
get_forward().set("is_extend_in_batch", False) get_forward().set("is_extend_in_batch", False)
with get_forward().scoped(
fuse_mlp_allreduce=True,
mlp_reduce_scatter=True,
flashinfer_trtllm_bypass=True,
):
self.assertEqual(probe(torch.zeros(())).item(), 28)
self.assertEqual(probe(torch.zeros(())).item(), 0)
def test_graph_visible_flags_are_process_visible_across_threads(self): def test_graph_visible_flags_are_process_visible_across_threads(self):
# Documented divergence from the contextvar-backed flags: plain slots # Documented divergence from the contextvar-backed flags: plain slots
@@ -812,6 +825,38 @@ class TestForwardFlags(_IsolatedServerArgs):
self.assertIs(get_forward().moe_output_buffer, sentinel) self.assertIs(get_forward().moe_output_buffer, sentinel)
self.assertIsNone(get_forward().moe_output_buffer) self.assertIsNone(get_forward().moe_output_buffer)
def test_mlp_comm_forward_flags(self):
"""Decoder-published MLP collective flags: scoped restore + skip helpers."""
from sglang.srt.layers.moe.utils import (
should_skip_mlp_all_reduce,
should_skip_post_experts_all_reduce,
)
from sglang.srt.runtime_context import get_forward
reset_context()
fwd = get_forward()
self.assertFalse(fwd.fuse_mlp_allreduce)
self.assertFalse(fwd.mlp_reduce_scatter)
self.assertFalse(fwd.flashinfer_trtllm_bypass)
self.assertFalse(should_skip_mlp_all_reduce())
with fwd.scoped(fuse_mlp_allreduce=True):
self.assertTrue(fwd.fuse_mlp_allreduce)
self.assertTrue(should_skip_mlp_all_reduce())
# Fusion alone is enough to skip post-experts AR.
self.assertTrue(should_skip_post_experts_all_reduce(is_tp_path=True))
self.assertFalse(fwd.fuse_mlp_allreduce)
self.assertFalse(should_skip_mlp_all_reduce())
with fwd.scoped(mlp_reduce_scatter=True):
self.assertTrue(fwd.mlp_reduce_scatter)
self.assertTrue(should_skip_mlp_all_reduce())
self.assertFalse(fwd.mlp_reduce_scatter)
with fwd.scoped(flashinfer_trtllm_bypass=True):
self.assertTrue(fwd.flashinfer_trtllm_bypass)
self.assertFalse(fwd.flashinfer_trtllm_bypass)
class TestPublishLifecycle(_IsolatedServerArgs): class TestPublishLifecycle(_IsolatedServerArgs):
"""Publish installs the resolved server_args and seeds the capture tier.""" """Publish installs the resolved server_args and seeds the capture tier."""