[Qwen3.8 CP 4/4] Prefill context parallel E2E support of Qwen3.8 by Shunkangz · Pull Request #39724 · sgl-project/sglang

Shunkangz · 2026-09-16T05:46:18Z

Motivation

Modifications

Accuracy Tests

Speed Tests and Profiling

Checklist

Format your code according to the Format code with pre-commit.
Add unit tests according to the Run and add unit tests.
Update documentation according to Write documentations.
Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
Follow the SGLang code style guidance.

Review and Merge Process

Ping Merge Oncalls to start the process. See the PR Merge Process.
Get approvals from CODEOWNERS and other reviewers.
Trigger CI tests with comments or contact authorized users to do so.
- Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

CI States

Latest PR Test (Base): ⏳ Run #38520129531
Latest PR Test (Extra): ❌ Run #38520129446
Latest PR Test (AMD ROCm 10): ⏳ Run #38520129509

Manifest of the nine upstream PRs and the local commits this branch carries on top of upstream main, in build order, with the reasoning for the CP series (sgl-project#39721-sgl-project#39724) that was evaluated and left out, the HiCache caveat (sgl-project#39830), and the editable-install trap that makes git worktrees useless for baselining library changes here. Generated from git, so the SHAs match the branch exactly. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

Qwen4-Exp runs prefill CP with CP-TP group sharing (--enable-prefill-cp --attn-cp-size N --tp-size N --cp-strategy zigzag): the residual stream, QSA attention and the QSA indexer are zigzag CP-sharded, and the GDN and the MoE keep their TP partition. - Model: shard the residual after the embedding and gather it before the logits. The GDN gathers the sequence, sums its TP partials and keeps this rank's rows; the PLE gathers, runs and shards; the QSA indexer gets the global RoPE positions. The MoE needs no model code: the boundary's MoE-CP gather -> TP experts -> all-reduce -> take-back already does it, so the two MHC construction guards are lifted under the flag. - PD: a sharing prefill rank r holds every KV head, the CP-replicated QSA/PLE state and GDN head shard r. It registers as TP rank r of a CP-1 prefill and sends its own state, and the mamba path slices by the linear-attention TP width. Before, every decode rank got rank 0's shard. - Hook: register Qwen4-Exp; allow a prefill node (mooncake only); reject a decode node, EP, MoE a2a and PP. - Test: PD E2E, prefill TP4/CP4 -> decode TP4 on 8x B200, with a /server_info guard, a 20K-token needle at three depths, and GSM8K. Measured on 8x B200 (NVFP4): GSM8K 0.975 (200, thinking) and needle 3/3, the same as a plain-TP4 prefill node (0.975, 3/3). With the PD transfer change reverted, the needle fails at two of three depths. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

…he non-CP path - Hook: reject --enable-waterfill (it switches the MoE to DeepEP + EP after this hook) and SGLANG_DISAGG_STAGING_BUFFER on a prefill node (decode sees this prefill as TP cp_size while its sender stages by attention TP 1). - Model: shard with the existing cp_split_before_forward. - Tests: a short-prompt E2E check of the non-CP path under the flag, plus CPU rows for the CP-rank-as-TP-rank registration and the state-skip rule. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

The rank pairing already lives in CommonKVManager; what stayed per backend was the TP width the state is compared and sliced by. CommonKVManager now names it state_tp_size (attention TP, or the TP group under CP-TP group sharing), and mooncake, nixl and mori use it for the mamba state and the QSA state, while KV heads keep following attention TP. The hook no longer requires mooncake. Verified prefill TP4/CP4 -> decode TP4 on 8x B200 over both backends: nixl GSM8K 0.96, mooncake 0.97, needle 3/3 on each. Mori is not installed in the test container and was only read. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Under CP-TP group sharing the GDN layer hand-wrote its collectives: gather the zigzag shard into token order, sum the TP-partitioned heads, shard back. A stage can now declare both facts and the boundary does the rest: - declare_attn(sum_group=SumGroup.TP): the output owes a sum over the group the heads are partitioned over, here the TP group while attention TP is 1. An ordinary batch completes it at the next input like any attention sum. - declare_attn(in_token_order=True): on a CP extend the stage reads the whole sequence in token order. Its entry reads on this rank's shard and then gathers in token order; the next input sums over TP, takes this rank's shard back, writes it into the residual there, and gathers for the MoE. Qwen4-Exp declares both for its linear-attention layers, and its GDN layer is back to prepare -> linear_attn -> finish. The static sum check no longer needs the sharing exemption. Other models resolve the same plans: the new fields default to today's behaviour, and the new entry branch only takes an edge that raised before. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

The QSA indexer's CP path takes the whole sequence's RoPE positions. The model threaded them through every layer as cp_global_positions, but the batch keeps them whole under CP-TP group sharing: forward_batch.positions, or forward_batch.mrope_positions when the model runs mrope. The attention layer now reads them there, and the model keeps a plain flag for its final gather. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

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>

Shunkangz · 2026-10-12T08:44:02Z

/tag-and-rerun-ci

Fridge003 · 2026-10-13T01:51:49Z

/rerun-test test/registered/e2e/disaggregation/test_qwen4_exp_pd_prefill_cp.py

github-actions · 2026-10-13T01:52:14Z

🚀 8-gpu-b200 (1 test): ❌ View workflow run

cd test/ && python3 registered/e2e/disaggregation/test_qwen4_exp_pd_prefill_cp.py

Shunkangz · 2026-10-13T04:07:11Z

/rerun-test test/registered/e2e/disaggregation/test_qwen4_exp_pd_prefill_cp.py

github-actions · 2026-10-13T04:07:39Z

⛔ test/registered/e2e/disaggregation/test_qwen4_exp_pd_prefill_cp.py: File not found: test/registered/e2e/disaggregation/test_qwen4_exp_pd_prefill_cp.py

Resolve the MoRI Mamba-state transfer conflict by retaining upstream layer-aware transfer blocks and passing the branch state TP size to the slicing helper. Validation: 455 unit tests and 7,334 subtests passed across layer-boundary, disaggregation wire/bootstrap, and MoRI deferred-release tests; 97 focused state-transfer cases passed. Python syntax, isort, and git diff checks passed.

Include upstream 7ff3a11, which landed during merge validation. This update only changes JIT activation fast-math architecture selection and merges without conflicts.

Fridge003 · 2026-10-13T08:36:18Z

/rerun-test test/registered/e2e/disaggregation/test_qwen4_exp_pd_prefill_cp.py

github-actions · 2026-10-13T08:36:49Z

🚀 8-gpu-b200 (1 test): ❌ View workflow run

cd test/ && python3 registered/e2e/disaggregation/test_qwen4_exp_pd_prefill_cp.py

Fridge003 · 2026-10-14T05:06:19Z

/rerun-test test/registered/e2e/disaggregation/test_qwen4_exp_pd_prefill_cp.py

github-actions · 2026-10-14T05:06:49Z

⚠️ Rebase Required Before Re-run

A major update has landed on main. Your PR is diverged relative to required base commit c053fc71d3bd.

Re-run was not dispatched. What to do:

Rebase your branch onto the latest main and push again
Follow issue #21065 for context
CI-fix PRs may request the bypass-maintenance label to skip this check

Fridge003 · 2026-10-14T06:04:51Z

/rerun-test test/registered/e2e/disaggregation/test_qwen4_exp_pd_prefill_cp.py

github-actions · 2026-10-14T06:05:16Z

🚀 8-gpu-b200 (1 test): ✅ View workflow run

cd test/ && python3 registered/e2e/disaggregation/test_qwen4_exp_pd_prefill_cp.py

Conflict in qwen4_exp.py _compute_qsa_topk_indices_eager. Main's cp/4 (sgl-project#39724) added the cp_global_rope_positions indexer kwargs next to the use_host_prefill_lengths override that this branch drops. Kept main's CP kwargs and dropped the override: the QSA backend already fills prefill_sequence_lengths_cpu for every non-speculative extend. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Shunkangz requested review from BBuf, Edwardf0t1, Fridge003, HaiShaw, Qiaolin-Yu, Ying1123, ch-wan, hebiao064, hnyls2002, ispobock and merrymercy as code owners September 16, 2026 05:46

Shunkangz mentioned this pull request Sep 16, 2026

Prefill Context Parallel Support for Qwen3.8-Next-Flash #39062

Open

8 of 11 tasks

Fridge003 added the release-highlight Candidate PR for release note highlight label Oct 8, 2026

Shunkangz force-pushed the cp/4-qwen4exp branch from fc9cbc3 to 49dfc80 Compare October 9, 2026 03:34

Shunkangz requested review from ByronHsu, Duyi-Wang, ShangmingCai and sogalin as code owners October 9, 2026 03:34

Shunkangz and others added 2 commits October 8, 2026 20:46

Shunkangz requested review from Lzy17, billishyahao and kkHuang-amd as code owners October 9, 2026 05:54

Shunkangz and others added 3 commits October 9, 2026 01:16

github-actions Bot added the run-ci CI: run the baseline test suite on this PR label Oct 12, 2026

AliceChenyy mentioned this pull request Oct 12, 2026

[RFC] Prefill sequence parallelism with replicated MoE for small-TP long-prompt prefill (Qwen3.5 MoE) #43820

Open

Fridge003 mentioned this pull request Oct 13, 2026

[Roadmap] Prefill Context Parallelism (2026 Q4) #21788

Open

15 of 21 tasks

Shunkangz added 2 commits October 13, 2026 00:17

Merge latest main into cp/4-qwen4exp

903b3cb

Include upstream 7ff3a11, which landed during merge validation. This update only changes JIT activation fast-math architecture selection and merges without conflicts.

test: remove Qwen4-Exp long-context needle test [skip ci]

8e875b5

Fridge003 approved these changes Oct 14, 2026

View reviewed changes

Merge branch 'main' into cp/4-qwen4exp

19378f8

Fridge003 merged commit 0d130f1 into sgl-project:main Oct 14, 2026
94 of 119 checks passed

Navigation Menu

Repository navigation

[Qwen3.8 CP 4/4] Prefill context parallel E2E support of Qwen3.8 - #39724

[Qwen3.8 CP 4/4] Prefill context parallel E2E support of Qwen3.8#39724
Fridge003 merged 10 commits into
sgl-project:mainfrom
Shunkangz:cp/4-qwen4exp

Shunkangz commented Sep 16, 2026 •

edited by github-actions Bot

Loading

Uh oh!

Shunkangz commented Oct 12, 2026

Uh oh!

Fridge003 commented Oct 13, 2026

Uh oh!

github-actions Bot commented Oct 13, 2026 •

edited

Loading

Uh oh!

Shunkangz commented Oct 13, 2026

Uh oh!

github-actions Bot commented Oct 13, 2026 •

edited

Loading

Uh oh!

Fridge003 commented Oct 13, 2026

Uh oh!

github-actions Bot commented Oct 13, 2026 •

edited

Loading

Uh oh!

Fridge003 commented Oct 14, 2026

Uh oh!

github-actions Bot commented Oct 14, 2026

Uh oh!

Fridge003 commented Oct 14, 2026

Uh oh!

github-actions Bot commented Oct 14, 2026 •

edited

Loading

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

Shunkangz commented Sep 16, 2026 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Modifications

Accuracy Tests

Speed Tests and Profiling

Checklist

Review and Merge Process

CI States

Uh oh!

Shunkangz commented Oct 12, 2026

Uh oh!

Fridge003 commented Oct 13, 2026

Uh oh!

github-actions Bot commented Oct 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Shunkangz commented Oct 13, 2026

Uh oh!

github-actions Bot commented Oct 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Fridge003 commented Oct 13, 2026

Uh oh!

github-actions Bot commented Oct 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Fridge003 commented Oct 14, 2026

Uh oh!

github-actions Bot commented Oct 14, 2026

⚠️ Rebase Required Before Re-run

Uh oh!

Fridge003 commented Oct 14, 2026

Uh oh!

github-actions Bot commented Oct 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Shunkangz commented Sep 16, 2026 •

edited by github-actions Bot

Loading

github-actions Bot commented Oct 13, 2026 •

edited

Loading

github-actions Bot commented Oct 13, 2026 •

edited

Loading

github-actions Bot commented Oct 13, 2026 •

edited

Loading

github-actions Bot commented Oct 14, 2026 •

edited

Loading