Repository navigation
[Qwen3.8 CP 1/4] Context parallelism for QSA attention and the sparse indexer - #39721
Merged
Merged
Conversation
Shunkangz
requested review from
BBuf,
Edwardf0t1,
Fridge003,
HaiShaw,
Qiaolin-Yu,
Ying1123,
ch-wan,
hebiao064,
ispobock and
merrymercy
as code owners
September 16, 2026 05:41
Shunkangz
changed the title
8 of 9 tasks
detain
added a commit
to detain/sglang
that referenced
this pull request
Sep 17, 2026
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>
Fridge003
reviewed
Sep 29, 2026
Collaborator
|
Please be aware that the layer communicator is under refactor (see this PR list #41557) It would be great if the communication operations between Attention and FFN can be rebased on the latest implementation of layer communicator |
Adds the prefill-CP path to the Qwen sparse attention backend and its indexer, plus the shared block-layout helpers the CP collectives need. QwenSparseAttnBackend._forward_extend_cp all-gathers the local K/V into global token order, writes the complete sequence into every rank's pool (out_cache_loc is trimmed to the full extend count by the runner), and issues one sparse_gqa call per zigzag block. QSAIndexer.forward_cuda_cp projects q/k on local rows with local positions, all-gathers the raw index-K into global order for the unchanged ring store and block compression, then scores and selects top-k on local rows -- so the output rows already line up with the local q rows and need no gather after. layers/cp/collocated.py holds the row-layout helpers shared by any collocated-CP model: the rank-major block gather/scatter and the reduce-scatter that lands each rank's zigzag rows. The CPU unit test pins the invariant they rest on -- RS(cp_rank_major_blocks(x))[rank] must equal what ZigzagCPStrategy.shard_hidden_states gives that rank. Depends only on the existing prefill-CP infrastructure (is_cp_active, attn_cp_metadata, ZigzagCPStrategy); no model is wired to it yet. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Contributor
Author
For this PR, I only support the context parallel inside qsa. There is no change of communication between Attention and FFN. I will pay more attention about this when I integrate the E2E cp in the latter PR. Thank you. |
Contributor
Author
|
/tag-run-ci-label |
Collaborator
|
/rerun-test test/registered/unit/cp/test_qsa_cp_unit.py |
Contributor
|
Results for 🚀 |
Fridge003
approved these changes
Oct 1, 2026
alexnails
added a commit
to alexnails/sglang
that referenced
this pull request
Oct 3, 2026
The merge broke every CPU, GPU and jit-kernel shard with
ImportError: cannot import name 'MultiPlatformOp'
from 'sglang.srt.layers.utils'
This branch drops that re-export on purpose: it was the only reason
sglang.srt.layers.utils -- imported by dozens of model files for
PPMissingLayer / get_layer_id -- pulled in sglang.kernels.fused_op. The
supported spelling is the module path, and the alias module stays.
Since then main grew two importers of the old spelling, dsa_indexer_kpool
(sgl-project#39273) and qsa_indexer (sgl-project#39721). Neither conflicts textually, so the merge
took both sides cleanly and the break only showed up in CI.
The blast radius was wider than the two files because
import_model_classes catches every exception and logs a warning, so any
model whose chain reached either one silently vanished from the arch table
and surfaced much later as "KimiLinearForCausalLM has no SGlang
implementation" in the unified-memory e2e test.
Verified: every sglang.srt module now imports except for absent optional
deps (polars, numba, ray, vllm, cutlass), and all 36 attributes that
sglang.srt.platforms consumers read off current_platform resolve on
SRTPlatform.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
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
2 participants
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.
Motivation
Adds the prefill-CP path to the Qwen sparse attention backend and its indexer, plus the shared block-layout helpers the CP collectives need.
QwenSparseAttnBackend._forward_extend_cp all-gathers the local K/V into global token order, writes the complete sequence into every rank's pool (out_cache_loc is trimmed to the full extend count by the runner), and issues one sparse_gqa call per zigzag block. QSAIndexer.forward_cuda_cp projects q/k on local rows with local positions, all-gathers the raw index-K into global order for the unchanged ring store and block compression, then scores and selects top-k on local rows -- so the output rows already line up with the local q rows and need no gather after.
layers/cp/collocated.py holds the row-layout helpers shared by any collocated-CP model: the rank-major block gather/scatter and the reduce-scatter that lands each rank's zigzag rows. The CPU unit test pins the invariant they rest on -- RS(cp_rank_major_blocks(x))[rank] must equal what ZigzagCPStrategy.shard_hidden_states gives that rank.
Depends only on the existing prefill-CP infrastructure (is_cp_active, attn_cp_metadata, ZigzagCPStrategy); no model is wired to it yet.
Modifications
Accuracy Tests
Speed Tests and Profiling
Checklist
Review and Merge Process
/tag-and-rerun-ci,/tag-run-ci-label,/rerun-failed-ciCI States
Latest PR Test (Base): 🚫 Run #36677787006
Latest PR Test (Extra): ❌ Run #36677786667
Latest PR Test (AMD ROCm 10): ❌ Run #36677786943