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,
mup_vector: Optional[torch.Tensor] = None,
use_triton_causal_conv: bool = False,
should_allreduce_fusion: bool = False,
):
assert isinstance(self.forward_metadata, Mamba2Metadata)
# Page-major stores state strided; only the stride-aware Triton causal-conv
@@ -759,7 +758,6 @@ class Mamba2AttnBackend(MambaAttnBackendBase):
forward_batch=forward_batch,
mup_vector=mup_vector,
use_triton_causal_conv=use_triton_causal_conv,
should_allreduce_fusion=should_allreduce_fusion,
)
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,
mup_vector: Optional[torch.Tensor] = None,
use_triton_causal_conv: bool = False,
should_allreduce_fusion: bool = False,
):
# Returns the projected result. When `output` is given it is also
# written into that buffer (required by the cuda-graph split ops, which
@@ -761,9 +760,7 @@ class MambaMixer2(torch.nn.Module):
# norm usage
hidden_states = self.norm(preallocated_ssm_out, gate)
mixer_out, _ = self.out_proj(
hidden_states, skip_all_reduce=should_allreduce_fusion
)
mixer_out, _ = self.out_proj(hidden_states)
if output is not None:
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.
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
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 (
is_allocation_symmetric,
)
from sglang.srt.layers.moe.utils import should_skip_mlp_all_reduce
from sglang.srt.layers.parameter import (
BasevLLMParameter,
BlockQuantScaleParameter,
@@ -1537,7 +1538,15 @@ class RowParallelLinear(LinearBase):
with symm_ctx:
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:
output = get_parallel().attn_tp_group.all_reduce(output_parallel)
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,
initialize_moe_config,
is_tbo_enabled,
should_skip_mlp_all_reduce,
should_skip_post_experts_all_reduce,
should_use_dp_reduce_scatterv,
should_use_flashinfer_cutlass_moe_fp4_allgather,
@@ -25,6 +26,7 @@ __all__ = [
"get_moe_a2a_backend",
"get_moe_runner_backend",
"get_deepep_mode",
"should_skip_mlp_all_reduce",
"should_skip_post_experts_all_reduce",
"should_use_dp_reduce_scatterv",
"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 (
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
_is_npu = is_npu()
@@ -411,20 +411,27 @@ def should_use_dp_reduce_scatterv():
)
def should_skip_post_experts_all_reduce(
*,
is_tp_path: bool,
use_reduce_scatter: bool = False,
should_allreduce_fusion: bool = False,
) -> bool:
def should_skip_mlp_all_reduce() -> bool:
"""Whether dense MLP / row-parallel projections should skip their all-reduce.
True when the decoder published ``fuse_mlp_allreduce`` (next residual+LN
absorbs the AR) or ``mlp_reduce_scatter`` (postprocess will reduce-scatter)
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
downstream component will fuse, replace, or absorb it.
Skip reasons, in order:
- ``should_allreduce_fusion``: LayerCommunicator will fuse the all-reduce
with the next layer's residual all-reduce.
- ``use_reduce_scatter``: LayerCommunicator's post-attention scatter will
do reduce-scatter, which would double-reduce on top of an all-reduce.
- ``get_forward().fuse_mlp_allreduce``: LayerCommunicator will fuse the
all-reduce with the next layer's residual all-reduce.
- ``get_forward().mlp_reduce_scatter``: LayerCommunicator's post-attention
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
path replaces the all-reduce with a reduce-scatterv.
- ``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
(``tensorrt_llm/_torch/modules/fused_moe/interface.py:879``).
The first two args are layer-context flags from ``LayerCommunicator`` and
default to ``False`` for models that don't use it. Pass ``is_tp_path=True``
The first two reasons come from per-layer ``ForwardFlags`` published by
the decoder via ``get_forward().scoped(...)``. Pass ``is_tp_path=True``
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
if should_use_dp_reduce_scatterv():
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.topk import TopKOutput
from sglang.srt.layers.moe.utils import should_skip_mlp_all_reduce
from sglang.srt.layers.vocab_parallel_embedding import (
ParallelLMHead,
VocabParallelEmbedding,
@@ -732,6 +733,7 @@ class RowParallelLinearWithLoRA(BaseLayerWithLoRA):
self.base_layer.reduce_results
and self.base_layer.tp_size > 1
and not skip_all_reduce
and not should_skip_mlp_all_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_reduce,
)
from sglang.srt.layers.moe.utils import should_skip_mlp_all_reduce
from sglang.srt.lora.trtllm_lora_temp import (
get_lora_side_stream,
get_original_column_forward,
@@ -125,6 +126,7 @@ def row_parallel_lora_forward(
self.base_layer.reduce_results
and self.base_layer.tp_size > 1
and not skip_all_reduce
and not should_skip_mlp_all_reduce()
)
if should_reduce:
python/sglang/srt/models/apertus.py
+1 -5
View File
@@ -96,15 +96,11 @@ class ApertusMLP(nn.Module):
self,
x,
forward_batch=None,
use_reduce_scatter: bool = False,
):
# note: with xielu, there's no gate_proj
x, _ = self.up_proj(x)
x = self.act_fn(x)
x, _ = self.down_proj(
x,
skip_all_reduce=use_reduce_scatter,
)
x, _ = self.down_proj(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,
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
LoraConfig = None
@@ -127,17 +132,13 @@ class BailingMoEMLP(nn.Module):
self,
hidden_states: torch.Tensor,
forward_batch: Optional[ForwardBatch] = None,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
) -> torch.Tensor:
if (self.tp_size == 1) and hidden_states.shape[0] == 0:
return hidden_states
gate_up, _ = self.gate_up_proj(hidden_states)
hidden_states = self.act_fn(gate_up)
hidden_states, _ = self.down_proj(
hidden_states, skip_all_reduce=should_allreduce_fusion or use_reduce_scatter
)
hidden_states, _ = self.down_proj(hidden_states)
return hidden_states
@@ -305,15 +306,9 @@ class BailingMoESparseMoeBlock(nn.Module):
self,
hidden_states: torch.Tensor,
forward_batch: Optional[ForwardBatch] = None,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
) -> torch.Tensor:
if not get_moe_a2a_backend().is_deepep():
return self.forward_normal(
hidden_states,
should_allreduce_fusion,
use_reduce_scatter,
)
return self.forward_normal(hidden_states)
else:
return self.forward_deepep(hidden_states, forward_batch)
@@ -356,8 +351,6 @@ class BailingMoESparseMoeBlock(nn.Module):
def forward_normal(
self,
hidden_states: torch.Tensor,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
) -> torch.Tensor:
num_tokens, hidden_size = hidden_states.shape
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(
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)
return final_hidden_states.view(num_tokens, hidden_size)
@@ -672,22 +663,24 @@ class BailingMoEBlock(nn.Module):
forward_batch=forward_batch,
)
should_allreduce_fusion = (
fuse_mlp_allreduce = (
self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer(
forward_batch
)
)
# 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
)
hidden_states = self.mlp(
hidden_states, forward_batch, should_allreduce_fusion, use_reduce_scatter
)
with get_forward().scoped(
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
else:
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.models.deepseek_v2 import DeepseekV2AttentionMLA, DeepseekV2MLP, _is_hip
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 (
BumpAllocator,
add_prefix,
@@ -189,15 +194,10 @@ class BailingMLP(nn.Module):
def forward(
self,
x,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
):
x, _ = self.gate_up_proj(x)
x = self.act_fn(x)
x, _ = self.down_proj(
x,
skip_all_reduce=use_reduce_scatter or should_allreduce_fusion,
)
x, _ = self.down_proj(x)
return x
@@ -332,8 +332,6 @@ class BailingMoE(nn.Module):
def forward(
self,
hidden_states: torch.Tensor,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
) -> torch.Tensor:
num_tokens, hidden_size = hidden_states.shape
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(
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)
return final_hidden_states
@@ -890,20 +886,25 @@ class BailingMoELinearDecoderLayer(nn.Module):
# logger.warning(
# 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(
forward_batch
)
)
use_reduce_scatter = self.layer_communicator.should_use_reduce_scatter(
mlp_reduce_scatter = self.layer_communicator.should_use_reduce_scatter(
forward_batch
)
hidden_states = self.mlp(
hidden_states, should_allreduce_fusion, use_reduce_scatter
)
hidden_states, residual = self.layer_communicator.postprocess_layer(
hidden_states, residual, forward_batch
)
with get_forward().scoped(
fuse_mlp_allreduce=fuse_mlp_allreduce,
mlp_reduce_scatter=mlp_reduce_scatter,
):
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
@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_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.utils import (
BumpAllocator,
@@ -292,8 +297,6 @@ class DeepseekV2MLP(nn.Module):
self,
x,
forward_batch=None,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
gemm_output_zero_allocator: BumpAllocator = None,
):
if (self.tp_size == 1) and x.shape[0] == 0:
@@ -321,10 +324,7 @@ class DeepseekV2MLP(nn.Module):
self.down_proj.input_scale_inv,
enable_pdl=True,
)
out, _ = self.down_proj(
(out_fp4, out_scale),
skip_all_reduce=should_allreduce_fusion or use_reduce_scatter,
)
out, _ = self.down_proj((out_fp4, out_scale))
return out
if (
@@ -426,10 +426,7 @@ class DeepseekV2MLP(nn.Module):
silu_and_mul_clamp(gate_up, x, float(self.swiglu_limit))
else:
x = self.act_fn(gate_up)
x, _ = self.down_proj(
x,
skip_all_reduce=should_allreduce_fusion or use_reduce_scatter,
)
x, _ = self.down_proj(x)
return x
@@ -857,8 +854,6 @@ class DeepseekV2MoE(nn.Module):
self,
hidden_states: torch.Tensor,
forward_batch: Optional[ForwardBatch] = None,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
gemm_output_zero_allocator: BumpAllocator = None,
input_ids: 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(
hidden_states,
self.layer_id,
should_allreduce_fusion,
use_reduce_scatter,
)
elif (
self.alt_stream is not None
@@ -897,8 +890,6 @@ class DeepseekV2MoE(nn.Module):
):
return self.forward_normal_dual_stream(
hidden_states,
should_allreduce_fusion,
use_reduce_scatter,
gemm_output_zero_allocator,
input_ids,
input_ids_global=input_ids_global,
@@ -906,8 +897,6 @@ class DeepseekV2MoE(nn.Module):
else:
return self.forward_normal(
hidden_states,
should_allreduce_fusion,
use_reduce_scatter,
gemm_output_zero_allocator,
input_ids,
input_ids_global=input_ids_global,
@@ -921,19 +910,16 @@ class DeepseekV2MoE(nn.Module):
def forward_normal_dual_stream(
self,
hidden_states: torch.Tensor,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
gemm_output_zero_allocator: BumpAllocator = None,
input_ids: Optional[torch.Tensor] = None,
input_ids_global: Optional[torch.Tensor] = None,
*,
use_flashinfer_trtllm_bypass: bool = False,
) -> torch.Tensor:
# Note(kpham-sgl): issue order satisfies 3 constraints:
# - 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);
# - 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.
use_flashinfer_trtllm_bypass = get_forward().flashinfer_trtllm_bypass
current_stream = torch.cuda.current_stream()
self.alt_stream.wait_stream(current_stream)
has_shared_output = (
@@ -1014,8 +1000,6 @@ class DeepseekV2MoE(nn.Module):
if self.tp_size > 1 and not should_skip_post_experts_all_reduce(
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)
# TP1 shared experts are replicated, so add them after all-reduce to
@@ -1027,8 +1011,6 @@ class DeepseekV2MoE(nn.Module):
def forward_normal(
self,
hidden_states: torch.Tensor,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
gemm_output_zero_allocator: BumpAllocator = None,
input_ids: 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(
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()
dispatch_info = (
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(
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)
# TP1 shared experts are replicated, so add them after all-reduce to
@@ -1153,7 +1133,6 @@ class DeepseekV2MoE(nn.Module):
def forward_cpu(
self,
hidden_states: torch.Tensor,
should_allreduce_fusion: bool = False,
) -> torch.Tensor:
# router_logits: (num_tokens, n_experts)
router_logits = self.gate(hidden_states)
@@ -1202,7 +1181,7 @@ class DeepseekV2MoE(nn.Module):
), # block_size
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)
return final_hidden_states
@@ -2259,14 +2238,14 @@ class DeepseekV2DecoderLayer(nn.Module):
hidden_states, residual, forward_batch
)
should_allreduce_fusion = (
fuse_mlp_allreduce = (
self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer(
forward_batch
)
)
# 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
)
@@ -2284,22 +2263,24 @@ class DeepseekV2DecoderLayer(nn.Module):
else:
_mlp_ctx = nullcontext()
with _mlp_ctx:
hidden_states = self.mlp(
hidden_states,
forward_batch,
should_allreduce_fusion,
use_reduce_scatter,
gemm_output_zero_allocator,
)
with get_forward().scoped(
fuse_mlp_allreduce=fuse_mlp_allreduce,
mlp_reduce_scatter=mlp_reduce_scatter,
):
with _mlp_ctx:
hidden_states = self.mlp(
hidden_states,
forward_batch,
gemm_output_zero_allocator,
)
if (
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
if not should_allreduce_fusion:
if not fuse_mlp_allreduce:
hidden_states, residual = self.layer_communicator.postprocess_layer(
hidden_states, residual, forward_batch
)
@@ -2996,8 +2977,6 @@ class DeepseekV32ForCausalLM(DeepseekV2ForCausalLM):
def dsv2_flashinfer_moe_dual_stream_graph(
hidden_states: torch.Tensor,
layer_id: int,
should_allreduce_fusion: bool,
use_reduce_scatter: bool,
) -> torch.Tensor:
forward_context = get_tc_piecewise_forward_context()
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]
assert moe_fusion is not None
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,
)
with get_forward().scoped(flashinfer_trtllm_bypass=True):
return moe_fusion.forward_normal_dual_stream(hidden_states)
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_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:
from sglang.srt.layers.utils.cp_utils import (
@@ -134,7 +134,6 @@ if _is_xpu:
else:
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 (
LazyValue,
add_prefix,
@@ -1639,7 +1638,7 @@ class DeepseekV4DecoderLayer(nn.Module):
# 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
# 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.
_use_reduce_scatterv = (
_use_tp_moe_gather
@@ -1663,6 +1662,7 @@ class DeepseekV4DecoderLayer(nn.Module):
and forward_batch.dp_padding_mode.is_max_len()
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
# 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.
@@ -1702,17 +1702,17 @@ class DeepseekV4DecoderLayer(nn.Module):
hidden_states = _a2a_scatter_chunks[r].contiguous()
input_ids = input_ids.tensor_split(s)[r].contiguous()
input_ids_global = input_ids_global.tensor_split(s)[r].contiguous()
hidden_states = self.mlp(
hidden_states,
forward_batch,
input_ids=input_ids,
input_ids_global=input_ids_global,
# Skip the MoE-internal post-experts all_reduce when we will do the
# reduce via reduce_scatterv/reduce_scatter at the combine below
# (else double-reduce).
use_reduce_scatter=_use_cp or _use_reduce_scatterv or _use_reduce_scatter,
skip_shared_experts=_do_shared_local,
)
# Skip the MoE-internal post-experts all_reduce when we will do the
# reduce via reduce_scatterv/reduce_scatter at the combine below
# (else double-reduce).
with get_forward().scoped(mlp_reduce_scatter=mlp_reduce_scatter):
hidden_states = self.mlp(
hidden_states,
forward_batch,
input_ids=input_ids,
input_ids_global=input_ids_global,
skip_shared_experts=_do_shared_local,
)
if _use_cp and get_moe_a2a_backend().is_none():
hidden_states = dsa_cp_reduce_scatter_hidden_states(hidden_states)
elif _use_tp_moe_gather:
@@ -1723,7 +1723,7 @@ class DeepseekV4DecoderLayer(nn.Module):
if should_use_dp_reduce_scatterv() or _use_reduce_scatterv:
# SUM the TP-sharded per-rank partial expert outputs AND scatter
# 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.
get_tp_group().reduce_scatterv(
global_hidden_states,
@@ -1735,7 +1735,7 @@ class DeepseekV4DecoderLayer(nn.Module):
# expert outputs AND scatter each rank its own (MAX_LEN-padded)
# token chunk in one op (symmetric inverse of the MAX_LEN
# 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
# variable-length reduce_scatterv branch is gated by
# is_dp_gatherv_active(), which is False under MAX_LEN), which in
@@ -1920,19 +1920,19 @@ class DeepseekV4DecoderLayer(nn.Module):
state.pop("gather_keepalive")
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).
fb = state.forward_batch
global_hidden = state.pop("global_hidden")
global_ids = fb._tbo_global_input_ids
state.global_expert_out = self.mlp(
global_hidden,
fb,
use_reduce_scatter=True,
input_ids=global_ids,
input_ids_global=global_ids,
skip_shared_experts=state.do_shared_local,
)
with get_forward().scoped(mlp_reduce_scatter=True):
state.global_expert_out = self.mlp(
global_hidden,
fb,
input_ids=global_ids,
input_ids_global=global_ids,
skip_shared_experts=state.do_shared_local,
)
def op_combine_a(self, state):
# 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,
x,
forward_batch=None,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
):
gate_up, _ = self.gate_up_proj(x)
x = self.act_fn(gate_up)
x, _ = self.down_proj(
x,
skip_all_reduce=should_allreduce_fusion or use_reduce_scatter,
)
x, _ = self.down_proj(x)
return x
@@ -276,7 +271,6 @@ class ExaoneMoESparseMoEBlock(nn.Module):
self,
hidden_states: torch.Tensor,
forward_batch: Optional[ForwardBatch] = None,
use_reduce_scatter: bool = False,
) -> torch.Tensor:
num_tokens, hidden_dim = hidden_states.shape
hidden_states = hidden_states.view(-1, hidden_dim)
@@ -300,7 +294,6 @@ class ExaoneMoESparseMoEBlock(nn.Module):
final_hidden_states = final_hidden_states + shared_output
if self.tp_size > 1 and not should_skip_post_experts_all_reduce(
is_tp_path=True,
use_reduce_scatter=use_reduce_scatter,
):
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_context import get_attn_backend
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
logger = logging.getLogger(__name__)
@@ -85,16 +90,12 @@ class FalconH1MLP(nn.Module):
self,
x,
forward_batch=None,
use_reduce_scatter: bool = False,
):
gate_up, _ = self.gate_up_proj(x)
gate_up[:, : self.intermediate_size // self.tp_size] *= self.gate_multiplier
x = self.act_fn(gate_up)
x, _ = self.down_proj(
x,
skip_all_reduce=use_reduce_scatter,
)
x, _ = self.down_proj(x)
x = x * self.down_multiplier
return x
@@ -358,12 +359,11 @@ class FalconH1HybridAttentionDecoderLayer(nn.Module):
hidden_states, residual = self.layer_communicator.prepare_mlp(
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
)
hidden_states = self.feed_forward(
hidden_states, forward_batch, use_reduce_scatter
)
with get_forward().scoped(mlp_reduce_scatter=mlp_reduce_scatter):
hidden_states = self.feed_forward(hidden_states, forward_batch)
hidden_states, residual = self.layer_communicator.postprocess_layer(
hidden_states, residual, forward_batch
python/sglang/srt/models/glm4.py
+1 -5
View File
@@ -94,14 +94,10 @@ class Glm4MLP(nn.Module):
self,
x,
forward_batch=None,
use_reduce_scatter: bool = False,
):
gate_up, _ = self.gate_up_proj(x)
x = self.act_fn(gate_up)
x, _ = self.down_proj(
x,
skip_all_reduce=use_reduce_scatter,
)
x, _ = self.down_proj(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.utils import WeightsMapper, apply_qk_norm
from sglang.srt.runtime_context import (
get_forward,
get_parallel,
get_server_args,
get_stream,
@@ -167,17 +168,13 @@ class Glm4MoeMLP(nn.Module):
self,
x,
forward_batch=None,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
):
if (self.tp_size == 1) and x.shape[0] == 0:
return x
gate_up, _ = self.gate_up_proj(x)
x = self.act_fn(gate_up)
x, _ = self.down_proj(
x, skip_all_reduce=should_allreduce_fusion or use_reduce_scatter
)
x, _ = self.down_proj(x)
return x
@@ -564,8 +561,6 @@ class Glm4MoeSparseMoeBlock(nn.Module):
self,
hidden_states: torch.Tensor,
forward_batch: Optional[ForwardBatch] = None,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
) -> torch.Tensor:
if not self._enable_a2a_moe:
if (
@@ -574,25 +569,15 @@ class Glm4MoeSparseMoeBlock(nn.Module):
and hidden_states.shape[0] > 0
and get_is_capture_mode()
):
return self.forward_normal_dual_stream(
hidden_states,
should_allreduce_fusion,
use_reduce_scatter,
)
return self.forward_normal_dual_stream(hidden_states)
else:
return self.forward_normal(
hidden_states,
should_allreduce_fusion,
use_reduce_scatter,
)
return self.forward_normal(hidden_states)
else:
return self.forward_deepep(hidden_states, forward_batch)
def forward_normal_dual_stream(
self,
hidden_states: torch.Tensor,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
) -> torch.Tensor:
current_stream = torch.cuda.current_stream()
self.alt_stream.wait_stream(current_stream)
@@ -610,8 +595,6 @@ class Glm4MoeSparseMoeBlock(nn.Module):
final_hidden_states += shared_output
if self.tp_size > 1 and not should_skip_post_experts_all_reduce(
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)
return final_hidden_states
@@ -619,8 +602,6 @@ class Glm4MoeSparseMoeBlock(nn.Module):
def forward_normal(
self,
hidden_states: torch.Tensor,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
) -> torch.Tensor:
if hidden_states.shape[0] > 0:
shared_output = self._forward_shared_experts(hidden_states)
@@ -643,8 +624,6 @@ class Glm4MoeSparseMoeBlock(nn.Module):
final_hidden_states = final_hidden_states_out
if self.tp_size > 1 and not should_skip_post_experts_all_reduce(
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)
return final_hidden_states
@@ -958,22 +937,24 @@ class Glm4MoeDecoderLayer(nn.Module):
hidden_states, residual, forward_batch
)
should_allreduce_fusion = (
fuse_mlp_allreduce = (
self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer(
forward_batch
)
)
# 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
)
hidden_states = self.mlp(
hidden_states, forward_batch, should_allreduce_fusion, use_reduce_scatter
)
with get_forward().scoped(
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
else:
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_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 (
BumpAllocator,
LazyValue,
@@ -132,17 +137,13 @@ class Glm4MoeLiteMLP(nn.Module):
self,
x,
forward_batch=None,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
):
if (self.tp_size == 1) and x.shape[0] == 0:
return x
gate_up, _ = self.gate_up_proj(x)
x = self.act_fn(gate_up)
x, _ = self.down_proj(
x, skip_all_reduce=should_allreduce_fusion or use_reduce_scatter
)
x, _ = self.down_proj(x)
return x
@@ -313,8 +314,6 @@ class Glm4MoeLiteSparseMoeBlock(nn.Module):
self,
hidden_states: torch.Tensor,
forward_batch: Optional[ForwardBatch] = None,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
) -> torch.Tensor:
if not self._enable_a2a_moe:
if (
@@ -323,21 +322,15 @@ class Glm4MoeLiteSparseMoeBlock(nn.Module):
and hidden_states.shape[0] > 0
and get_is_capture_mode()
):
return self.forward_normal_dual_stream(
hidden_states, should_allreduce_fusion, use_reduce_scatter
)
return self.forward_normal_dual_stream(hidden_states)
else:
return self.forward_normal(
hidden_states, should_allreduce_fusion, use_reduce_scatter
)
return self.forward_normal(hidden_states)
else:
return self.forward_deepep(hidden_states, forward_batch)
def forward_normal_dual_stream(
self,
hidden_states: torch.Tensor,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
) -> torch.Tensor:
current_stream = torch.cuda.current_stream()
self.alt_stream.wait_stream(current_stream)
@@ -355,8 +348,6 @@ class Glm4MoeLiteSparseMoeBlock(nn.Module):
final_hidden_states += shared_output
if self.tp_size > 1 and not should_skip_post_experts_all_reduce(
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)
return final_hidden_states
@@ -364,8 +355,6 @@ class Glm4MoeLiteSparseMoeBlock(nn.Module):
def forward_normal(
self,
hidden_states: torch.Tensor,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
) -> torch.Tensor:
if hidden_states.shape[0] > 0:
shared_output = self._forward_shared_experts(hidden_states)
@@ -388,8 +377,6 @@ class Glm4MoeLiteSparseMoeBlock(nn.Module):
final_hidden_states = final_hidden_states_out
if self.tp_size > 1 and not should_skip_post_experts_all_reduce(
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)
return final_hidden_states
@@ -679,22 +666,24 @@ class Glm4MoeLiteDecoderLayer(nn.Module):
hidden_states, residual, forward_batch
)
should_allreduce_fusion = (
fuse_mlp_allreduce = (
self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer(
forward_batch
)
)
# 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
)
hidden_states = self.mlp(
hidden_states, forward_batch, should_allreduce_fusion, use_reduce_scatter
)
with get_forward().scoped(
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
else:
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,
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 (
LazyValue,
add_prefix,
@@ -254,10 +254,9 @@ class GptOssSparseMoeBlock(nn.Module):
self,
hidden_states: torch.Tensor,
forward_batch: Optional[ForwardBatch] = None,
should_allreduce_fusion: bool = False,
) -> torch.Tensor:
if not get_moe_a2a_backend().is_deepep():
return self.forward_normal(hidden_states, should_allreduce_fusion)
return self.forward_normal(hidden_states)
else:
raise Exception("forward_deepep branch not implemented yet")
@@ -274,7 +273,6 @@ class GptOssSparseMoeBlock(nn.Module):
def forward_normal(
self,
hidden_states: torch.Tensor,
should_allreduce_fusion: bool = False,
) -> torch.Tensor:
# `hidden_states` may arrive pre-padded along the last dim when the
# 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)
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)
# When input was pre-padded, FusedMoE.forward_impl captured the
@@ -603,18 +601,19 @@ class GptOssDecoderLayer(nn.Module):
hidden_states, residual, forward_batch
)
should_allreduce_fusion = (
fuse_mlp_allreduce = (
self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer(
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
if not should_allreduce_fusion:
if not fuse_mlp_allreduce:
hidden_states, residual = self.layer_communicator.postprocess_layer(
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_loader.weight_utils import default_weight_loader
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
logger = logging.getLogger(__name__)
@@ -101,17 +101,12 @@ class LagunaMLP(nn.Module):
self,
x: torch.Tensor,
forward_batch: Optional[ForwardBatch] = None,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
) -> torch.Tensor:
gate_up, _ = self.gate_up_proj(x)
x = self.act_fn(gate_up)
# Skip the in-block reduce when LayerCommunicator will fuse it or when
# the next layer expects reduce-scatter — otherwise we'd double-reduce.
x, _ = self.down_proj(
x,
skip_all_reduce=should_allreduce_fusion or use_reduce_scatter,
)
# RowParallelLinear honors ForwardFlags (fuse_mlp_allreduce /
# mlp_reduce_scatter) published by the decoder via scoped().
x, _ = self.down_proj(x)
return x
@@ -204,8 +199,6 @@ class LagunaMoE(nn.Module):
self,
hidden_states: torch.Tensor,
forward_batch: Optional[ForwardBatch] = None,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
) -> torch.Tensor:
if hidden_states.shape[0] == 0:
return hidden_states
@@ -233,8 +226,6 @@ class LagunaMoE(nn.Module):
if self.tp_size > 1 and not should_skip_post_experts_all_reduce(
is_tp_path=True,
use_reduce_scatter=use_reduce_scatter,
should_allreduce_fusion=should_allreduce_fusion,
):
final = tensor_model_parallel_all_reduce(final)
if self._shared_expert_tp1:
@@ -502,23 +493,25 @@ class LagunaDecoderLayer(nn.Module):
hidden_states, residual, forward_batch
)
should_allreduce_fusion = (
fuse_mlp_allreduce = (
self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer(
forward_batch
)
)
use_reduce_scatter = self.layer_communicator.should_use_reduce_scatter(
mlp_reduce_scatter = self.layer_communicator.should_use_reduce_scatter(
forward_batch
)
hidden_states = self.mlp(
hidden_states,
forward_batch=forward_batch,
should_allreduce_fusion=should_allreduce_fusion,
use_reduce_scatter=use_reduce_scatter,
)
with get_forward().scoped(
fuse_mlp_allreduce=fuse_mlp_allreduce,
mlp_reduce_scatter=mlp_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
else:
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,
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,
@@ -134,16 +139,13 @@ class LLaDA2MoeMLP(nn.Module):
self,
hidden_states: torch.Tensor,
forward_batch: Optional[ForwardBatch] = None,
use_reduce_scatter: bool = False,
) -> torch.Tensor:
if (self.tp_size == 1) and hidden_states.shape[0] == 0:
return hidden_states
gate_up, _ = self.gate_up_proj(hidden_states)
hidden_states = self.act_fn(gate_up)
hidden_states, _ = self.down_proj(
hidden_states, skip_all_reduce=use_reduce_scatter
)
hidden_states, _ = self.down_proj(hidden_states)
return hidden_states
@@ -317,10 +319,9 @@ class LLaDA2MoeSparseMoeBlock(nn.Module):
self,
hidden_states: torch.Tensor,
forward_batch: Optional[ForwardBatch] = None,
use_reduce_scatter: bool = False,
) -> torch.Tensor:
if not get_moe_a2a_backend().is_deepep():
return self.forward_normal(hidden_states, use_reduce_scatter)
return self.forward_normal(hidden_states)
else:
return self.forward_deepep(hidden_states, forward_batch)
@@ -360,7 +361,6 @@ class LLaDA2MoeSparseMoeBlock(nn.Module):
def forward_normal(
self,
hidden_states: torch.Tensor,
use_reduce_scatter: bool = False,
) -> torch.Tensor:
num_tokens, hidden_size = hidden_states.shape
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(
is_tp_path=True,
use_reduce_scatter=use_reduce_scatter,
):
final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states)
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.
use_reduce_scatter = self.layer_communicator.should_use_reduce_scatter(
mlp_reduce_scatter = self.layer_communicator.should_use_reduce_scatter(
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=hidden_states,
python/sglang/srt/models/llama.py
+1 -5
View File
@@ -112,14 +112,10 @@ class LlamaMLP(nn.Module):
self,
x,
forward_batch=None,
use_reduce_scatter: bool = False,
):
gate_up, _ = self.gate_up_proj(x)
x = self.act_fn(gate_up)
x, _ = self.down_proj(
x,
skip_all_reduce=use_reduce_scatter,
)
x, _ = self.down_proj(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.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 (
add_prefix,
fast_topk,
@@ -139,7 +139,6 @@ class Llama4MoE(nn.Module):
self,
hidden_states,
forward_batch: ForwardBatch,
use_reduce_scatter: bool = False,
):
shared_out, routed_out = self._forward_core(
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(
is_tp_path=True,
use_reduce_scatter=use_reduce_scatter,
):
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.
use_reduce_scatter = self.layer_communicator.should_use_reduce_scatter(
mlp_reduce_scatter = self.layer_communicator.should_use_reduce_scatter(
forward_batch
)
# Fully Connected
hidden_states = self.feed_forward(
hidden_states, forward_batch, use_reduce_scatter
)
with get_forward().scoped(mlp_reduce_scatter=mlp_reduce_scatter):
hidden_states = self.feed_forward(hidden_states, forward_batch)
hidden_states, residual = self.layer_communicator.postprocess_layer(
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_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 (
LazyValue,
add_prefix,
@@ -166,17 +166,13 @@ class MiMoV2MLP(nn.Module):
self,
x,
forward_batch: ForwardBatch = None,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
):
if (self.tp_size == 1) and x.shape[0] == 0:
return x
gate_up, _ = self.gate_up_proj(x)
x = self.act_fn(gate_up)
x, _ = self.down_proj(
x, skip_all_reduce=should_allreduce_fusion or use_reduce_scatter
)
x, _ = self.down_proj(x)
return x
@@ -315,23 +311,15 @@ class MiMoV2MoE(nn.Module):
self,
hidden_states: torch.Tensor,
forward_batch: Optional[ForwardBatch] = None,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
) -> torch.Tensor:
if not self._enable_a2a_moe:
return self.forward_normal(
hidden_states,
should_allreduce_fusion,
use_reduce_scatter,
)
return self.forward_normal(hidden_states)
else:
return self.forward_deepep(hidden_states, forward_batch)
def forward_normal(
self,
hidden_states: torch.Tensor,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
) -> torch.Tensor:
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(
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)
@@ -746,22 +732,24 @@ class MiMoV2DecoderLayer(nn.Module):
hidden_states, residual, forward_batch
)
should_allreduce_fusion = (
fuse_mlp_allreduce = (
self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer(
forward_batch
)
)
# 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
)
hidden_states = self.mlp(
hidden_states, forward_batch, should_allreduce_fusion, use_reduce_scatter
)
with get_forward().scoped(
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
else:
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,
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.
# Importing from the wrong module causes this file to fail import, which prevents the
@@ -553,24 +553,18 @@ class MiniMaxM2MoE(nn.Module):
self,
hidden_states: torch.Tensor,
forward_batch: Optional[ForwardBatch] = None,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
) -> torch.Tensor:
if (
not get_moe_a2a_backend().is_deepep()
and not get_moe_a2a_backend().is_ascend_fuseep()
):
return self.forward_normal(
hidden_states, should_allreduce_fusion, use_reduce_scatter
)
return self.forward_normal(hidden_states)
else:
return self.forward_deepep(hidden_states, forward_batch)
def forward_normal(
self,
hidden_states: torch.Tensor,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
) -> torch.Tensor:
num_tokens, hidden_dim = hidden_states.shape
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)
if self.tp_size > 1 and not should_skip_post_experts_all_reduce(
is_tp_path=True,
use_reduce_scatter=use_reduce_scatter,
should_allreduce_fusion=should_allreduce_fusion,
is_tp_path=True
):
final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states)
@@ -1012,21 +1004,23 @@ class MiniMaxM2DecoderLayer(nn.Module):
hidden_states, residual, forward_batch
)
should_allreduce_fusion = (
fuse_mlp_allreduce = (
self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer(
forward_batch
)
)
use_reduce_scatter = self.layer_communicator.should_use_reduce_scatter(
mlp_reduce_scatter = self.layer_communicator.should_use_reduce_scatter(
forward_batch
)
hidden_states = self.block_sparse_moe(
hidden_states, forward_batch, should_allreduce_fusion, use_reduce_scatter
)
with get_forward().scoped(
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
else:
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,
)
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 (
add_prefix,
get_current_device_stream_fast,
@@ -132,14 +132,10 @@ class NemotronHMLP(nn.Module):
def forward(
self,
x: torch.Tensor,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
):
x, _ = self.up_proj(x)
x = self.act_fn(x)
x, _ = self.down_proj(
x, skip_all_reduce=should_allreduce_fusion or use_reduce_scatter
)
x, _ = self.down_proj(x)
return x
@@ -306,8 +302,6 @@ class NemotronHMoE(nn.Module):
def forward(
self,
hidden_states: torch.Tensor,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
) -> torch.Tensor:
num_tokens, hidden_dim = hidden_states.shape
# 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(
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)
@@ -344,20 +336,20 @@ class NemotronHMLPLikeDecoderLayer(nn.Module):
hidden_states, residual = self.layer_communicator.prepare_mlp(
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
)
should_allreduce_fusion = (
fuse_mlp_allreduce = (
self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer(
forward_batch
)
)
hidden_states = self.mixer.forward(
hidden_states,
should_allreduce_fusion=should_allreduce_fusion,
use_reduce_scatter=use_reduce_scatter,
)
if should_allreduce_fusion:
with get_forward().scoped(
fuse_mlp_allreduce=fuse_mlp_allreduce,
mlp_reduce_scatter=mlp_reduce_scatter,
):
hidden_states = self.mixer.forward(hidden_states)
if fuse_mlp_allreduce:
hidden_states._sglang_needs_allreduce_fusion = True
else:
hidden_states, residual = self.layer_communicator.postprocess_layer(
@@ -369,15 +361,14 @@ class NemotronHMLPLikeDecoderLayer(nn.Module):
self.norm, hidden_states, residual
)
should_allreduce_fusion = (
fuse_mlp_allreduce = (
self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer(
forward_batch
)
)
hidden_states = self.mixer.forward(
hidden_states, should_allreduce_fusion=should_allreduce_fusion
)
if should_allreduce_fusion:
with get_forward().scoped(fuse_mlp_allreduce=fuse_mlp_allreduce):
hidden_states = self.mixer.forward(hidden_states)
if fuse_mlp_allreduce:
hidden_states._sglang_needs_allreduce_fusion = True
return hidden_states, residual
@@ -505,7 +496,6 @@ class NemotronHMambaDecoderLayer(NemotronHAttnLikeDecoderLayer):
self,
hidden_states: torch.Tensor,
forward_batch: ForwardBatch,
should_allreduce_fusion: bool = False,
) -> torch.Tensor:
"""Core Mamba forward logic, called directly or via split op."""
original_num_tokens = hidden_states.shape[0]
@@ -523,7 +513,6 @@ class NemotronHMambaDecoderLayer(NemotronHAttnLikeDecoderLayer):
output=None,
forward_batch=forward_batch,
use_triton_causal_conv=True,
should_allreduce_fusion=should_allreduce_fusion,
)
return pad_to_original_num_tokens(output, original_num_tokens)
@@ -551,28 +540,25 @@ class NemotronHMambaDecoderLayer(NemotronHAttnLikeDecoderLayer):
self.norm, hidden_states, residual
)
should_allreduce_fusion = (
fuse_mlp_allreduce = (
self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer(
forward_batch
)
)
if is_in_breakable_cuda_graph():
output = torch.empty_like(hidden_states)
breakable_nemotron_mamba2_with_output(
hidden_states, output, self.layer_id, should_allreduce_fusion
)
elif is_in_tc_piecewise_cuda_graph():
output = torch.empty_like(hidden_states)
nemotron_mamba2_with_output(
hidden_states, output, self.layer_id, should_allreduce_fusion
)
else:
output = self._forward_mamba(
hidden_states, forward_batch, should_allreduce_fusion
)
with get_forward().scoped(fuse_mlp_allreduce=fuse_mlp_allreduce):
if is_in_breakable_cuda_graph():
output = torch.empty_like(hidden_states)
breakable_nemotron_mamba2_with_output(
hidden_states, output, self.layer_id
)
elif is_in_tc_piecewise_cuda_graph():
output = torch.empty_like(hidden_states)
nemotron_mamba2_with_output(hidden_states, output, self.layer_id)
else:
output = self._forward_mamba(hidden_states, forward_batch)
if should_allreduce_fusion:
if fuse_mlp_allreduce:
output._sglang_needs_allreduce_fusion = True
return output, residual
@@ -647,15 +633,12 @@ class NemotronHAttention(nn.Module):
self,
hidden_states: torch.Tensor,
forward_batch: ForwardBatch,
should_allreduce_fusion: bool = False,
) -> torch.Tensor:
if not is_dp_attention_enabled():
qkv, _ = self.qkv_proj(hidden_states)
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)
output, _ = self.o_proj(
attn_output, skip_all_reduce=should_allreduce_fusion
)
output, _ = self.o_proj(attn_output)
return output
padded_shape = hidden_states.shape[0]
@@ -733,18 +716,18 @@ class NemotronHAttentionDecoderLayer(NemotronHAttnLikeDecoderLayer):
self.norm, hidden_states, residual
)
should_allreduce_fusion = (
fuse_mlp_allreduce = (
self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer(
forward_batch
)
)
hidden_states = self.mixer.forward(
hidden_states=hidden_states,
forward_batch=forward_batch,
should_allreduce_fusion=should_allreduce_fusion,
)
if should_allreduce_fusion:
with get_forward().scoped(fuse_mlp_allreduce=fuse_mlp_allreduce):
hidden_states = self.mixer.forward(
hidden_states=hidden_states,
forward_batch=forward_batch,
)
if fuse_mlp_allreduce:
hidden_states._sglang_needs_allreduce_fusion = True
return hidden_states, residual
@@ -1206,7 +1189,6 @@ def nemotron_mamba2_with_output(
hidden_states: torch.Tensor,
output: torch.Tensor,
layer_id: int,
should_allreduce_fusion: bool = False,
) -> None:
"""Split op for Mamba2 forward in piecewise CUDA graph mode."""
context = get_tc_piecewise_forward_context()
@@ -1226,9 +1208,7 @@ def nemotron_mamba2_with_output(
if hidden_states.shape[0] != num_actual_tokens:
hidden_states = hidden_states[:num_actual_tokens]
ret = mamba_layer._forward_mamba(
hidden_states, forward_batch, should_allreduce_fusion
)
ret = mamba_layer._forward_mamba(hidden_states, forward_batch)
# Copy result back; output may be larger (padded) so only fill actual tokens
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.runner import get_is_capture_mode
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 (
add_prefix,
cpu_has_amx_support,
@@ -209,14 +209,10 @@ class Qwen2MoeMLP(nn.Module):
def forward(
self,
x,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
):
gate_up, _ = self.gate_up_proj(x)
x = self.act_fn(gate_up)
x, _ = self.down_proj(
x, skip_all_reduce=should_allreduce_fusion or use_reduce_scatter
)
x, _ = self.down_proj(x)
return x
@@ -545,8 +541,6 @@ class Qwen2MoeSparseMoeBlock(nn.Module):
self,
hidden_states: torch.Tensor,
forward_batch: Optional[ForwardBatch] = None,
use_reduce_scatter: bool = False,
should_allreduce_fusion: bool = False,
) -> torch.Tensor:
num_tokens, hidden_dim = hidden_states.shape
hidden_states = hidden_states.view(-1, hidden_dim)
@@ -591,8 +585,6 @@ class Qwen2MoeSparseMoeBlock(nn.Module):
self.tp_size > 1
and not should_skip_post_experts_all_reduce(
is_tp_path=True,
use_reduce_scatter=use_reduce_scatter,
should_allreduce_fusion=should_allreduce_fusion,
)
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.
use_reduce_scatter = self.layer_communicator.should_use_reduce_scatter(
mlp_reduce_scatter = self.layer_communicator.should_use_reduce_scatter(
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, 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_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
from sglang.srt.utils import (
@@ -682,27 +687,27 @@ class Qwen3_5LinearDecoderLayer(nn.Module):
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
)
should_allreduce_fusion = (
fuse_mlp_allreduce = (
self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer(
forward_batch
)
)
if isinstance(self.mlp, Qwen2MoeSparseMoeBlock):
hidden_states = self.mlp(
hidden_states,
forward_batch,
use_reduce_scatter,
should_allreduce_fusion,
)
else:
hidden_states = self.mlp(
hidden_states, should_allreduce_fusion, use_reduce_scatter
)
if should_allreduce_fusion:
with get_forward().scoped(
fuse_mlp_allreduce=fuse_mlp_allreduce,
mlp_reduce_scatter=mlp_reduce_scatter,
):
if isinstance(self.mlp, Qwen2MoeSparseMoeBlock):
hidden_states = self.mlp(
hidden_states,
forward_batch,
)
else:
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(
@@ -1067,27 +1072,27 @@ class Qwen3_5AttentionDecoderLayer(nn.Module):
hidden_states, residual = self.layer_communicator.prepare_mlp(
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
)
should_allreduce_fusion = (
fuse_mlp_allreduce = (
self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer(
forward_batch
)
)
if isinstance(self.mlp, Qwen2MoeSparseMoeBlock):
hidden_states = self.mlp(
hidden_states,
forward_batch,
use_reduce_scatter,
should_allreduce_fusion,
)
else:
hidden_states = self.mlp(
hidden_states, should_allreduce_fusion, use_reduce_scatter
)
if should_allreduce_fusion:
with get_forward().scoped(
fuse_mlp_allreduce=fuse_mlp_allreduce,
mlp_reduce_scatter=mlp_reduce_scatter,
):
if isinstance(self.mlp, Qwen2MoeSparseMoeBlock):
hidden_states = self.mlp(
hidden_states,
forward_batch,
)
else:
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(
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,
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,
@@ -286,17 +291,13 @@ class Qwen3MoeSparseMoeBlock(nn.Module):
self,
hidden_states: torch.Tensor,
forward_batch: Optional[ForwardBatch] = None,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
) -> torch.Tensor:
if (
not get_moe_a2a_backend().is_deepep()
and not get_moe_a2a_backend().is_ascend_fuseep()
):
return self.forward_normal(
hidden_states, should_allreduce_fusion, use_reduce_scatter
)
return self.forward_normal(hidden_states)
else:
return self.forward_deepep(hidden_states, forward_batch)
@@ -313,8 +314,6 @@ class Qwen3MoeSparseMoeBlock(nn.Module):
def forward_normal(
self,
hidden_states: torch.Tensor,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
) -> torch.Tensor:
num_tokens, hidden_dim = hidden_states.shape
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)
if self.ep_size > 1 and not should_skip_post_experts_all_reduce(
is_tp_path=False,
use_reduce_scatter=use_reduce_scatter,
should_allreduce_fusion=should_allreduce_fusion,
is_tp_path=False
):
final_hidden_states = moe_expert_parallel_all_reduce(final_hidden_states)
if self.tp_size > 1 and not should_skip_post_experts_all_reduce(
is_tp_path=True,
use_reduce_scatter=use_reduce_scatter,
should_allreduce_fusion=should_allreduce_fusion,
is_tp_path=True
):
final_hidden_states = moe_tensor_model_parallel_all_reduce(
final_hidden_states
@@ -828,22 +823,24 @@ class Qwen3MoeDecoderLayer(nn.Module):
hidden_states, residual, forward_batch
)
should_allreduce_fusion = (
fuse_mlp_allreduce = (
self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer(
forward_batch
)
)
# 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
)
hidden_states = self.mlp(
hidden_states, forward_batch, should_allreduce_fusion, use_reduce_scatter
)
with get_forward().scoped(
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
else:
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,
)
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 (
LazyValue,
add_prefix,
@@ -471,30 +476,28 @@ def _apply_qwen3_next_mlp(
hidden_states, residual = layer.layer_communicator.prepare_mlp(
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
)
should_allreduce_fusion = (
fuse_mlp_allreduce = (
layer.layer_communicator.should_fuse_mlp_allreduce_with_next_layer(
forward_batch
)
)
if isinstance(layer.mlp, Qwen2MoeSparseMoeBlock):
hidden_states = layer.mlp(
hidden_states,
forward_batch=forward_batch,
use_reduce_scatter=use_reduce_scatter,
should_allreduce_fusion=should_allreduce_fusion,
)
else:
hidden_states = layer.mlp(
hidden_states,
should_allreduce_fusion=should_allreduce_fusion,
use_reduce_scatter=use_reduce_scatter,
)
with get_forward().scoped(
fuse_mlp_allreduce=fuse_mlp_allreduce,
mlp_reduce_scatter=mlp_reduce_scatter,
):
if isinstance(layer.mlp, Qwen2MoeSparseMoeBlock):
hidden_states = layer.mlp(
hidden_states,
forward_batch=forward_batch,
)
else:
hidden_states = layer.mlp(hidden_states)
if should_allreduce_fusion:
if fuse_mlp_allreduce:
hidden_states._sglang_needs_allreduce_fusion = True
else:
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 (
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 (
BumpAllocator,
add_prefix,
@@ -205,16 +210,12 @@ class SarvamMoEMLP(nn.Module):
self,
x,
forward_batch: ForwardBatch = None,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
):
if x.shape[0] == 0:
return x
gate_up, _ = self.gate_up_proj(x)
x = self.act_fn(gate_up)
x, _ = self.down_proj(
x, skip_all_reduce=should_allreduce_fusion or use_reduce_scatter
)
x, _ = self.down_proj(x)
return x
@@ -315,8 +316,6 @@ class SarvamMoESparseMoeBlock(nn.Module):
self,
hidden_states: torch.Tensor,
forward_batch: Optional[ForwardBatch] = None,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
gemm_output_zero_allocator: Optional[BumpAllocator] = None,
) -> torch.Tensor:
del gemm_output_zero_allocator
@@ -327,13 +326,9 @@ class SarvamMoESparseMoeBlock(nn.Module):
and hidden_states.shape[0] > 0
and get_is_capture_mode()
):
return self.forward_normal_dual_stream(
hidden_states, should_allreduce_fusion, use_reduce_scatter
)
return self.forward_normal_dual_stream(hidden_states)
else:
return self.forward_normal(
hidden_states, should_allreduce_fusion, use_reduce_scatter
)
return self.forward_normal(hidden_states)
def get_moe_weights(self):
return [
@@ -359,8 +354,6 @@ class SarvamMoESparseMoeBlock(nn.Module):
def forward_normal_dual_stream(
self,
hidden_states: torch.Tensor,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
) -> torch.Tensor:
num_tokens, hidden_dim = hidden_states.shape
current_stream = torch.cuda.current_stream()
@@ -374,8 +367,6 @@ class SarvamMoESparseMoeBlock(nn.Module):
final_hidden_states = final_hidden_states + shared_out
if self.tp_size > 1 and not should_skip_post_experts_all_reduce(
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)
return final_hidden_states.view(num_tokens, hidden_dim)
@@ -383,8 +374,6 @@ class SarvamMoESparseMoeBlock(nn.Module):
def forward_normal(
self,
hidden_states: torch.Tensor,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
) -> torch.Tensor:
if hidden_states.shape[0] == 0:
return hidden_states
@@ -416,8 +405,6 @@ class SarvamMoESparseMoeBlock(nn.Module):
if self.tp_size > 1 and not should_skip_post_experts_all_reduce(
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)
@@ -1111,25 +1098,27 @@ class SarvamMoEMLADecoderLayer(nn.Module):
hidden_states, residual = self.layer_communicator.prepare_mlp(
hidden_states, residual, forward_batch
)
should_allreduce_fusion = (
fuse_mlp_allreduce = (
self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer(
forward_batch
)
)
use_reduce_scatter = self.layer_communicator.should_use_reduce_scatter(
mlp_reduce_scatter = self.layer_communicator.should_use_reduce_scatter(
forward_batch
)
hidden_states = self.mlp(
hidden_states, forward_batch, should_allreduce_fusion, use_reduce_scatter
)
with get_forward().scoped(
fuse_mlp_allreduce=fuse_mlp_allreduce,
mlp_reduce_scatter=mlp_reduce_scatter,
):
hidden_states = self.mlp(hidden_states, forward_batch)
if (
not self.is_layer_sparse
and self.attn_tp_size > 1
and not use_reduce_scatter
and not should_allreduce_fusion
and not mlp_reduce_scatter
and not fuse_mlp_allreduce
):
hidden_states = tensor_model_parallel_all_reduce(hidden_states)
if should_allreduce_fusion:
if fuse_mlp_allreduce:
hidden_states._sglang_needs_allreduce_fusion = True
else:
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,
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
logger = logging.getLogger(__name__)
@@ -74,12 +79,10 @@ class SDARMLP(nn.Module):
)
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)
hidden_states = self.act_fn(gate_up)
hidden_states, _ = self.down_proj(
hidden_states, skip_all_reduce=use_reduce_scatter
)
hidden_states, _ = self.down_proj(hidden_states)
return hidden_states
@@ -333,10 +336,11 @@ class SDARBlock(nn.Module):
forward_batch,
)
use_reduce_scatter = self.layer_communicator.should_use_reduce_scatter(
mlp_reduce_scatter = self.layer_communicator.should_use_reduce_scatter(
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, 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,
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
logger = logging.getLogger(__name__)
@@ -126,18 +131,12 @@ class SDARMoeSparseMoeBlock(nn.Module):
self,
hidden_states: torch.Tensor,
forward_batch: Optional[ForwardBatch] = None,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
) -> torch.Tensor:
if (
not get_moe_a2a_backend().is_deepep()
and not get_moe_a2a_backend().is_ascend_fuseep()
):
return self.forward_normal(
hidden_states,
should_allreduce_fusion=should_allreduce_fusion,
use_reduce_scatter=use_reduce_scatter,
)
return self.forward_normal(hidden_states)
else:
assert forward_batch is not None, "deepep/fuseep MoE needs forward_batch"
return self.forward_deepep(hidden_states, forward_batch)
@@ -145,8 +144,6 @@ class SDARMoeSparseMoeBlock(nn.Module):
def forward_normal(
self,
hidden_states: torch.Tensor,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
) -> torch.Tensor:
num_tokens, hidden_dim = hidden_states.shape
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(
is_tp_path=True,
use_reduce_scatter=use_reduce_scatter,
should_allreduce_fusion=should_allreduce_fusion,
):
out = tensor_model_parallel_all_reduce(out)
@@ -408,23 +403,25 @@ class SDARMoeBlock(nn.Module):
hidden_states, residual, forward_batch
)
should_allreduce_fusion = (
fuse_mlp_allreduce = (
self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer(
forward_batch
)
)
use_reduce_scatter = self.layer_communicator.should_use_reduce_scatter(
mlp_reduce_scatter = self.layer_communicator.should_use_reduce_scatter(
forward_batch
)
hidden_states = self.mlp(
hidden_states,
forward_batch=forward_batch,
should_allreduce_fusion=should_allreduce_fusion,
use_reduce_scatter=use_reduce_scatter,
)
with get_forward().scoped(
fuse_mlp_allreduce=fuse_mlp_allreduce,
mlp_reduce_scatter=mlp_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
else:
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_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
Step3p5Config = None
@@ -179,17 +184,13 @@ class Step3p5MoEMLP(nn.Module):
self,
hidden_states: torch.Tensor,
forward_batch: Optional[ForwardBatch] = None,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
) -> torch.Tensor:
if (
not get_moe_a2a_backend().is_deepep()
and not get_moe_a2a_backend().is_ascend_fuseep()
):
return self.forward_normal(
hidden_states, should_allreduce_fusion, use_reduce_scatter
)
return self.forward_normal(hidden_states)
else:
return self.forward_deepep(hidden_states, forward_batch)
@@ -206,8 +207,6 @@ class Step3p5MoEMLP(nn.Module):
def forward_normal(
self,
hidden_states: torch.Tensor,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
) -> torch.Tensor:
num_tokens, hidden_dim = hidden_states.shape
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)
if self.tp_size > 1 and not should_skip_post_experts_all_reduce(
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)
@@ -617,12 +614,12 @@ class Step3p5DecoderLayer(nn.Module):
forward_batch,
)
should_allreduce_fusion = (
fuse_mlp_allreduce = (
self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer(
forward_batch
)
)
use_reduce_scatter = self.layer_communicator.should_use_reduce_scatter(
mlp_reduce_scatter = self.layer_communicator.should_use_reduce_scatter(
forward_batch
)
@@ -630,24 +627,24 @@ class Step3p5DecoderLayer(nn.Module):
# Both share_expert and MoE return unreduced (TP-partial) outputs.
# Combine them first, then do a single all-reduce — saving one
# full-TP all-reduce per layer.
# Force fuse_mlp_allreduce=True so MoE skips its internal AR.
share_output = self.share_expert(hidden_states)
moe_output = self.moe(
hidden_states,
forward_batch,
should_allreduce_fusion=True,
use_reduce_scatter=use_reduce_scatter,
)
with get_forward().scoped(
fuse_mlp_allreduce=True,
mlp_reduce_scatter=mlp_reduce_scatter,
):
moe_output = self.moe(hidden_states, forward_batch)
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)
else:
hidden_states = self.mlp(hidden_states)
# Dense MLP uses reduce_results=True, so the output is already
# all-reduced. Do NOT set the fusion flag — otherwise the next
# 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
else:
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;
# graph runners force False around capture.
"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,
# EP dispatch, DP gather/scatter): plain-slot backed. Before moving a
# flag out of this set, prove no read/write site sits under
# torch.compile.
# EP dispatch, DP gather/scatter, MLP/MoE skip-AR): plain-slot backed.
# Before moving a flag out of this set, prove no read/write site sits
# under torch.compile.
_GRAPH_VISIBLE = frozenset(
{
"attn_input_scattered",
"attn_inputs",
"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
)
append: "dumper.dump('pre_mlp_residual', hidden_states, dims='t h # tp:replicated')"
- match: |
hidden_states = self.mlp(
hidden_states, forward_batch, should_allreduce_fusion, use_reduce_scatter
)
- match: "hidden_states = self.mlp(hidden_states, forward_batch)"
append: "dumper.dump('mlp_output', hidden_states, dims='t h[moe_tp:partial] # tp:replicated')"
# --- attention internals ---
@@ -150,10 +147,7 @@ patches:
hidden_states, residual, forward_batch
)
append: "dumper.dump('pre_mlp_residual', hidden_states, dims='t h # tp:replicated')"
- match: |
hidden_states = self.mlp(
hidden_states, forward_batch, should_allreduce_fusion, use_reduce_scatter
)
- match: "hidden_states = self.mlp(hidden_states, forward_batch)"
append: "dumper.dump('mlp_output', hidden_states, dims='t h[moe_tp:partial] # tp:replicated')"
# --- attention internals ---
test/registered/unit/test_runtime_context.py
+45
View File
@@ -696,6 +696,12 @@ class TestForwardFlags(_IsolatedServerArgs):
x = x + 1
if fwd.is_extend_in_batch:
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
self.assertEqual(probe(torch.zeros(())).item(), 0)
@@ -704,6 +710,13 @@ class TestForwardFlags(_IsolatedServerArgs):
get_forward().set("is_extend_in_batch", True)
self.assertEqual(probe(torch.zeros(())).item(), 2)
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):
# 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.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):
"""Publish installs the resolved server_args and seeds the capture tier."""