[dghosef]

Context / Sequence Parallel stack: 1/5 #1124 → 2/5 #1125 → 3/5 #1138 → 4/5 #1141 → 5/5 #1147

Stacked on #1138.

What this does

• A node can now split a long prefill's tokens across its SP GPUs and keep its KV stored by position (Context / Sequence Parallel 1/5: split the KV cache by position across GPUs (DCP) #1124's context_parallel). Before, you had to pick one: split work (Ulysses, KV by head) or split memory (DCP).
• Two ways to do the split over position-stored KV:
  • all-gather: gather the step's q/k/v, run Context / Sequence Parallel 1/5: split the KV cache by position across GPUs (DCP) #1124's attention unchanged, keep this GPU's share of the output;
  • ring: queries go round the GPUs while each GPU's K/V stays put, so no GPU ever holds the whole step's K/V.
• You pick the strategy once per node, with exchange on its sequence shard (SeqShardSpec, a resources: override beside min_tokens): auto (the default), ulysses, all_gather, ring or none. auto follows the KV storage: Ulysses over KV stored by head, and over KV stored by position the node runs every step copied (DCP). So existing configs don't change, and splitting over position-stored KV is something you opt into. A value the storage can't take fails at load (table below).
• Ragged (cacheless) attentions keep the model's own sp_exchange from Context / Sequence Parallel 3/5: move Ulysses into the engine and pick a strategy per resource #1138, now "none" (the default) or "ulysses". A shard that never splits stops them splitting too, and then their heads stay over TP.
• The Qwen3-Omni thinker uses it: configs/qwen3omni_thinker_tp4_sp2.yaml runs it on TP4×SP2 with KV stored by position, thinker_shard: {min_tokens: 8192, exchange: all_gather}, and min_tokens is a placeholder.

How it works

• The shard checks exchange against its node's KV storage when the deployment is derived, and each paged attention picks up the resolved value through its resolve_exchange when it's built. So there is one choice per node, and one attention can't split while another doesn't.
• All-gather is just a wrapper around the existing placed attention: one gather of q, k and v over the SP group, the copied path as it is, then a slice.
• Ring stripes the split by KV position: GPU r gets the step's tokens at positions x % P == r, which are exactly the slots it stores, so every K/V write is local. Then P rounds pass queries round the group; each round is one paged call over this GPU's slots, causal or not as the step declares, and the partials are merged in fp32. Hops are synchronous for now (overlap is a follow-up), only the FlashInfer backend plans it, and ring runs eager only: a captured step on a ring node runs copied.
• CommGroup.all_to_all_rows sends rows to each group rank in counts the caller gives. It replaces Context / Sequence Parallel 1/5: split the KV cache by position across GPUs (DCP) #1124's to_partners, and the ring uses it for its hops.
• When to pick ring: peak memory is the resident memory plus the bigger of attention and the MLP/MoE phase. Ring only lowers it once attention is the bigger one: from 4 SP GPUs, or on dense models. For the thinker at 2 SP GPUs the MoE sets the peak, so its config uses all-gather. (This is arithmetic, not measured.)
• Chunked prefill is still off on SP nodes, so a long thinker prompt prefills in one piece.
• The thinker declares a sequence shard over its one label and brackets its layer stack; the MoE padding mask and layer-0 embeds are computed over the whole step first, and the outputs are gathered after the final norm.

Shard exchange	KV storage	Split steps	Copied steps (decode, short prompts)
auto (default)	by head	Ulysses	head slice
auto (default)	by position	never split	DCP merge
ulysses	by head (refused by position)	Ulysses	head slice
all_gather	by position (refused by head)	gather q/k/v, placed attention, slice	DCP merge
ring	by position (refused by head)	striped ring, eager only	DCP merge
none	by position (refused by head on SP GPUs)	never split	DCP merge

Testing

The CPU rows ran on 7ecc854. The two commits on top change only docs and a docstring.

Test	Where	Result
ruff check .		clean
Full CPU Core suite vs main (eceea1b)	CPU box	10376 passed; 0 failures only on this branch (the 38 failures are all main's too)
Rust tests (test/rust/), built from this head	CPU box	666 passed
Targeted SP / CP modules	CPU box	487 passed, 37 skipped (skips need 2 GPUs or FlashInfer)
Thinker toy model from the real layers, 2 ranks, split vs copied, eager and compiled: by-head Ulysses, by-position all-gather, by-position ring	CPU, 2 ranks (gloo)	passed
Ring vs all-gather vs one rank, causal and not; exchange values the storage can't take refused at load, and LoRA refused where steps split	CPU, 2+ ranks (gloo)	passed
A shard that never splits (auto over KV by position, none) plans every step copied, and a ragged Ulysses there plans its heads over TP	CPU	passed
Thinker TP4×SP2 config passes startup validation	CPU	passed
GPU: thinker TP4×SP2 parity vs TP4 (48 needle/summary prompts at 1K–60K, gsm8k 250, speech out)	8×H100	needles 24/24 on all-gather and ring; 0 diverge under vLLM's top-5 rule, and 23–24/48 outputs differ vs a same-code floor of 8/48, all near-ties; the rows below split that up; gsm8k 96.4% vs 96.4%; speech WER 1.0% vs 1.1%
GPU: Thinker TP1×SP2 all-gather vs the copied path, 48 greedy prompts (1K/30K/60K), c=1	2×H200	0/48 differ (copied vs itself 0/48); prefill hidden states bit-identical to TP1; needles 24/24
GPU: Thinker TP1×SP2 ring vs the copied path, same prompts	2×H200	15/48 differ, all 30K/60K summaries (floor 0/48); needles 24/24; prefill drifts from layer 0 (relative 2.2e-3) to 3.5e-2 at the final norm. Hypothesis: rounding from merging bf16 partials
GPU: copied context parallel vs TP1, same prompts	2×H200	22/48 differ, summaries only (TP1 vs TP1 0/48; 1024- and 4096-token chunks vs whole 0/48), so most of the TP4 difference above comes from the base, not this PR
GPU: TTFT p50 at 64K / 128K c=1, 2 fresh-server rounds	8×H100	ring 0.87 / 2.52 s, all-gather 1.00 / 2.79 s, DCP-only 1.25 / 3.25 s, TP8 2.37 / 7.20 s
GPU: all-gather vs ring at 128K, peak memory per GPU	8×H100	TP4×SP2: ring 58.8 GB, all-gather 57.9 GB (DCP-only 63.3); TP2×SP4: ring 2.69 s at 66.8 GB, all-gather 3.26 s at 70.5 GB, DCP-only OOMs
GPU: long/short mix (8 × 1K streams + a 32K prompt every 10 s) and the min_tokens sweep	8×H100	short TPOT p50 10.5–11.0 ms on every TP4×SP2 arm (TP8 7.3), 32K TTFT ring 0.39 s, all-gather 0.47, DCP-only 0.58; all-gather beats DCP-only from 16K tokens and is slower at 4K–8K
GPU: unit suites, 17 fresh SP boots with a 32K first request, aborts and bad requests, SP = 1 regression check (Omni TP2, Qwen3.5-0.8B) vs #1138	8×H100	183 passed, 0 failed; 0 hangs; no regression

🤖 Generated with Claude Code