diff --git a/docs/docs/developer_guide/layer_boundary.mdx b/docs/docs/developer_guide/layer_boundary.mdx index 8e7018989767..326443b6e12a 100644 --- a/docs/docs/developer_guide/layer_boundary.mdx +++ b/docs/docs/developer_guide/layer_boundary.mdx @@ -14,8 +14,8 @@ A boundary connects a **producer output** to a **consumer input**. An attention- | Responsibility | Owner | | --- | --- | | Attention, projections, activation functions, expert routing, EP dispatch/combine | Model compute modules | -| How a producer contributes to the residual | Producer's `StageUpdate` | -| How a consumer reads the residual | Consumer's `StageRead` and norm | +| How a producer contributes to the residual | Producer's `ResidualUpdate` | +| How a consumer reads the residual | Consumer's `ResidualReadout` and norm | | Supported row transitions and reduction ordering | Bound boundary paths | | Actual outstanding work and residual tensors | Per-forward `ResidualStream` | | Hardware-specific fused completion | Fusion providers | @@ -31,7 +31,7 @@ EP all-to-all is an operation inside MoE computation, not an input preparation s - `ATTN_DP`: different attention data-parallel token sets. - `ATTN_CP`: context-parallel (CP) token partitions for an active CP batch. -- `ATTN_TP_SCATTER`: token slices distributed across attention-TP ranks. +- `ATTN_TP`: token slices distributed across attention-TP ranks. Size-one axes are omitted. Layout is not a `DTensor` placement or a general cross-mesh redistribution engine. It does not describe arbitrary hidden-dimension sharding. The actual row ordering, padding, and per-rank counts come from batch metadata and the paired movement operations. @@ -43,14 +43,14 @@ A declaration states how its compute cooperates with the boundary's reduction de | Contract or state | Meaning | | --- | --- | -| `ProducerReduction.PARTIAL` | Attention only (the `declare_attn()` default): compute always returns a contribution requiring a sum and publishes it with `finish()`. Build its output projection so that it never reduces (for example `reduce_results=False`). `declare_ffn()` rejects it | -| `ProducerReduction.SCOPED` | Compute obeys the exit's reduction decision and publishes through `exit()`: the `declare_ffn()` default, and `declare_attn(reduction=ProducerReduction.SCOPED)` for a mixer | -| `ProducerReduction.LOCAL_TAIL` | FFN only: compute adds a replicated component after its internal sum. The boundary never defers that sum to the next layer, but reduce-scatter permissions remain, so compute that skips its sum for `mlp_reduce_scatter` must add the component exactly once. `declare_attn()` rejects it | -| `StageOutput.leaves_for_*` | Construction-time permission to select a particular handoff | +| `ProducerReduction.ALWAYS_PARTIAL` | Attention only (the `declare_attn()` default): compute always returns a contribution requiring a sum and publishes it with `finish()`. Build its output projection so that it never reduces (for example `reduce_results=False`). `declare_ffn()` rejects it | +| `ProducerReduction.EXIT_SCOPED` | Compute obeys the exit's reduction decision and publishes through `exit()`: the `declare_ffn()` default, and `declare_attn(reduction=ProducerReduction.EXIT_SCOPED)` for a mixer | +| `ProducerReduction.TAIL_AFTER_SUM` | FFN only: compute adds a replicated component after its internal sum. The boundary never defers that sum to the next layer, but reduce-scatter permissions remain, so compute that skips its sum for `mlp_reduce_scatter` must add the component exactly once. `declare_attn()` rejects it | +| `OutputContract.may_*` | Construction-time permission to leave the sum to a particular later step | | `Contribution.owed` | Work this specific output still requires | | `OwedOutput` | Model-facing handle preventing ordinary tensor use of that unfinished output | -`UnreducedOutput` is an internal adapter form. `ResidualStream.leave()` converts it into an owned contribution and an opaque model handle. `HandoffOutput` represents producer-specific work such as a deferred MoE finalize; its `complete()` supplies an unfused completion path. A handoff is consumed by the next layer's input or, at the end of the stack, by a final norm that accepts it. +`UnreducedOutput` is an internal adapter form. `ResidualStream.record()` converts it into an owned contribution and an opaque model handle. `DeferredFinalize` represents producer-specific work such as a deferred MoE finalize; its `complete()` supplies an unfused completion path. A handoff is consumed by the next layer's input or, at the end of the stack, by a final norm that accepts it. ### Residual stream @@ -58,7 +58,7 @@ A declaration states how its compute cooperates with the boundary's reduction de | State | Stored data | Next preparation | | --- | --- | --- | -| Initial | Neither residual nor contribution | Initialize residual through the consumer's `enter` | +| Initial | Neither residual nor contribution | Initialize residual through the consumer's `init_residual` | | Written | Residual, no pending contribution | Read input without applying another update | | Pending | Contribution, optionally an existing residual | Complete owed work, apply the producer update, then read | @@ -70,19 +70,19 @@ Communication completion and residual update are distinct steps. The `residual.b | --- | --- | | `complete_output(hidden_states, forward_batch)` | Finishes outstanding communication/finalize work and leaves the update pending | | `fold(hidden_states, forward_batch)` | Completes a plain contribution and folds it into the residual, leaving a written stream | -| `written(hidden_states, forward_batch)` | Replaces the stream with a written one whose residual is `hidden_states`. Use it after compute that produced the full residual itself. It discards any pending contribution, so call `fold()` first when one may be pending | +| `set_written(hidden_states, forward_batch)` | Replaces the stream with a written one whose residual is `hidden_states`. Use it after compute that produced the full residual itself. It discards any pending contribution, so call `fold()` first when one may be pending | | `take_output(hidden_states, forward_batch)` | Releases the stream at a terminal exit whose output is already written, such as Nemotron-H MTP after its internal fold and norm. It rejects a pending contribution | -Always retain the returned tensor or handle after each of these calls. A stack-entry boundary accepts a written stream and skips its `enter` operation. +Always retain the returned tensor or handle after each of these calls. A stack-entry boundary accepts a written stream and skips its `init_residual` operation. ### Residual update and input read -The producer owns `StageUpdate`; the consumer owns `StageRead` and its norm. The actual update travels with the contribution. Bound paths use declared update capabilities to select valid ordering: +The producer owns `ResidualUpdate`; the consumer owns `ResidualReadout` and its norm. The actual update travels with the contribution. Bound paths use declared update capabilities to select valid ordering: - A plain add can participate in compatible fused add+norm kernels or add the residual on one rank before summation. - A nonlinear update cannot generally move before reduction. The post operation of manifold-constrained hyper-connections (MHC) is one such case. -- `at_producer` requests an update at the producer exit. A deferred update must guarantee that its parameters and state outlive that producer. -- `StageRead.before_gather` preserves reads that must execute on the source rows before a DP gather. +- `applied_at_exit` requests an update at the producer exit. A deferred update must guarantee that its parameters and state outlive that producer. +- `ResidualReadout.reads_before_dp_gather` preserves reads that must execute on the source rows before a DP gather. A read returns `(compute_input, residual)`. `update_and_read()` receives the producer's update together with the read, so an implementation can fuse them: the plain add+norm read does, and keeps its FP32 accumulation instead of normalizing a separately rounded snapshot. Fusion is not guaranteed. An implementation such as MHC may run the update and the read as separate steps, provided it preserves its existing rounding order. @@ -111,7 +111,7 @@ self.attn_boundary, self.ffn_boundary = make_stages( ) ``` -`previous` describes the actual external producer. You do not need an executable object from the previous layer or another pipeline rank. For a sparse or custom preceding stage, build `previous` with the same arguments the producing layer uses for its own declaration, including `next_sparse`, for example `declare_ffn(sparse=prev_sparse, next_sparse=this_sparse)`; the two-batch overlap (TBO) handoff rows derive from that pair. +`previous` describes the actual external producer. You do not need an executable object from the previous layer or another pipeline rank. For a sparse or custom preceding stage, build `previous` with the same arguments the producing layer uses for its own declaration, including `next_layer_sparse`, for example `declare_ffn(sparse=prev_sparse, next_layer_sparse=this_sparse)`; the two-batch overlap (TBO) exit rows derive from that pair. `terminal=True` marks the end of the model's layer stack, not every Python decoder layer. A pipeline partition that has more layers downstream is not terminal merely because its local module list ends. A NextN/multi-token-prediction (MTP) draft module is its own stack: its layer enters with `previous=None` and is terminal even when it reuses the target's decoder class. In that case, use the draft's layer count, for example `terminal=layer_id == (1 if is_nextn else config.num_hidden_layers) - 1`. A terminal stage leaves no work for a following layer, except the final-norm finalize handoff described in [Enter and leave the layer stack](#enter-and-leave-the-layer-stack). @@ -123,9 +123,9 @@ For a single independently constructed stage, use `make_attn_stage()` or `make_f For heterogeneous layer stacks, declare the actual neighbouring stage kinds; the framework does not require an alternating attention/FFN pattern. `make_stages(..., following=declaration)` describes the external consumer after the last local stage, just as `previous` describes the producer before the first. -For a single-stage `ProducerReduction.SCOPED` mixer, `following` decides whether the mixer leaves its attention-TP sum to the next stage: always before an FFN, and per batch before another mixer. The next stage makes the same decision from its own `previous`, so both must describe the same neighbours. Otherwise a pipeline handoff reduces the mixer's output twice: `to_pp()` exports it complete and the receiver's `from_pp()` declares the sum again. +For a single-stage `ProducerReduction.EXIT_SCOPED` mixer, `following` decides whether the mixer leaves its attention-TP sum to the next stage: always before an FFN, and per batch before another mixer. The next stage makes the same decision from its own `previous`, so both must describe the same neighbours. Otherwise a pipeline handoff reduces the mixer's output twice: `to_pp()` exports it complete and the receiver's `from_pp()` declares the sum again. -`models/nemotron_h_utils.py` (used by `models/nemotron_h.py`) builds each Nemotron-H stage as `make_stages((decl, norm), previous=..., following=..., terminal=...)`, with mixers declared as `declare_attn(reduction=ProducerReduction.SCOPED, gathers_tp_input=False)`. +`models/nemotron_h_utils.py` (used by `models/nemotron_h.py`) builds each Nemotron-H stage as `make_stages((decl, norm), previous=..., following=..., terminal=...)`, with mixers declared as `declare_attn(reduction=ProducerReduction.EXIT_SCOPED, gathers_attn_tp_input=False)`. ### More than two stages and branches @@ -142,7 +142,7 @@ first_ffn, second_attn, second_ffn = make_stages( This fragment assumes the source boundary and the three norms already exist. `prepared_from` means the source has already performed its read. Enter through `branch_input(source, hidden_states, forward_batch)`, not `prepare()`, to move that input and fork the residual stream without normalizing twice. The first stage's norm is therefore never applied; LongCat passes the MoE stage's norm there. -The branch adapter also provides `branch_output()` to place a complete contribution on handoff rows and `merge_branch()` to combine it with another branch's output. These methods encode row movement and residual ownership; the model still determines its compute schedule. Branch transport currently supports only ordinary token rows: construction rejects context-parallel, input-scattered, and sequence-parallel variants. See `python/sglang/srt/models/longcat_flash.py` for a complete integration. +The branch adapter also provides `branch_output()` to place a complete contribution on the branch's exit rows and `merge_branch()` to combine it with another branch's output. These methods encode row movement and residual ownership; the model still determines its compute schedule. Branch transport currently supports only ordinary token rows: construction rejects context-parallel, input-scattered, and sequence-parallel variants. See `python/sglang/srt/models/longcat_flash.py` for a complete integration. ### Execute a layer @@ -166,11 +166,11 @@ hidden_states = output.finish(hidden_states) This example uses Qwen3's compute member name `mlp`; boundary construction does not require renaming compute classes or checkpoint parameters. -`prepare()` returns input in the consumer read's format, which can include quantized tensors. An attention with a `PARTIAL` output contract uses `finish()` to register its contribution. An FFN or single-stage mixer uses `exit()` to publish flags during compute, then calls the returned object's `finish()` once. The exit context restores the previous runtime flags even when compute raises; do not call `finish()` on a failed computation. +`prepare()` returns input in the consumer read's format, which can include quantized tensors. An attention with an `ALWAYS_PARTIAL` output contract uses `finish()` to register its contribution. An FFN or single-stage mixer uses `exit()` to publish flags during compute, then calls the returned object's `finish()` once. The exit context restores the previous runtime flags even when compute raises; do not call `finish()` on a failed computation. Compute cooperates through the flags the scope publishes on `get_forward()`: -- `FfnExit` publishes `fuse_mlp_allreduce` (skip the output all-reduce), `mlp_reduce_scatter` (leave the sum to the boundary's reduce-scatter), and `defer_moe_finalize` (a MoE may return a `HandoffOutput`). +- `FfnExit` publishes `fuse_mlp_allreduce` (skip the output all-reduce), `mlp_reduce_scatter` (leave the sum to the boundary's reduce-scatter), and `defer_moe_finalize` (a MoE may return a `DeferredFinalize`). - `MixerExit` publishes `fuse_mlp_allreduce` and exposes the same decision as `skips_reduction`. Before an FFN, the mixer always leaves its attention-TP sum to the FFN's input, including under attention DP. Before another mixer it may defer the sum to that mixer's input, but never under attention DP. `RowParallelLinear` and the MoE output reduction honour these flags automatically; a custom reduction reads them from the exit object. Returning a handoff while `defer_moe_finalize` is false fails in `finish()`. @@ -188,20 +188,20 @@ residual_batch.start(forward_batch) hidden_states = self.embed_tokens(input_ids) for layer in self.layers: hidden_states = layer(positions, hidden_states, forward_batch) -hidden_states = residual_batch.norm( +hidden_states = residual_batch.final_norm( hidden_states, forward_batch, self.norm, skip_empty=True ) ``` -`skip_empty=True` still completes owed work on an empty batch, which may be a collective that other ranks join, and skips only the norm kernel. `norm()` and `to_pp()` release the batch's stream after they hand on its output. +`skip_empty=True` still completes owed work on an empty batch, which may be a collective that other ranks join, and skips only the norm kernel. `final_norm()` and `to_pp()` release the batch's stream after they hand on its output. This is a non-pipeline stack fragment. Pipeline reception uses the first boundary's `from_pp(tensors, forward_batch)` to reconstruct a stream; the sender uses `residual_batch.to_pp()`. Dynamic completion work is resolved before transport. By default, `to_pp()` preserves a statically declared partial sum for the receiver's incoming contract to reconstruct and complete. Do not call `complete_output()` before that handoff: it would complete a sum that the receiver still expects to perform. -A terminal FFN normally completes its output before the final norm. The exception is a MoE finalize handoff. When the stage's fusion provider can defer the finalize at the terminal layer (for CuteDSL, `install_cutedsl_fusion(..., terminal_finalize=True)`), the exit may leave a `HandoffOutput` for the final norm. Pass the handle straight to `residual_batch.norm(hidden_states, forward_batch, self.norm, handoff_norm=service)`, where `service` is the object `install_cutedsl_fusion()` returned; its `finalize(handoff, residual, gamma)` fuses finalize, all-reduce, residual add, and norm. The fused path reads `layernorm.gemma_weight`, so it currently expects a `GemmaRMSNorm` final norm. Without `handoff_norm`, `norm()` first completes the handoff unfused, and `capture_output` is then allowed. With `handoff_norm`, `norm()` rejects `capture_output` when a handoff arrives. Calling `complete_output()` first runs the producer's unfused finalize, and `snapshot()` raises on a handoff. See `python/sglang/srt/models/qwen3_5.py`. +A terminal FFN normally completes its output before the final norm. The exception is a MoE finalize handoff. When the stage's fusion provider can defer the finalize at the terminal layer (for CuteDSL, `install_cutedsl_fusion(..., terminal_finalize=True)`), the exit may leave a `DeferredFinalize` for the final norm. Pass the handle straight to `residual_batch.final_norm(hidden_states, forward_batch, self.norm, finalize_norm=service)`, where `service` is the object `install_cutedsl_fusion()` returned; its `finalize(handoff, residual, gamma)` fuses finalize, all-reduce, residual add, and norm. The fused path reads `layernorm.gemma_weight`, so it currently expects a `GemmaRMSNorm` final norm. Without `finalize_norm`, `final_norm()` first completes the handoff unfused, and `capture` is then allowed. With `finalize_norm`, `final_norm()` rejects `capture` when a handoff arrives. Calling `complete_output()` first runs the producer's unfused finalize, and `snapshot()` raises on a handoff. See `python/sglang/srt/models/qwen3_5.py`. -`model_forward_stages()` in `python/sglang/srt/batch_overlap/two_batch_overlap.py` gives each TBO microbatch its own stream: it splits with `ResidualStream.arrive()` and merges the results afterwards. Do not share one mutable stream between interleaved microbatches or hold an `exit()` scope open across a TBO yield point. +`model_forward_stages()` in `python/sglang/srt/batch_overlap/two_batch_overlap.py` gives each TBO microbatch its own stream: it splits with `ResidualStream.from_handoff()` and merges the results afterwards. Do not share one mutable stream between interleaved microbatches or hold an `exit()` scope open across a TBO yield point. -In the operation-scheduled TBO path, the FFN runs without an `exit()` scope. No skip flags are published, so compute returns a fully reduced output. `boundary.postprocess(hidden_states, forward_batch)` applies any output transform, moves that output back to the layer's handoff rows (a plain scatter under attention DP), and leaves nothing owed. +In the operation-scheduled TBO path, the FFN runs without an `exit()` scope. No skip flags are published, so compute returns a fully reduced output. `boundary.finish_complete_output(hidden_states, forward_batch)` applies any output transform, moves that output back to the layer's exit rows (a plain scatter under attention DP), and leaves nothing owed. ### Auxiliary capture and extra contributions @@ -210,17 +210,17 @@ Choose the access method according to what the consumer needs: | Operation | Effect on main output | | --- | --- | | `residual_batch.snapshot()` | Produces an owned plain residual snapshot; a sum can run on a copy without consuming the main contribution | -| Attention `boundary.prepare(..., capture_output=collector.capture)` | Completes the main output's owed communication (unfused; the add+norm kernel still runs), then captures the updated residual this read produces, moved back to the producer's rows. At stack entry it captures the input before `enter`; with `post_residual_addition` it captures output plus residual before the addition. FFN `prepare()` does not accept it | -| Attention `boundary.prepare(..., captured_last_layer_outputs=collector)` | Captures the same updated residual on the attention's compute-input rows, after its input gather | +| Attention `boundary.prepare(..., capture=collector.capture)` | Completes the main output's owed communication (unfused; the add+norm kernel still runs), then captures the updated residual this read produces, moved back to the producer's rows. At stack entry it captures the input before `init_residual`; with `post_residual_addition` it captures output plus residual before the addition. FFN `prepare()` does not accept it | +| Attention `boundary.prepare(..., capture_gathered=collector)` | Captures the same updated residual on the attention's compute-input rows, after its input gather | | `boundary.capture_output()` | Preserves its explicit capture convention: completes dynamic work on the main contribution, but snapshots a statically declared sum on a copy; returns both updated handle and snapshot | | `residual_batch.add_to_output(..., extra)` | Completes the contribution before adding extra once, keeping the residual update pending | -| `residual_batch.norm(..., capture_output=collector.capture)` | Final norm and capture of its updated residual | +| `residual_batch.final_norm(..., capture=collector.capture)` | Final norm and capture of its updated residual | Snapshots currently require a plain residual update. Producer-specific finalize handoffs need an explicit main-output capture adapter. A separate BF16 snapshot followed by norm can change results relative to a fused add+norm kernel's FP32 accumulation. Capture from the existing update/read operation when that ordering matters. -A capture callback has the signature `capture(value, *, owned=False)`, as `AuxHiddenStateList.capture` does. With `owned=False` the value may alias the mutable residual or a reusable communication buffer, and the collector must copy it before retaining it. The boundary passes `owned=True` only for storage nothing else will write: a capture move that gathers rows (`StageEntry.capture_move_allocates`), a residual sum computed only for the capture, or a residual that the following local stage's input path certifies it leaves untouched. `make_stages()` copies that certificate into the producer's entry at construction; no plan is consulted during forward. `residual_batch.norm(..., capture_output=...)` always passes a borrowed value. +A capture callback has the signature `capture(value, *, owned=False)`, as `AuxHiddenStateList.capture` does. With `owned=False` the value may alias the mutable residual or a reusable communication buffer, and the collector must copy it before retaining it. The boundary passes `owned=True` only for storage nothing else will write: a capture move that gathers rows (`EntryPath.capture_move_allocates`), a residual sum computed only for the capture, or a residual that the following local stage's input path certifies it leaves untouched. `make_stages()` copies that certificate into the producer's entry at construction; no plan is consulted during forward. `residual_batch.final_norm(..., capture=...)` always passes a borrowed value. ### Choose a reduction policy @@ -240,14 +240,14 @@ Required attention and MoE collectives are unaffected. These are permissions, no Before you open a pull request for a new model, check each item: -- Every layer passes the real producer as `previous`, including `sparse`/`next_sparse`, and `previous=None` only where the stack starts. +- Every layer passes the real producer as `previous`, including `sparse`/`next_layer_sparse`, and `previous=None` only where the stack starts. - `terminal` is set only on the last layer of the model's stack, using the draft's own layer count for NextN/MTP modules. - A single-stage mixer or FFN also passes `following`, and both neighbours describe each other consistently. -- The model forward calls `residual_batch.start()` before the first layer, and leaves the stack through `residual_batch.norm()`, `to_pp()`, or `take_output()`. +- The model forward calls `residual_batch.start()` before the first layer, and leaves the stack through `residual_batch.final_norm()`, `to_pp()`, or `take_output()`. - Pipeline reception calls `from_pp(tensors, forward_batch)` on the first local boundary. - Split prefill starts the stream only in the segment that runs the first layer and completes it only in the final segment. - Extra contributions such as deepstack embeddings go through `residual_batch.add_to_output()`, never an in-place add on the layer output. -- Auxiliary capture uses `prepare(..., capture_output=...)` or `norm(..., capture_output=...)` with a callback that honours `owned`. +- Auxiliary capture uses `prepare(..., capture=...)` or `final_norm(..., capture=...)` with a callback that honours `owned`. Common errors and their usual cause: @@ -258,7 +258,7 @@ Common errors and their usual cause: | `write the residual update before taking the final output` | `take_output()` ran while a contribution was still pending; fold or norm it first | | `no stage boundary path for the active batch` | The active batch variant (CP, input-scattered, sequence parallel) is not supported by this stage; add a server-argument check that rejects the configuration | | `a prepared branch must enter through branch_input, not prepare` | A `prepared_from` stage was entered with `prepare()` | -| `PARTIAL is not supported for ffn stages` / `LOCAL_TAIL is not supported for attention stages` | The declaration's `reduction` does not match the stage kind | +| `ALWAYS_PARTIAL is not supported for ffn stages` / `TAIL_AFTER_SUM is not supported for attention stages` | The declaration's `reduction` does not match the stage kind | | `snapshot requires a plain residual update` | A snapshot was taken on a nonlinear (for example MHC) update | ## Boundary internals @@ -275,7 +275,7 @@ flowchart TD B --> I["prepare()
complete,
update, read,
move rows"] I -->|"FFN or mixer"| X["exit()
choose
completion,
publish flags"] X --> M["Model
computation"] - I -->|"PARTIAL"| M + I -->|"ALWAYS_PARTIAL"| M M --> O["finish()
record the
contribution"] O --> S ``` @@ -283,11 +283,11 @@ flowchart TD Construction and execution have different lifetimes: 1. **At model initialization**, `declare_attn()` and `declare_ffn()` describe compute behavior. `make_stages()` connects those declarations and binds each stage's own norm and hooks. -2. **During construction**, factories resolve token rows and reduction contracts for supported batch variants. `make_boundary()` binds consumer work; `make_output_boundary()` binds output transport. `StagePlan` retains those paths. +2. **During construction**, factories resolve token rows and reduction contracts for supported batch variants. `bind_entry()` binds consumer work; `bind_exit()` binds output transport. `StagePlan` retains those paths. 3. **At each forward**, batch facts select an existing path. An exit can make a batch-dependent completion decision, but it does not rebuild the declarations. The decision publishes compute flags and supplies the matching completion action. 4. **Across stages**, `ForwardBatch.residual_stream` carries residual state and the producer's actual contribution. Stages do not store forward tensors on their reusable plans. -The resolved `StageDecl` is a lower-level type than `StageDeclaration`: one `StageInput` and one `StageOutput` after topology and batch variant have been applied. An `EdgeDecl` combines a producer output, consumer requirement, and the residual's source/destination layouts. +The resolved `StageContract` is a lower-level type than `StageDeclaration`: one `InputContract` and one `OutputContract` after topology and batch variant have been applied. An `EdgeContract` combines a producer output, consumer requirement, and the residual's source/destination layouts. ### Module map @@ -304,8 +304,8 @@ Paths in this table are relative to `python/sglang/srt/layers/layer_boundary/`. | `ops.py` | Collective and token-row movement primitives | | `stage.py` | Model-facing `StageBoundary` methods | | `exit.py` | Select output completion and publish scoped compute flags | -| `output.py` | Internal reduction/finalize handoffs and output transforms | -| `residual/__init__.py` | `StageUpdate`/`StageRead` protocols and `LayerResidual` | +| `output.py` | Unreduced outputs, deferred finalizes and output transforms | +| `residual/__init__.py` | `ResidualUpdate`/`ResidualReadout` protocols and `LayerResidualOps` | | `residual/stream.py` | Per-forward state and ownership of outstanding work | | `residual/batch.py` | Model-stack entry, final norm, pipeline-parallel (PP) transport, snapshots, and extra contributions | | `residual/access.py` | Tensor access and materialization helpers | @@ -327,13 +327,13 @@ An output decision is made once. Its flags control whether compute performs redu A backend fusion provider is passed per stage as the `fusions` option, `(declaration, norm, {"fusions": provider})`; give a layer's attention and FFN stages the same provider. The boundary calls it as follows: -- `attention_input(plan)` and `ffn_input(plan)` return ordered consumer candidates, tried before the built-in all-reduce fusion. -- `can_defer_finalize(plan, forward_batch)` decides whether an FFN may return a `HandoffOutput`. It is called on every FFN exit that has a provider. +- `attn_input_fusions(plan)` and `ffn_input_fusions(plan)` return ordered consumer candidates, tried before the built-in all-reduce fusion. +- `can_defer_finalize(plan, forward_batch)` decides whether an FFN may return a `DeferredFinalize`. It is called on every FFN exit that has a provider. - `can_defer_all_reduce(plan, forward_batch)` lets a LoRA or `SGLANG_SHARED_EXPERT_TP1` FFN leave its all-reduce to the next input. Such an FFN defers only without an expert all-to-all backend or quantized communication, when its sum is over the TP group, and when the provider, FlashInfer fusion, or AITER fusion accepts the batch. Other FFNs defer when the post-expert sum is a single all-reduce. `install_cutedsl_fusion()` and `prepare_cutedsl_fusion()` find providers through the decoder layer's `attn_boundary` and `ffn_boundary` attributes. -Consumer fusion candidates are ordered. A candidate returning `None` must do so before modifying inputs or starting a collective, so fallback remains valid. `FusedMlpInput` declares the sum group it completes with residual add and norm. The consumer only tries compatible candidates for the declared update/read capabilities. +Consumer fusion candidates are ordered. A candidate returning `None` must do so before modifying inputs or starting a collective, so fallback remains valid. `FfnInputFusion` declares the sum group it completes with residual add and norm. The consumer only tries compatible candidates for the declared update/read capabilities. An `OutputTransform` declares contribution processing before residual update. Its explicit `before_reduce_scatter` setting preserves a supported implementation's order; it is not permission to commute arbitrary operations across a reduction. @@ -344,7 +344,7 @@ Producer GEMM plus reduce-scatter fusion is not a generic implemented API here. - **Attention/DeepSeek Sparse Attention (DSA) adapters:** retain `AttnTpContext` and its input-scattered/latent-input hooks. They are not a general asynchronous dependency scheduler. - **CP adapters:** pair gather, return-to-local-rows, and eligible reduce-scatter operations using the same token mapping. - **LoRA adapter:** publishes the selected compute input's rows for LoRA execution. -- **TBO and branch adapters:** manage handoffs that cannot be represented as a normal sequential prepare call. +- **TBO and branch adapters:** manage layer exits and entries that cannot be represented as a normal sequential prepare call. ### Extend and validate the module @@ -353,7 +353,7 @@ Place changes according to their responsibility: 1. Add model compute facts to a declaration only when they describe compute behavior. Put user policy and unsupported configuration checks in server-argument handling. 2. Add new update/read semantics in `residual/`; declare only capabilities the implementation really supports. 3. Add a token transition in boundary selection and movement primitives, with its padding and row-order contract. Do not duplicate that logic in each model. -4. Put backend completion in `fusions/` and specialized handoff behavior in `adapters/`. +4. Put backend completion in `fusions/` and specialized exit and entry behavior in `adapters/`. 5. Keep models on stage factories and boundary methods. Avoid reading a neighbour's execution plan or reconstructing a collective decision during forward. Focused tests live in `test/registered/unit/layer_boundary/`. Cover numerical output and residual state, collective group/row correctness, and supported fallback behavior for the changed path. Relevant examples include `test_ffn_exit.py`, `test_reduction_fusion.py`, and `test_terminal_stages.py`. A declaration-only assertion cannot establish that a new parallel path moves and reduces the correct values. diff --git a/python/sglang/srt/batch_overlap/two_batch_overlap.py b/python/sglang/srt/batch_overlap/two_batch_overlap.py index dc9854b0985d..8f9750ea9e6d 100644 --- a/python/sglang/srt/batch_overlap/two_batch_overlap.py +++ b/python/sglang/srt/batch_overlap/two_batch_overlap.py @@ -19,7 +19,7 @@ tbo_split_moves, ) from sglang.srt.layers.layer_boundary.residual import batch as residual_batch -from sglang.srt.layers.layer_boundary.residual.access import finish_layer_stack +from sglang.srt.layers.layer_boundary.residual.access import export_output from sglang.srt.layers.layer_boundary.residual.stream import ResidualStream from sglang.srt.layers.moe import ( get_deepep_mode, @@ -954,9 +954,9 @@ def model_forward_stages( if not enable_tbo: return execute_operations(inputs, strategy.operations)["hidden_states"] - stream = residual_batch.current(forward_batch) + stream = residual_batch.stream_of(forward_batch) pending = stream.pending - hidden_states, residual = stream.finish(hidden_states) + hidden_states, residual = stream.export(hidden_states) inputs["hidden_states"] = hidden_states parts = _model_forward_tbo_split_inputs( **inputs, @@ -964,7 +964,7 @@ def model_forward_stages( layer_input_rows=layers[0].attn_boundary.incoming_residual_rows, ) for part in parts: - part["hidden_states"], child_stream = ResidualStream.arrive( + part["hidden_states"], child_stream = ResidualStream.from_handoff( part["hidden_states"], part.pop("residual"), pending.update if pending is not None else None, @@ -987,7 +987,7 @@ def model_forward_stages( delta_stages=[0, strategy.tbo_delta_stages], ) for output in outputs: - output["residual"] = residual_batch.current(output["forward_batch"]) + output["residual"] = residual_batch.stream_of(output["forward_batch"]) hidden_states, forward_batch.residual_stream = _model_forward_tbo_merge_outputs( *outputs, original_len ) @@ -1110,7 +1110,7 @@ def _model_forward_tbo_merge_outputs(output_a, output_b, original_len): assert pending_a.update is pending_b.update update = pending_a.update for output in (output_a, output_b): - output["hidden_states"], output["residual"] = finish_layer_stack( + output["hidden_states"], output["residual"] = export_output( output["hidden_states"], output["residual"], output["forward_batch"] ) @@ -1133,7 +1133,7 @@ def _handle_key(name): hidden, residual = _handle_key("hidden_states"), _handle_key("residual") return ( - ResidualStream.arrive(hidden, residual, update) + ResidualStream.from_handoff(hidden, residual, update) if has_stream else (hidden, residual) ) diff --git a/python/sglang/srt/hardware_backend/npu/modules/deepseek_v2_attention_mla_npu.py b/python/sglang/srt/hardware_backend/npu/modules/deepseek_v2_attention_mla_npu.py index 768cec6d13c9..7b2ec443bc24 100644 --- a/python/sglang/srt/hardware_backend/npu/modules/deepseek_v2_attention_mla_npu.py +++ b/python/sglang/srt/hardware_backend/npu/modules/deepseek_v2_attention_mla_npu.py @@ -35,7 +35,7 @@ def forward_mha_prepare_npu( hidden_states: torch.Tensor, forward_batch: "ForwardBatch", zero_allocator: "BumpAllocator", - input_on_attention_tp_slices: bool, + input_on_attn_tp_slices: bool, ): if m.q_lora_rank is not None: q, latent_cache = ( @@ -63,7 +63,7 @@ def forward_mha_prepare_npu( else: q = m.q_a_layernorm(q) - if _use_ag_after_qlora and input_on_attention_tp_slices: + if _use_ag_after_qlora and input_on_attn_tp_slices: q = scattered_to_tp_attn_full(q, forward_batch) latent_cache = scattered_to_tp_attn_full(latent_cache, forward_batch) q = m.q_b_proj(q)[0].view(-1, m.num_local_heads, m.qk_head_dim) @@ -157,7 +157,7 @@ def forward_mla_prepare_npu( hidden_states: torch.Tensor, forward_batch: "ForwardBatch", zero_allocator: "BumpAllocator", - input_on_attention_tp_slices: bool, + input_on_attn_tp_slices: bool, ): if is_mla_preprocess_enabled(): if not hasattr(m, "mla_preprocess"): @@ -190,7 +190,7 @@ def forward_mla_prepare_npu( q_lora = None if m.q_lora_rank is not None: qkv_latent = get_attn_tp_context().fetch_qkv_latent() - if _use_ag_after_qlora and input_on_attention_tp_slices: + if _use_ag_after_qlora and input_on_attn_tp_slices: q, latent_cache = qkv_latent.split( [m.q_lora_rank, m.kv_lora_rank + m.qk_rope_head_dim], dim=-1, @@ -362,7 +362,7 @@ def forward_dsa_prepare_npu( hidden_states: torch.Tensor, forward_batch: "ForwardBatch", zero_allocator: "BumpAllocator", - input_on_attention_tp_slices: bool, + input_on_attn_tp_slices: bool, prev_topk_indices: torch.Tensor = None, ): dynamic_scale = None @@ -396,7 +396,7 @@ def forward_dsa_prepare_npu( ) # overlap qk norm q = m.q_a_layernorm(q) - if _use_ag_after_qlora and input_on_attention_tp_slices: + if _use_ag_after_qlora and input_on_attn_tp_slices: q = scattered_to_tp_attn_full(q, forward_batch) latent_cache = scattered_to_tp_attn_full(latent_cache, forward_batch) q_lora = q.clone() # required for topk_indices @@ -485,7 +485,7 @@ def forward_dsa_prepare_npu( positions, forward_batch, m.layer_id, - input_on_attention_tp_slices, + input_on_attn_tp_slices, dynamic_scale, ) else: diff --git a/python/sglang/srt/layers/attention/dsa/dsa_npu_indexer.py b/python/sglang/srt/layers/attention/dsa/dsa_npu_indexer.py index b607e72218a0..da832b28d72e 100644 --- a/python/sglang/srt/layers/attention/dsa/dsa_npu_indexer.py +++ b/python/sglang/srt/layers/attention/dsa/dsa_npu_indexer.py @@ -60,7 +60,7 @@ def forward_npu( positions: torch.Tensor, forward_batch: ForwardBatch, layer_id: int, - input_on_attention_tp_slices: bool = False, + input_on_attn_tp_slices: bool = False, dynamic_scale: torch.Tensor = None, ) -> torch.Tensor: if get_attn_backend().forward_metadata.seq_lens_cpu_int is None: @@ -140,7 +140,7 @@ def forward_npu( k_proj = self.wk(x)[0] # [b, s, 7168] @ [7168, 128] = [b, s, 128] k = self.k_norm(k_proj) - if _use_ag_after_qlora and input_on_attention_tp_slices: + if _use_ag_after_qlora and input_on_attn_tp_slices: k = scattered_to_tp_attn_full(k, forward_batch) k_pe, k_nope = torch.split( k, @@ -284,7 +284,7 @@ def forward_npu( torch.npu.current_stream().wait_event(q_rope_event) if envs.SGLANG_NPU_USE_MULTI_STREAM.get(): torch.npu.current_stream().wait_event(weights_event) - if _use_ag_after_qlora and input_on_attention_tp_slices: + if _use_ag_after_qlora and input_on_attn_tp_slices: weights = scattered_to_tp_attn_full(weights, forward_batch) block_table = get_attn_backend().forward_metadata.block_tables if ( diff --git a/python/sglang/srt/layers/flashinfer_comm_fusion.py b/python/sglang/srt/layers/flashinfer_comm_fusion.py index 2cf739d0837e..c99da881e2d5 100644 --- a/python/sglang/srt/layers/flashinfer_comm_fusion.py +++ b/python/sglang/srt/layers/flashinfer_comm_fusion.py @@ -941,7 +941,7 @@ def can_use_flashinfer_allreduce( # Size checks stay last: they read the token dim, which is symbolic under # Dynamo, so statically-off configs must short-circuit before reaching them - # (same ordering rule as apply_flashinfer_allreduce_fusion). + # (same ordering rule as flashinfer_ar_fusion_applies). token_num, hidden_dim = input_.shape # MNNVL hard-fails instead of falling back when the width is not float4-aligned diff --git a/python/sglang/srt/layers/layer_boundary/__init__.py b/python/sglang/srt/layers/layer_boundary/__init__.py index d5d0ca7a853d..dd857d8c169d 100644 --- a/python/sglang/srt/layers/layer_boundary/__init__.py +++ b/python/sglang/srt/layers/layer_boundary/__init__.py @@ -18,27 +18,27 @@ get_attn_tp_context, ) from sglang.srt.layers.layer_boundary.boundary import ( - make_boundary, - make_output_boundary, + bind_entry, + bind_exit, tbo_split_moves, ) from sglang.srt.layers.layer_boundary.construction import ( BatchVariant, - StageEdges, + VariantEdges, ) from sglang.srt.layers.layer_boundary.contracts import ( - EdgeDecl, - FusedMlpInput, - HandoffRows, + EdgeContract, + EntryPath, + ExitRows, + FfnInputFusion, + InputContract, + OutputContract, ProducerReduction, - StageDecl, - StageEntry, - StageInput, + StageContract, StageKind, - StageOutput, - StageSteps, + StagePath, ) -from sglang.srt.layers.layer_boundary.exit import FfnCompletion, FfnExit, MixerExit +from sglang.srt.layers.layer_boundary.exit import ExitDecision, FfnExit, MixerExit from sglang.srt.layers.layer_boundary.factories import ( declare_attn, declare_ffn, @@ -50,9 +50,9 @@ Layout, SumGroup, TokenAxis, - enable_moe_dense_fully_dp, - moe_cp_gathers_sparse_moe_input, - sparse_moe_gathers_over_moe_cp, + batch_gathers_over_moe_cp, + is_dense_ffn_fully_dp, + moe_gathers_over_moe_cp, token_axis_sizes, ) from sglang.srt.layers.layer_boundary.ops import ( @@ -60,17 +60,17 @@ tp_reduce_scatter, ) from sglang.srt.layers.layer_boundary.output import ( - HandoffOutput, + DeferredFinalize, UnreducedOutput, - reduce_output, + complete_owed, ) -from sglang.srt.layers.layer_boundary.residual import LayerResidual +from sglang.srt.layers.layer_boundary.residual import LayerResidualOps from sglang.srt.layers.layer_boundary.residual.add_norm import ( - ADD, FUSE_ALLREDUCE_MAX_BATCH_SIZE, - NORM_QUANT_READ, - NORM_READ, - PLAIN_RESIDUAL, + NORM_QUANT_READOUT, + NORM_READOUT, + PLAIN_ADD, + PLAIN_RESIDUAL_OPS, ) from sglang.srt.layers.layer_boundary.residual.mhc import ( MHCState, @@ -82,40 +82,40 @@ "make_attn_stage", "make_ffn_stage", "make_stages", - "ADD", + "PLAIN_ADD", "AttentionInputs", - "StageSteps", - "EdgeDecl", - "HandoffRows", + "StagePath", + "EdgeContract", + "ExitRows", "ProducerReduction", "FUSE_ALLREDUCE_MAX_BATCH_SIZE", - "FfnCompletion", + "ExitDecision", "FfnExit", - "FusedMlpInput", - "HandoffOutput", - "LayerResidual", + "FfnInputFusion", + "DeferredFinalize", + "LayerResidualOps", "Layout", "MHCState", "MixerExit", - "NORM_QUANT_READ", - "NORM_READ", - "PLAIN_RESIDUAL", - "StageDecl", - "StageEntry", - "StageInput", + "NORM_QUANT_READOUT", + "NORM_READOUT", + "PLAIN_RESIDUAL_OPS", + "StageContract", + "EntryPath", + "InputContract", "StageKind", - "StageOutput", + "OutputContract", "SumGroup", "TokenAxis", "UnreducedOutput", - "enable_moe_dense_fully_dp", + "is_dense_ffn_fully_dp", "get_attn_tp_context", - "make_boundary", - "make_output_boundary", - "moe_cp_gathers_sparse_moe_input", + "bind_entry", + "bind_exit", + "batch_gathers_over_moe_cp", "move_rows", - "reduce_output", - "sparse_moe_gathers_over_moe_cp", + "complete_owed", + "moe_gathers_over_moe_cp", "tbo_split_moves", "token_axis_sizes", "tp_reduce_scatter", diff --git a/python/sglang/srt/layers/layer_boundary/adapters/attention.py b/python/sglang/srt/layers/layer_boundary/adapters/attention.py index 3767408fadea..798e6d93c2ee 100644 --- a/python/sglang/srt/layers/layer_boundary/adapters/attention.py +++ b/python/sglang/srt/layers/layer_boundary/adapters/attention.py @@ -26,7 +26,7 @@ is_dp_attention_enabled, ) from sglang.srt.layers.layer_boundary.layout import ( - enable_moe_dense_fully_dp, + is_dense_ffn_fully_dp, ) from sglang.srt.layers.moe import get_moe_a2a_backend from sglang.srt.model_executor.cuda_graph_config import ( @@ -107,7 +107,7 @@ def init_context(self, q_lora_rank, is_dsa, is_mhc=False): and get_parallel().tp_size > 1 and not is_dp_attention_enabled() and get_moe_a2a_backend().is_none() - and not enable_moe_dense_fully_dp() + and not is_dense_ffn_fully_dp() and not check_cuda_graph_backend(Phase.PREFILL, Backend.TC_PIECEWISE) and get_spec().speculative_algorithm != "EAGLE3" ) @@ -168,7 +168,7 @@ def get_attn_tp_context(): return ATTN_TP_CONTEXT -def _redistribute_from_attn_tp_shards(tensor: torch.Tensor) -> torch.Tensor: +def attn_tp_gather(tensor: torch.Tensor) -> torch.Tensor: gathered = get_local_dp_buffer( get_parallel().attn_tp_group, hidden_size=tensor.shape[-1] ) @@ -176,6 +176,6 @@ def _redistribute_from_attn_tp_shards(tensor: torch.Tensor) -> torch.Tensor: return gathered -def _redistribute_to_attn_tp_shards(tensor: torch.Tensor) -> torch.Tensor: +def attn_tp_slice(tensor: torch.Tensor) -> torch.Tensor: parallel = get_parallel() return tensor.tensor_split(parallel.attn_tp_size)[parallel.attn_tp_rank] diff --git a/python/sglang/srt/layers/layer_boundary/adapters/branch.py b/python/sglang/srt/layers/layer_boundary/adapters/branch.py index c1b13ec2b744..4098d603a772 100644 --- a/python/sglang/srt/layers/layer_boundary/adapters/branch.py +++ b/python/sglang/srt/layers/layer_boundary/adapters/branch.py @@ -37,8 +37,8 @@ def branch_input( """The FFN input and residual that ``source``'s boundary read for its own FFN, for this layer's FFN, which branches from the same input: moved to the rows this FFN needs and its residual's rows.""" - rows, residual_rows, _ = source._branch_rows(forward_batch) - to, residual_to, _ = plan._branch_rows(forward_batch) + rows, residual_rows, _ = source.branch_rows(forward_batch) + to, residual_to, _ = plan.branch_rows(forward_batch) if stream.pending is not None: raise RuntimeError("a branch must start from a prepared stage input") residual = stream.residual @@ -56,7 +56,7 @@ def branch_output( """This layer's complete FFN output as a branch's contribution, which adds to the layer's output without writing the residual: moved to the rows the layer hands on.""" - rows, _, to = plan._branch_rows(forward_batch) + rows, _, to = plan.branch_rows(forward_batch) return move_rows(hidden_states, rows, to, forward_batch) @@ -71,15 +71,15 @@ def merge_branch( """A contribution from ``branch_output`` summed with what ``source``'s layer hands on, ``hidden_states`` and ``residual``, moved to the rows this layer hands on.""" - _, _, rows = source._branch_rows(forward_batch) - _, _, to = plan._branch_rows(forward_batch) + _, _, rows = source.branch_rows(forward_batch) + _, _, to = plan.branch_rows(forward_batch) stream.check(hidden_states) if stream.pending is None: raise RuntimeError("branch merge requires a pending producer contribution") update = stream.pending.update - hidden_states, residual = stream.finish(hidden_states) + hidden_states, residual = stream.export(hidden_states) hidden_states = move_rows(hidden_states, rows, to, forward_batch) residual = move_rows(residual, rows, to, forward_batch) output = contribution + hidden_states stream.write(residual) - return stream.leave(output, update), stream + return stream.record(output, update), stream diff --git a/python/sglang/srt/layers/layer_boundary/adapters/context_parallel.py b/python/sglang/srt/layers/layer_boundary/adapters/context_parallel.py index 0b74f4290b52..3cae0b06877b 100644 --- a/python/sglang/srt/layers/layer_boundary/adapters/context_parallel.py +++ b/python/sglang/srt/layers/layer_boundary/adapters/context_parallel.py @@ -49,7 +49,7 @@ def maybe_prefetch_next_full_attention_kv( prefetch_kv_buffer(next_full_attention_layer_id) -def dsa_cp_gather_hidden_states(hidden_states: torch.Tensor): +def attn_cp_gather(hidden_states: torch.Tensor): attn_dp_size = get_parallel().attn_dp_size attn_tp_size = get_parallel().attn_tp_size assert attn_dp_size == 1 and attn_tp_size == 1 @@ -61,7 +61,7 @@ def dsa_cp_gather_hidden_states(hidden_states: torch.Tensor): return hidden_states -def dsa_cp_reduce_scatter_hidden_states(hidden_states: torch.Tensor): +def attn_cp_reduce_scatter(hidden_states: torch.Tensor): attn_dp_size = get_parallel().attn_dp_size attn_tp_size = get_parallel().attn_tp_size assert attn_dp_size == 1 and attn_tp_size == 1 diff --git a/python/sglang/srt/layers/layer_boundary/adapters/lora.py b/python/sglang/srt/layers/layer_boundary/adapters/lora.py index b2648ef2ec4e..a2ab2076d96c 100644 --- a/python/sglang/srt/layers/layer_boundary/adapters/lora.py +++ b/python/sglang/srt/layers/layer_boundary/adapters/lora.py @@ -17,12 +17,12 @@ from sglang.srt.runtime_context import LoRABatchLayout, get_forward -def publish_attention(enabled): +def publish_attn_lora_rows(enabled): if enabled: get_forward().set("lora_batch_layout", LoRABatchLayout.DP_LOCAL) -def publish_ffn(enabled, input_rows): +def publish_ffn_lora_rows(enabled, input_rows): if enabled: get_forward().set( "lora_batch_layout", diff --git a/python/sglang/srt/layers/layer_boundary/adapters/overlap.py b/python/sglang/srt/layers/layer_boundary/adapters/overlap.py index 312fd537454b..fdb154cf20d1 100644 --- a/python/sglang/srt/layers/layer_boundary/adapters/overlap.py +++ b/python/sglang/srt/layers/layer_boundary/adapters/overlap.py @@ -1,26 +1,25 @@ """Row requirements at the entrance to two-batch overlap.""" -from sglang.srt.layers.layer_boundary.contracts import HandoffRows -from sglang.srt.layers.layer_boundary.layout import enable_moe_dense_fully_dp +from sglang.srt.layers.layer_boundary.contracts import ExitRows +from sglang.srt.layers.layer_boundary.layout import is_dense_ffn_fully_dp from sglang.srt.runtime_context import get_exec -def tbo_handoff(sparse, next_sparse): +def tbo_exit_rows(sparse, next_layer_sparse): """Record a potential split between a dense layer and a sparse layer.""" - return HandoffRows.TBO_SPLIT if not sparse and next_sparse else None + return ExitRows.TBO_SPLIT if not sparse and next_layer_sparse else None -def resolve_handoff_rows(requirement): - """Bind a handoff requirement with the construction-time parallel config. +def resolve_exit_rows(requirement): + """Bind an exit-row requirement to the construction-time parallel config. Dense-local output must reach attention rows before the TBO split. Moving that gather after the split changes padding and collective sizes. """ - if requirement is HandoffRows.TBO_SPLIT: + if requirement is ExitRows.TBO_SPLIT: return ( - HandoffRows.ATTENTION - if enable_moe_dense_fully_dp() - and get_exec().overlap.enable_two_batch_overlap + ExitRows.ATTENTION + if is_dense_ffn_fully_dp() and get_exec().overlap.enable_two_batch_overlap else None ) return requirement diff --git a/python/sglang/srt/layers/layer_boundary/boundary.py b/python/sglang/srt/layers/layer_boundary/boundary.py index 2bedf1eaaba9..b870f9e79292 100644 --- a/python/sglang/srt/layers/layer_boundary/boundary.py +++ b/python/sglang/srt/layers/layer_boundary/boundary.py @@ -21,47 +21,47 @@ from sglang.srt.layers.layer_boundary.contracts import ( CpMoves, - EdgeDecl, - FusedMlpInput, - StageInput, - StageOutput, + EdgeContract, + FfnInputFusion, + InputContract, + OutputContract, ) from sglang.srt.layers.layer_boundary.layout import ( Layout, SumGroup, TokenAxis, - _gathers_over_attention_cp, + _cp_gathers_over_attn_cp, _same_ranks, _sum_group, token_axis_sizes, ) from sglang.srt.layers.layer_boundary.ops import ( - gather_attention_tp, - identity_output, + attn_cp_reduce_scatter_output, + attn_cp_take_back_output, + attn_tp_gather_input, + attn_tp_slice_output, + dp_cp_take_back_output, + keep_output, + moe_cp_take_back_output, move_rows, - output_on_residual_shard, - reduce_scatter_over_cp, - scatter_moe_cp_output, - scatter_output, - take_back_attention_cp_shard, - take_back_cp_shard, + residual_slice_output, tp_reduce_scatter, tp_slice, - update_and_gather, + update_attn_tp_gather_output, ) from sglang.srt.layers.layer_boundary.prepare import ( - _consumer_step, - _dispatch_consumer, - _mlp_input_dp_partial, - _mlp_input_dp_replicate, - _mlp_input_gather_attention_cp, - _mlp_input_gather_moe_cp, - _mlp_input_on_residual_shard, - _mlp_input_residual_into_sum, - _mlp_input_scatter, - _mlp_input_slice, - _mlp_input_without_dp, - _read_input, + _attn_tp_reduce_scatter_update_read, + _attn_tp_slice_update_read, + _dispatch_by_update, + _dp_gather_sum_read, + _reduce_update_read, + _reduce_update_read_dp_gather, + _run_entry, + _then_attn_cp_gather, + _then_moe_cp_gather, + _tp_reduce_scatter_update_read_gather, + _tp_sum_with_residual_read, + _update_read, ) from sglang.srt.runtime_context import get_parallel @@ -75,11 +75,11 @@ def tbo_split_moves(layer_input_rows: Layout) -> Tuple[Callable, Callable]: ) if layer_input_rows == attention: - return identity_output, identity_output - if layer_input_rows.sharded - attention.sharded == {TokenAxis.ATTN_TP_SCATTER}: + return keep_output, keep_output + if layer_input_rows.sharded - attention.sharded == {TokenAxis.ATTN_TP}: # Each rank's slice: write the residual in and gather over attention # TP, then take the slice of each half. - return update_and_gather, scatter_output + return update_attn_tp_gather_output, attn_tp_slice_output raise NotImplementedError(f"{layer_input_rows=}") @@ -87,26 +87,26 @@ def _cp_moves() -> CpMoves: """DSA and MLA CP gather equal shards over the attention-CP group and can complete a sum over it. GQA prefill CP gathers blocks padded to the longest over the MoE-CP group and takes back only a complete output.""" - if _gathers_over_attention_cp(): + if _cp_gathers_over_attn_cp(): return CpMoves( - gather=_mlp_input_gather_attention_cp, - take_back=take_back_attention_cp_shard, - reduce_scatter=reduce_scatter_over_cp, + gather=_then_attn_cp_gather, + take_back=attn_cp_take_back_output, + reduce_scatter=attn_cp_reduce_scatter_output, reduce_scatter_group=lambda: get_parallel().attn_cp_group, ) return CpMoves( - gather=_mlp_input_gather_moe_cp, - take_back=scatter_moe_cp_output, + gather=_then_moe_cp_gather, + take_back=moe_cp_take_back_output, ) -class Boundary(msgspec.Struct, frozen=True): +class EdgeBinding(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 + edge: EdgeContract # The consumer's half: completing what the input owes, the add and the # norm, and the moves onto the rows it needs that come before the read. prepare: Optional[Callable] = None @@ -141,21 +141,22 @@ def capture_move(self) -> Optional[Callable]: ) -def input_rows(edge: EdgeDecl) -> Layout: +def input_rows(edge: EdgeContract) -> Layout: """Consumer rows, retaining the axes it gathers inside its computation.""" return Layout( - edge.need.layout.sharded | (edge.residual_to.sharded & edge.need.gathers_itself) + edge.need.layout.sharded + | (edge.residual_to.sharded & edge.need.gathered_by_compute) ) -def make_boundary( - edge: EdgeDecl, +def bind_entry( + edge: EdgeContract, *, - fusions: Tuple["FusedMlpInput", ...] = (), + fusions: Tuple["FfnInputFusion", ...] = (), carried_fusions: Tuple[Callable, ...] = (), cp_moves: Optional[CpMoves] = None, enters_stack: bool = False, -) -> Boundary: +) -> EdgeBinding: """Bind the consumer half of an edge at construction time. Args: @@ -170,15 +171,15 @@ def make_boundary( Multiple update capabilities select among preconstructed paths; a single capability needs no runtime dispatcher. """ - capabilities = edge.update_capabilities + capabilities = edge.arriving_plain_add if not capabilities: if edge.produced.update is None: raise ValueError("an arrival must declare its update capabilities") - capabilities = (edge.produced.update.adds_plainly,) + capabilities = (edge.produced.update.is_plain_add,) paths = { - capability: _bind_consumer( + capability: _bind_entry_path( edge, - adds_plainly=capability, + is_plain_add=capability, fusions=fusions, carried_fusions=carried_fusions, cp_moves=cp_moves, @@ -195,21 +196,21 @@ def make_boundary( first, preserves_residual=None, prepare=partial( - _dispatch_consumer, + _dispatch_by_update, paths={capability: path.prepare for capability, path in paths.items()}, ), ) -def _bind_consumer( - edge: EdgeDecl, +def _bind_entry_path( + edge: EdgeContract, *, - adds_plainly: bool, - fusions: Tuple["FusedMlpInput", ...] = (), + is_plain_add: bool, + fusions: Tuple["FfnInputFusion", ...] = (), carried_fusions: Tuple[Callable, ...] = (), cp_moves: Optional[CpMoves] = None, enters_stack: bool = False, -) -> Boundary: +) -> EdgeBinding: """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 @@ -221,26 +222,26 @@ def _bind_consumer( first stage. The steps read only this edge's declarations, never what the producer chose for a batch.""" # A fused kernel runs the add and the norm itself. - plain = adds_plainly and edge.need.read.norms_plainly - step, input_move = _select_input_steps( + plain = is_plain_add and edge.need.read.is_plain_norm + step, input_move = _select_entry_step( edge.produced, residual=edge.residual, residual_to=edge.residual_to, need=edge.need, - adds_plainly=adds_plainly, + is_plain_add=is_plain_add, fusions=fusions if plain else (), residual_joins_sum=edge.residual_joins_sum, cp_moves=cp_moves, enters_stack=enters_stack, ) completed_step = None - if edge.produced.always_leaves: - completed_step, _ = _select_input_steps( - msgspec.structs.replace(edge.produced, always_leaves=False), + if edge.produced.always_partial: + completed_step, _ = _select_entry_step( + msgspec.structs.replace(edge.produced, always_partial=False), residual=edge.residual, residual_to=edge.residual_to, need=edge.need, - adds_plainly=adds_plainly, + is_plain_add=is_plain_add, fusions=(), residual_joins_sum=False, cp_moves=cp_moves, @@ -250,12 +251,12 @@ def _bind_consumer( # again (e.g. MHC expansion). Both alternatives are bound at construction. written_step = None if enters_stack: - written_step, _ = _select_input_steps( + written_step, _ = _select_entry_step( edge.produced, residual=edge.residual, residual_to=edge.residual_to, need=edge.need, - adds_plainly=adds_plainly, + is_plain_add=is_plain_add, fusions=fusions if plain else (), residual_joins_sum=edge.residual_joins_sum, cp_moves=cp_moves, @@ -263,10 +264,10 @@ def _bind_consumer( ) preserves_residual = None if ( - edge.produced.always_leaves + edge.produced.always_partial and edge.produced.group is SumGroup.ATTN_TP and isinstance(step, partial) - and step.func is _mlp_input_without_dp + and step.func is _reduce_update_read and not step.keywords["gathers_residual"] ): selected = step.keywords["fusions"] @@ -276,15 +277,15 @@ def _bind_consumer( preserves_residual = next( f.preserves_residual for f in fusions if f.run == selected[0] ) - return Boundary( + return EdgeBinding( edge, preserves_residual=preserves_residual, prepare=partial( - _consumer_step, + _run_entry, step=step, - adds_plainly=adds_plainly, + is_plain_add=is_plain_add, carried_fusions=carried_fusions if plain else (), - expected_sum=edge.produced.group if edge.produced.always_leaves else None, + expected_sum=edge.produced.group if edge.produced.always_partial else None, completed_step=completed_step, written_step=written_step, ), @@ -292,10 +293,8 @@ def _bind_consumer( ) -def make_output_boundary( - edge: EdgeDecl, *, cp_moves: Optional[CpMoves] = None -) -> Boundary: - """Bind the producer half of a layer/branch handoff. +def bind_exit(edge: EdgeContract, *, cp_moves: Optional[CpMoves] = None) -> EdgeBinding: + """Bind the producer half of a layer or branch exit. Args: edge: Output contract and destination/residual rows. Deferred updates @@ -308,24 +307,24 @@ def make_output_boundary( transport. The receiver binds its read independently. """ update = edge.produced.update - if not getattr(update, "at_producer", False): - if not getattr(update, "can_defer_across_layers", False): + if not getattr(update, "applied_at_exit", False): + if not getattr(update, "outlives_layer", False): raise NotImplementedError( "a deferred update must guarantee its lifetime across layers" ) - if not update.adds_plainly and get_parallel().pp_size > 1: + if not update.is_plain_add and get_parallel().pp_size > 1: raise NotImplementedError( "pipeline boundaries require a plain add or a producer-written residual" ) if edge.need.layout != edge.residual_to: raise NotImplementedError(f"{edge=}") - returns_over_dp, output_move, completes_sum = _select_ffn_output_move( + returns_over_dp, output_move, completes_sum = _select_exit_move( edge.produced, residual=edge.residual, to=edge.residual_to, cp_moves=cp_moves, ) - return Boundary( + return EdgeBinding( edge, output_move=output_move, returns_over_dp=returns_over_dp, @@ -333,14 +332,14 @@ def make_output_boundary( ) -def _select_input_steps( - produced: StageOutput, +def _select_entry_step( + produced: OutputContract, *, residual: Layout, residual_to: Layout, - need: StageInput, - adds_plainly: bool, - fusions: Tuple[FusedMlpInput, ...], + need: InputContract, + is_plain_add: bool, + fusions: Tuple[FfnInputFusion, ...], residual_joins_sum: bool, cp_moves: Optional[CpMoves], enters_stack: bool, @@ -357,17 +356,17 @@ def _select_input_steps( 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: + owes = produced.group if produced.always_partial else None + if produced.always_partial and owes is None: raise NotImplementedError(f"{produced=}") - gathered = produced.layout.sharded - need.layout.sharded - need.gathers_itself + gathered = produced.layout.sharded - need.layout.sharded - need.gathered_by_compute sliced = need.layout.sharded - produced.layout.sharded if sliced: # Each attention-TP rank takes its own slice: the reduce-scatter # completes the attention-TP sum and slices in one collective; a # complete value is only sliced. if ( - sliced != {TokenAxis.ATTN_TP_SCATTER} + sliced != {TokenAxis.ATTN_TP} or gathered or owes not in (None, SumGroup.ATTN_TP) or residual_to != need.layout @@ -376,13 +375,15 @@ def _select_input_steps( raise NotImplementedError(f"{produced=} {residual=} {need=}") return ( partial( - _mlp_input_scatter if owes is SumGroup.ATTN_TP else _mlp_input_slice, + _attn_tp_reduce_scatter_update_read + if owes is SumGroup.ATTN_TP + else _attn_tp_slice_update_read, scatters_residual=residual != residual_to, read=read, ), None, ) - if residual_to.sharded - produced.layout.sharded == {TokenAxis.ATTN_TP_SCATTER}: + if residual_to.sharded - produced.layout.sharded == {TokenAxis.ATTN_TP}: 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). @@ -390,7 +391,9 @@ def _select_input_steps( raise NotImplementedError(f"{produced=} {residual=} {need=}") return ( partial( - _mlp_input_on_residual_shard, read=read, reduces=owes is not None + _tp_reduce_scatter_update_read_gather, + read=read, + reduces=owes is not None, ), None, ) @@ -399,14 +402,14 @@ def _select_input_steps( if ( owes not in (None, SumGroup.TP) or residual.sharded - or TokenAxis.ATTN_TP_SCATTER not in need.gathers_itself - or residual_to.sharded != {TokenAxis.ATTN_TP_SCATTER} + or TokenAxis.ATTN_TP not in need.gathered_by_compute + or residual_to.sharded != {TokenAxis.ATTN_TP} ): raise NotImplementedError(f"{produced=} {residual=} {need=}") return ( partial( - _read_input, - layer_input=tp_reduce_scatter if owes is not None else tp_slice, + _update_read, + pre_move=tp_reduce_scatter if owes is not None else tp_slice, enters_stack=enters_stack, read=read, ), @@ -416,19 +419,19 @@ def _select_input_steps( # 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, _ = _select_input_steps( + on_chunk, _ = _select_entry_step( produced, residual=residual, residual_to=residual_to, - need=StageInput(produced.layout, read=read), - adds_plainly=adds_plainly, + need=InputContract(produced.layout, read=read), + is_plain_add=is_plain_add, fusions=fusions, residual_joins_sum=residual_joins_sum, cp_moves=cp_moves, enters_stack=enters_stack, ) return (partial(cp_moves.gather, gather=on_chunk), None) - if gathered == {TokenAxis.ATTN_TP_SCATTER}: + if gathered == {TokenAxis.ATTN_TP}: # A complete input on each rank's slice, gathered over attention TP # once it is read. if ( @@ -439,12 +442,12 @@ def _select_input_steps( raise NotImplementedError(f"{produced=} {residual=} {need=}") return ( partial( - _read_input, - layer_input=None, + _update_read, + pre_move=None, enters_stack=enters_stack, read=read, ), - gather_attention_tp, + attn_tp_gather_input, ) if ( residual_to != produced.layout @@ -455,7 +458,7 @@ def _select_input_steps( {TokenAxis.ATTN_DP, TokenAxis.ATTN_CP}, ) or residual.sharded - residual_to.sharded - not in (frozenset(), {TokenAxis.ATTN_TP_SCATTER}) + not in (frozenset(), {TokenAxis.ATTN_TP}) ): raise NotImplementedError(f"{produced=} {residual=} {need=}") # A residual arriving on each rank's slice is gathered back first. @@ -465,7 +468,7 @@ def _select_input_steps( if gathers_residual: return ( partial( - _mlp_input_without_dp, + _reduce_update_read, gathers_residual=True, fusions=(), group=None, @@ -475,8 +478,8 @@ def _select_input_steps( ) return ( partial( - _read_input, - layer_input=None, + _update_read, + pre_move=None, enters_stack=enters_stack, read=read, ), @@ -485,9 +488,9 @@ def _select_input_steps( 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 owes is not SumGroup.ATTN_TP or not adds_plainly: + if owes is not SumGroup.ATTN_TP or not is_plain_add: raise NotImplementedError(f"{produced=} {residual=} {need=}") - return (partial(_mlp_input_residual_into_sum, read=read), None) + return (partial(_tp_sum_with_residual_read, 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 ( @@ -497,7 +500,7 @@ def _select_input_steps( fused = tuple(f for f in fusions if f.completes is owes) return ( partial( - _mlp_input_without_dp, + _reduce_update_read, gathers_residual=gathers_residual, fusions=tuple(f.run for f in fused), group=owes, @@ -516,13 +519,13 @@ def _select_input_steps( places_cp_shards = TokenAxis.ATTN_CP in gathered if ( owes_attention_tp - and not read.before_gather - and adds_plainly - and read.norms_plainly + and not read.reads_before_dp_gather + and is_plain_add + and read.is_plain_norm ): return ( partial( - _mlp_input_dp_partial, + _dp_gather_sum_read, gathers_residual=gathers_residual, places_cp_shards=places_cp_shards, read=read, @@ -531,7 +534,7 @@ def _select_input_steps( ) return ( partial( - _mlp_input_dp_replicate, + _reduce_update_read_dp_gather, gathers_residual=gathers_residual, reduces_attention_tp=owes_attention_tp, places_cp_shards=places_cp_shards, @@ -541,8 +544,8 @@ def _select_input_steps( ) -def _select_ffn_output_move( - produced: StageOutput, +def _select_exit_move( + produced: OutputContract, *, residual: Layout, to: Layout, @@ -557,22 +560,22 @@ def _select_ffn_output_move( update = produced.update if produced.layout == residual: if to == residual: - return False, identity_output, False - if to.sharded == residual.sharded - {TokenAxis.ATTN_TP_SCATTER}: + return False, keep_output, False + if to.sharded == residual.sharded - {TokenAxis.ATTN_TP}: # Each rank's slice back to the attention's rows: write the output # into the residual, then gather over attention TP. - return False, partial(update_and_gather, update=update), False + return False, partial(update_attn_tp_gather_output, update=update), False raise NotImplementedError(f"{produced=} {residual=} {to=}") returned = residual.sharded - produced.layout.sharded - if returned == {TokenAxis.ATTN_TP_SCATTER} and to in (residual, produced.layout): + if returned == {TokenAxis.ATTN_TP} and to in (residual, produced.layout): # The residual stays on each rank's slice (MHC on an input-scattered # batch): a reduce-scatter onto the slice completes the sum the FFN # leaves; a complete output is only sliced. - sums = produced.leaves_for_reduce_scatter + sums = produced.may_reduce_scatter return ( False, partial( - output_on_residual_shard, + residual_slice_output, sums=sums, gathers_back=to != residual, update=update, @@ -584,7 +587,7 @@ def _select_ffn_output_move( if returned == {TokenAxis.ATTN_CP}: if cp_moves is None: raise NotImplementedError(f"{produced=} {residual=} {to=}") - if not produced.leaves_for_reduce_scatter: + if not produced.may_reduce_scatter: # A complete output: this rank's block of it, nothing summed. return False, cp_moves.take_back, False # The FFN leaves its sum: only a take-back that sums over the same @@ -596,7 +599,7 @@ def _select_ffn_output_move( return False, cp_moves.reduce_scatter, True if returned == {TokenAxis.ATTN_DP, TokenAxis.ATTN_CP}: # This rank's CP shard, from where the DP gather put it. - return False, take_back_cp_shard, False + return False, dp_cp_take_back_output, False if returned != {TokenAxis.ATTN_DP}: raise NotImplementedError(f"{produced=} {residual=} {to=}") return True, None, False diff --git a/python/sglang/srt/layers/layer_boundary/construction.py b/python/sglang/srt/layers/layer_boundary/construction.py index 61907661abab..d0045c342979 100644 --- a/python/sglang/srt/layers/layer_boundary/construction.py +++ b/python/sglang/srt/layers/layer_boundary/construction.py @@ -21,33 +21,33 @@ from sglang.srt.layers.layer_boundary.adapters.attention import get_attn_tp_context from sglang.srt.layers.layer_boundary.boundary import ( _cp_moves, - make_boundary, - make_output_boundary, + bind_entry, + bind_exit, ) from sglang.srt.layers.layer_boundary.contracts import ( BatchVariant, CpMoves, - EdgeDecl, + EdgeContract, + EntryPath, ProducerReduction, - StageEntry, StageKind, - StageSteps, + StagePath, ) -from sglang.srt.layers.layer_boundary.exit import OutputBoundary +from sglang.srt.layers.layer_boundary.exit import ExitPolicy from sglang.srt.layers.layer_boundary.fusions.allreduce import ( - attention_fusions, - ffn_fusions, + attn_input_fusions, + ffn_input_fusions, ) from sglang.srt.layers.layer_boundary.layout import ( TokenAxis, _batch_shards_over_cp, - _gathers_over_attention_cp, - _generic_prefill_cp_shards_tokens, - enable_moe_dense_fully_dp, + _cp_gathers_over_attn_cp, + _prefill_cp_shards_tokens, + is_dense_ffn_fully_dp, ) from sglang.srt.layers.layer_boundary.prepare import ( - _hand_qkv_hook_its_input, - _hand_scattered_input_to_attention, + _attn_input_default, + _attn_input_scattered, ) from sglang.srt.layers.layer_boundary.stage import StageBoundary from sglang.srt.layers.moe import ( @@ -64,14 +64,14 @@ _use_ag_after_qlora = envs.SGLANG_USE_AG_AFTER_QLORA.get() -def _refuse_uncovered_cp_moe(moe_on_local_rows: bool, cp_shards: bool) -> None: +def _reject_unsupported_cp_moe(moe_on_local_rows: bool, cp_shards: bool) -> None: """A MoE layer under attention CP whose tokens the steps cannot bring to it: under a prefill CP, one dispatched per DP shard under attention DP and GQA CP, and one on the TP group whose data-parallel groups are the CP ranks under DSA or MLA CP; and under attention DP, one on the TP group whose data-parallel groups are the CP ranks.""" parallel = get_parallel() - gqa = not _gathers_over_attention_cp() + gqa = not _cp_gathers_over_attn_cp() if moe_on_local_rows: if cp_shards and gqa and parallel.attn_dp_size > 1: raise NotImplementedError( @@ -90,7 +90,7 @@ def _refuse_uncovered_cp_moe(moe_on_local_rows: bool, cp_shards: bool) -> None: ) -def _input_can_be_scattered() -> bool: +def _input_scattered_possible() -> bool: """Whether a batch may run this layer with input-scattered attention: configured, on TP without attention DP, a prefill CP, an a2a backend or a dense MLP on every rank. The rest of what ``AttnTpContext.init_context`` requires is only @@ -100,28 +100,28 @@ def _input_can_be_scattered() -> bool: parallel.enable_attn_tp_input_scattered and parallel.tp_size > 1 and parallel.attn_dp_size == 1 - and not _generic_prefill_cp_shards_tokens() + and not _prefill_cp_shards_tokens() and get_moe_a2a_backend().is_none() - and not enable_moe_dense_fully_dp() + and not is_dense_ffn_fully_dp() ) @dataclass(frozen=True) -class StageEdges: +class VariantEdges: """One stage's incoming and outgoing contracts for one batch variant. Fields: incoming: Consumer-side contract used to bind prepare. outgoing: Producer-side contract used to bind exit transport. - handoff: Optional attention adapter run after input preparation/movement. + attn_input_adapter: Optional attention adapter run after input preparation/movement. cp_moves: Strategy-specific context-parallel gather and return operations. No executable neighbouring stage or neighbouring norm is required. """ - incoming: EdgeDecl - outgoing: EdgeDecl - handoff: Optional[Callable] = None + incoming: EdgeContract + outgoing: EdgeContract + attn_input_adapter: Optional[Callable] = None cp_moves: Optional[CpMoves] = None @@ -135,13 +135,13 @@ class StagePlan: Args: kind: Boundary role selecting attention or FFN fusion adapters. norm: This stage's consumer normalization module. - variants: Mapping from BatchVariant to StageEdges. Construction binds + variants: Mapping from BatchVariant to VariantEdges. Construction binds each entry once; a forward selects an existing path by batch facts. enters_stack: Whether this stage initializes the residual from embeddings. - prepared_input: Branch reusing an already-read input; requires branch_input + is_branch: Branch reusing an already-read input; requires branch_input instead of a normal prepare to avoid repeating the read/update. terminal: Whether the stage ends the model's layer stack. - direct_handoff: Attention can publish its output with finish instead of + finishes_directly: Attention can publish its output with finish instead of an exit scope and output transport. qkv_latent_func: Optional hook for prepared attention input. fusions: Optional backend provider of ordered consumer fusion candidates @@ -158,17 +158,17 @@ def __init__( variants, *, enters_stack=False, - prepared_input=False, + is_branch=False, terminal=False, - direct_handoff=False, + finishes_directly=False, qkv_latent_func=None, fusions=None, ): self.norm = norm - self.variants = dict(variants) + self.edges = dict(variants) self.enters_stack = enters_stack self.terminal = terminal - self.direct_handoff = direct_handoff + self.finishes_directly = finishes_directly self.qkv_latent_func = qkv_latent_func self.fusions = fusions self._speculative_algo = SpeculativeAlgorithm.from_string( @@ -177,50 +177,48 @@ def __init__( self._publish_lora_layout = get_parallel().enable_dp_attention and bool( get_lora().enable_lora ) - self._fusion_rows = None + self._next_input_rows = None carried = ( - attention_fusions( - self, next(iter(self.variants.values())).incoming.need.read - ) + attn_input_fusions(self, next(iter(self.edges.values())).incoming.need.read) if kind is StageKind.ATTENTION else () ) - fused = ffn_fusions(self) if kind is StageKind.FFN else () - self._paths = {} - for variant, edges in self.variants.items(): - if prepared_input: + fused = ffn_input_fusions(self) if kind is StageKind.FFN else () + self.paths = {} + for variant, edges in self.edges.items(): + if is_branch: # The branch adapter moves an already-read input and forks its # stream. Binding a normal prepare would add or norm it twice. - entry = StageEntry( + entry = EntryPath( prepare=_requires_branch_input, input_rows=edges.incoming.need.layout, ) else: - into = make_boundary( + into = bind_entry( edges.incoming, fusions=fused, carried_fusions=carried, cp_moves=edges.cp_moves, enters_stack=enters_stack, ) - entry = StageEntry( + entry = EntryPath( prepare=into.prepare, input_rows=into.input_rows, input_move=into.input_move, - handoff=edges.handoff, + attn_input_adapter=edges.attn_input_adapter, capture_move=into.capture_move, capture_move_allocates=into.capture_move_allocates, preserves_residual=into.preserves_residual, - input_sum=edges.incoming.produced.group - if edges.incoming.produced.always_leaves + declared_sum=edges.incoming.produced.group + if edges.incoming.produced.always_partial else None, ) out = ( None - if direct_handoff - else make_output_boundary(edges.outgoing, cp_moves=edges.cp_moves) + if finishes_directly + else bind_exit(edges.outgoing, cp_moves=edges.cp_moves) ) - self._paths[variant] = StageSteps( + self.paths[variant] = StagePath( entry=entry, output=edges.outgoing.produced, output_move=None if out is None else out.output_move, @@ -232,13 +230,13 @@ def __init__( @property def incoming_residual_rows(self): - return self.variants[BatchVariant.ORDINARY].incoming.residual + return self.edges[BatchVariant.ORDINARY].incoming.residual @property - def input_on_attention_tp_slices(self): - return TokenAxis.ATTN_TP_SCATTER in self.incoming_residual_rows.sharded + def input_on_attn_tp_slices(self): + return TokenAxis.ATTN_TP in self.incoming_residual_rows.sharded - def _variant(self, forward_batch): + def variant_for(self, forward_batch): # The batch determines its rows; missing paths must not change them. if get_forward().sp_active: return BatchVariant.SEQUENCE_PARALLEL @@ -248,32 +246,32 @@ def _variant(self, forward_batch): return BatchVariant.CONTEXT_PARALLEL return BatchVariant.ORDINARY - def _batch_steps(self, forward_batch): - variant = self._variant(forward_batch) + def path_for(self, forward_batch): + variant = self.variant_for(forward_batch) try: - return self._paths[variant] + return self.paths[variant] except KeyError: raise NotImplementedError( f"no stage boundary path for the active {variant.name} batch" ) from None - def fusion_rows(self, forward_batch): - if self._fusion_rows is not None: - return self._fusion_rows[self._variant(forward_batch)] - entry = self._batch_steps(forward_batch) + def fused_input_rows(self, forward_batch): + if self._next_input_rows is not None: + return self._next_input_rows[self.variant_for(forward_batch)] + entry = self.path_for(forward_batch) return entry.entry.input_rows def produced(self, forward_batch): - return self._batch_steps(forward_batch).output + return self.path_for(forward_batch).output @cached_property def output(self): - return OutputBoundary(self) + return ExitPolicy(self) - def _branch_rows(self, forward_batch): - if self._variant(forward_batch) is not BatchVariant.ORDINARY: + def branch_rows(self, forward_batch): + if self.variant_for(forward_batch) is not BatchVariant.ORDINARY: raise NotImplementedError("branch transport requires ordinary batch rows") - edges = self.variants[BatchVariant.ORDINARY] + edges = self.edges[BatchVariant.ORDINARY] return ( edges.incoming.need.layout, edges.incoming.residual_to, @@ -289,30 +287,32 @@ def _bind_stage(declaration, norm, incoming, outgoing, **options): variants = {} for variant, edge in incoming.entries.items(): if ( - declaration.update.at_producer + declaration.update.applied_at_exit and TokenAxis.ATTN_CP in edge.produced.layout.sharded - edge.need.layout.sharded ): raise NotImplementedError("MHC with a gather over attention CP") - handoff = None + attn_input_adapter = None if declaration.kind is StageKind.ATTENTION: - handoff = ( - _hand_scattered_input_to_attention + attn_input_adapter = ( + _attn_input_scattered if variant is BatchVariant.INPUT_SCATTERED - and not declaration.update.adds_plainly - else _hand_qkv_hook_its_input + and not declaration.update.is_plain_add + else _attn_input_default ) moves = _cp_moves() if variant is BatchVariant.CONTEXT_PARALLEL else None - variants[variant] = StageEdges(edge, outgoing.exits[variant], handoff, moves) + variants[variant] = VariantEdges( + edge, outgoing.exits[variant], attn_input_adapter, moves + ) plan = StagePlan( declaration.kind, norm, variants, enters_stack=incoming.producer is None, - prepared_input=declaration.prepared_from is not None, + is_branch=declaration.prepared_from is not None, terminal=declaration.terminal, - direct_handoff=declaration.kind is StageKind.ATTENTION - and declaration.reduction is ProducerReduction.PARTIAL, + finishes_directly=declaration.kind is StageKind.ATTENTION + and declaration.reduction is ProducerReduction.ALWAYS_PARTIAL, **options, ) if declaration.kind is StageKind.ATTENTION and outgoing.consumer is not None: @@ -320,7 +320,7 @@ def _bind_stage(declaration, norm, incoming, outgoing, **options): # link to its execution plan. Kernel binding remains consumer-owned. from sglang.srt.layers.layer_boundary.boundary import input_rows - plan._fusion_rows = { + plan._next_input_rows = { v: input_rows(edge) for v, edge in outgoing.entries.items() } diff --git a/python/sglang/srt/layers/layer_boundary/contracts.py b/python/sglang/srt/layers/layer_boundary/contracts.py index 33843557740e..aa72c5dafbd8 100644 --- a/python/sglang/srt/layers/layer_boundary/contracts.py +++ b/python/sglang/srt/layers/layer_boundary/contracts.py @@ -22,28 +22,29 @@ from sglang.srt.distributed import GroupCoordinator from sglang.srt.layers.layer_boundary.layout import Layout, SumGroup, TokenAxis from sglang.srt.layers.layer_boundary.output import OutputTransform -from sglang.srt.layers.layer_boundary.residual import StageRead, StageUpdate -from sglang.srt.layers.layer_boundary.residual.add_norm import ADD, NORM_READ +from sglang.srt.layers.layer_boundary.residual import ResidualReadout, ResidualUpdate +from sglang.srt.layers.layer_boundary.residual.add_norm import NORM_READOUT, PLAIN_ADD class ProducerReduction(Enum): """Describe compute's cooperation with its output boundary. - PARTIAL is attention-only: finish() publishes its partial sum. SCOPED - follows exit() flags for an FFN or single-stage mixer. LOCAL_TAIL is FFN-only: + ALWAYS_PARTIAL is attention-only: finish() publishes its partial sum. EXIT_SCOPED + follows exit() flags for an FFN or single-stage mixer. TAIL_AFTER_SUM is FFN-only: compute adds a replicated component after its internal sum, so the boundary never defers that sum to the next layer. A selected reduce-scatter still applies, so compute must add the tail only when it completes the sum. """ - PARTIAL = auto() - SCOPED = auto() + ALWAYS_PARTIAL = auto() + EXIT_SCOPED = auto() # A replicated component is added after the sum inside compute. - LOCAL_TAIL = auto() + TAIL_AFTER_SUM = auto() -class HandoffRows(Enum): - """A required layer/branch handoff, independent of compute's input rows.""" +class ExitRows(Enum): + """The rows a layer or branch must leave its output on, independent of + compute's input rows.""" ATTENTION = auto() TBO_SPLIT = auto() @@ -56,12 +57,12 @@ class BatchVariant(Enum): SEQUENCE_PARALLEL = auto() -class StageInput(msgspec.Struct, frozen=True): +class InputContract(msgspec.Struct, frozen=True): """Declare the consumer's input rows and read operation. Fields: layout: Token sharding required by compute. - gathers_itself: Token axes compute can gather internally; the boundary + gathered_by_compute: Token axes compute can gather internally; the boundary may hand it input still sharded over those axes. read: Operation that derives input from the updated residual. """ @@ -69,23 +70,23 @@ class StageInput(msgspec.Struct, frozen=True): layout: Layout # Token axes the consumer gathers over itself when its input arrives # sharded over them. - gathers_itself: FrozenSet[TokenAxis] = frozenset() + gathered_by_compute: FrozenSet[TokenAxis] = frozenset() # How the consumer reads its input from the residual. - read: StageRead = NORM_READ + read: ResidualReadout = NORM_READOUT -class StageOutput(msgspec.Struct, frozen=True): +class OutputContract(msgspec.Struct, frozen=True): """Declare producer rows and permitted reduction handoffs at construction. Fields: layout: Token sharding of the producer contribution. group: Named sum group, or None when there is no reduction. - always_leaves: Compute always returns a partial sum. - leaves_for_next_layer: Compute can skip reduction under the exit scope + always_partial: Compute always returns a partial sum. + may_defer_to_next: Compute can skip reduction under the exit scope and let the following layer complete it. - leaves_for_reduce_scatter: Compute can leave reduction to a fixed-size + may_reduce_scatter: Compute can leave reduction to a fixed-size reduce-scatter selected by the boundary. - leaves_for_reduce_scatterv: Compute can leave reduction to the selected + may_reduce_scatterv: Compute can leave reduction to the selected variable-size attention-DP combine. update: Producer residual operation; None on an arrival contract that declares capabilities and obtains the actual update from the stream. @@ -100,20 +101,20 @@ class StageOutput(msgspec.Struct, frozen=True): group: Optional[SumGroup] = None # The producer never reduces its output, e.g. a row-parallel projection # built with reduce_results=False. - always_leaves: bool = False + always_partial: bool = False # Otherwise it leaves the sum only when the boundary publishes the flag for # it: fuse_mlp_allreduce, or mlp_reduce_scatter. - leaves_for_next_layer: bool = False - leaves_for_reduce_scatter: bool = False + may_defer_to_next: bool = False + may_reduce_scatter: bool = False # Whether the selected attention-DP combine may complete this sum. - leaves_for_reduce_scatterv: bool = False + may_reduce_scatterv: bool = False # How the producer's output is written into the residual. An arrival # description has no producer object; its edge declares capabilities only. - update: Optional[StageUpdate] = ADD + update: Optional[ResidualUpdate] = PLAIN_ADD transform: Optional[OutputTransform] = None -class StageDecl(msgspec.Struct, frozen=True): +class StageContract(msgspec.Struct, frozen=True): """Resolved input and output contracts for one stage and batch variant. Fields: @@ -121,11 +122,11 @@ class StageDecl(msgspec.Struct, frozen=True): output: Producer row, reduction and residual-update contract. """ - input: StageInput - output: StageOutput + input: InputContract + output: OutputContract -class EdgeDecl(msgspec.Struct, frozen=True): +class EdgeContract(msgspec.Struct, frozen=True): """Describe both sides of a boundary and the residual's row movement. Fields: @@ -135,23 +136,23 @@ class EdgeDecl(msgspec.Struct, frozen=True): residual_to: Residual layout after the boundary. residual_joins_sum: Whether one rank may add the residual into a partial before reduction; valid only for an eligible plain-add update. - update_capabilities: Allowed values of StageUpdate.adds_plainly for + arriving_plain_add: Allowed values of ResidualUpdate.is_plain_add for arriving contributions. Empty means use produced.update's capability. The actual update object travels with the residual stream. """ - produced: StageOutput - need: StageInput + produced: OutputContract + need: InputContract 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 # Capabilities allowed to arrive from another layer, not its update object. - update_capabilities: Tuple[bool, ...] = () + arriving_plain_add: Tuple[bool, ...] = () -class FusedMlpInput(msgspec.Struct, frozen=True): +class FfnInputFusion(msgspec.Struct, frozen=True): """A kernel that completes the sum the attention output owes together with the residual add and the post-attention norm, in that order. @@ -181,7 +182,7 @@ class CpMoves(msgspec.Struct, frozen=True): reduce_scatter_group: Optional[Callable[[], GroupCoordinator]] = None -class StageEntry(msgspec.Struct, frozen=True): +class EntryPath(msgspec.Struct, frozen=True): """Bound consumer operations for one batch variant. Fields: @@ -189,15 +190,15 @@ class StageEntry(msgspec.Struct, frozen=True): returning compute input and updated residual; finishes owed work, applies the producer update and performs the consumer read. input_rows: Layout handed to compute after preparation and input_move. - input_move: Optional movement after prepare, before the compute handoff. - handoff: Optional callable(input, forward_batch, qkv_latent_func) that + input_move: Optional movement after prepare, before compute takes the input. + attn_input_adapter: Optional callable(input, forward_batch, qkv_latent_func) that adapts already-placed input for attention. capture_move: Optional movement of the updated residual back onto the producer's rows for auxiliary capture. capture_move_allocates: Whether capture_move returns fresh storage (a gather) rather than a view (a cut); capture then retains it without copying. - input_sum: Statically owed sum group for an otherwise raw input tensor. + declared_sum: Statically owed sum group for an otherwise raw input tensor. preserves_residual: Optional (residual, batch) predicate certifying this input path leaves residual untouched, including backend fallback. capture_preserves_residual: Same guarantee for the local following stage, @@ -215,21 +216,21 @@ class StageEntry(msgspec.Struct, frozen=True): input_move: Optional[Callable] = None # Hands the stage its input once it is on its rows: # (hidden_states, forward_batch, qkv_latent_func) -> hidden_states. - handoff: Optional[Callable] = None + attn_input_adapter: Optional[Callable] = None # Return the updated residual to the producer's rows for aux capture. capture_move: Optional[Callable] = None capture_move_allocates: bool = False - input_sum: Optional[SumGroup] = None + declared_sum: Optional[SumGroup] = None preserves_residual: Optional[Callable] = None # Bound from a local successor at construction; no runtime plan lookup. capture_preserves_residual: Optional[Callable] = None -class StageSteps(msgspec.Struct, frozen=True): +class StagePath(msgspec.Struct, frozen=True): """Precomputed entry and exit work for one stage and batch variant. Fields: - entry: Bound consumer preparation and handoff. + entry: Bound consumer half: preparation, input move and input adapter. output: Producer contract used by the exit decision. output_move: Fixed output transport, or None when absent or chosen per batch by the attention-DP exit path. @@ -237,8 +238,8 @@ class StageSteps(msgspec.Struct, frozen=True): returns_over_dp: Whether output uses batch-dependent attention-DP transport. """ - entry: StageEntry - output: StageOutput + entry: EntryPath + output: OutputContract # None means no fixed move. returns_over_dp selects batch-dependent DP transport. output_move: Optional[Callable] output_move_completes_sum: bool = False diff --git a/python/sglang/srt/layers/layer_boundary/exit.py b/python/sglang/srt/layers/layer_boundary/exit.py index 9528802f2f69..44c8a187e5ab 100644 --- a/python/sglang/srt/layers/layer_boundary/exit.py +++ b/python/sglang/srt/layers/layer_boundary/exit.py @@ -35,18 +35,18 @@ _sum_group, ) from sglang.srt.layers.layer_boundary.ops import ( - _all_reduce_then_to_local_tokens, - _redistribute_output, - _reduce_and_redistribute_output_max_len, - _reduce_and_redistribute_output_varlen, - _to_local_tokens, + _dp_scatter_step, + all_reduce_to_dp_local, + dp_reduce_scatter, + dp_reduce_scatterv, + to_dp_local, ) from sglang.srt.layers.layer_boundary.output import ( UnreducedOutput, ) from sglang.srt.layers.layer_boundary.residual.add_norm import ( - apply_aiter_all_reduce_fusion, - apply_flashinfer_allreduce_fusion, + aiter_ar_fusion_applies, + flashinfer_ar_fusion_applies, ) from sglang.srt.layers.layer_boundary.residual.stream import ResidualStream from sglang.srt.layers.moe import ( @@ -67,49 +67,49 @@ ) -def _reduce_and_redistribute_output_step( +def _select_dp_reduce_scatter( forward_batch: ForwardBatch, *, - leaves_for_reduce_scatter: bool, - leaves_for_reduce_scatterv: bool, + may_reduce_scatter: bool, + may_reduce_scatterv: bool, ) -> Optional[Callable[[torch.Tensor, torch.Tensor, ForwardBatch], None]]: """The reduce-scatter that brings an FFN output gathered over attention DP back to this rank's tokens when the FFN leaves its sum to it (see - StageOutput); None when the FFN reduces the output and only a scatter + OutputContract); None when the FFN reduces the output and only a scatter remains.""" - if should_use_dp_reduce_scatterv() and leaves_for_reduce_scatterv: - return _reduce_and_redistribute_output_varlen + if should_use_dp_reduce_scatterv() and may_reduce_scatterv: + return dp_reduce_scatterv if ( - leaves_for_reduce_scatter + may_reduce_scatter and forward_batch.dp_padding_mode.is_max_len() and can_use_dp_reduce_scatter() ): - return _reduce_and_redistribute_output_max_len + return dp_reduce_scatter return None -class OutputBoundary: +class ExitPolicy: """Choose and finish one producer's output; the batch owns its residual.""" def __init__(self, plan): self.plan = plan - def postprocess_layer( + def finish_complete_output( self, hidden_states: torch.Tensor, stream: ResidualStream, forward_batch: ForwardBatch, ): """Move a complete output from the operation-scheduled producer path.""" - steps = self.plan._batch_steps(forward_batch) - hidden_states, residual = self._complete_ffn_output_now( + steps = self.plan.path_for(forward_batch) + hidden_states, residual = self._complete_now( hidden_states, stream.residual, forward_batch=forward_batch, dp_step=None, steps=steps, ) - return self._leave_ffn_output( + return self._record_output( hidden_states, residual, stream, @@ -117,14 +117,14 @@ def postprocess_layer( ) @staticmethod - def _leave_ffn_output(hidden_states, residual, stream, update, declared_sum=None): + def _record_output(hidden_states, residual, stream, update, declared_sum=None): stream.residual = residual - if update.at_producer: + if update.applied_at_exit: return stream.write(hidden_states) - return stream.leave(hidden_states, update, declared_sum=declared_sum) + return stream.record(hidden_states, update, declared_sum=declared_sum) @staticmethod - def _declared_ffn_sum(steps, skipped_reduction): + def _sum_owed_after_skip(steps, skipped_reduction): # The input-scattered path leaves the sum for the next input's TP # reduce-scatter. Other skip paths complete it in the output move. if ( @@ -135,13 +135,13 @@ def _declared_ffn_sum(steps, skipped_reduction): return SumGroup.TP return None - def _local_token_move_can_go_to_next_layer(self, steps) -> bool: + def _next_input_can_scatter(self, steps) -> bool: """Whether the next layer's input can run this layer's move of its FFN output back to this rank's tokens: the base postprocess scatter, when the next layer's input also writes the output into the residual.""" - return steps.returns_over_dp and not steps.output.update.at_producer + return steps.returns_over_dp and not steps.output.update.applied_at_exit - def _postprocess_dp_step( + def _dp_reduce_scatter_step( self, forward_batch: ForwardBatch, steps ) -> Optional[Callable[[torch.Tensor, torch.Tensor, ForwardBatch], None]]: """The reduce-scatter that brings this layer's FFN output back to this @@ -149,21 +149,19 @@ def _postprocess_dp_step( only scatter, or does not move tokens.""" if not steps.returns_over_dp: return None - return _reduce_and_redistribute_output_step( + return _select_dp_reduce_scatter( forward_batch, - leaves_for_reduce_scatter=steps.output.leaves_for_reduce_scatter, - leaves_for_reduce_scatterv=steps.output.leaves_for_reduce_scatterv, + may_reduce_scatter=steps.output.may_reduce_scatter, + may_reduce_scatterv=steps.output.may_reduce_scatterv, ) - def _ffn_leaves_sum_to_reduce_scatter( - self, steps, dp_step: Optional[Callable] - ) -> bool: + def _skips_sum_for_reduce_scatter(self, steps, dp_step: Optional[Callable]) -> bool: """Whether the FFN leaves its sum out because a reduce-scatter completes it: the attention-DP one ``dp_step`` names, or the CP / input-scattered one.""" if dp_step is not None: return True - if not steps.output.leaves_for_reduce_scatter: + if not steps.output.may_reduce_scatter: return False if steps.output_move_completes_sum: return True @@ -174,15 +172,13 @@ def ffn_reduction_group(self, steps) -> GroupCoordinator: group on a sparse layer, the TP group a dense MLP reduces over.""" return _sum_group(steps.output.group) - def _select_ffn_completion( - self, forward_batch: ForwardBatch, steps - ) -> FfnCompletion: + def _decide(self, forward_batch: ForwardBatch, steps) -> ExitDecision: """Decide once, before the FFN runs, what it skips and what completes its output: the next layer's input, or this layer's postprocess step.""" - dp_step = self._postprocess_dp_step(forward_batch, steps) - mlp_reduce_scatter = self._ffn_leaves_sum_to_reduce_scatter(steps, dp_step) + dp_step = self._dp_reduce_scatter_step(forward_batch, steps) + mlp_reduce_scatter = self._skips_sum_for_reduce_scatter(steps, dp_step) complete_now = partial( - self._complete_ffn_output_now, + self._complete_now, forward_batch=forward_batch, dp_step=dp_step, steps=steps, @@ -191,10 +187,10 @@ def _select_ffn_completion( self.plan.fusions is not None and self.plan.fusions.can_defer_finalize(self.plan, forward_batch) ) - if not steps.output.leaves_for_next_layer and not ( + if not steps.output.may_defer_to_next and not ( self.plan.terminal and defer_moe_finalize ): - return FfnCompletion( + return ExitDecision( defer_moe_finalize=False, fuse_mlp_allreduce=False, mlp_reduce_scatter=mlp_reduce_scatter, @@ -202,7 +198,7 @@ def _select_ffn_completion( ) # Producers declare remaining work independently of the kernel chosen # by the consumer. Every handoff also carries an unfused completion. - fuse_mlp_allreduce = defer_moe_finalize or self._ffn_sum_moves_to_next_layer( + fuse_mlp_allreduce = defer_moe_finalize or self._defers_sum( forward_batch, steps, mlp_reduce_scatter=mlp_reduce_scatter, dp_step=dp_step ) if fuse_mlp_allreduce: @@ -212,37 +208,35 @@ def _select_ffn_completion( # this rank's tokens. wrap = partial( UnreducedOutput, - reduce_and_redistribute=partial( - _all_reduce_then_to_local_tokens, group, forward_batch + reduce_to_dp_local=partial( + all_reduce_to_dp_local, group, forward_batch ), ) else: wrap = partial(UnreducedOutput, group=group) - complete = partial(_leave_to_next_layer, wrap) + complete = partial(_defer, wrap) elif ( dp_step is not None and not self.plan.terminal - and self._local_token_move_can_go_to_next_layer(steps) + and self._next_input_can_scatter(steps) ): complete = partial( - _leave_to_next_layer, + _defer, partial( UnreducedOutput, - reduce_and_redistribute=partial( - _to_local_tokens, dp_step, forward_batch - ), + reduce_to_dp_local=partial(to_dp_local, dp_step, forward_batch), ), ) else: complete = complete_now - return FfnCompletion( + return ExitDecision( defer_moe_finalize=defer_moe_finalize, fuse_mlp_allreduce=fuse_mlp_allreduce, mlp_reduce_scatter=mlp_reduce_scatter, complete=complete, ) - def _complete_ffn_output_now( + def _complete_now( self, hidden_states: torch.Tensor, residual: torch.Tensor, @@ -257,8 +251,8 @@ def _complete_ffn_output_now( if steps.output.transform is not None: hidden_states = steps.output.transform.apply(hidden_states) if steps.returns_over_dp: - hidden_states = _to_local_tokens( - dp_step or _redistribute_output, forward_batch, hidden_states + hidden_states = to_dp_local( + dp_step or _dp_scatter_step, forward_batch, hidden_states ) else: hidden_states, residual = steps.output_move( @@ -267,7 +261,7 @@ def _complete_ffn_output_now( forward_batch=forward_batch, ) update = steps.output.update - if residual is not None and update.at_producer: + if residual is not None and update.applied_at_exit: hidden_states = update.update(hidden_states, residual) residual = None return hidden_states, residual @@ -287,8 +281,8 @@ def ffn_exit( result as a context manager around the FFN call, then call ``finish``.""" return FfnExit(self, forward_batch, stream=stream) - def _ffn_sum_can_move_to_next_layer(self, steps) -> bool: - # Under the MoE-CP all-gather the fusion path would skip postprocess_layer + def _sum_deferral_allowed(self, steps) -> bool: + # Under the MoE-CP all-gather the fusion path would skip finish_complete_output # and its MoE-CP scatter, leaving hidden_states longer than the residual. if is_enable_moe_cp_allgather() or steps.output.group is None: return False @@ -314,7 +308,7 @@ def _ffn_sum_can_move_to_next_layer(self, steps) -> bool: return not get_attn_tp_context().input_scattered - def _ffn_sum_moves_to_next_layer( + def _defers_sum( self, forward_batch: ForwardBatch, steps, @@ -328,22 +322,19 @@ def _ffn_sum_moves_to_next_layer( return ( get_parallel().tp_size > 1 and not self.plan.terminal - and self._ffn_sum_can_move_to_next_layer(steps) - and _can_defer_ffn_reduction(forward_batch, self.plan) + and self._sum_deferral_allowed(steps) + and _batch_allows_deferred_sum(forward_batch, self.plan) and not mlp_reduce_scatter # Under attention DP the next layer must also run postprocess's # scatter back to this rank's tokens, and nothing more. and ( not is_dp_attention_enabled() - or ( - self._local_token_move_can_go_to_next_layer(steps) - and dp_step is None - ) + or (self._next_input_can_scatter(steps) and dp_step is None) ) ) -def _can_defer_ffn_reduction(forward_batch: ForwardBatch, boundary=None) -> bool: +def _batch_allows_deferred_sum(forward_batch: ForwardBatch, boundary=None) -> bool: """Admit ordinary single-sum outputs, plus LoRA and TP1 shared-expert outputs when a fused consumer (the backend's can_defer_all_reduce, FlashInfer or aiter) can take them. @@ -371,16 +362,14 @@ def _can_defer_ffn_reduction(forward_batch: ForwardBatch, boundary=None) -> bool and boundary.fusions.can_defer_all_reduce(boundary, forward_batch) ): return True - if apply_flashinfer_allreduce_fusion(forward_batch.input_ids.shape[0]): + if flashinfer_ar_fusion_applies(forward_batch.input_ids.shape[0]): return True # Aiter also checks width and bytes, so use the actual residual storage. residual = forward_batch.residual_stream.residual - return residual is not None and apply_aiter_all_reduce_fusion( - residual, forward_batch - ) + return residual is not None and aiter_ar_fusion_applies(residual, forward_batch) -class FfnCompletion(msgspec.Struct, frozen=True): +class ExitDecision(msgspec.Struct, frozen=True): """One decision shared by compute flags and the matching output completion. Fields: @@ -400,7 +389,7 @@ class FfnCompletion(msgspec.Struct, frozen=True): complete: Callable[[torch.Tensor, torch.Tensor], Tuple] -def _leave_to_next_layer( +def _defer( wrap: Callable[[torch.Tensor], UnreducedOutput], hidden_states: torch.Tensor, residual: torch.Tensor, @@ -420,7 +409,7 @@ class MixerExit: __slots__ = ( "skips_reduction", - "_hands_on", + "_defers", "_scope", "_update", "_declared_sum", @@ -429,23 +418,23 @@ class MixerExit: def __init__( self, - boundary: OutputBoundary, + boundary: ExitPolicy, forward_batch: ForwardBatch, *, stream: ResidualStream, ): - steps = boundary.plan._batch_steps(forward_batch) + steps = boundary.plan.path_for(forward_batch) produced = steps.output self._stream = stream self._update = produced.update - self._declared_sum = produced.group if produced.always_leaves else None - self._hands_on = ( - produced.leaves_for_next_layer + self._declared_sum = produced.group if produced.always_partial else None + self._defers = ( + produced.may_defer_to_next and get_parallel().tp_size > 1 and not is_dp_attention_enabled() - and boundary._ffn_sum_can_move_to_next_layer(steps) + and boundary._sum_deferral_allowed(steps) ) - self.skips_reduction = produced.always_leaves or self._hands_on + self.skips_reduction = produced.always_partial or self._defers self._scope = get_forward().scoped(fuse_mlp_allreduce=self.skips_reduction) def __enter__(self) -> MixerExit: @@ -457,11 +446,11 @@ def __exit__(self, *exc_info): def finish(self, hidden_states: torch.Tensor): """Hand the actual mixer result to the next boundary through the stream.""" - if self._hands_on: + if self._defers: hidden_states = UnreducedOutput( hidden_states, group=get_parallel().tp_group ) - return self._stream.leave( + return self._stream.record( hidden_states, self._update, declared_sum=self._declared_sum ) @@ -470,7 +459,7 @@ class FfnExit: """Scope a selected FFN output decision and retain its completion action. Args: - boundary: OutputBoundary that owns the producer's bound paths. + boundary: ExitPolicy that owns the producer's bound paths. forward_batch: Batch selecting reduction/finalize eligibility. stream: This invocation's residual stream, updated by finish(). @@ -498,21 +487,21 @@ class FfnExit: def __init__( self, - boundary: OutputBoundary, + boundary: ExitPolicy, forward_batch: ForwardBatch, *, stream: ResidualStream, ): self._stream = stream self.boundary = boundary - steps = boundary.plan._batch_steps(forward_batch) - completion = boundary._select_ffn_completion(forward_batch, steps) + steps = boundary.plan.path_for(forward_batch) + completion = boundary._decide(forward_batch, steps) self.defer_moe_finalize = completion.defer_moe_finalize self.fuse_mlp_allreduce = completion.fuse_mlp_allreduce self.mlp_reduce_scatter = completion.mlp_reduce_scatter self._complete = completion.complete self._update = steps.output.update - self._declared_sum = boundary._declared_ffn_sum( + self._declared_sum = boundary._sum_owed_after_skip( steps, completion.mlp_reduce_scatter ) self._scope = get_forward().scoped( @@ -536,6 +525,6 @@ def finish(self, hidden_states: torch.Tensor): assert self.defer_moe_finalize, "unrequested deferred MoE handoff" else: hidden_states, residual = self._complete(hidden_states, residual) - return self.boundary._leave_ffn_output( + return self.boundary._record_output( hidden_states, residual, self._stream, self._update, self._declared_sum ) diff --git a/python/sglang/srt/layers/layer_boundary/factories.py b/python/sglang/srt/layers/layer_boundary/factories.py index 5efd7ce9dff0..53b0da79cf25 100644 --- a/python/sglang/srt/layers/layer_boundary/factories.py +++ b/python/sglang/srt/layers/layer_boundary/factories.py @@ -10,78 +10,78 @@ from sglang.srt.layers import layernorm_sp from sglang.srt.layers.layer_boundary.adapters.overlap import ( - resolve_handoff_rows, - tbo_handoff, + resolve_exit_rows, + tbo_exit_rows, ) from sglang.srt.layers.layer_boundary.boundary import _cp_moves from sglang.srt.layers.layer_boundary.construction import ( BatchVariant, _bind_stage, - _input_can_be_scattered, - _refuse_uncovered_cp_moe, + _input_scattered_possible, + _reject_unsupported_cp_moe, _use_ag_after_qlora, ) from sglang.srt.layers.layer_boundary.contracts import ( - EdgeDecl, - HandoffRows, + EdgeContract, + ExitRows, + InputContract, + OutputContract, ProducerReduction, - StageDecl, - StageInput, + StageContract, StageKind, - StageOutput, ) from sglang.srt.layers.layer_boundary.layout import ( Layout, SumGroup, TokenAxis, - _gathers_over_attention_cp, - _generic_prefill_cp_shards_tokens, - enable_moe_dense_fully_dp, - sparse_moe_gathers_over_moe_cp, + _cp_gathers_over_attn_cp, + _prefill_cp_shards_tokens, + is_dense_ffn_fully_dp, + moe_gathers_over_moe_cp, token_axis_sizes, ) from sglang.srt.layers.layer_boundary.output import OutputTransform -from sglang.srt.layers.layer_boundary.residual import StageRead, StageUpdate +from sglang.srt.layers.layer_boundary.residual import ResidualReadout, ResidualUpdate from sglang.srt.layers.layer_boundary.residual.add_norm import ( - ADD, - NORM_QUANT_READ, - NORM_READ, + NORM_QUANT_READOUT, + NORM_READOUT, + PLAIN_ADD, ) from sglang.srt.layers.moe import is_moe_input_scattered_across_dp_ranks from sglang.srt.runtime_context import get_exec, get_parallel -def _variants(): +def _active_variants(): yield BatchVariant.ORDINARY - if _generic_prefill_cp_shards_tokens(): + if _prefill_cp_shards_tokens(): yield BatchVariant.CONTEXT_PARALLEL - if _input_can_be_scattered(): + if _input_scattered_possible(): yield BatchVariant.INPUT_SCATTERED if layernorm_sp.layernorm_sp_enabled(): yield BatchVariant.SEQUENCE_PARALLEL -def _rows(variant): +def _row_layouts(variant): axes = token_axis_sizes(cp_active=variant is BatchVariant.CONTEXT_PARALLEL) attention = Layout.sharded_over( TokenAxis.ATTN_DP, TokenAxis.ATTN_CP, axis_sizes=axes ) local = Layout.sharded_over( - TokenAxis.ATTN_DP, TokenAxis.ATTN_CP, TokenAxis.ATTN_TP_SCATTER, axis_sizes=axes + TokenAxis.ATTN_DP, TokenAxis.ATTN_CP, TokenAxis.ATTN_TP, axis_sizes=axes ) return axes, attention, local, Layout.sharded_over(axis_sizes=axes) -def _ffn_decl( +def _resolve_ffn( variant, *, sparse, terminal=False, output_transform=None, - read=NORM_READ, - update=ADD, + read=NORM_READOUT, + update=PLAIN_ADD, dense_tp_size=None, - reduction=ProducerReduction.SCOPED, + reduction=ProducerReduction.EXIT_SCOPED, ): parallel = get_parallel() if dense_tp_size not in (None, 1, parallel.tp_size): @@ -94,48 +94,46 @@ def _ffn_decl( can_move_output = output_transform is None or output_transform.before_reduce_scatter use_reduce_scatter = strategy in ("rs", "rs+rsv") and can_move_output use_reduce_scatterv = strategy in ("rsv", "rs+rsv") and can_move_output - cp_shards = _generic_prefill_cp_shards_tokens() - axes, attention, local, full = _rows(variant) - on_local = ( + cp_shards = _prefill_cp_shards_tokens() + axes, attention, local, full = _row_layouts(variant) + on_rank_rows = ( is_moe_input_scattered_across_dp_ranks() if sparse - else ( - enable_moe_dense_fully_dp() if dense_tp_size is None else dense_tp_size == 1 - ) + else (is_dense_ffn_fully_dp() if dense_tp_size is None else dense_tp_size == 1) ) if parallel.attn_cp_size > 1 and sparse: - _refuse_uncovered_cp_moe(on_local, cp_shards) - on_cp = ( + _reject_unsupported_cp_moe(on_rank_rows, cp_shards) + on_cp_shards = ( sparse and parallel.attn_cp_size > 1 and parallel.moe_dp_size == parallel.attn_cp_size - and not _gathers_over_attention_cp() + and not _cp_gathers_over_attn_cp() ) group = SumGroup.MOE_OUTPUT if sparse else SumGroup.TP if variant is BatchVariant.SEQUENCE_PARALLEL: return ( - StageDecl( - StageInput( + StageContract( + InputContract( full, - gathers_itself=frozenset({TokenAxis.ATTN_TP_SCATTER}), + gathered_by_compute=frozenset({TokenAxis.ATTN_TP}), read=read, ), - StageOutput(local, update=update, transform=output_transform), + OutputContract(local, update=update, transform=output_transform), ), local, local, ) if variant is BatchVariant.INPUT_SCATTERED: - scattered_residual = update.at_producer + scattered_residual = update.applied_at_exit residual = local if scattered_residual else attention returned = local if scattered_residual and not terminal else attention return ( - StageDecl( - StageInput(full, read=read), - StageOutput( + StageContract( + InputContract(full, read=read), + OutputContract( attention, group=group, - leaves_for_reduce_scatter=use_reduce_scatter + may_reduce_scatter=use_reduce_scatter and (scattered_residual or not terminal), update=update, transform=output_transform, @@ -148,38 +146,38 @@ def _ffn_decl( may_leave = variant is not BatchVariant.CONTEXT_PARALLEL may_scatter = True elif cp_shards or parallel.attn_dp_size > 1: - may_leave = may_scatter = parallel.attn_cp_size == 1 or on_cp + may_leave = may_scatter = parallel.attn_cp_size == 1 or on_cp_shards else: may_leave = may_scatter = True rows = ( local - if on_local + if on_rank_rows else Layout.sharded_over( - *((TokenAxis.ATTN_CP,) if on_cp else ()), axis_sizes=axes + *((TokenAxis.ATTN_CP,) if on_cp_shards else ()), axis_sizes=axes ) ) produced = ( - StageOutput(rows, update=update, transform=output_transform) - if on_local - else StageOutput( + OutputContract(rows, update=update, transform=output_transform) + if on_rank_rows + else OutputContract( rows, group=group, - leaves_for_next_layer=may_leave + may_defer_to_next=may_leave and not terminal - and not update.at_producer - and update.can_defer_across_layers + and not update.applied_at_exit + and update.outlives_layer and output_transform is None - and reduction is ProducerReduction.SCOPED, - leaves_for_reduce_scatter=use_reduce_scatter and may_scatter, - leaves_for_reduce_scatterv=use_reduce_scatterv and may_leave, + and reduction is ProducerReduction.EXIT_SCOPED, + may_reduce_scatter=use_reduce_scatter and may_scatter, + may_reduce_scatterv=use_reduce_scatterv and may_leave, update=update, transform=output_transform, ) ) - returned = local if on_local and not terminal else attention + returned = local if on_rank_rows and not terminal else attention return ( - StageDecl(StageInput(rows, read=read), produced), - local if on_local else attention, + StageContract(InputContract(rows, read=read), produced), + local if on_rank_rows else attention, returned, ) @@ -202,10 +200,10 @@ class StageDeclaration: residual update, with an explicit reduction-order contract. reduction: Whether compute always leaves a partial sum, obeys the exit scope, or adds a replicated component after its own sum. - gathers_tp_input: Whether attention gathers TP-sharded input itself. + gathers_attn_tp_input: Whether attention gathers TP-sharded input itself. dense_tp_size: Dense FFN compute width: None uses the configured width, 1 means local compute, and the full TP size means TP compute. - handoff_rows: Required FFN output rows at the layer or branch handoff. + exit_rows: Required FFN output rows at the layer or branch exit. previous: Declaration whose output this stage consumes. It may be reconstructed locally, including across pipeline ranks. prepared_from: Declaration whose already-read input a branch reuses. @@ -215,15 +213,15 @@ class StageDeclaration: """ kind: StageKind - read: StageRead - update: StageUpdate + read: ResidualReadout + update: ResidualUpdate sparse: bool = False terminal: bool = False output_transform: Optional[OutputTransform] = None - reduction: ProducerReduction = ProducerReduction.SCOPED - gathers_tp_input: bool = False + reduction: ProducerReduction = ProducerReduction.EXIT_SCOPED + gathers_attn_tp_input: bool = False dense_tp_size: Optional[int] = None - handoff_rows: Optional[HandoffRows] = None + exit_rows: Optional[ExitRows] = None # Only declarations participate in construction, never executable stages. previous: Optional[StageDeclaration] = None prepared_from: Optional[StageDeclaration] = None @@ -242,9 +240,9 @@ class StageConnection: Fields: producer: Source declaration, or None at stack entry. - consumer: Destination declaration, or None at a layer/stack handoff. - exits: Producer-side EdgeDecl for each supported BatchVariant. - entries: Consumer-side EdgeDecl for each supported BatchVariant. + consumer: Destination declaration, or None at a layer or stack exit. + exits: Producer-side EdgeContract for each supported BatchVariant. + entries: Consumer-side EdgeContract for each supported BatchVariant. Equivalent declarations can be reconstructed across pipeline partitions; adjacent layers need not share the same connection object. @@ -252,8 +250,8 @@ class StageConnection: producer: Optional[StageDeclaration] consumer: Optional[StageDeclaration] - exits: Mapping[BatchVariant, EdgeDecl] - entries: Mapping[BatchVariant, EdgeDecl] + exits: Mapping[BatchVariant, EdgeContract] + entries: Mapping[BatchVariant, EdgeContract] def __post_init__(self): object.__setattr__(self, "exits", MappingProxyType(dict(self.exits))) @@ -264,11 +262,11 @@ def declare_attn( *, previous: Optional[StageDeclaration] = None, prepared_from: Optional[StageDeclaration] = None, - read=NORM_QUANT_READ, - update=ADD, + read=NORM_QUANT_READOUT, + update=PLAIN_ADD, terminal=False, - reduction=ProducerReduction.PARTIAL, - gathers_tp_input=True, + reduction=ProducerReduction.ALWAYS_PARTIAL, + gathers_attn_tp_input=True, ): """Declare attention or a mixer; construct its executable boundary later. @@ -278,16 +276,16 @@ def declare_attn( read: Input operation; defaults to normalization with quantization support. update: Operation that adds this stage's output to the residual. terminal: Whether this stage ends the model's layer stack. - reduction: PARTIAL for an output projection that always skips reduction; - SCOPED for a mixer that follows its exit scope's reduction decision. - LOCAL_TAIL is rejected for attention stages. - gathers_tp_input: Whether compute gathers attention-TP input slices itself. + reduction: ALWAYS_PARTIAL for an output projection that always skips reduction; + EXIT_SCOPED for a mixer that follows its exit scope's reduction decision. + TAIL_AFTER_SUM is rejected for attention stages. + gathers_attn_tp_input: Whether compute gathers attention-TP input slices itself. Returns: A StageDeclaration with no norm, tensors or execution plan. """ - if reduction is ProducerReduction.LOCAL_TAIL: - raise ValueError("LOCAL_TAIL is not supported for attention stages") + if reduction is ProducerReduction.TAIL_AFTER_SUM: + raise ValueError("TAIL_AFTER_SUM is not supported for attention stages") return StageDeclaration( StageKind.ATTENTION, read, @@ -296,7 +294,7 @@ def declare_attn( prepared_from=prepared_from, terminal=terminal, reduction=reduction, - gathers_tp_input=gathers_tp_input, + gathers_attn_tp_input=gathers_attn_tp_input, ) @@ -305,14 +303,14 @@ def declare_ffn( previous: Optional[StageDeclaration] = None, prepared_from: Optional[StageDeclaration] = None, sparse=False, - read=NORM_READ, - update=ADD, + read=NORM_READOUT, + update=PLAIN_ADD, terminal=False, output_transform=None, - next_sparse=False, + next_layer_sparse=False, dense_tp_size=None, - reduction=ProducerReduction.SCOPED, - handoff_rows=None, + reduction=ProducerReduction.EXIT_SCOPED, + exit_rows=None, ): """Declare a dense or MoE FFN independently of its compute module. @@ -324,21 +322,21 @@ def declare_ffn( update: Operation writing FFN output into the residual. terminal: Whether this stage ends the model's layer stack. output_transform: Optional contribution transform before residual update. - next_sparse: Whether the next decoder layer's FFN is sparse; used only - to derive the TBO handoff when handoff_rows is not supplied. + next_layer_sparse: Whether the next decoder layer's FFN is sparse; used only + to derive the TBO exit rows when exit_rows is not supplied. dense_tp_size: Dense compute width: None for configuration, 1 for local compute, or the full TP size. reduction: How compute cooperates with the exit's reduction decision: - SCOPED (default) follows the exit scope; LOCAL_TAIL marks a - replicated tail after the sum. PARTIAL is rejected for FFN stages. - handoff_rows: Explicit output-row requirement; otherwise derived from + EXIT_SCOPED (default) follows the exit scope; TAIL_AFTER_SUM marks a + replicated tail after the sum. ALWAYS_PARTIAL is rejected for FFN stages. + exit_rows: Explicit output-row requirement; otherwise derived from the adjacent FFN kinds and TBO configuration. Returns: A StageDeclaration with no norm, tensors or execution plan. """ - if reduction is ProducerReduction.PARTIAL: - raise ValueError("PARTIAL is not supported for ffn stages") + if reduction is ProducerReduction.ALWAYS_PARTIAL: + raise ValueError("ALWAYS_PARTIAL is not supported for ffn stages") return StageDeclaration( StageKind.FFN, read, @@ -350,23 +348,23 @@ def declare_ffn( output_transform=output_transform, dense_tp_size=dense_tp_size, reduction=reduction, - handoff_rows=handoff_rows or tbo_handoff(sparse, next_sparse), + exit_rows=exit_rows or tbo_exit_rows(sparse, next_layer_sparse), ) -def _resolve(stage, variant, following=None): - axes, attention, local, full = _rows(variant) - if stage.update.at_producer: - if stage.sparse and sparse_moe_gathers_over_moe_cp(): +def _resolve_stage(stage, variant, following=None): + axes, attention, local, full = _row_layouts(variant) + if stage.update.applied_at_exit: + if stage.sparse and moe_gathers_over_moe_cp(): raise NotImplementedError( "MHC does not support a MoE gathered over the MoE-CP group" ) - if get_parallel().attn_cp_size > 1 and _input_can_be_scattered(): + if get_parallel().attn_cp_size > 1 and _input_scattered_possible(): raise NotImplementedError( "MHC with input-scattered attention under attention CP" ) if stage.kind is StageKind.FFN: - declaration, residual, returned = _ffn_decl( + declaration, residual, returned = _resolve_ffn( variant, sparse=stage.sparse, terminal=stage.terminal, @@ -376,29 +374,29 @@ def _resolve(stage, variant, following=None): dense_tp_size=stage.dense_tp_size, reduction=stage.reduction, ) - if resolve_handoff_rows(stage.handoff_rows) is HandoffRows.ATTENTION: + if resolve_exit_rows(stage.exit_rows) is ExitRows.ATTENTION: returned = attention return declaration, residual, returned sp = variant is BatchVariant.SEQUENCE_PARALLEL scattered = variant is BatchVariant.INPUT_SCATTERED gathers = ( - frozenset({TokenAxis.ATTN_TP_SCATTER}) - if stage.gathers_tp_input and (sp or scattered or _use_ag_after_qlora) + frozenset({TokenAxis.ATTN_TP}) + if stage.gathers_attn_tp_input and (sp or scattered or _use_ag_after_qlora) else frozenset() ) - owes = not sp and axes[TokenAxis.ATTN_TP_SCATTER] > 1 - declaration = StageDecl( - StageInput(attention, gathers_itself=gathers, read=stage.read), - StageOutput( + owes = not sp and axes[TokenAxis.ATTN_TP] > 1 + declaration = StageContract( + InputContract(attention, gathered_by_compute=gathers, read=stage.read), + OutputContract( local if sp else attention, group=SumGroup.ATTN_TP if owes else None, - always_leaves=owes + always_partial=owes and ( - stage.reduction is ProducerReduction.PARTIAL + stage.reduction is ProducerReduction.ALWAYS_PARTIAL or (following is not None and following.kind is StageKind.FFN) ), - leaves_for_next_layer=owes - and stage.reduction is ProducerReduction.SCOPED + may_defer_to_next=owes + and stage.reduction is ProducerReduction.EXIT_SCOPED and following is not None and following.kind is StageKind.ATTENTION, update=stage.update, @@ -410,7 +408,7 @@ def _resolve(stage, variant, following=None): def _connect(producer, consumer, *, residual_from=None): """Resolve producer and consumer declarations; None denotes an external endpoint. - A missing producer is the stack input; a missing consumer is a handoff + A missing producer is the stack input; a missing consumer is a layer or stack exit whose next read is bound separately (or the terminal output). ``residual_from`` names the boundary that placed the producer's residual. It is needed for attention, whose residual can stay on finer rows than its @@ -420,21 +418,21 @@ def _connect(producer, consumer, *, residual_from=None): raise ValueError("a boundary needs at least one declared side") before, after = producer, consumer exits, entries = {}, {} - for variant in _variants(): - _, attention, local, _ = _rows(variant) + for variant in _active_variants(): + _, attention, local, _ = _row_layouts(variant) if before is None: rows = local if variant is BatchVariant.SEQUENCE_PARALLEL else attention owes = variant is BatchVariant.INPUT_SCATTERED - arrived = StageOutput( + arrived = OutputContract( rows, group=SumGroup.TP if owes else None, - always_leaves=owes, + always_partial=owes, update=None, ) residual = rows capabilities = (True,) else: - decl, during, returned = _resolve(before, variant, following=after) + decl, during, returned = _resolve_stage(before, variant, following=after) if during is None: if residual_from is not None: if residual_from.consumer != producer: @@ -442,59 +440,58 @@ def _connect(producer, consumer, *, residual_from=None): during = residual_from.entries[variant].residual_to elif ( before.kind is StageKind.ATTENTION - and before.reduction is ProducerReduction.SCOPED + and before.reduction is ProducerReduction.EXIT_SCOPED ): during = attention else: raise ValueError( "attention output needs its incoming residual placement" ) - exits[variant] = EdgeDecl( - decl.output, StageInput(returned), during, returned + exits[variant] = EdgeContract( + decl.output, InputContract(returned), during, returned ) if ( before.kind is StageKind.ATTENTION - and before.reduction is ProducerReduction.PARTIAL + and before.reduction is ProducerReduction.ALWAYS_PARTIAL ): arrived, residual, capabilities = decl.output, during, () else: owes = ( variant is BatchVariant.INPUT_SCATTERED - and not before.update.at_producer + and not before.update.applied_at_exit ) carries = ( ( before.kind is StageKind.ATTENTION - and before.reduction is ProducerReduction.SCOPED + and before.reduction is ProducerReduction.EXIT_SCOPED ) - or resolve_handoff_rows(before.handoff_rows) - is HandoffRows.ATTENTION - ) and (decl.output.always_leaves or decl.output.leaves_for_next_layer) - arrived = StageOutput( + or resolve_exit_rows(before.exit_rows) is ExitRows.ATTENTION + ) and (decl.output.always_partial or decl.output.may_defer_to_next) + arrived = OutputContract( returned, group=decl.output.group if carries else (SumGroup.TP if owes else None), - always_leaves=decl.output.always_leaves if carries else owes, - leaves_for_next_layer=decl.output.leaves_for_next_layer + always_partial=decl.output.always_partial if carries else owes, + may_defer_to_next=decl.output.may_defer_to_next if carries else False, update=None, ) - residual, capabilities = returned, (before.update.adds_plainly,) + residual, capabilities = returned, (before.update.is_plain_add,) if after is None: continue - decl, during, _ = _resolve(after, variant) + decl, during, _ = _resolve_stage(after, variant) if after.kind is StageKind.ATTENTION: during = ( - Layout(residual.sharded | {TokenAxis.ATTN_TP_SCATTER}) + Layout(residual.sharded | {TokenAxis.ATTN_TP}) if variant is BatchVariant.INPUT_SCATTERED - and arrived.always_leaves + and arrived.always_partial and arrived.update is None - and after.reduction is ProducerReduction.PARTIAL + and after.reduction is ProducerReduction.ALWAYS_PARTIAL else residual ) - elif resolve_handoff_rows(after.handoff_rows) is HandoffRows.ATTENTION: + elif resolve_exit_rows(after.exit_rows) is ExitRows.ATTENTION: during = ( decl.input.layout if residual.sharded <= decl.input.layout.sharded @@ -503,22 +500,22 @@ def _connect(producer, consumer, *, residual_from=None): joins = ( variant is BatchVariant.INPUT_SCATTERED and arrived.update is not None - and arrived.update.adds_plainly + and arrived.update.is_plain_add and after.kind is StageKind.FFN ) - edge = EdgeDecl( + edge = EdgeContract( arrived, decl.input, residual, during, residual_joins_sum=joins, - update_capabilities=capabilities, + arriving_plain_add=capabilities, ) entries[variant] = edge if ( before is not None and before.kind is StageKind.ATTENTION - and before.reduction is ProducerReduction.PARTIAL + and before.reduction is ProducerReduction.ALWAYS_PARTIAL ): exits[variant] = edge return StageConnection(producer, consumer, exits, entries) @@ -531,19 +528,19 @@ def _fork_input(prepared, consumer): execute the branch's input norm again. """ entries = {} - for variant in _variants(): + for variant in _active_variants(): if variant is not BatchVariant.ORDINARY: raise NotImplementedError( "prepared branch transport requires ordinary token rows" ) source = prepared.entries[variant] - declaration, during, _ = _resolve(consumer, variant) - entries[variant] = EdgeDecl( - StageOutput(source.need.layout, update=None), + declaration, during, _ = _resolve_stage(consumer, variant) + entries[variant] = EdgeContract( + OutputContract(source.need.layout, update=None), declaration.input, source.residual_to, during, - update_capabilities=(True,), + arriving_plain_add=(True,), ) return StageConnection(prepared.consumer, consumer, {}, entries) @@ -558,7 +555,7 @@ def _incoming(stage): _incoming(previous) if previous is not None and previous.kind is StageKind.ATTENTION - and previous.reduction is ProducerReduction.PARTIAL + and previous.reduction is ProducerReduction.ALWAYS_PARTIAL else None ) return _connect( @@ -594,7 +591,7 @@ def make_attn_stage( declaration: Attention StageDeclaration, including its input source. norm: This consumer's normalization module, never its neighbour's norm. following: Local consumer declaration whose previous is declaration. - None leaves a layer/stack handoff for an independently bound reader. + None leaves a layer or stack exit for an independently bound reader. qkv_latent_func: Optional attention input hook, invoked after preparation and movement onto the compute input rows. fusions: Optional backend provider. Consumer side: ordered @@ -632,7 +629,7 @@ def make_ffn_stage( declaration: FFN StageDeclaration, including its input source. norm: This FFN's input normalization module. following: Local consumer declaration whose previous is declaration; - None leaves a layer/stack handoff for an independently bound reader. + None leaves a layer or stack exit for an independently bound reader. fusions: Optional backend fusion provider, as in make_attn_stage. Returns: @@ -663,7 +660,7 @@ def make_stages( prepared_from: Already-read input declaration reused by the first stage of a branch; mutually exclusive with previous. following: External consumer declaration after the last local stage. - None denotes a layer/stack handoff with an independently bound read. + None denotes a layer or stack exit with an independently bound read. terminal: Marks only the final stage as the end of the model's stack. Returns: @@ -717,15 +714,15 @@ def make_stages( for producer, consumer in zip(boundaries, boundaries[1:]): if not ( producer.kind is StageKind.ATTENTION - and producer.declaration.reduction is ProducerReduction.PARTIAL + and producer.declaration.reduction is ProducerReduction.ALWAYS_PARTIAL and consumer.kind is StageKind.FFN ): continue - for variant, steps in producer.plan._paths.items(): - next_steps = consumer.plan._paths[variant] + for variant, steps in producer.plan.paths.items(): + next_steps = consumer.plan.paths[variant] predicate = next_steps.entry.preserves_residual if predicate is not None: - producer.plan._paths[variant] = msgspec.structs.replace( + producer.plan.paths[variant] = msgspec.structs.replace( steps, entry=msgspec.structs.replace( steps.entry, capture_preserves_residual=predicate diff --git a/python/sglang/srt/layers/layer_boundary/fusions/allreduce.py b/python/sglang/srt/layers/layer_boundary/fusions/allreduce.py index 1b09a73abf3c..f3a2341d2660 100644 --- a/python/sglang/srt/layers/layer_boundary/fusions/allreduce.py +++ b/python/sglang/srt/layers/layer_boundary/fusions/allreduce.py @@ -19,15 +19,15 @@ import torch -from sglang.srt.layers.layer_boundary.contracts import FusedMlpInput +from sglang.srt.layers.layer_boundary.contracts import FfnInputFusion from sglang.srt.layers.layer_boundary.layout import SumGroup from sglang.srt.layers.layer_boundary.output import UnreducedOutput from sglang.srt.layers.layer_boundary.residual.add_norm import ( - NORM_QUANT_READ, + NORM_QUANT_READOUT, Fp8Input, - NormQuantRead, - apply_aiter_all_reduce_fusion, - apply_flashinfer_allreduce_fusion, + NormQuantReadout, + aiter_ar_fusion_applies, + flashinfer_ar_fusion_applies, ) from sglang.srt.layers.layernorm import GemmaRMSNorm, RMSNorm from sglang.srt.layers.moe import post_experts_reduction_group @@ -38,18 +38,18 @@ _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and is_hip() -def attention_fusions(plan, read=NORM_QUANT_READ) -> Tuple[Callable, ...]: +def attn_input_fusions(plan, read=NORM_QUANT_READOUT) -> Tuple[Callable, ...]: """The fused kernels that complete what the previous layer left together 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; the boundary tries them only when the update it writes in is a plain add. A backend's come first.""" - given = plan.fusions.attention_input(plan) if plan.fusions else () + given = plan.fusions.attn_input_fusions(plan) if plan.fusions else () if not hasattr(plan.norm, "forward_with_allreduce_fusion"): return given fuses_quant = ( - isinstance(read, NormQuantRead) + isinstance(read, NormQuantReadout) and read.fp8_input is not None and _use_aiter and not get_bool_env_var("SGLANG_DISABLE_FUSED_AR_QUANT", default="false") @@ -59,7 +59,7 @@ def attention_fusions(plan, read=NORM_QUANT_READ) -> Tuple[Callable, ...]: return ( *given, partial( - complete_attention_input, + fused_attn_input, plan, fuses_quant=fuses_quant, keep_bf16=fuses_quant and read.fp8_input is Fp8Input.TUPLE_AND_BF16, @@ -67,7 +67,7 @@ def attention_fusions(plan, read=NORM_QUANT_READ) -> Tuple[Callable, ...]: ) -def complete_attention_input( +def fused_attn_input( plan, owed: UnreducedOutput, residual: torch.Tensor, @@ -91,8 +91,8 @@ def complete_attention_input( return None hidden_states = owed.partial if not ( - apply_aiter_all_reduce_fusion(hidden_states, forward_batch) - or apply_flashinfer_allreduce_fusion(hidden_states.shape[0]) + aiter_ar_fusion_applies(hidden_states, forward_batch) + or flashinfer_ar_fusion_applies(hidden_states.shape[0]) ): return None if fuses_quant: @@ -111,19 +111,19 @@ def complete_attention_input( ) -def ffn_fusions(plan) -> Tuple["FusedMlpInput", ...]: +def ffn_input_fusions(plan) -> Tuple["FfnInputFusion", ...]: """The fused kernels that can take the attention -> FFN steps, in the order they are tried. They add the residual plainly; the boundary tries them only when the update it writes in is a plain add. A backend's come first.""" - given = plan.fusions.ffn_input(plan) if plan.fusions else () + given = plan.fusions.ffn_input_fusions(plan) if plan.fusions else () if not hasattr(plan.norm, "forward_with_allreduce_fusion"): return given return ( *given, - FusedMlpInput( + FfnInputFusion( completes=SumGroup.ATTN_TP, - run=partial(complete_ffn_input, plan), + run=partial(fused_ffn_input, plan), preserves_residual=( flashinfer_preserves_residual if isinstance(plan.norm, (RMSNorm, GemmaRMSNorm)) @@ -143,11 +143,11 @@ def flashinfer_preserves_residual(value, forward_batch): return ( not _use_aiter and get_parallel().attn_tp_size > 1 - and apply_flashinfer_allreduce_fusion(value.shape[0]) + and flashinfer_ar_fusion_applies(value.shape[0]) ) -def complete_ffn_input( +def fused_ffn_input( plan, hidden_states: torch.Tensor, residual: torch.Tensor, @@ -156,8 +156,8 @@ def complete_ffn_input( """The attention-TP all-reduce, residual add and norm in one aiter or flashinfer kernel; None when neither takes the batch.""" if not ( - apply_aiter_all_reduce_fusion(hidden_states, forward_batch) - or apply_flashinfer_allreduce_fusion(hidden_states.shape[0]) + aiter_ar_fusion_applies(hidden_states, forward_batch) + or flashinfer_ar_fusion_applies(hidden_states.shape[0]) ): return None return plan.norm.forward_with_allreduce_fusion( diff --git a/python/sglang/srt/layers/layer_boundary/fusions/cutedsl.py b/python/sglang/srt/layers/layer_boundary/fusions/cutedsl.py index 1ea53005d3ac..ed32bff45e95 100644 --- a/python/sglang/srt/layers/layer_boundary/fusions/cutedsl.py +++ b/python/sglang/srt/layers/layer_boundary/fusions/cutedsl.py @@ -3,7 +3,7 @@ Two patterns share one workspace, both consumed at the next layer's input RMSNorm (and, with terminal_finalize, at the final norm): AR + residual + RMSNorm, and the same with the MoE finalize and the shared-expert add folded in -when the runner hands back a MoeFinalizeHandoff. +when the runner hands back a MoeDeferredFinalize. """ from __future__ import annotations @@ -16,8 +16,8 @@ from sglang.srt.layers.dp_attention import is_dp_attention_enabled from sglang.srt.layers.layer_boundary import ( - FusedMlpInput, - HandoffOutput, + DeferredFinalize, + FfnInputFusion, SumGroup, TokenAxis, UnreducedOutput, @@ -75,7 +75,7 @@ def _resolve_max_m(*, max_running_requests: int | None) -> int: return max(positive) -class MoeFinalizeHandoff(HandoffOutput, frozen=True): +class MoeDeferredFinalize(DeferredFinalize, frozen=True): """Unfinalized routed output plus the separately gated shared contribution. The next layer's fused finalize + AR + add + norm takes it; ``finish`` is the MoE's own unfused tail, for any other reader.""" @@ -98,7 +98,7 @@ def from_flashinfer( gated_shared_output: torch.Tensor, m: int, reduce: Callable[[torch.Tensor], torch.Tensor], - ) -> MoeFinalizeHandoff: + ) -> MoeDeferredFinalize: """``reduce`` is the MoE's own sum of its finalized output.""" from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import ( finalize_flashinfer_trtllm_deferred_output, @@ -123,7 +123,7 @@ def from_flashinfer( ) -class CuteDSLFusionService: +class CuteDSLWorkspace: def __init__( self, *, @@ -167,7 +167,7 @@ def supports(self, m: int) -> bool: def finalize( self, - handoff: MoeFinalizeHandoff, + handoff: MoeDeferredFinalize, residual: torch.Tensor, gamma: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: @@ -208,73 +208,71 @@ class CuteDSLFusion: FFN when terminal_finalize is installed.""" def __init__(self) -> None: - self.service: CuteDSLFusionService | None = None + self.service: CuteDSLWorkspace | None = None # Producer capabilities; consumer kernel selection is independent. - self.hands_off_finalize = False + self.defers_finalize = False self.terminal_finalize = False def install( self, - service: CuteDSLFusionService, + service: CuteDSLWorkspace, *, - hands_off_finalize: bool, + defers_finalize: bool, terminal_finalize: bool = False, ) -> None: """Install the shared workspace and this producer's capabilities.""" self.service = service - self.hands_off_finalize = hands_off_finalize + self.defers_finalize = defers_finalize self.terminal_finalize = terminal_finalize - def attention_input(self, layer) -> tuple: + def attn_input_fusions(self, plan) -> tuple: return ( - partial(self._finalize_output_and_update_and_read_residual, layer), - partial(self._reduce_output_and_update_and_read_residual, layer), + partial(self._finalize_add_norm, plan), + partial(self._all_reduce_add_norm, plan), ) - def ffn_input(self, layer) -> tuple: + def ffn_input_fusions(self, plan) -> tuple: parallel = get_parallel() if ( - TokenAxis.ATTN_TP_SCATTER not in layer.incoming_residual_rows.sharded + TokenAxis.ATTN_TP not in plan.incoming_residual_rows.sharded # The workspace sums over TP, which is then the attention-TP group. and parallel.attn_tp_size == parallel.tp_size - and _fused_norm_gamma(layer.norm) is not None + and _fused_norm_gamma(plan.norm) is not None ): return ( - FusedMlpInput( + FfnInputFusion( completes=SumGroup.ATTN_TP, run=partial( - self._mlp_input_reduce_output_and_update_and_read_residual, - layer, + self._ffn_input_all_reduce_add_norm, + plan, ), ), ) return () - def _finalize_output_and_update_and_read_residual( - self, layer, owed, residual, forward_batch, post_residual_addition + def _finalize_add_norm( + self, plan, owed, residual, forward_batch, post_residual_addition ): """Finish a deferred MoE finalize with the all-reduce, residual add and input norm, when this layer's kernel takes the batch; otherwise the handoff is completed by its producer's own tail.""" - if not isinstance(owed, MoeFinalizeHandoff): + if not isinstance(owed, MoeDeferredFinalize): return None - gamma = _fused_norm_gamma(layer.norm) - if gamma is None or not self._should_use_finalize(layer, forward_batch, owed.m): + gamma = _fused_norm_gamma(plan.norm) + if gamma is None or not self._finalize_eligible(plan, forward_batch, owed.m): return None if post_residual_addition is not None: residual = residual + post_residual_addition assert self.service is not None return self.service.finalize(handoff=owed, residual=residual, gamma=gamma) - def _reduce_output_and_update_and_read_residual( - self, layer, owed, residual, forward_batch, post_residual_addition + def _all_reduce_add_norm( + self, plan, owed, residual, forward_batch, post_residual_addition ): """Complete the all-reduce the previous layer left with the residual add and input norm.""" if not isinstance(owed, UnreducedOutput) or not ( - self._can_consume_post_moe_all_reduce( - layer, forward_batch, int(owed.partial.shape[0]) - ) + self._all_reduce_eligible(plan, forward_batch, int(owed.partial.shape[0])) ): return None if post_residual_addition is not None: @@ -283,16 +281,16 @@ def _reduce_output_and_update_and_read_residual( return self.service.all_reduce_residual_rms_norm( local_contribution=owed.partial, residual=residual, - gamma=_fused_norm_gamma(layer.norm), + gamma=_fused_norm_gamma(plan.norm), ) - def _mlp_input_reduce_output_and_update_and_read_residual( - self, layer, hidden_states, residual, forward_batch + def _ffn_input_all_reduce_add_norm( + self, plan, hidden_states, residual, forward_batch ): """The attention output's all-reduce with the residual add and the post-attention norm.""" if not ( - self._common_eligible(layer, forward_batch, int(hidden_states.shape[0])) + self._eligible(plan, forward_batch, int(hidden_states.shape[0])) and residual is not None and not get_exec().comm.enable_quant_communications ): @@ -301,49 +299,44 @@ def _mlp_input_reduce_output_and_update_and_read_residual( return self.service.all_reduce_residual_rms_norm( local_contribution=hidden_states, residual=residual, - gamma=_fused_norm_gamma(layer.norm), + gamma=_fused_norm_gamma(plan.norm), ) - def _should_use_finalize(self, layer, forward_batch: ForwardBatch, m: int) -> bool: + def _finalize_eligible(self, plan, forward_batch: ForwardBatch, m: int) -> bool: return ( - self._common_eligible(layer, forward_batch, m) - and get_parallel().moe_ep_size == 1 + self._eligible(plan, forward_batch, m) and get_parallel().moe_ep_size == 1 ) - def _can_consume_post_moe_all_reduce( - self, layer, forward_batch: ForwardBatch, m: int - ) -> bool: + def _all_reduce_eligible(self, plan, forward_batch: ForwardBatch, m: int) -> bool: """Incoming, and independent of this layer's own successor.""" return ( - self._common_eligible(layer, forward_batch, m) - and _fused_norm_gamma(layer.norm) is not None + self._eligible(plan, forward_batch, m) + and _fused_norm_gamma(plan.norm) is not None and not get_exec().comm.enable_quant_communications ) - def can_defer_all_reduce(self, layer, forward_batch: ForwardBatch) -> bool: + def can_defer_all_reduce(self, plan, forward_batch: ForwardBatch) -> bool: """Preserve the producer's fused sum path when this workspace is usable. The next consumer still selects its kernel and completes the sum with an ordinary all-reduce if that kernel declines the actual input. """ return ( - self._common_eligible( - layer, forward_batch, int(forward_batch.input_ids.shape[0]) - ) + self._eligible(plan, forward_batch, int(forward_batch.input_ids.shape[0])) and not get_exec().comm.enable_quant_communications ) - def can_defer_finalize(self, layer, forward_batch: ForwardBatch) -> bool: + def can_defer_finalize(self, plan, forward_batch: ForwardBatch) -> bool: """Whether this producer can emit a handoff with an unfused fallback.""" return ( - self.hands_off_finalize - and (not layer.terminal or self.terminal_finalize) - and self._should_use_finalize( - layer, forward_batch, int(forward_batch.input_ids.shape[0]) + self.defers_finalize + and (not plan.terminal or self.terminal_finalize) + and self._finalize_eligible( + plan, forward_batch, int(forward_batch.input_ids.shape[0]) ) ) - def _common_eligible(self, layer, forward_batch: ForwardBatch, m: int) -> bool: + def _eligible(self, plan, forward_batch: ForwardBatch, m: int) -> bool: parallel = get_parallel() return bool( self.service is not None @@ -357,8 +350,7 @@ def _common_eligible(self, layer, forward_batch: ForwardBatch, m: int) -> bool: and get_moe_a2a_backend().is_none() and parallel.tp_size > 1 # The FFN runs on the full rows, not each rank's own slice. - and TokenAxis.ATTN_TP_SCATTER - not in layer.fusion_rows(forward_batch).sharded + and TokenAxis.ATTN_TP not in plan.fused_input_rows(forward_batch).sharded ) @@ -376,13 +368,13 @@ def install_cutedsl_fusion( can_defer_finalize: _LayerPredicate, label: str, terminal_finalize: bool = False, -) -> CuteDSLFusionService | None: +) -> CuteDSLWorkspace | None: """One shared workspace handle per fusion-enabled layer, or None. Every entry of ``layers`` carries ``attn_boundary`` and ``ffn_boundary``. terminal_finalize lets the last layer hand its finalize to the final norm; the model must then pass the returned service as - residual_batch.norm(handoff_norm=...). + residual_batch.final_norm(finalize_norm=...). """ if get_flags().moe.in_speculative_scope: # A draft shares the target's process, which holds one workspace. @@ -407,7 +399,7 @@ def install_cutedsl_fusion( f"{norm.variance_epsilon}" ) - service = CuteDSLFusionService( + service = CuteDSLWorkspace( hidden_size=hidden_size, top_k=top_k, rms_epsilon=rms_epsilon, @@ -417,11 +409,11 @@ def install_cutedsl_fusion( fusion = _fusion_of(layer) if fusion is None: continue - hands_off_finalize = bool(can_defer_finalize(layer)) - hands_off += hands_off_finalize + defers_finalize = bool(can_defer_finalize(layer)) + hands_off += defers_finalize fusion.install( service, - hands_off_finalize=hands_off_finalize, + defers_finalize=defers_finalize, terminal_finalize=terminal_finalize, ) logger.info( diff --git a/python/sglang/srt/layers/layer_boundary/layout.py b/python/sglang/srt/layers/layer_boundary/layout.py index 13f964fd5451..d98f8e9074fb 100644 --- a/python/sglang/srt/layers/layer_boundary/layout.py +++ b/python/sglang/srt/layers/layer_boundary/layout.py @@ -44,7 +44,7 @@ class TokenAxis(Enum): ATTN_DP = auto() ATTN_CP = auto() # Each attention-TP rank holds a slice of its group's tokens. - ATTN_TP_SCATTER = auto() + ATTN_TP = auto() class Layout(msgspec.Struct, frozen=True): @@ -78,11 +78,11 @@ class SumGroup(Enum): MOE_OUTPUT = auto() -def enable_moe_dense_fully_dp(): +def is_dense_ffn_fully_dp(): return get_parallel().moe_dense_tp_size == 1 -def _generic_prefill_cp_shards_tokens() -> bool: +def _prefill_cp_shards_tokens() -> bool: """Whether the strategy prefill CP path shards prefill tokens across CP ranks.""" parallel = get_parallel() return parallel.attn_cp_size > 1 and parallel.enable_prefill_cp @@ -116,7 +116,7 @@ def moe_cp_gathered_rows(forward_batch: ForwardBatch) -> Optional[List[int]]: return None -def sparse_moe_gathers_over_moe_cp() -> bool: +def moe_gathers_over_moe_cp() -> bool: """Whether a sparse MoE's input is gathered over the MoE-CP group on a CP extend: a MoE on the TP group under GQA CP whose MoE-CP group is wider than the MoE's data-parallel groups. DSA and MLA CP gather over attention CP @@ -124,17 +124,14 @@ def sparse_moe_gathers_over_moe_cp() -> bool: return ( not is_moe_input_scattered_across_dp_ranks() and is_enable_moe_cp_allgather() - and not _gathers_over_attention_cp() + and not _cp_gathers_over_attn_cp() ) -def moe_cp_gathers_sparse_moe_input(forward_batch: ForwardBatch) -> bool: +def batch_gathers_over_moe_cp(forward_batch: ForwardBatch) -> bool: """Whether a sparse MoE's input is gathered over the MoE-CP group on this batch: a CP extend, for a MoE that gathers there.""" - return ( - sparse_moe_gathers_over_moe_cp() - and moe_cp_gathered_rows(forward_batch) is not None - ) + return moe_gathers_over_moe_cp() and moe_cp_gathered_rows(forward_batch) is not None def _cp_shard_token_rows(forward_batch: ForwardBatch) -> List[int]: @@ -162,11 +159,11 @@ def token_axis_sizes(*, cp_active: bool = False) -> Dict[TokenAxis, int]: return { TokenAxis.ATTN_DP: parallel.attn_dp_size, TokenAxis.ATTN_CP: parallel.attn_cp_size if cp_active else 1, - TokenAxis.ATTN_TP_SCATTER: parallel.attn_tp_size, + TokenAxis.ATTN_TP: parallel.attn_tp_size, } -def _gathers_over_attention_cp() -> bool: +def _cp_gathers_over_attn_cp() -> bool: """Whether a CP extend gathers the FFN input over the attention-CP group in equal shards and takes the output back with a reduce-scatter there: DSA and MLA CP. GQA prefill CP gathers over the MoE-CP group instead.""" @@ -179,7 +176,7 @@ def _batch_shards_over_cp(forward_batch: ForwardBatch) -> bool: never read the CP predicates.""" if not forward_batch.forward_mode.is_context_parallel_extend(): return False - if _gathers_over_attention_cp(): + if _cp_gathers_over_attn_cp(): return dsa_use_prefill_cp(forward_batch) or is_mla_cp_active(forward_batch) return moe_cp_gathered_rows(forward_batch) is not None diff --git a/python/sglang/srt/layers/layer_boundary/ops.py b/python/sglang/srt/layers/layer_boundary/ops.py index 89faea4e08d6..7d10a802772c 100644 --- a/python/sglang/srt/layers/layer_boundary/ops.py +++ b/python/sglang/srt/layers/layer_boundary/ops.py @@ -41,11 +41,11 @@ moe_cp_all_gather_into_tensor, ) from sglang.srt.layers.layer_boundary.adapters.attention import ( - _redistribute_from_attn_tp_shards, - _redistribute_to_attn_tp_shards, + attn_tp_gather, + attn_tp_slice, ) from sglang.srt.layers.layer_boundary.adapters.context_parallel import ( - dsa_cp_reduce_scatter_hidden_states, + attn_cp_reduce_scatter, ) from sglang.srt.layers.layer_boundary.layout import ( Layout, @@ -53,8 +53,8 @@ _cp_shard_token_rows, moe_cp_gathered_rows, ) -from sglang.srt.layers.layer_boundary.residual import StageUpdate -from sglang.srt.layers.layer_boundary.residual.add_norm import ADD +from sglang.srt.layers.layer_boundary.residual import ResidualUpdate +from sglang.srt.layers.layer_boundary.residual.add_norm import PLAIN_ADD from sglang.srt.model_executor.forward_batch_info import ForwardBatch from sglang.srt.runtime_context import get_exec, get_parallel @@ -89,7 +89,7 @@ def tp_slice(hidden_states, residual): return hidden_states, residual -def gather_attention_tp( +def attn_tp_gather_input( hidden_states: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], forward_batch: ForwardBatch, ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: @@ -116,10 +116,10 @@ def gather_attention_tp( gathered_hidden_states.append(output) return tuple(gathered_hidden_states) - return _redistribute_from_attn_tp_shards(hidden_states) + return attn_tp_gather(hidden_states) -def _reduce_and_redistribute_output_to_attn_tp_shards( +def attn_tp_reduce_scatter( hidden_states: torch.Tensor, ) -> torch.Tensor: parallel = get_parallel() @@ -130,7 +130,7 @@ def _reduce_and_redistribute_output_to_attn_tp_shards( return local_hidden_states -def _redistribute_input_to_moe_cp( +def moe_cp_gather( hidden_states: torch.Tensor, rows: List[int], moe_cp_size: int ) -> torch.Tensor: # Zigzag split can produce unequal token counts across CP ranks @@ -150,7 +150,7 @@ def _redistribute_input_to_moe_cp( return output -def _mlp_input_reduce_output( +def attn_tp_all_reduce( hidden_states: torch.Tensor, forward_batch: ForwardBatch, may_quantize: bool = True ) -> torch.Tensor: if ( @@ -162,7 +162,7 @@ def _mlp_input_reduce_output( return attention_tensor_model_parallel_all_reduce(hidden_states) -def _redistribute_input_to_dp( +def dp_gather( hidden_states: torch.Tensor, forward_batch: ForwardBatch, cp_shard_counts: Optional[List[int]] = None, @@ -174,7 +174,7 @@ def _redistribute_input_to_dp( return global_hidden_states -def _reduce_and_redistribute_output_to_dp( +def dp_gather_sum( hidden_states: torch.Tensor, forward_batch: ForwardBatch, cp_shard_counts: Optional[List[int]] = None, @@ -198,15 +198,13 @@ def move_rows( gathered, cut = rows.sharded - to.sharded, to.sharded - rows.sharded if (gathered and cut) or TokenAxis.ATTN_CP in gathered | cut: raise NotImplementedError(f"{rows=} {to=}") - if TokenAxis.ATTN_TP_SCATTER in gathered: - hidden_states = _redistribute_from_attn_tp_shards(hidden_states) + if TokenAxis.ATTN_TP in gathered: + hidden_states = attn_tp_gather(hidden_states) if TokenAxis.ATTN_DP in gathered: - hidden_states = _redistribute_input_to_dp(hidden_states, forward_batch) + hidden_states = dp_gather(hidden_states, forward_batch) if TokenAxis.ATTN_DP in cut: - hidden_states = _to_local_tokens( - _redistribute_output, forward_batch, hidden_states - ) - if TokenAxis.ATTN_TP_SCATTER in cut: + hidden_states = to_dp_local(_dp_scatter_step, forward_batch, hidden_states) + if TokenAxis.ATTN_TP in cut: parallel = get_parallel() hidden_states = hidden_states.tensor_split(parallel.attn_tp_size)[ parallel.attn_tp_rank @@ -214,9 +212,7 @@ def move_rows( return hidden_states -def _tp_all_gather_scattered_rows( - hidden_states: torch.Tensor, forward_batch: ForwardBatch -) -> torch.Tensor: +def tp_gather(hidden_states: torch.Tensor, forward_batch: ForwardBatch) -> torch.Tensor: # Input-scattered attention keeps the same number of tokens on every TP rank. total_tokens = forward_batch.input_ids.shape[0] output = hidden_states.new_empty((total_tokens, hidden_states.shape[-1])) @@ -231,7 +227,7 @@ def _dp_scatter_group() -> GroupCoordinator: return parallel.attn_tp_group -def _reduce_and_redistribute_output_varlen( +def dp_reduce_scatterv( local_hidden_states: torch.Tensor, hidden_states: torch.Tensor, forward_batch: ForwardBatch, @@ -243,7 +239,7 @@ def _reduce_and_redistribute_output_varlen( ) -def _reduce_and_redistribute_output_max_len( +def dp_reduce_scatter( local_hidden_states: torch.Tensor, hidden_states: torch.Tensor, forward_batch: ForwardBatch, @@ -251,7 +247,7 @@ def _reduce_and_redistribute_output_max_len( dp_reduce_scatter_tensor(local_hidden_states, hidden_states) -def _redistribute_output( +def _dp_scatter_step( local_hidden_states: torch.Tensor, hidden_states: torch.Tensor, forward_batch: ForwardBatch, @@ -259,9 +255,7 @@ def _redistribute_output( dp_scatter(local_hidden_states, hidden_states, forward_batch) -def _redistribute_output_from_moe_cp( - hidden_states: torch.Tensor, rows: List[int] -) -> torch.Tensor: +def moe_cp_take_back(hidden_states: torch.Tensor, rows: List[int]) -> torch.Tensor: moe_cp_rank = get_moe_cp_rank() # The allgather was padded to max_tokens_per_rank (equal chunks). # Extract this rank's actual (non-padded) tokens from its chunk. @@ -272,7 +266,7 @@ def _redistribute_output_from_moe_cp( ).contiguous() -def _to_local_tokens( +def to_dp_local( step: Callable[[torch.Tensor, torch.Tensor, ForwardBatch], None], forward_batch: ForwardBatch, hidden_states: torch.Tensor, @@ -282,15 +276,13 @@ def _to_local_tokens( return local_hidden_states -def _all_reduce_then_to_local_tokens( +def all_reduce_to_dp_local( group: GroupCoordinator, forward_batch: ForwardBatch, hidden_states: torch.Tensor ) -> torch.Tensor: - return _to_local_tokens( - _redistribute_output, forward_batch, group.all_reduce(hidden_states) - ) + return to_dp_local(_dp_scatter_step, forward_batch, group.all_reduce(hidden_states)) -def identity_output( +def keep_output( hidden_states: torch.Tensor, residual: torch.Tensor, forward_batch: ForwardBatch, @@ -299,25 +291,25 @@ def identity_output( return hidden_states, residual -def update_and_gather( +def update_attn_tp_gather_output( hidden_states: torch.Tensor, residual: torch.Tensor, forward_batch: ForwardBatch, - update: StageUpdate = ADD, + update: ResidualUpdate = PLAIN_ADD, **kwargs, ): hidden_states = update.update(hidden_states, residual) - return _redistribute_from_attn_tp_shards(hidden_states), None + return attn_tp_gather(hidden_states), None -def output_on_residual_shard( +def residual_slice_output( hidden_states: torch.Tensor, residual: torch.Tensor, forward_batch: ForwardBatch, *, sums: bool, gathers_back: bool, - update: StageUpdate, + update: ResidualUpdate, **kwargs, ): """Bring the FFN output onto the slice of the rows the residual is on: @@ -339,16 +331,16 @@ def output_on_residual_shard( return hidden_states, None -def scatter_output( +def attn_tp_slice_output( hidden_states: torch.Tensor, residual: torch.Tensor, forward_batch: ForwardBatch, ): assert residual is None, "not yet handled residual!=None" - return _redistribute_to_attn_tp_shards(hidden_states), None + return attn_tp_slice(hidden_states), None -def take_back_attention_cp_shard( +def attn_cp_take_back_output( hidden_states: torch.Tensor, residual: torch.Tensor, forward_batch: ForwardBatch, @@ -361,7 +353,7 @@ def take_back_attention_cp_shard( return shard, residual -def reduce_scatter_over_cp( +def attn_cp_reduce_scatter_output( hidden_states: torch.Tensor, residual: torch.Tensor, forward_batch: ForwardBatch, @@ -369,10 +361,10 @@ def reduce_scatter_over_cp( ): """DSA and MLA CP: sum the FFN output over the attention-CP group and keep this rank's shard.""" - return dsa_cp_reduce_scatter_hidden_states(hidden_states), residual + return attn_cp_reduce_scatter(hidden_states), residual -def take_back_cp_shard( +def dp_cp_take_back_output( hidden_states: torch.Tensor, residual: torch.Tensor, forward_batch: ForwardBatch, @@ -393,7 +385,7 @@ def take_back_cp_shard( return local_hidden_states, residual -def scatter_moe_cp_output( +def moe_cp_take_back_output( hidden_states: torch.Tensor, residual: torch.Tensor, forward_batch: ForwardBatch, @@ -409,15 +401,13 @@ def scatter_moe_cp_output( If DP>1, further scatter back to the local DP slice. """ # Only scatter back during prefill; decode was never allgathered so no-op. - # Safe w.r.t. empty tensors: same reasoning as _mlp_input_gather_moe_cp + # Safe w.r.t. empty tensors: same reasoning as _then_moe_cp_gather # — CP extend always has non-zero tokens per rank, and decode skips this path. rows = moe_cp_gathered_rows(forward_batch) if rows is not None: - hidden_states = _redistribute_output_from_moe_cp(hidden_states, rows) + hidden_states = moe_cp_take_back(hidden_states, rows) if get_parallel().attn_dp_size > 1: - hidden_states = _to_local_tokens( - _redistribute_output, forward_batch, hidden_states - ) + hidden_states = to_dp_local(_dp_scatter_step, forward_batch, hidden_states) return hidden_states, residual diff --git a/python/sglang/srt/layers/layer_boundary/output.py b/python/sglang/srt/layers/layer_boundary/output.py index 4966f60981de..f2c9c1532222 100644 --- a/python/sglang/srt/layers/layer_boundary/output.py +++ b/python/sglang/srt/layers/layer_boundary/output.py @@ -43,42 +43,42 @@ class UnreducedOutput(msgspec.Struct, frozen=True): Fields: partial: Tensor containing this rank's contribution to the sum. group: All-reduce group when no redistribution callable is supplied. - reduce_and_redistribute: Callable(partial) completing the sum and moving + reduce_to_dp_local: Callable(partial) completing the sum and moving it to destination rows; takes precedence over group. - Exits hand this form to ResidualStream.leave(), which exposes an opaque - OwedOutput to models. Low-level adapters use reduce_output() before reading - it. A group is required when reduce_and_redistribute is absent. + Exits hand this form to ResidualStream.record(), which exposes an opaque + OwedOutput to models. Low-level adapters use complete_owed() before reading + it. A group is required when reduce_to_dp_local is absent. """ partial: torch.Tensor group: Optional[GroupCoordinator] = None # Under attention DP: the reduction that also brings ``partial`` back to this # rank's tokens (a reduce-scatter, or an all-reduce then a scatter). - reduce_and_redistribute: Optional[Callable[[torch.Tensor], torch.Tensor]] = None + reduce_to_dp_local: Optional[Callable[[torch.Tensor], torch.Tensor]] = None -class HandoffOutput(msgspec.Struct, frozen=True): +class DeferredFinalize(msgspec.Struct, frozen=True): """A layer output that still owes work only its producer knows how to do (a MoE's finalize and sum), left for the next layer's input or for a terminal - norm that accepts it (residual_batch.norm(handoff_norm=...)). A fused kernel + norm that accepts it (residual_batch.final_norm(finalize_norm=...)). A fused kernel there may do that work together with its own; anything else passes it - through reduce_output(), which calls ``complete()``.""" + through complete_owed(), which calls ``complete()``.""" def complete(self) -> torch.Tensor: """Do the owed work, unfused, and return the complete output.""" raise NotImplementedError -def reduce_output( - hidden_states: Union[torch.Tensor, UnreducedOutput, HandoffOutput, None], +def complete_owed( + hidden_states: Union[torch.Tensor, UnreducedOutput, DeferredFinalize, None], ) -> Optional[torch.Tensor]: - """Run the work an UnreducedOutput or a HandoffOutput still owes; pass + """Run the work an UnreducedOutput or a DeferredFinalize still owes; pass anything else through.""" if isinstance(hidden_states, UnreducedOutput): - if hidden_states.reduce_and_redistribute is not None: - return hidden_states.reduce_and_redistribute(hidden_states.partial) + if hidden_states.reduce_to_dp_local is not None: + return hidden_states.reduce_to_dp_local(hidden_states.partial) return hidden_states.group.all_reduce(hidden_states.partial) - if isinstance(hidden_states, HandoffOutput): + if isinstance(hidden_states, DeferredFinalize): return hidden_states.complete() return hidden_states diff --git a/python/sglang/srt/layers/layer_boundary/prepare.py b/python/sglang/srt/layers/layer_boundary/prepare.py index e923bcd6a872..20d35e9b6247 100644 --- a/python/sglang/srt/layers/layer_boundary/prepare.py +++ b/python/sglang/srt/layers/layer_boundary/prepare.py @@ -32,12 +32,12 @@ ) from sglang.srt.layers.layer_boundary.adapters.attention import ( AttentionInputs, - _redistribute_from_attn_tp_shards, - _redistribute_to_attn_tp_shards, + attn_tp_gather, + attn_tp_slice, get_attn_tp_context, ) from sglang.srt.layers.layer_boundary.adapters.context_parallel import ( - dsa_cp_gather_hidden_states, + attn_cp_gather, ) from sglang.srt.layers.layer_boundary.layout import ( SumGroup, @@ -45,12 +45,12 @@ moe_cp_gathered_rows, ) from sglang.srt.layers.layer_boundary.output import ( - HandoffOutput, + DeferredFinalize, UnreducedOutput, - reduce_output, + complete_owed, ) -from sglang.srt.layers.layer_boundary.residual import StageRead, StageUpdate -from sglang.srt.layers.layer_boundary.residual.add_norm import ADD, NORM_READ +from sglang.srt.layers.layer_boundary.residual import ResidualReadout, ResidualUpdate +from sglang.srt.layers.layer_boundary.residual.add_norm import NORM_READOUT, PLAIN_ADD from sglang.srt.layers.layer_boundary.residual.stream import DeclaredSum from sglang.srt.model_executor.forward_batch_info import ForwardBatch from sglang.srt.runtime_context import get_parallel @@ -65,68 +65,68 @@ from sglang.srt.hardware_backend.npu.cmo import prepare_weight_cache from sglang.srt.layers.layer_boundary.ops import ( - _mlp_input_reduce_output, - _redistribute_input_to_dp, - _redistribute_input_to_moe_cp, - _reduce_and_redistribute_output_to_attn_tp_shards, - _reduce_and_redistribute_output_to_dp, - _tp_all_gather_scattered_rows, + attn_tp_all_reduce, + attn_tp_reduce_scatter, + dp_gather, + dp_gather_sum, + moe_cp_gather, + tp_gather, ) -def _mlp_input_without_dp( +def _reduce_update_read( hidden_states: torch.Tensor, residual: torch.Tensor, forward_batch: ForwardBatch, - layernorm: torch.nn.Module, + norm: torch.nn.Module, *, cache=None, gathers_residual: bool, fusions: Tuple[Callable, ...], group: SumGroup = SumGroup.ATTN_TP, - read: StageRead = NORM_READ, - update: StageUpdate = ADD, + read: ResidualReadout = NORM_READOUT, + update: ResidualUpdate = PLAIN_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) + residual = update.gather_residual_attn_tp(residual) for fused in fusions: result = fused(hidden_states, residual, forward_batch) if result is not None: return result 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 + hidden_states = attn_tp_all_reduce( + hidden_states, forward_batch, may_quantize=update.is_plain_add ) elif group is SumGroup.TP: hidden_states = tensor_model_parallel_all_reduce(hidden_states) if _is_npu and cache is not None: _ = prepare_weight_cache(hidden_states, cache) - return read.update_and_read(update, hidden_states, residual, layernorm) + return read.update_and_read(update, hidden_states, residual, norm) -def _mlp_input_dp_replicate( +def _reduce_update_read_dp_gather( hidden_states: torch.Tensor, residual: torch.Tensor, forward_batch: ForwardBatch, - layernorm: torch.nn.Module, + norm: torch.nn.Module, *, cache=None, gathers_residual: bool, reduces_attention_tp: bool, places_cp_shards: bool = False, - read: StageRead = NORM_READ, - update: StageUpdate = ADD, + read: ResidualReadout = NORM_READOUT, + update: ResidualUpdate = PLAIN_ADD, ): """Attention DP: complete the attention-TP sum if it is owed, write the output into the residual and read the FFN input locally, then gather. With ``places_cp_shards`` each CP rank puts its shard of the DP group's tokens beside the others' in the group's slot.""" if gathers_residual: - residual = update.residual_from_attn_tp_shards(residual) + residual = update.gather_residual_attn_tp(residual) if hidden_states.shape[0] != 0: if reduces_attention_tp: hidden_states = attention_tensor_model_parallel_all_reduce(hidden_states) @@ -135,48 +135,46 @@ def _mlp_input_dp_replicate( disabled=not is_allocation_symmetric(), ): hidden_states, residual = read.update_and_read( - update, hidden_states, residual, layernorm + update, hidden_states, residual, norm ) else: hidden_states, residual = read.update_and_read( - update, hidden_states, residual, layernorm + update, hidden_states, residual, norm ) cp_shard_counts = _cp_shard_token_rows(forward_batch) if places_cp_shards else None return ( - _redistribute_input_to_dp(hidden_states, forward_batch, cp_shard_counts), + dp_gather(hidden_states, forward_batch, cp_shard_counts), residual, ) -def _mlp_input_dp_partial( +def _dp_gather_sum_read( hidden_states: torch.Tensor, residual: torch.Tensor, forward_batch: ForwardBatch, - layernorm: torch.nn.Module, + norm: torch.nn.Module, *, cache=None, gathers_residual: bool, places_cp_shards: bool = False, - read: StageRead = NORM_READ, - update: StageUpdate = ADD, + read: ResidualReadout = NORM_READOUT, + update: ResidualUpdate = PLAIN_ADD, ): """Attention DP: one rank adds the residual, the gather sums it, then read the input from the sum. With ``places_cp_shards`` each CP rank puts its shard of the DP group's tokens beside the others' in the group's slot.""" if gathers_residual: - residual = _redistribute_from_attn_tp_shards(residual) + residual = attn_tp_gather(residual) if get_parallel().attn_tp_rank == 0: hidden_states += residual cp_shard_counts = _cp_shard_token_rows(forward_batch) if places_cp_shards else None - hidden_states = _reduce_and_redistribute_output_to_dp( - hidden_states, forward_batch, cp_shard_counts - ) + hidden_states = dp_gather_sum(hidden_states, forward_batch, cp_shard_counts) dp_scatter(residual, hidden_states, forward_batch, cp_shard_counts) - hidden_states, _ = read.read(hidden_states, layernorm) + hidden_states, _ = read.read(hidden_states, norm) return hidden_states, residual -def _hand_qkv_hook_its_input( +def _attn_input_default( hidden_states: torch.Tensor, forward_batch: ForwardBatch, qkv_latent_func: Optional[Callable], @@ -189,7 +187,7 @@ def _hand_qkv_hook_its_input( return hidden_states -def _hand_scattered_input_to_attention( +def _attn_input_scattered( hidden_states: torch.Tensor, forward_batch: ForwardBatch, qkv_latent_func: Optional[Callable], @@ -200,7 +198,7 @@ def _hand_scattered_input_to_attention( told they are.""" ctx = get_attn_tp_context() if ctx.is_dsa or qkv_latent_func is None: - hidden_states = _tp_all_gather_scattered_rows(hidden_states, forward_batch) + hidden_states = tp_gather(hidden_states, forward_batch) if qkv_latent_func is not None: ctx.set_attn_inputs( AttentionInputs( @@ -213,36 +211,36 @@ def _hand_scattered_input_to_attention( return hidden_states -def _dispatch_consumer( - hidden_states, residual, forward_batch, norm, *, paths, update=ADD, **call +def _dispatch_by_update( + hidden_states, residual, forward_batch, norm, *, paths, update=PLAIN_ADD, **call ): if residual is None: # A missing residual cannot join a partial sum. The non-plain path - # performs enter/read without that optimization. + # performs init_residual/read without that optimization. prepare = paths[False] else: try: - prepare = paths[update.adds_plainly] + prepare = paths[update.is_plain_add] except KeyError: raise RuntimeError("producer update has no bound input path") from None return prepare(hidden_states, residual, forward_batch, norm, update=update, **call) -def _consumer_step( - hidden_states: Union[torch.Tensor, UnreducedOutput, HandoffOutput], +def _run_entry( + hidden_states: Union[torch.Tensor, UnreducedOutput, DeferredFinalize], residual: Optional[torch.Tensor], forward_batch: ForwardBatch, norm: torch.nn.Module, *, step: Callable, - adds_plainly: bool, + is_plain_add: bool, carried_fusions: Tuple[Callable, ...], expected_sum: Optional[SumGroup] = None, completed_step: Optional[Callable] = None, pending=None, written_step: Optional[Callable] = None, written: bool = False, - update: StageUpdate = ADD, + update: ResidualUpdate = PLAIN_ADD, **call, ): """A boundary's half into a stage: complete what the value carries, a sum @@ -253,7 +251,7 @@ def _consumer_step( whose declared sum is already completed uses ``completed_step`` instead; a pending declared sum must match ``expected_sum``. ``call`` is what the stage's read takes (the attention's ``quant_format`` and ``post_residual_addition``).""" - if residual is not None and update.adds_plainly != adds_plainly: + if residual is not None and update.is_plain_add != is_plain_add: raise RuntimeError("producer update does not match the boundary's capability") if pending is not None: if isinstance(pending.owed, DeclaredSum): @@ -274,11 +272,8 @@ def _consumer_step( 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) + if isinstance(owed, UnreducedOutput) and owed.reduce_to_dp_local is not None: + hidden_states = complete_owed(owed) else: for fused in carried_fusions: result = fused( @@ -286,36 +281,36 @@ def _consumer_step( ) if result is not None: return result - if isinstance(owed, HandoffOutput) or owed.partial.shape[0] != 0: - hidden_states = reduce_output(owed) + if isinstance(owed, DeferredFinalize) or owed.partial.shape[0] != 0: + hidden_states = complete_owed(owed) else: hidden_states = owed.partial return step(hidden_states, residual, forward_batch, norm, update=update, **call) -def _read_input( +def _update_read( hidden_states: torch.Tensor, residual: Optional[torch.Tensor], forward_batch: ForwardBatch, norm: torch.nn.Module, *, cache=None, - layer_input: Optional[Callable], + pre_move: Optional[Callable], enters_stack: bool, - read: StageRead, - update: StageUpdate = ADD, + read: ResidualReadout, + update: ResidualUpdate = PLAIN_ADD, quant_format: str = "", post_residual_addition: Optional[torch.Tensor] = None, ): - """Complete what the input owes by construction (``layer_input``), then + """Complete what the input owes by construction (``pre_move``), 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) + if pre_move is not None: + hidden_states, residual = pre_move(hidden_states, residual) if enters: - hidden_states, residual = read.enter(hidden_states), None + hidden_states, residual = read.init_residual(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) @@ -329,50 +324,50 @@ def _read_input( ) -def _mlp_input_scatter( +def _attn_tp_reduce_scatter_update_read( hidden_states: torch.Tensor, residual: torch.Tensor, forward_batch: ForwardBatch, - layernorm: torch.nn.Module, + norm: torch.nn.Module, *, cache=None, scatters_residual: bool, - read: StageRead = NORM_READ, - update: StageUpdate = ADD, + read: ResidualReadout = NORM_READOUT, + update: ResidualUpdate = PLAIN_ADD, ): - hidden_states = _reduce_and_redistribute_output_to_attn_tp_shards(hidden_states) + hidden_states = attn_tp_reduce_scatter(hidden_states) if scatters_residual: - residual = update.residual_to_attn_tp_shard(residual) - return read.update_and_read(update, hidden_states, residual, layernorm) + residual = update.slice_residual_attn_tp(residual) + return read.update_and_read(update, hidden_states, residual, norm) -def _mlp_input_slice( +def _attn_tp_slice_update_read( hidden_states: torch.Tensor, residual: Optional[torch.Tensor], forward_batch: ForwardBatch, - layernorm: torch.nn.Module, + norm: torch.nn.Module, *, cache=None, scatters_residual: bool, - read: StageRead = NORM_READ, - update: StageUpdate = ADD, + read: ResidualReadout = NORM_READOUT, + update: ResidualUpdate = PLAIN_ADD, ): - hidden_states = _redistribute_to_attn_tp_shards(hidden_states).clone() + hidden_states = attn_tp_slice(hidden_states).clone() if scatters_residual and residual is not None: - residual = update.residual_to_attn_tp_shard(residual) - return read.update_and_read(update, hidden_states, residual, layernorm) + residual = update.slice_residual_attn_tp(residual) + return read.update_and_read(update, hidden_states, residual, norm) -def _mlp_input_on_residual_shard( +def _tp_reduce_scatter_update_read_gather( hidden_states: torch.Tensor, residual: torch.Tensor, forward_batch: ForwardBatch, - layernorm: torch.nn.Module, + norm: torch.nn.Module, *, cache=None, reduces: bool = True, - read: StageRead, - update: StageUpdate, + read: ResidualReadout, + update: ResidualUpdate, ): """The residual stays on each rank's slice while the FFN takes the full rows: reduce-scatter the attention output onto the slice, which completes @@ -386,54 +381,52 @@ def _mlp_input_on_residual_shard( parallel.tp_group.reduce_scatter_tensor(shard, hidden_states) else: shard = shard.clone() - shard, residual = read.update_and_read(update, shard, residual, layernorm) + shard, residual = read.update_and_read(update, shard, residual, norm) attn_tp_all_gather_into_tensor(hidden_states, shard) return hidden_states, residual -def _mlp_input_residual_into_sum( +def _tp_sum_with_residual_read( hidden_states: torch.Tensor, residual: torch.Tensor, forward_batch: ForwardBatch, - layernorm: torch.nn.Module, + norm: torch.nn.Module, *, cache=None, - read: StageRead = NORM_READ, - update: StageUpdate = ADD, + read: ResidualReadout = NORM_READOUT, + update: ResidualUpdate = PLAIN_ADD, ): - return _tp_all_reduce_with_scattered_residual( - hidden_states, residual, layernorm, read - ) + return _tp_all_reduce_with_scattered_residual(hidden_states, residual, norm, read) -def _mlp_input_gather_attention_cp( +def _then_attn_cp_gather( hidden_states: torch.Tensor, residual: torch.Tensor, forward_batch: ForwardBatch, - layernorm: torch.nn.Module, + norm: torch.nn.Module, *, cache=None, gather: Callable, - update: StageUpdate = ADD, + update: ResidualUpdate = PLAIN_ADD, ): """DSA and MLA CP: complete this rank's shard, then gather the shards, of equal length, over the attention-CP group. The residual stays on the shard.""" hidden_states, residual = gather( - hidden_states, residual, forward_batch, layernorm, update=update, cache=cache + hidden_states, residual, forward_batch, norm, update=update, cache=cache ) - return dsa_cp_gather_hidden_states(hidden_states), residual + return attn_cp_gather(hidden_states), residual -def _mlp_input_gather_moe_cp( +def _then_moe_cp_gather( hidden_states: torch.Tensor, residual: torch.Tensor, forward_batch: ForwardBatch, - layernorm: torch.nn.Module, + norm: torch.nn.Module, *, cache=None, gather: Callable, - update: StageUpdate = ADD, + update: ResidualUpdate = PLAIN_ADD, ): """Gather for the FFN, then over the MoE-CP group so each rank holds all tokens of its MoE group (moe_dp_size < attn_cp_size). The residual stays on @@ -447,14 +440,12 @@ def _mlp_input_gather_moe_cp( return hidden_states, residual hidden_states, residual = gather( - hidden_states, residual, forward_batch, layernorm, update=update, cache=cache + hidden_states, residual, forward_batch, norm, update=update, cache=cache ) rows = moe_cp_gathered_rows(forward_batch) if rows is not None and hidden_states.shape[0] > 0: - hidden_states = _redistribute_input_to_moe_cp( - hidden_states, rows, get_moe_cp_size() - ) + hidden_states = moe_cp_gather(hidden_states, rows, get_moe_cp_size()) return hidden_states, residual @@ -462,8 +453,8 @@ def _mlp_input_gather_moe_cp( def _tp_all_reduce_with_scattered_residual( hidden_states: torch.Tensor, residual: torch.Tensor, - layernorm: torch.nn.Module, - read: StageRead = NORM_READ, + norm: torch.nn.Module, + read: ResidualReadout = NORM_READOUT, ): parallel = get_parallel() if hidden_states.shape[0] == 0: @@ -472,4 +463,4 @@ def _tp_all_reduce_with_scattered_residual( scattered_states = hidden_states.tensor_split(parallel.tp_size)[parallel.tp_rank] scattered_states += residual residual = tensor_model_parallel_all_reduce(hidden_states) - return read.read(residual, layernorm) + return read.read(residual, norm) diff --git a/python/sglang/srt/layers/layer_boundary/residual/__init__.py b/python/sglang/srt/layers/layer_boundary/residual/__init__.py index 19c12d9882d0..c24705c6f650 100644 --- a/python/sglang/srt/layers/layer_boundary/residual/__init__.py +++ b/python/sglang/srt/layers/layer_boundary/residual/__init__.py @@ -21,31 +21,31 @@ import torch -class StageUpdate(Protocol): +class ResidualUpdate(Protocol): """Producer-owned operation writing a contribution into the residual. Fields: - adds_plainly: Whether update is plain addition, permitting compatible + is_plain_add: Whether update is plain addition, permitting compatible add+norm fusion and adding the residual on one rank before a sum. - at_producer: Whether the producer writes the residual at its exit rather + applied_at_exit: Whether the producer writes the residual at its exit rather than leaving the operation for the next prepare. - can_defer_across_layers: Whether parameters/state remain valid after + outlives_layer: Whether parameters/state remain valid after leaving the producer layer, including any offload or state reuse. The shard conversion methods must move any update-associated state together - with the residual. Nonlinear updates must not claim adds_plainly. + with the residual. Nonlinear updates must not claim is_plain_add. """ # A plain add, which one rank may run before the sum it adds into # completes (the DP partial order, a residual that joins the attention # output's sum), and which a fused add + norm kernel may run. - adds_plainly: bool + is_plain_add: bool # The stage writes its output into the residual itself at its end, instead # of leaving that to the next stage's read. - at_producer: bool + applied_at_exit: bool # Its parameters and state remain valid after leaving the producer layer. # Stateful or offloaded implementations must not opt in without that guarantee. - can_defer_across_layers: bool + outlives_layer: bool def update(self, hidden_states, residual) -> torch.Tensor: """Write hidden_states, the producer contribution, into residual. @@ -55,36 +55,37 @@ def update(self, hidden_states, residual) -> torch.Tensor: This operation does not perform the consumer's normalization. """ - def residual_to_attn_tp_shard(self, residual) -> torch.Tensor: + def slice_residual_attn_tp(self, residual) -> torch.Tensor: """This attention-TP rank's slice of the residual, with whatever the update reads along with it.""" - def residual_from_attn_tp_shards(self, residual) -> torch.Tensor: + def gather_residual_attn_tp(self, residual) -> torch.Tensor: """The residual gathered from every attention-TP rank's slice.""" -class StageRead(Protocol): +class ResidualReadout(Protocol): """Consumer-owned operation deriving compute input from the residual. Fields: - norms_plainly: Read is normalization/optional quantization without changing + is_plain_norm: Read is normalization/optional quantization without changing the residual, allowing compatible fused add+norm implementations. - before_gather: Preserve this read on source rows before a DP gather. + reads_before_dp_gather: Preserve this read on source rows before a DP gather. - enter initializes the stack residual. read consumes an already-written - residual; update_and_read first applies the actual producer's update. - Both reads return (compute_input, residual), preserving their kernel's - rounding order rather than normalizing a separately rounded snapshot. + init_residual initializes the stack residual. read consumes an + already-written residual; update_and_read first applies the actual + producer's update. Both reads return (compute_input, residual), preserving + their kernel's rounding order rather than normalizing a separately rounded + snapshot. """ # The input is the residual's norm, in the quantization the call asks for, # and the residual is left as it is: what a fused add + norm kernel computes, # and what may run on the rows a sum completes onto. - norms_plainly: bool + is_plain_norm: bool # Preserve the read on the source rows before a DP gather. - before_gather: bool + reads_before_dp_gather: bool - def enter(self, hidden_states) -> torch.Tensor: + def init_residual(self, hidden_states) -> torch.Tensor: """The residual the layer stack starts from, given its input.""" def read( @@ -110,7 +111,7 @@ def read( def update_and_read( self, - update: StageUpdate, + update: ResidualUpdate, hidden_states, residual, norm, @@ -120,7 +121,7 @@ def update_and_read( """Apply a producer update and read input while preserving kernel ordering. Args: - update: Actual producer's StageUpdate, not a consumer-inferred operation. + update: Actual producer's ResidualUpdate, not a consumer-inferred operation. hidden_states: Producer contribution after required communication. residual: Previous residual, or None at entry where supported. norm: Consumer normalization module. @@ -133,11 +134,11 @@ def update_and_read( """ -class LayerResidual(NamedTuple): +class LayerResidualOps(NamedTuple): """The reads and updates a decoder layer declares for its attention and its FFN.""" - attention_read: StageRead - attention_update: StageUpdate - ffn_read: StageRead - ffn_update: StageUpdate + attn_readout: ResidualReadout + attn_update: ResidualUpdate + ffn_readout: ResidualReadout + ffn_update: ResidualUpdate diff --git a/python/sglang/srt/layers/layer_boundary/residual/access.py b/python/sglang/srt/layers/layer_boundary/residual/access.py index 10ab220c96bb..9210e6f1323a 100644 --- a/python/sglang/srt/layers/layer_boundary/residual/access.py +++ b/python/sglang/srt/layers/layer_boundary/residual/access.py @@ -19,16 +19,16 @@ from sglang.srt.distributed import GroupCoordinator from sglang.srt.layers.layer_boundary.output import ( - HandoffOutput, + DeferredFinalize, UnreducedOutput, - reduce_output, + complete_owed, ) from sglang.srt.layers.layer_boundary.residual.stream import OwedOutput, ResidualStream from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors def buffer( - hidden_states: Union[torch.Tensor, UnreducedOutput, HandoffOutput], + hidden_states: Union[torch.Tensor, UnreducedOutput, DeferredFinalize], ) -> Optional[torch.Tensor]: """Storage that can be reused after prepare consumes the input. A finalize handoff has no reusable layer-output tensor. This does not complete or read @@ -37,7 +37,7 @@ def buffer( return hidden_states.contribution.value if isinstance(hidden_states, UnreducedOutput): return hidden_states.partial - if isinstance(hidden_states, HandoffOutput): + if isinstance(hidden_states, DeferredFinalize): return None return hidden_states @@ -48,7 +48,7 @@ def add_to_output(hidden_states, residual, extra: torch.Tensor) -> Tuple: hidden_states = ( residual.complete(hidden_states) if isinstance(residual, ResidualStream) - else reduce_output(hidden_states) + else complete_owed(hidden_states) ) hidden_states.add_(extra) return hidden_states, residual @@ -59,10 +59,10 @@ def fold(hidden_states, residual): The following stage receives it with no outstanding residual addition.""" stream = residual if isinstance(residual, ResidualStream) else None if stream is not None: - if stream.pending is not None and not stream.pending.update.adds_plainly: + if stream.pending is not None and not stream.pending.update.is_plain_add: raise NotImplementedError("fold requires a plain residual update") - hidden_states, residual = stream.finish(hidden_states) - hidden_states = reduce_output(hidden_states) + hidden_states, residual = stream.export(hidden_states) + hidden_states = complete_owed(hidden_states) if residual is not None: hidden_states = hidden_states + residual return ( @@ -77,33 +77,33 @@ def written(hidden_states): return hidden_states, ResidualStream(hidden_states) -def finish_layer_stack( - hidden_states: Union[torch.Tensor, UnreducedOutput, HandoffOutput], +def export_output( + hidden_states: Union[torch.Tensor, UnreducedOutput, DeferredFinalize], residual: Optional[torch.Tensor], forward_batch: ForwardBatch, *, final_norm_takes_handoff: bool = False, preserve_declared: bool = False, -) -> Tuple[Union[torch.Tensor, HandoffOutput], Optional[torch.Tensor]]: +) -> Tuple[Union[torch.Tensor, DeferredFinalize], Optional[torch.Tensor]]: """Complete what this layer left for a next layer, for callers that carry an explicit residual tensor or merge TBO microbatches. Stage-boundary - decoders use residual_batch.norm(), to_pp() or take_output() instead: + decoders use residual_batch.final_norm(), to_pp() or take_output() instead: to_pp() keeps a declared partial sum for the receiving from_pp(), which this helper would complete early. A final norm that does a producer's handoff together with its own work (``final_norm_takes_handoff``) receives it as it is.""" if isinstance(residual, ResidualStream): - return residual.finish( + return residual.export( hidden_states, takes_handoff=final_norm_takes_handoff, preserve_declared=preserve_declared, ) - if final_norm_takes_handoff and isinstance(hidden_states, HandoffOutput): + if final_norm_takes_handoff and isinstance(hidden_states, DeferredFinalize): return hidden_states, residual - return reduce_output(hidden_states), residual + return complete_owed(hidden_states), residual -def norm_output(hidden_states, residual, norm, capture_output=None, **read_kwargs): +def final_norm_pair(hidden_states, residual, norm, capture=None, **read_kwargs): """The final norm and an optional capture of the same updated residual. Keep add and norm together: normalizing a separately rounded residual is @@ -111,12 +111,12 @@ def norm_output(hidden_states, residual, norm, capture_output=None, **read_kwarg capture callback owns retention of the borrowed residual storage. """ if residual is None: - if capture_output is not None: - capture_output(hidden_states) + if capture is not None: + capture(hidden_states) return norm(hidden_states) hidden_states, residual = norm(hidden_states, residual, **read_kwargs) - if capture_output is not None: - capture_output(residual) + if capture is not None: + capture(residual) return hidden_states diff --git a/python/sglang/srt/layers/layer_boundary/residual/add_norm.py b/python/sglang/srt/layers/layer_boundary/residual/add_norm.py index d0710625e680..c1378d05e39f 100644 --- a/python/sglang/srt/layers/layer_boundary/residual/add_norm.py +++ b/python/sglang/srt/layers/layer_boundary/residual/add_norm.py @@ -26,11 +26,11 @@ uses_cutedsl_ar_fusion, ) from sglang.srt.layers.layer_boundary.adapters.attention import ( - _redistribute_from_attn_tp_shards, - _redistribute_to_attn_tp_shards, + attn_tp_gather, + attn_tp_slice, get_attn_tp_context, ) -from sglang.srt.layers.layer_boundary.residual import LayerResidual +from sglang.srt.layers.layer_boundary.residual import LayerResidualOps from sglang.srt.layers.quantization.fp8_utils import ( _use_aiter_bpreshuffle_gfx95, materialize_bpreshuffle_fp8_scale_tuple, @@ -139,7 +139,7 @@ def _fused_rmsnorm_fp8_per_token_quant( FUSE_ALLREDUCE_MAX_BATCH_SIZE = 2048 -def apply_flashinfer_allreduce_fusion(batch_size: int): +def flashinfer_ar_fusion_applies(batch_size: int): return ( # NOTE: flashinfer 0.6.1 caused performance regression on sm100 for allreduce fusion # Ref: https://github.com/sgl-project/sglang/issues/17237 @@ -158,7 +158,7 @@ def apply_flashinfer_allreduce_fusion(batch_size: int): ) -def aiter_all_reduce_fusion_enabled_for(forward_mode: ForwardMode) -> bool: +def aiter_ar_fusion_enabled_for(forward_mode: ForwardMode) -> bool: comm = get_exec().comm if not comm.enable_aiter_allreduce_fusion: return False @@ -167,9 +167,7 @@ def aiter_all_reduce_fusion_enabled_for(forward_mode: ForwardMode) -> bool: return not comm.disable_aiter_allreduce_fusion_in_decode -def apply_aiter_all_reduce_fusion( - input_tensor: torch.Tensor, forward_batch: ForwardBatch -): +def aiter_ar_fusion_applies(input_tensor: torch.Tensor, forward_batch: ForwardBatch): n = input_tensor.shape[-1] total_bytes = input_tensor.numel() * input_tensor.element_size() # Aiter's should_custom_ar uses <= max_size/2 (64 MB); match that boundary. @@ -180,7 +178,7 @@ def apply_aiter_all_reduce_fusion( and total_bytes <= 8 * 1024 * 8192 and get_parallel().tp_size != 6 and not is_dp_attention_enabled() - and aiter_all_reduce_fusion_enabled_for(forward_batch.forward_mode) + and aiter_ar_fusion_enabled_for(forward_batch.forward_mode) ) @@ -250,7 +248,7 @@ def _update_and_read_residual_aiter_fp8_per_token( ) -def _attn_input_update_and_read_residual(quant_format: str): +def _norm_quant_kernel(quant_format: str): """Add the previous layer's output to the residual and read the attention input from it: the input norm, fused with the quantization this format wants. Without a residual (the first layer, or one already folded into @@ -264,22 +262,22 @@ def _attn_input_update_and_read_residual(quant_format: str): return _update_and_read_residual_plain -class Add: +class PlainAdd: """A plain residual: the output is added into it.""" - adds_plainly = True - at_producer = False - can_defer_across_layers = True + is_plain_add = True + applied_at_exit = False + outlives_layer = True def update(self, hidden_states, residual): hidden_states += residual return hidden_states - def residual_to_attn_tp_shard(self, residual): - return _redistribute_to_attn_tp_shards(residual) + def slice_residual_attn_tp(self, residual): + return attn_tp_slice(residual) - def residual_from_attn_tp_shards(self, residual): - return _redistribute_from_attn_tp_shards(residual) + def gather_residual_attn_tp(self, residual): + return attn_tp_gather(residual) class Fp8Input(Enum): @@ -290,25 +288,23 @@ class Fp8Input(Enum): @dataclass(frozen=True) -class NormQuantRead: +class NormQuantReadout: """The input is the residual's norm, fused with the quantization the consumer asks for (``quant_format``) and a post-residual addition; a plain add of the previous output runs in the same kernel. An empty batch skips the norm.""" - norms_plainly = True - before_gather: bool = False + is_plain_norm = True + reads_before_dp_gather: bool = False fp8_input: Optional[Fp8Input] = None - def enter(self, hidden_states): + def init_residual(self, hidden_states): return hidden_states def read(self, residual, norm, quant_format="", post_residual_addition=None): if residual.shape[0] == 0: return residual, residual - return _attn_input_update_and_read_residual(quant_format)( - norm, residual, None, None - ) + return _norm_quant_kernel(quant_format)(norm, residual, None, None) def update_and_read( self, @@ -319,24 +315,24 @@ def update_and_read( quant_format="", post_residual_addition=None, ): - if not update.adds_plainly: + if not update.is_plain_add: return self.read(update.update(hidden_states, residual), norm, quant_format) if hidden_states.shape[0] == 0: return hidden_states, hidden_states - return _attn_input_update_and_read_residual(quant_format)( + return _norm_quant_kernel(quant_format)( norm, hidden_states, residual, post_residual_addition ) @dataclass(frozen=True) -class NormRead: +class NormReadout: """The input is the residual's norm; a plain add of the previous output runs in the same kernel. An empty batch skips the norm.""" - norms_plainly = True - before_gather: bool = False + is_plain_norm = True + reads_before_dp_gather: bool = False - def enter(self, hidden_states): + def init_residual(self, hidden_states): return hidden_states def read(self, residual, norm, quant_format="", post_residual_addition=None): @@ -360,7 +356,7 @@ def update_and_read( if residual is None: # The layer stack starts at this stage: its input is the residual. return self.read(hidden_states, norm, quant_format, post_residual_addition) - if not update.adds_plainly: + if not update.is_plain_add: return self.read(update.update(hidden_states, residual), norm, quant_format) if quant_format or post_residual_addition is not None: raise NotImplementedError( @@ -371,14 +367,14 @@ def update_and_read( return norm(hidden_states, residual) -ADD = Add() -NORM_QUANT_READ = NormQuantRead() -NORM_READ = NormRead() +PLAIN_ADD = PlainAdd() +NORM_QUANT_READOUT = NormQuantReadout() +NORM_READOUT = NormReadout() # A plain residual: the attention reads with its input norm and the quantization # it wants, the FFN with its norm, and each stage's output is added. -PLAIN_RESIDUAL = LayerResidual( - attention_read=NORM_QUANT_READ, - attention_update=ADD, - ffn_read=NORM_READ, - ffn_update=ADD, +PLAIN_RESIDUAL_OPS = LayerResidualOps( + attn_readout=NORM_QUANT_READOUT, + attn_update=PLAIN_ADD, + ffn_readout=NORM_READOUT, + ffn_update=PLAIN_ADD, ) diff --git a/python/sglang/srt/layers/layer_boundary/residual/batch.py b/python/sglang/srt/layers/layer_boundary/residual/batch.py index 06a1f59e29a2..83929dcb6485 100644 --- a/python/sglang/srt/layers/layer_boundary/residual/batch.py +++ b/python/sglang/srt/layers/layer_boundary/residual/batch.py @@ -27,7 +27,7 @@ def start(forward_batch): forward_batch.residual_stream = ResidualStream() -def current(forward_batch): +def stream_of(forward_batch): stream = forward_batch.residual_stream if stream is None: raise RuntimeError("start the layer stack before entering a stage") @@ -36,12 +36,12 @@ def current(forward_batch): def complete_output(hidden_states, forward_batch): """Complete the contribution while retaining its pending residual update.""" - return current(forward_batch).complete(hidden_states) + return stream_of(forward_batch).complete(hidden_states) def take_output(hidden_states, forward_batch): """Release a terminal output whose residual update has already been written.""" - stream = current(forward_batch) + stream = stream_of(forward_batch) stream.check(hidden_states) if stream.pending is not None: raise RuntimeError("write the residual update before taking the final output") @@ -49,13 +49,13 @@ def take_output(hidden_states, forward_batch): return hidden_states -def norm( +def final_norm( hidden_states, forward_batch, layernorm, - capture_output=None, + capture=None, *, - handoff_norm=None, + finalize_norm=None, skip_empty=False, **read_kwargs, ): @@ -65,9 +65,9 @@ def norm( hidden_states: Current stream output or opaque owed handle. forward_batch: Batch owning the residual stream. layernorm: Final norm supporting the model's output/residual pair. - capture_output: Optional callback retaining the same updated residual; + capture: Optional callback retaining the same updated residual; it must copy borrowed storage when retention requires ownership. - handoff_norm: Optional adapter with finalize(handoff, residual, gamma) + finalize_norm: Optional adapter with finalize(handoff, residual, gamma) for a producer-specific finalize handoff; when a handoff arrives it cannot be combined with capture. Without it, a handoff is completed unfused first. Its gamma is layernorm.gemma_weight, so the final @@ -79,29 +79,29 @@ def norm( Normalized tensor. Add and norm remain together to preserve the kernel's accumulation/rounding order; a snapshot is not used as the norm input. """ - hidden_states, residual = current(forward_batch).finish( - hidden_states, takes_handoff=handoff_norm is not None + hidden_states, residual = stream_of(forward_batch).export( + hidden_states, takes_handoff=finalize_norm is not None ) # The terminal consumer now owns the pair. Do not keep layer buffers alive # through logits processing or the next forward on this batch. forward_batch.residual_stream = None - from sglang.srt.layers.layer_boundary.output import HandoffOutput + from sglang.srt.layers.layer_boundary.output import DeferredFinalize - if isinstance(hidden_states, HandoffOutput): + if isinstance(hidden_states, DeferredFinalize): if residual is None: raise RuntimeError("invalid final deferred MoE handoff") - if capture_output is not None: + if capture is not None: raise RuntimeError( "final handoff capture requires an explicit capture adapter" ) - hidden_states, _ = handoff_norm.finalize( + hidden_states, _ = finalize_norm.finalize( handoff=hidden_states, residual=residual, gamma=layernorm.gemma_weight ) return hidden_states if skip_empty and hidden_states.shape[0] == 0: return hidden_states - return access.norm_output( - hidden_states, residual, layernorm, capture_output, **read_kwargs + return access.final_norm_pair( + hidden_states, residual, layernorm, capture, **read_kwargs ) @@ -117,7 +117,7 @@ def to_pp(hidden_states, forward_batch, *, preserve_declared=True): Runtime-selected completion work is finished before transport. """ - hidden_states, residual = current(forward_batch).finish( + hidden_states, residual = stream_of(forward_batch).export( hidden_states, preserve_declared=preserve_declared ) forward_batch.residual_stream = None @@ -125,19 +125,19 @@ def to_pp(hidden_states, forward_batch, *, preserve_declared=True): def snapshot(hidden_states, forward_batch): - return current(forward_batch).snapshot(hidden_states) + return stream_of(forward_batch).snapshot(hidden_states) def add_to_output(hidden_states, forward_batch, extra): - output, _ = access.add_to_output(hidden_states, current(forward_batch), extra) + output, _ = access.add_to_output(hidden_states, stream_of(forward_batch), extra) return output def fold(hidden_states, forward_batch): - output, _ = access.fold(hidden_states, current(forward_batch)) + output, _ = access.fold(hidden_states, stream_of(forward_batch)) return output -def written(hidden_states, forward_batch): +def set_written(hidden_states, forward_batch): forward_batch.residual_stream = ResidualStream(hidden_states) return hidden_states diff --git a/python/sglang/srt/layers/layer_boundary/residual/mhc.py b/python/sglang/srt/layers/layer_boundary/residual/mhc.py index d6916c267dd7..6aa41f9d80b5 100644 --- a/python/sglang/srt/layers/layer_boundary/residual/mhc.py +++ b/python/sglang/srt/layers/layer_boundary/residual/mhc.py @@ -19,7 +19,7 @@ import torch from sglang.kernels.ops.layernorm.mhc import hc_contract, hc_expand -from sglang.srt.layers.layer_boundary.residual import LayerResidual +from sglang.srt.layers.layer_boundary.residual import LayerResidualOps from sglang.srt.runtime_context import get_parallel @@ -28,7 +28,7 @@ class MHCState: """A layer's residual as hyper-connection streams: the residual is hc_mult streams, a stage's output is written in with hc_post and the next stage's input read with hc_pre and the norm. A read produces the h_res / h_post the - next write-back consumes; they move with the tokens. ``layer_residual()`` + next write-back consumes; they move with the tokens. ``residual_ops()`` gives the reads and updates the layer's two stages declare. Parameters belong to the owning layer; this state only holds scratch shared across communication stages.""" @@ -50,7 +50,9 @@ def _resolve_out_norm(out_norm): return None, None return out_norm.weight.data, out_norm.variance_epsilon - def attn_split(self, hidden_states, out_norm: Optional[torch.nn.Module] = None): + def read_attn_input( + self, hidden_states, out_norm: Optional[torch.nn.Module] = None + ): residual = hidden_states out_norm_weight, out_norm_eps = self._resolve_out_norm(out_norm) hidden_states, self.h_res, self.h_post, norm_fused = self.hc_attn_pre( @@ -60,7 +62,7 @@ def attn_split(self, hidden_states, out_norm: Optional[torch.nn.Module] = None): hidden_states = out_norm(hidden_states) return hidden_states, residual - def attn_to_mlp( + def update_and_read_ffn_input( self, hidden_states, residual, out_norm: Optional[torch.nn.Module] = None ): out_norm_weight, out_norm_eps = self._resolve_out_norm(out_norm) @@ -90,97 +92,99 @@ def attn_to_mlp( hidden_states = out_norm(hidden_states) return hidden_states, residual - def mlp_combine(self, hidden_states, residual): + def apply_post(self, hidden_states, residual): return self.hc_post(hidden_states, residual, self.h_res, self.h_post) - def reset_aux(self): + def clear_coefficients(self): self.h_res = None self.h_post = None - def residual_to_attn_tp_shard(self, residual): + def slice_residual_attn_tp(self, residual): parallel = get_parallel() rank, size = parallel.attn_tp_rank, parallel.attn_tp_size self.h_res = self.h_res.tensor_split(size)[rank] self.h_post = self.h_post.tensor_split(size)[rank] return residual.tensor_split(size)[rank] - def residual_from_attn_tp_shards(self, residual): + def gather_residual_attn_tp(self, residual): raise NotImplementedError( "Unsupported: h_res/h_post allgather not implemented." ) - def layer_residual(self) -> LayerResidual: - return LayerResidual( - attention_read=_AttentionRead(self), - attention_update=_AttentionUpdate(self), - ffn_read=_FfnRead(self), + def residual_ops(self) -> LayerResidualOps: + return LayerResidualOps( + attn_readout=_AttnReadout(self), + attn_update=_AttnUpdate(self), + ffn_readout=_FfnReadout(self), ffn_update=_FfnUpdate(self), ) -class _AttentionRead: +class _AttnReadout: """hc_pre and the input norm, from streams that already hold the previous layer's output: an MHC layer takes its input written back.""" - norms_plainly = False - before_gather = False + is_plain_norm = False + reads_before_dp_gather = False def __init__(self, state: MHCState): self.state = state - def enter(self, hidden_states): + def init_residual(self, hidden_states): return hc_expand(hidden_states, self.state.hc_mult) def read(self, residual, norm, quant_format="", post_residual_addition=None): if quant_format: raise NotImplementedError(f"an MHC attention input in {quant_format=}") - return self.state.attn_split(residual, out_norm=norm) + return self.state.read_attn_input(residual, out_norm=norm) def update_and_read(self, update, hidden_states, residual, norm, **kwargs): raise NotImplementedError("an MHC layer takes its input written back") -class _AttentionUpdate: +class _AttnUpdate: """hc_post with the coefficients the attention's read produced. It is not a plain add, so it runs only once the sum it writes in is complete.""" - adds_plainly = False - at_producer = False - can_defer_across_layers = False + is_plain_add = False + applied_at_exit = False + outlives_layer = False def __init__(self, state: MHCState): self.state = state def update(self, hidden_states, residual): - return self.state.mlp_combine(hidden_states, residual) + return self.state.apply_post(hidden_states, residual) - def residual_to_attn_tp_shard(self, residual): - return self.state.residual_to_attn_tp_shard(residual) + def slice_residual_attn_tp(self, residual): + return self.state.slice_residual_attn_tp(residual) - def residual_from_attn_tp_shards(self, residual): - return self.state.residual_from_attn_tp_shards(residual) + def gather_residual_attn_tp(self, residual): + return self.state.gather_residual_attn_tp(residual) -class _FfnRead: +class _FfnReadout: """The attention output's hc_post and the FFN input's hc_pre and norm, fused in hc_ffn_post_pre when it takes the batch.""" - norms_plainly = False - before_gather = False + is_plain_norm = False + reads_before_dp_gather = False def __init__(self, state: MHCState): self.state = state - def enter(self, hidden_states): + def init_residual(self, hidden_states): return hc_expand(hidden_states, self.state.hc_mult) def read(self, residual, norm, quant_format="", post_residual_addition=None): raise NotImplementedError("an MHC FFN input read without its attention") def update_and_read(self, update, hidden_states, residual, norm, **kwargs): - if not isinstance(update, _AttentionUpdate) or update.state is not self.state: + if not isinstance(update, _AttnUpdate) or update.state is not self.state: raise NotImplementedError(f"an MHC FFN input after {update=}") - return self.state.attn_to_mlp(hidden_states, residual, out_norm=norm) + return self.state.update_and_read_ffn_input( + hidden_states, residual, out_norm=norm + ) class _FfnUpdate: @@ -188,22 +192,22 @@ class _FfnUpdate: layer runs itself; the last layer also contracts the streams into the hidden states the layer stack hands on.""" - adds_plainly = False - at_producer = True - can_defer_across_layers = False + is_plain_add = False + applied_at_exit = True + outlives_layer = False def __init__(self, state: MHCState): self.state = state def update(self, hidden_states, residual): - hidden_states = self.state.mlp_combine(hidden_states, residual) - self.state.reset_aux() + hidden_states = self.state.apply_post(hidden_states, residual) + self.state.clear_coefficients() if self.state.is_last_layer: hidden_states = hc_contract(hidden_states, self.state.hc_mult) return hidden_states - def residual_to_attn_tp_shard(self, residual): - return self.state.residual_to_attn_tp_shard(residual) + def slice_residual_attn_tp(self, residual): + return self.state.slice_residual_attn_tp(residual) - def residual_from_attn_tp_shards(self, residual): - return self.state.residual_from_attn_tp_shards(residual) + def gather_residual_attn_tp(self, residual): + return self.state.gather_residual_attn_tp(residual) diff --git a/python/sglang/srt/layers/layer_boundary/residual/stream.py b/python/sglang/srt/layers/layer_boundary/residual/stream.py index 34638f56a8f8..7a8c47979109 100644 --- a/python/sglang/srt/layers/layer_boundary/residual/stream.py +++ b/python/sglang/srt/layers/layer_boundary/residual/stream.py @@ -20,25 +20,25 @@ from sglang.srt.distributed import GroupCoordinator from sglang.srt.layers.layer_boundary.layout import SumGroup, _sum_group -from sglang.srt.layers.layer_boundary.output import HandoffOutput, UnreducedOutput -from sglang.srt.layers.layer_boundary.residual import StageUpdate +from sglang.srt.layers.layer_boundary.output import DeferredFinalize, UnreducedOutput +from sglang.srt.layers.layer_boundary.residual import ResidualUpdate -class CarriedSum(msgspec.Struct, frozen=True): +class DeferredSum(msgspec.Struct, frozen=True): """Runtime reduction left by an exit decision. Fields: group: All-reduce group when no redistribution callable is supplied. - reduce_and_redistribute: Optional callable(value) returning a reduced + reduce_to_dp_local: Optional callable(value) returning a reduced value on destination rows; takes precedence over group. """ group: Optional[GroupCoordinator] = None - reduce_and_redistribute: Optional[Callable] = None + reduce_to_dp_local: Optional[Callable] = None def complete(self, value): - if self.reduce_and_redistribute is not None: - return self.reduce_and_redistribute(value) + if self.reduce_to_dp_local is not None: + return self.reduce_to_dp_local(value) return self.group.all_reduce(value) @@ -63,23 +63,23 @@ class Contribution(msgspec.Struct): """ value: Optional[torch.Tensor] - update: StageUpdate - owed: Union[CarriedSum, DeclaredSum, HandoffOutput, None] = None + update: ResidualUpdate + owed: Union[DeferredSum, DeclaredSum, DeferredFinalize, None] = None def for_boundary(self): # These forms are private inputs to the existing fused-kernel adapters. - if isinstance(self.owed, CarriedSum): + if isinstance(self.owed, DeferredSum): return UnreducedOutput( - self.value, self.owed.group, self.owed.reduce_and_redistribute + self.value, self.owed.group, self.owed.reduce_to_dp_local ) - if isinstance(self.owed, HandoffOutput): + if isinstance(self.owed, DeferredFinalize): return self.owed return self.value def complete(self): - if isinstance(self.owed, (CarriedSum, DeclaredSum)): + if isinstance(self.owed, (DeferredSum, DeclaredSum)): value = self.owed.complete(self.value) - elif isinstance(self.owed, HandoffOutput): + elif isinstance(self.owed, DeferredFinalize): value = self.owed.complete() else: return self.value @@ -123,7 +123,7 @@ def __init__(self, residual=None): self.pending = None @classmethod - def arrive(cls, hidden, residual, update, *, declared_sum=None): + def from_handoff(cls, hidden, residual, update, *, declared_sum=None): """Reconstruct a stream from a tensor handoff, including a TBO microbatch. Args: @@ -139,7 +139,7 @@ def arrive(cls, hidden, residual, update, *, declared_sum=None): stream = cls(residual) if residual is None: return stream.write(hidden), stream - return stream.leave(hidden, update, declared_sum=declared_sum), stream + return stream.record(hidden, update, declared_sum=declared_sum), stream def write(self, residual): if self.pending is not None: @@ -150,11 +150,11 @@ def write(self, residual): self.pending = None return residual - def leave(self, output, update, *, declared_sum=None): + def record(self, output, update, *, declared_sum=None): """Record one producer contribution without applying its residual update. Args: - output: Complete tensor, UnreducedOutput, or producer-specific HandoffOutput. + output: Complete tensor, UnreducedOutput, or producer-specific DeferredFinalize. update: Producer operation that the next boundary must apply. declared_sum: SumGroup statically owed by a raw tensor, or None. Cannot be combined with another output wrapper's completion contract. @@ -173,9 +173,9 @@ def leave(self, output, update, *, declared_sum=None): contribution = Contribution( output.partial, update, - CarriedSum(output.group, output.reduce_and_redistribute), + DeferredSum(output.group, output.reduce_to_dp_local), ) - elif isinstance(output, HandoffOutput): + elif isinstance(output, DeferredFinalize): contribution = Contribution(None, update, output) else: contribution = Contribution(output, update) @@ -207,7 +207,7 @@ def input(self, hidden): if self.pending is not None: return self.pending.for_boundary(), self.residual # An initialized, written stream needs read only. An uninitialized - # stream still uses the first stage's enter operation. + # stream still uses the first stage's init_residual. return hidden, None def complete(self, hidden): @@ -219,17 +219,17 @@ def snapshot(self, hidden): pending = self.pending if pending is None: return hidden.clone() - if not pending.update.adds_plainly: + if not pending.update.is_plain_add: raise NotImplementedError("snapshot requires a plain residual update") if pending.owed is None: value = pending.value - elif isinstance(pending.owed, (CarriedSum, DeclaredSum)): + elif isinstance(pending.owed, (DeferredSum, DeclaredSum)): value = pending.owed.complete(pending.value.clone()) else: raise NotImplementedError("a finalize handoff requires main-output capture") return value.clone() if self.residual is None else value + self.residual - def finish(self, hidden, *, takes_handoff=False, preserve_declared=False): + def export(self, hidden, *, takes_handoff=False, preserve_declared=False): """Export the output/residual pair without closing or consuming the stream. Args: @@ -248,6 +248,6 @@ def finish(self, hidden, *, takes_handoff=False, preserve_declared=False): return hidden, None if preserve_declared and isinstance(self.pending.owed, DeclaredSum): return self.pending.value, self.residual - if takes_handoff and isinstance(self.pending.owed, HandoffOutput): + if takes_handoff and isinstance(self.pending.owed, DeferredFinalize): return self.pending.owed, self.residual return self.pending.complete(), self.residual diff --git a/python/sglang/srt/layers/layer_boundary/stage.py b/python/sglang/srt/layers/layer_boundary/stage.py index e8a2829e9862..fb44b3181b91 100644 --- a/python/sglang/srt/layers/layer_boundary/stage.py +++ b/python/sglang/srt/layers/layer_boundary/stage.py @@ -21,21 +21,21 @@ from sglang.srt.layers.aux_hidden_states import AuxHiddenStateAccumulator from sglang.srt.layers.layer_boundary.adapters import branch from sglang.srt.layers.layer_boundary.adapters.lora import ( - publish_attention, - publish_ffn, + publish_attn_lora_rows, + publish_ffn_lora_rows, ) from sglang.srt.layers.layer_boundary.contracts import ( BatchVariant, - StageEntry, + EntryPath, StageKind, - StageSteps, + StagePath, ) -from sglang.srt.layers.layer_boundary.ops import gather_attention_tp +from sglang.srt.layers.layer_boundary.ops import attn_tp_gather_input from sglang.srt.layers.layer_boundary.residual.access import buffer, from_pp from sglang.srt.layers.layer_boundary.residual.add_norm import ( - ADD, + PLAIN_ADD, ) -from sglang.srt.layers.layer_boundary.residual.batch import current +from sglang.srt.layers.layer_boundary.residual.batch import stream_of from sglang.srt.layers.layer_boundary.residual.stream import DeclaredSum, ResidualStream from sglang.srt.model_executor.forward_batch_info import ForwardBatch from sglang.srt.runtime_context import get_forward @@ -75,10 +75,10 @@ def norm(self): return self.plan.norm def entry( - self, forward_batch: ForwardBatch, steps: Optional[StageSteps] = None - ) -> StageEntry: + self, forward_batch: ForwardBatch, steps: Optional[StagePath] = None + ) -> EntryPath: if steps is None: - steps = self.plan._batch_steps(forward_batch) + steps = self.plan.path_for(forward_batch) return steps.entry def _prepare( @@ -86,7 +86,7 @@ def _prepare( hidden_states, stream: ResidualStream, forward_batch: ForwardBatch, - steps: Optional[StageSteps] = None, + steps: Optional[StagePath] = None, **call, ): """The stage's input and the residual, from the previous stage's output @@ -96,10 +96,10 @@ def _prepare( if ( stream.pending is None and stream.residual is None - and entry.input_sum is not None + and entry.declared_sum is not None ): - hidden_states = stream.leave( - hidden_states, ADD, declared_sum=entry.input_sum + hidden_states = stream.record( + hidden_states, PLAIN_ADD, declared_sum=entry.declared_sum ) if stream.pending is not None: call["update"] = stream.pending.update @@ -115,8 +115,8 @@ def _prepare( hidden_states=hidden_states, forward_batch=forward_batch, ) - if entry.handoff is not None: - hidden_states = entry.handoff( + if entry.attn_input_adapter is not None: + hidden_states = entry.attn_input_adapter( hidden_states, forward_batch, self.plan.qkv_latent_func ) stream.write(residual) @@ -131,8 +131,8 @@ def incoming_residual_rows(self): return self.plan.incoming_residual_rows @property - def input_on_attention_tp_slices(self): - return self.plan.input_on_attention_tp_slices + def input_on_attn_tp_slices(self): + return self.plan.input_on_attn_tp_slices def prepare(self, hidden_states, forward_batch, *, cache=None, **call): """Finish the previous contribution and produce this stage's compute input. @@ -143,9 +143,11 @@ def prepare(self, hidden_states, forward_batch, *, cache=None, **call): forward_batch: Batch owning the stream and selecting the prepared path. cache: FFN-only NPU weight cache prefetched on the FFN input path. **call: Attention-only read/adapter options: quant_format, - post_residual_addition, captured_last_layer_outputs and capture_output. - FFN stages do not accept these options. A capture callback is - called as capture_output(value, owned=bool): owned=True means + post_residual_addition, capture_gathered and capture. FFN + stages do not accept these options. capture is a callback on + the producer's rows; capture_gathered is a collector on the + attention's compute-input rows. A capture callback is + called as capture(value, owned=bool): owned=True means value is fresh storage it may keep, and owned=False means it must copy before retaining. @@ -161,21 +163,21 @@ def prepare(self, hidden_states, forward_batch, *, cache=None, **call): else self._prepare_input ) hidden_states, forward_batch.residual_stream = prepare( - hidden_states, current(forward_batch), forward_batch, **call + hidden_states, stream_of(forward_batch), forward_batch, **call ) return hidden_states def finish(self, hidden_states, forward_batch): """Hand this attention's actual contribution to the following stage.""" - if self.kind is not StageKind.ATTENTION or not self.plan.direct_handoff: + if self.kind is not StageKind.ATTENTION or not self.plan.finishes_directly: raise RuntimeError( "use the stage exit scope for an FFN or single-stage mixer" ) produced = self.plan.produced(forward_batch) - return current(forward_batch).leave( + return stream_of(forward_batch).record( hidden_states, produced.update, - declared_sum=produced.group if produced.always_leaves else None, + declared_sum=produced.group if produced.always_partial else None, ) def exit(self, forward_batch): @@ -191,10 +193,10 @@ def exit(self, forward_batch): using the same decision rather than selecting a second path. """ result = ( - self.plan.output.ffn_exit(forward_batch, stream=current(forward_batch)) + self.plan.output.ffn_exit(forward_batch, stream=stream_of(forward_batch)) if self.kind is StageKind.FFN else self.plan.output.mixer_exit( - forward_batch, stream=current(forward_batch) + forward_batch, stream=stream_of(forward_batch) ) ) return result @@ -216,18 +218,18 @@ def from_pp(self, tensors, forward_batch, *, allow_missing_residual: bool = Fals tensors, residual_in_hidden=( self.declaration.previous is not None - and self.declaration.previous.update.at_producer + and self.declaration.previous.update.applied_at_exit ), allow_missing_residual=allow_missing_residual, ) - declared_sum = self.entry(forward_batch).input_sum - hidden_states, forward_batch.residual_stream = ResidualStream.arrive( - hidden_states, residual, ADD, declared_sum=declared_sum + declared_sum = self.entry(forward_batch).declared_sum + hidden_states, forward_batch.residual_stream = ResidualStream.from_handoff( + hidden_states, residual, PLAIN_ADD, declared_sum=declared_sum ) return hidden_states def snapshot(self, hidden_states, forward_batch): - return current(forward_batch).snapshot(hidden_states) + return stream_of(forward_batch).snapshot(hidden_states) def capture_output(self, hidden_states, forward_batch, *, skip_empty=False): """Capture an output while preserving the caller's reduction timing. @@ -242,7 +244,7 @@ def capture_output(self, hidden_states, forward_batch, *, skip_empty=False): contribution; a declared sum is reduced only on the snapshot copy. Always use updated_handle afterwards. Plain residual updates are required. """ - stream = current(forward_batch) + stream = stream_of(forward_batch) if skip_empty: storage = buffer(hidden_states) if storage is not None and storage.shape[0] == 0: @@ -251,12 +253,12 @@ def capture_output(self, hidden_states, forward_batch, *, skip_empty=False): hidden_states = stream.complete(hidden_states) return hidden_states, stream.snapshot(hidden_states) - def postprocess(self, hidden_states, forward_batch): + def finish_complete_output(self, hidden_states, forward_batch): """Move an FFN output that compute already completed outside exit() (the operation-scheduled TBO path) onto the rows the layer hands on; it chooses no reduction step.""" - return self.plan.output.postprocess_layer( - hidden_states, current(forward_batch), forward_batch + return self.plan.output.finish_complete_output( + hidden_states, stream_of(forward_batch), forward_batch ) def branch_input(self, source, hidden_states, forward_batch): @@ -264,7 +266,7 @@ def branch_input(self, source, hidden_states, forward_batch): self.plan, source.plan, hidden_states, - current(forward_batch), + stream_of(forward_batch), forward_batch, ) forward_batch.residual_stream = stream @@ -278,7 +280,7 @@ def merge_branch(self, contribution, hidden_states, source, forward_batch): self.plan, contribution, hidden_states, - current(forward_batch), + stream_of(forward_batch), source.plan, forward_batch, ) @@ -288,9 +290,9 @@ def merge_branch(self, contribution, hidden_states, source, forward_batch): def _prepare_input(self, hidden_states, stream, forward_batch, **call): plan = self.plan if self.kind is StageKind.ATTENTION: - publish_attention(plan._publish_lora_layout) + publish_attn_lora_rows(plan._publish_lora_layout) if ( - plan._paths.get(BatchVariant.SEQUENCE_PARALLEL) is not None + plan.paths.get(BatchVariant.SEQUENCE_PARALLEL) is not None and plan.enters_stack ): get_forward().set( @@ -299,7 +301,9 @@ def _prepare_input(self, hidden_states, stream, forward_batch, **call): if get_forward().sp_active: hidden_states = layernorm_sp.sp_entry_scatter(hidden_states) else: - publish_ffn(plan._publish_lora_layout, self.entry(forward_batch).input_rows) + publish_ffn_lora_rows( + plan._publish_lora_layout, self.entry(forward_batch).input_rows + ) return self._prepare(hidden_states, stream, forward_batch, **call) def _prepare_attention( @@ -307,28 +311,28 @@ def _prepare_attention( hidden_states: torch.Tensor, stream: ResidualStream, forward_batch: ForwardBatch, - captured_last_layer_outputs: Optional[AuxHiddenStateAccumulator] = None, + capture_gathered: Optional[AuxHiddenStateAccumulator] = None, post_residual_addition: Optional[torch.Tensor] = None, quant_format: str = "", - capture_output: Optional[Callable] = None, + capture: Optional[Callable] = None, ): # Aux consumers need a materialized output before the input norm. # Complete communication here, preserving the existing add+norm kernel # and its FP32 accumulation. Its residual result also supplies capture. - capture_before_read = capture_output is not None and ( + capture_before_read = capture is not None and ( (stream.residual is None and stream.pending is None) or post_residual_addition is not None ) - if capture_output is not None: + if capture is not None: hidden_states = stream.complete(hidden_states) if capture_before_read: # Embeddings precede enter; HF deepstack capture precedes the # extra addition. Neither is the residual returned by the read. - value, previous = stream.finish(hidden_states) + value, previous = stream.export(hidden_states) if previous is None: - capture_output(value) + capture(value) else: - capture_output(value + previous, owned=True) + capture(value + previous, owned=True) hidden_states, stream = self._prepare_input( hidden_states, stream, @@ -337,13 +341,13 @@ def _prepare_attention( post_residual_addition=post_residual_addition, ) entry = self.entry(forward_batch) - if capture_output is not None and not capture_before_read: + if capture is not None and not capture_before_read: value = stream.residual move = entry.capture_move if move is not None: value = move(value, forward_batch=forward_batch) keeps_input = entry.capture_preserves_residual - capture_output( + capture( value, owned=entry.capture_move_allocates or ( @@ -352,7 +356,7 @@ def _prepare_attention( and keeps_input(value, forward_batch) ), ) - if captured_last_layer_outputs is not None: + if capture_gathered is not None: residual_value = stream.residual move = entry.input_move gathered_last_layer_output = ( @@ -363,11 +367,11 @@ def _prepare_attention( forward_batch=forward_batch, ) ) - # Input gathers allocate fresh DP buffers (see gather_attention_tp). + # Input gathers allocate fresh DP buffers (see attn_tp_gather_input). # Without a gather this is the mutable residual, so retention copies. - captured_last_layer_outputs.capture( + capture_gathered.capture( gathered_last_layer_output, - owned=move is gather_attention_tp + owned=move is attn_tp_gather_input or ( move is None and entry.capture_preserves_residual is not None diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index 0349758734b4..7c874a39483b 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/python/sglang/srt/model_executor/forward_batch_info.py @@ -1260,7 +1260,7 @@ def set_local_num_token_non_padded(self, *, sharded: bool) -> None: def moe_num_token_non_padded(self) -> Optional[torch.Tensor]: """Bound for masking a sparse MoE's padded rows, or None when the MoE input is a gathered buffer whose real rows are not a prefix of it.""" - from sglang.srt.layers.layer_boundary import moe_cp_gathers_sparse_moe_input + from sglang.srt.layers.layer_boundary import batch_gathers_over_moe_cp from sglang.srt.layers.moe.utils import is_moe_input_scattered_across_dp_ranks if self.num_token_non_padded is None: @@ -1269,7 +1269,7 @@ def moe_num_token_non_padded(self) -> Optional[torch.Tensor]: if is_moe_input_scattered_across_dp_ranks(): # a2a dispatch, FP4 all-gather and dwdp all route the local shard. return self.num_token_non_padded - if moe_cp_gathers_sparse_moe_input(self): + if batch_gathers_over_moe_cp(self): return None # DSA / MLA CP all-gather across attention CP on a prefill, which leaves # the real rows zigzag-permuted rather than in a prefix. diff --git a/python/sglang/srt/models/bailing_moe.py b/python/sglang/srt/models/bailing_moe.py index b485de6bd156..bc2de5dbbdf9 100644 --- a/python/sglang/srt/models/bailing_moe.py +++ b/python/sglang/srt/models/bailing_moe.py @@ -41,7 +41,7 @@ from sglang.srt.layers.layer_boundary import ( declare_attn, declare_ffn, - enable_moe_dense_fully_dp, + is_dense_ffn_fully_dp, make_stages, ) from sglang.srt.layers.layer_boundary.residual import batch as residual_batch @@ -684,7 +684,7 @@ def __init__( prefix=add_prefix("mlp", prefix), ) else: - if enable_moe_dense_fully_dp(): + if is_dense_ffn_fully_dp(): mlp_tp_rank, mlp_tp_size = 0, 1 else: mlp_tp_rank, mlp_tp_size = None, None @@ -704,12 +704,12 @@ def __init__( ( declare_ffn( sparse=self.is_layer_sparse, - next_sparse=is_next_layer_sparse, + next_layer_sparse=is_next_layer_sparse, ), self.post_attention_layernorm, ), previous=declare_ffn( - sparse=is_previous_layer_sparse, next_sparse=self.is_layer_sparse + sparse=is_previous_layer_sparse, next_layer_sparse=self.is_layer_sparse ) if layer_id != 0 else None, @@ -734,8 +734,8 @@ def forward( hidden_states = self.attn_boundary.prepare( hidden_states, forward_batch, - captured_last_layer_outputs=captured_last_layer_outputs, - capture_output=capture_output, + capture_gathered=captured_last_layer_outputs, + capture=capture_output, ) if hidden_states.shape[0] != 0: @@ -847,7 +847,7 @@ def forward( else: hidden_states = residual_batch.complete_output(hidden_states, forward_batch) if not forward_batch.forward_mode.is_idle(): - hidden_states = residual_batch.norm( + hidden_states = residual_batch.final_norm( hidden_states, forward_batch, self.norm ) diff --git a/python/sglang/srt/models/bailing_moe_linear.py b/python/sglang/srt/models/bailing_moe_linear.py index 097b42586a85..8d87dcc3429e 100644 --- a/python/sglang/srt/models/bailing_moe_linear.py +++ b/python/sglang/srt/models/bailing_moe_linear.py @@ -830,11 +830,13 @@ def __init__( ( declare_ffn( sparse=is_moe_layer, - next_sparse=is_next_layer_moe_layer, + next_layer_sparse=is_next_layer_moe_layer, ), self.post_attention_layernorm, ), - previous=declare_ffn(sparse=is_previous_moe_layer, next_sparse=is_moe_layer) + previous=declare_ffn( + sparse=is_previous_moe_layer, next_layer_sparse=is_moe_layer + ) if layer_id != 0 else None, terminal=layer_id == (1 if is_nextn else config.num_hidden_layers) - 1, @@ -1010,7 +1012,7 @@ def forward( else: hidden_states = residual_batch.complete_output(hidden_states, forward_batch) if not forward_batch.forward_mode.is_idle(): - hidden_states = residual_batch.norm( + hidden_states = residual_batch.final_norm( hidden_states, forward_batch, self.norm ) return hidden_states diff --git a/python/sglang/srt/models/bailing_moe_nextn.py b/python/sglang/srt/models/bailing_moe_nextn.py index 82191e3aede2..583e869936ff 100644 --- a/python/sglang/srt/models/bailing_moe_nextn.py +++ b/python/sglang/srt/models/bailing_moe_nextn.py @@ -183,7 +183,7 @@ def forward( hidden_states = residual_batch.complete_output(hidden_states, forward_batch) if not forward_batch.forward_mode.is_idle(): - hidden_states = residual_batch.norm( + hidden_states = residual_batch.final_norm( hidden_states, forward_batch, self.final_layernorm ) diff --git a/python/sglang/srt/models/bailing_moe_v3.py b/python/sglang/srt/models/bailing_moe_v3.py index e238082e491d..279a11bf3a7c 100644 --- a/python/sglang/srt/models/bailing_moe_v3.py +++ b/python/sglang/srt/models/bailing_moe_v3.py @@ -33,7 +33,7 @@ from sglang.srt.layers.layer_boundary import ( declare_attn, declare_ffn, - enable_moe_dense_fully_dp, + is_dense_ffn_fully_dp, make_stages, ) from sglang.srt.layers.layer_boundary.residual import batch as residual_batch @@ -1110,7 +1110,7 @@ def __init__( is_next_layer_sparse = not (self.expert_num == 1) and ( self.layer_id + 1 >= config.first_k_dense_replace ) - if enable_moe_dense_fully_dp(): + if is_dense_ffn_fully_dp(): mlp_tp_rank, mlp_tp_size = 0, 1 else: mlp_tp_rank, mlp_tp_size = None, None @@ -1162,11 +1162,13 @@ def __init__( ( declare_ffn( sparse=is_moe_layer, - next_sparse=is_next_layer_sparse, + next_layer_sparse=is_next_layer_sparse, ), self.post_attention_layernorm, ), - previous=declare_ffn(sparse=is_previous_moe_layer, next_sparse=is_moe_layer) + previous=declare_ffn( + sparse=is_previous_moe_layer, next_layer_sparse=is_moe_layer + ) if layer_id != 0 else None, terminal=layer_id == (1 if is_nextn else config.num_hidden_layers) - 1, @@ -1183,7 +1185,7 @@ def forward( **kwargs, ) -> torch.Tensor: hidden_states = self.attn_boundary.prepare( - hidden_states, forward_batch, capture_output=capture_output + hidden_states, forward_batch, capture=capture_output ) if not forward_batch.forward_mode.is_idle(): @@ -1213,7 +1215,7 @@ def forward( with self.ffn_boundary.exit(forward_batch) as ffn_exit: if not ( - enable_moe_dense_fully_dp() + is_dense_ffn_fully_dp() and (not self.is_layer_sparse) and hidden_states.shape[0] == 0 ): @@ -1397,11 +1399,11 @@ def capture_output(value, *, owned=False): else: hidden_states = residual_batch.complete_output(hidden_states, forward_batch) if not forward_batch.forward_mode.is_idle(): - hidden_states = residual_batch.norm( + hidden_states = residual_batch.final_norm( hidden_states, forward_batch, self.norm, - capture_output=capture_output + capture=capture_output if capture_aux and self.end_layer - 1 in self.layers_to_capture else None, ) diff --git a/python/sglang/srt/models/deepseek_nextn.py b/python/sglang/srt/models/deepseek_nextn.py index 7bc345ac2ca1..fc8e3d428b49 100644 --- a/python/sglang/srt/models/deepseek_nextn.py +++ b/python/sglang/srt/models/deepseek_nextn.py @@ -228,7 +228,7 @@ def forward( ) hidden_states = residual_batch.complete_output(hidden_states, forward_batch) if not forward_batch.forward_mode.is_idle(): - hidden_states = residual_batch.norm( + hidden_states = residual_batch.final_norm( hidden_states, forward_batch, self.shared_head.norm ) diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 6ce92987a5f1..fbebff26c90f 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -75,8 +75,8 @@ from sglang.srt.layers.layer_boundary import ( declare_attn, declare_ffn, - enable_moe_dense_fully_dp, get_attn_tp_context, + is_dense_ffn_fully_dp, make_stages, ) from sglang.srt.layers.layer_boundary.adapters.context_parallel import ( @@ -1112,10 +1112,10 @@ def forward_normal_dual_stream( # deferred_finalize excludes _shared_expert_tp1, so the shared add folds in. assert shared_output is not None from sglang.srt.layers.layer_boundary.fusions.cutedsl import ( - MoeFinalizeHandoff, + MoeDeferredFinalize, ) - return MoeFinalizeHandoff.from_flashinfer( + return MoeDeferredFinalize.from_flashinfer( final_hidden_states, gated_shared_output=shared_output, m=hidden_states.shape[0], @@ -2294,7 +2294,7 @@ def forward( hidden_states: torch.Tensor, forward_batch: ForwardBatch, zero_allocator: BumpAllocator, - input_on_attention_tp_slices: bool = False, + input_on_attn_tp_slices: bool = False, llama_4_scaling: Optional[torch.Tensor] = None, prev_topk_indices: Optional[torch.Tensor] = None, ): @@ -2303,7 +2303,7 @@ def forward( hidden_states=hidden_states, forward_batch=forward_batch, zero_allocator=zero_allocator, - input_on_attention_tp_slices=input_on_attention_tp_slices, + input_on_attn_tp_slices=input_on_attn_tp_slices, llama_4_scaling=llama_4_scaling, prev_topk_indices=prev_topk_indices, ) @@ -2315,7 +2315,7 @@ def forward_prepare( hidden_states: torch.Tensor, forward_batch: ForwardBatch, zero_allocator: BumpAllocator, - input_on_attention_tp_slices: bool = False, + input_on_attn_tp_slices: bool = False, llama_4_scaling: Optional[torch.Tensor] = None, prev_topk_indices: Optional[torch.Tensor] = None, ): @@ -2400,7 +2400,7 @@ def forward_prepare( hidden_states, forward_batch, zero_allocator, - input_on_attention_tp_slices, + input_on_attn_tp_slices, ) elif attn_forward_method == AttnForwardMethod.MLA_NPU: inner_state = forward_mla_prepare_npu( @@ -2409,7 +2409,7 @@ def forward_prepare( hidden_states, forward_batch, zero_allocator, - input_on_attention_tp_slices, + input_on_attn_tp_slices, ) elif attn_forward_method == AttnForwardMethod.DSA_NPU: inner_state = forward_dsa_prepare_npu( @@ -2418,7 +2418,7 @@ def forward_prepare( hidden_states, forward_batch, zero_allocator, - input_on_attention_tp_slices, + input_on_attn_tp_slices, prev_topk_indices, ) else: @@ -2599,7 +2599,7 @@ def __init__( is_nextn=is_nextn, ) else: - if enable_moe_dense_fully_dp(): + if is_dense_ffn_fully_dp(): mlp_tp_rank, mlp_tp_size = 0, 1 else: mlp_tp_rank, mlp_tp_size = None, None @@ -2656,14 +2656,15 @@ def _build_stages( ( declare_ffn( sparse=self.is_layer_sparse, - next_sparse=self._stage_next_sparse, + next_layer_sparse=self._stage_next_sparse, output_transform=output, ), post_attention_layernorm, {"fusions": fusions}, ), previous=declare_ffn( - sparse=self._stage_previous_sparse, next_sparse=self.is_layer_sparse + sparse=self._stage_previous_sparse, + next_layer_sparse=self.is_layer_sparse, ) if not self._stage_enters_stack else None, @@ -2726,7 +2727,7 @@ def forward( hidden_states = self.attn_boundary.prepare( hidden_states, forward_batch, - captured_last_layer_outputs=captured_last_layer_outputs, + capture_gathered=captured_last_layer_outputs, quant_format=self._resolve_gfx95_quant_format(), ) @@ -2737,9 +2738,7 @@ def forward( forward_batch=forward_batch, zero_allocator=zero_allocator, llama_4_scaling=llama_4_scaling, - input_on_attention_tp_slices=( - self.attn_boundary.input_on_attention_tp_slices - ), + input_on_attn_tp_slices=(self.attn_boundary.input_on_attn_tp_slices), prev_topk_indices=prev_topk_indices, ) if isinstance(hidden_states, tuple): @@ -2813,7 +2812,7 @@ def op_comm_prepare_mlp(self, state): ) def op_comm_postprocess_layer(self, state): - hidden_states = self.ffn_boundary.postprocess( + hidden_states = self.ffn_boundary.finish_complete_output( state.pop("hidden_states_mlp_output"), state.forward_batch ) @@ -3158,7 +3157,7 @@ def forward( else: hidden_states = residual_batch.complete_output(hidden_states, forward_batch) if not forward_batch.forward_mode.is_idle(): - hidden_states = residual_batch.norm( + hidden_states = residual_batch.final_norm( hidden_states, forward_batch, self.norm ) @@ -3465,7 +3464,7 @@ def dsv2_flashinfer_moe_dual_stream_graph( mlp_reduce_scatter=mlp_reduce_scatter, flashinfer_trtllm_bypass=True, lora_batch_layout=LoRABatchLayout.TP_GLOBAL, - # The op's Tensor schema cannot carry a MoeFinalizeHandoff. + # The op's Tensor schema cannot carry a MoeDeferredFinalize. defer_moe_finalize=False, ): return moe_fusion.forward_normal_dual_stream(hidden_states) diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index b6b31eb5cbc6..2413876c61de 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/python/sglang/srt/models/deepseek_v4.py @@ -100,8 +100,8 @@ from sglang.srt.layers.engram import Engram, EngramHasher, EngramLayout from sglang.srt.layers.layer_boundary import get_attn_tp_context from sglang.srt.layers.layer_boundary.adapters.context_parallel import ( - dsa_cp_gather_hidden_states, - dsa_cp_reduce_scatter_hidden_states, + attn_cp_gather, + attn_cp_reduce_scatter, ) from sglang.srt.layers.layernorm import RMSNorm from sglang.srt.layers.linear import ColumnParallelLinear, RowParallelLinear @@ -3819,7 +3819,7 @@ def _run_moe_ffn_dp_sync( if _use_cp: moe_a2a_backend = get_moe_a2a_backend() if moe_a2a_backend.is_none(): - hidden_states = dsa_cp_gather_hidden_states(hidden_states) + hidden_states = attn_cp_gather(hidden_states) else: assert ( moe_a2a_backend.is_deepep() @@ -3870,7 +3870,7 @@ def _run_moe_ffn_dp_sync( skip_shared_experts=_do_shared_local, ) if _use_cp and get_moe_a2a_backend().is_none(): - hidden_states = dsa_cp_reduce_scatter_hidden_states(hidden_states) + hidden_states = attn_cp_reduce_scatter(hidden_states) elif _use_tp_moe_gather: hidden_states, global_hidden_states = ( get_local_dp_buffer(get_parallel().tp_group), diff --git a/python/sglang/srt/models/dots3_common/modeling.py b/python/sglang/srt/models/dots3_common/modeling.py index ed4ff8ed925b..181ec5a30b1b 100644 --- a/python/sglang/srt/models/dots3_common/modeling.py +++ b/python/sglang/srt/models/dots3_common/modeling.py @@ -61,7 +61,7 @@ from sglang.srt.layers.layer_boundary import ( declare_attn, declare_ffn, - enable_moe_dense_fully_dp, + is_dense_ffn_fully_dp, make_stages, ) from sglang.srt.layers.layer_boundary.residual import batch as residual_batch @@ -1553,7 +1553,7 @@ def __init__( is_nextn=is_nextn, ) else: - if enable_moe_dense_fully_dp(): + if is_dense_ffn_fully_dp(): mlp_tp_rank, mlp_tp_size = 0, 1 else: mlp_tp_rank, mlp_tp_size = None, None @@ -1577,12 +1577,12 @@ def __init__( ( declare_ffn( sparse=self.is_layer_sparse, - next_sparse=is_next_layer_sparse, + next_layer_sparse=is_next_layer_sparse, ), self.post_attention_layernorm, ), previous=declare_ffn( - sparse=is_previous_layer_sparse, next_sparse=self.is_layer_sparse + sparse=is_previous_layer_sparse, next_layer_sparse=self.is_layer_sparse ) if layer_id != 0 else None, @@ -1660,7 +1660,7 @@ def op_comm_prepare_mlp(self, state): def op_mlp(self, state): hidden_states = state.pop("hidden_states_mlp_input") if not ( - enable_moe_dense_fully_dp() + is_dense_ffn_fully_dp() and (not self.is_layer_sparse) and hidden_states.shape[0] == 0 ): @@ -1671,7 +1671,7 @@ def op_mlp(self, state): state.hidden_states_mlp_output = hidden_states def op_comm_postprocess_layer(self, state): - hidden_states = self.ffn_boundary.postprocess( + hidden_states = self.ffn_boundary.finish_complete_output( state.pop("hidden_states_mlp_output"), state.forward_batch ) @@ -1800,7 +1800,7 @@ def forward( return residual_batch.to_pp(hidden_states, forward_batch) else: if not forward_batch.forward_mode.is_idle(): - hidden_states = residual_batch.norm( + hidden_states = residual_batch.final_norm( hidden_states, forward_batch, self.norm ) return hidden_states diff --git a/python/sglang/srt/models/dots3_common/nextn.py b/python/sglang/srt/models/dots3_common/nextn.py index 47c6ecf1bbca..af8854c49b7e 100644 --- a/python/sglang/srt/models/dots3_common/nextn.py +++ b/python/sglang/srt/models/dots3_common/nextn.py @@ -122,7 +122,7 @@ def forward( hidden_states = residual_batch.complete_output(hidden_states, forward_batch) if not forward_batch.forward_mode.is_idle(): - hidden_states = residual_batch.norm( + hidden_states = residual_batch.final_norm( hidden_states, forward_batch, head.shared_head.norm ) return hidden_states diff --git a/python/sglang/srt/models/ernie4.py b/python/sglang/srt/models/ernie4.py index 405b6683555f..71d71b48195b 100644 --- a/python/sglang/srt/models/ernie4.py +++ b/python/sglang/srt/models/ernie4.py @@ -26,7 +26,7 @@ from sglang.srt.distributed import ( tensor_model_parallel_all_reduce, ) -from sglang.srt.layers.layer_boundary import enable_moe_dense_fully_dp +from sglang.srt.layers.layer_boundary import is_dense_ffn_fully_dp from sglang.srt.layers.layernorm import RMSNorm from sglang.srt.layers.logits_processor import LogitsProcessor from sglang.srt.layers.moe.ep_moe.layer import get_moe_impl_class @@ -198,7 +198,7 @@ def __init__( prefix=add_prefix("mlp", prefix), ) else: - if enable_moe_dense_fully_dp(): + if is_dense_ffn_fully_dp(): mlp_tp_rank, mlp_tp_size = 0, 1 else: mlp_tp_rank, mlp_tp_size = None, None diff --git a/python/sglang/srt/models/falcon_h1.py b/python/sglang/srt/models/falcon_h1.py index ecfbd97a403e..76fb83516b15 100644 --- a/python/sglang/srt/models/falcon_h1.py +++ b/python/sglang/srt/models/falcon_h1.py @@ -230,12 +230,12 @@ def __init__( self.feed_forward.scale_output, before_reduce_scatter=True ), sparse=self.is_layer_sparse, - next_sparse=is_next_layer_sparse, + next_layer_sparse=is_next_layer_sparse, ), self.pre_ff_layernorm, ), previous=declare_ffn( - sparse=is_previous_layer_sparse, next_sparse=self.is_layer_sparse + sparse=is_previous_layer_sparse, next_layer_sparse=self.is_layer_sparse ) if layer_id != 0 else None, @@ -442,7 +442,7 @@ def forward( hidden_states = residual_batch.complete_output(hidden_states, forward_batch) if not forward_batch.forward_mode.is_idle(): - hidden_states = residual_batch.norm( + hidden_states = residual_batch.final_norm( hidden_states, forward_batch, self.final_layernorm ) diff --git a/python/sglang/srt/models/gigachat35.py b/python/sglang/srt/models/gigachat35.py index b256bc6ad840..22d77c3e1d62 100644 --- a/python/sglang/srt/models/gigachat35.py +++ b/python/sglang/srt/models/gigachat35.py @@ -484,9 +484,7 @@ def forward( hidden_states=hidden_states, forward_batch=forward_batch, zero_allocator=zero_allocator, - input_on_attention_tp_slices=( - self.attn_boundary.input_on_attention_tp_slices - ), + input_on_attn_tp_slices=(self.attn_boundary.input_on_attn_tp_slices), ) get_attn_tp_context().clear_attn_inputs() @@ -586,7 +584,9 @@ def forward( hidden_states = residual_batch.complete_output(hidden_states, forward_batch) if hidden_states.shape[0] != 0: - hidden_states = residual_batch.norm(hidden_states, forward_batch, self.norm) + hidden_states = residual_batch.final_norm( + hidden_states, forward_batch, self.norm + ) return hidden_states diff --git a/python/sglang/srt/models/gigachat35_mtp.py b/python/sglang/srt/models/gigachat35_mtp.py index d7f22e2eb845..3d376f9651d0 100644 --- a/python/sglang/srt/models/gigachat35_mtp.py +++ b/python/sglang/srt/models/gigachat35_mtp.py @@ -122,7 +122,7 @@ def forward( hidden_states_before_norm = residual_batch.snapshot( hidden_states, forward_batch ) - hidden_states = residual_batch.norm( + hidden_states = residual_batch.final_norm( hidden_states, forward_batch, self.shared_head.norm ) diff --git a/python/sglang/srt/models/glm4_moe.py b/python/sglang/srt/models/glm4_moe.py index e2a2875f80a1..b09c3e703178 100644 --- a/python/sglang/srt/models/glm4_moe.py +++ b/python/sglang/srt/models/glm4_moe.py @@ -46,7 +46,7 @@ from sglang.srt.layers.layer_boundary import ( declare_attn, declare_ffn, - enable_moe_dense_fully_dp, + is_dense_ffn_fully_dp, make_stages, ) from sglang.srt.layers.layer_boundary.residual import batch as residual_batch @@ -847,7 +847,7 @@ def __init__( alt_stream=alt_stream, ) else: - if enable_moe_dense_fully_dp(): + if is_dense_ffn_fully_dp(): mlp_tp_rank, mlp_tp_size = 0, 1 else: mlp_tp_rank, mlp_tp_size = None, None @@ -871,12 +871,12 @@ def __init__( ( declare_ffn( sparse=self.is_layer_sparse, - next_sparse=is_next_layer_sparse, + next_layer_sparse=is_next_layer_sparse, ), self.post_attention_layernorm, ), previous=declare_ffn( - sparse=is_previous_layer_sparse, next_sparse=self.is_layer_sparse + sparse=is_previous_layer_sparse, next_layer_sparse=self.is_layer_sparse ) if layer_id != 0 else None, @@ -946,7 +946,7 @@ def forward( hidden_states, forward_batch, quant_format=self.attn_quant_format, - capture_output=capture_output, + capture=capture_output, ) hidden_states = self.self_attn( @@ -992,7 +992,7 @@ def op_comm_prepare_mlp(self, state): ) def op_comm_postprocess_layer(self, state): - hidden_states = self.ffn_boundary.postprocess( + hidden_states = self.ffn_boundary.finish_complete_output( state.pop("hidden_states_mlp_output"), state.forward_batch ) @@ -1123,7 +1123,7 @@ def forward( else: hidden_states = residual_batch.complete_output(hidden_states, forward_batch) if not forward_batch.forward_mode.is_idle(): - hidden_states = residual_batch.norm( + hidden_states = residual_batch.final_norm( hidden_states, forward_batch, self.norm ) if len(aux_hidden_states) == 0: diff --git a/python/sglang/srt/models/glm4_moe_lite.py b/python/sglang/srt/models/glm4_moe_lite.py index e910e693da03..7e9fc19e8afd 100644 --- a/python/sglang/srt/models/glm4_moe_lite.py +++ b/python/sglang/srt/models/glm4_moe_lite.py @@ -40,8 +40,8 @@ from sglang.srt.layers.layer_boundary import ( declare_attn, declare_ffn, - enable_moe_dense_fully_dp, get_attn_tp_context, + is_dense_ffn_fully_dp, make_stages, ) from sglang.srt.layers.layer_boundary.residual import batch as residual_batch @@ -560,7 +560,7 @@ def __init__( is_nextn=is_nextn, ) else: - if enable_moe_dense_fully_dp(): + if is_dense_ffn_fully_dp(): mlp_tp_rank, mlp_tp_size = 0, 1 else: mlp_tp_rank, mlp_tp_size = None, None @@ -590,12 +590,12 @@ def __init__( ( declare_ffn( sparse=self.is_layer_sparse, - next_sparse=is_next_layer_sparse, + next_layer_sparse=is_next_layer_sparse, ), self.post_attention_layernorm, ), previous=declare_ffn( - sparse=is_previous_layer_sparse, next_sparse=self.is_layer_sparse + sparse=is_previous_layer_sparse, next_layer_sparse=self.is_layer_sparse ) if layer_id != 0 else None, @@ -637,7 +637,7 @@ def forward( hidden_states, forward_batch, quant_format=getattr(self, "_gfx95_quant_format", ""), - capture_output=capture_output, + capture=capture_output, ) hidden_states = self.self_attn( @@ -645,9 +645,7 @@ def forward( hidden_states=hidden_states, forward_batch=forward_batch, zero_allocator=zero_allocator, - input_on_attention_tp_slices=( - self.attn_boundary.input_on_attention_tp_slices - ), + input_on_attn_tp_slices=(self.attn_boundary.input_on_attn_tp_slices), ) if isinstance(hidden_states, tuple): hidden_states = hidden_states[0] @@ -694,7 +692,7 @@ def op_comm_prepare_mlp(self, state): ) def op_comm_postprocess_layer(self, state): - hidden_states = self.ffn_boundary.postprocess( + hidden_states = self.ffn_boundary.finish_complete_output( state.pop("hidden_states_mlp_output"), state.forward_batch ) @@ -830,7 +828,7 @@ def forward( else: hidden_states = residual_batch.complete_output(hidden_states, forward_batch) if not forward_batch.forward_mode.is_idle(): - hidden_states = residual_batch.norm( + hidden_states = residual_batch.final_norm( hidden_states, forward_batch, self.norm ) diff --git a/python/sglang/srt/models/glm4_moe_lite_nextn.py b/python/sglang/srt/models/glm4_moe_lite_nextn.py index 25e0b7c93a86..9ecd58cc5d31 100644 --- a/python/sglang/srt/models/glm4_moe_lite_nextn.py +++ b/python/sglang/srt/models/glm4_moe_lite_nextn.py @@ -123,7 +123,7 @@ def forward( hidden_states = residual_batch.complete_output(hidden_states, forward_batch) if not forward_batch.forward_mode.is_idle(): - hidden_states = residual_batch.norm( + hidden_states = residual_batch.final_norm( hidden_states, forward_batch, self.shared_head.norm ) diff --git a/python/sglang/srt/models/glm4_moe_nextn.py b/python/sglang/srt/models/glm4_moe_nextn.py index 83750d24328d..f6b3ab051066 100644 --- a/python/sglang/srt/models/glm4_moe_nextn.py +++ b/python/sglang/srt/models/glm4_moe_nextn.py @@ -107,7 +107,7 @@ def forward( hidden_states = residual_batch.complete_output(hidden_states, forward_batch) if not forward_batch.forward_mode.is_idle(): - hidden_states = residual_batch.norm( + hidden_states = residual_batch.final_norm( hidden_states, forward_batch, self.shared_head.norm ) diff --git a/python/sglang/srt/models/glm5_next.py b/python/sglang/srt/models/glm5_next.py index 5fc8c9bc5b49..4a9f605c1869 100644 --- a/python/sglang/srt/models/glm5_next.py +++ b/python/sglang/srt/models/glm5_next.py @@ -27,12 +27,12 @@ from sglang.srt.layers.attention.vision import VisionAttention from sglang.srt.layers.aux_hidden_states import AuxHiddenStateList from sglang.srt.layers.layer_boundary import ( - PLAIN_RESIDUAL, + PLAIN_RESIDUAL_OPS, MHCState, declare_attn, declare_ffn, - enable_moe_dense_fully_dp, get_attn_tp_context, + is_dense_ffn_fully_dp, make_stages, ) from sglang.srt.layers.layer_boundary.residual import access as residual_access @@ -740,7 +740,7 @@ def __init__( is_nextn=is_nextn, ) else: - if enable_moe_dense_fully_dp(): + if is_dense_ffn_fully_dp(): mlp_tp_rank, mlp_tp_size = 0, 1 else: mlp_tp_rank, mlp_tp_size = None, None @@ -780,7 +780,7 @@ def __init__( ) terminal = layer_id == (1 if is_nextn else config.num_hidden_layers) - 1 - residual = PLAIN_RESIDUAL + residual = PLAIN_RESIDUAL_OPS if self.config.mhc: residual = MHCState( hc_mult=config.hc_mult, @@ -793,12 +793,10 @@ def __init__( else None ), is_last_layer=terminal, - ).layer_residual() + ).residual_ops() self.attn_boundary, self.ffn_boundary = make_stages( ( - declare_attn( - read=residual.attention_read, update=residual.attention_update - ), + declare_attn(read=residual.attn_readout, update=residual.attn_update), self.input_layernorm, { "qkv_latent_func": self.self_attn.prepare_qkv_latent @@ -809,15 +807,15 @@ def __init__( ( declare_ffn( sparse=self.is_layer_sparse, - next_sparse=is_next_layer_sparse, - read=residual.ffn_read, + next_layer_sparse=is_next_layer_sparse, + read=residual.ffn_readout, update=residual.ffn_update, ), self.post_attention_layernorm, ), previous=declare_ffn( sparse=is_previous_layer_sparse, - next_sparse=self.is_layer_sparse, + next_layer_sparse=self.is_layer_sparse, update=residual.ffn_update, ) if layer_id != 0 @@ -933,7 +931,7 @@ def forward( hidden_states_orig = residual_access.buffer(hidden_states) hidden_states = self.attn_boundary.prepare( - hidden_states, forward_batch, capture_output=capture_output + hidden_states, forward_batch, capture=capture_output ) hidden_states = self.self_attn( @@ -941,9 +939,7 @@ def forward( hidden_states=hidden_states, forward_batch=forward_batch, zero_allocator=zero_allocator, - input_on_attention_tp_slices=( - self.attn_boundary.input_on_attention_tp_slices - ), + input_on_attn_tp_slices=(self.attn_boundary.input_on_attn_tp_slices), prev_topk_indices=prev_topk_indices, ) if isinstance(hidden_states, tuple): @@ -1195,7 +1191,7 @@ def capture_output(aux_hidden_state, *, owned=False): else: hidden_states = residual_batch.complete_output(hidden_states, forward_batch) if not forward_batch.forward_mode.is_idle(): - hidden_states = residual_batch.norm( + hidden_states = residual_batch.final_norm( hidden_states, forward_batch, self.norm ) diff --git a/python/sglang/srt/models/gpt_oss.py b/python/sglang/srt/models/gpt_oss.py index 5218e5d76dae..8537db66cdc4 100644 --- a/python/sglang/srt/models/gpt_oss.py +++ b/python/sglang/srt/models/gpt_oss.py @@ -604,12 +604,12 @@ def __init__( ( declare_ffn( sparse=self.is_layer_sparse, - next_sparse=is_next_layer_sparse, + next_layer_sparse=is_next_layer_sparse, ), self.post_attention_layernorm, ), previous=declare_ffn( - sparse=is_previous_layer_sparse, next_sparse=self.is_layer_sparse + sparse=is_previous_layer_sparse, next_layer_sparse=self.is_layer_sparse ) if layer_id != 0 else None, @@ -624,7 +624,7 @@ def forward( capture_output=None, ) -> torch.Tensor: hidden_states = self.attn_boundary.prepare( - hidden_states, forward_batch, capture_output=capture_output + hidden_states, forward_batch, capture=capture_output ) if hidden_states.shape[0] != 0: @@ -730,11 +730,11 @@ def forward( else: hidden_states = residual_batch.complete_output(hidden_states, forward_batch) if hidden_states.shape[0] != 0: - hidden_states = residual_batch.norm( + hidden_states = residual_batch.final_norm( hidden_states, forward_batch, self.norm, - capture_output=aux_hidden_states.capture + capture=aux_hidden_states.capture if self.end_layer in self.layers_to_capture else None, ) diff --git a/python/sglang/srt/models/interns2_mobius.py b/python/sglang/srt/models/interns2_mobius.py index 229daa8e8017..c6b52c08f7bc 100644 --- a/python/sglang/srt/models/interns2_mobius.py +++ b/python/sglang/srt/models/interns2_mobius.py @@ -21,7 +21,10 @@ make_stages, ) from sglang.srt.layers.layer_boundary.residual import batch as residual_batch -from sglang.srt.layers.layer_boundary.residual.add_norm import Fp8Input, NormQuantRead +from sglang.srt.layers.layer_boundary.residual.add_norm import ( + Fp8Input, + NormQuantReadout, +) from sglang.srt.layers.layernorm import GemmaRMSNorm from sglang.srt.layers.linear import ( QKVParallelLinear, @@ -492,7 +495,7 @@ def __init__( self.attn_boundary, self.ffn_boundary = make_stages( ( declare_attn( - read=NormQuantRead( + read=NormQuantReadout( fp8_input=Fp8Input.TUPLE_AND_BF16 if accepts_fp8_input else None ) ), @@ -501,11 +504,11 @@ def __init__( ( declare_ffn( sparse=True, - next_sparse=True, + next_layer_sparse=True, ), self.post_attention_layernorm, ), - previous=declare_ffn(sparse=True, next_sparse=True) + previous=declare_ffn(sparse=True, next_layer_sparse=True) if layer_id != 0 else None, terminal=layer_id == config.num_hidden_layers - 1, @@ -521,7 +524,7 @@ def forward( hidden_states = self.attn_boundary.prepare( hidden_states, forward_batch, - captured_last_layer_outputs=kwargs.get("captured_last_layer_outputs"), + capture_gathered=kwargs.get("captured_last_layer_outputs"), ) if not forward_batch.forward_mode.is_idle(): hidden_states = self.linear_attn(hidden_states, forward_batch) @@ -624,7 +627,7 @@ def __init__( self.attn_boundary, self.ffn_boundary = make_stages( ( declare_attn( - read=NormQuantRead( + read=NormQuantReadout( fp8_input=Fp8Input.TUPLE if accepts_fp8_input else None ) ), @@ -633,11 +636,11 @@ def __init__( ( declare_ffn( sparse=True, - next_sparse=True, + next_layer_sparse=True, ), self.post_attention_layernorm, ), - previous=declare_ffn(sparse=True, next_sparse=True) + previous=declare_ffn(sparse=True, next_layer_sparse=True) if layer_id != 0 else None, terminal=layer_id == config.num_hidden_layers - 1, @@ -657,7 +660,7 @@ def forward( hidden_states = self.attn_boundary.prepare( hidden_states, forward_batch, - captured_last_layer_outputs=captured_last_layer_outputs, + capture_gathered=captured_last_layer_outputs, ) if not forward_batch.forward_mode.is_idle(): hidden_states = self.self_attention( @@ -791,7 +794,9 @@ def forward( hidden_states = residual_batch.complete_output(hidden_states, forward_batch) if hidden_states.shape[0] != 0: - hidden_states = residual_batch.norm(hidden_states, forward_batch, self.norm) + hidden_states = residual_batch.final_norm( + hidden_states, forward_batch, self.norm + ) return ( hidden_states if not aux_hidden_states diff --git a/python/sglang/srt/models/laguna.py b/python/sglang/srt/models/laguna.py index 22ce045c4929..e9707856199e 100644 --- a/python/sglang/srt/models/laguna.py +++ b/python/sglang/srt/models/laguna.py @@ -462,12 +462,12 @@ def __init__( ( declare_ffn( sparse=self.is_layer_sparse, - next_sparse=is_next_layer_sparse, + next_layer_sparse=is_next_layer_sparse, ), self.post_attention_layernorm, ), previous=declare_ffn( - sparse=is_previous_layer_sparse, next_sparse=self.is_layer_sparse + sparse=is_previous_layer_sparse, next_layer_sparse=self.is_layer_sparse ) if layer_id != 0 else None, @@ -482,7 +482,7 @@ def forward( capture_output=None, ) -> torch.Tensor: hidden_states = self.attn_boundary.prepare( - hidden_states, forward_batch, capture_output=capture_output + hidden_states, forward_batch, capture=capture_output ) if hidden_states.shape[0] != 0: hidden_states = self.self_attn( @@ -610,11 +610,11 @@ def forward( hidden_states = residual_batch.complete_output(hidden_states, forward_batch) if hidden_states.shape[0] != 0: - hidden_states = residual_batch.norm( + hidden_states = residual_batch.final_norm( hidden_states, forward_batch, self.norm, - capture_output=aux_hidden_states.append + capture=aux_hidden_states.append if self.end_layer in self.layers_to_capture else None, ) diff --git a/python/sglang/srt/models/llada2.py b/python/sglang/srt/models/llada2.py index 5e10a9e2956a..a36e3708c0e6 100644 --- a/python/sglang/srt/models/llada2.py +++ b/python/sglang/srt/models/llada2.py @@ -37,7 +37,7 @@ from sglang.srt.layers.layer_boundary import ( declare_attn, declare_ffn, - enable_moe_dense_fully_dp, + is_dense_ffn_fully_dp, make_stages, ) from sglang.srt.layers.layer_boundary.residual import batch as residual_batch @@ -624,7 +624,7 @@ def __init__( prefix=add_prefix("mlp", prefix), ) else: - if enable_moe_dense_fully_dp(): + if is_dense_ffn_fully_dp(): mlp_tp_rank, mlp_tp_size = 0, 1 else: mlp_tp_rank, mlp_tp_size = None, None @@ -644,12 +644,12 @@ def __init__( ( declare_ffn( sparse=self.is_layer_sparse, - next_sparse=is_next_layer_sparse, + next_layer_sparse=is_next_layer_sparse, ), self.post_attention_layernorm, ), previous=declare_ffn( - sparse=is_previous_layer_sparse, next_sparse=self.is_layer_sparse + sparse=is_previous_layer_sparse, next_layer_sparse=self.is_layer_sparse ) if layer_id != 0 else None, @@ -764,7 +764,7 @@ def forward( else: hidden_states = residual_batch.complete_output(hidden_states, forward_batch) if not forward_batch.forward_mode.is_idle(): - hidden_states = residual_batch.norm( + hidden_states = residual_batch.final_norm( hidden_states, forward_batch, self.norm ) return hidden_states diff --git a/python/sglang/srt/models/llama4.py b/python/sglang/srt/models/llama4.py index 48d1e007db35..97ac09ce48b3 100644 --- a/python/sglang/srt/models/llama4.py +++ b/python/sglang/srt/models/llama4.py @@ -432,11 +432,13 @@ def __init__( ( declare_ffn( sparse=is_moe_layer, - next_sparse=is_next_moe_layer, + next_layer_sparse=is_next_moe_layer, ), self.post_attention_layernorm, ), - previous=declare_ffn(sparse=is_previous_moe_layer, next_sparse=is_moe_layer) + previous=declare_ffn( + sparse=is_previous_moe_layer, next_layer_sparse=is_moe_layer + ) if layer_id != 0 else None, terminal=layer_id == config.num_hidden_layers - 1, @@ -461,7 +463,7 @@ def forward( capture_output=None, ) -> torch.Tensor: hidden_states = self.attn_boundary.prepare( - hidden_states, forward_batch, capture_output=capture_output + hidden_states, forward_batch, capture=capture_output ) if hidden_states.shape[0] != 0: @@ -538,7 +540,9 @@ def forward( hidden_states = residual_batch.complete_output(hidden_states, forward_batch) if not forward_batch.forward_mode.is_idle(): - hidden_states = residual_batch.norm(hidden_states, forward_batch, self.norm) + hidden_states = residual_batch.final_norm( + hidden_states, forward_batch, self.norm + ) if len(aux_hidden_states) == 0: return hidden_states @@ -578,7 +582,7 @@ def forward_split_prefill( ) if end != self.model.config.num_hidden_layers: return None - forward_batch.hidden_states = residual_batch.norm( + forward_batch.hidden_states = residual_batch.final_norm( forward_batch.hidden_states, forward_batch, self.model.norm ) return self.logits_processor( diff --git a/python/sglang/srt/models/longcat_flash.py b/python/sglang/srt/models/longcat_flash.py index 1ed7c296e4d6..a176ca84df79 100644 --- a/python/sglang/srt/models/longcat_flash.py +++ b/python/sglang/srt/models/longcat_flash.py @@ -438,10 +438,10 @@ def __init__( {"qkv_latent_func": self.self_attn[0].prepare_qkv_latent}, ), ( - declare_ffn(sparse=True, next_sparse=True), + declare_ffn(sparse=True, next_layer_sparse=True), self.post_attention_layernorm[0], ), - previous=declare_ffn(sparse=True, next_sparse=True) + previous=declare_ffn(sparse=True, next_layer_sparse=True) if self.layer_id != 0 else None, terminal=self.layer_id == config.num_hidden_layers - 1, @@ -642,7 +642,7 @@ def forward( topk_indices, ) - hidden_states = residual_batch.norm( + hidden_states = residual_batch.final_norm( hidden_states, forward_batch, self.norm, skip_empty=True ) diff --git a/python/sglang/srt/models/longcat_flash_nextn.py b/python/sglang/srt/models/longcat_flash_nextn.py index 90f71e26474a..28ab746ed23b 100644 --- a/python/sglang/srt/models/longcat_flash_nextn.py +++ b/python/sglang/srt/models/longcat_flash_nextn.py @@ -168,7 +168,7 @@ def __init__( self.attn_boundary, self.ffn_boundary = make_stages( (declare_attn(), self.input_layernorm), ( - declare_ffn(sparse=False, next_sparse=False), + declare_ffn(sparse=False, next_layer_sparse=False), self.post_attention_layernorm, ), previous=declare_ffn() if self.layer_id != 0 else None, @@ -195,7 +195,9 @@ def forward( hidden_states = self.attn_boundary.finish(hidden_states, forward_batch) hidden_states = self.ffn_boundary.prepare(hidden_states, forward_batch) hidden_states = self.mlp(hidden_states) - hidden_states = self.ffn_boundary.postprocess(hidden_states, forward_batch) + hidden_states = self.ffn_boundary.finish_complete_output( + hidden_states, forward_batch + ) return hidden_states @@ -275,7 +277,7 @@ def forward( hidden_states = residual_batch.complete_output(hidden_states, forward_batch) if not forward_batch.forward_mode.is_idle(): - hidden_states = residual_batch.norm( + hidden_states = residual_batch.final_norm( hidden_states, forward_batch, self.final_layernorm ) return hidden_states diff --git a/python/sglang/srt/models/mellum.py b/python/sglang/srt/models/mellum.py index 5a7d39498074..607e089b5c39 100644 --- a/python/sglang/srt/models/mellum.py +++ b/python/sglang/srt/models/mellum.py @@ -454,12 +454,12 @@ def _is_sparse(lid: int) -> bool: ( declare_ffn( sparse=self.is_layer_sparse, - next_sparse=is_next_layer_sparse, + next_layer_sparse=is_next_layer_sparse, ), self.post_attention_layernorm, ), previous=declare_ffn( - sparse=is_previous_layer_sparse, next_sparse=self.is_layer_sparse + sparse=is_previous_layer_sparse, next_layer_sparse=self.is_layer_sparse ) if layer_id != 0 else None, diff --git a/python/sglang/srt/models/mimo_v2.py b/python/sglang/srt/models/mimo_v2.py index 836c2625ffc1..886c363a8d4d 100644 --- a/python/sglang/srt/models/mimo_v2.py +++ b/python/sglang/srt/models/mimo_v2.py @@ -38,7 +38,7 @@ from sglang.srt.layers.layer_boundary import ( declare_attn, declare_ffn, - enable_moe_dense_fully_dp, + is_dense_ffn_fully_dp, make_stages, ) from sglang.srt.layers.layer_boundary.residual import batch as residual_batch @@ -840,7 +840,7 @@ def __init__( layer_id=layer_id, ) else: - if enable_moe_dense_fully_dp(): + if is_dense_ffn_fully_dp(): mlp_tp_rank, mlp_tp_size = 0, 1 else: mlp_tp_rank, mlp_tp_size = None, None @@ -864,12 +864,12 @@ def __init__( ( declare_ffn( sparse=self.is_layer_sparse, - next_sparse=is_next_layer_sparse, + next_layer_sparse=is_next_layer_sparse, ), self.post_attention_layernorm, ), previous=declare_ffn( - sparse=is_previous_layer_sparse, next_sparse=self.is_layer_sparse + sparse=is_previous_layer_sparse, next_layer_sparse=self.is_layer_sparse ) if layer_id != 0 else None, @@ -887,7 +887,7 @@ def forward( hidden_states = self.attn_boundary.prepare( hidden_states, forward_batch, - captured_last_layer_outputs=captured_last_layer_outputs, + capture_gathered=captured_last_layer_outputs, ) if hidden_states.shape[0] != 0: @@ -945,7 +945,7 @@ def op_comm_prepare_mlp(self, state): ) def op_comm_postprocess_layer(self, state): - hidden_states = self.ffn_boundary.postprocess( + hidden_states = self.ffn_boundary.finish_complete_output( state.pop("hidden_states_mlp_output"), state.forward_batch ) @@ -1092,7 +1092,7 @@ def forward( hidden_states_before_norm = residual_batch.snapshot( hidden_states, forward_batch ) - hidden_states = residual_batch.norm( + hidden_states = residual_batch.final_norm( hidden_states, forward_batch, self.norm ) diff --git a/python/sglang/srt/models/mimo_v2_nextn.py b/python/sglang/srt/models/mimo_v2_nextn.py index 95f48457ac1f..791610334966 100644 --- a/python/sglang/srt/models/mimo_v2_nextn.py +++ b/python/sglang/srt/models/mimo_v2_nextn.py @@ -27,7 +27,7 @@ from sglang.srt.layers.layer_boundary import ( declare_attn, declare_ffn, - enable_moe_dense_fully_dp, + is_dense_ffn_fully_dp, make_stages, ) from sglang.srt.layers.layer_boundary.residual import batch as residual_batch @@ -101,7 +101,7 @@ def __init__( is_previous_layer_sparse = True is_next_layer_sparse = False - if enable_moe_dense_fully_dp(): + if is_dense_ffn_fully_dp(): mlp_tp_rank, mlp_tp_size = 0, 1 else: mlp_tp_rank, mlp_tp_size = None, None @@ -124,12 +124,12 @@ def __init__( ( declare_ffn( sparse=self.is_layer_sparse, - next_sparse=is_next_layer_sparse, + next_layer_sparse=is_next_layer_sparse, ), self.post_attention_layernorm, ), previous=declare_ffn( - sparse=is_previous_layer_sparse, next_sparse=self.is_layer_sparse + sparse=is_previous_layer_sparse, next_layer_sparse=self.is_layer_sparse ) if layer_id != 0 else None, @@ -235,7 +235,7 @@ def forward( hidden_states_before_norm = residual_batch.snapshot( hidden_states, forward_batch ) - hidden_states = residual_batch.norm( + hidden_states = residual_batch.final_norm( hidden_states, forward_batch, self.final_layernorm ) diff --git a/python/sglang/srt/models/minimax_m2.py b/python/sglang/srt/models/minimax_m2.py index 39a3539cecad..7a418a6a74f2 100644 --- a/python/sglang/srt/models/minimax_m2.py +++ b/python/sglang/srt/models/minimax_m2.py @@ -988,12 +988,12 @@ def __init__( ( declare_ffn( sparse=self.is_layer_sparse, - next_sparse=is_next_layer_sparse, + next_layer_sparse=is_next_layer_sparse, ), self.post_attention_layernorm, ), previous=declare_ffn( - sparse=is_previous_layer_sparse, next_sparse=self.is_layer_sparse + sparse=is_previous_layer_sparse, next_layer_sparse=self.is_layer_sparse ) if layer_id != 0 else None, @@ -1011,7 +1011,7 @@ def forward( hidden_states = self.attn_boundary.prepare( hidden_states, forward_batch, - captured_last_layer_outputs=captured_last_layer_outputs, + capture_gathered=captured_last_layer_outputs, ) if not forward_batch.forward_mode.is_idle(): hidden_states = self.self_attn( @@ -1065,7 +1065,7 @@ def op_comm_prepare_mlp(self, state): def op_comm_postprocess_layer(self, state): """Communication postprocess for layer - TBO operation""" - hidden_states = self.ffn_boundary.postprocess( + hidden_states = self.ffn_boundary.finish_complete_output( state.pop("hidden_states_mlp_output"), state.forward_batch ) @@ -1180,7 +1180,9 @@ def forward( hidden_states = residual_batch.complete_output(hidden_states, forward_batch) if hidden_states.shape[0] != 0: - hidden_states = residual_batch.norm(hidden_states, forward_batch, self.norm) + hidden_states = residual_batch.final_norm( + hidden_states, forward_batch, self.norm + ) if len(aux_hidden_states) == 0: return hidden_states diff --git a/python/sglang/srt/models/minimax_m3.py b/python/sglang/srt/models/minimax_m3.py index f08b2c1a8d7e..2cb86c4d5f9b 100644 --- a/python/sglang/srt/models/minimax_m3.py +++ b/python/sglang/srt/models/minimax_m3.py @@ -45,7 +45,7 @@ from sglang.srt.layers.layer_boundary import ( declare_attn, declare_ffn, - enable_moe_dense_fully_dp, + is_dense_ffn_fully_dp, make_stages, ) from sglang.srt.layers.layer_boundary.residual import batch as residual_batch @@ -1317,7 +1317,7 @@ def __init__( prefix=add_prefix("mlp", prefix), ) else: - if enable_moe_dense_fully_dp(): + if is_dense_ffn_fully_dp(): mlp_tp_rank, mlp_tp_size = 0, 1 else: mlp_tp_rank, mlp_tp_size = None, None @@ -1359,12 +1359,12 @@ def _is_layer_sparse(lid): ( declare_ffn( sparse=self.is_layer_sparse, - next_sparse=is_next_layer_sparse, + next_layer_sparse=is_next_layer_sparse, ), self.post_attention_layernorm, ), previous=declare_ffn( - sparse=is_previous_layer_sparse, next_sparse=self.is_layer_sparse + sparse=is_previous_layer_sparse, next_layer_sparse=self.is_layer_sparse ) if layer_id != 0 else None, @@ -1382,7 +1382,7 @@ def forward( hidden_states = self.attn_boundary.prepare( hidden_states, forward_batch, - captured_last_layer_outputs=captured_last_layer_outputs, + capture_gathered=captured_last_layer_outputs, **kwargs, ) @@ -1519,7 +1519,9 @@ def forward( return residual_batch.to_pp(hidden_states, forward_batch) hidden_states = residual_batch.complete_output(hidden_states, forward_batch) if hidden_states.shape[0] != 0: - hidden_states = residual_batch.norm(hidden_states, forward_batch, self.norm) + hidden_states = residual_batch.final_norm( + hidden_states, forward_batch, self.norm + ) if len(aux_hidden_states) == 0: return hidden_states diff --git a/python/sglang/srt/models/moss_vl.py b/python/sglang/srt/models/moss_vl.py index a6b54f31b391..4f8a03c442f9 100644 --- a/python/sglang/srt/models/moss_vl.py +++ b/python/sglang/srt/models/moss_vl.py @@ -1020,7 +1020,7 @@ def __init__( self.attn_boundary, self.ffn_boundary = make_stages( (declare_attn(), self.input_layernorm), ( - declare_ffn(sparse=False, next_sparse=False), + declare_ffn(sparse=False, next_layer_sparse=False), self.post_attention_layernorm, ), previous=declare_ffn() if layer_id != 0 else None, @@ -1120,7 +1120,9 @@ def forward( positions=positions, vision_position_ids=vision_position_ids, ) - hidden_states = residual_batch.written(hidden_states, forward_batch) + hidden_states = residual_batch.set_written( + hidden_states, forward_batch + ) elif isinstance(decoder_layer, MossVLSelfAttentionDecoderLayer): hidden_states = decoder_layer( positions=positions, @@ -1130,7 +1132,9 @@ def forward( else: raise ValueError(f"Unknown decoder layer type {type(decoder_layer)}") - hidden_states = residual_batch.norm(hidden_states, forward_batch, self.norm) + hidden_states = residual_batch.final_norm( + hidden_states, forward_batch, self.norm + ) return hidden_states diff --git a/python/sglang/srt/models/nemotron_h.py b/python/sglang/srt/models/nemotron_h.py index f361910552ee..3440b01f0d7e 100644 --- a/python/sglang/srt/models/nemotron_h.py +++ b/python/sglang/srt/models/nemotron_h.py @@ -759,7 +759,9 @@ def forward( aux_hidden_states.append( residual_batch.snapshot(hidden_states, forward_batch) ) - hidden_states = residual_batch.norm(hidden_states, forward_batch, self.norm_f) + hidden_states = residual_batch.final_norm( + hidden_states, forward_batch, self.norm_f + ) if aux_hidden_states: return hidden_states, aux_hidden_states return hidden_states diff --git a/python/sglang/srt/models/nemotron_h_mtp.py b/python/sglang/srt/models/nemotron_h_mtp.py index 1392624d74cf..14a77ac25260 100644 --- a/python/sglang/srt/models/nemotron_h_mtp.py +++ b/python/sglang/srt/models/nemotron_h_mtp.py @@ -105,7 +105,7 @@ def forward( ) if self.has_start_projections: - hidden_states = residual_batch.written(hidden_states, forward_batch) + hidden_states = residual_batch.set_written(hidden_states, forward_batch) hidden_states = super().forward( hidden_states=hidden_states, @@ -115,7 +115,7 @@ def forward( if self.has_end_norm: hidden_states = residual_batch.fold(hidden_states, forward_batch) - hidden_states = residual_batch.written( + hidden_states = residual_batch.set_written( self.final_layernorm(hidden_states), forward_batch ) @@ -187,7 +187,7 @@ def forward( ) if self.has_start_projections: - hidden_states = residual_batch.written(hidden_states, forward_batch) + hidden_states = residual_batch.set_written(hidden_states, forward_batch) hidden_states = super().forward( hidden_states=hidden_states, @@ -197,7 +197,7 @@ def forward( if self.has_end_norm: hidden_states = residual_batch.fold(hidden_states, forward_batch) - hidden_states = residual_batch.written( + hidden_states = residual_batch.set_written( self.final_layernorm(hidden_states), forward_batch ) diff --git a/python/sglang/srt/models/nemotron_h_utils.py b/python/sglang/srt/models/nemotron_h_utils.py index bce3b6ac0065..afed7add0277 100644 --- a/python/sglang/srt/models/nemotron_h_utils.py +++ b/python/sglang/srt/models/nemotron_h_utils.py @@ -2,13 +2,16 @@ from sglang.srt.configs.nemotron_h import ATTENTION, MAMBA, MOE from sglang.srt.layers.layer_boundary import ( - HandoffRows, + ExitRows, ProducerReduction, declare_attn, declare_ffn, make_stages, ) -from sglang.srt.layers.layer_boundary.residual.add_norm import NormQuantRead, NormRead +from sglang.srt.layers.layer_boundary.residual.add_norm import ( + NormQuantReadout, + NormReadout, +) from sglang.srt.layers.layernorm import RMSNorm from sglang.srt.runtime_context import get_parallel @@ -22,15 +25,15 @@ def is_attn_layer(layer_type: str) -> bool: def _declaration(pattern: str, layer_idx: int): if is_attn_layer(pattern[layer_idx]): return declare_attn( - read=NormQuantRead(before_gather=True), - reduction=ProducerReduction.SCOPED, - gathers_tp_input=False, + read=NormQuantReadout(reads_before_dp_gather=True), + reduction=ProducerReduction.EXIT_SCOPED, + gathers_attn_tp_input=False, ) return declare_ffn( sparse=pattern[layer_idx] == MOE, - read=NormRead(before_gather=True), + read=NormReadout(reads_before_dp_gather=True), dense_tp_size=get_parallel().tp_size if pattern[layer_idx] != MOE else None, - handoff_rows=HandoffRows.ATTENTION, + exit_rows=ExitRows.ATTENTION, ) diff --git a/python/sglang/srt/models/qwen2.py b/python/sglang/srt/models/qwen2.py index 588aecf7939a..793b326e4c87 100644 --- a/python/sglang/srt/models/qwen2.py +++ b/python/sglang/srt/models/qwen2.py @@ -28,7 +28,7 @@ from sglang.srt.layers.activation import SiluAndMul from sglang.srt.layers.dp_attention import is_dp_attention_enabled from sglang.srt.layers.layer_boundary.residual.access import ( - finish_layer_stack, + export_output, from_pp, snapshot, ) @@ -427,9 +427,7 @@ def forward( residual, ) - hidden_states, residual = finish_layer_stack( - hidden_states, residual, forward_batch - ) + hidden_states, residual = export_output(hidden_states, residual, forward_batch) if not self.pp_group.is_last_rank: return PPProxyTensors( { diff --git a/python/sglang/srt/models/qwen2_moe.py b/python/sglang/srt/models/qwen2_moe.py index 8755df797ef8..3eff6d9de4c7 100644 --- a/python/sglang/srt/models/qwen2_moe.py +++ b/python/sglang/srt/models/qwen2_moe.py @@ -866,10 +866,10 @@ def forward( if shared_output is None: raise RuntimeError("Qwen deferred finalize requires shared output") from sglang.srt.layers.layer_boundary.fusions.cutedsl import ( - MoeFinalizeHandoff, + MoeDeferredFinalize, ) - return MoeFinalizeHandoff.from_flashinfer( + return MoeDeferredFinalize.from_flashinfer( final_hidden_states, gated_shared_output=shared_output, m=num_tokens, @@ -1058,12 +1058,12 @@ def __init__( ( declare_ffn( sparse=self.is_layer_sparse, - next_sparse=is_next_layer_sparse, + next_layer_sparse=is_next_layer_sparse, ), self.post_attention_layernorm, ), previous=declare_ffn( - sparse=is_previous_layer_sparse, next_sparse=self.is_layer_sparse + sparse=is_previous_layer_sparse, next_layer_sparse=self.is_layer_sparse ) if layer_id != 0 else None, @@ -1082,7 +1082,7 @@ def forward( hidden_states = self.attn_boundary.prepare( hidden_states, forward_batch, - captured_last_layer_outputs=captured_last_layer_outputs, + capture_gathered=captured_last_layer_outputs, **kwargs, ) @@ -1216,7 +1216,7 @@ def forward( else: hidden_states = residual_batch.complete_output(hidden_states, forward_batch) if hidden_states.shape[0] != 0: - hidden_states = residual_batch.norm( + hidden_states = residual_batch.final_norm( hidden_states, forward_batch, self.norm ) @@ -1314,7 +1314,7 @@ def forward_split_prefill( forward_batch.hidden_states, forward_batch ) # norm - hidden_states = residual_batch.norm( + hidden_states = residual_batch.final_norm( forward_batch.hidden_states, forward_batch, self.model.norm ) forward_batch.hidden_states = hidden_states diff --git a/python/sglang/srt/models/qwen3.py b/python/sglang/srt/models/qwen3.py index 2939fc608d82..94ad78b95b0a 100644 --- a/python/sglang/srt/models/qwen3.py +++ b/python/sglang/srt/models/qwen3.py @@ -379,7 +379,7 @@ def __init__( self.attn_boundary, self.ffn_boundary = make_stages( (declare_attn(), self.input_layernorm), ( - declare_ffn(sparse=False, next_sparse=False), + declare_ffn(sparse=False, next_layer_sparse=False), self.post_attention_layernorm, ), previous=declare_ffn() if layer_id != 0 else None, @@ -399,7 +399,7 @@ def forward( hidden_states, forward_batch, post_residual_addition=post_residual_addition, - capture_output=capture_output, + capture=capture_output, ) if hidden_states.shape[0] != 0: hidden_states = self.self_attn( @@ -479,7 +479,9 @@ def forward( return residual_batch.to_pp(hidden_states, forward_batch) hidden_states = residual_batch.complete_output(hidden_states, forward_batch) if hidden_states.shape[0] != 0: - hidden_states = residual_batch.norm(hidden_states, forward_batch, self.norm) + hidden_states = residual_batch.final_norm( + hidden_states, forward_batch, self.norm + ) return ( (hidden_states, aux_hidden_states) if aux_hidden_states else hidden_states ) @@ -607,7 +609,7 @@ def forward_split_prefill( ) if end == self.model.config.num_hidden_layers: - forward_batch.hidden_states = residual_batch.norm( + forward_batch.hidden_states = residual_batch.final_norm( forward_batch.hidden_states, forward_batch, self.model.norm ) # logits process diff --git a/python/sglang/srt/models/qwen3_5.py b/python/sglang/srt/models/qwen3_5.py index fac16f30a38a..c1d17ee1221a 100644 --- a/python/sglang/srt/models/qwen3_5.py +++ b/python/sglang/srt/models/qwen3_5.py @@ -60,7 +60,10 @@ from sglang.srt.layers.layer_boundary.residual import batch as residual_batch # Layers - Others -from sglang.srt.layers.layer_boundary.residual.add_norm import Fp8Input, NormQuantRead +from sglang.srt.layers.layer_boundary.residual.add_norm import ( + Fp8Input, + NormQuantReadout, +) from sglang.srt.layers.layernorm import GemmaRMSNorm # Layers - Linear @@ -1097,7 +1100,7 @@ def __init__( self.attn_boundary, self.ffn_boundary = make_stages( ( declare_attn( - read=NormQuantRead( + read=NormQuantReadout( fp8_input=Fp8Input.TUPLE_AND_BF16 if accepts_fp8_input else None ) ), @@ -1109,13 +1112,13 @@ def __init__( ( declare_ffn( sparse=is_layer_sparse, - next_sparse=is_next_layer_sparse, + next_layer_sparse=is_next_layer_sparse, ), self.post_attention_layernorm, {"fusions": boundary_fusions}, ), previous=declare_ffn( - sparse=is_previous_layer_sparse, next_sparse=is_layer_sparse + sparse=is_previous_layer_sparse, next_layer_sparse=is_layer_sparse ) if layer_id != 0 else None, @@ -1132,7 +1135,7 @@ def forward( hidden_states = self.attn_boundary.prepare( hidden_states, forward_batch, - captured_last_layer_outputs=kwargs.get("captured_last_layer_outputs", None), + capture_gathered=kwargs.get("captured_last_layer_outputs", None), ) # fused AR+quant hands down a (fp8, scale) / (bf16, fp8, scale) tuple @@ -1303,7 +1306,7 @@ def __init__( self.attn_boundary, self.ffn_boundary = make_stages( ( declare_attn( - read=NormQuantRead( + read=NormQuantReadout( fp8_input=Fp8Input.TUPLE if accepts_fp8_input else None ) ), @@ -1315,13 +1318,13 @@ def __init__( ( declare_ffn( sparse=is_layer_sparse, - next_sparse=is_next_layer_sparse, + next_layer_sparse=is_next_layer_sparse, ), self.post_attention_layernorm, {"fusions": boundary_fusions}, ), previous=declare_ffn( - sparse=is_previous_layer_sparse, next_sparse=is_layer_sparse + sparse=is_previous_layer_sparse, next_layer_sparse=is_layer_sparse ) if layer_id != 0 else None, @@ -1536,7 +1539,7 @@ def forward( hidden_states = self.attn_boundary.prepare( hidden_states, forward_batch, - captured_last_layer_outputs=captured_last_layer_outputs, + capture_gathered=captured_last_layer_outputs, ) # fused AR+quant hands down a (fp8, scale) / (bf16, fp8, scale) tuple @@ -1843,11 +1846,11 @@ def forward( if not self.pp_group.is_last_rank: return residual_batch.to_pp(hidden_states, forward_batch) - hidden_states = residual_batch.norm( + hidden_states = residual_batch.final_norm( hidden_states, forward_batch, self.norm, - handoff_norm=self.flashinfer_mnnvl_cutedsl_fusion, + finalize_norm=self.flashinfer_mnnvl_cutedsl_fusion, skip_empty=True, ) diff --git a/python/sglang/srt/models/qwen3_moe.py b/python/sglang/srt/models/qwen3_moe.py index d99f61687b5a..ff033ea5dd65 100644 --- a/python/sglang/srt/models/qwen3_moe.py +++ b/python/sglang/srt/models/qwen3_moe.py @@ -871,12 +871,12 @@ def __init__( ( declare_ffn( sparse=self.is_layer_sparse, - next_sparse=is_next_layer_sparse, + next_layer_sparse=is_next_layer_sparse, ), self.post_attention_layernorm, ), previous=declare_ffn( - sparse=is_previous_layer_sparse, next_sparse=self.is_layer_sparse + sparse=is_previous_layer_sparse, next_layer_sparse=self.is_layer_sparse ) if layer_id != 0 else None, @@ -896,8 +896,8 @@ def forward( hidden_states = self.attn_boundary.prepare( hidden_states, forward_batch, - captured_last_layer_outputs=captured_last_layer_outputs, - capture_output=capture_output, + capture_gathered=captured_last_layer_outputs, + capture=capture_output, **kwargs, ) @@ -945,7 +945,7 @@ def op_comm_prepare_mlp(self, state): ) def op_comm_postprocess_layer(self, state): - hidden_states = self.ffn_boundary.postprocess( + hidden_states = self.ffn_boundary.finish_complete_output( state.pop("hidden_states_mlp_output"), state.forward_batch ) @@ -1103,7 +1103,7 @@ def forward_split_prefill( forward_batch.hidden_states, forward_batch ) # norm - hidden_states = residual_batch.norm( + hidden_states = residual_batch.final_norm( forward_batch.hidden_states, forward_batch, self.model.norm ) forward_batch.hidden_states = hidden_states diff --git a/python/sglang/srt/models/qwen3_next.py b/python/sglang/srt/models/qwen3_next.py index ccfd7c5656f4..20ba669bbc7b 100644 --- a/python/sglang/srt/models/qwen3_next.py +++ b/python/sglang/srt/models/qwen3_next.py @@ -542,12 +542,12 @@ def __init__( ( declare_ffn( sparse=self.is_layer_sparse, - next_sparse=is_next_layer_sparse, + next_layer_sparse=is_next_layer_sparse, ), self.post_attention_layernorm, ), previous=declare_ffn( - sparse=is_previous_layer_sparse, next_sparse=self.is_layer_sparse + sparse=is_previous_layer_sparse, next_layer_sparse=self.is_layer_sparse ) if layer_id != 0 else None, @@ -565,7 +565,7 @@ def forward( hidden_states = self.attn_boundary.prepare( hidden_states, forward_batch, - captured_last_layer_outputs=captured_last_layer_outputs, + capture_gathered=captured_last_layer_outputs, ) if not forward_batch.forward_mode.is_idle(): @@ -710,12 +710,12 @@ def __init__( ( declare_ffn( sparse=self.is_layer_sparse, - next_sparse=is_next_layer_sparse, + next_layer_sparse=is_next_layer_sparse, ), self.post_attention_layernorm, ), previous=declare_ffn( - sparse=is_previous_layer_sparse, next_sparse=self.is_layer_sparse + sparse=is_previous_layer_sparse, next_layer_sparse=self.is_layer_sparse ) if layer_id != 0 else None, @@ -835,7 +835,7 @@ def forward( hidden_states = self.attn_boundary.prepare( hidden_states, forward_batch, - captured_last_layer_outputs=captured_last_layer_outputs, + capture_gathered=captured_last_layer_outputs, ) if not forward_batch.forward_mode.is_idle(): @@ -945,7 +945,9 @@ def forward( hidden_states = residual_batch.complete_output(hidden_states, forward_batch) if not forward_batch.forward_mode.is_idle(): - hidden_states = residual_batch.norm(hidden_states, forward_batch, self.norm) + hidden_states = residual_batch.final_norm( + hidden_states, forward_batch, self.norm + ) if len(aux_hidden_states) == 0: return hidden_states diff --git a/python/sglang/srt/models/qwen3_vl.py b/python/sglang/srt/models/qwen3_vl.py index 6f9faffd5684..d3076ca6e67c 100644 --- a/python/sglang/srt/models/qwen3_vl.py +++ b/python/sglang/srt/models/qwen3_vl.py @@ -1269,7 +1269,7 @@ def forward( return residual_batch.to_pp(hidden_states, forward_batch) hidden_states = residual_batch.complete_output(hidden_states, forward_batch) if hidden_states.shape[0] != 0: - hidden_states = residual_batch.norm( + hidden_states = residual_batch.final_norm( hidden_states, forward_batch, self.norm, diff --git a/python/sglang/srt/models/qwen3_vl_moe.py b/python/sglang/srt/models/qwen3_vl_moe.py index 7f709e7d10ba..5f4d31c6a2be 100644 --- a/python/sglang/srt/models/qwen3_vl_moe.py +++ b/python/sglang/srt/models/qwen3_vl_moe.py @@ -159,7 +159,7 @@ def forward( else: hidden_states = residual_batch.complete_output(hidden_states, forward_batch) if hidden_states.shape[0] != 0: - hidden_states = residual_batch.norm( + hidden_states = residual_batch.final_norm( hidden_states, forward_batch, self.norm, diff --git a/python/sglang/srt/models/sarvam_moe.py b/python/sglang/srt/models/sarvam_moe.py index a58b19ecdbe1..2a21eb3123e9 100644 --- a/python/sglang/srt/models/sarvam_moe.py +++ b/python/sglang/srt/models/sarvam_moe.py @@ -25,7 +25,7 @@ from sglang.srt.layers.layer_boundary import ( declare_attn, declare_ffn, - enable_moe_dense_fully_dp, + is_dense_ffn_fully_dp, make_stages, ) from sglang.srt.layers.layer_boundary.residual import batch as residual_batch @@ -300,7 +300,7 @@ def __init__( and config.num_shared_experts > 0 ): intermediate_size = config.moe_intermediate_size * config.num_shared_experts - if enable_moe_dense_fully_dp(): + if is_dense_ffn_fully_dp(): shared_tp_rank, shared_tp_size = 0, 1 else: shared_tp_rank, shared_tp_size = None, None @@ -1028,7 +1028,7 @@ def __init__( alt_stream=alt_stream, ) else: - if enable_moe_dense_fully_dp(): + if is_dense_ffn_fully_dp(): mlp_tp_rank, mlp_tp_size = 0, 1 else: mlp_tp_rank, mlp_tp_size = None, None @@ -1059,12 +1059,12 @@ def __init__( ( declare_ffn( sparse=self.is_layer_sparse, - next_sparse=is_next_layer_sparse, + next_layer_sparse=is_next_layer_sparse, ), self.post_attention_layernorm, ), previous=declare_ffn( - sparse=is_previous_layer_sparse, next_sparse=self.is_layer_sparse + sparse=is_previous_layer_sparse, next_layer_sparse=self.is_layer_sparse ) if layer_id != 0 else None, @@ -1172,7 +1172,9 @@ def forward( hidden_states = residual_batch.complete_output(hidden_states, forward_batch) if hidden_states.shape[0] != 0: - hidden_states = residual_batch.norm(hidden_states, forward_batch, self.norm) + hidden_states = residual_batch.final_norm( + hidden_states, forward_batch, self.norm + ) return hidden_states @@ -1275,7 +1277,7 @@ def forward_split_prefill( forward_batch.hidden_states = residual_batch.complete_output( forward_batch.hidden_states, forward_batch ) - forward_batch.hidden_states = residual_batch.norm( + forward_batch.hidden_states = residual_batch.final_norm( forward_batch.hidden_states, forward_batch, self.model.norm ) return self.logits_processor( @@ -1454,7 +1456,7 @@ def forward_split_prefill( forward_batch.hidden_states = residual_batch.complete_output( forward_batch.hidden_states, forward_batch ) - forward_batch.hidden_states = residual_batch.norm( + forward_batch.hidden_states = residual_batch.final_norm( forward_batch.hidden_states, forward_batch, self.model.norm ) diff --git a/python/sglang/srt/models/sdar.py b/python/sglang/srt/models/sdar.py index 6e8d510e0e32..e4b8bef30a7a 100644 --- a/python/sglang/srt/models/sdar.py +++ b/python/sglang/srt/models/sdar.py @@ -300,7 +300,7 @@ def __init__( ( declare_ffn( sparse=False, - next_sparse=False, + next_layer_sparse=False, ), self.post_attention_layernorm, ), @@ -414,7 +414,7 @@ def forward( else: hidden_states = residual_batch.complete_output(hidden_states, forward_batch) if not forward_batch.forward_mode.is_idle(): - hidden_states = residual_batch.norm( + hidden_states = residual_batch.final_norm( hidden_states, forward_batch, self.norm ) return hidden_states diff --git a/python/sglang/srt/models/sdar_moe.py b/python/sglang/srt/models/sdar_moe.py index f4cbd6a29bfd..d541d1513294 100644 --- a/python/sglang/srt/models/sdar_moe.py +++ b/python/sglang/srt/models/sdar_moe.py @@ -360,10 +360,10 @@ def __init__( self.attn_boundary, self.ffn_boundary = make_stages( (declare_attn(), self.input_layernorm), ( - declare_ffn(sparse=True, next_sparse=True), + declare_ffn(sparse=True, next_layer_sparse=True), self.post_attention_layernorm, ), - previous=declare_ffn(sparse=True, next_sparse=True) + previous=declare_ffn(sparse=True, next_layer_sparse=True) if layer_id != 0 else None, terminal=layer_id == config.num_hidden_layers - 1, @@ -481,7 +481,9 @@ def forward( hidden_states = residual_batch.complete_output(hidden_states, forward_batch) if not forward_batch.forward_mode.is_idle(): - hidden_states = residual_batch.norm(hidden_states, forward_batch, self.norm) + hidden_states = residual_batch.final_norm( + hidden_states, forward_batch, self.norm + ) return hidden_states diff --git a/python/sglang/srt/models/step3_vl.py b/python/sglang/srt/models/step3_vl.py index 98ccda9cd5bf..0d0fe098ba3f 100644 --- a/python/sglang/srt/models/step3_vl.py +++ b/python/sglang/srt/models/step3_vl.py @@ -379,12 +379,13 @@ def __init__( ( declare_ffn( sparse=self.is_layer_sparse, - next_sparse=self.is_next_layer_sparse, + next_layer_sparse=self.is_next_layer_sparse, ), self.post_attention_layernorm, ), previous=declare_ffn( - sparse=self.is_previous_layer_sparse, next_sparse=self.is_layer_sparse + sparse=self.is_previous_layer_sparse, + next_layer_sparse=self.is_layer_sparse, ) if layer_id != 0 else None, @@ -481,7 +482,9 @@ def forward( hidden_states = residual_batch.complete_output(hidden_states, forward_batch) if hidden_states.shape[0] != 0: - hidden_states = residual_batch.norm(hidden_states, forward_batch, self.norm) + hidden_states = residual_batch.final_norm( + hidden_states, forward_batch, self.norm + ) return hidden_states diff --git a/python/sglang/srt/models/step3p5.py b/python/sglang/srt/models/step3p5.py index c640e094dd65..166c64ba636f 100644 --- a/python/sglang/srt/models/step3p5.py +++ b/python/sglang/srt/models/step3p5.py @@ -571,12 +571,13 @@ def __init__( ( declare_ffn( sparse=self.is_moe_layer, - next_sparse=self.is_next_layer_sparse, + next_layer_sparse=self.is_next_layer_sparse, ), self.post_attention_layernorm, ), previous=declare_ffn( - sparse=self.is_previous_layer_sparse, next_sparse=self.is_moe_layer + sparse=self.is_previous_layer_sparse, + next_layer_sparse=self.is_moe_layer, ) if not (0 if is_nextn else layer_id) == 0 else None, @@ -721,7 +722,7 @@ def forward( hidden_states_before_norm = residual_batch.snapshot( hidden_states, forward_batch ) - hidden_states = residual_batch.norm( + hidden_states = residual_batch.final_norm( hidden_states, forward_batch, self.norm ) return hidden_states, hidden_states_before_norm diff --git a/python/sglang/srt/models/step3p5_mtp.py b/python/sglang/srt/models/step3p5_mtp.py index fd4696eb6b83..72411024bd98 100644 --- a/python/sglang/srt/models/step3p5_mtp.py +++ b/python/sglang/srt/models/step3p5_mtp.py @@ -123,7 +123,7 @@ def forward( hidden_states_before_norm = residual_batch.snapshot( hidden_states, forward_batch ) - hidden_states = residual_batch.norm( + hidden_states = residual_batch.final_norm( hidden_states, forward_batch, self.shared_head.norm ) diff --git a/python/sglang/srt/models/xllm.py b/python/sglang/srt/models/xllm.py index 59a5e77a09b0..79717bf6b82f 100644 --- a/python/sglang/srt/models/xllm.py +++ b/python/sglang/srt/models/xllm.py @@ -43,7 +43,7 @@ from sglang.srt.layers.layer_boundary import ( declare_attn, declare_ffn, - enable_moe_dense_fully_dp, + is_dense_ffn_fully_dp, make_stages, ) from sglang.srt.layers.layer_boundary.residual import batch as residual_batch @@ -1512,7 +1512,7 @@ def _is_sparse(lid): prefix=add_prefix("mlp", prefix), ) else: - if enable_moe_dense_fully_dp(): + if is_dense_ffn_fully_dp(): mlp_tp_rank, mlp_tp_size = 0, 1 else: mlp_tp_rank, mlp_tp_size = None, None @@ -1533,12 +1533,12 @@ def _is_sparse(lid): ( declare_ffn( sparse=self.is_layer_sparse, - next_sparse=is_next_layer_sparse, + next_layer_sparse=is_next_layer_sparse, ), self.post_attention_layernorm, ), previous=declare_ffn( - sparse=is_previous_layer_sparse, next_sparse=self.is_layer_sparse + sparse=is_previous_layer_sparse, next_layer_sparse=self.is_layer_sparse ) if layer_id != 0 else None, @@ -1648,7 +1648,7 @@ def forward( else: hidden_states = residual_batch.complete_output(hidden_states, forward_batch) if hidden_states.shape[0] != 0: - hidden_states = residual_batch.norm( + hidden_states = residual_batch.final_norm( hidden_states, forward_batch, self.norm ) diff --git a/python/sglang/test/boundary_fixtures.py b/python/sglang/test/boundary_fixtures.py index bce61d7abab5..2314ff84b480 100644 --- a/python/sglang/test/boundary_fixtures.py +++ b/python/sglang/test/boundary_fixtures.py @@ -8,8 +8,8 @@ declare_ffn, make_stages, ) -from sglang.srt.layers.layer_boundary.ops import identity_output -from sglang.srt.layers.layer_boundary.residual.add_norm import PLAIN_RESIDUAL +from sglang.srt.layers.layer_boundary.ops import keep_output +from sglang.srt.layers.layer_boundary.residual.add_norm import PLAIN_RESIDUAL_OPS from sglang.srt.layers.layer_boundary.stage import StageBoundary @@ -21,8 +21,8 @@ def make_test_stages( last=False, sparse=False, previous_sparse=False, - next_sparse=False, - residual=PLAIN_RESIDUAL, + next_layer_sparse=False, + residual=PLAIN_RESIDUAL_OPS, output=None, **options, ): @@ -30,16 +30,14 @@ def make_test_stages( None if first else declare_ffn( - sparse=previous_sparse, next_sparse=sparse, update=residual.ffn_update + sparse=previous_sparse, next_layer_sparse=sparse, update=residual.ffn_update ) ) - attention = declare_attn( - read=residual.attention_read, update=residual.attention_update - ) + attention = declare_attn(read=residual.attn_readout, update=residual.attn_update) ffn = declare_ffn( sparse=sparse, - next_sparse=next_sparse, - read=residual.ffn_read, + next_layer_sparse=next_layer_sparse, + read=residual.ffn_readout, update=residual.ffn_update, output_transform=output, ) @@ -63,8 +61,8 @@ def stub_plan(): plan.norm = None plan.fusions = None plan._publish_lora_layout = False - plan._fusion_rows = None - plan._paths = {} + plan._next_input_rows = None + plan.paths = {} plan.enters_stack = False return plan @@ -83,29 +81,29 @@ def stub_stage(plan, kind): def sp_region_steps(): """Local SP rows with no owed sum, for tests of activation and exits.""" from sglang.srt.layers.layer_boundary import ( - NORM_QUANT_READ, - EdgeDecl, + NORM_QUANT_READOUT, + EdgeContract, + EntryPath, + InputContract, Layout, - StageEntry, - StageInput, - StageOutput, - StageSteps, + OutputContract, + StagePath, TokenAxis, - make_boundary, + bind_entry, ) - from sglang.srt.layers.layer_boundary.prepare import _hand_qkv_hook_its_input + from sglang.srt.layers.layer_boundary.prepare import _attn_input_default - rows = Layout(frozenset({TokenAxis.ATTN_TP_SCATTER})) - output = StageOutput(rows) + rows = Layout(frozenset({TokenAxis.ATTN_TP})) + output = OutputContract(rows) - def entry(read, handoff=None): - selected = make_boundary( - EdgeDecl(output, StageInput(rows, read=read), rows, rows) + def entry(read, attn_input_adapter=None): + selected = bind_entry( + EdgeContract(output, InputContract(rows, read=read), rows, rows) ) - return StageEntry(selected.prepare, rows, handoff=handoff) + return EntryPath(selected.prepare, rows, attn_input_adapter=attn_input_adapter) - return StageSteps( - entry(NORM_QUANT_READ, _hand_qkv_hook_its_input), output, identity_output + return StagePath( + entry(NORM_QUANT_READOUT, _attn_input_default), output, keep_output ) @@ -120,16 +118,16 @@ def prepare_attention(stage, hidden, residual, forward_batch, *args, **call): def prepare_raw(stage, method, hidden, residual, forward_batch, *args, **call): - from sglang.srt.layers.layer_boundary.residual.add_norm import ADD + from sglang.srt.layers.layer_boundary.residual.add_norm import PLAIN_ADD from sglang.srt.layers.layer_boundary.residual.stream import ResidualStream if not isinstance(residual, ResidualStream): stream = ResidualStream(residual) if residual is not None: - hidden = stream.leave( + hidden = stream.record( hidden, - call.get("update", ADD), - declared_sum=stage.entry(forward_batch).input_sum, + call.get("update", PLAIN_ADD), + declared_sum=stage.entry(forward_batch).declared_sum, ) else: stream = residual @@ -155,7 +153,7 @@ def postprocess_output(boundary, hidden, residual, forward_batch): stream = ( residual if isinstance(residual, ResidualStream) else ResidualStream(residual) ) - hidden = boundary.postprocess_layer(hidden, stream, forward_batch) + hidden = boundary.finish_complete_output(hidden, stream, forward_batch) if isinstance(residual, ResidualStream): return hidden, stream return (hidden, None) if stream.pending is None else stream.input(hidden) diff --git a/test/registered/kernels/ops/layernorm/test_mhc_kernels.py b/test/registered/kernels/ops/layernorm/test_mhc_kernels.py index 2a5c6fdf4844..922b607fcee0 100644 --- a/test/registered/kernels/ops/layernorm/test_mhc_kernels.py +++ b/test/registered/kernels/ops/layernorm/test_mhc_kernels.py @@ -174,7 +174,7 @@ def _check_glm_boundary(x, residual, post, comb, fn, scale, base, *, use_norm): ] # Literal, not derived from the cutoff constant: deriving it makes this a # mirror that stays green when the cutoff moves. None is the empty batch, - # which attn_to_mlp short-circuits before reaching the callback. + # which update_and_read_ffn_input short-circuits before reaching the callback. fused_expected = {1: True, 6: True, 17: False}[x.shape[0]] if x.shape[0] else None with envs.SGLANG_OPT_FUSE_MHC_POST_PRE.override(True): if x.shape[0] > 0: @@ -193,15 +193,17 @@ def _check_glm_boundary(x, residual, post, comb, fn, scale, base, *, use_norm): f"num_tokens={x.shape[0]} fused={not declined}, " f"expected fused={fused_expected}" ) - outputs = [s.attn_to_mlp(x, residual.flatten(1), norm) for s in states] + outputs = [ + s.update_and_read_ffn_input(x, residual.flatten(1), norm) for s in states + ] torch.testing.assert_close(outputs[0][0], outputs[1][0], atol=2e-2, rtol=2e-2) torch.testing.assert_close(outputs[0][1], outputs[1][1], atol=0, rtol=0) torch.testing.assert_close(states[0].h_res, states[1].h_res, atol=1e-3, rtol=1e-3) torch.testing.assert_close(states[0].h_post, states[1].h_post, atol=1e-3, rtol=1e-3) # The next combine must consume the FFN mixing matrices, not attention's. torch.testing.assert_close( - states[0].mlp_combine(x, outputs[0][1]), - states[1].mlp_combine(x, outputs[1][1]), + states[0].apply_post(x, outputs[0][1]), + states[1].apply_post(x, outputs[1][1]), atol=2e-3, rtol=2e-2, ) diff --git a/test/registered/ops/test_aiter_allreduce_fusion_amd.py b/test/registered/ops/test_aiter_allreduce_fusion_amd.py index 727255d2069d..ec6dd7d1a7f8 100755 --- a/test/registered/ops/test_aiter_allreduce_fusion_amd.py +++ b/test/registered/ops/test_aiter_allreduce_fusion_amd.py @@ -11,7 +11,7 @@ import torch from sglang.srt.layers.layer_boundary import UnreducedOutput -from sglang.srt.layers.layer_boundary.fusions.allreduce import complete_attention_input +from sglang.srt.layers.layer_boundary.fusions.allreduce import fused_attn_input from sglang.srt.layers.moe import utils as moe_utils from sglang.srt.model_executor.forward_batch_info import ForwardMode from sglang.srt.runtime_context import get_context, get_flags, get_parallel @@ -386,7 +386,7 @@ def _evaluate_gate( hidden = torch.ones(8, 32) owed = UnreducedOutput(hidden, group=group if matching_group else object()) with ( - patch_communicator("apply_flashinfer_allreduce_fusion", return_value=False), + patch_communicator("flashinfer_ar_fusion_applies", return_value=False), patch_communicator("post_experts_reduction_group", return_value=group), patch_communicator("_use_aiter", use_aiter), get_parallel().override(tp_size=tp_world_size), @@ -397,7 +397,7 @@ def _evaluate_gate( ), get_flags().dp.override(enabled=dp_attention), ): - result = complete_attention_input( + result = fused_attn_input( plan, owed, hidden.clone(), diff --git a/test/registered/unit/layer_boundary/test_aux_capture_deferred_allreduce.py b/test/registered/unit/layer_boundary/test_aux_capture_deferred_allreduce.py index 847e78022aae..21517125012e 100644 --- a/test/registered/unit/layer_boundary/test_aux_capture_deferred_allreduce.py +++ b/test/registered/unit/layer_boundary/test_aux_capture_deferred_allreduce.py @@ -11,7 +11,7 @@ from sglang.srt.layers.layer_boundary import StageKind from sglang.srt.layers.layer_boundary import prepare as comm_ops from sglang.srt.layers.layer_boundary.contracts import BatchVariant -from sglang.srt.layers.layer_boundary.ops import identity_output +from sglang.srt.layers.layer_boundary.ops import keep_output from sglang.srt.layers.layer_boundary.residual.stream import ResidualStream from sglang.srt.model_executor.forward_batch_info import ForwardMode, PPProxyTensors from sglang.srt.models.bailing_moe import BailingMoEModel @@ -68,43 +68,42 @@ def __init__(self, defer, return_topk=False, stage_api=False): self.layer_communicator = stub_plan() self.layer_communicator.norm = None self.attn_boundary = stub_stage(self.layer_communicator, StageKind.ATTENTION) - self.layer_communicator._paths[BatchVariant.ORDINARY] = comm.StageSteps( - entry=comm.StageEntry( + self.layer_communicator.paths[BatchVariant.ORDINARY] = comm.StagePath( + entry=comm.EntryPath( prepare=partial( - comm_ops._consumer_step, + comm_ops._run_entry, step=partial( - comm_ops._read_input, - layer_input=None, + comm_ops._update_read, + pre_move=None, enters_stack=False, - read=comm.NORM_QUANT_READ, - update=comm.ADD, + read=comm.NORM_QUANT_READOUT, + update=comm.PLAIN_ADD, ), carried_fusions=(), ), input_rows=comm.Layout(frozenset()), input_move=identity_input, - handoff=comm_ops._hand_qkv_hook_its_input, + attn_input_adapter=comm_ops._attn_input_default, ), - output=comm.StageOutput( + output=comm.OutputContract( comm.Layout(frozenset()), group=comm.SumGroup.TP, - leaves_for_next_layer=True, + may_defer_to_next=True, ), - output_move=identity_output, - ) - self.layer_communicator.output._ffn_sum_moves_to_next_layer = ( - lambda batch, steps, **_: defer + output_move=keep_output, ) + self.layer_communicator.output._defers_sum = lambda batch, steps, **_: defer self.layer_communicator.terminal = False - self.layer_communicator._paths[BatchVariant.SEQUENCE_PARALLEL] = None - self.layer_communicator._paths[BatchVariant.INPUT_SCATTERED] = None - self.layer_communicator._paths[BatchVariant.CONTEXT_PARALLEL] = None + self.layer_communicator.paths[BatchVariant.SEQUENCE_PARALLEL] = None + self.layer_communicator.paths[BatchVariant.INPUT_SCATTERED] = None + self.layer_communicator.paths[BatchVariant.CONTEXT_PARALLEL] = None self.layer_communicator.output.ffn_reduction_group = lambda forward_batch: GROUP - self.layer_communicator.output._ffn_leaves_sum_to_reduce_scatter = ( + self.layer_communicator.output._skips_sum_for_reduce_scatter = ( lambda batch, dp_step: False ) - self.layer_communicator.output._complete_ffn_output_now = ( - lambda hidden, residual, **_: (all_reduce(hidden), residual) + self.layer_communicator.output._complete_now = lambda hidden, residual, **_: ( + all_reduce(hidden), + residual, ) def forward( @@ -121,7 +120,7 @@ def forward( residual = stream if isinstance(residual, ResidualStream): hidden_states, residual = residual.input(hidden_states) - hidden_states = comm.reduce_output(hidden_states) + hidden_states = comm.complete_owed(hidden_states) if residual is None: residual = hidden_states.clone() else: @@ -199,8 +198,8 @@ def test_skipping_an_empty_capture_does_not_complete_its_output(self): boundary.norm = None partial = torch.empty(0, 4) batch = SimpleNamespace(residual_stream=ResidualStream(partial)) - owed = batch.residual_stream.leave( - comm.UnreducedOutput(partial, group=GROUP), comm.ADD + owed = batch.residual_stream.record( + comm.UnreducedOutput(partial, group=GROUP), comm.PLAIN_ADD ) with patch.object(GROUP, "all_reduce") as reduce: hidden, captured = stub_stage(boundary, StageKind.ATTENTION).capture_output( @@ -393,8 +392,8 @@ def test_bailing_capture_does_not_include_a_previous_pipeline_ranks_layer(self): def test_written_streams_do_not_read_a_separate_residual(self): comm_instance = stub_plan() comm_instance.norm = None - comm_instance._batch_steps = lambda batch: SimpleNamespace( - entry=SimpleNamespace(input_sum=None) + comm_instance.path_for = lambda batch: SimpleNamespace( + entry=SimpleNamespace(declared_sum=None) ) batch = SimpleNamespace(residual_stream=None) streams = torch.randn(2, 4, 3) @@ -405,7 +404,7 @@ def test_written_streams_do_not_read_a_separate_residual(self): stage = stub_stage(comm_instance, StageKind.ATTENTION) stage.declaration = replace( stage.declaration, - previous=declare_ffn(update=SimpleNamespace(at_producer=True)), + previous=declare_ffn(update=SimpleNamespace(applied_at_exit=True)), ) hidden = stage.from_pp(PPProxyTensors({"hidden_states": streams}), batch) residual = batch.residual_stream @@ -416,8 +415,8 @@ def test_written_streams_do_not_read_a_separate_residual(self): def test_optional_residual_and_declared_partial_keep_the_wire_values(self): comm_instance = stub_plan() comm_instance.norm = None - comm_instance._batch_steps = lambda batch: SimpleNamespace( - entry=SimpleNamespace(input_sum=None) + comm_instance.path_for = lambda batch: SimpleNamespace( + entry=SimpleNamespace(declared_sum=None) ) batch = SimpleNamespace(residual_stream=None) partial = torch.randn(2, 4) diff --git a/test/registered/unit/layer_boundary/test_aux_hidden_states.py b/test/registered/unit/layer_boundary/test_aux_hidden_states.py index b3e62306cb22..8b836d893817 100644 --- a/test/registered/unit/layer_boundary/test_aux_hidden_states.py +++ b/test/registered/unit/layer_boundary/test_aux_hidden_states.py @@ -8,7 +8,7 @@ from sglang.srt.layers.aux_hidden_states import AuxHiddenStateList, AuxHiddenStatePacker from sglang.srt.layers.layer_boundary import StageKind -from sglang.srt.layers.layer_boundary.residual.access import norm_output +from sglang.srt.layers.layer_boundary.residual.access import final_norm_pair from sglang.srt.layers.layer_boundary.residual.stream import ResidualStream from sglang.test.boundary_fixtures import prepare_attention, stub_plan, stub_stage from sglang.test.ci.ci_register import register_cpu_ci @@ -35,22 +35,22 @@ def test_breakable_graph_copies_aux_collector_into_plain_list(self): backend._copy_output_to_buffer(source, tuple(buffers), 2) def test_capture_move_owns_gather_but_not_slice(self): - from sglang.srt.layers.layer_boundary.boundary import Boundary + from sglang.srt.layers.layer_boundary.boundary import EdgeBinding from sglang.srt.layers.layer_boundary.layout import Layout, TokenAxis from sglang.test.communicator_patch import patch_communicator full = Layout(frozenset()) - sharded = Layout(frozenset({TokenAxis.ATTN_TP_SCATTER})) + sharded = Layout(frozenset({TokenAxis.ATTN_TP})) source = torch.ones(2, 3) for rows, target, owns in ((sharded, full, True), (full, sharded, False)): edge = SimpleNamespace( residual_to=rows, produced=SimpleNamespace(layout=target) ) - boundary = Boundary(edge=edge, prepare=None) + boundary = EdgeBinding(edge=edge, prepare=None) self.assertEqual(boundary.capture_move_allocates, owns) with ( patch_communicator( - "_redistribute_from_attn_tp_shards", + "attn_tp_gather", side_effect=lambda x: torch.cat((x, x)), ), patch_communicator( @@ -111,9 +111,9 @@ def test_packer_receives_borrowed_boundary_value_without_intermediate_clone(self stage = self.boundary(stream) outputs = AuxHiddenStatePacker(1) kwargs = ( - {"capture_output": outputs.capture} + {"capture": outputs.capture} if callback - else {"captured_last_layer_outputs": outputs} + else {"capture_gathered": outputs} ) with patch.object( torch.Tensor, "clone", side_effect=AssertionError("extra clone") @@ -147,7 +147,7 @@ def test_embedding_and_final_norm_capture_write_directly_to_packer(self): with patch.object( torch.Tensor, "clone", side_effect=AssertionError("extra clone") ): - norm_output(hidden, residual, norm, outputs.capture) + final_norm_pair(hidden, residual, norm, outputs.capture) hidden.zero_() updated.zero_() torch.testing.assert_close( @@ -170,9 +170,7 @@ def prepare(value, stream, batch, **kwargs): with patch.object( torch.Tensor, "clone", side_effect=AssertionError("extra clone") ): - prepare_attention( - stage, hidden, stream, None, capture_output=outputs.capture - ) + prepare_attention(stage, hidden, stream, None, capture=outputs.capture) torch.testing.assert_close(outputs.finalize(), torch.full((2, 3), 4.0)) @@ -182,8 +180,8 @@ def exercise(self, *, enabled, backend_accepts=True, custom=False, callback=Fals from sglang.srt.layers import layernorm from sglang.srt.layers.layer_boundary import ( - ADD, - FusedMlpInput, + PLAIN_ADD, + FfnInputFusion, SumGroup, declare_attn, declare_ffn, @@ -212,10 +210,10 @@ def custom_fused(value, residual, batch): custom_fusions = ( SimpleNamespace( - ffn_input=lambda plan: ( - FusedMlpInput(completes=SumGroup.ATTN_TP, run=custom_fused), + ffn_input_fusions=lambda plan: ( + FfnInputFusion(completes=SumGroup.ATTN_TP, run=custom_fused), ), - attention_input=lambda plan: (), + attn_input_fusions=lambda plan: (), ) if custom else None @@ -224,15 +222,13 @@ def custom_fused(value, residual, batch): forward_mode=ForwardMode.DECODE, residual_stream=ResidualStream(torch.full((2, 4), 2.0)), ) - hidden = fb.residual_stream.leave(torch.full((2, 4), 3.0), ADD) + hidden = fb.residual_stream.record(torch.full((2, 4), 3.0), PLAIN_ADD) outputs = AuxHiddenStateList() with ( fixture.planning(parallel), patch_communicator("_use_aiter", False), - patch_communicator("apply_aiter_all_reduce_fusion", return_value=False), - patch_communicator( - "apply_flashinfer_allreduce_fusion", return_value=enabled - ), + patch_communicator("aiter_ar_fusion_applies", return_value=False), + patch_communicator("flashinfer_ar_fusion_applies", return_value=enabled), patch_communicator( "attention_tensor_model_parallel_all_reduce", side_effect=lambda x: x * 2, @@ -258,9 +254,9 @@ def custom_fused(value, residual, batch): predicate is not None and predicate(hidden, fb), enabled and not custom ) kwargs = ( - {"capture_output": outputs.capture} + {"capture": outputs.capture} if callback - else {"captured_last_layer_outputs": outputs} + else {"capture_gathered": outputs} ) if enabled and not custom: # Only capture is forbidden from cloning. The real FlashInfer diff --git a/test/registered/unit/layer_boundary/test_batch_owned_residual.py b/test/registered/unit/layer_boundary/test_batch_owned_residual.py index 47875bc33dd8..41ac27622e80 100644 --- a/test/registered/unit/layer_boundary/test_batch_owned_residual.py +++ b/test/registered/unit/layer_boundary/test_batch_owned_residual.py @@ -10,7 +10,7 @@ import torch from sglang.srt.batch_overlap.two_batch_overlap import TboForwardBatchPreparer -from sglang.srt.layers.layer_boundary import ADD +from sglang.srt.layers.layer_boundary import PLAIN_ADD from sglang.srt.layers.layer_boundary.output import UnreducedOutput from sglang.srt.layers.layer_boundary.residual import batch from sglang.srt.model_executor.forward_batch_info import ( @@ -28,17 +28,17 @@ class TestBatchOwnedResidual(CustomTestCase): def test_start_releases_previous_call_and_resets_state(self): fb = SimpleNamespace(residual_stream=None) with self.assertRaisesRegex(RuntimeError, "start"): - batch.current(fb) + batch.stream_of(fb) batch.start(fb) - old = batch.current(fb) - self.assertIs(old, batch.current(fb)) + old = batch.stream_of(fb) + self.assertIs(old, batch.stream_of(fb)) value = torch.ones(2, 4) ref = weakref.ref(value) old.write(value) del value, old batch.start(fb) - fresh = batch.current(fb) - self.assertIs(fresh, batch.current(fb)) + fresh = batch.stream_of(fb) + self.assertIs(fresh, batch.stream_of(fb)) self.assertIsNone(fresh.pending) self.assertIsNone(fresh.residual) self.assertIsNone(ref()) @@ -46,27 +46,29 @@ def test_start_releases_previous_call_and_resets_state(self): def test_terminal_completion_releases_batch_before_final_norm(self): fb = SimpleNamespace(residual_stream=None) batch.start(fb) - stream = batch.current(fb) + stream = batch.stream_of(fb) residual = torch.full((2, 4), 3.0) partial = torch.ones(2, 4) group = SimpleNamespace(all_reduce=Mock(side_effect=lambda x: x * 2)) stream.write(residual) - hidden = stream.leave(UnreducedOutput(partial, group=group), ADD) - result = batch.norm(hidden, fb, lambda value, prior: (value + prior, prior)) + hidden = stream.record(UnreducedOutput(partial, group=group), PLAIN_ADD) + result = batch.final_norm( + hidden, fb, lambda value, prior: (value + prior, prior) + ) self.assertIsNone(fb.residual_stream) group.all_reduce.assert_called_once() torch.testing.assert_close(result, torch.full((2, 4), 5.0)) with self.assertRaisesRegex(RuntimeError, "start"): - batch.current(fb) + batch.stream_of(fb) def test_pp_export_releases_the_batch_owner(self): fb = SimpleNamespace(residual_stream=None) batch.start(fb) hidden = torch.full((2, 4), 2.0) residual = torch.ones(2, 4) - stream = batch.current(fb) + stream = batch.stream_of(fb) stream.write(residual) - output = stream.leave(hidden, ADD) + output = stream.record(hidden, PLAIN_ADD) proxy = batch.to_pp(output, fb) self.assertIs(proxy["hidden_states"], hidden) self.assertIs(proxy["residual"], residual) @@ -108,10 +110,10 @@ def test_tbo_metadata_and_reused_batch_remain_stream_free_between_calls(self): ] self.assertTrue(all(child.residual_stream is None for child in children)) batch.start(parent) - owner = batch.current(parent) + owner = batch.stream_of(parent) self.assertTrue(all(child.residual_stream is None for child in children)) hidden = owner.write(torch.ones(size, 4)) - batch.norm(hidden, parent, lambda value: value) + batch.final_norm(hidden, parent, lambda value: value) cloned = replace(parent) self.assertIsNone(cloned.residual_stream) self.assertIsNone(parent.residual_stream) @@ -154,12 +156,12 @@ def _check_split_prefill(self, cls): class Layer: def __call__(self, positions, hidden, forward_batch, **kwargs): - stream = batch.current(forward_batch) + stream = batch.stream_of(forward_batch) owner_ids.append(id(stream)) - hidden, old = stream.finish(hidden) + hidden, old = stream.export(hidden) stream.write(hidden if old is None else hidden + old) - return stream.leave( - UnreducedOutput(torch.ones_like(hidden), group=group), ADD + return stream.record( + UnreducedOutput(torch.ones_like(hidden), group=group), PLAIN_ADD ) self_outer = self @@ -196,12 +198,12 @@ def test_mtp_returns_written_normalized_output_and_releases_batch(self): from sglang.srt.models.nemotron_h_mtp import NemotronHMultiTokenPredictor def terminal_layer(*, inputs_embeds, hidden_states, forward_batch): - stream = batch.current(forward_batch) + stream = batch.stream_of(forward_batch) stream.write(hidden_states) - hidden_states = stream.leave(torch.ones_like(hidden_states), ADD) + hidden_states = stream.record(torch.ones_like(hidden_states), PLAIN_ADD) hidden_states = batch.fold(hidden_states, forward_batch) normalized = hidden_states * 3 - return batch.written(normalized, forward_batch) + return batch.set_written(normalized, forward_batch) predictor = SimpleNamespace(pattern_len=1, layers={"0": terminal_layer}) fb = SimpleNamespace( @@ -217,13 +219,13 @@ def terminal_layer(*, inputs_embeds, hidden_states, forward_batch): def test_take_output_rejects_pending_or_mismatched_outputs(self): fb = SimpleNamespace(residual_stream=None) batch.start(fb) - stream = batch.current(fb) + stream = batch.stream_of(fb) value = torch.ones(2, 4) stream.write(value) - pending = stream.leave(value * 2, ADD) + pending = stream.record(value * 2, PLAIN_ADD) with self.assertRaises(RuntimeError): batch.take_output(pending, fb) - self.assertIs(batch.current(fb), stream) + self.assertIs(batch.stream_of(fb), stream) stream.write(value) with self.assertRaises(RuntimeError): batch.take_output(value.clone(), fb) diff --git a/test/registered/unit/layer_boundary/test_boundary_edges.py b/test/registered/unit/layer_boundary/test_boundary_edges.py index dc07f20aa63b..3651033cbd36 100644 --- a/test/registered/unit/layer_boundary/test_boundary_edges.py +++ b/test/registered/unit/layer_boundary/test_boundary_edges.py @@ -8,19 +8,19 @@ from sglang.srt.layers import layer_boundary as comm from sglang.srt.layers.layer_boundary import ( - ADD, - EdgeDecl, - FusedMlpInput, + PLAIN_ADD, + EdgeContract, + FfnInputFusion, + InputContract, Layout, - StageInput, - StageOutput, + OutputContract, SumGroup, TokenAxis, - make_boundary, - make_output_boundary, + bind_entry, + bind_exit, ) from sglang.srt.layers.layer_boundary import prepare as comm_ops -from sglang.srt.layers.layer_boundary.ops import identity_output +from sglang.srt.layers.layer_boundary.ops import keep_output from sglang.srt.runtime_context import get_parallel from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.communicator_patch import patch_communicator @@ -33,7 +33,7 @@ def sizes(*, dp=1, cp=1, tp=1): return { TokenAxis.ATTN_DP: dp, TokenAxis.ATTN_CP: cp, - TokenAxis.ATTN_TP_SCATTER: tp, + TokenAxis.ATTN_TP: tp, } @@ -56,20 +56,20 @@ def test_what_the_producer_may_leave_does_not_change_the_consumer_half(self): (False, True), (False, True), ): - produced = StageOutput( + produced = OutputContract( attention, group=group, - leaves_for_next_layer=next_layer, - leaves_for_reduce_scatter=reduce_scatter, - leaves_for_reduce_scatterv=reduce_scatterv, + may_defer_to_next=next_layer, + may_reduce_scatter=reduce_scatter, + may_reduce_scatterv=reduce_scatterv, ) - edge = EdgeDecl( + edge = EdgeContract( produced=produced, - need=StageInput(attention), + need=InputContract(attention), residual=attention, residual_to=attention, ) - boundary = make_boundary(edge) + boundary = bind_entry(edge) step = boundary.prepare.keywords["step"] halves.add( ( @@ -81,9 +81,9 @@ def test_what_the_producer_may_leave_does_not_change_the_consumer_half(self): ) self.assertEqual(len(halves), 1) ((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(func, comm_ops._run_entry) + self.assertIs(step, comm_ops._update_read) + self.assertIsNone(dict(keywords)["pre_move"]) self.assertIsNone(move) @@ -94,79 +94,81 @@ class TestNonAlternatingEdges(CustomTestCase): 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 + produced = OutputContract( + attention, group=SumGroup.ATTN_TP, may_defer_to_next=True ) - edge = EdgeDecl( + edge = EdgeContract( produced=produced, - need=StageInput(attention), + need=InputContract(attention), residual=attention, residual_to=attention, ) - producer = make_output_boundary(edge) - self.assertIs(producer.output_move, identity_output) - consumer = make_boundary(edge) - self.assertIs(consumer.prepare.func, comm_ops._consumer_step) - self.assertIsNone(consumer.prepare.keywords["step"].keywords["layer_input"]) + producer = bind_exit(edge) + self.assertIs(producer.output_move, keep_output) + consumer = bind_entry(edge) + self.assertIs(consumer.prepare.func, comm_ops._run_entry) + self.assertIsNone(consumer.prepare.keywords["step"].keywords["pre_move"]) self.assertIsNone(consumer.input_move) 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), + complete = EdgeContract( + produced=OutputContract(attention), + need=InputContract(full), residual=attention, residual_to=attention, ) - step = make_boundary(complete).prepare.keywords["step"] - self.assertIs(step.func, comm_ops._mlp_input_dp_replicate) + step = bind_entry(complete).prepare.keywords["step"] + self.assertIs(step.func, comm_ops._reduce_update_read_dp_gather) 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 + leaves = EdgeContract( + produced=OutputContract( + attention, group=SumGroup.TP, may_defer_to_next=True ), - need=StageInput(full), + need=InputContract(full), residual=attention, residual_to=attention, ) - 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) + prepare = bind_entry(leaves).prepare + self.assertIs(prepare.func, comm_ops._run_entry) + self.assertIs( + prepare.keywords["step"].func, comm_ops._reduce_update_read_dp_gather + ) # 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), + always = EdgeContract( + produced=OutputContract(attention, group=SumGroup.TP, always_partial=True), + need=InputContract(full), residual=attention, residual_to=attention, ) with self.assertRaises(NotImplementedError): - make_boundary(always) + bind_entry(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), + edge = EdgeContract( + produced=OutputContract(full, group=SumGroup.TP, always_partial=True), + need=InputContract(full), residual=full, residual_to=full, ) - over_tp = comm.FusedMlpInput(completes=SumGroup.TP, run=MagicMock()) - over_attention_tp = comm.FusedMlpInput( + over_tp = comm.FfnInputFusion(completes=SumGroup.TP, run=MagicMock()) + over_attention_tp = comm.FfnInputFusion( completes=SumGroup.ATTN_TP, run=MagicMock() ) - boundary = make_boundary(edge, fusions=(over_tp, over_attention_tp)) + boundary = bind_entry(edge, fusions=(over_tp, over_attention_tp)) self.assertEqual( boundary.prepare.keywords["step"].keywords.get("fusions", ()), (over_tp.run,), ) - step = make_boundary(edge).prepare - self.assertIs(step.keywords["step"].func, comm_ops._mlp_input_without_dp) + step = bind_entry(edge).prepare + self.assertIs(step.keywords["step"].func, comm_ops._reduce_update_read) all_reduce = MagicMock(side_effect=lambda h: h * 2) with patch_communicator("tensor_model_parallel_all_reduce", all_reduce): hidden, residual = step( @@ -183,16 +185,16 @@ 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 + is_plain_add = False + applied_at_exit = False def update(self, hidden_states, residual): return 2 * hidden_states + residual - def residual_to_attn_tp_shard(self, residual): + def slice_residual_attn_tp(self, residual): return residual - def residual_from_attn_tp_shards(self, residual): + def gather_residual_attn_tp(self, residual): return residual @@ -204,54 +206,54 @@ class TestTheProducersUpdateChoosesTheOrder(CustomTestCase): 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 + return EdgeContract( + produced=OutputContract( + attention, group=SumGroup.ATTN_TP, always_partial=True, update=update ), - need=StageInput(rows(axis_sizes)), + need=InputContract(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) + plain = bind_entry(self.edge(comm.PLAIN_ADD)).prepare.keywords["step"] + self.assertIs(plain.func, comm_ops._dp_gather_sum_read) written = _WrittenIn() - other = make_boundary(self.edge(written)).prepare.keywords["step"] - self.assertIs(other.func, comm_ops._mlp_input_dp_replicate) + other = bind_entry(self.edge(written)).prepare.keywords["step"] + self.assertIs(other.func, comm_ops._reduce_update_read_dp_gather) self.assertTrue(other.keywords["reduces_attention_tp"]) self.assertNotIn("update", other.keywords) def test_capabilities_prebind_both_orders_without_capturing_an_update(self): edge = msgspec.structs.replace( - self.edge(None), update_capabilities=(True, False) + self.edge(None), arriving_plain_add=(True, False) ) - boundary = make_boundary(edge) + boundary = bind_entry(edge) paths = boundary.prepare.keywords["paths"] self.assertEqual(set(paths), {True, False}) - self.assertIs(paths[True].keywords["step"].func, comm_ops._mlp_input_dp_partial) + self.assertIs(paths[True].keywords["step"].func, comm_ops._dp_gather_sum_read) self.assertIs( - paths[False].keywords["step"].func, comm_ops._mlp_input_dp_replicate + paths[False].keywords["step"].func, comm_ops._reduce_update_read_dp_gather ) for path in paths.values(): self.assertNotIn("update", path.keywords["step"].keywords) def test_actual_update_selects_a_prebound_path(self): full = rows(sizes(tp=1)) - edge = EdgeDecl( - StageOutput(full, update=None), - StageInput(full), + edge = EdgeContract( + OutputContract(full, update=None), + InputContract(full), full, full, - update_capabilities=(True, False), + arriving_plain_add=(True, False), ) - boundary = make_boundary(edge) + boundary = bind_entry(edge) def norm(value, residual=None): return value if residual is None else (value + residual, value + residual) - for update, expected in ((comm.ADD, 4.0), (_WrittenIn(), 5.0)): - with self.subTest(plain=update.adds_plainly): + for update, expected in ((comm.PLAIN_ADD, 4.0), (_WrittenIn(), 5.0)): + with self.subTest(plain=update.is_plain_add): hidden, residual = boundary.prepare( torch.ones(2, 4), torch.full((2, 4), 3.0), @@ -264,19 +266,19 @@ def norm(value, residual=None): def test_fused_kernels_take_only_a_plain_add(self): full = rows(sizes(tp=2)) - fused = comm.FusedMlpInput(completes=SumGroup.ATTN_TP, run=MagicMock()) + fused = comm.FfnInputFusion(completes=SumGroup.ATTN_TP, run=MagicMock()) 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 + for update, offered in ((comm.PLAIN_ADD, True), (_WrittenIn(), False)): + with self.subTest(is_plain_add=offered): + edge = EdgeContract( + produced=OutputContract( + full, group=SumGroup.ATTN_TP, always_partial=True, update=update ), - need=StageInput(full), + need=InputContract(full), residual=full, residual_to=full, ) - boundary = make_boundary( + boundary = bind_entry( edge, fusions=(fused,), carried_fusions=(carried,) ) self.assertEqual( @@ -290,13 +292,13 @@ def test_fused_kernels_take_only_a_plain_add(self): 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), + edge = EdgeContract( + produced=OutputContract(full, group=SumGroup.TP, may_defer_to_next=True), + need=InputContract(full), residual=full, residual_to=full, ) - step = make_boundary(edge).prepare + step = bind_entry(edge).prepare group = SimpleNamespace(all_reduce=MagicMock(side_effect=lambda h: h * 2)) def norm(hidden, residual): @@ -318,36 +320,39 @@ def norm(hidden, residual): def test_cross_layer_update_requires_a_lifetime_guarantee(self): local = Layout(frozenset()) - stateful_update = SimpleNamespace(adds_plainly=False, at_producer=False) - edge = EdgeDecl( - StageOutput(local, update=stateful_update), StageInput(local), local, local + stateful_update = SimpleNamespace(is_plain_add=False, applied_at_exit=False) + edge = EdgeContract( + OutputContract(local, update=stateful_update), + InputContract(local), + local, + local, ) with self.assertRaisesRegex(NotImplementedError, "lifetime"): - make_output_boundary(edge) + bind_exit(edge) def test_pipeline_cannot_reconstruct_a_non_add_update(self): local = Layout(frozenset()) update = SimpleNamespace( - adds_plainly=False, at_producer=False, can_defer_across_layers=True + is_plain_add=False, applied_at_exit=False, outlives_layer=True ) - edge = EdgeDecl( - StageOutput(local, update=update), StageInput(local), local, local + edge = EdgeContract( + OutputContract(local, update=update), InputContract(local), local, local ) with get_parallel().override(pp_size=2): with self.assertRaisesRegex(NotImplementedError, "pipeline boundaries"): - make_output_boundary(edge) + bind_exit(edge) class ProbeRead: - before_gather = False + reads_before_dp_gather = False """A read that marks what it reads: the input is the residual plus 100. - ``norms_plainly`` says whether it may stand in for a norm.""" + ``is_plain_norm`` says whether it may stand in for a norm.""" - def __init__(self, norms_plainly): - self.norms_plainly = norms_plainly + def __init__(self, is_plain_norm): + self.is_plain_norm = is_plain_norm self.reads = 0 - def enter(self, hidden_states): + def init_residual(self, hidden_states): return hidden_states def read(self, residual, norm, quant_format="", post_residual_addition=None): @@ -376,17 +381,20 @@ def _ffn_input(self, read, *, dp=1, residual_joins_sum=False, fusions=()): axis_sizes = sizes(dp=dp, tp=2) attention = rows(axis_sizes, TokenAxis.ATTN_DP) residual = ( - rows(axis_sizes, TokenAxis.ATTN_DP, TokenAxis.ATTN_TP_SCATTER) + rows(axis_sizes, TokenAxis.ATTN_DP, TokenAxis.ATTN_TP) if residual_joins_sum else attention ) need = rows(axis_sizes) if dp > 1 else attention - boundary = make_boundary( - EdgeDecl( - produced=StageOutput( - attention, group=SumGroup.ATTN_TP, always_leaves=True, update=ADD + boundary = bind_entry( + EdgeContract( + produced=OutputContract( + attention, + group=SumGroup.ATTN_TP, + always_partial=True, + update=PLAIN_ADD, ), - need=StageInput(need, read=read), + need=InputContract(need, read=read), residual=residual, residual_to=attention, residual_joins_sum=residual_joins_sum, @@ -398,15 +406,15 @@ def _ffn_input(self, read, *, dp=1, residual_joins_sum=False, fusions=()): ].keywords.get("fusions", ()) def test_the_dp_partial_reads_the_gathered_sum_with_the_declared_read(self): - read = ProbeRead(norms_plainly=True) + read = ProbeRead(is_plain_norm=True) step, _ = self._ffn_input(read, dp=2) - self.assertIs(step.func, comm_ops._mlp_input_dp_partial) + self.assertIs(step.func, comm_ops._dp_gather_sum_read) hidden, residual = torch.ones(2, 4), torch.full((2, 4), 3.0) with ( get_parallel().override(attn_tp_rank=0), patch.object( comm_ops, - "_reduce_and_redistribute_output_to_dp", + "dp_gather_sum", lambda h, forward_batch, cp_shard_counts: h * 2, ), patch.object(comm_ops, "dp_scatter", lambda *args: None), @@ -417,13 +425,13 @@ def test_the_dp_partial_reads_the_gathered_sum_with_the_declared_read(self): self.assertIs(out_residual, residual) def test_a_read_that_changes_the_residual_replicates_instead(self): - step, _ = self._ffn_input(ProbeRead(norms_plainly=False), dp=2) - self.assertIs(step.func, comm_ops._mlp_input_dp_replicate) + step, _ = self._ffn_input(ProbeRead(is_plain_norm=False), dp=2) + self.assertIs(step.func, comm_ops._reduce_update_read_dp_gather) def test_the_residual_joined_sum_is_read_with_the_declared_read(self): - read = ProbeRead(norms_plainly=True) + read = ProbeRead(is_plain_norm=True) step, _ = self._ffn_input(read, residual_joins_sum=True) - self.assertIs(step.func, comm_ops._mlp_input_residual_into_sum) + self.assertIs(step.func, comm_ops._tp_sum_with_residual_read) hidden, residual = torch.ones(4, 4), torch.full((2, 4), 3.0) with ( get_parallel().override(tp_size=2, tp_rank=0), @@ -441,14 +449,14 @@ def test_the_residual_joined_sum_is_read_with_the_declared_read(self): torch.testing.assert_close(out, expected + 100) def test_a_fused_kernel_runs_only_for_a_read_that_leaves_the_residual(self): - fused = FusedMlpInput(completes=SumGroup.ATTN_TP, run=MagicMock()) - for norms_plainly in (True, False): - with self.subTest(norms_plainly=norms_plainly): + fused = FfnInputFusion(completes=SumGroup.ATTN_TP, run=MagicMock()) + for is_plain_norm in (True, False): + with self.subTest(is_plain_norm=is_plain_norm): step, tried = self._ffn_input( - ProbeRead(norms_plainly), fusions=(fused,) + ProbeRead(is_plain_norm), fusions=(fused,) ) - self.assertIs(step.func, comm_ops._mlp_input_without_dp) - expected = (fused.run,) if norms_plainly else () + self.assertIs(step.func, comm_ops._reduce_update_read) + expected = (fused.run,) if is_plain_norm else () self.assertEqual(tried, expected) self.assertEqual(step.keywords["fusions"], expected) @@ -456,22 +464,22 @@ def test_the_attention_input_tries_its_kernels_only_for_such_a_read(self): axis_sizes = sizes(tp=2) attention = rows(axis_sizes) kernel = MagicMock() - for norms_plainly in (True, False): - with self.subTest(norms_plainly=norms_plainly): - read = ProbeRead(norms_plainly) - edge = EdgeDecl( - produced=StageOutput( - attention, group=SumGroup.TP, leaves_for_next_layer=True + for is_plain_norm in (True, False): + with self.subTest(is_plain_norm=is_plain_norm): + read = ProbeRead(is_plain_norm) + edge = EdgeContract( + produced=OutputContract( + attention, group=SumGroup.TP, may_defer_to_next=True ), - need=StageInput(attention, read=read), + need=InputContract(attention, read=read), residual=attention, residual_to=attention, ) - boundary = make_boundary(edge, carried_fusions=(kernel,)) + boundary = bind_entry(edge, carried_fusions=(kernel,)) self.assertIs(boundary.prepare.keywords["step"].keywords["read"], read) self.assertEqual( boundary.prepare.keywords["carried_fusions"], - (kernel,) if norms_plainly else (), + (kernel,) if is_plain_norm else (), ) diff --git a/test/registered/unit/layer_boundary/test_boundary_output_contracts.py b/test/registered/unit/layer_boundary/test_boundary_output_contracts.py index 58958c890771..243bf8512f8b 100644 --- a/test/registered/unit/layer_boundary/test_boundary_output_contracts.py +++ b/test/registered/unit/layer_boundary/test_boundary_output_contracts.py @@ -9,7 +9,7 @@ import torch from sglang.srt.layers.layer_boundary import ( - ADD, + PLAIN_ADD, ProducerReduction, declare_attn, declare_ffn, @@ -21,7 +21,7 @@ from sglang.srt.layers.layer_boundary.contracts import BatchVariant from sglang.srt.layers.layer_boundary.fusions.cutedsl import CuteDSLFusion from sglang.srt.layers.layer_boundary.layout import SumGroup -from sglang.srt.layers.layer_boundary.prepare import _dispatch_consumer +from sglang.srt.layers.layer_boundary.prepare import _dispatch_by_update from sglang.srt.layers.layer_boundary.residual import batch from sglang.srt.layers.layer_boundary.residual.stream import ResidualStream from sglang.srt.model_executor.forward_batch_info import ForwardMode @@ -35,12 +35,12 @@ class TestBoundaryIntegrations(unittest.TestCase): def test_dp_mixer_keeps_attention_group_reduction(self): parallel = fixture.parallel_of(attn_dp=2, attn_tp=2) rows = fixture.comm.Layout(frozenset()) - produced = fixture.comm.StageOutput( - rows, group=SumGroup.ATTN_TP, leaves_for_next_layer=True + produced = fixture.comm.OutputContract( + rows, group=SumGroup.ATTN_TP, may_defer_to_next=True ) - plan = SimpleNamespace(_batch_steps=lambda _: SimpleNamespace(output=produced)) - boundary = exits.OutputBoundary(plan) - boundary._ffn_sum_can_move_to_next_layer = lambda _: True + plan = SimpleNamespace(path_for=lambda _: SimpleNamespace(output=produced)) + boundary = exits.ExitPolicy(plan) + boundary._sum_deferral_allowed = lambda _: True with ( fixture.planning(parallel), patch.object(exits, "is_dp_attention_enabled", return_value=True), @@ -62,12 +62,10 @@ def test_pipeline_preserves_declared_sum_for_one_receiver_completion(self): (declare_ffn(), fixture.Norm()), previous=declare_ffn(), ) - attn.plan._batch_steps = lambda _: attn.plan._paths[ - BatchVariant.INPUT_SCATTERED - ] + attn.plan.path_for = lambda _: attn.plan.paths[BatchVariant.INPUT_SCATTERED] fb = SimpleNamespace(residual_stream=ResidualStream(torch.full((4, 4), 3.0))) - hidden = fb.residual_stream.leave( - torch.ones(4, 4), ADD, declared_sum=SumGroup.TP + hidden = fb.residual_stream.record( + torch.ones(4, 4), PLAIN_ADD, declared_sum=SumGroup.TP ) wire = batch.to_pp(hidden, fb) self.assertIsNone(fb.residual_stream) @@ -89,9 +87,7 @@ def reduce_scatter(value, residual): (declare_ffn(), fixture.Norm()), previous=declare_ffn(), ) - attn.plan._batch_steps = lambda _: attn.plan._paths[ - BatchVariant.INPUT_SCATTERED - ] + attn.plan.path_for = lambda _: attn.plan.paths[BatchVariant.INPUT_SCATTERED] attn.plan.qkv_latent_func = None entry = attn.entry(fb) hidden, residual = fb.residual_stream.input(hidden) @@ -101,7 +97,7 @@ def reduce_scatter(value, residual): fb, fixture.Norm(), pending=fb.residual_stream.pending, - update=ADD, + update=PLAIN_ADD, ) self.assertEqual(reductions, ["RS"]) torch.testing.assert_close(output, torch.full((2, 4), 10.0)) @@ -116,11 +112,11 @@ def test_terminal_finalize_requires_explicit_final_consumer(self): for accepted in (False, True): fusion = CuteDSLFusion() fusion.install( - SimpleNamespace(), hands_off_finalize=True, terminal_finalize=accepted + SimpleNamespace(), defers_finalize=True, terminal_finalize=accepted ) with ( fixture.planning(parallel), - patch.object(fusion, "_should_use_finalize", return_value=True), + patch.object(fusion, "_finalize_eligible", return_value=True), patch( "sglang.srt.layers.layer_boundary.fusions.cutedsl.get_parallel", return_value=parallel, @@ -133,7 +129,7 @@ def test_terminal_finalize_requires_explicit_final_consumer(self): ) fb.residual_stream = ResidualStream(torch.ones(2, 4)) with patch.object( - ffn.plan.output, "_postprocess_dp_step", return_value=None + ffn.plan.output, "_dp_reduce_scatter_step", return_value=None ): decision = ffn.exit(fb) self.assertEqual(decision.defer_moe_finalize, accepted) @@ -142,7 +138,7 @@ def test_absent_residual_uses_non_partial_update_path(self): plain = Mock(side_effect=AssertionError("cannot add a missing residual")) generic = Mock(return_value=("input", None)) self.assertEqual( - _dispatch_consumer( + _dispatch_by_update( None, None, None, None, paths={True: plain, False: generic} ), ("input", None), @@ -191,21 +187,20 @@ def test_lora_or_shared_tp1_defers_only_when_fusion_is_eligible(self): exits, "post_experts_reduction_group", return_value=group ), patch.object( - exits, "apply_flashinfer_allreduce_fusion", return_value=enabled - ), - patch.object( - exits, "apply_aiter_all_reduce_fusion", return_value=False + exits, "flashinfer_ar_fusion_applies", return_value=enabled ), + patch.object(exits, "aiter_ar_fusion_applies", return_value=False), ): self.assertEqual( - exits._can_defer_ffn_reduction(fb), enabled and (lora or shared) + exits._batch_allows_deferred_sum(fb), + enabled and (lora or shared), ) def test_unsupported_producer_contracts_fail_at_declaration(self): with self.assertRaises(ValueError): - declare_ffn(reduction=ProducerReduction.PARTIAL) + declare_ffn(reduction=ProducerReduction.ALWAYS_PARTIAL) with self.assertRaises(ValueError): - declare_attn(reduction=ProducerReduction.LOCAL_TAIL) + declare_attn(reduction=ProducerReduction.TAIL_AFTER_SUM) def test_dense_decoder_forwards_capture_callback(self): from sglang.srt.models.llama4 import Llama4DecoderLayer @@ -214,9 +209,9 @@ def test_dense_decoder_forwards_capture_callback(self): for cls in (Qwen3DecoderLayer, Llama4DecoderLayer): captures = [] - def prepare(hidden, fb, capture_output=None, **kwargs): - if capture_output is not None: - capture_output(hidden) + def prepare(hidden, fb, capture=None, **kwargs): + if capture is not None: + capture(hidden) return hidden output = SimpleNamespace(finish=lambda hidden: hidden) diff --git a/test/registered/unit/layer_boundary/test_boundary_output_policy.py b/test/registered/unit/layer_boundary/test_boundary_output_policy.py index cbfd7d5dcada..d00944d79fd5 100644 --- a/test/registered/unit/layer_boundary/test_boundary_output_policy.py +++ b/test/registered/unit/layer_boundary/test_boundary_output_policy.py @@ -148,24 +148,24 @@ def move(step, batch, value): return value[:1] selected = ( - "_reduce_and_redistribute_output_varlen" + "dp_reduce_scatterv" if varlen and policy in ("rsv", "rs+rsv") - else "_reduce_and_redistribute_output_max_len" + else "dp_reduce_scatter" if max_len and can_rs and policy in ("rs", "rs+rsv") - else "_redistribute_output" + else "_dp_scatter_step" ) with ( fixture.planning(parallel), patch_communicator("should_use_dp_reduce_scatterv", lambda: varlen), patch_communicator("can_use_dp_reduce_scatter", lambda: can_rs), - patch_communicator("_to_local_tokens", move), + patch_communicator("to_dp_local", move), ): with layer.ffn.plan.output.ffn_exit( self.batch(max_len), stream=ResidualStream() ) as output: self.assertEqual( get_forward().mlp_reduce_scatter, - selected != "_redistribute_output", + selected != "_dp_scatter_step", ) result, _ = finish_exit(output, torch.ones(2, 4), torch.zeros(1, 4)) self.assertEqual(calls, [selected]) @@ -197,11 +197,11 @@ def move(step, batch, value): "should_use_dp_reduce_scatterv", lambda: varlen ), patch_communicator("can_use_dp_reduce_scatter", lambda: True), - patch_communicator("_to_local_tokens", move), + patch_communicator("to_dp_local", move), ): - output = layer.ffn.postprocess(hidden, batch) - value, saved_residual = batch.residual_stream.finish(output) - self.assertEqual(calls, ["_redistribute_output"]) + output = layer.ffn.finish_complete_output(hidden, batch) + value, saved_residual = batch.residual_stream.export(output) + self.assertEqual(calls, ["_dp_scatter_step"]) torch.testing.assert_close(value, hidden[:1]) self.assertIs(saved_residual, residual) @@ -252,7 +252,7 @@ def move(step, batch, value): fixture.planning(parallel), patch_communicator("should_use_dp_reduce_scatterv", lambda: False), patch_communicator("can_use_dp_reduce_scatter", lambda: True), - patch_communicator("_to_local_tokens", move), + patch_communicator("to_dp_local", move), ): with layer.ffn.plan.output.ffn_exit( self.batch(), stream=ResidualStream() @@ -275,10 +275,10 @@ def test_transform_without_pre_rs_implementation_uses_ar(self): parallel, output=OutputTransform(lambda x: x.square()), ) - output = layer.ffn.plan._paths.get(BatchVariant.ORDINARY).output - self.assertFalse(output.leaves_for_next_layer) - self.assertFalse(output.leaves_for_reduce_scatter) - self.assertFalse(output.leaves_for_reduce_scatterv) + output = layer.ffn.plan.paths.get(BatchVariant.ORDINARY).output + self.assertFalse(output.may_defer_to_next) + self.assertFalse(output.may_reduce_scatter) + self.assertFalse(output.may_reduce_scatterv) if __name__ == "__main__": diff --git a/test/registered/unit/layer_boundary/test_cutedsl_ar_fusion.py b/test/registered/unit/layer_boundary/test_cutedsl_ar_fusion.py index 6abfb5ae291f..c918c4cff80b 100644 --- a/test/registered/unit/layer_boundary/test_cutedsl_ar_fusion.py +++ b/test/registered/unit/layer_boundary/test_cutedsl_ar_fusion.py @@ -12,30 +12,30 @@ _with_early_finalize_shared_load, ) from sglang.srt.layers.layer_boundary import ( - ADD, - NORM_QUANT_READ, + NORM_QUANT_READOUT, + PLAIN_ADD, Layout, StageKind, SumGroup, TokenAxis, UnreducedOutput, - reduce_output, + complete_owed, ) from sglang.srt.layers.layer_boundary.construction import BatchVariant from sglang.srt.layers.layer_boundary.fusions.allreduce import ( - attention_fusions, - complete_attention_input, - complete_ffn_input, - ffn_fusions, + attn_input_fusions, + ffn_input_fusions, + fused_attn_input, + fused_ffn_input, ) from sglang.srt.layers.layer_boundary.fusions.cutedsl import ( CuteDSLFusion, - MoeFinalizeHandoff, + MoeDeferredFinalize, install_cutedsl_fusion, prepare_cutedsl_fusion, ) -from sglang.srt.layers.layer_boundary.prepare import _consumer_step, _read_input -from sglang.srt.layers.layer_boundary.residual.access import finish_layer_stack +from sglang.srt.layers.layer_boundary.prepare import _run_entry, _update_read +from sglang.srt.layers.layer_boundary.residual.access import export_output from sglang.srt.layers.layernorm import RMSNorm from sglang.srt.model_executor.cuda_graph_config import CudaGraphConfig, PhaseConfig from sglang.srt.model_executor.forward_batch_info import ForwardMode @@ -56,28 +56,28 @@ def _communicator(): comm.terminal = False comm.fusions = CuteDSLFusion() comm.norm = RMSNorm(8, eps=1e-6) - comm._attn_input_fusions = attention_fusions(comm) + comm._attn_input_fusions = attn_input_fusions(comm) # Only the ordinary batches' attention input half, with these entries. - comm._paths[BatchVariant.SEQUENCE_PARALLEL] = comm._paths[ + comm.paths[BatchVariant.SEQUENCE_PARALLEL] = comm.paths[ BatchVariant.INPUT_SCATTERED - ] = comm._paths[BatchVariant.CONTEXT_PARALLEL] = None - comm._paths[BatchVariant.ORDINARY] = SimpleNamespace( + ] = comm.paths[BatchVariant.CONTEXT_PARALLEL] = None + comm.paths[BatchVariant.ORDINARY] = SimpleNamespace( entry=SimpleNamespace( prepare=partial( - _consumer_step, - adds_plainly=True, + _run_entry, + is_plain_add=True, step=partial( - _read_input, - layer_input=None, + _update_read, + pre_move=None, enters_stack=False, - read=NORM_QUANT_READ, - update=ADD, + read=NORM_QUANT_READOUT, + update=PLAIN_ADD, ), carried_fusions=comm._attn_input_fusions, ), input_move=None, - input_sum=None, - handoff=None, + declared_sum=None, + attn_input_adapter=None, ) ) return comm @@ -115,7 +115,7 @@ def _install(layers, **kwargs): def eligible(): """Every shared gate open, so a case varies only the predicate it names.""" with ( - patch.object(CuteDSLFusion, "_common_eligible", return_value=True), + patch.object(CuteDSLFusion, "_eligible", return_value=True), patch( f"{_MODULE}.get_exec", return_value=SimpleNamespace( @@ -137,13 +137,13 @@ def test_last_layer_consumes_but_does_not_skip_the_pending_all_reduce(eligible): residual + 1, ) ), - hands_off_finalize=False, + defers_finalize=False, ) hidden_states = UnreducedOutput(torch.zeros(8, 8)) with ( patch_communicator( - "reduce_output", + "complete_owed", lambda *a, **k: pytest.fail("fell through to the unfused path"), ), ): @@ -163,40 +163,40 @@ def test_cutedsl_entries_come_before_the_base_fused_kernel(): fusion = comm.fusions finalize, reduce, base_input = comm._attn_input_fusions assert _bound(finalize) == ( - fusion._finalize_output_and_update_and_read_residual, + fusion._finalize_add_norm, (comm,), ) assert _bound(reduce) == ( - fusion._reduce_output_and_update_and_read_residual, + fusion._all_reduce_add_norm, (comm,), ) - assert _bound(base_input) == (complete_attention_input, (comm,)) - base = (complete_ffn_input, (comm,)) - cutedsl = (fusion._mlp_input_reduce_output_and_update_and_read_residual, (comm,)) + assert _bound(base_input) == (fused_attn_input, (comm,)) + base = (fused_ffn_input, (comm,)) + cutedsl = (fusion._ffn_input_all_reduce_add_norm, (comm,)) with patch( f"{_MODULE}.get_parallel", return_value=SimpleNamespace(attn_tp_size=2, tp_size=2), ): - comm.variants = { + comm.edges = { BatchVariant.ORDINARY: SimpleNamespace( incoming=SimpleNamespace(residual=Layout(frozenset())) ) } - fusions = ffn_fusions(comm) + fusions = ffn_input_fusions(comm) assert _bound(fusions[0].run) == cutedsl assert _bound(fusions[1].run) == base # The workspace reduces over the TP group, which is the attention-TP group here. assert [f.completes for f in fusions] == [SumGroup.ATTN_TP] * 2 # A residual on each rank's slice is gathered first, which the # workspace does not do. - comm.variants = { + comm.edges = { BatchVariant.ORDINARY: SimpleNamespace( incoming=SimpleNamespace( - residual=Layout(frozenset({TokenAxis.ATTN_TP_SCATTER})) + residual=Layout(frozenset({TokenAxis.ATTN_TP})) ) ) } - assert [_bound(f.run) for f in ffn_fusions(comm)] == [base] + assert [_bound(f.run) for f in ffn_input_fusions(comm)] == [base] def test_the_fusion_runs_only_on_the_ffn_full_rows(): @@ -223,20 +223,20 @@ def test_the_fusion_runs_only_on_the_ffn_full_rows(): ): for sharded, eligible in ( (frozenset(), True), - (frozenset({TokenAxis.ATTN_TP_SCATTER}), False), + (frozenset({TokenAxis.ATTN_TP}), False), ): - comm._batch_steps = lambda fb, rows=Layout(sharded): SimpleNamespace( + comm.path_for = lambda fb, rows=Layout(sharded): SimpleNamespace( entry=SimpleNamespace(input_rows=rows) ) - assert fusion._common_eligible(comm, _DECODE, 8) is eligible + assert fusion._eligible(comm, _DECODE, 8) is eligible def test_finalize_handoff_is_a_producer_capability(eligible): comm = _communicator() fusion = comm.fusions assert not fusion.can_defer_finalize(comm, _DECODE) - fusion.install(SimpleNamespace(), hands_off_finalize=True) - with patch.object(CuteDSLFusion, "_should_use_finalize", return_value=True): + fusion.install(SimpleNamespace(), defers_finalize=True) + with patch.object(CuteDSLFusion, "_finalize_eligible", return_value=True): assert fusion.can_defer_finalize(comm, _DECODE) @@ -263,8 +263,8 @@ def test_install_does_not_require_a_fused_successor(): can_defer_finalize=lambda layer: True, label="test", ) - assert first.fusions.hands_off_finalize - assert last.fusions.hands_off_finalize + assert first.fusions.defers_finalize + assert last.fusions.defers_finalize def test_a_service_nested_under_a_wrapper_is_prepared(): @@ -339,7 +339,7 @@ def test_the_handoff_views_the_producer_storage(): gemm2_out = torch.empty(m * top_k + 4, 16, dtype=torch.bfloat16) expert_weights = torch.empty(m + 1, top_k, dtype=torch.bfloat16) permuted_indices = torch.empty(m + 1, top_k, dtype=torch.int32) - handoff = MoeFinalizeHandoff.from_flashinfer( + handoff = MoeDeferredFinalize.from_flashinfer( SimpleNamespace( gemm2_out=gemm2_out, expert_weights=expert_weights, @@ -358,7 +358,7 @@ def test_the_handoff_views_the_producer_storage(): def _handoff(finish): rows = torch.zeros(2, 8) - return MoeFinalizeHandoff( + return MoeDeferredFinalize( routed_output=rows, expert_weights=rows, permuted_indices=rows, @@ -373,7 +373,7 @@ def test_its_producer_completes_a_handoff_for_any_other_reader(): the MoE's own unfused tail, run once.""" finished = torch.full((2, 8), 3.0) finish = MagicMock(return_value=finished) - assert reduce_output(_handoff(finish)) is finished + assert complete_owed(_handoff(finish)) is finished finish.assert_called_once_with() @@ -390,7 +390,7 @@ def test_from_flashinfer_completes_with_the_moe_s_own_sum(): "finalize_flashinfer_trtllm_deferred_output", side_effect=lambda d, s: s + 1, ) as finalize: - handoff = MoeFinalizeHandoff.from_flashinfer( + handoff = MoeDeferredFinalize.from_flashinfer( deferred, gated_shared_output=shared, m=2, reduce=lambda h: h * 10 ) finalize.assert_not_called() @@ -416,8 +416,8 @@ def forward(self, x, residual=None, post_residual_addition=None): with ( # A norm the kernel could fold in, on a batch it does not take. patch(f"{_MODULE}._fused_norm_gamma", return_value=torch.ones(8)), - patch.object(CuteDSLFusion, "_should_use_finalize", return_value=False), - patch.object(CuteDSLFusion, "_common_eligible", return_value=False), + patch.object(CuteDSLFusion, "_finalize_eligible", return_value=False), + patch.object(CuteDSLFusion, "_eligible", return_value=False), ): hidden, residual = prepare_input( stub_stage(comm, StageKind.ATTENTION), @@ -434,7 +434,7 @@ def test_the_layer_stack_hands_a_handoff_only_to_a_final_norm_that_takes_it(): for takes, completed in ((True, 0), (False, 1)): finish = MagicMock(return_value=torch.ones(2, 8)) handoff = _handoff(finish) - hidden, _ = finish_layer_stack( + hidden, _ = export_output( handoff, torch.zeros(2, 8), _DECODE, final_norm_takes_handoff=takes ) assert finish.call_count == completed @@ -596,16 +596,14 @@ def test_installed_cutedsl_provider_restores_deferred_sum(self): exits, "get_parallel", return_value=SimpleNamespace(tp_group=group) ), patch.object(exits, "post_experts_reduction_group", return_value=group), - patch.object( - exits, "apply_flashinfer_allreduce_fusion", return_value=False - ), + patch.object(exits, "flashinfer_ar_fusion_applies", return_value=False), patch.object( provider, "can_defer_all_reduce", return_value=eligible ) as gate, ): # Avoid platform-specific Aiter details: no residual cannot fuse. with patch.object(exits, "_use_aiter", False, create=True): - result = exits._can_defer_ffn_reduction(fb, boundary) + result = exits._batch_allows_deferred_sum(fb, boundary) self.assertEqual(result, eligible and not a2a) if a2a: gate.assert_not_called() diff --git a/test/registered/unit/layer_boundary/test_declared_attention_cp.py b/test/registered/unit/layer_boundary/test_declared_attention_cp.py index 1a6d945f276a..10c7e09911b8 100644 --- a/test/registered/unit/layer_boundary/test_declared_attention_cp.py +++ b/test/registered/unit/layer_boundary/test_declared_attention_cp.py @@ -177,7 +177,7 @@ def build(self, use_reduce_scatter): last=False, sparse=True, previous_sparse=False, - next_sparse=False, + next_layer_sparse=False, attention_norm=layernorm, ffn_norm=layernorm, ) diff --git a/test/registered/unit/layer_boundary/test_declared_decoder_boundary.py b/test/registered/unit/layer_boundary/test_declared_decoder_boundary.py index 384fb5277743..c244cc9296a3 100644 --- a/test/registered/unit/layer_boundary/test_declared_decoder_boundary.py +++ b/test/registered/unit/layer_boundary/test_declared_decoder_boundary.py @@ -33,20 +33,23 @@ from sglang.srt.layers.layer_boundary import prepare as comm_ops from sglang.srt.layers.layer_boundary.adapters import branch from sglang.srt.layers.layer_boundary.construction import BatchVariant -from sglang.srt.layers.layer_boundary.fusions.allreduce import complete_ffn_input +from sglang.srt.layers.layer_boundary.fusions.allreduce import fused_ffn_input from sglang.srt.layers.layer_boundary.ops import ( - gather_attention_tp, - identity_output, - output_on_residual_shard, - reduce_scatter_over_cp, - scatter_moe_cp_output, - scatter_output, - take_back_cp_shard, - update_and_gather, + attn_cp_reduce_scatter_output, + attn_tp_gather_input, + attn_tp_slice_output, + dp_cp_take_back_output, + keep_output, + moe_cp_take_back_output, + residual_slice_output, + update_attn_tp_gather_output, ) from sglang.srt.layers.layer_boundary.residual import mhc as mhc_module -from sglang.srt.layers.layer_boundary.residual.access import finish_layer_stack -from sglang.srt.layers.layer_boundary.residual.add_norm import NORM_READ, PLAIN_RESIDUAL +from sglang.srt.layers.layer_boundary.residual.access import export_output +from sglang.srt.layers.layer_boundary.residual.add_norm import ( + NORM_READOUT, + PLAIN_RESIDUAL_OPS, +) from sglang.srt.layers.layer_boundary.residual.stream import OwedOutput, ResidualStream from sglang.srt.layers.moe import utils as moe_utils from sglang.srt.model_executor.forward_batch_info import ForwardMode @@ -164,14 +167,19 @@ def forward_with_allreduce_fusion(self, *args, **kwargs): def layer_case( - layer_id, num_layers, *, sparse=False, previous_sparse=False, next_sparse=False + layer_id, + num_layers, + *, + sparse=False, + previous_sparse=False, + next_layer_sparse=False, ): return dict( sparse=sparse, first=layer_id == 0, last=layer_id == num_layers - 1, previous_sparse=previous_sparse, - next_sparse=next_sparse, + next_layer_sparse=next_layer_sparse, ) @@ -204,20 +212,20 @@ class TestStageLayoutSelection(CustomTestCase): def test_the_order_follows_the_attention_output(self): for attn_tp, force, order in ( - (2, False, comm_ops._mlp_input_dp_partial), - (2, True, comm_ops._mlp_input_dp_replicate), - (1, False, comm_ops._mlp_input_dp_replicate), + (2, False, comm_ops._dp_gather_sum_read), + (2, True, comm_ops._reduce_update_read_dp_gather), + (1, False, comm_ops._reduce_update_read_dp_gather), ): with self.subTest(attn_tp=attn_tp, force=force): communicator = build( layer_case(1, 3), parallel_of(attn_dp=2, attn_tp=attn_tp), - residual=PLAIN_RESIDUAL._replace( - ffn_read=replace(NORM_READ, before_gather=force) + residual=PLAIN_RESIDUAL_OPS._replace( + ffn_readout=replace(NORM_READOUT, reads_before_dp_gather=force) ), ) self.assertIs( - communicator.ffn.plan._paths.get(BatchVariant.ORDINARY) + communicator.ffn.plan.paths.get(BatchVariant.ORDINARY) .entry.prepare.keywords["step"] .func, order, @@ -231,8 +239,8 @@ def test_the_last_a2a_layer_folds_the_residual_back(self): a2a=True, ) self.assertIs( - last.ffn.plan._paths.get(BatchVariant.ORDINARY).output_move.func, - update_and_gather, + last.ffn.plan.paths.get(BatchVariant.ORDINARY).output_move.func, + update_attn_tp_gather_output, ) def test_an_attention_that_gathers_its_slice_itself(self): @@ -241,7 +249,7 @@ def test_an_attention_that_gathers_its_slice_itself(self): with patch_communicator("_use_ag_after_qlora", True): layer = build(facts, parallel, a2a=True) self.assertIsNone( - layer.attn.plan._paths.get(BatchVariant.ORDINARY).entry.input_move + layer.attn.plan.paths.get(BatchVariant.ORDINARY).entry.input_move ) def test_cp_without_prefill_cp(self): @@ -250,16 +258,14 @@ def test_cp_without_prefill_cp(self): so the FFN completes its own sum; without attention DP it may leave it and input-scattered attention runs, as without CP.""" under_dp = build(layer_case(1, 3), parallel_of(attn_dp=2, attn_tp=1, attn_cp=2)) - self.assertIsNone(under_dp.ffn.plan._paths.get(BatchVariant.CONTEXT_PARALLEL)) + self.assertIsNone(under_dp.ffn.plan.paths.get(BatchVariant.CONTEXT_PARALLEL)) self.assertFalse( - under_dp.ffn.plan._paths.get( - BatchVariant.ORDINARY - ).output.leaves_for_next_layer + under_dp.ffn.plan.paths.get(BatchVariant.ORDINARY).output.may_defer_to_next ) self.assertFalse( - under_dp.ffn.plan._paths.get( + under_dp.ffn.plan.paths.get( BatchVariant.ORDINARY - ).output.leaves_for_reduce_scatterv + ).output.may_reduce_scatterv ) without_dp = build( layer_case(1, 3), @@ -268,12 +274,12 @@ def test_cp_without_prefill_cp(self): ), ) self.assertTrue( - without_dp.ffn.plan._paths.get( + without_dp.ffn.plan.paths.get( BatchVariant.ORDINARY - ).output.leaves_for_next_layer + ).output.may_defer_to_next ) self.assertIsNotNone( - without_dp.ffn.plan._paths.get(BatchVariant.INPUT_SCATTERED) + without_dp.ffn.plan.paths.get(BatchVariant.INPUT_SCATTERED) ) @@ -301,7 +307,7 @@ def build_mhc( **facts, attention_norm=Norm(), ffn_norm=Norm(), - residual=mhc.layer_residual(), + residual=mhc.residual_ops(), ) stages.mhc = mhc return stages @@ -312,7 +318,7 @@ 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: + if step.func is comm_ops._run_entry: step = step.keywords["step"] self.assertIs(step.func, func) declaration = communicator.ffn.declaration @@ -330,17 +336,17 @@ def test_the_move_back_over_dp_stays_with_the_layer(self): parallel = parallel_of(attn_dp=2, attn_tp=2) communicator = build_mhc(layer_case(1, 3), parallel) self.assertTrue( - communicator.ffn.plan._paths.get(BatchVariant.ORDINARY).returns_over_dp + communicator.ffn.plan.paths.get(BatchVariant.ORDINARY).returns_over_dp ) self.assertFalse( - communicator.ffn.plan._paths.get( + communicator.ffn.plan.paths.get( BatchVariant.ORDINARY - ).output.leaves_for_next_layer + ).output.may_defer_to_next ) # The move back also writes the FFN output into the streams. self.assertFalse( - communicator.ffn.plan.output._local_token_move_can_go_to_next_layer( - communicator.ffn.plan.output.plan._batch_steps( + communicator.ffn.plan.output._next_input_can_scatter( + communicator.ffn.plan.output.plan.path_for( SimpleNamespace(forward_mode=ForwardMode.DECODE) ) ) @@ -358,12 +364,12 @@ def test_the_move_back_over_dp_stays_with_the_layer(self): patch_communicator("should_use_dp_reduce_scatterv", lambda: False), patch_communicator("can_use_dp_reduce_scatter", lambda: True), ): - step = communicator.ffn.plan.output._postprocess_dp_step( - max_len, communicator.ffn.plan.output.plan._batch_steps(max_len) + step = communicator.ffn.plan.output._dp_reduce_scatter_step( + max_len, communicator.ffn.plan.output.plan.path_for(max_len) ) self.assertTrue( - communicator.ffn.plan.output._ffn_leaves_sum_to_reduce_scatter( - communicator.ffn.plan.output.plan._batch_steps(max_len), step + communicator.ffn.plan.output._skips_sum_for_reduce_scatter( + communicator.ffn.plan.output.plan.path_for(max_len), step ) ) @@ -374,13 +380,13 @@ def test_the_exit_runs_the_move_back_over_dp_it_chose(self): "SUM_LEN", True, False, - transport_ops._reduce_and_redistribute_output_varlen, + transport_ops.dp_reduce_scatterv, ), ( "MAX_LEN", False, True, - transport_ops._reduce_and_redistribute_output_max_len, + transport_ops.dp_reduce_scatter, ), ): with self.subTest(name): @@ -395,7 +401,7 @@ def test_the_exit_runs_the_move_back_over_dp_it_chose(self): dp_padding_mode=SimpleNamespace(is_max_len=lambda: is_max_len), forward_mode=ForwardMode.DECODE, ) - choose = MagicMock(wraps=comm_exit._reduce_and_redistribute_output_step) + choose = MagicMock(wraps=comm_exit._select_dp_reduce_scatter) to_local_tokens = MagicMock(side_effect=lambda step, fb, h: h[:2]) with ( planning(parallel), @@ -403,8 +409,8 @@ def test_the_exit_runs_the_move_back_over_dp_it_chose(self): "should_use_dp_reduce_scatterv", lambda: reduce_scatterv ), patch_communicator("can_use_dp_reduce_scatter", lambda: True), - patch_communicator("_reduce_and_redistribute_output_step", choose), - patch_communicator("_to_local_tokens", to_local_tokens), + patch_communicator("_select_dp_reduce_scatter", choose), + patch_communicator("to_dp_local", to_local_tokens), ): with communicator.ffn.plan.output.ffn_exit( forward_batch, stream=ResidualStream() @@ -434,27 +440,27 @@ def test_a2a_layers_take_their_slice_with_the_coefficients(self): ] first_a2a, last = layers[1], layers[2] self.assert_step( - first_a2a.ffn.plan._paths.get(BatchVariant.ORDINARY).entry.prepare, - comm_ops._mlp_input_scatter, + first_a2a.ffn.plan.paths.get(BatchVariant.ORDINARY).entry.prepare, + comm_ops._attn_tp_reduce_scatter_update_read, first_a2a, scatters_residual=True, ) self.assert_step( - last.ffn.plan._paths.get(BatchVariant.ORDINARY).entry.prepare, - comm_ops._mlp_input_scatter, + last.ffn.plan.paths.get(BatchVariant.ORDINARY).entry.prepare, + comm_ops._attn_tp_reduce_scatter_update_read, last, scatters_residual=False, ) self.assert_step( - last.ffn.plan._paths.get(BatchVariant.ORDINARY).output_move, - update_and_gather, + last.ffn.plan.paths.get(BatchVariant.ORDINARY).output_move, + update_attn_tp_gather_output, last, ) state = first_a2a.mhc state.h_res, state.h_post = torch.arange(4.0), torch.arange(4.0) + 10 context = SimpleNamespace(attn_tp_rank=1, attn_tp_size=2) with patch_communicator("get_parallel", lambda: context): - residual = state.residual_to_attn_tp_shard(torch.arange(4.0) + 20) + residual = state.slice_residual_attn_tp(torch.arange(4.0) + 20) self.assertEqual(residual.tolist(), [22.0, 23.0]) self.assertEqual(state.h_res.tolist(), [2.0, 3.0]) self.assertEqual(state.h_post.tolist(), [12.0, 13.0]) @@ -466,15 +472,15 @@ def test_an_input_scattered_batch_keeps_the_residual_on_the_slice(self): for layer_id in range(3): with self.subTest(layer_id=layer_id): communicator = build_mhc(layer_case(layer_id, 3), parallel) - steps = communicator.ffn.plan._paths.get(BatchVariant.INPUT_SCATTERED) + steps = communicator.ffn.plan.paths.get(BatchVariant.INPUT_SCATTERED) self.assert_step( steps.entry.prepare, - comm_ops._mlp_input_on_residual_shard, + comm_ops._tp_reduce_scatter_update_read_gather, communicator, ) self.assert_step( steps.output_move, - output_on_residual_shard, + residual_slice_output, communicator, sums=True, gathers_back=layer_id == 2, @@ -484,31 +490,31 @@ def test_an_input_scattered_batch_keeps_the_residual_on_the_slice(self): layer_case(layer_id, 3), parallel, boundary_reduction="ar", - ).ffn.plan._paths.get(BatchVariant.INPUT_SCATTERED) + ).ffn.plan.paths.get(BatchVariant.INPUT_SCATTERED) self.assertFalse(completes.output_move.keywords["sums"]) self.assertFalse(completes.output_move_completes_sum) # Only the first layer's input, the embedding's partial sum, # is completed onto the slice. self.assertIs( - communicator.attn.plan._paths.get(BatchVariant.INPUT_SCATTERED) + communicator.attn.plan.paths.get(BatchVariant.INPUT_SCATTERED) .entry.prepare.keywords["step"] - .keywords["layer_input"], + .keywords["pre_move"], comm.tp_reduce_scatter if layer_id == 0 else None, ) self.assertIsNone( - communicator.attn.plan._paths.get( + communicator.attn.plan.paths.get( BatchVariant.INPUT_SCATTERED ).entry.input_move ) self.assertIs( - communicator.attn.plan._paths.get( + communicator.attn.plan.paths.get( BatchVariant.INPUT_SCATTERED - ).entry.handoff, - comm_ops._hand_scattered_input_to_attention, + ).entry.attn_input_adapter, + comm_ops._attn_input_scattered, ) self.assertIs( - communicator.ffn.plan._paths.get(BatchVariant.ORDINARY).output_move, - identity_output, + communicator.ffn.plan.paths.get(BatchVariant.ORDINARY).output_move, + keep_output, ) def test_two_batch_overlap_gathers_on_the_declarations(self): @@ -516,16 +522,16 @@ def test_two_batch_overlap_gathers_on_the_declarations(self): # rows: MHC writes its output into the streams, then gathers them. parallel = parallel_of(attn_dp=2, attn_tp=2, moe_dense_tp_size=1) communicator = build_mhc( - layer_case(1, 3, next_sparse=True), parallel, two_batch_overlap=True + layer_case(1, 3, next_layer_sparse=True), parallel, two_batch_overlap=True ) self.assert_step( - communicator.ffn.plan._paths.get(BatchVariant.ORDINARY).entry.prepare, - comm_ops._mlp_input_scatter, + communicator.ffn.plan.paths.get(BatchVariant.ORDINARY).entry.prepare, + comm_ops._attn_tp_reduce_scatter_update_read, communicator, ) self.assert_step( - communicator.ffn.plan._paths.get(BatchVariant.ORDINARY).output_move, - update_and_gather, + communicator.ffn.plan.paths.get(BatchVariant.ORDINARY).output_move, + update_attn_tp_gather_output, communicator, ) @@ -637,7 +643,7 @@ def layers(self, *, tbo): num_layers=4, sparse=sparse, previous_sparse=i > 2, - next_sparse=i >= 1, + next_layer_sparse=i >= 1, ) layers.append( make_test_stages(**facts, attention_norm=Norm(), ffn_norm=Norm()) @@ -646,40 +652,40 @@ def layers(self, *, tbo): def test_the_dense_layer_before_the_split_hands_on_the_attention_rows(self): attention = comm.Layout(frozenset({TokenAxis.ATTN_DP})) - local = comm.Layout(frozenset({TokenAxis.ATTN_DP, TokenAxis.ATTN_TP_SCATTER})) + local = comm.Layout(frozenset({TokenAxis.ATTN_DP, TokenAxis.ATTN_TP})) for tbo, gathered in ((True, True), (False, False)): with self.subTest(two_batch_overlap=tbo): first, before, after, last = self.layers(tbo=tbo) for layer in (first, before, after, last): - self.assertIsNotNone(layer.attn.plan.variants) + self.assertIsNotNone(layer.attn.plan.edges) self.assertIs( - first.ffn.plan._paths.get(BatchVariant.ORDINARY).output_move, - identity_output, + first.ffn.plan.paths.get(BatchVariant.ORDINARY).output_move, + keep_output, ) if gathered: self.assertIs( - before.ffn.plan._paths.get( + before.ffn.plan.paths.get( BatchVariant.ORDINARY ).output_move.func, - update_and_gather, + update_attn_tp_gather_output, ) self.assertEqual(after.attn.incoming_residual_rows, attention) self.assertIsNone( - after.attn.plan._paths.get( + after.attn.plan.paths.get( BatchVariant.ORDINARY ).entry.input_move ) else: self.assertIs( - before.ffn.plan._paths.get(BatchVariant.ORDINARY).output_move, - identity_output, + before.ffn.plan.paths.get(BatchVariant.ORDINARY).output_move, + keep_output, ) self.assertEqual(after.attn.incoming_residual_rows, local) self.assertIs( - after.attn.plan._paths.get( + after.attn.plan.paths.get( BatchVariant.ORDINARY ).entry.input_move, - gather_attention_tp, + attn_tp_gather_input, ) def test_the_split_moves_from_the_first_layer_s_rows(self): @@ -688,11 +694,11 @@ def test_the_split_moves_from_the_first_layer_s_rows(self): with planning(parallel): self.assertEqual( comm.tbo_split_moves(comm.Layout(dp)), - (identity_output, identity_output), + (keep_output, keep_output), ) self.assertEqual( - comm.tbo_split_moves(comm.Layout(dp | {TokenAxis.ATTN_TP_SCATTER})), - (update_and_gather, scatter_output), + comm.tbo_split_moves(comm.Layout(dp | {TokenAxis.ATTN_TP})), + (update_attn_tp_gather_output, attn_tp_slice_output), ) with self.assertRaises(NotImplementedError): comm.tbo_split_moves(comm.Layout(frozenset())) @@ -705,31 +711,32 @@ class TestTheAttentionOutputDecidesItsSum(CustomTestCase): def test_the_attention_output_owes_the_attention_tp_sum(self): for attn_tp in (1, 2): stages = build(layer_case(1, 3), parallel_of(attn_dp=2, attn_tp=attn_tp)) - produced = stages.attn.plan._paths.get(BatchVariant.ORDINARY).output + produced = stages.attn.plan.paths.get(BatchVariant.ORDINARY).output self.assertIs(produced.group, SumGroup.ATTN_TP if attn_tp > 1 else None) - self.assertIs(produced.always_leaves, attn_tp > 1) + self.assertIs(produced.always_partial, attn_tp > 1) self.assertIs( - stages.ffn.plan._paths.get(BatchVariant.ORDINARY).output.group, + stages.ffn.plan.paths.get(BatchVariant.ORDINARY).output.group, SumGroup.TP, ) self.assertFalse( - stages.ffn.plan._paths.get(BatchVariant.ORDINARY).output.always_leaves + stages.ffn.plan.paths.get(BatchVariant.ORDINARY).output.always_partial ) def run_steps(self, produced, *, force=False): sizes = { TokenAxis.ATTN_DP: 2, TokenAxis.ATTN_CP: 1, - TokenAxis.ATTN_TP_SCATTER: 2, + TokenAxis.ATTN_TP: 2, } - steps, _ = comm_boundary._select_input_steps( + steps, _ = comm_boundary._select_entry_step( produced, residual=comm.Layout.sharded_over(TokenAxis.ATTN_DP, axis_sizes=sizes), residual_to=comm.Layout.sharded_over(TokenAxis.ATTN_DP, axis_sizes=sizes), - need=comm.StageInput( - comm.Layout(frozenset()), read=replace(NORM_READ, before_gather=force) + need=comm.InputContract( + comm.Layout(frozenset()), + read=replace(NORM_READOUT, reads_before_dp_gather=force), ), - adds_plainly=True, + is_plain_add=True, fusions=(), residual_joins_sum=False, cp_moves=None, @@ -741,9 +748,7 @@ def run_steps(self, produced, *, force=False): "attention_tensor_model_parallel_all_reduce", lambda x: reduced.append(x) or 2 * x, ), - patch_communicator( - "_redistribute_input_to_dp", lambda h, fb, cp_shard_counts=None: h - ), + patch_communicator("dp_gather", lambda h, fb, cp_shard_counts=None: h), patch_communicator( "get_parallel", lambda: parallel_of(attn_dp=2, attn_tp=2) ), @@ -762,7 +767,7 @@ def run_steps(self, produced, *, force=False): def test_a_complete_output_is_not_summed_again(self): layout = comm.Layout(frozenset({TokenAxis.ATTN_DP})) - complete = comm.StageOutput(layout) + complete = comm.OutputContract(layout) for force in (False, True): with self.subTest(force=force): reductions, hidden, residual = self.run_steps(complete, force=force) @@ -772,7 +777,7 @@ def test_a_complete_output_is_not_summed_again(self): def test_an_owed_output_is_summed_once(self): layout = comm.Layout(frozenset({TokenAxis.ATTN_DP})) - owed = comm.StageOutput(layout, group=SumGroup.ATTN_TP, always_leaves=True) + owed = comm.OutputContract(layout, group=SumGroup.ATTN_TP, always_partial=True) reductions, hidden, residual = self.run_steps(owed, force=True) self.assertEqual(reductions, 1) # The stand-in all-reduce doubles the one rank's value. @@ -782,7 +787,7 @@ 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 = comm.Layout(frozenset({TokenAxis.ATTN_DP})) - asked = comm.StageOutput(layout, group=SumGroup.ATTN_TP) + asked = comm.OutputContract(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)) @@ -794,7 +799,7 @@ class TestFusedKernelsTakeOnlyTheStepsTheyComplete(CustomTestCase): attention output owes.""" def fused(self, completes): - return comm.FusedMlpInput( + return comm.FfnInputFusion( completes=completes, run=lambda h, r, fb: None, ) @@ -803,15 +808,15 @@ def select(self, *, attn_dp, fusions): sizes = { TokenAxis.ATTN_DP: attn_dp, TokenAxis.ATTN_CP: 1, - TokenAxis.ATTN_TP_SCATTER: 2, + TokenAxis.ATTN_TP: 2, } attention = comm.Layout.sharded_over(TokenAxis.ATTN_DP, axis_sizes=sizes) - steps, _ = comm_boundary._select_input_steps( - comm.StageOutput(attention, group=SumGroup.ATTN_TP, always_leaves=True), + steps, _ = comm_boundary._select_entry_step( + comm.OutputContract(attention, group=SumGroup.ATTN_TP, always_partial=True), residual=comm.Layout.sharded_over(TokenAxis.ATTN_DP, axis_sizes=sizes), residual_to=comm.Layout.sharded_over(TokenAxis.ATTN_DP, axis_sizes=sizes), - need=comm.StageInput(comm.Layout(frozenset())), - adds_plainly=True, + need=comm.InputContract(comm.Layout(frozenset())), + is_plain_add=True, fusions=fusions, residual_joins_sum=False, cp_moves=None, @@ -823,13 +828,13 @@ def test_a_kernel_over_another_group_is_not_chosen(self): over_tp = self.fused(SumGroup.TP) over_attention_tp = self.fused(SumGroup.ATTN_TP) steps, chosen = self.select(attn_dp=1, fusions=(over_tp, over_attention_tp)) - self.assertIs(steps.func, comm_ops._mlp_input_without_dp) + self.assertIs(steps.func, comm_ops._reduce_update_read) self.assertEqual(steps.keywords["fusions"], (over_attention_tp.run,)) self.assertEqual(chosen, (over_attention_tp.run,)) def test_no_kernel_is_chosen_where_rows_are_gathered(self): steps, chosen = self.select(attn_dp=2, fusions=(self.fused(SumGroup.ATTN_TP),)) - self.assertIs(steps.func, comm_ops._mlp_input_dp_partial) + self.assertIs(steps.func, comm_ops._dp_gather_sum_read) self.assertEqual(chosen, ()) @@ -839,7 +844,7 @@ class TestTheSequenceParallelRegion(CustomTestCase): themselves, so every boundary stays on this rank's slice. Other batches take the layer's ordinary declared steps.""" - SIZES = {TokenAxis.ATTN_DP: 1, TokenAxis.ATTN_CP: 1, TokenAxis.ATTN_TP_SCATTER: 2} + SIZES = {TokenAxis.ATTN_DP: 1, TokenAxis.ATTN_CP: 1, TokenAxis.ATTN_TP: 2} def test_what_the_ffn_exit_reads_comes_from_the_batch_s_steps(self): # Inside the region the FFN output is complete: no sum to move. @@ -862,8 +867,8 @@ def test_what_the_ffn_exit_reads_comes_from_the_batch_s_steps(self): patch_communicator("is_dp_attention_enabled", lambda: False), ): self.assertIs( - communicator.ffn.plan.output._ffn_sum_can_move_to_next_layer( - communicator.ffn.plan.output.plan._batch_steps( + communicator.ffn.plan.output._sum_deferral_allowed( + communicator.ffn.plan.output.plan.path_for( SimpleNamespace(forward_mode=ForwardMode.DECODE) ) ), @@ -877,7 +882,7 @@ class TestInputScatteredAttention(CustomTestCase): partial that a reduce-scatter completes onto each rank's slice, and the residual comes back to every row inside the attention output's sum.""" - SIZES = {TokenAxis.ATTN_DP: 1, TokenAxis.ATTN_CP: 1, TokenAxis.ATTN_TP_SCATTER: 2} + SIZES = {TokenAxis.ATTN_DP: 1, TokenAxis.ATTN_CP: 1, TokenAxis.ATTN_TP: 2} def test_the_order_of_a_residual_on_the_slice_is_declared(self): # The same layouts take two orders: with a scattered input the residual @@ -885,19 +890,21 @@ def test_the_order_of_a_residual_on_the_slice_is_declared(self): # gathered first. sizes = self.SIZES attention = comm.Layout.sharded_over(axis_sizes=sizes) - local = comm.Layout(frozenset({TokenAxis.ATTN_TP_SCATTER})) - owed = comm.StageOutput(attention, group=SumGroup.ATTN_TP, always_leaves=True) + local = comm.Layout(frozenset({TokenAxis.ATTN_TP})) + owed = comm.OutputContract( + attention, group=SumGroup.ATTN_TP, always_partial=True + ) for joins, step in ( - (True, comm_ops._mlp_input_residual_into_sum), - (False, comm_ops._mlp_input_without_dp), + (True, comm_ops._tp_sum_with_residual_read), + (False, comm_ops._reduce_update_read), ): with self.subTest(residual_joins_sum=joins): - steps, _ = comm_boundary._select_input_steps( + steps, _ = comm_boundary._select_entry_step( owed, residual=local, residual_to=attention, - need=comm.StageInput(attention), - adds_plainly=True, + need=comm.InputContract(attention), + is_plain_add=True, fusions=(), residual_joins_sum=joins, cp_moves=None, @@ -910,19 +917,19 @@ def test_a_complete_value_is_sliced_onto_each_rank(self): # layer stack, takes a complete input. sizes = self.SIZES attention = comm.Layout.sharded_over(axis_sizes=sizes) - local = comm.Layout(frozenset({TokenAxis.ATTN_TP_SCATTER})) - step, _ = comm_boundary._select_input_steps( - comm.StageOutput(attention), + local = comm.Layout(frozenset({TokenAxis.ATTN_TP})) + step, _ = comm_boundary._select_entry_step( + comm.OutputContract(attention), residual=attention, residual_to=local, - need=comm.StageInput(local), - adds_plainly=True, + need=comm.InputContract(local), + is_plain_add=True, fusions=(), residual_joins_sum=False, cp_moves=None, enters_stack=False, ) - self.assertIs(step.func, comm_ops._mlp_input_slice) + self.assertIs(step.func, comm_ops._attn_tp_slice_update_read) hidden = torch.arange(4.0)[:, None].expand(4, HIDDEN).clone() residual = torch.ones(4, HIDDEN) context = SimpleNamespace(attn_tp_size=2, attn_tp_rank=1) @@ -956,7 +963,7 @@ def test_which_layers_can_scatter_their_input(self): a2a=a2a, ) self.assertIs( - communicator.ffn.plan._paths.get(BatchVariant.INPUT_SCATTERED) + communicator.ffn.plan.paths.get(BatchVariant.INPUT_SCATTERED) is not None, expected, ) @@ -981,12 +988,12 @@ def test_a_batch_runs_them_only_while_its_input_is_scattered(self): ), ): self.assertIs( - communicator.ffn.plan._batch_steps( + communicator.ffn.plan.path_for( SimpleNamespace(forward_mode=ForwardMode.DECODE) ), - communicator.ffn.plan._paths.get(BatchVariant.INPUT_SCATTERED) + communicator.ffn.plan.paths.get(BatchVariant.INPUT_SCATTERED) if scattered - else communicator.ffn.plan._paths.get(BatchVariant.ORDINARY), + else communicator.ffn.plan.paths.get(BatchVariant.ORDINARY), ) def test_prepare_attn_completes_the_scattered_input(self): @@ -1040,8 +1047,8 @@ def test_the_last_layer_completes_its_own_sum(self): parallel, boundary_reduction="rs+rsv" if allow else "ar", ) - steps = communicator.ffn.plan._paths.get(BatchVariant.INPUT_SCATTERED) - self.assertIs(steps.output.leaves_for_reduce_scatter, hands_on) + steps = communicator.ffn.plan.paths.get(BatchVariant.INPUT_SCATTERED) + self.assertIs(steps.output.may_reduce_scatter, hands_on) class TestPrefillCP(CustomTestCase): @@ -1071,34 +1078,34 @@ def test_a_dsa_cp_extend_leaves_its_sum_to_the_reduce_scatter(self): facts = layer_case(1, 3, sparse=True, previous_sparse=False) communicator = build(facts, parallel, dsa_cp=True) cp, ordinary = ( - communicator.ffn.plan._paths.get(BatchVariant.CONTEXT_PARALLEL), - communicator.ffn.plan._paths.get(BatchVariant.ORDINARY), + communicator.ffn.plan.paths.get(BatchVariant.CONTEXT_PARALLEL), + communicator.ffn.plan.paths.get(BatchVariant.ORDINARY), ) self.assertIs( cp.entry.prepare.keywords["step"].func, - comm_ops._mlp_input_gather_attention_cp, + comm_ops._then_attn_cp_gather, ) self.assertIs( cp.entry.prepare.keywords["step"].keywords["gather"].func, - comm_ops._read_input, + comm_ops._update_read, ) self.assertIs( cp.output_move, - reduce_scatter_over_cp, + attn_cp_reduce_scatter_output, ) - self.assertTrue(cp.output.leaves_for_reduce_scatter) - self.assertFalse(cp.output.leaves_for_next_layer) + self.assertTrue(cp.output.may_reduce_scatter) + self.assertFalse(cp.output.may_defer_to_next) # The dense layer before it ran on the same shard: nothing to move. self.assertIsNone( - communicator.attn.plan._paths.get( + communicator.attn.plan.paths.get( BatchVariant.CONTEXT_PARALLEL ).entry.input_move ) # Other batches hold every token on each CP rank: the MoE sums itself. self.assertIs( - ordinary.entry.prepare.keywords["step"].func, comm_ops._read_input + ordinary.entry.prepare.keywords["step"].func, comm_ops._update_read ) - self.assertIs(ordinary.output_move, identity_output) + self.assertIs(ordinary.output_move, keep_output) with ( patch_communicator( "get_forward", @@ -1115,8 +1122,8 @@ def test_a_dsa_cp_extend_leaves_its_sum_to_the_reduce_scatter(self): patch_communicator("_batch_shards_over_cp", lambda fb: shards), ): self.assertIs( - communicator.ffn.plan.output._ffn_leaves_sum_to_reduce_scatter( - communicator.ffn.plan.output.plan._batch_steps( + communicator.ffn.plan.output._skips_sum_for_reduce_scatter( + communicator.ffn.plan.output.plan.path_for( SimpleNamespace(forward_mode=ForwardMode.DECODE) ), None, @@ -1152,42 +1159,42 @@ def test_dsa_cp_dense_layers_run_on_their_shard(self): facts = layer_case(1, 3, sparse=False, previous_sparse=False) communicator = build(facts, parallel, dsa_cp=True) for steps in ( - communicator.ffn.plan._paths.get(BatchVariant.ORDINARY), - communicator.ffn.plan._paths.get(BatchVariant.CONTEXT_PARALLEL), + communicator.ffn.plan.paths.get(BatchVariant.ORDINARY), + communicator.ffn.plan.paths.get(BatchVariant.CONTEXT_PARALLEL), ): self.assertIs( - steps.entry.prepare.keywords["step"].func, comm_ops._read_input + steps.entry.prepare.keywords["step"].func, comm_ops._update_read ) self.assertIsNone(steps.output.group) - self.assertIs(steps.output_move, identity_output) + self.assertIs(steps.output_move, keep_output) def test_a_cp_extend_gathers_over_cp_and_takes_its_chunk_back(self): communicator = build(layer_case(1, 3), self.cp_parallel()) - cp = communicator.ffn.plan._paths.get(BatchVariant.CONTEXT_PARALLEL) + cp = communicator.ffn.plan.paths.get(BatchVariant.CONTEXT_PARALLEL) self.assertIs( - cp.entry.prepare.keywords["step"].func, comm_ops._mlp_input_gather_moe_cp + cp.entry.prepare.keywords["step"].func, comm_ops._then_moe_cp_gather ) # Each rank completes its own chunk before the gather. self.assertIs( cp.entry.prepare.keywords["step"].keywords["gather"].func, - comm_ops._mlp_input_without_dp, + comm_ops._reduce_update_read, ) self.assertIs( cp.output_move, - scatter_moe_cp_output, + moe_cp_take_back_output, ) self.assertFalse(cp.returns_over_dp) - self.assertFalse(cp.output.leaves_for_next_layer) - self.assertFalse(cp.output.leaves_for_reduce_scatter) + self.assertFalse(cp.output.may_defer_to_next) + self.assertFalse(cp.output.may_reduce_scatter) self.assertIs( - communicator.ffn.plan._paths.get(BatchVariant.ORDINARY) + communicator.ffn.plan.paths.get(BatchVariant.ORDINARY) .entry.prepare.keywords["step"] .func, - comm_ops._mlp_input_without_dp, + comm_ops._reduce_update_read, ) self.assertIs( - communicator.ffn.plan._paths.get(BatchVariant.ORDINARY).output_move, - identity_output, + communicator.ffn.plan.paths.get(BatchVariant.ORDINARY).output_move, + keep_output, ) def test_a_moe_on_its_own_cp_shard_hands_on_its_sum(self): @@ -1198,27 +1205,27 @@ def test_a_moe_on_its_own_cp_shard_hands_on_its_sum(self): parallel, ) for steps in ( - communicator.ffn.plan._paths.get(BatchVariant.CONTEXT_PARALLEL), - communicator.ffn.plan._paths.get(BatchVariant.ORDINARY), + communicator.ffn.plan.paths.get(BatchVariant.CONTEXT_PARALLEL), + communicator.ffn.plan.paths.get(BatchVariant.ORDINARY), ): # The attention's rows are the MoE's: nothing to gather or take back. self.assertIs( steps.entry.prepare.keywords["step"].func, - comm_ops._mlp_input_without_dp, + comm_ops._reduce_update_read, ) - self.assertIs(steps.output_move, identity_output) - self.assertTrue(steps.output.leaves_for_next_layer) + self.assertIs(steps.output_move, keep_output) + self.assertTrue(steps.output.may_defer_to_next) dense = build(layer_case(1, 3), parallel) self.assertIs( - dense.ffn.plan._paths.get(BatchVariant.CONTEXT_PARALLEL) + dense.ffn.plan.paths.get(BatchVariant.CONTEXT_PARALLEL) .entry.prepare.keywords["step"] .func, - comm_ops._mlp_input_gather_moe_cp, + comm_ops._then_moe_cp_gather, ) self.assertFalse( - dense.ffn.plan._paths.get( + dense.ffn.plan.paths.get( BatchVariant.CONTEXT_PARALLEL - ).output.leaves_for_next_layer + ).output.may_defer_to_next ) def test_the_fused_kernels_run_on_each_chunk(self): @@ -1227,12 +1234,12 @@ def test_the_fused_kernels_run_on_each_chunk(self): **layer_case(1, 3), attention_norm=FusableNorm(), ffn_norm=FusableNorm() ) chunk = ( - communicator.ffn.plan._paths.get(BatchVariant.CONTEXT_PARALLEL) + communicator.ffn.plan.paths.get(BatchVariant.CONTEXT_PARALLEL) .entry.prepare.keywords["step"] .keywords["gather"] ) self.assertEqual(len(chunk.keywords["fusions"]), 1) - self.assertIs(chunk.keywords["fusions"][0].func, complete_ffn_input) + self.assertIs(chunk.keywords["fusions"][0].func, fused_ffn_input) self.assertEqual(chunk.keywords["fusions"][0].args, (communicator.ffn.plan,)) def cp_outcome(self, parallel, *, sparse=True, a2a=False, dsa_cp=False): @@ -1245,7 +1252,7 @@ def cp_outcome(self, parallel, *, sparse=True, a2a=False, dsa_cp=False): return "refused" return ( "ordinary" - if layer.ffn.plan._paths.get(BatchVariant.CONTEXT_PARALLEL) is None + if layer.ffn.plan.paths.get(BatchVariant.CONTEXT_PARALLEL) is None else "cp steps" ) @@ -1325,25 +1332,25 @@ 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_case(1, 3), parallel) cp, ordinary = ( - communicator.ffn.plan._paths.get(BatchVariant.CONTEXT_PARALLEL), - communicator.ffn.plan._paths.get(BatchVariant.ORDINARY), + communicator.ffn.plan.paths.get(BatchVariant.CONTEXT_PARALLEL), + communicator.ffn.plan.paths.get(BatchVariant.ORDINARY), ) self.assertIs( - cp.entry.prepare.keywords["step"].func, comm_ops._mlp_input_dp_partial + cp.entry.prepare.keywords["step"].func, comm_ops._dp_gather_sum_read ) self.assertTrue(cp.entry.prepare.keywords["step"].keywords["places_cp_shards"]) - self.assertIs(cp.output_move, take_back_cp_shard) + self.assertIs(cp.output_move, dp_cp_take_back_output) self.assertIs( - ordinary.entry.prepare.keywords["step"].func, comm_ops._mlp_input_dp_partial + ordinary.entry.prepare.keywords["step"].func, comm_ops._dp_gather_sum_read ) self.assertFalse( ordinary.entry.prepare.keywords["step"].keywords["places_cp_shards"] ) self.assertTrue(ordinary.returns_over_dp) for steps in (cp, ordinary): - self.assertFalse(steps.output.leaves_for_next_layer) - self.assertFalse(steps.output.leaves_for_reduce_scatter) - self.assertFalse(steps.output.leaves_for_reduce_scatterv) + self.assertFalse(steps.output.may_defer_to_next) + self.assertFalse(steps.output.may_reduce_scatter) + self.assertFalse(steps.output.may_reduce_scatterv) def test_a_batch_runs_them_only_on_a_cp_extend(self): communicator = build(layer_case(1, 3), self.cp_parallel()) @@ -1362,23 +1369,23 @@ def unread(fb): ( batch(False), unread, - communicator.ffn.plan._paths.get(BatchVariant.ORDINARY), + communicator.ffn.plan.paths.get(BatchVariant.ORDINARY), ), ( batch(True), lambda fb: None, - communicator.ffn.plan._paths.get(BatchVariant.ORDINARY), + communicator.ffn.plan.paths.get(BatchVariant.ORDINARY), ), ( batch(True), lambda fb: [2, 1], - communicator.ffn.plan._paths.get(BatchVariant.CONTEXT_PARALLEL), + communicator.ffn.plan.paths.get(BatchVariant.CONTEXT_PARALLEL), ), ): with ( self.subTest( expected=expected - is communicator.ffn.plan._paths.get(BatchVariant.CONTEXT_PARALLEL) + is communicator.ffn.plan.paths.get(BatchVariant.CONTEXT_PARALLEL) ), planning(self.cp_parallel()), patch_communicator("moe_cp_gathered_rows", rows), @@ -1393,7 +1400,7 @@ def unread(fb): lambda: SimpleNamespace(input_scattered=False), ), ): - self.assertIs(communicator.ffn.plan._batch_steps(fb), expected) + self.assertIs(communicator.ffn.plan.path_for(fb), expected) class TestBranchRows(CustomTestCase): @@ -1402,7 +1409,7 @@ class TestBranchRows(CustomTestCase): of its residual while the FFN runs and of what it hands on, and a complete value moves between them.""" - local = comm.Layout(frozenset({TokenAxis.ATTN_DP, TokenAxis.ATTN_TP_SCATTER})) + local = comm.Layout(frozenset({TokenAxis.ATTN_DP, TokenAxis.ATTN_TP})) attention = comm.Layout(frozenset({TokenAxis.ATTN_DP})) full = comm.Layout(frozenset()) @@ -1413,9 +1420,9 @@ class Communicator(SimpleNamespace): branch_output = branch.branch_output merge_branch = branch.merge_branch - moe = Communicator(_branch_rows=lambda fb: (self.local, self.local, self.local)) + moe = Communicator(branch_rows=lambda fb: (self.local, self.local, self.local)) dense = Communicator( - _branch_rows=lambda fb: (self.full, self.attention, self.attention) + branch_rows=lambda fb: (self.full, self.attention, self.attention) ) moves = [] @@ -1437,7 +1444,7 @@ def recorded(value, rows, to, forward_batch): self.assertEqual(moves, [("shortcut", self.local, self.local)]) moves.clear() stream = ResidualStream("residual") - hidden = stream.leave(2, comm.ADD) + hidden = stream.record(2, comm.PLAIN_ADD) merged = moe.merge_branch(1, hidden, stream, dense, None) self.assertEqual( moves, @@ -1467,9 +1474,9 @@ def tp_cut(h): value = torch.arange(4.0).view(4, 1) with ( - patch_communicator("_redistribute_from_attn_tp_shards", tp_gather), - patch_communicator("_redistribute_input_to_dp", lambda h, fb: dp_gather(h)), - patch_communicator("_to_local_tokens", lambda step, fb, h: dp_scatter(h)), + patch_communicator("attn_tp_gather", tp_gather), + patch_communicator("dp_gather", lambda h, fb: dp_gather(h)), + patch_communicator("to_dp_local", lambda step, fb, h: dp_scatter(h)), patch_communicator( "get_parallel", lambda: SimpleNamespace(attn_tp_size=2, attn_tp_rank=1), @@ -1489,7 +1496,7 @@ def tp_cut(h): ) cp = comm.Layout(frozenset({TokenAxis.ATTN_CP})) for rows, to in ( - (comm.Layout(frozenset({TokenAxis.ATTN_TP_SCATTER})), self.attention), + (comm.Layout(frozenset({TokenAxis.ATTN_TP})), self.attention), (cp, self.full), ): with ( @@ -1655,7 +1662,7 @@ def running(*, reduce_scatterv, a2a=False, use_reduce_scatter=True): "should_use_dp_reduce_scatterv": lambda: reduce_scatterv, "can_use_dp_reduce_scatter": lambda: True, "is_dp_attention_enabled": lambda: state().parallel.attn_dp_size > 1, - "apply_flashinfer_allreduce_fusion": lambda n: False, + "flashinfer_ar_fusion_applies": lambda n: False, "post_experts_output_is_complete": lambda **kw: False, "post_experts_reduction_group": lambda: state().parallel.tp_group, "get_lora": lambda: SimpleNamespace(enable_lora=False), @@ -1824,12 +1831,12 @@ def forward(rank): layer.attn, hidden, residual, forward_batch ) hidden = attention(hidden, s) - hidden = residual.leave( + hidden = residual.record( hidden, - comm.ADD, - declared_sum=layer.ffn.plan._batch_steps( + comm.PLAIN_ADD, + declared_sum=layer.ffn.plan.path_for( forward_batch - ).entry.input_sum, + ).entry.declared_sum, ) hidden, residual = prepare_input( layer.ffn, hidden, residual, forward_batch @@ -1846,7 +1853,7 @@ def forward(rank): hidden = moe(hidden, s) hidden, residual = finish_exit(ffn_exit, hidden, residual) handed_on.append(type(hidden)) - hidden, residual = finish_layer_stack(hidden, residual, forward_batch) + hidden, residual = export_output(hidden, residual, forward_batch) # The last a2a layer folds the residual into its output. if residual is not None: residual = residual[: s.rows] @@ -1956,9 +1963,9 @@ def test_only_a_plain_residual_reduces_quantized(self): batch = SimpleNamespace( forward_mode=SimpleNamespace(is_decode_or_idle=lambda: False) ) - for adds_plainly, reduction in ((True, "quant"), (False, "full")): - with self.subTest(adds_plainly=adds_plainly): - update = SimpleNamespace(adds_plainly=adds_plainly) + for is_plain_add, reduction in ((True, "quant"), (False, "full")): + with self.subTest(is_plain_add=is_plain_add): + update = SimpleNamespace(is_plain_add=is_plain_add) read = SimpleNamespace( update_and_read=lambda update, h, r, norm: (h, r) ) @@ -1974,7 +1981,7 @@ def test_only_a_plain_residual_reduces_quantized(self): lambda h: calls.append("full") or h, ), ): - comm_ops._mlp_input_without_dp( + comm_ops._reduce_update_read( torch.ones(2, HIDDEN), torch.ones(2, HIDDEN), batch, diff --git a/test/registered/unit/layer_boundary/test_deepstack_deferred_allreduce.py b/test/registered/unit/layer_boundary/test_deepstack_deferred_allreduce.py index 9a6f5941adf9..232cd75c1bd5 100644 --- a/test/registered/unit/layer_boundary/test_deepstack_deferred_allreduce.py +++ b/test/registered/unit/layer_boundary/test_deepstack_deferred_allreduce.py @@ -10,7 +10,7 @@ from sglang.srt.layers.layer_boundary import UnreducedOutput from sglang.srt.layers.layer_boundary.residual import batch as residual_batch -from sglang.srt.layers.layer_boundary.residual.add_norm import ADD +from sglang.srt.layers.layer_boundary.residual.add_norm import PLAIN_ADD from sglang.srt.layers.layer_boundary.residual.stream import OwedOutput from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase @@ -58,8 +58,8 @@ def forward( class DeferringStageLayer(DeferringLayer): def forward(self, positions=None, hidden_states=None, forward_batch=None, **kwargs): - stream = residual_batch.current(forward_batch) - hidden_states, residual = stream.finish(hidden_states) + stream = residual_batch.stream_of(forward_batch) + hidden_states, residual = stream.export(hidden_states) residual = hidden_states if residual is None else hidden_states + residual stream.write(residual) output = ( @@ -67,7 +67,7 @@ def forward(self, positions=None, hidden_states=None, forward_batch=None, **kwar if self.is_last_layer else UnreducedOutput(torch.full_like(residual, 1 / TP_SIZE), group=GROUP) ) - return stream.leave(output, ADD) + return stream.record(output, PLAIN_ADD) class SumNorm(nn.Module): diff --git a/test/registered/unit/layer_boundary/test_dense_mlp_under_prefill_cp.py b/test/registered/unit/layer_boundary/test_dense_mlp_under_prefill_cp.py index 25306befca74..00b8bcb7033e 100644 --- a/test/registered/unit/layer_boundary/test_dense_mlp_under_prefill_cp.py +++ b/test/registered/unit/layer_boundary/test_dense_mlp_under_prefill_cp.py @@ -7,7 +7,7 @@ ) from sglang.srt.layers.layer_boundary import prepare as comm_ops from sglang.srt.layers.layer_boundary.contracts import BatchVariant -from sglang.srt.layers.layer_boundary.ops import scatter_moe_cp_output +from sglang.srt.layers.layer_boundary.ops import moe_cp_take_back_output from sglang.test.boundary_fixtures import make_test_stages from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase @@ -28,16 +28,16 @@ def test_dense_mlp_gathers_across_cp(self): stages = make_test_stages( attention_norm=fixture.Norm(), ffn_norm=fixture.Norm() ) - attention = stages.attn.plan._paths.get(BatchVariant.CONTEXT_PARALLEL).entry - steps = stages.ffn.plan._paths.get(BatchVariant.CONTEXT_PARALLEL) + attention = stages.attn.plan.paths.get(BatchVariant.CONTEXT_PARALLEL).entry + steps = stages.ffn.plan.paths.get(BatchVariant.CONTEXT_PARALLEL) self.assertIn(TokenAxis.ATTN_CP, attention.input_rows.sharded) self.assertNotIn(TokenAxis.ATTN_CP, steps.entry.input_rows.sharded) self.assertIs( - steps.entry.prepare.keywords["step"].func, comm_ops._mlp_input_gather_moe_cp + steps.entry.prepare.keywords["step"].func, comm_ops._then_moe_cp_gather ) self.assertIs( steps.output_move, - scatter_moe_cp_output, + moe_cp_take_back_output, ) diff --git a/test/registered/unit/layer_boundary/test_dp_attention_cp_gather.py b/test/registered/unit/layer_boundary/test_dp_attention_cp_gather.py index 7fa7f9685eb4..ac11493e79ba 100644 --- a/test/registered/unit/layer_boundary/test_dp_attention_cp_gather.py +++ b/test/registered/unit/layer_boundary/test_dp_attention_cp_gather.py @@ -1,7 +1,7 @@ """The dense-FFN gather and take-back under attention DP x CP x TP, rank by rank. Every rank builds its layer's communicator and runs its prepare_mlp and -postprocess_layer for real on CPU; the TP-group all-reduce is replaced by the +finish_complete_output for real on CPU; the TP-group all-reduce is replaced by the sum of the buffers the ranks hand to it. """ @@ -22,7 +22,10 @@ from sglang.srt.layers.dp_attention import DpPaddingMode from sglang.srt.layers.layer_boundary import prepare as comm_ops from sglang.srt.layers.layer_boundary.contracts import BatchVariant -from sglang.srt.layers.layer_boundary.residual.add_norm import NORM_READ, PLAIN_RESIDUAL +from sglang.srt.layers.layer_boundary.residual.add_norm import ( + NORM_READOUT, + PLAIN_RESIDUAL_OPS, +) from sglang.srt.model_executor.forward_batch_info import ForwardMode from sglang.test.boundary_fixtures import ( make_test_stages, @@ -143,7 +146,7 @@ def run_ranks( padding=DpPaddingMode.SUM_LEN, norm=layernorm, norm_rows=lambda rows: rows, - read=NORM_READ, + read=NORM_READOUT, ffn_input=None, sparse=False, ): @@ -267,10 +270,10 @@ def as_rank(dp, cp, tp, all_reduce): last=False, sparse=sparse, previous_sparse=sparse, - next_sparse=False, + next_layer_sparse=False, attention_norm=norm, ffn_norm=norm, - residual=PLAIN_RESIDUAL._replace(ffn_read=read), + residual=PLAIN_RESIDUAL_OPS._replace(ffn_readout=read), ) yield SimpleNamespace( communicator=communicator, @@ -290,11 +293,11 @@ def as_rank(dp, cp, tp, all_reduce): def gather(rank, all_reduce): with as_rank(*rank, all_reduce) as r: self.assertIsNotNone( - r.communicator.ffn.plan._paths.get(BatchVariant.CONTEXT_PARALLEL), + r.communicator.ffn.plan.paths.get(BatchVariant.CONTEXT_PARALLEL), "declared under CP", ) if ffn_input is not None: - steps = r.communicator.ffn.plan._batch_steps(r.forward_batch) + steps = r.communicator.ffn.plan.path_for(r.forward_batch) self.assertIs(steps.entry.prepare.keywords["step"].func, ffn_input) hidden_states, residual = rank_inputs(*rank) return prepare_input( @@ -410,8 +413,8 @@ def test_forced_norm_before_gather_under_attention_tp(self): attn_tp_size=2, norm=rms_norm, norm_rows=rms_rows, - read=replace(NORM_READ, before_gather=True), - ffn_input=comm_ops._mlp_input_dp_replicate, + read=replace(NORM_READOUT, reads_before_dp_gather=True), + ffn_input=comm_ops._reduce_update_read_dp_gather, ) diff --git a/test/registered/unit/layer_boundary/test_ffn_exit.py b/test/registered/unit/layer_boundary/test_ffn_exit.py index 9dbcf1dd0a8d..da19337cfe0a 100644 --- a/test/registered/unit/layer_boundary/test_ffn_exit.py +++ b/test/registered/unit/layer_boundary/test_ffn_exit.py @@ -8,18 +8,16 @@ from sglang.srt.layers import layer_boundary as comm from sglang.srt.layers.layer_boundary import ( Layout, - StageOutput, + OutputContract, SumGroup, UnreducedOutput, + complete_owed, ) from sglang.srt.layers.layer_boundary import ops as transport_ops from sglang.srt.layers.layer_boundary import prepare as comm_ops -from sglang.srt.layers.layer_boundary import ( - reduce_output, -) from sglang.srt.layers.layer_boundary.contracts import BatchVariant -from sglang.srt.layers.layer_boundary.ops import identity_output -from sglang.srt.layers.layer_boundary.residual.access import finish_layer_stack +from sglang.srt.layers.layer_boundary.ops import keep_output +from sglang.srt.layers.layer_boundary.residual.access import export_output from sglang.srt.layers.layer_boundary.residual.stream import ResidualStream from sglang.srt.runtime_context import get_forward from sglang.test.boundary_fixtures import ( @@ -43,26 +41,26 @@ def make_group(scale=3): def ordinary_steps(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.StageSteps( - entry=comm.StageEntry( + return comm.StagePath( + entry=comm.EntryPath( prepare=partial( - comm_ops._consumer_step, + comm_ops._run_entry, step=partial( - comm_ops._read_input, - layer_input=None, + comm_ops._update_read, + pre_move=None, enters_stack=False, - read=comm.NORM_QUANT_READ, - update=comm.ADD, + read=comm.NORM_QUANT_READOUT, + update=comm.PLAIN_ADD, ), carried_fusions=(), ), input_rows=comm.Layout(frozenset()), input_move=identity_input, - handoff=comm_ops._hand_qkv_hook_its_input, + attn_input_adapter=comm_ops._attn_input_default, ), output=output, returns_over_dp=returns_over_dp, - output_move=None if returns_over_dp else identity_output, + output_move=None if returns_over_dp else keep_output, ) @@ -80,32 +78,32 @@ def make_communicator( without the process-wide parallel state.""" communicator = stub_plan() communicator.fusions = fusions - communicator._paths[BatchVariant.ORDINARY] = ordinary_steps( - StageOutput( + communicator.paths[BatchVariant.ORDINARY] = ordinary_steps( + OutputContract( Layout(frozenset()), group=SumGroup.TP, - leaves_for_next_layer=leaves_sum, + may_defer_to_next=leaves_sum, ), returns_over_dp=scatters_to_local_tokens, ) if fusions is None: - communicator.output._ffn_sum_moves_to_next_layer = MagicMock(return_value=fuse) + communicator.output._defers_sum = MagicMock(return_value=fuse) else: - communicator.output._ffn_sum_moves_to_next_layer = MagicMock(return_value=False) - communicator.output._ffn_leaves_sum_to_reduce_scatter = MagicMock( + communicator.output._defers_sum = MagicMock(return_value=False) + communicator.output._skips_sum_for_reduce_scatter = MagicMock( return_value=reduce_scatter ) communicator.terminal = False - communicator._paths[BatchVariant.SEQUENCE_PARALLEL] = None - communicator._paths[BatchVariant.INPUT_SCATTERED] = None - communicator._paths[BatchVariant.CONTEXT_PARALLEL] = None - communicator.output._postprocess_dp_step = MagicMock( + communicator.paths[BatchVariant.SEQUENCE_PARALLEL] = None + communicator.paths[BatchVariant.INPUT_SCATTERED] = None + communicator.paths[BatchVariant.CONTEXT_PARALLEL] = None + communicator.output._dp_reduce_scatter_step = MagicMock( return_value=reduce_scatter_step ) communicator.output.ffn_reduction_group = MagicMock( return_value=group or make_group() ) - communicator.output._complete_ffn_output_now = MagicMock( + communicator.output._complete_now = MagicMock( side_effect=lambda hidden_states, residual, **_: (hidden_states + 1, residual) ) return communicator @@ -138,7 +136,7 @@ def test_reduction_left_to_next_layer_declares_its_group(self): communicator = make_communicator(fuse=True, reduce_scatter=False, group=group) _, (hidden_states, _) = self.run_exit(communicator) self.assertIs(hidden_states.group, group) - self.assertIsNone(hidden_states.reduce_and_redistribute) + self.assertIsNone(hidden_states.reduce_to_dp_local) def test_a_reduce_scatter_is_left_to_the_next_layer(self): """Under attention DP the reduce-scatter postprocess would run goes to the @@ -153,11 +151,11 @@ def test_a_reduce_scatter_is_left_to_the_next_layer(self): seen, (hidden_states, residual) = self.run_exit(communicator) self.assertEqual(seen, (False, True)) self.assertIsInstance(hidden_states, UnreducedOutput) - bound = hidden_states.reduce_and_redistribute - self.assertIs(bound.func, transport_ops._to_local_tokens) + bound = hidden_states.reduce_to_dp_local + self.assertIs(bound.func, transport_ops.to_dp_local) self.assertEqual(bound.args, (step, self.forward_batch)) self.assertIs(residual, self.residual) - communicator.output._complete_ffn_output_now.assert_not_called() + communicator.output._complete_now.assert_not_called() step.assert_not_called() def test_a_deferred_sum_carries_the_scatter_back_under_attention_dp(self): @@ -167,10 +165,10 @@ def test_a_deferred_sum_carries_the_scatter_back_under_attention_dp(self): ) _, (hidden_states, _) = self.run_exit(communicator) self.assertIsInstance(hidden_states, UnreducedOutput) - bound = hidden_states.reduce_and_redistribute - self.assertIs(bound.func, transport_ops._all_reduce_then_to_local_tokens) + bound = hidden_states.reduce_to_dp_local + self.assertIs(bound.func, transport_ops.all_reduce_to_dp_local) self.assertEqual(bound.args, (group, self.forward_batch)) - communicator.output._complete_ffn_output_now.assert_not_called() + communicator.output._complete_now.assert_not_called() group.all_reduce.assert_not_called() def test_postprocess_completes_other_exits(self): @@ -182,7 +180,7 @@ def test_postprocess_completes_other_exits(self): seen, (hidden_states, residual) = self.run_exit(communicator) self.assertEqual(seen, (False, reduce_scatter)) self.assertNotIsInstance(hidden_states, UnreducedOutput) - communicator.output._complete_ffn_output_now.assert_called_once() + communicator.output._complete_now.assert_called_once() torch.testing.assert_close(hidden_states, self.hidden_states * 2 + 1) self.assertIs(residual, self.residual) @@ -193,8 +191,8 @@ def test_a_layer_that_declares_no_deferral_runs_postprocess(self): seen, (hidden_states, residual) = self.run_exit(communicator) self.assertEqual(seen, (False, True)) - communicator.output._ffn_sum_moves_to_next_layer.assert_not_called() - communicator.output._complete_ffn_output_now.assert_called_once() + communicator.output._defers_sum.assert_not_called() + communicator.output._complete_now.assert_called_once() torch.testing.assert_close(hidden_states, self.hidden_states * 2 + 1) def test_finish_applies_the_decision_the_ffn_saw(self): @@ -207,25 +205,23 @@ def test_finish_applies_the_decision_the_ffn_saw(self): self.forward_batch, stream=ResidualStream() ) as ffn_exit: seen = published_flags() - decide = communicator.output._ffn_sum_moves_to_next_layer + decide = communicator.output._defers_sum decide.return_value = not fuse hidden_states = self.hidden_states * 2 hidden_states, _ = finish_exit(ffn_exit, hidden_states, self.residual) self.assertEqual(seen, (fuse, False)) self.assertEqual(isinstance(hidden_states, UnreducedOutput), fuse) - self.assertEqual( - communicator.output._complete_ffn_output_now.called, not fuse - ) + self.assertEqual(communicator.output._complete_now.called, not fuse) def test_exit_selects_its_batch_path_once(self): communicator = make_communicator(fuse=False, reduce_scatter=False) - selected = communicator._paths.get(BatchVariant.ORDINARY) + selected = communicator.paths.get(BatchVariant.ORDINARY) lookup = MagicMock(return_value=selected) - communicator._batch_steps = lookup + communicator.path_for = lookup self.run_exit(communicator) lookup.assert_called_once_with(self.forward_batch) self.assertIs( - communicator.output._complete_ffn_output_now.call_args.kwargs["steps"], + communicator.output._complete_now.call_args.kwargs["steps"], selected, ) @@ -253,7 +249,7 @@ def test_deferral_implies_fusion_and_passes_the_handoff_through(self): self.assertEqual( finish_exit(ffn_exit, handoff, self.residual), (handoff, self.residual) ) - communicator.output._complete_ffn_output_now.assert_not_called() + communicator.output._complete_now.assert_not_called() self.assertFalse(get_forward().defer_moe_finalize) # The MoE may decline per forward; a tensor then takes the fused exit. @@ -285,21 +281,19 @@ def test_reduction_left_by_the_last_layer_runs_once(self): hidden_states = torch.ones(3, 4) hidden_states, _ = finish_exit(ffn_exit, hidden_states, torch.zeros(3, 4)) - hidden_states = reduce_output(hidden_states) + hidden_states = complete_owed(hidden_states) self.all_reduce.assert_called_once() torch.testing.assert_close(hidden_states, torch.full((3, 4), 3.0)) self.assertNotIsInstance(hidden_states, UnreducedOutput) - reduce_output(hidden_states) + complete_owed(hidden_states) self.all_reduce.assert_called_once() def test_a_reduce_scatter_left_by_the_last_layer_runs_once(self): local = torch.full((1, 4), 7.0) step = MagicMock(return_value=local) partial = torch.ones(3, 4) - hidden_states = reduce_output( - UnreducedOutput(partial, reduce_and_redistribute=step) - ) + hidden_states = complete_owed(UnreducedOutput(partial, reduce_to_dp_local=step)) step.assert_called_once_with(partial) self.assertIs(hidden_states, local) self.all_reduce.assert_not_called() @@ -316,7 +310,7 @@ def test_finish_layer_stack_completes_the_last_layer(self): hidden_states = torch.ones(3, 4) hidden_states, _ = finish_exit(ffn_exit, hidden_states, residual) - hidden_states, residual_out = finish_layer_stack( + hidden_states, residual_out = export_output( hidden_states, residual, object() ) self.assertEqual(self.all_reduce.call_count, int(fuse)) @@ -339,38 +333,34 @@ def communicator( ): communicator = stub_plan() communicator.terminal = is_last_layer - communicator._paths[BatchVariant.SEQUENCE_PARALLEL] = ( + communicator.paths[BatchVariant.SEQUENCE_PARALLEL] = ( sp_region_steps() if sp_region else None ) - communicator._paths[BatchVariant.INPUT_SCATTERED] = None - communicator._paths[BatchVariant.CONTEXT_PARALLEL] = None - communicator._paths[BatchVariant.ORDINARY] = ordinary_steps( - StageOutput( + communicator.paths[BatchVariant.INPUT_SCATTERED] = None + communicator.paths[BatchVariant.CONTEXT_PARALLEL] = None + communicator.paths[BatchVariant.ORDINARY] = ordinary_steps( + OutputContract( Layout(frozenset()), group=SumGroup.MOE_OUTPUT, - leaves_for_next_layer=True, - leaves_for_reduce_scatter=True, - leaves_for_reduce_scatterv=True, + may_defer_to_next=True, + may_reduce_scatter=True, + may_reduce_scatterv=True, ), returns_over_dp=scatters, ) - communicator.output._ffn_sum_moves_to_next_layer = ( - lambda forward_batch, steps, **_: fuse - ) - communicator.output._ffn_leaves_sum_to_reduce_scatter = ( + communicator.output._defers_sum = lambda forward_batch, steps, **_: fuse + communicator.output._skips_sum_for_reduce_scatter = ( lambda forward_batch, dp_step: not fuse ) - communicator.output._complete_ffn_output_now = lambda h, r, **_: ("now", r) + communicator.output._complete_now = lambda h, r, **_: ("now", r) self.group = make_group() communicator.output.ffn_reduction_group = lambda forward_batch: self.group return communicator def left(self, communicator, step, forward_batch=None): - with patch_communicator( - "_reduce_and_redistribute_output_step", return_value=step - ): - completion = communicator.output._select_ffn_completion( - forward_batch, communicator.output.plan._batch_steps(forward_batch) + with patch_communicator("_select_dp_reduce_scatter", return_value=step): + completion = communicator.output._decide( + forward_batch, communicator.output.plan.path_for(forward_batch) ) hidden_states, _ = completion.complete(torch.ones(3, 4), None) return None if hidden_states == "now" else hidden_states @@ -378,17 +368,17 @@ def left(self, communicator, step, forward_batch=None): def test_a_reduce_scatter_is_bound_for_the_next_layer(self): forward_batch = object() for step in ( - transport_ops._reduce_and_redistribute_output_varlen, - transport_ops._reduce_and_redistribute_output_max_len, + transport_ops.dp_reduce_scatterv, + transport_ops.dp_reduce_scatter, ): with self.subTest(step=step.__name__): left = self.left(self.communicator(), step, forward_batch) - bound = left.reduce_and_redistribute - self.assertIs(bound.func, transport_ops._to_local_tokens) + bound = left.reduce_to_dp_local + self.assertIs(bound.func, transport_ops.to_dp_local) self.assertEqual(bound.args, (step, forward_batch)) def test_postprocess_keeps_everything_else(self): - reduce_scatter = transport_ops._reduce_and_redistribute_output_varlen + reduce_scatter = transport_ops.dp_reduce_scatterv for name, communicator, step in ( ("scatter only", self.communicator(), None), ("last layer", self.communicator(is_last_layer=True), reduce_scatter), @@ -400,21 +390,17 @@ def test_postprocess_keeps_everything_else(self): def test_postprocess_keeps_an_active_layernorm_sp_region(self): communicator = self.communicator(sp_region=True) with get_forward().scoped(sp_active=True): - self.assertIsNone( - self.left( - communicator, transport_ops._reduce_and_redistribute_output_varlen - ) - ) + self.assertIsNone(self.left(communicator, transport_ops.dp_reduce_scatterv)) def test_a_deferred_sum_keeps_its_layout_or_scatters_back(self): forward_batch = object() kept = self.left(self.communicator(fuse=True, scatters=False), None) self.assertIs(kept.group, self.group) - self.assertIsNone(kept.reduce_and_redistribute) + self.assertIsNone(kept.reduce_to_dp_local) moved = self.left(self.communicator(fuse=True), None, forward_batch) self.assertIsNone(moved.group) - bound = moved.reduce_and_redistribute - self.assertIs(bound.func, transport_ops._all_reduce_then_to_local_tokens) + bound = moved.reduce_to_dp_local + self.assertIs(bound.func, transport_ops.all_reduce_to_dp_local) self.assertEqual(bound.args, (self.group, forward_batch)) def test_the_all_reduce_runs_before_the_scatter(self): @@ -435,9 +421,7 @@ def test_the_all_reduce_runs_before_the_scatter(self): side_effect=lambda out, full, fb: calls.append(("scatter", out)), ), ): - result = transport_ops._all_reduce_then_to_local_tokens( - group, object(), partial - ) + result = transport_ops.all_reduce_to_dp_local(group, object(), partial) self.assertEqual([c[0] for c in calls], ["all_reduce", "buffer", "scatter"]) self.assertIs(calls[0][1], partial) self.assertIs(result, local) diff --git a/test/registered/unit/layer_boundary/test_independent_stage_construction.py b/test/registered/unit/layer_boundary/test_independent_stage_construction.py index 9fc5316aca58..aec92ca795cc 100644 --- a/test/registered/unit/layer_boundary/test_independent_stage_construction.py +++ b/test/registered/unit/layer_boundary/test_independent_stage_construction.py @@ -35,7 +35,7 @@ def test_active_variant_without_a_bound_path_never_falls_back(self): for variant in BatchVariant: plan = stub_plan() ordinary = object() - plan._paths[BatchVariant.ORDINARY] = ordinary + plan.paths[BatchVariant.ORDINARY] = ordinary with ( self.subTest(variant=variant), patch.object( @@ -59,13 +59,13 @@ def test_active_variant_without_a_bound_path_never_falls_back(self): ), ): if variant is BatchVariant.ORDINARY: - self.assertIs(plan._batch_steps(None), ordinary) + self.assertIs(plan.path_for(None), ordinary) else: with self.assertRaisesRegex(NotImplementedError, variant.name): - plan._batch_steps(None) + plan.path_for(None) selected = object() - plan._paths[variant] = selected - self.assertIs(plan._batch_steps(None), selected) + plan.paths[variant] = selected + self.assertIs(plan.path_for(None), selected) def test_sequences_of_one_to_four_stages(self): for count in range(1, 5): @@ -146,12 +146,10 @@ def test_sequence_and_independent_factories_agree_on_terminal_placement(self): self.assertEqual( independent.plan.terminal, declaration.terminal ) - self.assertEqual( - combined.plan.variants, independent.plan.variants - ) - for variant, steps in combined.plan._paths.items(): + self.assertEqual(combined.plan.edges, independent.plan.edges) + for variant, steps in combined.plan.paths.items(): self.assertEqual( - steps.output, independent.plan._paths[variant].output + steps.output, independent.plan.paths[variant].output ) def test_sequence_branch_keeps_source_lineage(self): @@ -171,7 +169,7 @@ def test_sequence_branch_keeps_source_lineage(self): (declare_ffn(sparse=True), fixture.Norm()), previous=declare_ffn(sparse=True), ) - before = dict(source.plan.variants) + before = dict(source.plan.edges) first, attn, last = make_stages( ( declare_ffn(), @@ -184,12 +182,12 @@ def test_sequence_branch_keeps_source_lineage(self): ), prepared_from=source.declaration, ) - self.assertEqual(source.plan.variants, before) + self.assertEqual(source.plan.edges, before) self.assertIs(first.declaration.prepared_from, source.declaration) self.assertIs(attn.declaration.previous, first.declaration) self.assertIs(last.declaration.previous, attn.declaration) - for v, edge in first.plan.variants.items(): - source_input = source.plan.variants[v].incoming + for v, edge in first.plan.edges.items(): + source_input = source.plan.edges[v].incoming self.assertEqual( edge.incoming.produced.layout, source_input.need.layout ) @@ -202,21 +200,23 @@ def test_a_replicated_tail_stays_in_compute_without_a_fusion_object(self): declaration=declare_ffn( previous=declare_attn(), sparse=True, - reduction=ProducerReduction.LOCAL_TAIL + reduction=ProducerReduction.TAIL_AFTER_SUM if local_tail - else ProducerReduction.SCOPED, + else ProducerReduction.EXIT_SCOPED, ), norm=fixture.Norm(), ) self.assertIsNone(boundary.plan.fusions) - produced = boundary.plan._paths[BatchVariant.ORDINARY].output - self.assertEqual(produced.leaves_for_next_layer, not local_tail) + produced = boundary.plan.paths[BatchVariant.ORDINARY].output + self.assertEqual(produced.may_defer_to_next, not local_tail) # A local compute tail does not change the existing RS alternative. - self.assertTrue(produced.leaves_for_reduce_scatter) + self.assertTrue(produced.may_reduce_scatter) def test_mixer_successor_comes_from_the_local_sequence(self): for following in ( - declare_attn(reduction=ProducerReduction.SCOPED, gathers_tp_input=False), + declare_attn( + reduction=ProducerReduction.EXIT_SCOPED, gathers_attn_tp_input=False + ), declare_ffn(), ): with ( @@ -226,23 +226,24 @@ def test_mixer_successor_comes_from_the_local_sequence(self): mixer, consumer = make_stages( ( declare_attn( - reduction=ProducerReduction.SCOPED, gathers_tp_input=False + reduction=ProducerReduction.EXIT_SCOPED, + gathers_attn_tp_input=False, ), fixture.Norm(), ), (following, fixture.Norm()), terminal=True, ) - for variant, edge in mixer.plan.variants.items(): + for variant, edge in mixer.plan.edges.items(): output = edge.outgoing.produced self.assertEqual( - output.always_leaves, following.kind is StageKind.FFN + output.always_partial, following.kind is StageKind.FFN ) self.assertEqual( - output.leaves_for_next_layer, + output.may_defer_to_next, following.kind is StageKind.ATTENTION, ) - incoming = consumer.plan.variants[variant].incoming.produced + incoming = consumer.plan.edges[variant].incoming.produced self.assertEqual(incoming.group, output.group) diff --git a/test/registered/unit/layer_boundary/test_layernorm_sp.py b/test/registered/unit/layer_boundary/test_layernorm_sp.py index 244cf507d3ce..468687a0c1b4 100644 --- a/test/registered/unit/layer_boundary/test_layernorm_sp.py +++ b/test/registered/unit/layer_boundary/test_layernorm_sp.py @@ -17,8 +17,8 @@ from sglang.srt.layers import layernorm_sp from sglang.srt.layers.layer_boundary import ( Layout, + OutputContract, StageKind, - StageOutput, SumGroup, ) from sglang.srt.layers.layer_boundary import prepare as comm_ops @@ -151,40 +151,40 @@ def test_sp_flag_remains_visible_to_traced_linears(self): def communicator(self, *, first_layer): c = stub_plan() - c._paths[BatchVariant.SEQUENCE_PARALLEL] = sp_region_steps() - c._paths[BatchVariant.INPUT_SCATTERED] = None - c._paths[BatchVariant.CONTEXT_PARALLEL] = None + c.paths[BatchVariant.SEQUENCE_PARALLEL] = sp_region_steps() + c.paths[BatchVariant.INPUT_SCATTERED] = None + c.paths[BatchVariant.CONTEXT_PARALLEL] = None c.enters_stack = first_layer c.is_sparse = False c._attn_input_fusions = () c.norm = _Norm() c.qkv_latent_func = None - c._paths[BatchVariant.ORDINARY] = self.ordinary_steps( + c.paths[BatchVariant.ORDINARY] = self.ordinary_steps( MagicMock(side_effect=AssertionError("moved")) ) return c def ordinary_steps(self, attention_input): """The layer's steps outside the region, which must not run inside it.""" - return comm.StageSteps( - entry=comm.StageEntry( + return comm.StagePath( + entry=comm.EntryPath( prepare=partial( - comm_ops._consumer_step, + comm_ops._run_entry, step=partial( - comm_ops._read_input, - layer_input=None, + comm_ops._update_read, + pre_move=None, enters_stack=False, - read=comm.NORM_QUANT_READ, - update=comm.ADD, + read=comm.NORM_QUANT_READOUT, + update=comm.PLAIN_ADD, ), carried_fusions=(), - adds_plainly=True, + is_plain_add=True, ), input_rows=comm.Layout(frozenset()), input_move=attention_input, - handoff=comm_ops._hand_qkv_hook_its_input, + attn_input_adapter=comm_ops._attn_input_default, ), - output=StageOutput(Layout(frozenset()), group=SumGroup.TP), + output=OutputContract(Layout(frozenset()), group=SumGroup.TP), output_move=MagicMock(side_effect=AssertionError("postprocess ran")), ) @@ -220,7 +220,7 @@ def test_the_first_layer_opens_the_region_on_prefill(self): def test_decode_leaves_the_region_closed(self): communicator = self.communicator(first_layer=True) move = MagicMock(side_effect=lambda **k: k["hidden_states"]) - communicator._paths[BatchVariant.ORDINARY] = self.ordinary_steps(move) + communicator.paths[BatchVariant.ORDINARY] = self.ordinary_steps(move) (h, _), active, scatter = self.run_prepare_attn( communicator, ForwardMode.DECODE, torch.ones(2, 4), None ) @@ -230,7 +230,7 @@ def test_decode_leaves_the_region_closed(self): def test_a_later_layer_does_not_reopen_the_region(self): communicator = self.communicator(first_layer=False) - communicator._paths[BatchVariant.ORDINARY] = self.ordinary_steps( + communicator.paths[BatchVariant.ORDINARY] = self.ordinary_steps( MagicMock(side_effect=lambda **k: k["hidden_states"]) ) _, active, scatter = self.run_prepare_attn( diff --git a/test/registered/unit/layer_boundary/test_longcat_flash_shortcut.py b/test/registered/unit/layer_boundary/test_longcat_flash_shortcut.py index 1cea7bda7c4e..505896abd1cd 100644 --- a/test/registered/unit/layer_boundary/test_longcat_flash_shortcut.py +++ b/test/registered/unit/layer_boundary/test_longcat_flash_shortcut.py @@ -23,7 +23,7 @@ def test_shortcut_does_not_add_the_residual_again(self): tp, rows = 2, 4 # The MoE runs on each attention-TP rank's slice; the last layer hands on # the attention's rows, where the dense branch ends. - local = comm.Layout(frozenset({comm.TokenAxis.ATTN_TP_SCATTER})) + local = comm.Layout(frozenset({comm.TokenAxis.ATTN_TP})) attention = comm.Layout(frozenset()) class Communicator(SimpleNamespace): @@ -35,10 +35,10 @@ class Communicator(SimpleNamespace): moe_communicator = Communicator( prepare_attn=lambda h, r, batch: (h, r), prepare_mlp=lambda h, r, batch: (fork_hidden, fork_residual), - _branch_rows=lambda batch: (local, local, attention), + branch_rows=lambda batch: (local, local, attention), ) dense_communicator = Communicator( - _branch_rows=lambda batch: (attention, attention, attention) + branch_rows=lambda batch: (attention, attention, attention) ) layer = LongcatFlashDecoderLayer.__new__(LongcatFlashDecoderLayer) nn.Module.__init__(layer) @@ -60,7 +60,7 @@ class Communicator(SimpleNamespace): def dense_branch(*args): stream = ResidualStream(torch.full((rows, 3), 11.0)) - hidden = stream.leave(torch.full((rows, 3), 3.0), comm.ADD) + hidden = stream.record(torch.full((rows, 3), 3.0), comm.PLAIN_ADD) batch.residual_stream = stream return hidden, None @@ -82,7 +82,7 @@ def dense_branch(*args): hidden, _ = layer( torch.arange(rows), torch.zeros(rows, 3), batch, None, None ) - hidden, residual = batch.residual_stream.finish(hidden) + hidden, residual = batch.residual_stream.export(hidden) torch.testing.assert_close(hidden, torch.full((rows, 3), 5.0)) torch.testing.assert_close(hidden + residual, torch.full((rows, 3), 16.0)) torch.testing.assert_close(fork_residual, torch.full_like(fork_residual, 5.0)) diff --git a/test/registered/unit/layer_boundary/test_nemotron_h_aux_capture.py b/test/registered/unit/layer_boundary/test_nemotron_h_aux_capture.py index 23bc3b77dba4..11ee2d7511e3 100644 --- a/test/registered/unit/layer_boundary/test_nemotron_h_aux_capture.py +++ b/test/registered/unit/layer_boundary/test_nemotron_h_aux_capture.py @@ -121,8 +121,8 @@ def reduce_in_place(x): moe_dp_size=1, ), patch_communicator("get_moe_cp_size", return_value=1), - patch_communicator("apply_flashinfer_allreduce_fusion", return_value=False), - patch_communicator("apply_aiter_all_reduce_fusion", return_value=False), + patch_communicator("flashinfer_ar_fusion_applies", return_value=False), + patch_communicator("aiter_ar_fusion_applies", return_value=False), ): baseline = _build(pattern, tp, False)( batch.input_ids, torch.arange(2), batch, inputs_embeds=inputs.clone() diff --git a/test/registered/unit/layer_boundary/test_nemotron_h_mtp_reduction.py b/test/registered/unit/layer_boundary/test_nemotron_h_mtp_reduction.py index b53704c03dd6..0ff3ad10b5ff 100644 --- a/test/registered/unit/layer_boundary/test_nemotron_h_mtp_reduction.py +++ b/test/registered/unit/layer_boundary/test_nemotron_h_mtp_reduction.py @@ -6,7 +6,7 @@ from torch import nn from sglang.srt.layers.layer_boundary.layout import SumGroup -from sglang.srt.layers.layer_boundary.residual.add_norm import ADD +from sglang.srt.layers.layer_boundary.residual.add_norm import PLAIN_ADD from sglang.srt.layers.layer_boundary.residual.stream import ResidualStream from sglang.srt.model_executor.forward_batch_info import ForwardMode from sglang.srt.models import nemotron_h_mtp @@ -58,11 +58,9 @@ def test_attention_partial_is_reduced_once(self): ), patch_communicator("get_moe_cp_size", return_value=1), patch_communicator( - "apply_flashinfer_allreduce_fusion", return_value=False - ), - patch_communicator( - "apply_aiter_all_reduce_fusion", return_value=False + "flashinfer_ar_fusion_applies", return_value=False ), + patch_communicator("aiter_ar_fusion_applies", return_value=False), ): layer = nemotron_h_mtp.NemotronHMTPMoEDecoderLayer.__new__( nemotron_h_mtp.NemotronHMTPMoEDecoderLayer @@ -79,8 +77,10 @@ def test_attention_partial_is_reduced_once(self): residual = torch.tensor([[7.0, 3.0], [5.0, 9.0]]) expected = partial * tp + residual stream = ResidualStream(residual) - hidden = stream.leave( - partial, ADD, declared_sum=SumGroup.ATTN_TP if tp > 1 else None + hidden = stream.record( + partial, + PLAIN_ADD, + declared_sum=SumGroup.ATTN_TP if tp > 1 else None, ) batch = SimpleNamespace( forward_mode=ForwardMode.DECODE, diff --git a/test/registered/unit/layer_boundary/test_nemotron_h_stages.py b/test/registered/unit/layer_boundary/test_nemotron_h_stages.py index 4bd676db3be6..ca8ef7bfc227 100644 --- a/test/registered/unit/layer_boundary/test_nemotron_h_stages.py +++ b/test/registered/unit/layer_boundary/test_nemotron_h_stages.py @@ -9,7 +9,7 @@ from sglang.srt.layers.layer_boundary import ( Layout, MixerExit, - StageOutput, + OutputContract, SumGroup, TokenAxis, UnreducedOutput, @@ -54,7 +54,7 @@ def sizes(*, dp=1, tp=1): return { TokenAxis.ATTN_DP: dp, TokenAxis.ATTN_CP: 1, - TokenAxis.ATTN_TP_SCATTER: tp, + TokenAxis.ATTN_TP: tp, } @@ -86,8 +86,8 @@ def stages(pattern, *, dp=1, tp=1, a2a=False): patch_communicator("get_parallel", return_value=parallel), patch.object(utils, "get_parallel", return_value=parallel), patch_communicator("is_moe_input_scattered_across_dp_ranks", return_value=a2a), - patch_communicator("enable_moe_dense_fully_dp", return_value=False), - patch_communicator("_generic_prefill_cp_shards_tokens", return_value=False), + patch_communicator("is_dense_ffn_fully_dp", return_value=False), + patch_communicator("_prefill_cp_shards_tokens", return_value=False), ): return [layer_stage(pattern, i) for i in range(len(pattern))] @@ -114,27 +114,27 @@ def test_adjacent_layers_agree_on_their_shared_boundary(self): self.assertEqual(into.residual, rows) self.assertEqual(into.produced.layout, rows) produced = out_of.produced - may_leave = produced.always_leaves or produced.leaves_for_next_layer + may_leave = produced.always_partial or produced.may_defer_to_next self.assertEqual( into.produced.group, produced.group if may_leave else None ) self.assertEqual( - into.produced.always_leaves, produced.always_leaves + into.produced.always_partial, produced.always_partial ) self.assertEqual( - into.produced.leaves_for_next_layer, - produced.leaves_for_next_layer, + into.produced.may_defer_to_next, + produced.may_defer_to_next, ) self.assertEqual( - layers[0].edges[0].produced, StageOutput(rows, update=None) + layers[0].edges[0].produced, OutputContract(rows, update=None) ) self.assertTrue(layers[0].enters_stack) self.assertFalse(any(layer.enters_stack for layer in layers[1:])) last = layers[-1].edges[1].produced - self.assertFalse(last.always_leaves or last.leaves_for_next_layer) + self.assertFalse(last.always_partial or last.may_defer_to_next) def test_what_each_kind_of_boundary_carries(self): - # (pattern, boundary after layer 0): group, always_leaves, leaves_for_next_layer + # (pattern, boundary after layer 0): group, always_partial, may_defer_to_next cases = { "M-": (SumGroup.ATTN_TP, True, False), "*E": (SumGroup.ATTN_TP, True, False), @@ -148,19 +148,19 @@ def test_what_each_kind_of_boundary_carries(self): with self.subTest(pattern=pattern): into = stages(pattern, tp=2)[1].edges[0].produced self.assertEqual( - (into.group, into.always_leaves, into.leaves_for_next_layer), + (into.group, into.always_partial, into.may_defer_to_next), expected, ) # Without attention TP a mixer's output is complete. into = stages("M-", tp=1)[1].edges[0].produced - self.assertEqual(into, StageOutput(into.layout, update=None)) + self.assertEqual(into, OutputContract(into.layout, update=None)) # A MoE on this rank's own rows hands on a complete output; an a2a # backend dispatches only the MoE, so an MLP still sums over TP. into = stages("EM", tp=2, a2a=True)[1].edges[0].produced - self.assertEqual(into, StageOutput(into.layout, update=None)) + self.assertEqual(into, OutputContract(into.layout, update=None)) into = stages("-M", tp=2, a2a=True)[1].edges[0].produced self.assertEqual( - (into.group, into.always_leaves, into.leaves_for_next_layer), + (into.group, into.always_partial, into.may_defer_to_next), (SumGroup.TP, False, True), ) @@ -177,19 +177,19 @@ def test_decision_table(self): if always and may: continue with self.subTest( - always_leaves=always, leaves_for_next_layer=may, movable=movable + always_partial=always, may_defer_to_next=may, movable=movable ): - produced = StageOutput( + produced = OutputContract( attention, group=SumGroup.ATTN_TP if always or may else None, - always_leaves=always, - leaves_for_next_layer=may, + always_partial=always, + may_defer_to_next=may, ) communicator = SimpleNamespace( plan=SimpleNamespace( - _batch_steps=lambda batch: SimpleNamespace(output=produced), + path_for=lambda batch: SimpleNamespace(output=produced), ), - _ffn_sum_can_move_to_next_layer=MagicMock(return_value=movable), + _sum_deferral_allowed=MagicMock(return_value=movable), ) hidden = torch.ones(2, 4) with get_parallel().override(tp_group=tp_group, tp_size=2): diff --git a/test/registered/unit/layer_boundary/test_prepare_attn_steps.py b/test/registered/unit/layer_boundary/test_prepare_attn_steps.py index de9d3d59eaa8..8ff489fa88dd 100644 --- a/test/registered/unit/layer_boundary/test_prepare_attn_steps.py +++ b/test/registered/unit/layer_boundary/test_prepare_attn_steps.py @@ -14,8 +14,8 @@ from sglang.srt.layers.layer_boundary import StageKind from sglang.srt.layers.layer_boundary import prepare as comm_ops from sglang.srt.layers.layer_boundary.contracts import BatchVariant -from sglang.srt.layers.layer_boundary.fusions.allreduce import attention_fusions -from sglang.srt.layers.layer_boundary.ops import identity_output +from sglang.srt.layers.layer_boundary.fusions.allreduce import attn_input_fusions +from sglang.srt.layers.layer_boundary.ops import keep_output from sglang.test.boundary_fixtures import prepare_input, stub_plan, stub_stage from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.communicator_patch import patch_communicator @@ -53,33 +53,33 @@ def communicator(norm): c.norm = norm # A layer inside the stack: an absent residual was written back before. c.enters_stack = False - c._paths[BatchVariant.SEQUENCE_PARALLEL] = None - c._paths[BatchVariant.INPUT_SCATTERED] = None - c._paths[BatchVariant.CONTEXT_PARALLEL] = None + c.paths[BatchVariant.SEQUENCE_PARALLEL] = None + c.paths[BatchVariant.INPUT_SCATTERED] = None + c.paths[BatchVariant.CONTEXT_PARALLEL] = None c.qkv_latent_func = None # Construction picks the fused entries; call it under the platform patches. - c._attn_input_fusions = attention_fusions(c) + c._attn_input_fusions = attn_input_fusions(c) # A layer whose attention takes its input as it is and owes nothing on it. - c._paths[BatchVariant.ORDINARY] = comm.StageSteps( - entry=comm.StageEntry( + c.paths[BatchVariant.ORDINARY] = comm.StagePath( + entry=comm.EntryPath( prepare=partial( - comm_ops._consumer_step, - adds_plainly=True, + comm_ops._run_entry, + is_plain_add=True, step=partial( - comm_ops._read_input, - layer_input=None, + comm_ops._update_read, + pre_move=None, enters_stack=False, - read=comm.NORM_QUANT_READ, - update=comm.ADD, + read=comm.NORM_QUANT_READOUT, + update=comm.PLAIN_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, + attn_input_adapter=comm_ops._attn_input_default, ), - output=comm.StageOutput(comm.Layout(frozenset())), - output_move=identity_output, + output=comm.OutputContract(comm.Layout(frozenset())), + output_move=keep_output, ) return c @@ -102,10 +102,10 @@ def platform(*, use_aiter=False, gfx95=False, fusion=False, kernel_group=True): patch_communicator("_use_aiter_bpreshuffle_gfx95", False, create=True) ) stack.enter_context( - patch_communicator("apply_aiter_all_reduce_fusion", return_value=False) + patch_communicator("aiter_ar_fusion_applies", return_value=False) ) stack.enter_context( - patch_communicator("apply_flashinfer_allreduce_fusion", return_value=fusion) + patch_communicator("flashinfer_ar_fusion_applies", return_value=fusion) ) # The group the fused kernel reduces over. stack.enter_context( @@ -248,7 +248,7 @@ def test_an_owed_reduce_scatter_reaches_local_tokens_before_the_norm(self): step = MagicMock(return_value=local) hidden_states, residual = prepare_input( stub_stage(communicator(norm), StageKind.ATTENTION), - comm.UnreducedOutput(partial, reduce_and_redistribute=step), + comm.UnreducedOutput(partial, reduce_to_dp_local=step), torch.zeros(local_rows, 4), None, ) @@ -265,19 +265,19 @@ 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._paths[BatchVariant.ORDINARY] = msgspec.structs.replace( - c._paths.get(BatchVariant.ORDINARY), + c.paths[BatchVariant.ORDINARY] = msgspec.structs.replace( + c.paths.get(BatchVariant.ORDINARY), entry=msgspec.structs.replace( - c._paths.get(BatchVariant.ORDINARY).entry, + c.paths.get(BatchVariant.ORDINARY).entry, prepare=partial( - comm_ops._consumer_step, - adds_plainly=True, + comm_ops._run_entry, + is_plain_add=True, step=partial( - comm_ops._read_input, - layer_input=None, + comm_ops._update_read, + pre_move=None, enters_stack=False, - read=comm.NORM_QUANT_READ, - update=comm.ADD, + read=comm.NORM_QUANT_READOUT, + update=comm.PLAIN_ADD, ), carried_fusions=(takes_anything,), ), @@ -287,7 +287,7 @@ def test_an_owed_reduce_scatter_is_not_offered_to_a_fused_kernel(self): step = MagicMock(return_value=torch.full((1, 4), 7.0)) prepare_input( stub_stage(c, StageKind.ATTENTION), - comm.UnreducedOutput(partial_sum, reduce_and_redistribute=step), + comm.UnreducedOutput(partial_sum, reduce_to_dp_local=step), torch.zeros(1, 4), None, ) @@ -300,7 +300,7 @@ class TestFusedReadForms(CustomTestCase): def test_consumer_form_and_backend_policy_are_both_required(self): from sglang.srt.layers.layer_boundary.residual.add_norm import ( Fp8Input, - NormQuantRead, + NormQuantReadout, ) for form in (None, Fp8Input.TUPLE, Fp8Input.TUPLE_AND_BF16): @@ -324,8 +324,8 @@ def test_consumer_form_and_backend_policy_are_both_required(self): ), ): plan = SimpleNamespace(norm=MagicMock(), fusions=None) - (candidate,) = attention_fusions( - plan, NormQuantRead(fp8_input=form) + (candidate,) = attn_input_fusions( + plan, NormQuantReadout(fp8_input=form) ) expected = form is not None and enabled and not disabled_by_env self.assertEqual(candidate.keywords["fuses_quant"], expected) diff --git a/test/registered/unit/layer_boundary/test_reduction_fusion.py b/test/registered/unit/layer_boundary/test_reduction_fusion.py index 8800b930f525..24bf97e0d90f 100644 --- a/test/registered/unit/layer_boundary/test_reduction_fusion.py +++ b/test/registered/unit/layer_boundary/test_reduction_fusion.py @@ -9,14 +9,14 @@ from sglang.srt.layers import layer_boundary as comm from sglang.srt.layers.layer_boundary import ( Layout, - StageOutput, + OutputContract, SumGroup, ) from sglang.srt.layers.layer_boundary import ops as transport_ops from sglang.srt.layers.layer_boundary import prepare as comm_ops from sglang.srt.layers.layer_boundary.contracts import BatchVariant -from sglang.srt.layers.layer_boundary.exit import OutputBoundary -from sglang.srt.layers.layer_boundary.ops import identity_output +from sglang.srt.layers.layer_boundary.exit import ExitPolicy +from sglang.srt.layers.layer_boundary.ops import keep_output from sglang.srt.layers.moe import ( can_merge_post_experts_all_reduce, deferred_post_experts_all_reduce, @@ -34,29 +34,29 @@ def _steps(*, output=None, returns_over_dp=False, ffn_sum_is_movable=True): - return comm.StageSteps( - entry=comm.StageEntry( + return comm.StagePath( + entry=comm.EntryPath( prepare=partial( - comm_ops._consumer_step, + comm_ops._run_entry, step=partial( - comm_ops._read_input, - layer_input=None, + comm_ops._update_read, + pre_move=None, enters_stack=False, - read=comm.NORM_QUANT_READ, - update=comm.ADD, + read=comm.NORM_QUANT_READOUT, + update=comm.PLAIN_ADD, ), carried_fusions=(), ), input_rows=comm.Layout(frozenset()), input_move=identity_input, - handoff=comm_ops._hand_qkv_hook_its_input, + attn_input_adapter=comm_ops._attn_input_default, ), output=output - or StageOutput( + or OutputContract( Layout(frozenset()), group=SumGroup.TP if ffn_sum_is_movable else None ), returns_over_dp=returns_over_dp, - output_move=None if returns_over_dp else identity_output, + output_move=None if returns_over_dp else keep_output, ) @@ -64,12 +64,12 @@ def _fake_communicator(ffn_sum_is_movable=True): communicator = stub_plan() communicator._speculative_algo = None # An FFN gathered over MoE-CP or on each rank's own rows has no movable sum. - communicator._paths[BatchVariant.ORDINARY] = _steps( + communicator.paths[BatchVariant.ORDINARY] = _steps( ffn_sum_is_movable=ffn_sum_is_movable ) - communicator._paths[BatchVariant.SEQUENCE_PARALLEL] = None - communicator._paths[BatchVariant.INPUT_SCATTERED] = None - communicator._paths[BatchVariant.CONTEXT_PARALLEL] = None + communicator.paths[BatchVariant.SEQUENCE_PARALLEL] = None + communicator.paths[BatchVariant.INPUT_SCATTERED] = None + communicator.paths[BatchVariant.CONTEXT_PARALLEL] = None communicator.terminal = False return communicator @@ -298,7 +298,7 @@ def _should_fuse( ) with ( patch_communicator("is_enable_moe_cp_allgather", return_value=False), - patch_communicator("apply_flashinfer_allreduce_fusion", return_value=True), + patch_communicator("flashinfer_ar_fusion_applies", return_value=True), patch_communicator( "get_attn_tp_context", return_value=types.SimpleNamespace(input_scattered=False), @@ -310,12 +310,12 @@ def _should_fuse( tp_size=moe_ep_size * moe_tp_size * moe_dp_size, ), ): - return OutputBoundary( + return ExitPolicy( _fake_communicator(ffn_sum_is_movable) - )._ffn_sum_can_move_to_next_layer( - OutputBoundary( - _fake_communicator(ffn_sum_is_movable) - ).plan._batch_steps(forward_batch) + )._sum_deferral_allowed( + ExitPolicy(_fake_communicator(ffn_sum_is_movable)).plan.path_for( + forward_batch + ) ) def test_hybrid_ep_tp_fuses_when_mergeable(self): @@ -334,7 +334,7 @@ def test_pure_ep_still_fuses(self): self.assertTrue(self._should_fuse(moe_ep_size=4, moe_tp_size=1)) def test_a_layer_gathered_over_moe_cp_does_not_fuse(self): - # Fusion skips postprocess_layer, which holds the CP scatter; a dense + # Fusion skips finish_complete_output, which holds the CP scatter; a dense # layer gathered over the MoE-CP group (moe_dp_size == attn_cp_size) is # not caught by the is_enable_moe_cp_allgather gate. self.assertFalse( @@ -369,17 +369,17 @@ def _should_defer( tp_group_object = object() communicator = _fake_communicator() communicator.terminal = is_last_layer - communicator._paths[BatchVariant.SEQUENCE_PARALLEL] = ( + communicator.paths[BatchVariant.SEQUENCE_PARALLEL] = ( sp_region_steps() if sp_active else None ) - communicator._paths[BatchVariant.INPUT_SCATTERED] = None - communicator._paths[BatchVariant.CONTEXT_PARALLEL] = None - communicator._paths[BatchVariant.ORDINARY] = _steps( - output=StageOutput( + communicator.paths[BatchVariant.INPUT_SCATTERED] = None + communicator.paths[BatchVariant.CONTEXT_PARALLEL] = None + communicator.paths[BatchVariant.ORDINARY] = _steps( + output=OutputContract( Layout(frozenset()), group=SumGroup.MOE_OUTPUT, - leaves_for_reduce_scatter=True, - leaves_for_reduce_scatterv=True, + may_reduce_scatter=True, + may_reduce_scatterv=True, ), returns_over_dp=scatters_to_local_tokens, ) @@ -390,12 +390,10 @@ def _should_defer( residual_stream=types.SimpleNamespace(residual=torch.ones(batch_size, 4)), ) with ( - patch_communicator( - "_reduce_and_redistribute_output_step", return_value=step - ), + patch_communicator("_select_dp_reduce_scatter", return_value=step), get_forward().scoped(sp_active=sp_active), patch_communicator("is_enable_moe_cp_allgather", return_value=False), - patch_communicator("apply_flashinfer_allreduce_fusion", return_value=fused), + patch_communicator("flashinfer_ar_fusion_applies", return_value=fused), patch_communicator("_use_aiter", False), patch_communicator( "post_experts_reduction_group", @@ -471,9 +469,9 @@ def _should_defer( return_value=tp_group_object if tp_group else object(), ), ): - return communicator.output._ffn_sum_moves_to_next_layer( + return communicator.output._defers_sum( forward_batch, - communicator.output.plan._batch_steps(forward_batch), + communicator.output.plan.path_for(forward_batch), mlp_reduce_scatter=reduce_scatter, dp_step=step, ) @@ -517,11 +515,11 @@ def test_keeps_what_the_next_layer_input_cannot_run(self): for name, condition in ( ( "reduce-scatter", - dict(step=transport_ops._reduce_and_redistribute_output_varlen), + dict(step=transport_ops.dp_reduce_scatterv), ), ( "MAX_LEN reduce-scatter", - dict(step=transport_ops._reduce_and_redistribute_output_max_len), + dict(step=transport_ops.dp_reduce_scatter), ), ("other postprocess", dict(scatters_to_local_tokens=False)), ("LayerNorm SP", dict(sp_active=True)), diff --git a/test/registered/unit/layer_boundary/test_residual_stream.py b/test/registered/unit/layer_boundary/test_residual_stream.py index 0cf4e556bff6..f2941cefc5d4 100644 --- a/test/registered/unit/layer_boundary/test_residual_stream.py +++ b/test/registered/unit/layer_boundary/test_residual_stream.py @@ -7,19 +7,19 @@ from sglang.srt.layers.aux_hidden_states import AuxHiddenStateList from sglang.srt.layers.layer_boundary import ( - EdgeDecl, + EdgeContract, + EntryPath, + InputContract, Layout, - StageEntry, - StageInput, - StageOutput, - StageSteps, + OutputContract, + StagePath, SumGroup, - make_boundary, + bind_entry, ) from sglang.srt.layers.layer_boundary.contracts import BatchVariant, StageKind from sglang.srt.layers.layer_boundary.output import UnreducedOutput from sglang.srt.layers.layer_boundary.residual.access import add_to_output -from sglang.srt.layers.layer_boundary.residual.add_norm import ADD +from sglang.srt.layers.layer_boundary.residual.add_norm import PLAIN_ADD from sglang.srt.layers.layer_boundary.residual.stream import ( DeclaredSum, OwedOutput, @@ -40,16 +40,16 @@ def setUp(self): self.residual = torch.full((2, 4), 3.0) self.partial = torch.ones(2, 4) self.stream = ResidualStream(self.residual) - self.hidden = self.stream.leave( - UnreducedOutput(self.partial, group=self.group), ADD + self.hidden = self.stream.record( + UnreducedOutput(self.partial, group=self.group), PLAIN_ADD ) def test_written_input_does_not_repeat_first_stage_initialization(self): class Read: - norms_plainly = False + is_plain_norm = False enters = 0 - def enter(self, value): + def init_residual(self, value): self.enters += 1 return value * 3 @@ -58,11 +58,13 @@ def read(self, value, norm, quant_format=""): read = Read() rows = Layout(frozenset()) - boundary = make_boundary( - EdgeDecl(StageOutput(rows), StageInput(rows, read=read), rows, rows), + boundary = bind_entry( + EdgeContract( + OutputContract(rows), InputContract(rows, read=read), rows, rows + ), enters_stack=True, ) - steps = StageSteps(StageEntry(boundary.prepare, rows), StageOutput(rows), None) + steps = StagePath(EntryPath(boundary.prepare, rows), OutputContract(rows), None) stage = StageBoundary( SimpleNamespace(norm=None), declaration=SimpleNamespace(kind=StageKind.ATTENTION), @@ -77,8 +79,8 @@ def read(self, value, norm, quant_format=""): def test_consumer_executes_the_carried_update_not_the_bound_producers(self): class Update: - adds_plainly = False - at_producer = False + is_plain_add = False + applied_at_exit = False def __init__(self, increment): self.increment = increment @@ -89,7 +91,7 @@ def update(self, value, residual): return value + residual + self.increment class Read: - norms_plainly = False + is_plain_norm = False def update_and_read(self, update, value, residual, norm, **kwargs): residual = update.update(value, residual) @@ -97,15 +99,15 @@ def update_and_read(self, update, value, residual, norm, **kwargs): declared = Update(100) rows = Layout(frozenset()) - boundary = make_boundary( - EdgeDecl( - StageOutput(rows, update=declared), - StageInput(rows, read=Read()), + boundary = bind_entry( + EdgeContract( + OutputContract(rows, update=declared), + InputContract(rows, read=Read()), rows, rows, ) ) - steps = StageSteps(StageEntry(boundary.prepare, rows), StageOutput(rows), None) + steps = StagePath(EntryPath(boundary.prepare, rows), OutputContract(rows), None) stage = StageBoundary( SimpleNamespace(norm=None), declaration=SimpleNamespace(kind=StageKind.ATTENTION), @@ -114,7 +116,7 @@ def update_and_read(self, update, value, residual, norm, **kwargs): with self.subTest(increment=increment): actual = Update(increment) stream = ResidualStream(torch.full((2, 4), 3.0)) - hidden = stream.leave(torch.ones(2, 4), actual) + hidden = stream.record(torch.ones(2, 4), actual) result, stream = stage._prepare(hidden, stream, None, steps) torch.testing.assert_close( result, torch.full((2, 4), 2.0 * (4 + increment)) @@ -125,8 +127,8 @@ def update_and_read(self, update, value, residual, norm, **kwargs): def test_unbound_update_capability_fails_before_consuming_the_sum(self): rows = Layout(frozenset()) - boundary = make_boundary( - EdgeDecl(StageOutput(rows), StageInput(rows), rows, rows) + boundary = bind_entry( + EdgeContract(OutputContract(rows), InputContract(rows), rows, rows) ) with self.assertRaisesRegex(RuntimeError, "capability"): boundary.prepare( @@ -134,12 +136,12 @@ def test_unbound_update_capability_fails_before_consuming_the_sum(self): self.residual, None, None, - update=SimpleNamespace(adds_plainly=False), + update=SimpleNamespace(is_plain_add=False), ) def test_declared_sum_snapshot_and_pipeline_handoff_preserve_the_partial(self): stream = ResidualStream(self.residual) - hidden = stream.leave(self.partial, ADD, declared_sum=SumGroup.TP) + hidden = stream.record(self.partial, PLAIN_ADD, declared_sum=SumGroup.TP) self.assertIsInstance(hidden, OwedOutput) with patch( "sglang.srt.layers.layer_boundary.residual.stream._sum_group", @@ -148,27 +150,27 @@ def test_declared_sum_snapshot_and_pipeline_handoff_preserve_the_partial(self): snapshot = stream.snapshot(hidden) torch.testing.assert_close(snapshot, torch.full((2, 4), 5.0)) torch.testing.assert_close(self.partial, torch.ones(2, 4)) - wire, residual = stream.finish(hidden, preserve_declared=True) + wire, residual = stream.export(hidden, preserve_declared=True) self.assertIs(wire, self.partial) self.assertIsInstance(stream.pending.owed, DeclaredSum) - received, rebuilt = ResidualStream.arrive( - wire, residual, ADD, declared_sum=SumGroup.TP + received, rebuilt = ResidualStream.from_handoff( + wire, residual, PLAIN_ADD, declared_sum=SumGroup.TP ) self.assertIsInstance(received, OwedOutput) self.assertIs(rebuilt.pending.owed.group, SumGroup.TP) def test_materialized_declared_sum_is_not_reduced_again(self): rows = Layout(frozenset()) - boundary = make_boundary( - EdgeDecl( - StageOutput(rows, group=SumGroup.TP, always_leaves=True), - StageInput(rows), + boundary = bind_entry( + EdgeContract( + OutputContract(rows, group=SumGroup.TP, always_partial=True), + InputContract(rows), rows, rows, ) ) stream = ResidualStream(self.residual) - hidden = stream.leave(self.partial, ADD, declared_sum=SumGroup.TP) + hidden = stream.record(self.partial, PLAIN_ADD, declared_sum=SumGroup.TP) with patch( "sglang.srt.layers.layer_boundary.residual.stream._sum_group", return_value=self.group, @@ -191,11 +193,11 @@ def norm(value, residual): def test_declared_sum_must_match_the_consumer_contract(self): rows = Layout(frozenset()) - boundary = make_boundary( - EdgeDecl(StageOutput(rows), StageInput(rows), rows, rows) + boundary = bind_entry( + EdgeContract(OutputContract(rows), InputContract(rows), rows, rows) ) stream = ResidualStream(self.residual) - stream.leave(self.partial, ADD, declared_sum=SumGroup.TP) + stream.record(self.partial, PLAIN_ADD, declared_sum=SumGroup.TP) with self.assertRaisesRegex(RuntimeError, "sum does not match"): boundary.prepare( self.partial, @@ -203,13 +205,13 @@ def test_declared_sum_must_match_the_consumer_contract(self): None, None, pending=stream.pending, - update=ADD, + update=PLAIN_ADD, ) def test_snapshot_rejects_stateful_update_without_touching_main_state(self): - update = SimpleNamespace(adds_plainly=False, update=Mock()) + update = SimpleNamespace(is_plain_add=False, update=Mock()) stream = ResidualStream(self.residual) - hidden = stream.leave(UnreducedOutput(self.partial, group=self.group), update) + hidden = stream.record(UnreducedOutput(self.partial, group=self.group), update) with self.assertRaisesRegex(NotImplementedError, "plain residual update"): stream.snapshot(hidden) self.group.all_reduce.assert_not_called() @@ -234,7 +236,7 @@ def test_main_capture_completes_once_and_leaves_update_pending(self): self.assertIs(self.stream.complete(hidden), hidden) self.group.all_reduce.assert_called_once() self.assertIs(self.stream.pending.value, hidden) - self.assertIs(self.stream.pending.update, ADD) + self.assertIs(self.stream.pending.update, PLAIN_ADD) self.assertIsNone(self.stream.pending.owed) self.assertIs(self.stream.residual, self.residual) torch.testing.assert_close(hidden, torch.full((2, 4), 2.0)) @@ -274,7 +276,7 @@ def capture(value, *, owned=False): self.hidden, self.stream, None, - capture_output=capture, + capture=capture, ) self.assertEqual(events, ["add_norm", "capture"]) self.group.all_reduce.assert_called_once() @@ -311,7 +313,7 @@ def prepare(hidden, stream, batch, **kwargs): self.stream, None, post_residual_addition=extra, - capture_output=outputs.capture, + capture=outputs.capture, ) torch.testing.assert_close(outputs[0], torch.full((2, 4), 5.0)) torch.testing.assert_close(output, torch.full((2, 4), 36.0)) @@ -323,14 +325,14 @@ def test_capture_restores_the_producers_rows_after_a_scattered_read(self): from sglang.srt.layers.layer_boundary import TokenAxis full = Layout(frozenset()) - local = Layout(frozenset({TokenAxis.ATTN_TP_SCATTER})) - boundary = make_boundary( - EdgeDecl(StageOutput(full), StageInput(local), full, local) + local = Layout(frozenset({TokenAxis.ATTN_TP})) + boundary = bind_entry( + EdgeContract(OutputContract(full), InputContract(local), full, local) ) shard = torch.full((2, 4), 3.0) expected = torch.cat([shard, torch.full_like(shard, 7.0)]) with patch( - "sglang.srt.layers.layer_boundary.ops._redistribute_from_attn_tp_shards", + "sglang.srt.layers.layer_boundary.ops.attn_tp_gather", return_value=expected, ) as gather: captured = boundary.capture_move(shard, forward_batch=None) @@ -339,13 +341,13 @@ def test_capture_restores_the_producers_rows_after_a_scattered_read(self): torch.testing.assert_close(shard, torch.full((2, 4), 3.0)) def test_final_capture_uses_the_norms_residual_result(self): - from sglang.srt.layers.layer_boundary.residual.access import norm_output + from sglang.srt.layers.layer_boundary.residual.access import final_norm_pair outputs = AuxHiddenStateList() updated = torch.full_like(self.partial, 5.0) normalized = torch.full_like(self.partial, 13.0) norm = Mock(return_value=(normalized, updated)) - result = norm_output(self.partial, self.residual, norm, outputs.capture) + result = final_norm_pair(self.partial, self.residual, norm, outputs.capture) self.assertIs(result, normalized) norm.assert_called_once_with(self.partial, self.residual) updated.zero_() @@ -363,11 +365,11 @@ def test_deepstack_adds_once_after_completing_the_sum(self): def test_handles_cannot_cross_microbatches_or_be_reused_after_take(self): other = ResidualStream(self.residual.clone()) - other.leave(UnreducedOutput(self.partial.clone(), group=self.group), ADD) + other.record(UnreducedOutput(self.partial.clone(), group=self.group), PLAIN_ADD) with self.assertRaises(RuntimeError): other.input(self.hidden) with self.assertRaises(RuntimeError): - self.stream.leave(torch.zeros_like(self.partial), ADD) + self.stream.record(torch.zeros_like(self.partial), PLAIN_ADD) self.stream.write(self.residual) with self.assertRaises(RuntimeError): self.stream.input(self.hidden) @@ -376,8 +378,8 @@ def test_handles_cannot_cross_microbatches_or_be_reused_after_take(self): def test_local_row_move_is_owed_even_without_an_all_reduce(self): move = Mock(side_effect=lambda value: value[:1]) stream = ResidualStream(self.residual[:1]) - hidden = stream.leave( - UnreducedOutput(self.partial, reduce_and_redistribute=move), ADD + hidden = stream.record( + UnreducedOutput(self.partial, reduce_to_dp_local=move), PLAIN_ADD ) self.assertIsInstance(hidden, OwedOutput) completed = stream.complete(hidden) @@ -387,17 +389,19 @@ def test_local_row_move_is_owed_even_without_an_all_reduce(self): def test_prepare_releases_the_consumed_contribution(self): class Read: - norms_plainly = True + is_plain_norm = True def update_and_read(self, update, value, residual, norm, **kwargs): residual = update.update(value, residual) return residual * 2, residual rows = Layout(frozenset()) - boundary = make_boundary( - EdgeDecl(StageOutput(rows), StageInput(rows, read=Read()), rows, rows) + boundary = bind_entry( + EdgeContract( + OutputContract(rows), InputContract(rows, read=Read()), rows, rows + ) ) - steps = StageSteps(StageEntry(boundary.prepare, rows), StageOutput(rows), None) + steps = StagePath(EntryPath(boundary.prepare, rows), OutputContract(rows), None) stage = StageBoundary( SimpleNamespace(norm=None), declaration=SimpleNamespace(kind=StageKind.ATTENTION), @@ -406,9 +410,9 @@ def update_and_read(self, update, value, residual, norm, **kwargs): partial_ref = weakref.ref(partial) stream = ResidualStream(torch.full((2, 4), 3.0)) # The deferred DP completion: sum, then this rank's rows. - hidden = stream.leave( - UnreducedOutput(partial, reduce_and_redistribute=lambda x: x[:2] * 2), - ADD, + hidden = stream.record( + UnreducedOutput(partial, reduce_to_dp_local=lambda x: x[:2] * 2), + PLAIN_ADD, ) del partial result, stream = stage._prepare(hidden, stream, None, steps) @@ -429,9 +433,9 @@ def test_terminal_norm_releases_layer_buffers(self): residual = torch.ones(2, 4) contribution = torch.full_like(residual, 2) refs = [weakref.ref(residual), weakref.ref(contribution)] - batch.current(fb).write(residual) - output = batch.current(fb).leave(contribution, ADD) - result = batch.norm(output, fb, lambda x, r: (x + r, r)) + batch.stream_of(fb).write(residual) + output = batch.stream_of(fb).record(contribution, PLAIN_ADD) + result = batch.final_norm(output, fb, lambda x, r: (x + r, r)) del output, residual, contribution torch.testing.assert_close(result, torch.full((2, 4), 3.0)) self.assertTrue(all(ref() is None for ref in refs)) @@ -445,8 +449,8 @@ def test_pp_export_transfers_buffer_ownership(self): residual = torch.ones(2, 4) contribution = torch.full_like(residual, 2) refs = [weakref.ref(residual), weakref.ref(contribution)] - batch.current(fb).write(residual) - output = batch.current(fb).leave(contribution, ADD) + batch.stream_of(fb).write(residual) + output = batch.stream_of(fb).record(contribution, PLAIN_ADD) proxy = batch.to_pp(output, fb) del output, residual, contribution self.assertTrue(all(ref() is not None for ref in refs)) @@ -458,9 +462,9 @@ def test_interleaved_batches_keep_independent_streams(self): from sglang.srt.layers.layer_boundary.residual import batch class Read: - norms_plainly = False + is_plain_norm = False - def enter(self, value): + def init_residual(self, value): return value def read(self, value, norm, quant_format="", **kwargs): @@ -471,31 +475,33 @@ def update_and_read(self, update, value, residual, norm, **kwargs): return updated * 2, updated rows = Layout(frozenset()) - boundary = make_boundary( - EdgeDecl(StageOutput(rows), StageInput(rows, read=Read()), rows, rows) + boundary = bind_entry( + EdgeContract( + OutputContract(rows), InputContract(rows, read=Read()), rows, rows + ) ) owner = stub_plan() owner.norm = None - owner._paths[BatchVariant.ORDINARY] = StageSteps( - StageEntry(boundary.prepare, rows), StageOutput(rows), None + owner.paths[BatchVariant.ORDINARY] = StagePath( + EntryPath(boundary.prepare, rows), OutputContract(rows), None ) - owner._paths[BatchVariant.SEQUENCE_PARALLEL] = owner._paths[ + owner.paths[BatchVariant.SEQUENCE_PARALLEL] = owner.paths[ BatchVariant.INPUT_SCATTERED - ] = owner._paths[BatchVariant.CONTEXT_PARALLEL] = None + ] = owner.paths[BatchVariant.CONTEXT_PARALLEL] = None stage = stub_stage(owner, StageKind.FFN) a, b = [SimpleNamespace(forward_mode=ForwardMode.DECODE) for _ in range(2)] batch.start(a) batch.start(b) first = stage.prepare(torch.ones(2, 4), a) second = stage.prepare(torch.full((2, 4), 10.0), b) - first = batch.current(a).leave(first * 3, ADD) - second = batch.current(b).leave(second * 5, ADD) + first = batch.stream_of(a).record(first * 3, PLAIN_ADD) + second = batch.stream_of(b).record(second * 5, PLAIN_ADD) torch.testing.assert_close(stage.prepare(first, a), torch.full((2, 4), 14.0)) torch.testing.assert_close(stage.prepare(second, b), torch.full((2, 4), 220.0)) - self.assertIsNot(batch.current(a), batch.current(b)) + self.assertIsNot(batch.stream_of(a), batch.stream_of(b)) batch.start(a) - self.assertIsNone(batch.current(a).pending) - self.assertIsNotNone(batch.current(b).residual) + self.assertIsNone(batch.stream_of(a).pending) + self.assertIsNotNone(batch.stream_of(b).residual) if __name__ == "__main__": diff --git a/test/registered/unit/layer_boundary/test_step3p5_dense_reduce_scatter.py b/test/registered/unit/layer_boundary/test_step3p5_dense_reduce_scatter.py index 62314a8eb7ea..96c6f85685e8 100644 --- a/test/registered/unit/layer_boundary/test_step3p5_dense_reduce_scatter.py +++ b/test/registered/unit/layer_boundary/test_step3p5_dense_reduce_scatter.py @@ -33,12 +33,12 @@ def mlp(hidden_states): complete = MagicMock(side_effect=lambda h, r: (h, r)) plan = stub_plan() - plan._batch_steps = lambda fb: SimpleNamespace( - output=comm.StageOutput(comm.Layout(frozenset())), + plan.path_for = lambda fb: SimpleNamespace( + output=comm.OutputContract(comm.Layout(frozenset())), returns_over_dp=False, output_move_completes_sum=False, ) - plan.output._select_ffn_completion = lambda fb, steps: comm.FfnCompletion( + plan.output._decide = lambda fb, steps: comm.ExitDecision( defer_moe_finalize=False, fuse_mlp_allreduce=False, mlp_reduce_scatter=reduce_scatter, diff --git a/test/registered/unit/layer_boundary/test_tbo_unreduced_input.py b/test/registered/unit/layer_boundary/test_tbo_unreduced_input.py index 69d7903dce98..10d338c0138c 100644 --- a/test/registered/unit/layer_boundary/test_tbo_unreduced_input.py +++ b/test/registered/unit/layer_boundary/test_tbo_unreduced_input.py @@ -9,7 +9,7 @@ import torch from sglang.srt.batch_overlap import two_batch_overlap as tbo -from sglang.srt.layers.layer_boundary import ADD, Layout, UnreducedOutput +from sglang.srt.layers.layer_boundary import PLAIN_ADD, Layout, UnreducedOutput from sglang.srt.layers.layer_boundary.residual.stream import ResidualStream from sglang.srt.utils import empty_context from sglang.test.ci.ci_register import register_cpu_ci @@ -23,7 +23,9 @@ def test_each_microbatch_owns_its_contribution_and_merge_keeps_the_update(self): group = SimpleNamespace(all_reduce=lambda x: x * 2) residual = torch.ones(4, 3) stream = ResidualStream(residual) - hidden = stream.leave(UnreducedOutput(torch.ones(4, 3), group=group), ADD) + hidden = stream.record( + UnreducedOutput(torch.ones(4, 3), group=group), PLAIN_ADD + ) batch = SimpleNamespace(residual_stream=stream, global_forward_mode=None) parts_seen = [] @@ -48,7 +50,7 @@ def execute(inputs_arr, **kwargs): state = part["forward_batch"].residual_stream value = part["hidden_states"] self.assertIsNot(state.pending, stream.pending) - self.assertIs(state.pending.update, ADD) + self.assertIs(state.pending.update, PLAIN_ADD) with self.assertRaises(RuntimeError): state.input( inputs_arr[1 if part is inputs_arr[0] else 0]["hidden_states"] @@ -56,7 +58,7 @@ def execute(inputs_arr, **kwargs): value, old_residual = state.input(value) torch.testing.assert_close(value, torch.full((2, 3), 2.0)) state.write(value + old_residual) - part["hidden_states"] = state.leave(value * 5, ADD) + part["hidden_states"] = state.record(value * 5, PLAIN_ADD) return inputs_arr with ( @@ -93,7 +95,7 @@ def execute(inputs_arr, **kwargs): self.assertTrue( all(part["forward_batch"].residual_stream is None for part in parts_seen) ) - self.assertIs(output.pending.update, ADD) + self.assertIs(output.pending.update, PLAIN_ADD) self.assertIs(output.pending.value, merged) torch.testing.assert_close(merged, torch.full((4, 3), 10.0)) torch.testing.assert_close(output.residual, torch.full((4, 3), 3.0)) 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 133e8125a3ed..2d22bf67f375 100644 --- a/test/registered/unit/lora/test_triton_dp_attention_unit.py +++ b/test/registered/unit/lora/test_triton_dp_attention_unit.py @@ -11,14 +11,14 @@ from sglang.srt.layers.dp_attention import DpPaddingMode from sglang.srt.layers.layer_boundary import ( - ADD, - NORM_QUANT_READ, + NORM_QUANT_READOUT, + PLAIN_ADD, Layout, StageKind, TokenAxis, ) -from sglang.srt.layers.layer_boundary.contracts import BatchVariant, StageEntry -from sglang.srt.layers.layer_boundary.prepare import _consumer_step, _read_input +from sglang.srt.layers.layer_boundary.contracts import BatchVariant, EntryPath +from sglang.srt.layers.layer_boundary.prepare import _run_entry, _update_read from sglang.srt.lora.backend.base_backend import BaseLoRABackend from sglang.srt.lora.backend.triton_backend import ( TritonLoRABackend, @@ -233,42 +233,42 @@ def test_communicator_publishes_layout_at_each_transition( communicator = stub_plan() communicator._publish_lora_layout = publish_lora_layout communicator.enters_stack = False - communicator._paths[BatchVariant.SEQUENCE_PARALLEL] = None + communicator.paths[BatchVariant.SEQUENCE_PARALLEL] = None communicator.norm = lambda x: x communicator.qkv_latent_func = None gathered, local = Layout(frozenset()), Layout(frozenset({TokenAxis.ATTN_DP})) # The rows of the steps the batch runs decide, not the ordinary steps'. - communicator._paths[BatchVariant.ORDINARY] = SimpleNamespace( + communicator.paths[BatchVariant.ORDINARY] = SimpleNamespace( entry=SimpleNamespace(input_rows=local if gathered_over_dp else gathered) ) - attention_entry = StageEntry( + attention_entry = EntryPath( input_rows=local, prepare=partial( - _consumer_step, + _run_entry, step=partial( - _read_input, - layer_input=None, + _update_read, + pre_move=None, enters_stack=False, - read=NORM_QUANT_READ, - update=ADD, + read=NORM_QUANT_READOUT, + update=PLAIN_ADD, ), carried_fusions=(), - adds_plainly=True, + is_plain_add=True, ), input_move=lambda hidden_states, **kwargs: hidden_states, - handoff=lambda hidden_states, *args: hidden_states, + attn_input_adapter=lambda hidden_states, *args: hidden_states, ) - ffn_entry = StageEntry( + ffn_entry = EntryPath( prepare=lambda hidden_states, residual, *args, **kwargs: ( hidden_states, residual, ), input_rows=gathered if gathered_over_dp else local, input_move=None, - handoff=None, + attn_input_adapter=None, ) selected = SimpleNamespace(entry=ffn_entry) - communicator._batch_steps = lambda forward_batch: selected + communicator.path_for = lambda forward_batch: selected hidden = torch.zeros(num_tokens, 4) with get_forward().scoped(lora_batch_layout=initial): for _ in range(2): diff --git a/test/registered/unit/model_executor/test_num_token_non_padded_localization.py b/test/registered/unit/model_executor/test_num_token_non_padded_localization.py index a51a5073f7e3..af93d3cd17bf 100644 --- a/test/registered/unit/model_executor/test_num_token_non_padded_localization.py +++ b/test/registered/unit/model_executor/test_num_token_non_padded_localization.py @@ -54,7 +54,7 @@ def sparse_moe_input(rows): "is_moe_input_scattered_across_dp_ranks", return_value=local ), patch_communicator("is_enable_moe_cp_allgather", return_value=rows == "moe_cp"), - patch_communicator("_gathers_over_attention_cp", return_value=False), + patch_communicator("_cp_gathers_over_attn_cp", return_value=False), ): yield diff --git a/test/registered/unit/test_runtime_context.py b/test/registered/unit/test_runtime_context.py index 12e6287ae6bd..48d7c6bbc4dc 100644 --- a/test/registered/unit/test_runtime_context.py +++ b/test/registered/unit/test_runtime_context.py @@ -1443,7 +1443,7 @@ def probe(x): self.assertEqual(probe(torch.zeros(())).item(), 0) def test_parallel_config_leaves_trace_under_torch_compile(self): - # Regression: gate helpers such as ``enable_moe_dense_fully_dp()`` read + # Regression: gate helpers such as ``is_dense_ffn_fully_dp()`` read # parallel config leaves inside compiled model forwards, which must # stay dynamo-traceable (``object.__getattribute__`` graph-breaks). # fullgraph=True turns any graph break back into a failure.