From a61ec34ed1ecb4a4e5f029640f6013cb07678d55 Mon Sep 17 00:00:00 2001 From: Cheng Wan Date: Sat, 26 Sep 2026 21:30:38 -0700 Subject: [PATCH] Build the boundary into any stage with one construction make_boundary no longer routes on whether its consumer is an attention or an FFN. One selection over the edge's declarations chooses the steps into a stage: - what a value carries for a batch (a sum or a handoff its producer left) is completed first, by a step every consumer runs; - a sum the producer always leaves is completed as the layouts say; - the input reaches the rows the stage needs before the read (a gather over attention DP or CP, each rank's own slice) or after it (a gather over attention TP). The producer's half after a layer's last stage is make_output_boundary. The producer's update chooses the order. The DP partial order, a residual that joins the sum, and the fused add + norm kernels run only for a plain add, so the boundary, not the layer, now decides whether a layer's fused kernels are tried. A sum the producer always leaves over the TP group is completed like one over attention TP where every TP rank holds the rows, and is rejected on rows sharded over attention DP or CP. InputRead is removed; a layer that is one stage says which of the two stages it stands for with StageKind. A stage entry with nothing to move after prepare has no input move. No behavior change. --- .../srt/layers/communicator/__init__.py | 6 +- .../srt/layers/communicator/boundary.py | 399 ++++++++---------- .../sglang/srt/layers/communicator/layer.py | 52 ++- python/sglang/srt/layers/communicator/ops.py | 107 ++++- python/sglang/srt/models/nemotron_h_utils.py | 21 +- .../unit/layers/moe/test_cutedsl_ar_fusion.py | 17 +- .../unit/layers/test_boundary_edges.py | 213 ++++++++-- .../unit/layers/test_communicator_ffn_exit.py | 24 +- .../layers/test_declared_decoder_boundary.py | 206 +++++---- .../layers/test_dp_attention_cp_gather.py | 2 +- .../test_layer_communicator_fusion_gate.py | 24 +- .../unit/layers/test_layernorm_sp.py | 15 +- .../test_postprocess_reduce_scatterv.py | 25 +- .../unit/layers/test_prepare_attn_steps.py | 40 +- .../lora/test_triton_dp_attention_unit.py | 17 +- .../test_aux_capture_deferred_allreduce.py | 25 +- 16 files changed, 742 insertions(+), 451 deletions(-) diff --git a/python/sglang/srt/layers/communicator/__init__.py b/python/sglang/srt/layers/communicator/__init__.py index ea64def84b73..0458a9894e0a 100644 --- a/python/sglang/srt/layers/communicator/__init__.py +++ b/python/sglang/srt/layers/communicator/__init__.py @@ -22,16 +22,17 @@ DecoderLayerSides, EdgeDecl, FusedMlpInput, - InputRead, LayerStage, StageDecl, StageEntry, StageInput, + StageKind, StageOutput, decoder_layer_edges, decoder_layer_sides, input_scattered_layer_sides, make_boundary, + make_output_boundary, sequence_parallel_layer_sides, stage_edges, tbo_split_moves, @@ -95,7 +96,6 @@ "FfnExitFusion", "FusedMlpInput", "HandoffOutput", - "InputRead", "LayerCommunicator", "LayerResidual", "LayerScatterModes", @@ -111,6 +111,7 @@ "StageDecl", "StageEntry", "StageInput", + "StageKind", "StageOutput", "SumGroup", "TokenAxis", @@ -122,6 +123,7 @@ "input_scattered_layer_sides", "layer_input_buffer", "make_boundary", + "make_output_boundary", "moe_cp_gathers_sparse_moe_input", "move_rows", "reduce_output", diff --git a/python/sglang/srt/layers/communicator/boundary.py b/python/sglang/srt/layers/communicator/boundary.py index 7940efc8bf22..dfcbef2ca8eb 100644 --- a/python/sglang/srt/layers/communicator/boundary.py +++ b/python/sglang/srt/layers/communicator/boundary.py @@ -16,14 +16,13 @@ from enum import Enum, auto from functools import partial -from typing import Callable, FrozenSet, Mapping, Optional, Tuple, Union +from typing import Callable, FrozenSet, Mapping, Optional, Tuple import msgspec import torch from sglang.srt.distributed import GroupCoordinator from sglang.srt.layers.communicator.layout import ( - CommunicateContext, Layout, SumGroup, TokenAxis, @@ -35,25 +34,20 @@ from sglang.srt.layers.communicator.ops import ( CommunicateSimpleFn, CommunicateSummableTensorPairFn, + _consumer_step, _hand_qkv_hook_its_input, - _mlp_input_completing_owed, _mlp_input_dp_partial, _mlp_input_dp_replicate, _mlp_input_gather_attention_cp, _mlp_input_gather_moe_cp, - _mlp_input_norm, _mlp_input_on_residual_shard, _mlp_input_residual_into_sum, _mlp_input_scatter, _mlp_input_slice, _mlp_input_without_dp, + _read_input, tp_reduce_scatter, ) -from sglang.srt.layers.communicator.output import ( - HandoffOutput, - UnreducedOutput, - reduce_output, -) from sglang.srt.layers.communicator.residual import ( LayerResidual, StageRead, @@ -64,7 +58,6 @@ NORM_QUANT_READ, NORM_READ, ) -from sglang.srt.model_executor.forward_batch_info import ForwardBatch from sglang.srt.runtime_context import get_parallel @@ -560,63 +553,10 @@ def returns_over_dp(self) -> bool: return self.ffn_output_move is None -def _attention_input_step( - hidden_states: Union[torch.Tensor, "UnreducedOutput"], - residual: Optional[torch.Tensor], - forward_batch: ForwardBatch, - norm: torch.nn.Module, - context: "CommunicateContext", - *, - quant_format: str, - post_residual_addition: Optional[torch.Tensor], - layer_input: Optional[Callable], - fusions: Tuple[Callable, ...], - enters_stack: bool, - read: StageRead, - update: StageUpdate, -) -> 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 - if isinstance(owed, HandoffOutput): - # No fused kernel took the handoff: its producer completes it. - hidden_states, owed = reduce_output(owed), None - else: - hidden_states = owed.partial - if layer_input is not None: - hidden_states, residual = layer_input(hidden_states, residual, context) - if enters: - hidden_states, residual = read.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 read.read(hidden_states, norm, quant_format) - return read.update_and_read( - update, 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).""" +class StageKind(Enum): + """Which of a decoder layer's two stages a layer that is one stage takes + the place of: its norm, its read and update, its fused kernels and its + entry in the layer's steps.""" ATTENTION = auto() FFN = auto() @@ -624,12 +564,12 @@ class InputRead(Enum): class LayerStage(msgspec.Struct, frozen=True): """A layer that is one stage of a sequence of stages, each an - attention-like mixer or an FFN: how it reads its input, its two boundaries + attention-like mixer or an FFN: which of the two it is, its two boundaries (into it, and out of it onto the rows every layer hands on, as ``stage_edges`` gives them), and whether the layer stack starts at it.""" - reads: InputRead + kind: StageKind edges: Tuple[EdgeDecl, EdgeDecl] enters_stack: bool = False @@ -642,11 +582,11 @@ class Boundary(msgspec.Struct, frozen=True): 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. + # norm, and the moves onto the rows it needs that come before the read. prepare: Optional[Callable] = None - # Into an attention, the move onto its rows after prepare. + # The move onto the consumer's rows after prepare; None when there is none. input_move: Optional[Callable] = None - # The fused kernels an FFN's prepare tries first. + # The fused kernels prepare tries first on a sum the input always owes. 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 @@ -668,79 +608,70 @@ def input_rows(self) -> Layout: def make_boundary( edge: EdgeDecl, *, - reads: Optional[InputRead], - fusions: Tuple = (), + fusions: Tuple["FusedMlpInput", ...] = (), + carried_fusions: Tuple[Callable, ...] = (), force_layernorm_before_gather: bool = False, cp_moves: Optional[CpMoves] = None, 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, - update=edge.produced.update, - ) - 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, - read=edge.need.read, - update=edge.produced.update, - ) - return Boundary(edge, prepare=input_step, fused=fused) - layer_input = 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 ( - 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"{edge=}") - layer_input = tp_reduce_scatter + """The consumer's half of ``edge``, chosen from the edge's declarations and + the producer's update: what a value carries for a batch (a sum or a handoff + its producer left) is completed first, trying ``carried_fusions``; a sum + the producer always leaves as the layouts say, trying ``fusions``. Fused + kernels run only for a plain add and a read that is the residual's norm + and leaves the residual as it is; orders that add the residual before the + sum completes only for a plain add. ``cp_moves`` for an edge that gathers + over attention CP; ``enters_stack`` for the edge into the layer stack's + first stage. The steps read only this edge's declarations, never what the + producer chose for a batch.""" + update = edge.produced.update + # A fused kernel runs the add and the norm itself. + plain = update.adds_plainly and edge.need.read.norms_plainly + step, fused, input_move = _select_input_steps( + edge.produced, + residual=edge.residual, + residual_to=edge.residual_to, + need=edge.need, + update=update, + fusions=fusions if plain else (), + force_layernorm_before_gather=force_layernorm_before_gather, + residual_joins_sum=edge.residual_joins_sum, + cp_moves=cp_moves, + enters_stack=enters_stack, + ) return Boundary( edge, prepare=partial( - _attention_input_step, - layer_input=layer_input, - # A fused kernel runs the add and the norm itself. - fusions=( - fusions - if edge.produced.update.adds_plainly and edge.need.read.norms_plainly - else () - ), - enters_stack=enters_stack, - read=edge.need.read, - update=edge.produced.update, + _consumer_step, + step=step, + carried_fusions=carried_fusions if plain else (), + owes_by_construction=edge.produced.always_leaves, ), - input_move=_select_attention_input_move(edge.residual_to, edge.need), + input_move=input_move, + fused=fused, + ) + + +def make_output_boundary( + edge: EdgeDecl, *, cp_moves: Optional[CpMoves] = None +) -> Boundary: + """The producer's half of ``edge``, out of a layer's last stage onto the + rows the layer hands on (``edge.need``), whose consumer runs in the next + layer: the postprocess that moves the output there. ``cp_moves`` for an + edge that returns across attention CP.""" + 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, + update=edge.produced.update, + ) + return Boundary( + edge, + output_move=None if returns_over_dp else output_move, + output_move_completes_sum=completes_sum, ) @@ -760,18 +691,18 @@ def _select_boundary_steps( 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) + out_of_ffn = make_output_boundary(edges.out_of_ffn, cp_moves=cp_moves) into_ffn = make_boundary( edges.into_ffn, - reads=InputRead.FFN, fusions=fusions, force_layernorm_before_gather=force_layernorm_before_gather, cp_moves=cp_moves, ) into_attention = make_boundary( edges.into_attention, - reads=InputRead.ATTENTION, - fusions=attention_fusions, + carried_fusions=attention_fusions, + force_layernorm_before_gather=force_layernorm_before_gather, + cp_moves=cp_moves, enters_stack=enters_stack, ) return BoundarySteps( @@ -784,6 +715,7 @@ def _select_boundary_steps( ffn=StageEntry( prepare=into_ffn.prepare, input_rows=into_ffn.input_rows, + input_move=into_ffn.input_move, fused=into_ffn.fused, ), ffn_output=edges.out_of_ffn.produced, @@ -793,69 +725,33 @@ def _select_boundary_steps( ) -def _select_attention_input_move(rows: Layout, need: StageInput) -> Callable: - """How the rows a layer takes become its attention's input: as they are, - or gathered over attention TP from each rank's slice, unless the attention - gathers them itself.""" - gathered = rows.sharded - need.layout.sharded - need.gathers_itself - if not need.layout.sharded <= rows.sharded or gathered not in ( - frozenset(), - {TokenAxis.ATTN_TP_SCATTER}, - ): - raise NotImplementedError(f"{rows=} {need=}") - if gathered: - return CommunicateSimpleFn._scattered_to_tp_attn_full - return CommunicateSimpleFn._trivial - - -def _select_ffn_input( +def _select_input_steps( produced: StageOutput, *, residual: Layout, residual_to: Layout, need: StageInput, - force_layernorm_before_gather: bool, + update: StageUpdate, fusions: Tuple[FusedMlpInput, ...], - residual_joins_sum: bool = False, - cp_moves: Optional[CpMoves] = None, - read: StageRead = NORM_READ, - update: StageUpdate = ADD, -) -> Tuple[Callable, Tuple[FusedMlpInput, ...]]: - """The steps from the attention output to the FFN input, and the fused - kernels they try first: complete the attention-TP sum, move the residual to - the rows it has while the FFN runs, write the output into it and read the - FFN input, and bring the rows to what the FFN's group needs: a gather over - attention DP, or each rank's own slice. A kernel in ``fusions`` is tried - 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, - read=read, - update=update, - ) - 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 - if owes_attention_tp != produced.always_leaves or produced.group not in ( - None, - SumGroup.ATTN_TP, - ): + force_layernorm_before_gather: bool, + residual_joins_sum: bool, + cp_moves: Optional[CpMoves], + enters_stack: bool, +) -> Tuple[Callable, Tuple[FusedMlpInput, ...], Optional[Callable]]: + """The steps into a stage for a value that is complete or owes the sum its + producer always leaves, the fused kernels they try first, and the move + after them: complete that sum, move the residual to the rows it has while + the stage runs, write the output into it with ``update`` and read the input, + and bring the input onto the rows the stage needs: a gather over attention + DP or CP, each rank's own slice, or after the read a gather over attention + TP. A kernel in ``fusions`` is tried only when nothing is gathered or + sliced, and only if it completes the sum the input owes. An update that is + not a plain add runs only after the sum completes.""" + read = need.read + # What the input owes by construction; a sum the producer leaves only for + # some batches comes with the value and is completed before these steps. + owes = produced.group if produced.always_leaves else None + if produced.always_leaves and owes is None: raise NotImplementedError(f"{produced=}") gathered = produced.layout.sharded - need.layout.sharded - need.gathers_itself sliced = need.layout.sharded - produced.layout.sharded @@ -866,44 +762,89 @@ def _select_ffn_input( if ( sliced != {TokenAxis.ATTN_TP_SCATTER} or gathered + or owes not in (None, SumGroup.ATTN_TP) or residual_to != need.layout or residual not in (produced.layout, need.layout) ): raise NotImplementedError(f"{produced=} {residual=} {need=}") return ( partial( - _mlp_input_scatter if owes_attention_tp else _mlp_input_slice, + _mlp_input_scatter if owes is SumGroup.ATTN_TP else _mlp_input_slice, scatters_residual=residual != residual_to, read=read, update=update, ), (), + None, ) if residual_to.sharded - produced.layout.sharded == {TokenAxis.ATTN_TP_SCATTER}: - # The residual stays on each rank's slice while the FFN takes the - # attention's rows (MHC on an input-scattered batch). - if gathered or not owes_attention_tp or residual != residual_to: + if residual == residual_to: + # The residual stays on each rank's slice while the stage takes the + # rows around it (MHC on an input-scattered batch). + if gathered or owes is not SumGroup.ATTN_TP: + raise NotImplementedError(f"{produced=} {residual=} {need=}") + return ( + partial(_mlp_input_on_residual_shard, read=read, update=update), + (), + None, + ) + # A reduce-scatter completes the TP sum onto each rank's slice, which + # the stage takes: its group is the TP group without attention DP or CP. + if ( + owes is not SumGroup.TP + or residual.sharded + or TokenAxis.ATTN_TP_SCATTER not in need.gathers_itself + or residual_to.sharded != {TokenAxis.ATTN_TP_SCATTER} + ): raise NotImplementedError(f"{produced=} {residual=} {need=}") return ( - partial(_mlp_input_on_residual_shard, read=read, update=update), + partial( + _read_input, + layer_input=tp_reduce_scatter, + enters_stack=enters_stack, + read=read, + update=update, + ), (), + None, ) if gathered == {TokenAxis.ATTN_CP}: # Each CP rank completes its own chunk, then the CP moves gather them. if cp_moves is None: raise NotImplementedError(f"{produced=} {need=}") - on_chunk, fused = _select_ffn_input( + on_chunk, fused, _ = _select_input_steps( produced, residual=residual, residual_to=residual_to, - need=StageInput(produced.layout), - force_layernorm_before_gather=force_layernorm_before_gather, + need=StageInput(produced.layout, read=read), + update=update, fusions=fusions, + force_layernorm_before_gather=force_layernorm_before_gather, residual_joins_sum=residual_joins_sum, - read=read, - update=update, + cp_moves=cp_moves, + enters_stack=enters_stack, + ) + return partial(cp_moves.gather, gather=on_chunk), fused, None + if gathered == {TokenAxis.ATTN_TP_SCATTER}: + # A complete input on each rank's slice, gathered over attention TP + # once it is read. + if ( + owes is not None + or residual != residual_to + or residual_to != produced.layout + ): + raise NotImplementedError(f"{produced=} {residual=} {need=}") + return ( + partial( + _read_input, + layer_input=None, + enters_stack=enters_stack, + read=read, + update=update, + ), + (), + CommunicateSimpleFn._scattered_to_tp_attn_full, ) - return partial(cp_moves.gather, gather=on_chunk), fused if ( residual_to != produced.layout or gathered @@ -919,30 +860,48 @@ def _select_ffn_input( # A residual arriving on each rank's slice is gathered back first. gathers_residual = residual != residual_to if not gathered: - if not owes_attention_tp: - return partial(_mlp_input_norm, read=read, update=update), () + if owes is None: + if gathers_residual: + raise NotImplementedError(f"{produced=} {residual=} {need=}") + return ( + partial( + _read_input, + layer_input=None, + enters_stack=enters_stack, + read=read, + update=update, + ), + (), + None, + ) if gathers_residual and residual_joins_sum: # Each rank adds its slice of the residual into its share of the # sum, so the all-reduce also brings the residual back to every row. - if not update.adds_plainly: + if owes is not SumGroup.ATTN_TP or not update.adds_plainly: raise NotImplementedError(f"{produced=} {residual=} {need=}") - return partial(_mlp_input_residual_into_sum, read=read), () - # A fused kernel runs the add and the norm itself. - fused = ( - tuple(f for f in fusions if f.completes is produced.group) - if update.adds_plainly and read.norms_plainly - else () - ) + return partial(_mlp_input_residual_into_sum, read=read), (), None + # A sum over TP completes only on rows every TP rank holds: the TP group + # spans attention DP and CP. + if owes is not SumGroup.ATTN_TP and ( + owes is not SumGroup.TP or produced.layout.sharded + ): + raise NotImplementedError(f"{produced=} {need=}") + fused = tuple(f for f in fusions if f.completes is owes) return ( partial( _mlp_input_without_dp, gathers_residual=gathers_residual, fusions=tuple(f.run for f in fused), + group=owes, read=read, update=update, ), fused, + None, ) + if owes not in (None, SumGroup.ATTN_TP): + raise NotImplementedError(f"{produced=} {need=}") + owes_attention_tp = owes is SumGroup.ATTN_TP # The partial order adds the residual on attention-TP rank 0 before the DP # gather's collective completes that sum, which only a plain residual add # allows. @@ -963,6 +922,7 @@ def _select_ffn_input( read=read, ), (), + None, ) return ( partial( @@ -974,6 +934,7 @@ def _select_ffn_input( update=update, ), (), + None, ) diff --git a/python/sglang/srt/layers/communicator/layer.py b/python/sglang/srt/layers/communicator/layer.py index 32c8359a55f3..2fc59409c017 100644 --- a/python/sglang/srt/layers/communicator/layer.py +++ b/python/sglang/srt/layers/communicator/layer.py @@ -31,14 +31,15 @@ BoundarySteps, DecoderLayerSides, FusedMlpInput, - InputRead, LayerStage, StageEntry, + StageKind, _cp_moves, _select_boundary_steps, decoder_layer_sides, input_scattered_layer_sides, make_boundary, + make_output_boundary, scattered_residual_layer_sides, sequence_parallel_layer_sides, with_residual, @@ -419,30 +420,26 @@ def _init_stage(self, stage: "LayerStage") -> None: self._declared = None self._cp_steps = self._input_scattered_steps = self._sp_steps = None into_edge, out_edge = stage.edges - reads_ffn = stage.reads is InputRead.FFN + is_ffn = stage.kind is StageKind.FFN self.stage_edges = stage.edges into = make_boundary( into_edge, - reads=stage.reads, - fusions=( - self._select_mlp_input_fusions() - if reads_ffn - else self._attn_input_fusions - ), + fusions=self._select_mlp_input_fusions() if is_ffn else (), + carried_fusions=() if is_ffn else self._attn_input_fusions, force_layernorm_before_gather=self.force_layernorm_before_dp_gather, enters_stack=stage.enters_stack, ) - out = make_boundary(out_edge, reads=None) + out = make_output_boundary(out_edge) entry = StageEntry( prepare=into.prepare, input_rows=into.input_rows, - input_move=None if reads_ffn else into.input_move, - handoff=None if reads_ffn else _hand_qkv_hook_its_input, + input_move=into.input_move, + handoff=None if is_ffn else _hand_qkv_hook_its_input, fused=into.fused, ) self._steps = BoundarySteps( - attention=None if reads_ffn else entry, - ffn=entry if reads_ffn else None, + attention=None if is_ffn else entry, + ffn=entry if is_ffn else None, ffn_output=out_edge.produced, ffn_output_move=out.output_move, ffn_output_move_completes_sum=out.output_move_completes_sum, @@ -662,10 +659,15 @@ def prepare_attn_and_capture_last_layer_outputs( post_residual_addition=post_residual_addition, ) if captured_last_layer_outputs is not None: - gathered_last_layer_output = self.attn.entry(forward_batch).input_move( - hidden_states=residual, - forward_batch=forward_batch, - context=self._context, + move = self.attn.entry(forward_batch).input_move + gathered_last_layer_output = ( + residual + if move is None + else move( + hidden_states=residual, + forward_batch=forward_batch, + context=self._context, + ) ) if ( gathered_last_layer_output is residual @@ -762,11 +764,9 @@ def _select_attn_input_fusions(self) -> Tuple[Callable, ...]: with the residual update and the input norm, in the order they are tried. Each takes (owed, residual, forward_batch, post_residual_addition) and returns None when it does not take the batch. They add the residual - plainly: what arrives is written in as this layer's FFN writes its own.""" - if not ( - self._residual.ffn_update.adds_plainly - and hasattr(self.input_layernorm, "forward_with_allreduce_fusion") - ): + plainly; the boundary tries them only when the update it writes in is a + plain add.""" + if not hasattr(self.input_layernorm, "forward_with_allreduce_fusion"): return () self._attn_input_fuses_quant = ( self.enable_fused_ar_quant @@ -821,11 +821,9 @@ def _reduce_output_and_update_and_read_residual( def _select_mlp_input_fusions(self) -> Tuple["FusedMlpInput", ...]: """The fused kernels that can take the attention -> FFN steps, in the - order they are tried. They add the residual plainly.""" - if not ( - self._residual.attention_update.adds_plainly - and hasattr(self.post_attention_layernorm, "forward_with_allreduce_fusion") - ): + order they are tried. They add the residual plainly; the boundary tries + them only when the update it writes in is a plain add.""" + if not hasattr(self.post_attention_layernorm, "forward_with_allreduce_fusion"): return () return ( FusedMlpInput( diff --git a/python/sglang/srt/layers/communicator/ops.py b/python/sglang/srt/layers/communicator/ops.py index a7aa7833ec75..d122af316a7d 100644 --- a/python/sglang/srt/layers/communicator/ops.py +++ b/python/sglang/srt/layers/communicator/ops.py @@ -36,11 +36,16 @@ from sglang.srt.layers.communicator.layout import ( CommunicateContext, Layout, + SumGroup, TokenAxis, _cp_shard_token_rows, moe_cp_gathered_rows, ) -from sglang.srt.layers.communicator.output import UnreducedOutput, reduce_output +from sglang.srt.layers.communicator.output import ( + HandoffOutput, + UnreducedOutput, + reduce_output, +) from sglang.srt.layers.communicator.residual import StageRead, StageUpdate from sglang.srt.layers.communicator.residual.add_norm import ADD, NORM_READ from sglang.srt.layers.communicator_dsa_cp import ( @@ -249,19 +254,26 @@ def _mlp_input_without_dp( *, gathers_residual: bool, fusions: Tuple[Callable, ...], + group: SumGroup = SumGroup.ATTN_TP, read: StageRead = NORM_READ, update: StageUpdate = ADD, ): + """Complete the sum the input owes over ``group`` on the rows it is on, + unless one of ``fusions`` does it with the residual add and the norm, then + write it into the residual and read the input.""" if gathers_residual: residual = update.residual_from_attn_tp_shards(residual) for fused in fusions: result = fused(hidden_states, residual, forward_batch) if result is not None: return result - # MHC sums its streams in full precision. - hidden_states = _mlp_input_reduce_output( - hidden_states, forward_batch, may_quantize=update.adds_plainly - ) + if group is SumGroup.ATTN_TP: + # MHC sums its streams in full precision. + hidden_states = _mlp_input_reduce_output( + hidden_states, forward_batch, may_quantize=update.adds_plainly + ) + else: + hidden_states = tensor_model_parallel_all_reduce(hidden_states) if _is_npu and context.cache is not None: _ = prepare_weight_cache(hidden_states, context.cache) return read.update_and_read(update, hidden_states, residual, layernorm) @@ -381,33 +393,88 @@ def _tp_all_gather_scattered_rows( return output -def _mlp_input_completing_owed( - hidden_states: Union[torch.Tensor, "UnreducedOutput"], - residual: torch.Tensor, +def _consumer_step( + hidden_states: Union[torch.Tensor, UnreducedOutput, HandoffOutput], + residual: Optional[torch.Tensor], forward_batch: ForwardBatch, - layernorm: torch.nn.Module, + norm: torch.nn.Module, context: CommunicateContext, *, step: Callable, + carried_fusions: Tuple[Callable, ...], + owes_by_construction: bool = False, + **call, ): - """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 - ) + """A boundary's half into a stage: complete what the value carries, a sum + or a handoff its producer left for this batch, then run ``step`` on the + complete value. Under attention DP the reduction back to this rank's tokens + comes first; otherwise one of ``carried_fusions`` may complete it with the + residual add and the read, and an empty batch has nothing to sum. An input + that owes its sum by construction (``owes_by_construction``) is that sum's + only carrier. ``call`` is what the stage's read takes (the attention's + ``quant_format`` and ``post_residual_addition``).""" + if not isinstance(hidden_states, torch.Tensor): + owed = hidden_states + if residual is None: + raise RuntimeError(f"{type(owed).__name__} requires residual input") + if owes_by_construction: + raise RuntimeError( + f"an input that owes its sum by construction arrived as " + f"{type(owed).__name__}" + ) + if ( + isinstance(owed, UnreducedOutput) + and owed.reduce_and_redistribute is not None + ): + hidden_states = reduce_output(owed) + else: + for fused in carried_fusions: + result = fused( + owed, residual, forward_batch, call.get("post_residual_addition") + ) + if result is not None: + return result + if isinstance(owed, HandoffOutput) or owed.partial.shape[0] != 0: + hidden_states = reduce_output(owed) + else: + hidden_states = owed.partial + return step(hidden_states, residual, forward_batch, norm, context, **call) -def _mlp_input_norm( +def _read_input( hidden_states: torch.Tensor, - residual: torch.Tensor, + residual: Optional[torch.Tensor], forward_batch: ForwardBatch, - layernorm: torch.nn.Module, + norm: torch.nn.Module, context: CommunicateContext, *, - read: StageRead = NORM_READ, - update: StageUpdate = ADD, + layer_input: Optional[Callable], + enters_stack: bool, + read: StageRead, + update: StageUpdate, + quant_format: str = "", + post_residual_addition: Optional[torch.Tensor] = None, ): - return read.update_and_read(update, hidden_states, residual, layernorm) + """Complete what the input owes by construction (``layer_input``), then + write the previous stage's output into the residual and read the stage's + input with ``norm``. The layer stack's first stage (``enters_stack``) + starts its residual from its input.""" + enters = residual is None and enters_stack + if layer_input is not None: + hidden_states, residual = layer_input(hidden_states, residual, context) + if enters: + hidden_states, residual = read.enter(hidden_states), None + if residual is None: + # The previous layer already wrote its output into the residual. + return read.read(hidden_states, norm, quant_format) + return read.update_and_read( + update, + hidden_states, + residual, + norm, + quant_format=quant_format, + post_residual_addition=post_residual_addition, + ) def _mlp_input_scatter( diff --git a/python/sglang/srt/models/nemotron_h_utils.py b/python/sglang/srt/models/nemotron_h_utils.py index 75bd4f92dfc6..6c36aee4ec83 100644 --- a/python/sglang/srt/models/nemotron_h_utils.py +++ b/python/sglang/srt/models/nemotron_h_utils.py @@ -5,7 +5,6 @@ from sglang.srt.configs.nemotron_h import ATTENTION, MAMBA, MOE from sglang.srt.layers.communicator import ( NORM_QUANT_READ, - InputRead, LayerCommunicator, LayerScatterModes, LayerStage, @@ -13,6 +12,7 @@ ScatterMode, StageDecl, StageInput, + StageKind, StageOutput, SumGroup, TokenAxis, @@ -29,12 +29,13 @@ def is_attn_layer(layer_type: str) -> bool: return layer_type in ATTN_LAYERS -def _stage_kind(pattern: str, layer_idx: int) -> Optional[InputRead]: - """How the stage at ``layer_idx`` reads its input: a Mamba or attention - mixer like an attention, an MLP or MoE like an FFN; None past either end.""" +def _stage_kind(pattern: str, layer_idx: int) -> Optional[StageKind]: + """Which stage of a decoder layer the stage at ``layer_idx`` stands for: a + Mamba or attention mixer the attention, an MLP or MoE the FFN; None past + either end.""" if not 0 <= layer_idx < len(pattern): return None - return InputRead.ATTENTION if is_attn_layer(pattern[layer_idx]) else InputRead.FFN + return StageKind.ATTENTION if is_attn_layer(pattern[layer_idx]) else StageKind.FFN def _stage_decl(pattern: str, layer_idx: int) -> StageDecl: @@ -56,8 +57,8 @@ def _stage_decl(pattern: str, layer_idx: int) -> StageDecl: StageOutput( attention, group=SumGroup.ATTN_TP if owes else None, - always_leaves=owes and following is InputRead.FFN, - leaves_for_next_layer=owes and following is InputRead.ATTENTION, + always_leaves=owes and following is StageKind.FFN, + leaves_for_next_layer=owes and following is StageKind.ATTENTION, ), ) sparse = pattern[layer_idx] == MOE @@ -75,7 +76,7 @@ def _stage_decl(pattern: str, layer_idx: int) -> StageDecl: StageOutput( full, group=SumGroup.MOE_OUTPUT if sparse else SumGroup.TP, - leaves_for_next_layer=following is InputRead.ATTENTION, + leaves_for_next_layer=following is StageKind.ATTENTION, leaves_for_reduce_scatter=True, leaves_for_reduce_scatterv=True, ), @@ -89,7 +90,7 @@ def layer_stage(pattern: str, layer_idx: int) -> LayerStage: TokenAxis.ATTN_DP, TokenAxis.ATTN_CP, axis_sizes=token_axis_sizes() ) return LayerStage( - reads=_stage_kind(pattern, layer_idx), + kind=_stage_kind(pattern, layer_idx), edges=stage_edges( previous=( _stage_decl(pattern, layer_idx - 1).output if layer_idx > 0 else None @@ -126,7 +127,7 @@ def make_layer_communicator( """The communicator of a layer that is one stage: only its own norm, and boundaries built from the stages next to it in the pattern.""" stage = layer_stage(pattern, layer_idx) - for_attn = stage.reads is InputRead.ATTENTION + for_attn = stage.kind is StageKind.ATTENTION return LayerCommunicator( layer_scatter_modes=_build_layer_scatter_modes( pattern[layer_idx] == MOE, is_last_layer=layer_idx == len(pattern) - 1 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 fc2375c6d835..bdbf70c1e7ad 100644 --- a/test/registered/unit/layers/moe/test_cutedsl_ar_fusion.py +++ b/test/registered/unit/layers/moe/test_cutedsl_ar_fusion.py @@ -16,7 +16,7 @@ UnreducedOutput, reduce_output, ) -from sglang.srt.layers.communicator.boundary import _attention_input_step +from sglang.srt.layers.communicator.ops import _consumer_step, _read_input from sglang.srt.layers.flashinfer_mnnvl_cutedsl import ( FlashInferMNNVLCuteDSLARFusion, _retargeted_config, @@ -54,12 +54,15 @@ def _communicator(): comm._steps = SimpleNamespace( attention=SimpleNamespace( prepare=partial( - _attention_input_step, - layer_input=None, - fusions=comm._attn_input_fusions, - enters_stack=False, - read=NORM_QUANT_READ, - update=ADD, + _consumer_step, + step=partial( + _read_input, + layer_input=None, + enters_stack=False, + read=NORM_QUANT_READ, + update=ADD, + ), + carried_fusions=comm._attn_input_fusions, ), input_move=None, handoff=None, diff --git a/test/registered/unit/layers/test_boundary_edges.py b/test/registered/unit/layers/test_boundary_edges.py index 83191166a5ad..d46a73f51951 100644 --- a/test/registered/unit/layers/test_boundary_edges.py +++ b/test/registered/unit/layers/test_boundary_edges.py @@ -3,6 +3,7 @@ from types import SimpleNamespace from unittest.mock import MagicMock, patch +import msgspec import torch from sglang.srt.layers import communicator as comm @@ -10,7 +11,6 @@ ADD, EdgeDecl, FusedMlpInput, - InputRead, Layout, StageInput, StageOutput, @@ -23,9 +23,11 @@ decoder_layer_sides, input_scattered_layer_sides, make_boundary, + make_output_boundary, ) from sglang.srt.layers.communicator import ops as comm_ops from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.communicator_patch import patch_communicator from sglang.test.test_utils import CustomTestCase register_cpu_ci(est_time=5, suite="base-a-test-cpu") @@ -88,8 +90,32 @@ def test_an_input_owed_by_construction_is_completed_onto_the_slice(self): 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) + step = make_boundary(edge).prepare.keywords["step"] + self.assertIs(step.keywords["layer_input"], comm.tp_reduce_scatter) + + def test_each_entry_keeps_the_move_after_its_read(self): + sides = decoder_layer_sides( + axis_sizes=sizes(dp=2, tp=2), + ffn_on_local_rows=False, + previous_on_local_rows=False, + is_last_layer=False, + 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) + moves = {edges.into_attention: MagicMock(), edges.into_ffn: MagicMock()} + real = comm_boundary.make_boundary + + def with_move(edge, **kwargs): + return msgspec.structs.replace(real(edge, **kwargs), input_move=moves[edge]) + + with patch.object(comm_boundary, "make_boundary", with_move): + steps = comm_boundary._select_boundary_steps(sides) + self.assertIs(steps.attention.input_move, moves[edges.into_attention]) + self.assertIs(steps.ffn.input_move, moves[edges.into_ffn]) class TestAMoeOnEachCpShard(CustomTestCase): @@ -147,19 +173,22 @@ def test_what_the_producer_may_leave_does_not_change_the_consumer_half(self): residual=attention, residual_to=attention, ) - boundary = make_boundary(edge, reads=InputRead.ATTENTION) + boundary = make_boundary(edge) + step = boundary.prepare.keywords["step"] halves.add( ( boundary.prepare.func, - tuple(sorted(boundary.prepare.keywords.items())), + step.func, + tuple(sorted(step.keywords.items())), boundary.input_move, ) ) self.assertEqual(len(halves), 1) - ((func, keywords, move),) = halves - self.assertIs(func, comm_boundary._attention_input_step) + ((func, step, keywords, move),) = halves + self.assertIs(func, comm_ops._consumer_step) + self.assertIs(step, comm_ops._read_input) self.assertIsNone(dict(keywords)["layer_input"]) - self.assertIs(move, comm.CommunicateSimpleFn._trivial) + self.assertIsNone(move) class TestNonAlternatingEdges(CustomTestCase): @@ -178,11 +207,12 @@ def test_a_mixer_into_a_mixer(self): residual=attention, residual_to=attention, ) - producer = make_boundary(edge, reads=None) + producer = make_output_boundary(edge) 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) + consumer = make_boundary(edge) + self.assertIs(consumer.prepare.func, comm_ops._consumer_step) + self.assertIsNone(consumer.prepare.keywords["step"].keywords["layer_input"]) + self.assertIsNone(consumer.input_move) def test_an_ffn_into_an_ffn(self): axis_sizes = sizes(dp=2, tp=2) @@ -194,7 +224,7 @@ def test_an_ffn_into_an_ffn(self): residual=attention, residual_to=attention, ) - step = make_boundary(complete, reads=InputRead.FFN).prepare + step = make_boundary(complete).prepare.keywords["step"] self.assertIs(step.func, comm_ops._mlp_input_dp_replicate) self.assertFalse(step.keywords["reduces_attention_tp"]) # A sum the value carries is completed first, then the same steps run. @@ -206,11 +236,12 @@ def test_an_ffn_into_an_ffn(self): residual=attention, residual_to=attention, ) - step = make_boundary(leaves, reads=InputRead.FFN).prepare - self.assertIs(step.func, comm_ops._mlp_input_completing_owed) - self.assertIs(step.keywords["step"].func, comm_ops._mlp_input_dp_replicate) - # A producer that always leaves a sum other than the attention TP's is - # not supported. + prepare = make_boundary(leaves).prepare + self.assertIs(prepare.func, comm_ops._consumer_step) + self.assertIs(prepare.keywords["step"].func, comm_ops._mlp_input_dp_replicate) + # A TP sum left on the rows of one attention-DP shard: the TP group spans + # the other shards, whose ranks hold other rows, so no all-reduce over it + # completes these. always = EdgeDecl( produced=StageOutput(attention, group=SumGroup.TP, always_leaves=True), need=StageInput(full), @@ -218,7 +249,119 @@ def test_an_ffn_into_an_ffn(self): residual_to=attention, ) with self.assertRaises(NotImplementedError): - make_boundary(always, reads=InputRead.FFN) + make_boundary(always) + + def test_an_ffn_that_always_leaves_its_tp_sum(self): + # Without attention DP every TP rank holds the rows, so TP completes the sum. + full = rows(sizes(tp=2)) + edge = EdgeDecl( + produced=StageOutput(full, group=SumGroup.TP, always_leaves=True), + need=StageInput(full), + residual=full, + residual_to=full, + ) + over_tp = comm.FusedMlpInput( + completes=SumGroup.TP, run=MagicMock(), may_return_new_residual=False + ) + over_attention_tp = comm.FusedMlpInput( + completes=SumGroup.ATTN_TP, run=MagicMock(), may_return_new_residual=False + ) + boundary = make_boundary(edge, fusions=(over_tp, over_attention_tp)) + self.assertEqual(boundary.fused, (over_tp,)) + step = make_boundary(edge).prepare + self.assertIs(step.keywords["step"].func, comm_ops._mlp_input_without_dp) + all_reduce = MagicMock(side_effect=lambda h: h * 2) + with patch_communicator("tensor_model_parallel_all_reduce", all_reduce): + hidden, residual = step( + torch.ones(3, 4), + torch.full((3, 4), 5.0), + None, + lambda hidden, residual: (hidden + residual, hidden + residual), + None, + ) + all_reduce.assert_called_once() + torch.testing.assert_close(hidden, torch.full((3, 4), 7.0)) + + +class _WrittenIn: + """An update that is not a plain add: the output is written into the + residual only once its sum is complete.""" + + adds_plainly = False + at_producer = False + + def update(self, hidden_states, residual): + return 2 * hidden_states + residual + + def residual_to_attn_tp_shard(self, residual, context): + return residual + + def residual_from_attn_tp_shards(self, residual): + return residual + + +class TestTheProducersUpdateChoosesTheOrder(CustomTestCase): + """On the same rows, the steps into a stage follow the update the producer + declares: a plain add may join a sum before it completes and run inside a + fused add + norm, another update runs after the sum completes.""" + + def edge(self, update): + axis_sizes = sizes(dp=2, tp=2) + attention = rows(axis_sizes, TokenAxis.ATTN_DP) + return EdgeDecl( + produced=StageOutput( + attention, group=SumGroup.ATTN_TP, always_leaves=True, update=update + ), + need=StageInput(rows(axis_sizes)), + residual=attention, + residual_to=attention, + ) + + def test_the_dp_gather_order(self): + plain = make_boundary(self.edge(comm.ADD)).prepare.keywords["step"] + self.assertIs(plain.func, comm_ops._mlp_input_dp_partial) + written = _WrittenIn() + other = make_boundary(self.edge(written)).prepare.keywords["step"] + self.assertIs(other.func, comm_ops._mlp_input_dp_replicate) + self.assertTrue(other.keywords["reduces_attention_tp"]) + self.assertIs(other.keywords["update"], written) + + def test_fused_kernels_take_only_a_plain_add(self): + full = rows(sizes(tp=2)) + fused = comm.FusedMlpInput( + completes=SumGroup.ATTN_TP, run=MagicMock(), may_return_new_residual=False + ) + carried = MagicMock() + for update, offered in ((comm.ADD, True), (_WrittenIn(), False)): + with self.subTest(adds_plainly=offered): + edge = EdgeDecl( + produced=StageOutput( + full, group=SumGroup.ATTN_TP, always_leaves=True, update=update + ), + need=StageInput(full), + residual=full, + residual_to=full, + ) + boundary = make_boundary( + edge, fusions=(fused,), carried_fusions=(carried,) + ) + self.assertEqual(boundary.fused, (fused,) if offered else ()) + self.assertEqual( + boundary.prepare.keywords["carried_fusions"], + (carried,) if offered else (), + ) + + def test_the_consumer_reads_as_it_declares(self): + full = rows(sizes(tp=2)) + read = MagicMock() + edge = EdgeDecl( + produced=StageOutput(full), + need=StageInput(full, read=read), + residual=full, + residual_to=full, + ) + step = make_boundary(edge).prepare.keywords["step"] + self.assertIs(step.keywords["read"], read) def test_the_next_ffn_completes_a_sum_left_for_it_once(self): full = rows(sizes(tp=2)) @@ -228,7 +371,7 @@ def test_the_next_ffn_completes_a_sum_left_for_it_once(self): residual=full, residual_to=full, ) - step = make_boundary(edge, reads=InputRead.FFN).prepare + step = make_boundary(edge).prepare group = SimpleNamespace(all_reduce=MagicMock(side_effect=lambda h: h * 2)) def norm(hidden, residual): @@ -259,7 +402,7 @@ def test_the_producer_half_ends_on_the_rows_it_hands_on(self): residual_to=attention, ) with self.assertRaises(NotImplementedError): - make_boundary(edge, reads=None) + make_output_boundary(edge) class ProbeRead: @@ -304,17 +447,19 @@ def _ffn_input(self, read, *, dp=1, residual_joins_sum=False, fusions=()): else attention ) need = rows(axis_sizes) if dp > 1 else attention - return comm_boundary._select_ffn_input( - StageOutput(attention, group=SumGroup.ATTN_TP, always_leaves=True), - residual=residual, - residual_to=attention, - need=StageInput(need, read=read), - force_layernorm_before_gather=False, + boundary = make_boundary( + EdgeDecl( + produced=StageOutput( + attention, group=SumGroup.ATTN_TP, always_leaves=True, update=ADD + ), + need=StageInput(need, read=read), + residual=residual, + residual_to=attention, + residual_joins_sum=residual_joins_sum, + ), fusions=fusions, - residual_joins_sum=residual_joins_sum, - read=read, - update=ADD, ) + return boundary.prepare.keywords["step"], boundary.fused def test_the_dp_partial_reads_the_gathered_sum_with_the_declared_read(self): read = ProbeRead(norms_plainly=True) @@ -391,12 +536,10 @@ def test_the_attention_input_tries_its_kernels_only_for_such_a_read(self): residual=attention, residual_to=attention, ) - boundary = make_boundary( - edge, reads=InputRead.ATTENTION, fusions=(kernel,) - ) - self.assertIs(boundary.prepare.keywords["read"], read) + boundary = make_boundary(edge, carried_fusions=(kernel,)) + self.assertIs(boundary.prepare.keywords["step"].keywords["read"], read) self.assertEqual( - boundary.prepare.keywords["fusions"], + boundary.prepare.keywords["carried_fusions"], (kernel,) if norms_plainly else (), ) diff --git a/test/registered/unit/layers/test_communicator_ffn_exit.py b/test/registered/unit/layers/test_communicator_ffn_exit.py index adaa4df2e428..630d32a35bd7 100644 --- a/test/registered/unit/layers/test_communicator_ffn_exit.py +++ b/test/registered/unit/layers/test_communicator_ffn_exit.py @@ -55,19 +55,29 @@ def ordinary_steps(ffn_output, *, returns_over_dp=False): return comm.BoundarySteps( attention=comm.StageEntry( prepare=partial( - comm_boundary._attention_input_step, - layer_input=None, - fusions=(), - enters_stack=False, - read=comm.NORM_QUANT_READ, - update=comm.ADD, + comm_ops._consumer_step, + step=partial( + comm_ops._read_input, + layer_input=None, + enters_stack=False, + read=comm.NORM_QUANT_READ, + update=comm.ADD, + ), + carried_fusions=(), ), input_rows=comm.Layout(frozenset()), input_move=comm.CommunicateSimpleFn._trivial, handoff=comm_ops._hand_qkv_hook_its_input, ), ffn=comm.StageEntry( - prepare=comm_ops._mlp_input_norm, input_rows=Layout(frozenset()) + prepare=partial( + comm_ops._read_input, + layer_input=None, + enters_stack=False, + read=comm.NORM_READ, + update=comm.ADD, + ), + input_rows=Layout(frozenset()), ), ffn_output=ffn_output, ffn_output_move=None diff --git a/test/registered/unit/layers/test_declared_decoder_boundary.py b/test/registered/unit/layers/test_declared_decoder_boundary.py index 733f3bd03a15..c3261523fad0 100644 --- a/test/registered/unit/layers/test_declared_decoder_boundary.py +++ b/test/registered/unit/layers/test_declared_decoder_boundary.py @@ -237,7 +237,9 @@ class TestWhichLayersUseDeclarations(CustomTestCase): declarations, or none do.""" def declared(self, communicator): - return getattr(communicator._steps.ffn.prepare, "func", None) in ( + return getattr( + communicator._steps.ffn.prepare.keywords["step"], "func", None + ) in ( comm_ops._mlp_input_dp_partial, comm_ops._mlp_input_dp_replicate, comm_ops._mlp_input_without_dp, @@ -250,10 +252,7 @@ def test_a_dense_model(self): with self.subTest(layer_id=layer_id): communicator = build(layer_facts(layer_id, 3), parallel) self.assertTrue(self.declared(communicator)) - self.assertIs( - communicator._steps.attention.input_move, - comm.CommunicateSimpleFn._trivial, - ) + self.assertIsNone(communicator._steps.attention.input_move) self.assertTrue(communicator._steps.returns_over_dp) def test_plain_tp(self): @@ -266,13 +265,16 @@ def test_plain_tp(self): post_attention_layernorm=norm(), ) self.assertIs( - communicator._steps.ffn.prepare.func, comm_ops._mlp_input_without_dp + communicator._steps.ffn.prepare.keywords["step"].func, + comm_ops._mlp_input_without_dp, ) entry = ( communicator._mlp_input_reduce_output_and_update_and_read_residual ) self.assertEqual( - communicator._steps.ffn.prepare.keywords["fusions"], + communicator._steps.ffn.prepare.keywords["step"].keywords[ + "fusions" + ], (entry,) if fuses else (), ) self.assertIs( @@ -291,7 +293,9 @@ def test_plain_tp(self): single = build(layer_facts(1, 3), one_rank) with planning(one_rank): self.assertIsNotNone(single._declared_sides()) - self.assertIs(single._steps.ffn.prepare.func, comm_ops._mlp_input_norm) + self.assertIs( + single._steps.ffn.prepare.keywords["step"].func, comm_ops._read_input + ) def test_the_order_follows_the_attention_output(self): for attn_tp, force, order in ( @@ -305,7 +309,9 @@ def test_the_order_follows_the_attention_output(self): parallel_of(attn_dp=2, attn_tp=attn_tp), force_layernorm_before_dp_gather=force, ) - self.assertIs(communicator._steps.ffn.prepare.func, order) + self.assertIs( + communicator._steps.ffn.prepare.keywords["step"].func, order + ) def test_a_dense_first_layer_before_sparse_ones(self): parallel = parallel_of(attn_dp=2, attn_tp=2) @@ -338,12 +344,17 @@ def test_a_dense_first_layer_before_sparse_ones(self): # The first a2a layer slices the residual it takes from a # dense layer; the next one takes it already sliced. self.assertEqual( - [layers[i]._steps.ffn.prepare.func for i in (1, 2)], + [ + layers[i]._steps.ffn.prepare.keywords["step"].func + for i in (1, 2) + ], [comm_ops._mlp_input_scatter] * 2, ) self.assertEqual( [ - layers[i]._steps.ffn.prepare.keywords["scatters_residual"] + layers[i] + ._steps.ffn.prepare.keywords["step"] + .keywords["scatters_residual"] for i in (1, 2) ], [True, False], @@ -354,10 +365,13 @@ def test_a_dense_first_layer_before_sparse_ones(self): comm.CommunicateSimpleFn._scattered_to_tp_attn_full, ) self.assertIs( - layers[3]._steps.ffn.prepare.func, comm_ops._mlp_input_dp_partial + layers[3]._steps.ffn.prepare.keywords["step"].func, + comm_ops._mlp_input_dp_partial, ) self.assertTrue( - layers[3]._steps.ffn.prepare.keywords["gathers_residual"] + layers[3] + ._steps.ffn.prepare.keywords["step"] + .keywords["gathers_residual"] ) def test_a_dense_mlp_on_every_rank(self): @@ -393,7 +407,8 @@ def test_a_dense_mlp_on_every_rank(self): self.assertEqual([self.declared(layer) for layer in layers], [True] * 4) for dense in layers[:2]: self.assertIs( - dense._steps.ffn.prepare.func, comm_ops._mlp_input_scatter + dense._steps.ffn.prepare.keywords["step"].func, + comm_ops._mlp_input_scatter, ) self.assertIsNone(dense._steps.ffn_output.group) for after_dense in layers[1:3]: @@ -423,9 +438,7 @@ def test_an_attention_that_gathers_its_slice_itself(self): ) with patch_communicator("_use_ag_after_qlora", True): layer = build(modes, parallel, a2a=True) - self.assertIs( - layer._steps.attention.input_move, comm.CommunicateSimpleFn._trivial - ) + self.assertIsNone(layer._steps.attention.input_move) def test_cp_without_prefill_cp(self): """No batch shards its tokens, so no CP steps. Under attention DP the @@ -485,6 +498,8 @@ class TestMhcOnTheDeclarations(CustomTestCase): other layer; only the residual operations the steps run are MHC's.""" def assert_step(self, step, func, communicator, **keywords): + if step.func is comm_ops._consumer_step: + step = step.keywords["step"] self.assertIs(step.func, func) residual = communicator._residual if "read" in step.keywords: @@ -502,7 +517,7 @@ def test_each_boundary_runs_the_shared_step(self): ( "attention TP 1", parallel_of(attn_dp=1, attn_tp=1), - comm_ops._mlp_input_norm, + comm_ops._read_input, comm.CommunicateSummableTensorPairFn._trivial, ), ( @@ -527,7 +542,10 @@ def test_each_boundary_runs_the_shared_step(self): self.assertIs(communicator._steps.ffn_output_move, output_move) # The fused add + RMSNorm kernels do not write hc_post. self.assertEqual(communicator._steps.ffn.fused, ()) - self.assertEqual(communicator._attn_input_fusions, ()) + self.assertEqual( + communicator._steps.attention.prepare.keywords["carried_fusions"], + (), + ) def test_the_move_back_over_dp_stays_with_the_layer(self): parallel = parallel_of(attn_dp=2, attn_tp=2) @@ -668,12 +686,10 @@ 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.attention.prepare.keywords["layer_input"], + steps.attention.prepare.keywords["step"].keywords["layer_input"], comm.tp_reduce_scatter if layer_id == 0 else None, ) - self.assertIs( - steps.attention.input_move, comm.CommunicateSimpleFn._trivial - ) + self.assertIsNone(steps.attention.input_move) self.assertIs( steps.attention.handoff, comm_ops._hand_scattered_input_to_attention ) @@ -829,10 +845,7 @@ def test_the_dense_layer_before_the_split_hands_on_the_attention_rows(self): if gathered: self.assertIs(before._steps.ffn_output_move.func, pair._gather) self.assertEqual(after.input_rows, attention) - self.assertIs( - after._steps.attention.input_move, - comm.CommunicateSimpleFn._trivial, - ) + self.assertIsNone(after._steps.attention.input_move) else: self.assertIs(before._steps.ffn_output_move, pair._trivial) self.assertEqual(after.input_rows, local) @@ -974,13 +987,17 @@ def run_steps(self, produced, *, force=False): TokenAxis.ATTN_TP_SCATTER: 2, } sides = sides_of(sizes) - steps, _ = comm_boundary._select_ffn_input( + steps, _, _ = comm_boundary._select_input_steps( produced, residual=sides.input_rows, residual_to=sides.ffn_residual_rows, need=sides.ffn, - force_layernorm_before_gather=force, + update=comm.ADD, fusions=(), + force_layernorm_before_gather=force, + residual_joins_sum=False, + cp_moves=None, + enters_stack=False, ) reduced = [] context = SimpleNamespace(attn_tp_size=2, attn_tp_rank=0) @@ -1039,6 +1056,21 @@ def test_an_owed_output_is_summed_once(self): # The stand-in all-reduce doubles the one rank's value. torch.testing.assert_close(residual, torch.full((1, HIDDEN), 17.0)) + def test_a_sum_left_only_for_some_batches_comes_with_the_value(self): + # Like a mixer exit: the steps take the output as complete, and a value + # that carries the sum is completed before them. + layout = sides_of( + { + TokenAxis.ATTN_DP: 2, + TokenAxis.ATTN_CP: 1, + TokenAxis.ATTN_TP_SCATTER: 2, + } + ).input_rows + asked = comm.StageOutput(layout, group=SumGroup.ATTN_TP) + reductions, hidden, residual = self.run_steps(asked) + self.assertEqual(reductions, 0) + torch.testing.assert_close(residual, torch.full((1, HIDDEN), 10.0)) + def test_declarations_the_steps_cannot_run_are_rejected(self): layout = sides_of( { @@ -1048,8 +1080,6 @@ def test_declarations_the_steps_cannot_run_are_rejected(self): } ).input_rows for produced in ( - # Leaves the attention-TP sum only when asked, like a mixer exit. - comm.StageOutput(layout, group=SumGroup.ATTN_TP), # Always leaves a sum over no group. comm.StageOutput(layout, always_leaves=True), # Owes a sum over a group the steps do not complete. @@ -1082,14 +1112,19 @@ def select(self, *, attn_dp, fusions): TokenAxis.ATTN_TP_SCATTER: 2, } ) - return comm_boundary._select_ffn_input( + steps, fused, _ = comm_boundary._select_input_steps( sides.attention_output, residual=sides.input_rows, residual_to=sides.ffn_residual_rows, need=sides.ffn, - force_layernorm_before_gather=False, + update=comm.ADD, fusions=fusions, + force_layernorm_before_gather=False, + residual_joins_sum=False, + cp_moves=None, + enters_stack=False, ) + return steps, fused def test_a_kernel_over_another_group_is_not_chosen(self): over_tp = self.fused(SumGroup.TP, True) @@ -1141,13 +1176,12 @@ class TestTheSequenceParallelRegion(CustomTestCase): SIZES = {TokenAxis.ATTN_DP: 1, TokenAxis.ATTN_CP: 1, TokenAxis.ATTN_TP_SCATTER: 2} SP_STEPS = comm.BoundarySteps( attention=comm.StageEntry( - prepare=comm_boundary._attention_input_step, + prepare=comm_ops._read_input, input_rows=comm.Layout(frozenset({TokenAxis.ATTN_TP_SCATTER})), - input_move=comm.CommunicateSimpleFn._trivial, handoff=comm_ops._hand_qkv_hook_its_input, ), ffn=comm.StageEntry( - prepare=comm_ops._mlp_input_norm, + prepare=comm_ops._read_input, input_rows=comm.Layout(frozenset({TokenAxis.ATTN_TP_SCATTER})), ), ffn_output=sequence_parallel_layer_sides(axis_sizes=SIZES).ffn_output, @@ -1158,18 +1192,17 @@ class TestTheSequenceParallelRegion(CustomTestCase): def unbound(self, steps): """``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.prepare.keywords, {"read": comm.NORM_READ, "update": comm.ADD} - ) - attention = steps.attention.prepare.keywords - self.assertIsNone(attention["layer_input"]) - self.assertIs(attention["read"], comm.NORM_QUANT_READ) - self.assertIs(attention["update"], comm.ADD) + ffn = steps.ffn.prepare.keywords["step"] + attention = steps.attention.prepare.keywords["step"] + for step, read in ((ffn, comm.NORM_READ), (attention, comm.NORM_QUANT_READ)): + self.assertIsNone(step.keywords["layer_input"]) + self.assertIs(step.keywords["read"], read) + self.assertIs(step.keywords["update"], comm.ADD) replace = msgspec.structs.replace return replace( steps, - ffn=replace(steps.ffn, prepare=steps.ffn.prepare.func), - attention=replace(steps.attention, prepare=steps.attention.prepare.func), + ffn=replace(steps.ffn, prepare=ffn.func), + attention=replace(steps.attention, prepare=attention.func), ) def test_the_region_declarations_choose_local_steps(self): @@ -1214,7 +1247,8 @@ def test_a_layer_under_sp_takes_both_sets_of_steps(self): ) self.assertEqual(self.unbound(communicator._sp_steps), self.SP_STEPS) self.assertIs( - communicator._steps.ffn.prepare.func, comm_ops._mlp_input_without_dp + communicator._steps.ffn.prepare.keywords["step"].func, + comm_ops._mlp_input_without_dp, ) def test_what_the_ffn_exit_reads_comes_from_the_batch_s_steps(self): @@ -1282,17 +1316,16 @@ def test_the_declarations_choose_the_steps(self): ) steps = comm_boundary._select_boundary_steps(sides) self.assertIs( - steps.attention.prepare.keywords["layer_input"], + steps.attention.prepare.keywords["step"].keywords["layer_input"], comm.tp_reduce_scatter, ) - self.assertIs( - steps.attention.input_move, comm.CommunicateSimpleFn._trivial - ) + self.assertIsNone(steps.attention.input_move) slice_ = comm.Layout(frozenset({TokenAxis.ATTN_TP_SCATTER})) self.assertEqual(steps.attention.input_rows, slice_) self.assertEqual(steps.ffn.input_rows, comm.Layout(frozenset())) self.assertIs( - steps.ffn.prepare.func, comm_ops._mlp_input_residual_into_sum + steps.ffn.prepare.keywords["step"].func, + comm_ops._mlp_input_residual_into_sum, ) self.assertIs( steps.ffn_output_move, @@ -1314,14 +1347,17 @@ def test_the_order_of_a_residual_on_the_slice_is_declared(self): (False, comm_ops._mlp_input_without_dp), ): with self.subTest(residual_joins_sum=joins): - steps, _ = comm_boundary._select_ffn_input( + steps, _, _ = comm_boundary._select_input_steps( owed, residual=local, residual_to=attention, need=comm.StageInput(attention), - force_layernorm_before_gather=False, + update=comm.ADD, fusions=(), + force_layernorm_before_gather=False, residual_joins_sum=joins, + cp_moves=None, + enters_stack=False, ) self.assertIs(getattr(steps, "func", steps), step) @@ -1331,13 +1367,17 @@ def test_a_complete_value_is_sliced_onto_each_rank(self): sizes = self.SIZES attention = comm.Layout.sharded_over(axis_sizes=sizes) local = comm.Layout(frozenset({TokenAxis.ATTN_TP_SCATTER})) - step, _ = comm_boundary._select_ffn_input( + step, _, _ = comm_boundary._select_input_steps( comm.StageOutput(attention), residual=attention, residual_to=local, need=comm.StageInput(local), - force_layernorm_before_gather=False, + update=comm.ADD, fusions=(), + force_layernorm_before_gather=False, + residual_joins_sum=False, + cp_moves=None, + enters_stack=False, ) self.assertIs(step.func, comm_ops._mlp_input_slice) hidden = torch.arange(4.0)[:, None].expand(4, HIDDEN).clone() @@ -1512,8 +1552,9 @@ def test_every_variant_takes_the_layers_fused_entries(self): ] self.assertEqual(len(variants), 2) for steps in variants: - bound = steps.attention.prepare.keywords - self.assertIs(bound["fusions"], fusions) + prepare = steps.attention.prepare + self.assertIs(prepare.keywords["carried_fusions"], fusions) + bound = prepare.keywords["step"].keywords self.assertIs(bound["enters_stack"], layer_id == 0) @@ -1546,8 +1587,14 @@ def test_a_dsa_cp_extend_leaves_its_sum_to_the_reduce_scatter(self): ) communicator = build(modes, parallel, dsa_cp=True, allow_reduce_scatter=True) cp, ordinary = communicator._cp_steps, communicator._steps - self.assertIs(cp.ffn.prepare.func, comm_ops._mlp_input_gather_attention_cp) - self.assertIs(cp.ffn.prepare.keywords["gather"].func, comm_ops._mlp_input_norm) + self.assertIs( + cp.ffn.prepare.keywords["step"].func, + comm_ops._mlp_input_gather_attention_cp, + ) + self.assertIs( + cp.ffn.prepare.keywords["step"].keywords["gather"].func, + comm_ops._read_input, + ) self.assertIs( cp.ffn_output_move, comm.CommunicateSummableTensorPairFn._reduce_scatter_over_cp, @@ -1555,9 +1602,9 @@ def test_a_dsa_cp_extend_leaves_its_sum_to_the_reduce_scatter(self): self.assertTrue(cp.ffn_output.leaves_for_reduce_scatter) self.assertFalse(cp.ffn_output.leaves_for_next_layer) # The dense layer before it ran on the same shard: nothing to move. - self.assertIs(cp.attention.input_move, comm.CommunicateSimpleFn._trivial) + self.assertIsNone(cp.attention.input_move) # Other batches hold every token on each CP rank: the MoE sums itself. - self.assertIs(ordinary.ffn.prepare.func, comm_ops._mlp_input_norm) + self.assertIs(ordinary.ffn.prepare.keywords["step"].func, comm_ops._read_input) self.assertIs( ordinary.ffn_output_move, comm.CommunicateSummableTensorPairFn._trivial ) @@ -1615,7 +1662,7 @@ def test_dsa_cp_dense_layers_run_on_their_shard(self): ) communicator = build(modes, parallel, dsa_cp=True, allow_reduce_scatter=True) for steps in (communicator._steps, communicator._cp_steps): - self.assertIs(steps.ffn.prepare.func, comm_ops._mlp_input_norm) + self.assertIs(steps.ffn.prepare.keywords["step"].func, comm_ops._read_input) self.assertIsNone(steps.ffn_output.group) self.assertIs( steps.ffn_output_move, comm.CommunicateSummableTensorPairFn._trivial @@ -1624,10 +1671,13 @@ def test_dsa_cp_dense_layers_run_on_their_shard(self): def test_a_cp_extend_gathers_over_cp_and_takes_its_chunk_back(self): communicator = build(layer_facts(1, 3), self.cp_parallel()) cp = communicator._cp_steps - self.assertIs(cp.ffn.prepare.func, comm_ops._mlp_input_gather_moe_cp) + self.assertIs( + cp.ffn.prepare.keywords["step"].func, comm_ops._mlp_input_gather_moe_cp + ) # Each rank completes its own chunk before the gather. self.assertIs( - cp.ffn.prepare.keywords["gather"].func, comm_ops._mlp_input_without_dp + cp.ffn.prepare.keywords["step"].keywords["gather"].func, + comm_ops._mlp_input_without_dp, ) self.assertIs( cp.ffn_output_move, @@ -1637,7 +1687,8 @@ def test_a_cp_extend_gathers_over_cp_and_takes_its_chunk_back(self): self.assertFalse(cp.ffn_output.leaves_for_next_layer) self.assertFalse(cp.ffn_output.leaves_for_reduce_scatter) self.assertIs( - communicator._steps.ffn.prepare.func, comm_ops._mlp_input_without_dp + communicator._steps.ffn.prepare.keywords["step"].func, + comm_ops._mlp_input_without_dp, ) self.assertIs( communicator._steps.ffn_output_move, @@ -1654,14 +1705,17 @@ def test_a_moe_on_its_own_cp_shard_hands_on_its_sum(self): ) for steps in (communicator._cp_steps, communicator._steps): # The attention's rows are the MoE's: nothing to gather or take back. - self.assertIs(steps.ffn.prepare.func, comm_ops._mlp_input_without_dp) + self.assertIs( + steps.ffn.prepare.keywords["step"].func, comm_ops._mlp_input_without_dp + ) self.assertIs( steps.ffn_output_move, comm.CommunicateSummableTensorPairFn._trivial ) self.assertTrue(steps.ffn_output.leaves_for_next_layer) dense = build(layer_facts(1, 3), parallel, allow_deferred_ffn_reduction=True) self.assertIs( - dense._cp_steps.ffn.prepare.func, comm_ops._mlp_input_gather_moe_cp + dense._cp_steps.ffn.prepare.keywords["step"].func, + comm_ops._mlp_input_gather_moe_cp, ) self.assertFalse(dense._cp_steps.ffn_output.leaves_for_next_layer) @@ -1672,7 +1726,7 @@ def test_the_fused_kernels_run_on_each_chunk(self): input_layernorm=FusableNorm(), post_attention_layernorm=FusableNorm(), ) - chunk = communicator._cp_steps.ffn.prepare.keywords["gather"] + chunk = communicator._cp_steps.ffn.prepare.keywords["step"].keywords["gather"] self.assertEqual( chunk.keywords["fusions"], (communicator._mlp_input_reduce_output_and_update_and_read_residual,), @@ -1774,13 +1828,19 @@ def test_under_attention_dp_one_dp_sum_gathers_both_axes(self): parallel = parallel_of(attn_dp=2, attn_tp=2, attn_cp=2, enable_prefill_cp=True) communicator = build(layer_facts(1, 3), parallel) cp, ordinary = communicator._cp_steps, communicator._steps - self.assertIs(cp.ffn.prepare.func, comm_ops._mlp_input_dp_partial) - self.assertTrue(cp.ffn.prepare.keywords["places_cp_shards"]) + self.assertIs( + cp.ffn.prepare.keywords["step"].func, comm_ops._mlp_input_dp_partial + ) + self.assertTrue(cp.ffn.prepare.keywords["step"].keywords["places_cp_shards"]) self.assertIs( cp.ffn_output_move, comm.CommunicateSummableTensorPairFn._take_back_cp_shard ) - self.assertIs(ordinary.ffn.prepare.func, comm_ops._mlp_input_dp_partial) - self.assertFalse(ordinary.ffn.prepare.keywords["places_cp_shards"]) + self.assertIs( + ordinary.ffn.prepare.keywords["step"].func, comm_ops._mlp_input_dp_partial + ) + self.assertFalse( + ordinary.ffn.prepare.keywords["step"].keywords["places_cp_shards"] + ) self.assertTrue(ordinary.returns_over_dp) for steps in (cp, ordinary): self.assertFalse(steps.ffn_output.leaves_for_next_layer) @@ -2461,7 +2521,7 @@ def update_and_read( def _stage(self, read, update): rows = comm.Layout.sharded_over(axis_sizes=self.SIZES) return comm.LayerStage( - reads=comm.InputRead.FFN, + kind=comm.StageKind.FFN, edges=comm.stage_edges( previous=None, stage=comm.StageDecl( diff --git a/test/registered/unit/layers/test_dp_attention_cp_gather.py b/test/registered/unit/layers/test_dp_attention_cp_gather.py index 08f5e4a9f329..cfadbbacf02f 100644 --- a/test/registered/unit/layers/test_dp_attention_cp_gather.py +++ b/test/registered/unit/layers/test_dp_attention_cp_gather.py @@ -276,7 +276,7 @@ def gather(rank, all_reduce): self.assertIsNotNone(r.communicator._cp_steps, "declared under CP") if ffn_input is not None: steps = r.communicator._batch_steps(r.forward_batch) - self.assertIs(steps.ffn.prepare.func, ffn_input) + self.assertIs(steps.ffn.prepare.keywords["step"].func, ffn_input) hidden_states, residual = rank_inputs(*rank) return r.communicator.prepare_mlp( hidden_states, residual, r.forward_batch 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 83d06ffe2096..85ac55b87265 100644 --- a/test/registered/unit/layers/test_layer_communicator_fusion_gate.py +++ b/test/registered/unit/layers/test_layer_communicator_fusion_gate.py @@ -51,19 +51,29 @@ def _steps(*, ffn_output=None, returns_over_dp=False, ffn_sum_is_movable=True): return comm.BoundarySteps( attention=comm.StageEntry( prepare=partial( - comm_boundary._attention_input_step, - layer_input=None, - fusions=(), - enters_stack=False, - read=comm.NORM_QUANT_READ, - update=comm.ADD, + comm_ops._consumer_step, + step=partial( + comm_ops._read_input, + layer_input=None, + enters_stack=False, + read=comm.NORM_QUANT_READ, + update=comm.ADD, + ), + carried_fusions=(), ), input_rows=comm.Layout(frozenset()), input_move=comm.CommunicateSimpleFn._trivial, handoff=comm_ops._hand_qkv_hook_its_input, ), ffn=comm.StageEntry( - prepare=comm_ops._mlp_input_norm, input_rows=Layout(frozenset()) + prepare=partial( + comm_ops._read_input, + layer_input=None, + enters_stack=False, + read=comm.NORM_READ, + update=comm.ADD, + ), + input_rows=Layout(frozenset()), ), ffn_output=ffn_output or StageOutput(Layout(frozenset()), group=SumGroup.TP), ffn_output_move=None diff --git a/test/registered/unit/layers/test_layernorm_sp.py b/test/registered/unit/layers/test_layernorm_sp.py index 12ea23d53acc..7546d4d547b5 100644 --- a/test/registered/unit/layers/test_layernorm_sp.py +++ b/test/registered/unit/layers/test_layernorm_sp.py @@ -194,12 +194,15 @@ def ordinary_steps(self, attention_input): return comm.BoundarySteps( attention=comm.StageEntry( prepare=partial( - comm_boundary._attention_input_step, - layer_input=None, - fusions=(), - enters_stack=False, - read=comm.NORM_QUANT_READ, - update=comm.ADD, + comm_ops._consumer_step, + step=partial( + comm_ops._read_input, + layer_input=None, + enters_stack=False, + read=comm.NORM_QUANT_READ, + update=comm.ADD, + ), + carried_fusions=(), ), input_rows=comm.Layout(frozenset()), input_move=attention_input, diff --git a/test/registered/unit/layers/test_postprocess_reduce_scatterv.py b/test/registered/unit/layers/test_postprocess_reduce_scatterv.py index 6496649cfb38..7243375a4693 100644 --- a/test/registered/unit/layers/test_postprocess_reduce_scatterv.py +++ b/test/registered/unit/layers/test_postprocess_reduce_scatterv.py @@ -18,7 +18,6 @@ LayerScatterModes, ScatterMode, ) -from sglang.srt.layers.communicator import boundary as comm_boundary from sglang.srt.layers.communicator import ops as comm_ops from sglang.srt.layers.moe import utils as moe_utils from sglang.srt.layers.moe.utils import ( @@ -39,19 +38,29 @@ def steps(*, ffn_output_move): return comm.BoundarySteps( attention=comm.StageEntry( prepare=partial( - comm_boundary._attention_input_step, - layer_input=None, - fusions=(), - enters_stack=False, - read=comm.NORM_QUANT_READ, - update=comm.ADD, + comm_ops._consumer_step, + step=partial( + comm_ops._read_input, + layer_input=None, + enters_stack=False, + read=comm.NORM_QUANT_READ, + update=comm.ADD, + ), + carried_fusions=(), ), input_rows=comm.Layout(frozenset()), input_move=comm.CommunicateSimpleFn._trivial, handoff=comm_ops._hand_qkv_hook_its_input, ), ffn=comm.StageEntry( - prepare=comm_ops._mlp_input_norm, input_rows=comm.Layout(frozenset()) + prepare=partial( + comm_ops._read_input, + layer_input=None, + enters_stack=False, + read=comm.NORM_READ, + update=comm.ADD, + ), + input_rows=comm.Layout(frozenset()), ), ffn_output=comm.StageOutput(comm.Layout(frozenset())), ffn_output_move=ffn_output_move, diff --git a/test/registered/unit/layers/test_prepare_attn_steps.py b/test/registered/unit/layers/test_prepare_attn_steps.py index b0bb9dd1e003..8c10d9e7cf3f 100644 --- a/test/registered/unit/layers/test_prepare_attn_steps.py +++ b/test/registered/unit/layers/test_prepare_attn_steps.py @@ -11,7 +11,6 @@ import torch from sglang.srt.layers import communicator as comm -from sglang.srt.layers.communicator import boundary as comm_boundary from sglang.srt.layers.communicator import ops as comm_ops from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.communicator_patch import patch_communicator @@ -62,19 +61,29 @@ def communicator(norm): c._steps = comm.BoundarySteps( attention=comm.StageEntry( prepare=partial( - comm_boundary._attention_input_step, - layer_input=None, - fusions=c._attn_input_fusions, - enters_stack=False, - read=comm.NORM_QUANT_READ, - update=comm.ADD, + comm_ops._consumer_step, + step=partial( + comm_ops._read_input, + layer_input=None, + enters_stack=False, + read=comm.NORM_QUANT_READ, + update=comm.ADD, + ), + carried_fusions=c._attn_input_fusions, ), input_rows=comm.Layout(frozenset()), input_move=lambda hidden_states, **_: hidden_states, handoff=comm_ops._hand_qkv_hook_its_input, ), ffn=comm.StageEntry( - prepare=comm_ops._mlp_input_norm, input_rows=comm.Layout(frozenset()) + prepare=partial( + comm_ops._read_input, + layer_input=None, + enters_stack=False, + read=comm.NORM_READ, + update=comm.ADD, + ), + input_rows=comm.Layout(frozenset()), ), ffn_output=comm.StageOutput(comm.Layout(frozenset())), ffn_output_move=comm.CommunicateSummableTensorPairFn._trivial, @@ -249,12 +258,15 @@ def test_an_owed_reduce_scatter_is_not_offered_to_a_fused_kernel(self): attention=msgspec.structs.replace( c._steps.attention, prepare=partial( - comm_boundary._attention_input_step, - layer_input=None, - fusions=(takes_anything,), - enters_stack=False, - read=comm.NORM_QUANT_READ, - update=comm.ADD, + comm_ops._consumer_step, + step=partial( + comm_ops._read_input, + layer_input=None, + enters_stack=False, + read=comm.NORM_QUANT_READ, + update=comm.ADD, + ), + carried_fusions=(takes_anything,), ), ), ) 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 6e3cbc991192..8a00033a6253 100644 --- a/test/registered/unit/lora/test_triton_dp_attention_unit.py +++ b/test/registered/unit/lora/test_triton_dp_attention_unit.py @@ -16,7 +16,7 @@ Layout, TokenAxis, ) -from sglang.srt.layers.communicator.boundary import _attention_input_step +from sglang.srt.layers.communicator.ops import _consumer_step, _read_input 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 ( @@ -244,12 +244,15 @@ def test_communicator_publishes_layout_at_each_transition( selected = SimpleNamespace( attention=SimpleNamespace( prepare=partial( - _attention_input_step, - layer_input=None, - fusions=(), - enters_stack=False, - read=NORM_QUANT_READ, - update=ADD, + _consumer_step, + step=partial( + _read_input, + layer_input=None, + enters_stack=False, + read=NORM_QUANT_READ, + update=ADD, + ), + carried_fusions=(), ), input_move=lambda hidden_states, **kwargs: hidden_states, handoff=lambda hidden_states, *args: hidden_states, 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 61f48f2e7db5..290d3fa4b972 100644 --- a/test/registered/unit/models/test_aux_capture_deferred_allreduce.py +++ b/test/registered/unit/models/test_aux_capture_deferred_allreduce.py @@ -7,7 +7,6 @@ from torch import nn from sglang.srt.layers import communicator as comm -from sglang.srt.layers.communicator import boundary as comm_boundary from sglang.srt.layers.communicator import ops as comm_ops from sglang.srt.model_executor.forward_batch_info import ForwardMode from sglang.srt.models.bailing_moe import BailingMoEModel @@ -52,19 +51,29 @@ def __init__(self, defer, return_topk=False): self.layer_communicator._steps = comm.BoundarySteps( attention=comm.StageEntry( prepare=partial( - comm_boundary._attention_input_step, - layer_input=None, - fusions=(), - enters_stack=False, - read=comm.NORM_QUANT_READ, - update=comm.ADD, + comm_ops._consumer_step, + step=partial( + comm_ops._read_input, + layer_input=None, + enters_stack=False, + read=comm.NORM_QUANT_READ, + update=comm.ADD, + ), + carried_fusions=(), ), input_rows=comm.Layout(frozenset()), input_move=comm.CommunicateSimpleFn._trivial, handoff=comm_ops._hand_qkv_hook_its_input, ), ffn=comm.StageEntry( - prepare=comm_ops._mlp_input_norm, input_rows=comm.Layout(frozenset()) + prepare=partial( + comm_ops._read_input, + layer_input=None, + enters_stack=False, + read=comm.NORM_READ, + update=comm.ADD, + ), + input_rows=comm.Layout(frozenset()), ), ffn_output=comm.StageOutput( comm.Layout(frozenset()),