[Perf] CP: QSA Context Parallel Optimization by Shunkangz · Pull Request #43792 · sgl-project/sglang

Shunkangz · 2026-10-12T07:30:50Z

Motivation

Under prefill context parallelism, each QSA layer runs two all-gathers: one for the index keys and one for the layer's K/V. The following changes reduce synchronization overhead and overlap communication with computation.

1. Remove Host–Device Synchronizations

The first change removes four host–device synchronizations per layer on this path, or five when using a cached prefix:

The positions.max().item() call in the indexer's RoPE.
The .tolist() conversion of sequence lengths.
The prefix's request-index .tolist() conversion, when using a cached prefix.
The per-layer torch.tensor copy that builds the sparse-attention kernel's K offsets.
The sparse-attention kernel's max_q.item() call.

Three of these synchronizations, or four with a cached prefix, occurred immediately after an all-gather. As a result, every rank's CPU waited for the slowest rank on each layer, and the GPU then sat idle while the CPU relaunched work. Removing these synchronizations lets the CPU run ahead.

This saves:

CP4: approximately 8 ms at both 8K and 32K, reducing latency by 27.5% and 11.7%, respectively.
CP8: approximately 3–8.5 ms at 8K and 32K, reducing latency by 12.1% and 19.4%, respectively.

There is no numeric change.

2. Overlap the K/V All-Gather with the Indexer

The second change overlaps the K/V all-gather with the indexer. A hook fires immediately after the indexer issues its own key gather and starts the K/V gather on a side stream.

The K/V gather then runs during:

Compression.
MQA scoring.
Top-k selection.

Attention waits for the gather only when it needs the result.

The index-key gather stays inline: it is short, and its long waits come from load imbalance, which overlap cannot hide.

This provides additional latency reductions:

At 32K: 1.5 ms on CP4 and 1.85 ms on CP8, or 2.5% and 5.3%, respectively.
At 128K: approximately 7 ms.
At the 8K first chunk: neutral overall, because the layer is limited by kernel launch speed.

Modifications

Accuracy Tests

Speed Tests and Profiling

CP4

Shape	Main	Commit 1	Commit 1 gain	Commit 2	Commit 2 gain (over Commit 1)	Total gain
8K	29.24	21.19	−8.05 (−27.5%)	21.49	+0.30 (+1.4%)¹	−7.75 (−26.5%)
32K	68.45	60.44	−8.02 (−11.7%)	58.94	−1.50 (−2.5%)	−9.52 (−13.9%)
128K	431.42	425.53	−5.89 (−1.4%)	418.46	−7.07 (−1.7%)	−12.96 (−3.0%)

CP8

Shape	Main	Commit 1	Commit 1 gain	Commit 2	Commit 2 gain (over Commit 1)	Total gain
8K	24.52	21.56	−2.96 (−12.1%)	21.25	−0.31 (−1.4%)	−3.27 (−13.3%)
32K	43.57	35.11	−8.46 (−19.4%)	33.26	−1.85 (−5.3%)	−10.32 (−23.7%)
128K	229.28	222.51	−6.77 (−3.0%)	215.77	−6.74 (−3.0%)	−13.51 (−5.9%)

This is the nsys timeline before this PR. There are two areas that we should pay attention. In the first area, we have some device and host sync overhead. In the second area, the kv cache allgather is serial function call.

After we merge this PR, the kv cache allgather for new generated kv can be issued earlier and overlap with indexer computation. However, the allgather takes longer time because we do not use the 0SM comm. This will be integrated after the 0SM #44029 merged.

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 #39539850826
Latest PR Test (Extra): ❌ Run #39539850792
Latest PR Test (AMD ROCm 10): ❌ Run #39539850831

Each QSA layer on the CP prefill path synchronized the host with the device four times (five with a cached prefix). Three of them sit after an all-gather, so every rank's CPU waited for the slowest rank's gather each layer and the GPU then idled while the CPU relaunched. Remove them: - S1 apply_rope's positions.max().item(): the CP indexer skips the per-call cos/sin cache sizing; the load-time reservation already covers the positions that this layer's attention reads without a check. - S2 get_prefill_mqa_inputs' .tolist(): extend metadata carries the host sequence lengths (prefill_sequence_lengths_cpu) from seq_lens_cpu. - S3 prefix pack's req_pool_indices.tolist(): use the scheduler's req_pool_indices_cpu. - S4 cu_k torch.tensor(host list): built on device from seq_lens. - S5 sparse_gqa_fwd_interface_triton_ck's max_q .item(): optional host max_q, fed from the zigzag metadata's max_seqlen_q_combined. Measured with a 12-QSA-layer CP4 benchmark on 8x B200 (real kernels and NCCL, Qwen3.8-Flash-Next shapes): 8K 29.7 -> 21.3 ms, 8K chunk on a 56K prefix 35.6 -> 26.2 ms, 4x8K 66.3 -> 57.5 ms, 32K 68.5 -> 60.4 ms. Host syncs per layer go from 4 (5 with a prefix) to 0; the written caches, every top-k selection input and the attention output match the old path bitwise. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

With the host syncs gone, the layer's K/V all-gather sits on the QSA CP critical path ahead of sparse attention. It does not feed compression, MQA scoring or top-k, so start it on a side stream right after the indexer's key gather and let the attention wait for it: - QSAIndexer.forward_cuda_cp takes on_key_gather_issued (plumbed through forward_cuda as cp_on_key_gather_issued), called right after the key gather is enqueued, so every rank issues the CP group's collectives in one order. - QwenSparseAttnBackend.prefetch_cp_kv issues the gather: every rank sends its padded rows, and one index_select per half restores global token order through a per-forward index built from the zigzag host metadata. All buffers are allocated on the consumer stream and kept alive until the wait, so no record_stream is needed. _forward_extend_cp consumes a prefetched gather and otherwise gathers inline; init_forward_metadata drains an unconsumed one. The key gather stays inline: it is short, and its waits are the previous layer's sparse-attention imbalance, which overlap cannot hide. The Qwen4-Exp call site (cp/4) enables it by passing cp_on_key_gather_issued=partial(backend.prefetch_cp_kv, self.attn, k, v, forward_batch) to the indexer. Measured with a 12-QSA-layer CP4 benchmark on 8x B200 against the previous commit: 8K chunk on a 56K prefix 26.25 -> 25.64 ms, 4x8K 57.52 -> 55.85 ms, 32K 60.43 -> 58.77 ms, 128K 425.8 -> 419.1 ms; the CPU-bound 8K first chunk is +0.2 ms. The written caches, every top-k selection input and the attention output match the inline gather bitwise, padded and ragged batches included. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Move the side-stream gather into CP utilities with generic contiguous output splits. Reuse strategy-owned token ordering for zigzag and interleave while preserving QSA launch timing, deferred waits, and buffer lifetime. Validated with 114 unit tests and 20 subtests, all pre-commit checks, and six two-GPU PyNccl cases covering ordering, output splits, stream dependencies, and asynchronous launch.

The prefetch had its own gather: a side-stream all_gather into buffers allocated on the main stream, an index_select by a per-forward gather index, and generic CPAsyncGather/CPGatherHandle/build_gather_index helpers in layers/cp. Run the inline gather path instead (cp_materialize_global_token_order) on the qsa_cp_kv stream. The side stream waits for the main stream when the indexer's key gather has been issued; the layer's _forward_extend_cp waits for the side stream before reading the gathered K/V, and the side stream, which allocated them, waits for those reads before they are released. The caller keeps k/v alive through the join. No record_stream, no keep-alive tuple, and layers/cp is back to main. This also removes the gather index's per-forward host-to-device copy, which synchronized the host once per forward. A single pending slot replaces the per-layer dict; prefetch_cp_kv asserts it is empty, and init_forward_metadata joins one a failed forward left behind. The cached-prefix pack reads slots from the per-forward token_slot_table and cu_k from the per-forward int32 sequence_lengths, instead of req_pool_indices(_cpu) and a per-layer cast of seq_lens. With --enable-symm-mem the gather's symmetric buffer now comes from the side stream's share of the NCCL pool, separate from the forward stream's preallocation: one more segment of about the largest K/V gather (2 KiB per token for Qwen3.8-Flash-Next). Measured with the 12-QSA-layer CP4 benchmark on 8x H200, before the release wait and the cu_k reuse were added: the written caches, every top-k selection input and the attention output match the inline gather bitwise (8193, 8K on a 56K prefix, 4x8K, 3000+5192 on a 64-token prefix); host syncs per layer stay at 0; with the join skipped and the gathers delayed, the check fails on every rank. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

The QSA backend now fills prefill_sequence_lengths_cpu for every non-speculative extend batch when it builds the forward metadata. Every breakable prefill CUDA graph capture and replay rebuilds that metadata from the live batch before the indexer's eager break runs, so the per-layer msgspec.structs.replace in _compute_qsa_topk_indices_eager rebuilt the same tuple. The backend's copy slices seq_lens_cpu by seq_lens' length, as the indexer's device fallback does; the override sliced by batch_size. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Shunkangz · 2026-10-15T02:47:37Z

/tag-run-ci-label

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 · 2026-10-16T00:59:52Z

/rerun-failed-ci

Fridge003 · 2026-10-20T22:56:28Z

Can you please paste a screenshot of profile traces, which shows the overlap of kv all gather and indexer computation?

Shunkangz · 2026-10-22T02:45:40Z

/rerun-failed-ci

Fridge003 · 2026-10-22T05:07:22Z

Can merge after CI pass

Shunkangz and others added 2 commits October 11, 2026 22:50

Shunkangz requested review from Fridge003, HaiShaw, Qiaolin-Yu, hebiao064, ispobock and merrymercy as code owners October 12, 2026 07:30

Shunkangz requested review from BBuf, Edwardf0t1, Ying1123 and ch-wan as code owners October 13, 2026 01:57

Fridge003 mentioned this pull request Oct 13, 2026

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

Open

15 of 21 tasks

Shunkangz and others added 2 commits October 14, 2026 02:10

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

Shunkangz mentioned this pull request Oct 17, 2026

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

Open

8 of 11 tasks

Shunkangz changed the title ~~QSA Context Parallel Optimization~~ [Perf][NVIDIA] QSA Context Parallel Optimization Oct 19, 2026

Shunkangz changed the title ~~[Perf][NVIDIA] QSA Context Parallel Optimization~~ [Perf] QSA Context Parallel Optimization Oct 19, 2026