Repository navigation
[Feature] Megatron LayerNorm sequence parallelism (--enable-layernorm-sp) - #30915
Merged
Fridge003 merged 2 commits intoSep 1, 2026
Merged
Conversation
Contributor
|
Warning You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again! |
MartinHua
force-pushed
the
feat/layernorm-sp-generalize
branch
from
July 12, 2026 05:31
694aec5 to
ba0f52f
Compare
MartinHua
marked this pull request as ready for review
August 2, 2026 06:12
MartinHua
requested review from
BBuf,
Edwardf0t1,
Fridge003,
HaiShaw,
Ying1123,
ch-wan,
ispobock and
merrymercy
as code owners
August 2, 2026 06:12
Contributor
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
Fridge003
reviewed
Aug 10, 2026
Collaborator
|
Please solve the conflicts |
MartinHua
force-pushed
the
feat/layernorm-sp-generalize
branch
7 times, most recently
from
August 12, 2026 03:31
bea9beb to
b9486c9
Compare
Fridge003
reviewed
Aug 25, 2026
MartinHua
force-pushed
the
feat/layernorm-sp-generalize
branch
6 times, most recently
from
August 29, 2026 03:59
f7333c4 to
0a4e746
Compare
Fridge003
approved these changes
Aug 29, 2026
MartinHua
force-pushed
the
feat/layernorm-sp-generalize
branch
2 times, most recently
from
August 29, 2026 19:56
1f7429a to
6bcba46
Compare
Generic Megatron-style LayerNorm SP under pure tensor parallelism (arXiv:2205.05198), fully decoupled from model code: no changes under python/sglang/srt/models/ and none to dp_attention.py. The row-parallel all-reduce is split into reduce-scatter (g-bar) + all-gather (g) so the LayerNorm/residual regions run on sequence-sharded activations, cutting prefill activation memory with no extra communication volume. Prefill (EXTEND) only; Qwen3 dense; off by default and byte-identical when off. All SP logic lives in the new layers/layernorm_sp.py: the Qwen3 allowlist, the config guards, the per-forward sp_active flag, the entry-scatter / exit-gather collectives, and the fused g / g-bar matmul fast-paths (torch symmetric memory, unquantized bf16/fp16, with plain collective fallbacks). Integration is entirely in the shared layers: - LayerCommunicator builds an all-SCATTERED sibling when the model is SP-enabled and delegates prepare_attn / prepare_mlp / postprocess_layer to it while the region is active (LayerNorm on the shard, no collectives). It also owns the region: on the first decoder layer it re-evaluates activation and entry-scatters, so the state is self-correcting each forward. The all-SCATTERED case is added to CommunicateWithAllReduceAndLayerNormFn: under pure TP attn_tp_size == tp_size > 1, so the existing same-group-size shortcut does not cover it. - Row/ColumnParallelLinear auto-participate while sp_active (no per-model marking): column does the g all-gather+matmul then narrows to the real token count; row does the matmul+g-bar reduce-scatter. - LogitsProcessor exit-gathers back to the full sequence before the LM head. It recomputes the decision from config + forward mode and takes the token count from input_ids rather than reading sp_active / module state, because it can sit outside the CUDA-graph-captured region, where Python writes made inside that region do not re-execute on replay and would read stale. ServerArgs.check_server_args validates the flag (supported architecture, tp_size > 1, not --enable-dp-attention, not speculative decoding) so an unsupported model fails loud instead of silently ignoring it.
Unit test (no server): the Qwen3 allowlist, the prefill-only activation rule, the four config guards, and that the exit-gather predicate stays independent of the sp_active flag (the flag is stale on CUDA-graph replay). E2E test: Qwen3-1.7B with --tp 2 --enable-layernorm-sp, GSM8K accuracy above a threshold -- SP must not regress model output.
MartinHua
force-pushed
the
feat/layernorm-sp-generalize
branch
from
August 30, 2026 17:39
6bcba46 to
fa37892
Compare
Fridge003
approved these changes
Sep 1, 2026
15 of 21 tasks
StevenChenSE
pushed a commit
to StevenChenSE/sglang
that referenced
this pull request
Sep 6, 2026
5 of 6 tasks
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
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.
[Feature] Megatron LayerNorm sequence parallelism (
--enable-layernorm-sp)What
Megatron-style LayerNorm sequence parallelism (arXiv:2205.05198) for prefill under pure tensor parallelism: the row-parallel
all_reducebecomesreduce_scatter(g-bar) +all_gather(g), so LayerNorm/residual run on1/tpof the tokens -- less activation memory, and no extra communication volume, since the two forms move the same bytes. Both collectives are fused with the adjacent GEMMs via torch symmetric memory. Off by default and byte-identical when off.Design
All SP logic lives in the new
layers/layernorm_sp.py. There are no changes undersrt/models/and none todp_attention.py.LayerCommunicatorowns the region: it entry-scatters on the first decoder layer, then delegatesprepare_attn/prepare_mlp/postprocess_layerto an all-SCATTEREDsibling.Row/ColumnParallelLinearauto-participate in g / g-bar while the region is active, so no per-model marking.LogitsProcessorexit-gathers before the LM head.flags.sp.enabledis materialized once per worker byinitialize_layernorm_sp, besideinitialize_dp_attention.SP_SUPPORTED_ARCHITECTURESplus validation on hardware.The exit gather recomputes its condition from config +
forward_moderather than reading the per-forwardsp_activeflag: that flag is written by Python inside the CUDA-graph-captured region, so it is stale on replay, and reading it there fed sequence-sharded hidden states to the LM head.Requirements
arg_groups/layernorm_sp_hook.pyrejects these at startup, from the resolution pipeline: an architecture outsideSP_SUPPORTED_ARCHITECTURES(Qwen3ForCausalLM; others lack theLayerCommunicatorand participant-linear wiring),tp_size <= 1,--enable-dp-attention(SP shards over the full TP group, which under DP attention spans groups holding different sequences), and speculative decoding (the captured aux hidden states would be sequence-sharded).Silent gates, not errors: prefill (
EXTEND) only, and the fused path needs unquantized bias-free bf16/fp16 linears, falling back to plain collectives otherwise. NVLink/NVSwitch is required in practice but not yet enforced.Tests
test/registered/unit/layers/test_layernorm_sp.py(CPU) covers the allowlist, the prefill-only rule, the four guards, and a regression case pinning that the exit-gather predicate ignoressp_active.test/registered/models_e2e/test_layernorm_sp.py(2 GPU) runs Qwen3-8B with--tp 2 --enable-layernorm-spagainst a GSM8K threshold.Results
GSM8K is Qwen3-8B, tp=2, 1319 questions; prefill latency is the same model at 16K. Accuracy deltas are fp-reordering noise.
The gain scales with TP degree.
Follow-ups
CI States
Latest PR Test (Base): ❌ Run #33325891316
Latest PR Test (Extra): ❌ Run #33325891165
Latest PR Test (AMD ROCm 7.2): ❌ Run #33325891318