diff --git a/python/sglang/srt/layers/boundary_layout.py b/python/sglang/srt/layers/boundary_layout.py index b18c55a9979e..8b30f64fd105 100644 --- a/python/sglang/srt/layers/boundary_layout.py +++ b/python/sglang/srt/layers/boundary_layout.py @@ -87,8 +87,9 @@ class StageOutput(msgspec.Struct, frozen=True): class DecoderLayerSides(msgspec.Struct, frozen=True): - """The declarations both sides of a decoder layer's boundaries are chosen - from.""" + """What a decoder layer declares about its attention, its FFN and the rows + between them; ``decoder_layer_edges`` turns it into the layer's + boundaries.""" # The rows the layer's input and residual arrive on. input_rows: Layout @@ -107,6 +108,68 @@ class DecoderLayerSides(msgspec.Struct, frozen=True): residual_joins_attention_sum: bool = False +class EdgeDecl(msgspec.Struct, frozen=True): + """One boundary between two stages, as the layer that runs one side of it + sees it: what arrives from the producer, what the consumer needs, and the + rows the residual is on before and after the boundary.""" + + produced: StageOutput + need: StageInput + residual: Layout + residual_to: Layout + # Whether the residual is added into one rank's share of the produced sum + # before that sum completes, instead of after it. + residual_joins_sum: bool = False + + +class DecoderLayerEdges(msgspec.Struct, frozen=True): + """A decoder layer's three boundaries: from the rows the previous layer + handed on into the attention, from the attention into the FFN, and from + the FFN onto the rows the layer hands on. The next layer runs the rest of + the last one.""" + + into_attention: EdgeDecl + into_ffn: EdgeDecl + out_of_ffn: EdgeDecl + + +def decoder_layer_edges(sides: DecoderLayerSides) -> DecoderLayerEdges: + """The boundaries of a decoder layer, each with the declarations of its two + sides. What arrives at the layer is complete unless it owes a sum by + construction (``input_owes``); a sum the previous layer leaves for a batch + comes with the value.""" + owes = sides.input_owes is not None + arrived = StageOutput(sides.input_rows, group=sides.input_owes, always_leaves=owes) + # Completing what the input owes leaves it and the residual on each + # attention-TP rank's slice. + attention_rows = ( + Layout(sides.input_rows.sharded | {TokenAxis.ATTN_TP_SCATTER}) + if owes + else sides.input_rows + ) + return DecoderLayerEdges( + into_attention=EdgeDecl( + produced=arrived, + need=sides.attention, + residual=sides.input_rows, + residual_to=attention_rows, + ), + into_ffn=EdgeDecl( + produced=sides.attention_output, + need=sides.ffn, + residual=attention_rows, + residual_to=sides.ffn_residual_rows, + residual_joins_sum=sides.residual_joins_attention_sum, + ), + out_of_ffn=EdgeDecl( + produced=sides.ffn_output, + need=StageInput(sides.output_rows), + residual=sides.ffn_residual_rows, + residual_to=sides.output_rows, + ), + ) + + def sequence_parallel_layer_sides( *, axis_sizes: Mapping[TokenAxis, int] ) -> DecoderLayerSides: diff --git a/python/sglang/srt/layers/communicator.py b/python/sglang/srt/layers/communicator.py index ef824db1f4e2..2412046c16b5 100644 --- a/python/sglang/srt/layers/communicator.py +++ b/python/sglang/srt/layers/communicator.py @@ -39,11 +39,13 @@ from sglang.srt.layers.aux_hidden_states import AuxHiddenStateAccumulator from sglang.srt.layers.boundary_layout import ( DecoderLayerSides, + EdgeDecl, Layout, StageInput, StageOutput, SumGroup, TokenAxis, + decoder_layer_edges, decoder_layer_sides, input_scattered_layer_sides, scattered_residual_layer_sides, @@ -835,6 +837,8 @@ def __init__( self._context.force_layernorm_before_dp_gather = ( force_layernorm_before_dp_gather ) + # The fused kernels every batch's attention input tries first. + self._attn_input_fusions = self._select_attn_input_fusions() # The steps the layer's ordinary batches run. sides = self._declared_sides() self._declared = sides @@ -866,7 +870,6 @@ def __init__( if sides is not None and self._input_can_be_scattered() else None ) - self._attn_input_fusions = self._select_attn_input_fusions() self._speculative_algo = SpeculativeAlgorithm.from_string( get_spec().speculative_algorithm ) @@ -882,6 +885,8 @@ def __init__( _select_boundary_steps( sequence_parallel_layer_sides(axis_sizes=_token_axis_sizes()), residual_ops=residual_ops, + attention_fusions=self._attn_input_fusions, + enters_stack=self.layer_scatter_modes.is_first_layer, ) if layernorm_sp.layernorm_sp_enabled() else None @@ -1033,7 +1038,13 @@ def _select_boundaries_from_scatter_modes(self) -> "BoundarySteps": ffn_sum_is_movable=modes.mlp_mode not in (ScatterMode.MOE_FULL, ScatterMode.SCATTERED), fused=fused, - layer_input=_complete_scattered_input, + attention_prepare=partial( + _attention_input_step, + layer_input=_complete_scattered_input, + fusions=self._attn_input_fusions, + enters_stack=modes.is_first_layer, + residual_ops=ops, + ), ) def _steps_for_input_scattered(self, sides: DecoderLayerSides) -> "BoundarySteps": @@ -1059,14 +1070,24 @@ def _steps_for_input_scattered(self, sides: DecoderLayerSides) -> "BoundarySteps # DSA and hook-less attention take the slice gathered. handoff = _hand_scattered_input_to_attention return _select_boundary_steps( - scattered, residual_ops=self._residual_ops, attention_handoff=handoff + scattered, + residual_ops=self._residual_ops, + attention_handoff=handoff, + attention_fusions=self._attn_input_fusions, + enters_stack=self.layer_scatter_modes.is_first_layer, ) def _steps_from_declarations( self, sides: DecoderLayerSides, **kwargs ) -> "BoundarySteps": """The steps this layer runs for a set of declarations.""" - return _select_boundary_steps(sides, residual_ops=self._residual_ops, **kwargs) + return _select_boundary_steps( + sides, + residual_ops=self._residual_ops, + attention_fusions=self._attn_input_fusions, + enters_stack=self.layer_scatter_modes.is_first_layer, + **kwargs, + ) def _post_init_communicate(self) -> Tuple[Callable, Callable]: """The attention input move and the postprocess the scatter modes @@ -1172,62 +1193,24 @@ def prepare_attn( post_residual_addition: Optional[torch.Tensor] = None, ): self.publish_attn_lora_layout() - # The layer stack's first layer starts its residual from its input. - enters = residual is None and self.layer_scatter_modes.is_first_layer - # What the previous layer left to complete: an UnreducedOutput, or a - # producer's handoff that one of the fused entries consumes. - owed = None if isinstance(hidden_states, torch.Tensor) else hidden_states - if owed is not None and residual is None: - raise RuntimeError(f"{type(owed).__name__} requires residual input") - if ( - isinstance(owed, UnreducedOutput) - and owed.reduce_and_redistribute is not None - ): - # No fused kernel runs under attention DP: the reduce-scatter back to - # this rank's tokens comes first. - hidden_states, owed = reduce_output(owed), None - if owed is not None: - for fused in self._attn_input_fusions: - result = fused(owed, residual, forward_batch, post_residual_addition) - if result is not None: - return self._finish_prepare_attn( - hidden_states=result[0], - residual=result[1], - forward_batch=forward_batch, - ) - hidden_states = owed.partial # The SP region opens at the first layer, re-evaluated per forward so a - # crash mid-loop cannot leak into the next one. + # crash mid-loop cannot leak into the next one. It sets what the batch's + # steps are chosen from, and the first layer's input owes nothing. if self._sp_steps is not None and self.layer_scatter_modes.is_first_layer: get_forward().set( "sp_active", layernorm_sp.runs_sp(forward_batch.forward_mode) ) if get_forward().sp_active: hidden_states = layernorm_sp.sp_entry_scatter(hidden_states) - layer_input = self._batch_steps(forward_batch).layer_input - if layer_input is not None: - hidden_states, residual = layer_input( - hidden_states, residual, self._context - ) - ops = self._residual_ops - if enters: - hidden_states, residual = ops.enter(hidden_states), None - if owed is not None and hidden_states.shape[0] != 0: - hidden_states = reduce_output(owed) - if residual is None: - # The previous layer already wrote its output into the residual. - hidden_states, residual = ops.read_attention_input( - hidden_states, self.input_layernorm, quant_format - ) - else: - hidden_states, residual = ops.update_and_read_attention_input( - hidden_states, - residual, - self.input_layernorm, - quant_format, - post_residual_addition, - ) - + hidden_states, residual = self._batch_steps(forward_batch).attention_prepare( + hidden_states, + residual, + forward_batch, + self.input_layernorm, + self._context, + quant_format=quant_format, + post_residual_addition=post_residual_addition, + ) return self._finish_prepare_attn( hidden_states=hidden_states, residual=residual, @@ -1398,7 +1381,7 @@ def _mlp_input_reduce_output_and_update_and_read_residual( def prepare_mlp( self, - hidden_states: torch.Tensor, + hidden_states: Union[torch.Tensor, UnreducedOutput], residual: torch.Tensor, forward_batch: ForwardBatch, cache=None, @@ -2294,6 +2277,10 @@ class BoundarySteps(msgspec.Struct, frozen=True): from the attention output to the FFN input, and the FFN output on to the rows the layer hands on.""" + # The half into the attention: completes what the input owes, writes the + # previous output into the residual and reads the attention input + # (_attention_input_step); attention_input then moves it. + attention_prepare: Callable attention_input: Callable ffn_input: Callable # The rows ffn_input hands the FFN. @@ -2309,9 +2296,6 @@ class BoundarySteps(msgspec.Struct, frozen=True): ffn_output_move_completes_sum: bool = False # The fused kernels ffn_input tries first. fused: Tuple["FusedMlpInput", ...] = () - # Completes what the layer's input owes before the input norm; None when it - # owes nothing. - layer_input: Optional[Callable] = None # Hands the attention its input once attention_input has moved it: # (hidden_states, forward_batch, qkv_latent_func) -> hidden_states. attention_handoff: Callable = _hand_qkv_hook_its_input @@ -2321,62 +2305,213 @@ def returns_over_dp(self) -> bool: return self.ffn_output_move is None -def _select_boundary_steps( - sides: DecoderLayerSides, +def _attention_input_step( + hidden_states: Union[torch.Tensor, "UnreducedOutput"], + residual: Optional[torch.Tensor], + forward_batch: ForwardBatch, + norm: torch.nn.Module, + context: "CommunicateContext", *, - fusions: Tuple["FusedMlpInput", ...] = (), + quant_format: str, + post_residual_addition: Optional[torch.Tensor], + layer_input: Optional[Callable], + fusions: Tuple[Callable, ...], + enters_stack: bool, + residual_ops: ResidualOps, +) -> Tuple[torch.Tensor, torch.Tensor]: + """A boundary's half into an attention: complete what the previous layer + left (a value that owes a sum, or a producer's handoff one of ``fusions`` + consumes with the add and norm), what the input owes by construction + (``layer_input``), then write the previous output into the residual and read + the attention input with ``norm``. The layer stack's first layer + (``enters_stack``) starts its residual from its input.""" + enters = residual is None and enters_stack + owed = None if isinstance(hidden_states, torch.Tensor) else hidden_states + if owed is not None and residual is None: + raise RuntimeError(f"{type(owed).__name__} requires residual input") + if isinstance(owed, UnreducedOutput) and owed.reduce_and_redistribute is not None: + # No fused kernel runs under attention DP: the reduce-scatter back to + # this rank's tokens comes first. + hidden_states, owed = reduce_output(owed), None + if owed is not None: + for fused in fusions: + result = fused(owed, residual, forward_batch, post_residual_addition) + if result is not None: + return result + hidden_states = owed.partial + if layer_input is not None: + hidden_states, residual = layer_input(hidden_states, residual, context) + if enters: + hidden_states, residual = residual_ops.enter(hidden_states), None + if owed is not None and hidden_states.shape[0] != 0: + hidden_states = reduce_output(owed) + if residual is None: + # The previous layer already wrote its output into the residual. + return residual_ops.read_attention_input(hidden_states, norm, quant_format) + return residual_ops.update_and_read_attention_input( + hidden_states, residual, norm, quant_format, post_residual_addition + ) + + +class InputRead(Enum): + """How a boundary's consumer reads its input from the residual: with the + attention input norm (prepare_attn) or with the FFN input norm and its + fused kernels (prepare_mlp).""" + + ATTENTION = auto() + FFN = auto() + + +class Boundary(msgspec.Struct, frozen=True): + """The steps one layer runs at one boundary, chosen from both sides' + declarations. A layer runs the consumer's half of a boundary into one of + its stages, and the producer's half of the boundary after its last stage; + the neighbouring layer runs the other half of that one.""" + + edge: EdgeDecl + # The consumer's half: completing what the input owes, the add and the + # norm; into the FFN also the move onto the rows it needs. + prepare: Optional[Callable] = None + # Into an attention, the move onto its rows after prepare. + input_move: Optional[Callable] = None + # The fused kernels an FFN's prepare tries first. + fused: Tuple["FusedMlpInput", ...] = () + # The producer's half: the postprocess that moves the output onto the rows + # the layer hands on; None when it goes back over attention DP, whose step + # the FFN exit and postprocess choose per batch. + output_move: Optional[Callable] = None + # Whether output_move also completes the sum the producer leaves. + output_move_completes_sum: bool = False + + @property + def input_rows(self) -> Layout: + """The rows the consumer is handed: what it needs, still sharded over + the axes it gathers itself.""" + need = self.edge.need + return Layout( + need.layout.sharded + | (self.edge.produced.layout.sharded & need.gathers_itself) + ) + + +def make_boundary( + edge: EdgeDecl, + *, + reads: Optional[InputRead], + fusions: Tuple = (), force_layernorm_before_gather: bool = False, cp_moves: Optional[CpMoves] = None, residual_ops: ResidualOps = ADD_AND_NORM, - attention_handoff: Callable = _hand_qkv_hook_its_input, -) -> BoundarySteps: - """The steps a set of declarations chooses, around the layer's residual - operations; ``cp_moves`` for the ones that gather over attention CP.""" - returns_over_dp, ffn_output_move, completes_sum = _select_ffn_output_move( - sides.ffn_output, - residual=sides.ffn_residual_rows, - to=sides.output_rows, - cp_moves=cp_moves, - residual_ops=residual_ops, - ) - rows = sides.input_rows + enters_stack: bool = False, +) -> Boundary: + """The steps a layer runs at ``edge``, around the residual operations; + ``cp_moves`` for an edge that gathers over or returns across attention CP. + ``reads`` is how the consumer reads its input, or None when the consumer + runs in the next layer and ``edge.need`` is the rows this layer hands on: + then only the producer's half runs here. ``fusions`` are the fused kernels + the consumer tries first (FusedMlpInput into an FFN, the attention input's + entries into an attention); ``enters_stack`` for the edge into the layer + stack's first attention. The consumer's half reads only this edge's + declarations, never what the producer chose for a batch; a sum left for a + batch arrives with the value.""" + if reads is None: + if edge.need.layout != edge.residual_to: + raise NotImplementedError(f"{edge=}") + returns_over_dp, output_move, completes_sum = _select_ffn_output_move( + edge.produced, + residual=edge.residual, + to=edge.residual_to, + cp_moves=cp_moves, + residual_ops=residual_ops, + ) + return Boundary( + edge, + output_move=None if returns_over_dp else output_move, + output_move_completes_sum=completes_sum, + ) + if reads is InputRead.FFN: + input_step, fused = _select_ffn_input( + edge.produced, + residual=edge.residual, + residual_to=edge.residual_to, + need=edge.need, + force_layernorm_before_gather=force_layernorm_before_gather, + fusions=fusions, + residual_joins_sum=edge.residual_joins_sum, + cp_moves=cp_moves, + residual_ops=residual_ops, + ) + return Boundary(edge, prepare=input_step, fused=fused) layer_input = None - if sides.input_owes is not None: + if edge.produced.always_leaves: # A reduce-scatter completes the TP sum onto each rank's slice, which the # attention takes: its group is the TP group without attention DP or CP. if ( - sides.input_owes is not SumGroup.TP - or rows.sharded - or TokenAxis.ATTN_TP_SCATTER not in sides.attention.gathers_itself + edge.produced.group is not SumGroup.TP + or edge.residual.sharded + or TokenAxis.ATTN_TP_SCATTER not in edge.need.gathers_itself + or edge.residual_to.sharded != {TokenAxis.ATTN_TP_SCATTER} ): - raise NotImplementedError(f"{sides=}") + raise NotImplementedError(f"{edge=}") layer_input = tp_reduce_scatter - rows = Layout(frozenset({TokenAxis.ATTN_TP_SCATTER})) - ffn_input, fused = _select_ffn_input( - sides.attention_output, - residual=rows, - residual_to=sides.ffn_residual_rows, - need=sides.ffn, - force_layernorm_before_gather=force_layernorm_before_gather, + return Boundary( + edge, + prepare=partial( + _attention_input_step, + layer_input=layer_input, + fusions=fusions, + enters_stack=enters_stack, + residual_ops=residual_ops, + ), + input_move=_select_attention_input_move(edge.residual_to, edge.need), + ) + + +def _select_boundary_steps( + sides: DecoderLayerSides, + *, + fusions: Tuple["FusedMlpInput", ...] = (), + force_layernorm_before_gather: bool = False, + cp_moves: Optional[CpMoves] = None, + residual_ops: ResidualOps = ADD_AND_NORM, + attention_handoff: Callable = _hand_qkv_hook_its_input, + attention_fusions: Tuple[Callable, ...] = (), + enters_stack: bool = False, +) -> BoundarySteps: + """The steps of a decoder layer: its three boundaries, each chosen from the + declarations of its two sides, around the layer's residual operations. Both + edges that cross attention CP take the same ``cp_moves``; the attention's + input tries ``attention_fusions``, and the layer stack's first layer + (``enters_stack``) starts its residual there.""" + edges = decoder_layer_edges(sides) + out_of_ffn = make_boundary( + edges.out_of_ffn, reads=None, cp_moves=cp_moves, residual_ops=residual_ops + ) + into_ffn = make_boundary( + edges.into_ffn, + reads=InputRead.FFN, fusions=fusions, - residual_joins_sum=sides.residual_joins_attention_sum, + force_layernorm_before_gather=force_layernorm_before_gather, cp_moves=cp_moves, residual_ops=residual_ops, ) + into_attention = make_boundary( + edges.into_attention, + reads=InputRead.ATTENTION, + fusions=attention_fusions, + residual_ops=residual_ops, + enters_stack=enters_stack, + ) return BoundarySteps( - attention_input=_select_attention_input_move(rows, sides.attention), - ffn_input=ffn_input, - # What the FFN needs, still sharded over the axes it gathers itself. - ffn_input_rows=Layout( - sides.ffn.layout.sharded - | (sides.attention_output.layout.sharded & sides.ffn.gathers_itself) - ), - ffn_output=sides.ffn_output, - ffn_output_move=None if returns_over_dp else ffn_output_move, - ffn_output_move_completes_sum=completes_sum, - ffn_sum_is_movable=sides.ffn_output.group is not None, - fused=fused, - layer_input=layer_input, + attention_prepare=into_attention.prepare, + attention_input=into_attention.input_move, + ffn_input=into_ffn.prepare, + ffn_input_rows=into_ffn.input_rows, + ffn_output=edges.out_of_ffn.produced, + ffn_output_move=out_of_ffn.output_move, + ffn_output_move_completes_sum=out_of_ffn.output_move_completes_sum, + ffn_sum_is_movable=edges.out_of_ffn.produced.group is not None, + fused=into_ffn.fused, attention_handoff=attention_handoff, ) @@ -2429,6 +2564,25 @@ def _select_ffn_input( only when nothing is gathered or sliced, and only if it completes the sum the attention output owes. A write-back that is not a plain add runs only after the sum completes.""" + if produced.leaves_for_next_layer and produced.group not in ( + None, + SumGroup.ATTN_TP, + ): + # A producer that leaves its sum only for some batches (an FFN before + # this one) hands that sum on with the value: complete it, then the + # input is a complete output. + step, fused = _select_ffn_input( + StageOutput(produced.layout), + residual=residual, + residual_to=residual_to, + need=need, + force_layernorm_before_gather=force_layernorm_before_gather, + fusions=fusions, + residual_joins_sum=residual_joins_sum, + cp_moves=cp_moves, + residual_ops=residual_ops, + ) + return partial(_mlp_input_completing_owed, step=step), fused # What the attention output owes decides the steps: the attention-TP sum, # always left by the output projection, or nothing. owes_attention_tp = produced.group is SumGroup.ATTN_TP @@ -2659,6 +2813,22 @@ def mlp_input_kind( ) +def _mlp_input_completing_owed( + hidden_states: Union[torch.Tensor, "UnreducedOutput"], + residual: torch.Tensor, + forward_batch: ForwardBatch, + layernorm: torch.nn.Module, + context: CommunicateContext, + *, + step: Callable, +): + """An FFN input whose producer may have left its sum: complete what the + value owes, then run ``step`` on the complete output.""" + return step( + reduce_output(hidden_states), residual, forward_batch, layernorm, context + ) + + def _mlp_input_norm( hidden_states: torch.Tensor, residual: torch.Tensor, diff --git a/test/registered/unit/layers/moe/test_cutedsl_ar_fusion.py b/test/registered/unit/layers/moe/test_cutedsl_ar_fusion.py index 2f28e1d60786..92225791d4a1 100644 --- a/test/registered/unit/layers/moe/test_cutedsl_ar_fusion.py +++ b/test/registered/unit/layers/moe/test_cutedsl_ar_fusion.py @@ -1,3 +1,4 @@ +from functools import partial from types import SimpleNamespace from unittest.mock import patch @@ -6,9 +7,11 @@ from sglang.srt.layers.boundary_layout import SumGroup from sglang.srt.layers.communicator import ( + ADD_AND_NORM, LayerCommunicator, ScatterMode, UnreducedOutput, + _attention_input_step, ) from sglang.srt.layers.flashinfer_mnnvl_cutedsl import ( FlashInferMNNVLCuteDSLARFusion, @@ -40,6 +43,18 @@ def _communicator(): comm.post_attention_layernorm = RMSNorm(8, eps=1e-6) comm.enable_fused_ar_quant = False comm._attn_input_fusions = comm._select_attn_input_fusions() + # Only the ordinary batches' attention input half, with these entries. + comm._context = SimpleNamespace() + comm._sp_steps = comm._input_scattered_steps = comm._cp_steps = None + comm._steps = SimpleNamespace( + attention_prepare=partial( + _attention_input_step, + layer_input=None, + fusions=comm._attn_input_fusions, + enters_stack=False, + residual_ops=ADD_AND_NORM, + ) + ) return comm diff --git a/test/registered/unit/layers/test_boundary_edges.py b/test/registered/unit/layers/test_boundary_edges.py new file mode 100644 index 000000000000..b13b38a8523c --- /dev/null +++ b/test/registered/unit/layers/test_boundary_edges.py @@ -0,0 +1,232 @@ +import itertools +import unittest +from types import SimpleNamespace +from unittest.mock import MagicMock + +import torch + +from sglang.srt.layers import communicator as comm +from sglang.srt.layers.boundary_layout import ( + EdgeDecl, + Layout, + StageInput, + StageOutput, + SumGroup, + TokenAxis, + decoder_layer_edges, + decoder_layer_sides, + input_scattered_layer_sides, +) +from sglang.srt.layers.communicator import InputRead, make_boundary +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + +Pair = comm.CommunicateSummableTensorPairFn + + +def sizes(*, dp=1, cp=1, tp=1): + return { + TokenAxis.ATTN_DP: dp, + TokenAxis.ATTN_CP: cp, + TokenAxis.ATTN_TP_SCATTER: tp, + } + + +def rows(axis_sizes, *axes): + return Layout.sharded_over(*axes, axis_sizes=axis_sizes) + + +class TestDecoderLayerEdges(CustomTestCase): + """A decoder layer's three boundaries hand the residual on from one to the + next: each edge starts on the rows the one before it ends on.""" + + def test_each_edge_starts_where_the_one_before_ends(self): + for dp, tp, ffn_local, previous_local, last in itertools.product( + (1, 2), (1, 2), (False, True), (False, True), (False, True) + ): + with self.subTest( + dp=dp, tp=tp, ffn_local=ffn_local, previous_local=previous_local + ): + sides = decoder_layer_sides( + axis_sizes=sizes(dp=dp, tp=tp), + ffn_on_local_rows=ffn_local, + previous_on_local_rows=previous_local, + is_last_layer=last, + attention_gathers_local_rows=False, + ffn_group=SumGroup.TP, + leaves_for_next_layer=True, + leaves_for_reduce_scatter=True, + leaves_for_reduce_scatterv=True, + ) + edges = decoder_layer_edges(sides) + self.assertEqual(edges.into_attention.residual, sides.input_rows) + self.assertEqual( + edges.into_attention.residual_to, edges.into_ffn.residual + ) + self.assertEqual(edges.into_ffn.residual_to, edges.out_of_ffn.residual) + self.assertEqual(edges.out_of_ffn.residual_to, sides.output_rows) + self.assertEqual(edges.out_of_ffn.need.layout, sides.output_rows) + # Nothing owed by construction: a sum left for a batch comes + # with the value. + self.assertFalse(edges.into_attention.produced.always_leaves) + + def test_an_input_owed_by_construction_is_completed_onto_the_slice(self): + axis_sizes = sizes(tp=2) + sides = input_scattered_layer_sides( + axis_sizes=axis_sizes, ffn_group=SumGroup.TP, hands_on_partial=True + ) + edge = decoder_layer_edges(sides).into_attention + self.assertIs(edge.produced.group, SumGroup.TP) + self.assertTrue(edge.produced.always_leaves) + self.assertEqual(edge.residual_to, rows(axis_sizes, TokenAxis.ATTN_TP_SCATTER)) + boundary = make_boundary(edge, reads=InputRead.ATTENTION) + self.assertIs(boundary.prepare.keywords["layer_input"], comm.tp_reduce_scatter) + + +class TestTheConsumerHalfReadsOnlyItsOwnSide(CustomTestCase): + """Two layers build the two halves of one boundary without sharing it; the + consumer's half never depends on what the producer may leave for a batch, + which arrives with the value.""" + + def test_what_the_producer_may_leave_does_not_change_the_consumer_half(self): + axis_sizes = sizes(dp=2, tp=2) + attention = rows(axis_sizes, TokenAxis.ATTN_DP) + halves = set() + for group, next_layer, reduce_scatter, reduce_scatterv in itertools.product( + (None, SumGroup.TP, SumGroup.MOE_OUTPUT), + (False, True), + (False, True), + (False, True), + ): + produced = StageOutput( + attention, + group=group, + leaves_for_next_layer=next_layer, + leaves_for_reduce_scatter=reduce_scatter, + leaves_for_reduce_scatterv=reduce_scatterv, + ) + edge = EdgeDecl( + produced=produced, + need=StageInput(attention), + residual=attention, + residual_to=attention, + ) + boundary = make_boundary(edge, reads=InputRead.ATTENTION) + halves.add( + ( + boundary.prepare.func, + tuple(sorted(boundary.prepare.keywords.items())), + boundary.input_move, + ) + ) + self.assertEqual(len(halves), 1) + ((func, keywords, move),) = halves + self.assertIs(func, comm._attention_input_step) + self.assertIsNone(dict(keywords)["layer_input"]) + self.assertIs(move, comm.CommunicateSimpleFn._trivial) + + +class TestNonAlternatingEdges(CustomTestCase): + """Stages that do not alternate attention and FFN build their boundaries + through the same entry.""" + + def test_a_mixer_into_a_mixer(self): + axis_sizes = sizes(dp=2, tp=2) + attention = rows(axis_sizes, TokenAxis.ATTN_DP) + produced = StageOutput( + attention, group=SumGroup.ATTN_TP, leaves_for_next_layer=True + ) + edge = EdgeDecl( + produced=produced, + need=StageInput(attention), + residual=attention, + residual_to=attention, + ) + producer = make_boundary(edge, reads=None) + self.assertIs(producer.output_move, Pair._trivial) + consumer = make_boundary(edge, reads=InputRead.ATTENTION) + self.assertIsNone(consumer.prepare.keywords["layer_input"]) + self.assertIs(consumer.input_move, comm.CommunicateSimpleFn._trivial) + + def test_an_ffn_into_an_ffn(self): + axis_sizes = sizes(dp=2, tp=2) + attention = rows(axis_sizes, TokenAxis.ATTN_DP) + full = rows(axis_sizes) + complete = EdgeDecl( + produced=StageOutput(attention), + need=StageInput(full), + residual=attention, + residual_to=attention, + ) + step = make_boundary(complete, reads=InputRead.FFN).prepare + self.assertIs(step.func, comm._mlp_input_dp_replicate) + self.assertFalse(step.keywords["reduces_attention_tp"]) + # A sum the value carries is completed first, then the same steps run. + leaves = EdgeDecl( + produced=StageOutput( + attention, group=SumGroup.TP, leaves_for_next_layer=True + ), + need=StageInput(full), + residual=attention, + residual_to=attention, + ) + step = make_boundary(leaves, reads=InputRead.FFN).prepare + self.assertIs(step.func, comm._mlp_input_completing_owed) + self.assertIs(step.keywords["step"].func, comm._mlp_input_dp_replicate) + # A producer that always leaves a sum other than the attention TP's is + # not supported. + always = EdgeDecl( + produced=StageOutput(attention, group=SumGroup.TP, always_leaves=True), + need=StageInput(full), + residual=attention, + residual_to=attention, + ) + with self.assertRaises(NotImplementedError): + make_boundary(always, reads=InputRead.FFN) + + def test_the_next_ffn_completes_a_sum_left_for_it_once(self): + full = rows(sizes(tp=2)) + edge = EdgeDecl( + produced=StageOutput(full, group=SumGroup.TP, leaves_for_next_layer=True), + need=StageInput(full), + residual=full, + residual_to=full, + ) + step = make_boundary(edge, reads=InputRead.FFN).prepare + group = SimpleNamespace(all_reduce=MagicMock(side_effect=lambda h: h * 2)) + + def norm(hidden, residual): + return hidden + residual, hidden + residual + + residual = torch.full((3, 4), 5.0) + hidden, _ = step( + comm.UnreducedOutput(torch.ones(3, 4), group=group), + residual, + None, + norm, + None, + ) + group.all_reduce.assert_called_once() + torch.testing.assert_close(hidden, torch.full((3, 4), 7.0)) + group.all_reduce.reset_mock() + hidden, _ = step(torch.full((3, 4), 2.0), residual, None, norm, None) + group.all_reduce.assert_not_called() + torch.testing.assert_close(hidden, torch.full((3, 4), 7.0)) + + def test_the_producer_half_ends_on_the_rows_it_hands_on(self): + axis_sizes = sizes(dp=2, tp=2) + attention = rows(axis_sizes, TokenAxis.ATTN_DP) + edge = EdgeDecl( + produced=StageOutput(attention), + need=StageInput(rows(axis_sizes)), + residual=attention, + residual_to=attention, + ) + with self.assertRaises(NotImplementedError): + make_boundary(edge, reads=None) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/layers/test_communicator_ffn_exit.py b/test/registered/unit/layers/test_communicator_ffn_exit.py index e70e8005f15e..71b324df3146 100644 --- a/test/registered/unit/layers/test_communicator_ffn_exit.py +++ b/test/registered/unit/layers/test_communicator_ffn_exit.py @@ -1,6 +1,7 @@ import ast import types import unittest +from functools import partial from pathlib import Path from unittest.mock import MagicMock, patch @@ -49,6 +50,13 @@ def ordinary_steps(ffn_output, *, returns_over_dp=False): """A layer's ordinary steps, carrying the FFN output declaration the exit reads and whether the output goes back over attention DP.""" return comm.BoundarySteps( + attention_prepare=partial( + comm._attention_input_step, + layer_input=None, + fusions=(), + enters_stack=False, + residual_ops=comm.ADD_AND_NORM, + ), attention_input=comm.CommunicateSimpleFn._trivial, ffn_input=comm._mlp_input_norm, ffn_input_rows=Layout(frozenset()), diff --git a/test/registered/unit/layers/test_declared_decoder_boundary.py b/test/registered/unit/layers/test_declared_decoder_boundary.py index cb9fb2e33862..a5afabc566f3 100644 --- a/test/registered/unit/layers/test_declared_decoder_boundary.py +++ b/test/registered/unit/layers/test_declared_decoder_boundary.py @@ -740,7 +740,7 @@ def test_an_input_scattered_batch_keeps_the_residual_on_the_slice(self): # Only the first layer's input, the embedding's partial sum, # is completed onto the slice. self.assertIs( - steps.layer_input, + steps.attention_prepare.keywords["layer_input"], comm.tp_reduce_scatter if layer_id == 0 else None, ) self.assertIs(steps.attention_input, comm.CommunicateSimpleFn._trivial) @@ -1199,6 +1199,7 @@ class TestTheSequenceParallelRegion(CustomTestCase): SIZES = {TokenAxis.ATTN_DP: 1, TokenAxis.ATTN_CP: 1, TokenAxis.ATTN_TP_SCATTER: 2} SP_STEPS = comm.BoundarySteps( + attention_prepare=comm._attention_input_step, attention_input=comm.CommunicateSimpleFn._trivial, ffn_input=comm._mlp_input_norm, # The linears gather each rank's slice themselves. @@ -1209,9 +1210,17 @@ class TestTheSequenceParallelRegion(CustomTestCase): ) def unbound(self, steps): - """``steps`` with the plain residual's binding taken off the FFN input.""" + """``steps`` with the plain residual's binding taken off the FFN input + and the attention input, whose input owes nothing in the region.""" self.assertEqual(steps.ffn_input.keywords, {"residual_ops": comm.ADD_AND_NORM}) - return msgspec.structs.replace(steps, ffn_input=steps.ffn_input.func) + attention = steps.attention_prepare.keywords + self.assertIsNone(attention["layer_input"]) + self.assertIs(attention["residual_ops"], comm.ADD_AND_NORM) + return msgspec.structs.replace( + steps, + ffn_input=steps.ffn_input.func, + attention_prepare=steps.attention_prepare.func, + ) def test_the_region_declarations_choose_local_steps(self): for attn_tp in (2, 4): @@ -1238,7 +1247,6 @@ def test_the_region_declarations_choose_local_steps(self): self.SP_STEPS.ffn_output_move, ), ) - self.assertIsNone(steps.layer_input) def test_a_layer_under_sp_takes_both_sets_of_steps(self): parallel = parallel_of(attn_dp=1, attn_tp=2) @@ -1318,7 +1326,10 @@ def test_the_declarations_choose_the_steps(self): hands_on_partial=hands_on, ) steps = comm._select_boundary_steps(sides) - self.assertIs(steps.layer_input, comm.tp_reduce_scatter) + self.assertIs( + steps.attention_prepare.keywords["layer_input"], + comm.tp_reduce_scatter, + ) self.assertIs(steps.attention_input, comm.CommunicateSimpleFn._trivial) self.assertIs(steps.ffn_input, comm._mlp_input_residual_into_sum) self.assertIs( @@ -1492,16 +1503,73 @@ def test_the_scatter_mode_steps_complete_a_scattered_input(self): ), parallel_of(attn_dp=1, attn_tp=2), )._steps - self.assertIs(steps.layer_input, comm._complete_scattered_input) + self.assertIs( + steps.attention_prepare.keywords["layer_input"], + comm._complete_scattered_input, + ) hidden, residual = torch.ones(2, HIDDEN), torch.ones(2, HIDDEN) with patch.object( comm, "get_attn_tp_context", lambda: SimpleNamespace(input_scattered=False) ): self.assertEqual( - steps.layer_input(hidden, residual, None), (hidden, residual) + steps.attention_prepare.keywords["layer_input"](hidden, residual, None), + (hidden, residual), ) +class TestTheAttentionInputHalf(CustomTestCase): + """Every batch variant's half into the attention tries the fused entries + the layer chose, and starts the residual only on the stack's first + layer.""" + + def build_fusable(self, modes, parallel, **planning_kwargs): + with planning(parallel, **planning_kwargs): + return LayerCommunicator( + layer_scatter_modes=modes, + input_layernorm=FusableNorm(), + post_attention_layernorm=Norm(), + ) + + def test_every_variant_takes_the_layers_fused_entries(self): + cp = parallel_of(attn_dp=1, attn_tp=2, attn_cp=2, enable_prefill_cp=True) + for name, parallel, kwargs in ( + ( + "input-scattered", + parallel_of(attn_dp=1, attn_tp=2, enable_attn_tp_input_scattered=True), + {}, + ), + ("LayerNorm SP", parallel_of(attn_dp=1, attn_tp=2), {"sp": True}), + ("prefill CP", cp, {}), + ): + for layer_id in (0, 1): + with self.subTest(name, layer_id=layer_id): + modes = planned_modes( + layer_id, + 3, + sparse=False, + previous_sparse=False, + parallel=parallel, + ) + communicator = self.build_fusable(modes, parallel, **kwargs) + fusions = communicator._attn_input_fusions + self.assertTrue(fusions) + variants = [ + steps + for steps in ( + communicator._steps, + communicator._sp_steps, + communicator._input_scattered_steps, + communicator._cp_steps, + ) + if steps is not None + ] + self.assertEqual(len(variants), 2) + for steps in variants: + bound = steps.attention_prepare.keywords + self.assertIs(bound["fusions"], fusions) + self.assertIs(bound["enters_stack"], layer_id == 0) + + class TestPrefillCP(CustomTestCase): """A prefill CP shards a batch's tokens over attention CP only on a CP extend; the layer runs its CP steps on those batches and its ordinary ones diff --git a/test/registered/unit/layers/test_layer_communicator_fusion_gate.py b/test/registered/unit/layers/test_layer_communicator_fusion_gate.py index 4dd680b3fda0..30fe2c1bb70f 100644 --- a/test/registered/unit/layers/test_layer_communicator_fusion_gate.py +++ b/test/registered/unit/layers/test_layer_communicator_fusion_gate.py @@ -1,6 +1,7 @@ import contextlib import types import unittest +from functools import partial from unittest.mock import MagicMock, patch import torch @@ -42,6 +43,13 @@ def sp_region_steps(): def _steps(*, ffn_output=None, returns_over_dp=False, ffn_sum_is_movable=True): return comm.BoundarySteps( + attention_prepare=partial( + comm._attention_input_step, + layer_input=None, + fusions=(), + enters_stack=False, + residual_ops=comm.ADD_AND_NORM, + ), attention_input=comm.CommunicateSimpleFn._trivial, ffn_input=comm._mlp_input_norm, ffn_input_rows=Layout(frozenset()), diff --git a/test/registered/unit/layers/test_layernorm_sp.py b/test/registered/unit/layers/test_layernorm_sp.py index 21bd2d6d5a3e..4a9bb8021c9d 100644 --- a/test/registered/unit/layers/test_layernorm_sp.py +++ b/test/registered/unit/layers/test_layernorm_sp.py @@ -6,6 +6,7 @@ """ import unittest +from functools import partial from types import SimpleNamespace from unittest.mock import MagicMock, patch @@ -186,6 +187,13 @@ def communicator(self, *, first_layer): def ordinary_steps(self, attention_input): """The layer's steps outside the region, which must not run inside it.""" return comm.BoundarySteps( + attention_prepare=partial( + comm._attention_input_step, + layer_input=None, + fusions=(), + enters_stack=False, + residual_ops=comm.ADD_AND_NORM, + ), attention_input=attention_input, ffn_input=MagicMock(side_effect=AssertionError("ordinary FFN input ran")), ffn_input_rows=Layout(frozenset()), diff --git a/test/registered/unit/layers/test_postprocess_reduce_scatterv.py b/test/registered/unit/layers/test_postprocess_reduce_scatterv.py index 3571479aa72f..2c848af16323 100644 --- a/test/registered/unit/layers/test_postprocess_reduce_scatterv.py +++ b/test/registered/unit/layers/test_postprocess_reduce_scatterv.py @@ -6,6 +6,7 @@ import itertools import types import unittest +from functools import partial from unittest.mock import patch import torch @@ -33,6 +34,13 @@ def steps(*, ffn_output_move): """A layer's steps with the given FFN output move; None sends the output back over attention DP.""" return comm.BoundarySteps( + attention_prepare=partial( + comm._attention_input_step, + layer_input=None, + fusions=(), + enters_stack=False, + residual_ops=comm.ADD_AND_NORM, + ), attention_input=comm.CommunicateSimpleFn._trivial, ffn_input=comm._mlp_input_norm, ffn_input_rows=comm.Layout(frozenset()), diff --git a/test/registered/unit/layers/test_prepare_attn_steps.py b/test/registered/unit/layers/test_prepare_attn_steps.py index 9eb2b2bd28de..9272b6276821 100644 --- a/test/registered/unit/layers/test_prepare_attn_steps.py +++ b/test/registered/unit/layers/test_prepare_attn_steps.py @@ -3,9 +3,11 @@ import contextlib import unittest +from functools import partial from types import SimpleNamespace from unittest.mock import MagicMock, patch +import msgspec import torch from sglang.srt.layers import communicator as comm @@ -51,8 +53,17 @@ def communicator(norm): c._context = None c.enable_fused_ar_quant = False c.fused_ar_quant_keep_bf16 = False + # Construction picks the fused entries; call it under the platform patches. + c._attn_input_fusions = c._select_attn_input_fusions() # A layer whose attention takes its input as it is and owes nothing on it. c._steps = comm.BoundarySteps( + attention_prepare=partial( + comm._attention_input_step, + layer_input=None, + fusions=c._attn_input_fusions, + enters_stack=False, + residual_ops=comm.ADD_AND_NORM, + ), attention_input=lambda hidden_states, **_: hidden_states, ffn_input=comm._mlp_input_norm, ffn_input_rows=comm.Layout(frozenset()), @@ -60,8 +71,6 @@ def communicator(norm): ffn_output_move=comm.CommunicateSummableTensorPairFn._trivial, ffn_sum_is_movable=False, ) - # Construction picks the fused entries; call it under the platform patches. - c._attn_input_fusions = c._select_attn_input_fusions() return c @@ -224,6 +233,31 @@ def test_an_owed_reduce_scatter_reaches_local_tokens_before_the_norm(self): self.assertEqual(norm.calls, []) self.assertIs(residual, local) + def test_an_owed_reduce_scatter_is_not_offered_to_a_fused_kernel(self): + with platform(fusion=True) as (_, all_reduce, _): + c = communicator(Norm()) + takes_anything = MagicMock(return_value=("fused", "fused")) + c._steps = msgspec.structs.replace( + c._steps, + attention_prepare=partial( + comm._attention_input_step, + layer_input=None, + fusions=(takes_anything,), + enters_stack=False, + residual_ops=comm.ADD_AND_NORM, + ), + ) + partial_sum = torch.ones(3, 4) + step = MagicMock(return_value=torch.full((1, 4), 7.0)) + c.prepare_attn( + comm.UnreducedOutput(partial_sum, reduce_and_redistribute=step), + torch.zeros(1, 4), + None, + ) + step.assert_called_once_with(partial_sum) + takes_anything.assert_not_called() + self.assertEqual(all_reduce.call_count, 0) + def test_an_unreduced_output_without_a_residual_is_rejected(self): for step in (None, MagicMock()): with ( diff --git a/test/registered/unit/lora/test_triton_dp_attention_unit.py b/test/registered/unit/lora/test_triton_dp_attention_unit.py index 54694cc543b5..6eed83c57317 100644 --- a/test/registered/unit/lora/test_triton_dp_attention_unit.py +++ b/test/registered/unit/lora/test_triton_dp_attention_unit.py @@ -2,6 +2,7 @@ import sys from contextlib import nullcontext +from functools import partial from types import SimpleNamespace from unittest.mock import patch @@ -9,7 +10,11 @@ import torch from sglang.srt.layers.boundary_layout import Layout, TokenAxis -from sglang.srt.layers.communicator import LayerCommunicator +from sglang.srt.layers.communicator import ( + ADD_AND_NORM, + LayerCommunicator, + _attention_input_step, +) from sglang.srt.layers.dp_attention import DpPaddingMode from sglang.srt.lora.backend.base_backend import BaseLoRABackend from sglang.srt.lora.backend.triton_backend import ( @@ -225,15 +230,20 @@ def test_communicator_publishes_layout_at_each_transition( communicator.post_attention_layernorm = None communicator.input_layernorm = lambda x: x communicator.qkv_latent_func = None - communicator._attn_input_fusions = () gathered, local = Layout(frozenset()), Layout(frozenset({TokenAxis.ATTN_DP})) # The rows of the steps the batch runs decide, not the ordinary steps'. communicator._steps = SimpleNamespace( ffn_input_rows=local if gathered_over_dp else gathered ) selected = SimpleNamespace( + attention_prepare=partial( + _attention_input_step, + layer_input=None, + fusions=(), + enters_stack=False, + residual_ops=ADD_AND_NORM, + ), attention_input=lambda hidden_states, **kwargs: hidden_states, - layer_input=None, attention_handoff=lambda hidden_states, *args: hidden_states, ffn_input=lambda hidden_states, residual, *args: (hidden_states, residual), ffn_input_rows=gathered if gathered_over_dp else local, diff --git a/test/registered/unit/models/test_aux_capture_deferred_allreduce.py b/test/registered/unit/models/test_aux_capture_deferred_allreduce.py index f6cb53cd4e02..b230861dfcb4 100644 --- a/test/registered/unit/models/test_aux_capture_deferred_allreduce.py +++ b/test/registered/unit/models/test_aux_capture_deferred_allreduce.py @@ -1,4 +1,5 @@ import unittest +from functools import partial from types import SimpleNamespace from unittest.mock import patch @@ -47,6 +48,13 @@ def __init__(self, defer, return_topk=False): self.return_topk = return_topk self.layer_communicator = comm.LayerCommunicator.__new__(comm.LayerCommunicator) self.layer_communicator._steps = comm.BoundarySteps( + attention_prepare=partial( + comm._attention_input_step, + layer_input=None, + fusions=(), + enters_stack=False, + residual_ops=comm.ADD_AND_NORM, + ), attention_input=comm.CommunicateSimpleFn._trivial, ffn_input=comm._mlp_input_norm, ffn_input_rows=comm.Layout(frozenset()),