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

python/sglang/srt/runtime_context.py

@@ -522,6 +522,15 @@ class DpFlags(_FlagGroupBase):
     buffer_device: Any = None
 
 
+@dataclasses.dataclass
+class SpFlags(_FlagGroupBase):
+    """LayerNorm sequence-parallelism flags, materialized by
+    ``initialize_layernorm_sp`` (after distributed setup; reads the model
+    config). See ``layers.layernorm_sp``."""
+
+    enabled: bool = False
+
+
 @dataclasses.dataclass
 class Flags(_FlagGroupBase):
     """Root of the runtime-flags tier.
@@ -529,12 +538,13 @@ class Flags(_FlagGroupBase):
     Resolved configuration lives in the config bags below (projected from the
     declarations at publish) — this tier only carries genuine runtime
     state whose value is not a function of the configuration alone, grouped
-    by lifecycle (``capture``) or subsystem (``moe`` / ``dp``).
+    by lifecycle (``capture``) or subsystem (``moe`` / ``dp`` / ``sp``).
     """
 
     capture: CaptureFlags = dataclasses.field(default_factory=CaptureFlags)
     moe: MoeFlags = dataclasses.field(default_factory=MoeFlags)
     dp: DpFlags = dataclasses.field(default_factory=DpFlags)
+    sp: SpFlags = dataclasses.field(default_factory=SpFlags)
 
 
 @dataclasses.dataclass
@@ -606,6 +616,8 @@ class ForwardFlags:
         "fuse_mlp_allreduce": False,
         "mlp_reduce_scatter": False,
         "flashinfer_trtllm_bypass": False,
+        # LayerNorm sequence parallelism region; see layers/layernorm_sp.py.
+        "sp_active": False,
     }
 
     # Read/written inside compiled graphs (vocab embedding, communicator,
@@ -620,6 +632,7 @@ class ForwardFlags:
             "fuse_mlp_allreduce",
             "mlp_reduce_scatter",
             "flashinfer_trtllm_bypass",
+            "sp_active",
         }
     )