Repository navigation
[Refactor] Declare every model's stages from a shared function and drop the meta-device neighbour build - #43460
Merged
Merged
Conversation
ch-wan
requested review from
BBuf,
Edwardf0t1,
Fridge003,
HaiShaw,
Jiminator,
JustinTong0323,
Ying1123,
fzyzcjy,
ispobock,
kpham-sgl,
merrymercy,
pyc96,
sogalin,
wisclmy0611 and
zijiexia
as code owners
October 10, 2026 00:47
github-actions
Bot
added
documentation
This was referenced Oct 10, 2026
ch-wan
added
bypass-fail-fast
ch-wan
force-pushed
the
cheng/refactor/pp-neighbour-declarations
branch
from
October 10, 2026 03:23
ca7d821 to
78ab7b5
Compare
ch-wan
force-pushed
the
cheng/refactor/model-stage-declarations
branch
from
October 10, 2026 03:23
3efc489 to
9cdb3b9
Compare
Collaborator
Author
|
/tag-and-rerun-ci |
An EXAONE MoE decoder layer reads whether its FFN and the next layer's are sparse from the config's is_moe_layer list. The layer class gains a stage_facts that returns its two stage declarations from the config and the layer index; the layer appends its stages from it and chooses its MoE or dense FFN by the declared FFN, and the model (also the MTP draft model, which reuses it) gives it to make_pp_layers, so under pipeline parallelism a rank takes its neighbours' stages from the function instead of building those layers on the meta device. The declaration's expressions move unchanged into stage_facts. The FFN choice read config.is_moe_layer[layer_id], the value the declaration records as sparse, so the declarations and the modules each layer builds are the same as before.
The GLM-4 MoE and GLM-4 MoE Lite decoder layers decide whether their FFN and the next layer's are sparse from the config, the layer index and whether it is the NextN layer. Each layer class gains a stage_facts that returns its two stage declarations from those inputs; the layer appends its stages from it and takes is_layer_sparse from the declared FFN, and the model (also the GLM-4V MoE text model, which reuses it) gives it to make_pp_layers, so under pipeline parallelism a rank takes its neighbours' stages from the function instead of building those layers on the meta device. The NextN layers build in a stack of their own and declare the same stages as before. The sparsity rules move unchanged: _is_layer_sparse becomes a static method that reads the config it is given where it read self.config, the same object. GLM-4 MoE Lite reads moe_layer_freq, which the causal LM and each layer set to 1 before any declaration is made, as before. The declarations and the modules each layer builds are the same as before.
Every GPT-OSS layer declares an attention and a sparse FFN whose next layer's FFN is sparse too. The declaration moves to a static stage_facts on the decoder layer; the layer appends its stages from it and takes is_layer_sparse from the declared FFN, and the model gives the stage_facts of the layer type it builds to make_pp_layers, so under pipeline parallelism a rank takes its neighbours' stages from the function instead of building those layers on the meta device. The declarations are the same calls with the same arguments as before, and every layer still builds the sparse MoE block.
A Kimi Linear decoder layer decides whether its FFN and the next layer's are sparse from the config and the layer index alone, whichever attention the layer runs. The layer class gains a stage_facts that returns its two stage declarations from those inputs; the layer appends its stages from it and chooses its MoE or dense FFN by the declared FFN, and the model gives it to make_pp_layers, so under pipeline parallelism a rank takes its neighbours' stages from the function instead of building those layers on the meta device. Both sparsity expressions move unchanged into stage_facts, so the declarations and the modules each layer builds are the same as before.
A Laguna decoder layer reads whether its FFN and the next layer's are sparse from the config's mlp_layer_types. The layer class gains a static stage_facts that returns its two stage declarations from the config and the layer index; the layer appends its stages from it and takes is_layer_sparse from the declared FFN, and the model gives the stage_facts of the layer type it builds to make_pp_layers, so under pipeline parallelism a rank takes its neighbours' stages from the function instead of building those layers on the meta device. Both sparsity expressions move unchanged into stage_facts, so the declarations and the modules each layer builds are the same as before.
…e list Apertus, Arcee, the three Bailing models, dots3, ERNIE 4.5 VL MoE, EXAONE MoE, GLM-4 MoE and its lite variant, GPT-OSS, Kimi Linear and Laguna now give make_layers their shared declaration function, so a pipeline rank no longer builds their neighbouring layers to read the stages, and they leave the transitional list.
LlamaDecoderLayer gains stage_facts, a classmethod that returns the attention and FFN stages a layer declares from the config and its index, and the layer declares its own stages with it. LlamaModel, which every LlamaForCausalLM-based model builds its layers through, and Ministral3Model pass it to make_pp_layers. Under pipeline parallelism a rank then reads the stages of a neighbouring layer from the function instead of building that layer again on the meta device, and checks each layer it builds against it. _attn_readout becomes a classmethod, so that a subclass's choice (the EAGLE draft layer reads layer 0 without a norm) is part of stage_facts too. The declarations are unchanged: the same read and the same dense FFN, built from the same expressions.
Qwen2Model builds the layers of the decoder layer type it is given, so each type it is used with gains a stage_facts that returns the stages a layer declares from the config and its index: Qwen2DecoderLayer, Qwen3DecoderLayer, and JetNemotronDecoderLayer, which adds its own residual and declares none. Each layer declares its own stages with it, and Qwen2Model passes the type's function to make_pp_layers, so under pipeline parallelism a neighbouring layer's stages come from it rather than from building that layer on the meta device. Qwen2DecoderLayer's _attn_readout becomes a classmethod, so that the EAGLE draft layer's plain read of layer 0 is part of stage_facts too. The declarations are unchanged.
Qwen2MoeModel builds the layers of the decoder layer type it is given, so each type it is used with gains a stage_facts that returns the stages a layer declares from the config and its index: Qwen2MoeDecoderLayer, Qwen3MoeDecoderLayer (which InternS1ProTextDecoderLayer inherits) and MellumDecoderLayer. Qwen2MoeModel passes the type's function to make_pp_layers, so under pipeline parallelism a neighbouring layer's stages come from it rather than from building that layer on the meta device. Each layer declares its stages with the function and takes is_layer_sparse from the FFN stage it returns. The sparsity rules move into the functions as they were: every Qwen2-MoE and Qwen3-MoE layer is sparse, and Mellum's mlp_layer_types check becomes the static _is_sparse(config, lid) with the same bounds and validation.
DeepseekV2DecoderLayer gains stage_facts, a classmethod that returns the attention and FFN stages a layer declares from the config, its index and whether it is a NextN layer. _build_stages declares with it, and DeepseekV2Model and GigaChat35Model pass their layer's function to make_pp_layers, so under pipeline parallelism a neighbouring layer's stages come from it rather than from building that layer on the meta device. _is_layer_sparse becomes a static method of the config, the layer index and is_nextn, with the same rule; the layer's is_layer_sparse and the declared sparsity of this and the next layer both come from it. GigaChat35DecoderLayer overrides stage_facts: its post-sublayer norms are the output transforms when layernorm_type puts a norm after the sublayers, which _norm_placement now decides for both the norms the layer builds and the transforms it declares. Its own sparsity rule becomes static the same way. The declarations are unchanged.
Glm4DecoderLayer gains stage_facts, which returns the stages a layer declares from the config and its index; given the layer's post-attention and post-MLP norms, it binds the output transforms to them. The layer declares its stages with it, and Glm4Model passes the layer type's function to make_pp_layers, so under pipeline parallelism a neighbouring layer's stages come from it rather than from building that layer on the meta device. A neighbour's transforms are bound to no norm; binding reads only what they declare. The declarations are unchanged.
Gemma4DecoderLayer gains stage_facts, which returns the stages a layer declares from the config and its index; given the layer's post-attention norm, it binds the attention's output transform to it. The layer declares its stages with it, and Gemma4TextModel passes it to make_pp_layers, so under pipeline parallelism a neighbouring layer's stages come from it rather than from building that layer on the meta device. Whether a layer has the parallel MoE block, which is also whether its FFN stage is sparse, is read by one static method, _enable_moe_block, for both the modules the layer builds and the stages it declares. The declarations are unchanged.
Both kinds of GraniteMoeHybrid layer declare the same two stages. A module-level stage_facts now returns them from the config, with the output scale bound to the layer's _scale_output when the layer declares itself, and GraniteMoeHybridModel passes it to make_pp_layers, so under pipeline parallelism a neighbouring layer's stages come from it rather than from building that layer on the meta device. The FFN stage was sparse when the layer had built a block_sparse_moe. That choice is now _has_moe(config), which both layers build the block by and stage_facts declares from, with the condition the layers built it under before. The declarations are unchanged.
Exaone4DecoderLayer gains stage_facts, which returns the stages a layer declares from the config and its index; given the layer's output norms, it binds the post-norm updates to them. The layer declares its stages with it, and Exaone4Model passes it to make_pp_layers, so under pipeline parallelism a neighbouring layer's stages come from it rather than from building that layer on the meta device. A neighbour's updates are bound to no norm; binding reads only what they declare. The declarations are unchanged.
Each ZAYA1 layer declares one stage: an attention on even layer ids, an MoE on odd ones. Each layer class gains a stage_facts that returns its stage, read through the layer's residual scaling when it has one, and the layer declares with it. _layer_class holds the even/odd rule for both building a layer and the module-level stage_facts, which ZayaModel passes to make_pp_layers, so under pipeline parallelism a neighbouring layer's stage comes from it rather than from building that layer on the meta device. The declarations are unchanged.
DeciLMDecoderLayer gains stage_facts, which returns the stages a layer declares from its block config: an attention unless the attention is a no-op, then an FFN unless the FFN is a no-op, and none when both are. The layer pairs them with its norms, and DeciModel passes the layer type's function to make_pp_layers, so under pipeline parallelism a neighbouring layer's stages come from it rather than from building that layer on the meta device. The declarations are unchanged.
Llama, Ministral 3, Qwen2 (and the layer types its model builds), Qwen2 MoE, DeepSeek-V2, GLM-4, Gemma 4, GraniteMoeHybrid, EXAONE 4.0, GigaChat 3.5, ZAYA1 and DeciLM now give make_layers their shared declaration function, so a pipeline rank no longer builds their neighbouring layers to read the stages, and they leave the transitional list.
A pipeline rank bound its first and last K3 stages by building the neighbouring layers again on the meta device, and each layer derived parts of its declarations from the modules it had just built: the dense MLP's TP width from down_proj's group, whether the FFN writes the next stream from the MoE's latent and a2a attributes. The selection rules those modules are built with are now functions of the config and the fixed parallel setup, shared by the modules and by _stage_facts, which returns the two stages a layer declares: _dense_mlp_parallel_group (the dense MLP's default group; its TP width is resolve_linear_parallel_group of it), _uses_latent_moe, _uses_ep_a2a and _latent_moe_writes_stream. The layer declares itself with _stage_facts, passing its own reads and updates of the attention-residual bank; without them the function gives what they declare, from AttnBankState.facts. The model hands the function to make_pp_layers. Whether K3's SP collective is enabled is settled by the model once before its layers are built and passed in, so the function never sets the collective up; it selects K3's tuned all-gather and the bank's slice mode as before. _carries_bank_slices takes it the same way. Every expression keeps its value: each moved rule is the expression it replaces, reading the config or the same server args instead of the built module, and enabled() is cached, so the layers and the model see one answer. The declared stages, and so the bound boundaries, are unchanged.
A pipeline rank bound its first and last GLM-5-Next stages by building the neighbouring layers again on the meta device. Glm5NextDecoderLayer.stage_facts returns the two stages the layer at a global index declares, from the config: the residual (a plain one, or hyper-connection streams under config.mhc), the attention's head partition, and the sparsity of the layer and the next one, by _is_layer_sparse, now a static method of the config that also chooses the layer's MLP. The layer declares itself with it, passing its own reads and updates of the streams; without them the function gives what they declare, MHCState.facts: the class constants of MHCState's reads and updates, which do not depend on the layer's parameters. The model hands the function to make_pp_layers. Each expression is the one it replaces, reading the config where it read the layer's copy of it, so the declared stages are unchanged.
A pipeline rank bound its first and last Qwen3.5 stages by building the neighbouring layers again on the meta device. Qwen3_5ForCausalLM picks each layer's class from decoder_layer_types, which Qwen4ExpModel replaces with its gated hyper-connection layers, so the declaration function belongs to the layer class. _qwen3_5_is_moe, the model-type rule that chooses a layer's MLP, now also gives its sparsity. Both Qwen3.5 layer classes expose _qwen3_5_stage_facts as stage_facts and declare themselves with it, passing their own attention read (the quantized input their projection takes, which binding does not read). The Qwen4-exp layers expose _qwen4_exp_stage_facts, which takes the sparsity by the same rule in place of the built MLP's type and whether the attention read needs every row from _has_ple, which also decides whether the layer builds its PLE; the layers pass their own gated reads and updates, and without them the function gives what they declare, from GatedResidualState.facts. The model hands make_pp_layers the function of the class it picks for each index. Each rule is the expression it replaces, reading the config where it read a built module whose construction the same expression chose, so the declared stages are unchanged.
Every model whose layers declare stage boundaries now gives make_layers its shared declaration function, so a pipeline rank no longer builds the layers another rank holds on the meta device to read their stages. make_layers passes the function, at each neighbouring index, to the layer stack, which reads the stages it returns; a model that declares stages under pipeline parallelism without one fails at construction. The neighbour-build context (building_neighbour_layer and is_building_neighbour_layer), the list of models that still used it, and the resource skips it needed in the MoE dispatcher, the KT expert wrapper, MiniMax-M2's fused QK norm and Qwen4-exp's PLE offload are removed: every layer built is now a layer this rank runs.
ch-wan
force-pushed
the
cheng/refactor/model-stage-declarations
branch
from
October 10, 2026 19:42
c1ab028 to
84eebdb
Compare
This was referenced Oct 10, 2026
Shunkangz
added a commit
to Shunkangz/sglang
that referenced
this pull request
Oct 12, 2026
Main's sgl-project#42044 (GLM-5.3-Flash prefill CP) added declare_attn(tp_group=) and an ATTN_CP entry for a head-parallel mixer on full rows: the head partition this branch declared as sum_group. The resolution uses main's names and keeps what zigzag needs: - The GDN declares tp_group plus the branch's in_token_order. Under sharing its sum is ATTN_CP, the same coordinator as TP. - The token-order take-back is a move in main's ATTN_CP entry (ops.cp_take_back_in_token_order: complete the sum, then take this rank's zigzag shard). Main's CP moves are rank-major and would take back the wrong rows. prepare._cp_take_back_update_read is gone. - _select_entry_step returns EntryStep (sgl-project#43453); the token-order gather is its input_move. - _bind_stage exempts both DSA/MLA CP and CP-TP group sharing. - _qwen4_exp_stage_facts (sgl-project#43460) declares the GDN options; the linear layer's stage_facts passes linear_attention=True. _ffn_exit_rows went with ExitRows (sgl-project#42936). - Hook: both models registered. The EP/a2a/waterfill/PP/decode rejections and the PD prefill allowance are Qwen4-Exp's; GLM keeps main's disaggregation rejection. Qwen4-Exp keeps the plain speculative rejection message. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
6 tasks done
This was referenced Oct 13, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This PR is part of a stack (oldest at bottom):
Motivation
#43459 lets a model give
make_layersits shared declaration function and keeps a transitional list of models that still had their pipeline neighbours built on the meta device. This PR gives every remaining model that declares stages such a function and removes the meta-device path.Modifications
AttnBankState.facts,MHCState.factsandGatedResidualState.factsgive what their reads and updates declare without the modules.building_neighbour_layer/is_building_neighbour_layer, the neighbour rebuild, and the resource skips it needed (MoE dispatcher, KT expert wrapper, MiniMax-M2's fused QK norm, Qwen4-exp's PLE offload) are removed: every layer a rank builds is a layer it runs. A model that declares stages under PP without a function fails at construction; every model that builds its layers withmake_layersunder PP has one.Every moved rule is the expression it replaces, reading the config or the same server args instead of the built module.
Compatibility
A model that declares stage boundaries and runs under pipeline parallelism must now give
make_layers/make_pp_layersits shared declaration function (stage_facts); without one, construction fails with an error naming the model, instead of building the neighbouring layers on the meta device. Every in-tree model that builds its layers withmake_layersunder PP gives one; an out-of-tree model integration that declares stages needs to do the same (see the layer boundary developer guide). Models that do not declare stage boundaries, or that do not run under PP, are unaffected.Accuracy Tests
How the stack is checked:
test_aiter_allreduce_fusion_amdcases that fail on NVIDIA on main too; the test files that changed most recently pass on every commit. The fulltest/registered/unitsuite (926 files, one process per file) shows no failure on the stack's head that main does not have.516fd1c1b5unless marked as on the merge candidate; rebasing onto current main had no conflicts, and the CPU checks were run again after it.961bcf481d, H200, asymmetric dummy weights unless noted): DeepSeek-V2-Lite input-scattered PP2 is bit-identical to PP1 (main crashes on it); Nemotron-H cut between a mixer and its FFN differs from PP1 only at the rounding level of the FFN-to-mixer cuts (main: 0.93 and 2.21); Llama, Qwen3-MoE, GPT-OSS and DeepSeek-V2-Lite are bit-identical to main at PP1 and PP2; GSM8K (real DeepSeek-V2-Lite weights, 250 questions, TP4 with attention DP2 and the dense FFN fully data-parallel) gives 0.360 at PP2, 0.364 at PP1 and 0.368 on main at PP1, medians of three runs whose run-to-run spread is 0.012–0.016. The last revision of the stack changed only comments, docstrings, an equivalent condition and explicitly set attributes after these runs; its bound plans are identical.For this PR:
SGLANG_K3_SP_ATTN_RES0 and 1; GB300 tuning tables copied under the B300 device name): every rank reports the collective enabled and the tuned kernels accept the calls; the outputs are bit-identical to main and to main before Kimi K3 moved onto layer boundaries (666da4f with [Fix] Shard Kimi K3's shared experts over the TP group they sum over #42895 applied), and per-layer prefill inputs match in all compared layers. A deliberately doubled sum in the tuned path changes the outputs, so the comparison can see a wrong sum.SGLANG_K3_SP_ATTN_RES0 and 1: all runs fall in one bitwise cluster across the three trees, apart from one main run 1 bf16 ulp off at layer 1 whose repeats rejoin it; call counts of every tuned kernel match.a2a,fi_a2a,ag_rs; DeepEP and the standard dispatcher), and DCP2 alone: bit-identical across the trees. DCP with prefill CP fails the same way on all three (KimiK3LinearModel.forward()does not takeinput_embeds).arandrs, decode graph on and off): bit-identical across the trees wherever it runs, and PP2 matches PP1 except with the decode graph on, where all three trees differ the same way. A 3+3 cut with the decode graph on, and dynamic chunking with PP2, fail identically on all three.Speed Tests and Profiling
A pipeline rank no longer builds its neighbours' layers on the meta device at startup. No forward-path change.
Kimi K3 on B200 on the merge candidate against main (TP4, DeepEP, custom all-reduce, tuned SP collective, attention-residual bank, CUDA graph decode, three runs per tree): prefill of 4,096 tokens takes 23.4 ms against 23.7 ms (
SGLANG_K3_SP_ATTN_RES=0) and 23.4 ms against 23.3 ms (=1); decode throughput at batch 1, 16 and 64 differs by less than the run-to-run spread; peak allocated (90.912 GiB per rank), steady-state memory and the KV pool size are identical.Checklist
CI States
Latest PR Test (Base): 🚫 Run #38080897600
Latest PR Test (Extra): 🚫 Run #38080897340
Latest PR Test (AMD ROCm 10): 🚫 Run #38080897283