[Feature] Megatron LayerNorm sequence parallelism (--enable-layernorm-sp)

python/sglang/srt/arg_groups/layernorm_sp_hook.py

@@ -0,0 +1,63 @@
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Optional
+
+if TYPE_CHECKING:
+    from sglang.srt.server_args import ServerArgs
+
+from sglang.srt.arg_groups.overrides import model_config_of, resolving_view
+
+logger = logging.getLogger(__name__)
+
+
+def handle_layernorm_sp(server_args: ServerArgs) -> None:
+    """Validate --enable-layernorm-sp against the resolved parallelism config.
+
+    Runs in the resolution pipeline rather than in the layers so a model that
+    never builds a LayerCommunicator rejects the flag instead of ignoring it.
+    """
+    cfg = resolving_view(server_args)
+    if not cfg.enable_layernorm_sp:
+        return
+    architectures = model_config_of(server_args).hf_config.architectures
+    validate_layernorm_sp(
+        architecture=architectures[0] if architectures else None,
+        tp_size=cfg.tp_size,
+        enable_dp_attention=cfg.enable_dp_attention,
+        speculative_algorithm=cfg.speculative_algorithm,
+    )
+
+
+def validate_layernorm_sp(
+    *,
+    architecture: Optional[str],
+    tp_size: int,
+    enable_dp_attention: bool,
+    speculative_algorithm: Optional[str],
+) -> None:
+    """Fail loud for unsupported / incompatible configs. Callers gate on the flag."""
+    from sglang.srt.layers.layernorm_sp import SP_SUPPORTED_ARCHITECTURES
+
+    if architecture not in SP_SUPPORTED_ARCHITECTURES:
+        raise ValueError(
+            "--enable-layernorm-sp is only supported for "
+            f"{sorted(SP_SUPPORTED_ARCHITECTURES)}; got {architecture}."
+        )
+    if tp_size <= 1:
+        raise ValueError(
+            "--enable-layernorm-sp requires tp_size > 1: there is no sequence to "
+            "shard across a single TP rank."
+        )
+    if enable_dp_attention:
+        raise ValueError(
+            "--enable-layernorm-sp is not compatible with --enable-dp-attention: "
+            "SP shards the sequence across the full TP group, which under DP "
+            "attention spans data-parallel groups holding different sequences."
+        )
+    if speculative_algorithm is not None:
+        raise ValueError(
+            "--enable-layernorm-sp is not compatible with speculative decoding "
+            "(EAGLE/EAGLE3): the captured aux hidden states would be "
+            "sequence-sharded."
+        )

python/sglang/srt/arg_groups/pipeline.py

@@ -318,6 +318,11 @@ def run_resolution_pipeline(server_args: Any) -> None:
 
     handle_speculative_decoding(server_args)
 
+    # After the speculative hook so speculative_algorithm is final.
+    from sglang.srt.arg_groups.layernorm_sp_hook import handle_layernorm_sp
+
+    handle_layernorm_sp(server_args)
+
     # Validate the CuteDSL A2A token budget now that num_tokens_per_req is final.
     validate_cutedsl_a2a_token_budget(server_args)
 

python/sglang/srt/distributed/bootstrap.py

@@ -27,6 +27,7 @@
 from sglang.srt.distributed.parallel_state_wrapper import ParallelState
 from sglang.srt.environ import envs
 from sglang.srt.layers.dp_attention import initialize_dp_attention
+from sglang.srt.layers.layernorm_sp import initialize_layernorm_sp
 from sglang.srt.platforms import current_platform
 from sglang.srt.runtime_context import (
     get_disagg,
@@ -280,6 +281,10 @@ def _init_parallel_groups(
         server_args=server_args,
         model_config=model_config,
     )
+    initialize_layernorm_sp(
+        server_args=server_args,
+        model_config=model_config,
+    )
     if is_npu():
         register_sgl_tp_rank(gpu_id)
 

python/sglang/srt/layers/communicator.py

@@ -31,6 +31,7 @@
     use_symmetric_memory,
 )
 from sglang.srt.environ import envs
+from sglang.srt.layers import layernorm_sp
 from sglang.srt.layers.attention.dsa.utils import (
     dsa_use_prefill_cp,
     is_dsa_enable_prefill_cp,
@@ -72,7 +73,7 @@
     Phase,
     check_cuda_graph_backend,
 )
-from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
 from sglang.srt.runtime_context import (
     get_exec,
     get_forward,
@@ -462,6 +463,7 @@ def __init__(
         force_layernorm_before_dp_gather: bool = False,
         enable_fused_ar_quant: bool = False,
         fused_ar_quant_keep_bf16: bool = False,
+        _is_sp_variant: bool = False,
     ):
         self.layer_scatter_modes = layer_scatter_modes
         self.input_layernorm = input_layernorm
@@ -482,6 +484,30 @@ def __init__(
             get_spec().speculative_algorithm
         )
 
+        # Under LayerNorm SP the norm/residual run on the sequence shard with no
+        # collectives, so delegate to an all-SCATTERED sibling while the region is
+        # active. _is_sp_variant stops the sibling from building its own.
+        self._sp_variant: Optional[LayerCommunicator] = None
+        if not _is_sp_variant and layernorm_sp.layernorm_sp_enabled():
+            self._sp_variant = LayerCommunicator(
+                layer_scatter_modes=LayerScatterModes(
+                    layer_input_mode=ScatterMode.SCATTERED,
+                    attn_mode=ScatterMode.SCATTERED,
+                    mlp_mode=ScatterMode.SCATTERED,
+                    middle_residual_mode=ScatterMode.SCATTERED,
+                    layer_output_mode=ScatterMode.SCATTERED,
+                ),
+                input_layernorm=input_layernorm,
+                post_attention_layernorm=post_attention_layernorm,
+                allow_reduce_scatter=allow_reduce_scatter,
+                is_last_layer=is_last_layer,
+                qkv_latent_func=qkv_latent_func,
+                force_layernorm_before_dp_gather=force_layernorm_before_dp_gather,
+                enable_fused_ar_quant=enable_fused_ar_quant,
+                fused_ar_quant_keep_bf16=fused_ar_quant_keep_bf16,
+                _is_sp_variant=True,
+            )
+
     def _post_init_communicate(self):
         self._communicate_simple_fn = CommunicateSimpleFn.get_fn(
             input_mode=self.layer_scatter_modes.layer_input_mode,
@@ -571,6 +597,24 @@ def prepare_attn(
         quant_format: str = "",
         post_residual_addition: Optional[torch.Tensor] = None,
     ):
+        # residual is None marks the first decoder layer, where the SP region
+        # opens: re-evaluated per forward so a crash mid-loop cannot leak into
+        # the next one.
+        if self._sp_variant is not None:
+            if residual is None:
+                get_forward().set(
+                    "sp_active", forward_batch.forward_mode == ForwardMode.EXTEND
+                )
+                if get_forward().sp_active:
+                    hidden_states = layernorm_sp.sp_entry_scatter(hidden_states)
+            if get_forward().sp_active:
+                return self._sp_variant.prepare_attn(
+                    hidden_states,
+                    residual,
+                    forward_batch,
+                    quant_format,
+                    post_residual_addition,
+                )
         if get_attn_tp_context().input_scattered:
             hidden_states, residual = self._tp_reduce_scatter(
                 hidden_states,
@@ -767,6 +811,10 @@ def prepare_mlp(
         forward_batch: ForwardBatch,
         cache=None,
     ):
+        if self._sp_variant is not None and get_forward().sp_active:
+            return self._sp_variant.prepare_mlp(
+                hidden_states, residual, forward_batch, cache
+            )
         if cache is not None:
             self._context.cache = cache
 
@@ -784,6 +832,10 @@ def postprocess_layer(
         residual: torch.Tensor,
         forward_batch: ForwardBatch,
     ):
+        if self._sp_variant is not None and get_forward().sp_active:
+            return self._sp_variant.postprocess_layer(
+                hidden_states, residual, forward_batch
+            )
         return self._communicate_summable_tensor_pair_fn(
             hidden_states=hidden_states,
             residual=residual,
@@ -1007,6 +1059,22 @@ def get_fn(
         ):
             return CommunicateWithAllReduceAndLayerNormFn._simple
 
+        if (
+            hidden_states_input_mode == ScatterMode.SCATTERED
+            and residual_input_mode == ScatterMode.SCATTERED
+            and hidden_states_output_mode == ScatterMode.SCATTERED
+            and residual_output_mode == ScatterMode.SCATTERED
+        ):
+            # Megatron LayerNorm sequence parallelism (layers/layernorm_sp.py):
+            # activations stay sequence-sharded across the attn->mlp boundary, so
+            # there is nothing to gather or scatter here -- just the residual add
+            # plus LayerNorm on the local shard. The row-parallel o_proj already
+            # issued the reduce-scatter (g-bar) that the all-reduce would have
+            # done, and the g all-gather is fused into the next column-parallel
+            # linear. Distinct from the branch above because under pure TP
+            # attn_tp_size == tp_size > 1, so that gate does not fire.
+            return CommunicateWithAllReduceAndLayerNormFn._simple
+
         if (
             (hidden_states_input_mode == ScatterMode.TP_ATTN_FULL)
             and (