[Refactor] Declare every model's stages from a shared function and drop the meta-device neighbour build

docs/docs/developer_guide/layer_boundary.mdx

@@ -122,7 +122,7 @@ with layer_stack():
 
 Every stage binds when the stack closes, once its producer and its consumer are known: the stage appended before it and the one appended after it. The last stage ends the model's layer stack and is marked terminal, unless a later layer declares a stage. A boundary's `declaration` is usable at once; when the stack closes, it is replaced by the declaration with its place in the stack (its `previous` and `terminal`), and the plan is bound. An error raised while binding carries a note naming the `append_stages()` call it came from. A terminal stage leaves no work for a following layer, except the final-norm finalize handoff described in [Enter and leave the layer stack](#enter-and-leave-the-layer-stack). A NextN/MTP draft module is its own stack: its layer enters the stack and ends it, even when it reuses the target's decoder class.
 
-Pipeline parallelism splits one layer stack across ranks, so the first and last local stages have neighbours that another rank builds. The model gives `make_layers()` (or `make_pp_layers()`) its shared declaration function, `stage_facts`: given a global layer index, it returns the stages that layer declares, in order (an empty tuple for a layer that declares none), from the config, the layer index, construction options known before the layer is built, and the fixed parallel topology, never from built modules or provider state. The layer declares its own stages with the same function, so both ranks derive a boundary from one source; under pipeline parallelism `make_layers()` also checks each layer it builds against the function. When the stack closes, it asks the function for the stages next to it, nearest layer first, passing over a layer that declares none, and keeps only what binding reads of them (`facts_of()`: the capabilities a read, update and output transform declare, not the objects). A model that does not give the function yet is on a transitional list: its neighbours are built again with the model's own layer function inside `building_neighbour_layer()` (on the meta device), which must then skip the host and device resources only a running layer needs when `is_building_neighbour_layer()` is true and must not depend on which pipeline rank builds it. A model off that list must give the function to run under pipeline parallelism.
+Pipeline parallelism splits one layer stack across ranks, so the first and last local stages have neighbours that another rank builds. The model gives `make_layers()` (or `make_pp_layers()`) its shared declaration function, `stage_facts`: given a global layer index, it returns the stages that layer declares, in order (an empty tuple for a layer that declares none), from the config, the layer index, construction options known before the layer is built, and the fixed parallel topology, never from built modules or provider state. The layer declares its own stages with the same function, so both ranks derive a boundary from one source; under pipeline parallelism `make_layers()` also checks each layer it builds against the function. When the stack closes, it asks the function for the stages next to it, nearest layer first, passing over a layer that declares none, and keeps only what binding reads of them (`facts_of()`: the capabilities a read, update and output transform declare, not the objects). A model whose layers declare stages must give the function to run under pipeline parallelism; no rank builds a layer another rank holds.
 
 Each item may include a third element containing constructor options: an attention stage accepts `qkv_latent_func` (for example `{"qkv_latent_func": self.self_attn.prepare_qkv_latent}`) and `fusions`; an FFN stage accepts only `fusions` (see [Fusion providers](#fusion-providers-and-output-decisions)). An option a stage does not take is rejected when appended. Pass declarations fresh from `declare_attn()` / `declare_ffn()`: one taken from a bound boundary already has its place in a stack, and is rejected.
 

python/sglang/srt/layers/layer_boundary/factories.py

@@ -7,7 +7,7 @@
 import sys
 from dataclasses import dataclass, fields, is_dataclass, replace
 from types import MappingProxyType
-from typing import Callable, Mapping, NamedTuple, Optional, Tuple
+from typing import Callable, Mapping, NamedTuple, Optional
 
 from sglang.srt.environ import envs
 from sglang.srt.layers import layernorm_sp
@@ -302,11 +302,10 @@ class StageDeclaration:
             contiguous slice of the rows, it returns them all, in rank order,
             or None to leave the gather to the boundary. Called on every batch,
             so it must be CUDA-graph safe. Across a pipeline boundary the
-            producer's rank takes it from the neighbouring stage's declaration
-            (the model's shared declaration function, or, for a model on the
-            transitional list, the layer built again on the meta device), so it
-            may use only the communication state of the rank it runs on, not
-            the declaring layer's weights or buffers.
+            producer's rank takes it from the model's shared declaration
+            function, without the declaring layer, so it may use only the
+            communication state of the rank it runs on, not that layer's
+            weights or buffers.
         previous: The stage whose output this stage consumes, as the stack
             records it: only the facts binding reads of it (see facts_of),
             on this rank or another.
@@ -849,12 +848,13 @@ def layer_stack(*, previous_layers=(), next_layers=(), final_read=None):
     Args:
         previous_layers: For the layers before this stack that another
             pipeline rank holds, nearest first, callables that return the
-            stages the layer declares as DeclaredStages (from the model's
-            shared declaration function), or else build the layer so that it
-            appends them. Called only if this stack appended stages, after its
-            own layers are built, until one of them declares a stage: its last
-            stage is the producer of this stack's first. Only what binding
-            reads of that stage is kept (see facts_of).
+            stages the layer declares, in order, without building it (the
+            model's shared declaration function at that layer's index; an
+            empty sequence for a layer that declares none). Called only if
+            this stack appended stages, after its own layers are built, until
+            one of them declares a stage: its last stage is the producer of
+            this stack's first. Only what binding reads of that stage is kept
+            (see facts_of).
         next_layers: Likewise for the layers after this stack, whose first
             declared stage is the consumer of this stack's last.
         final_read: The model's final read of the stack's output (a
@@ -881,13 +881,6 @@ def layer_stack(*, previous_layers=(), next_layers=(), final_read=None):
         _stack = outer
 
 
-class DeclaredStages(NamedTuple):
-    """The stages a neighbouring layer declares, from the model's shared
-    declaration function, which stand in for building the layer."""
-
-    stages: Tuple[StageDeclaration, ...]
-
-
 def check_declared_stages(stack, appended: int, expected, where: str) -> None:
     """Check that what a layer appended to ``stack`` since it held
     ``appended`` appends is what its shared declaration function says it
@@ -947,26 +940,9 @@ def _differing_fields(value, given, prefix=""):
 
 def _neighbour_stage(layers, *, last):
     """The stage a neighbouring layer declares next to this stack: the last
-    one of the nearest layer before it, or the first of the nearest after.
-    Branches are side paths, so they never stand next to the stack."""
-    global _stack
+    one of the nearest layer before it, or the first of the nearest after."""
     for layer in layers:
-        outer = _stack
-        _stack = _LayerStack()
-        try:
-            result = layer()
-            declared = (
-                result.stages
-                if isinstance(result, DeclaredStages)
-                else [
-                    declaration
-                    for append in _stack.appends
-                    if append.prepared_from is None
-                    for declaration in append.declarations
-                ]
-            )
-        finally:
-            _stack = outer
+        declared = tuple(layer())
         if declared:
             return facts_of(declared[-1] if last else declared[0])
     return None

python/sglang/srt/layers/layer_boundary/facts.py

@@ -81,6 +81,18 @@ def of(cls, transform) -> Optional["TransformFacts"]:
         return cls(before_reduce_scatter=transform.before_reduce_scatter)
 
 
+def residual_facts(ops):
+    """A layer's reads and updates (a LayerResidualOps) reduced to the
+    capabilities they declare, in the same shape: what its stages declare
+    without the modules and state the reads and updates hold."""
+    return type(ops)(
+        attn_readout=ReadFacts.of(ops.attn_readout),
+        attn_update=UpdateFacts.of(ops.attn_update),
+        ffn_readout=ReadFacts.of(ops.ffn_readout),
+        ffn_update=UpdateFacts.of(ops.ffn_update),
+    )
+
+
 def facts_of(declaration):
     """``declaration`` with its read, update and output transform reduced to
     the capabilities they declare, and without the declarations it was

python/sglang/srt/layers/layer_boundary/residual/attn_bank.py

@@ -21,6 +21,7 @@
 from sglang.srt.layers.attn_residual import AttnResidual
 from sglang.srt.layers.layer_boundary import ops
 from sglang.srt.layers.layer_boundary.contracts import ReadoutFusion
+from sglang.srt.layers.layer_boundary.facts import residual_facts
 from sglang.srt.layers.layer_boundary.layout import SumGroup
 from sglang.srt.layers.layer_boundary.residual import LayerResidualOps
 from sglang.srt.layers.layer_boundary.residual.add_norm import PLAIN_ADD
@@ -187,6 +188,22 @@ def residual_ops(self) -> LayerResidualOps:
             ffn_update=PLAIN_ADD,
         )
 
+    @classmethod
+    def facts(cls, *, reads_slices: bool = False) -> LayerResidualOps:
+        """What residual_ops() declares (see facts_of) for a layer whose bank
+        stays on each rank's attention-TP slice when ``reads_slices``. It does
+        not depend on the bank or the layer's modules, so the stages of a
+        layer that is not built declare it too."""
+        state = cls(
+            bank=None,
+            attn_score_proj=None,
+            attn_score_norm=None,
+            ffn_score_proj=None,
+            ffn_score_norm=None,
+            reads_slices=reads_slices,
+        )
+        return residual_facts(state.residual_ops())
+
 
 class _BankReadout:
     """A read that mixes the banked snapshots into the residual's norm, so it

python/sglang/srt/layers/layer_boundary/residual/gated.py

@@ -18,6 +18,7 @@
 
 import torch
 
+from sglang.srt.layers.layer_boundary.facts import residual_facts
 from sglang.srt.layers.layer_boundary.residual import LayerResidualOps
 from sglang.srt.runtime_context import get_parallel
 
@@ -111,6 +112,22 @@ def residual_ops(self) -> LayerResidualOps:
             ffn_update=_FfnUpdate(self),
         )
 
+    @classmethod
+    def facts(cls, *, attn_reads_every_row: bool = False) -> LayerResidualOps:
+        """What residual_ops() declares (see facts_of) for a layer whose
+        attention read needs every row when ``attn_reads_every_row``. It does
+        not depend on the layer's modules, so the stages of a layer that is
+        not built declare it too."""
+        state = cls(
+            expand=None,
+            attn_mix=None,
+            ffn_mix=None,
+            attn_combine=None,
+            ffn_combine=None,
+            attn_reads_every_row=attn_reads_every_row,
+        )
+        return residual_facts(state.residual_ops())
+
 
 class _AttnReadout:
     """The normalized, gated mix of the streams, from streams that already hold

python/sglang/srt/layers/layer_boundary/residual/mhc.py

@@ -19,6 +19,7 @@
 import torch
 
 from sglang.kernels.ops.layernorm.mhc import hc_contract, hc_expand
+from sglang.srt.layers.layer_boundary.facts import ReadFacts, UpdateFacts
 from sglang.srt.layers.layer_boundary.residual import LayerResidualOps
 from sglang.srt.runtime_context import get_parallel
 
@@ -119,6 +120,19 @@ def residual_ops(self) -> LayerResidualOps:
             ffn_update=_FfnUpdate(self),
         )
 
+    @staticmethod
+    def facts() -> LayerResidualOps:
+        """What residual_ops() declares (see facts_of): the class constants
+        of its reads and updates, the same for every layer whatever
+        parameters it holds, so the stages of a layer that is not built
+        declare them too."""
+        return LayerResidualOps(
+            attn_readout=ReadFacts.of(_AttnReadout),
+            attn_update=UpdateFacts.of(_AttnUpdate),
+            ffn_readout=ReadFacts.of(_FfnReadout),
+            ffn_update=UpdateFacts.of(_FfnUpdate),
+        )
+
 
 class _AttnReadout:
     """hc_pre and the input norm, from streams that already hold the previous

python/sglang/srt/layers/moe/fused_moe_triton/layer.py

@@ -87,7 +87,6 @@
     is_npu,
     round_up,
 )
-from sglang.srt.utils.common import is_building_neighbour_layer
 from sglang.srt.utils.custom_op import register_custom_op
 
 _is_hip = is_hip()
@@ -164,8 +163,6 @@ def create_moe_dispatcher(
     moe_runner_config: MoeRunnerConfig,
     quant_method: FusedMoEMethodBase,
 ) -> BaseDispatcher:
-    if is_building_neighbour_layer():
-        return StandardDispatcher(moe_runner_config)
     a2a_backend = get_moe_a2a_backend()
     if a2a_backend.is_none() and is_npu():
         return AscendTPDispatcher(moe_runner_config)

python/sglang/srt/layers/moe/kt_ep_wrapper.py

@@ -20,7 +20,6 @@
     get_schedule,
 )
 from sglang.srt.utils import get_compiler_backend
-from sglang.srt.utils.common import is_building_neighbour_layer
 
 if TYPE_CHECKING:
     from sglang.srt.layers.moe import MoeRunnerConfig
@@ -216,7 +215,7 @@ def create_weights(
 
         # 2. Initialize KT wrapper for CPU experts
         # CPU experts: num_gpu_experts to num_experts-1
-        if self.tp_rank == 0 and not is_building_neighbour_layer():
+        if self.tp_rank == 0:
             self.wrapper = KTMoEWrapper(
                 layer_idx=self.kt_config.layer_idx,
                 num_experts=num_experts,

python/sglang/srt/models/apertus.py

@@ -266,14 +266,18 @@ def __init__(
         self.feedforward_layernorm = RMSNorm(
             config.hidden_size, eps=config.rms_norm_eps
         )
+        attn, ffn = self.stage_facts(layer_id)
         self.attn_boundary, self.ffn_boundary = append_stages(
-            (declare_attn(), self.attention_layernorm),
-            (
-                declare_ffn(sparse=False, next_layer_sparse=False),
-                self.feedforward_layernorm,
-            ),
+            (attn, self.attention_layernorm),
+            (ffn, self.feedforward_layernorm),
         )
 
+    @staticmethod
+    def stage_facts(layer_id: int):
+        """The stages every Apertus layer declares, an attention and a dense
+        FFN: the model's shared declaration function (see make_layers)."""
+        return (declare_attn(), declare_ffn(sparse=False, next_layer_sparse=False))
+
     def forward(
         self,
         positions: torch.Tensor,
@@ -329,6 +333,7 @@ def __init__(
                 config=config, quant_config=quant_config, layer_id=idx, prefix=prefix
             ),
             prefix="model.layers",
+            stage_facts=ApertusDecoderLayer.stage_facts,
         )
 
         if self.pp_group.is_last_rank:

python/sglang/srt/models/arcee.py

@@ -242,14 +242,18 @@ def __init__(
         self.post_attention_layernorm = RMSNorm(
             config.hidden_size, eps=config.rms_norm_eps
         )
+        attn, ffn = self.stage_facts(layer_id)
         self.attn_boundary, self.ffn_boundary = append_stages(
-            (declare_attn(), self.input_layernorm),
-            (
-                declare_ffn(sparse=False, next_layer_sparse=False),
-                self.post_attention_layernorm,
-            ),
+            (attn, self.input_layernorm),
+            (ffn, self.post_attention_layernorm),
         )
 
+    @staticmethod
+    def stage_facts(layer_id: int):
+        """The stages every Arcee layer declares, an attention and a dense FFN:
+        the model's shared declaration function (see make_layers)."""
+        return (declare_attn(), declare_ffn(sparse=False, next_layer_sparse=False))
+
     def forward(
         self,
         positions: torch.Tensor,
@@ -303,6 +307,7 @@ def __init__(
                 config=config, quant_config=quant_config, layer_id=idx, prefix=prefix
             ),
             prefix="model.layers",
+            stage_facts=ArceeDecoderLayer.stage_facts,
         )
 
         if self.pp_group.is_last_rank:

python/sglang/srt/models/bailing_moe.py

@@ -655,12 +655,8 @@ def __init__(
         )
         self.layer_id = layer_id
 
-        self.is_layer_sparse = self._is_layer_sparse(
-            config, layer_id=layer_id, is_nextn=is_nextn
-        )
-        is_next_layer_sparse = self._is_layer_sparse(
-            config, layer_id=layer_id + 1, is_nextn=False
-        )
+        attn, ffn = self.stage_facts(config, layer_id, is_nextn=is_nextn)
+        self.is_layer_sparse = ffn.sparse
 
         if self.is_layer_sparse:
             self.mlp = BailingMoESparseMoeBlock(
@@ -684,18 +680,31 @@ def __init__(
         self.post_attention_layernorm = RMSNorm(hidden_size, eps=config.rms_norm_eps)
 
         self.attn_boundary, self.ffn_boundary = append_stages(
-            (declare_attn(), self.input_layernorm),
-            (
-                declare_ffn(
-                    sparse=self.is_layer_sparse,
-                    next_layer_sparse=is_next_layer_sparse,
+            (attn, self.input_layernorm),
+            (ffn, self.post_attention_layernorm),
+        )
+
+    @classmethod
+    def stage_facts(
+        cls, config: PretrainedConfig, layer_id: int, is_nextn: bool = False
+    ):
+        """The stages the layer at ``layer_id`` declares, from the config: the
+        model's shared declaration function (see make_layers)."""
+        return (
+            declare_attn(),
+            declare_ffn(
+                sparse=cls._is_layer_sparse(
+                    config, layer_id=layer_id, is_nextn=is_nextn
+                ),
+                next_layer_sparse=cls._is_layer_sparse(
+                    config, layer_id=layer_id + 1, is_nextn=False
                 ),
-                self.post_attention_layernorm,
             ),
         )
 
+    @staticmethod
     def _is_layer_sparse(
-        self, config: PretrainedConfig, layer_id: int, is_nextn: bool
+        config: PretrainedConfig, layer_id: int, is_nextn: bool
     ) -> bool:
         return is_nextn or (
             config.num_experts is not None and layer_id >= config.first_k_dense_replace
@@ -774,6 +783,7 @@ def __init__(
                 alt_stream=alt_stream,
             ),
             prefix=add_prefix("layers", prefix),
+            stage_facts=lambda idx: BailingMoEBlock.stage_facts(config, idx),
         )
         if self.pp_group.is_last_rank:
             self.norm = RMSNorm(self.embed_dim, eps=config.rms_norm_eps)