minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[diffusion] feat: support K/V-gather style sequence parallel (CP-like) attention (#32667)

Browse Source 
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Mick
2026-08-07 09:39:28 +08:00
committed by GitHub
co-authored by Claude Fable 5
parent 9ee658d4f6
commit 1e08b865f9
13 changed files with 1105 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/docs/sglang-diffusion/api/cli.mdx
+1
View File
@@ -85,6 +85,7 @@ Use `sglang generate --help` and `sglang serve --help` for the full argument lis
- `--sp-degree {N}`: sequence parallelism size
- `--dp-size {N}` (alias `--data-parallel-size`): number of data-parallel replicas. Each replica is a full copy of the engine on `num_gpus / N` GPUs with its own ingress; generation requests round-robin across replicas, realtime sessions stick to the replica holding their state, and control operations (weights, LoRA, memory occupation, shutdown) apply to every replica. Combines with the other parallelism axes (`num_gpus = dp × cfg × tp × sp`); monolithic serving only.
- `--ulysses-degree {N}` and `--ring-degree {N}`: USP parallelism controls
- `--kv-gather-degree {N}`: sequence-parallel degree that splits rows inside attention and exchanges with one K/V all-gather (queries stay local) instead of Ulysses all-to-all. Non-causal attention only; does not compose with `--ulysses-degree`/`--ring-degree` yet. When no SP degree is set explicitly, `sp_degree=2` defaults to `kv_gather_degree=2` (its measured-win zone) and higher degrees default to Ulysses; under that auto assignment, attention calls the gather path cannot take fall back to the Ulysses exchange, while an explicit degree fails instead of degrading.
- `--enable-cfg-parallel {true|false}`: enable or explicitly disable CFG parallelism
- `--encoder-parallel {auto|fold|dp|replicate}`: how the text/image encoders use the GPUs the DiT replica leaves idle during encoding. `auto` (the default for both `generate` and `serve`) TP-folds an encoder wide enough to pay for the per-layer all-reduce, selects DP for a server batch when it can engage, and otherwise replicates; `fold` forces the shard whenever the dims allow it; `dp` splits a batched encode across ranks and needs `--batching-max-size > 1` to engage; `replicate` encodes redundantly on every rank. `fold` and `replicate` are bitwise-identical to single-GPU encoding. See [Encoder Parallelism](/docs/sglang-diffusion/encoder_parallel).
- `--warmup-mode {off|request|server}`: control startup warmup for `sglang serve`; `off` skips warmup, `request` primes the request path, and `server` runs a full synthetic server warmup before serving traffic
docs/docs/sglang-diffusion/ring_sp_performance.mdx
+164 -2
View File
@@ -2,7 +2,7 @@
title: "Sequence Parallelism"
tag: "preserve"
metatags:
description: "Configure sequence parallelism, Ulysses, and ring-based sequence splitting for SGLang Diffusion workloads."
description: "Configure sequence parallelism, TP plus SP, Ulysses, K/V gather, and ring-based sequence splitting for SGLang Diffusion workloads."
---
Sequence parallelism splits long image or video latent sequences across GPUs. In SGLang Diffusion, the public controls are:
@@ -10,6 +10,7 @@ Sequence parallelism splits long image or video latent sequences across GPUs. In
- `--sp-degree`: total sequence parallel degree
- `--ulysses-degree`: Ulysses parallel degree
- `--ring-degree`: ring parallel degree
- `--sp-attention-mode`: attention exchange used inside each SP group
The degrees must satisfy:
@@ -17,11 +18,172 @@ The degrees must satisfy:
sp_degree = ulysses_degree * ring_degree
```
The default `--sp-attention-mode ulysses` uses all-to-all to redistribute
sequence shards over attention heads. `--sp-attention-mode kv_gather` keeps
queries sequence-sharded and all-gathers keys and values, then computes each
rank's local output directly. The K/V-gather mode currently supports
non-causal attention with `--ring-degree 1`. Varlen calls through the legacy
`UlyssesAttention` adapter and video sparse attention are not supported.
Use SP when sequence length or video shape makes the DiT forward pass the bottleneck and the model supports sequence sharding. For latency-oriented multi-GPU Qwen/Wan deployments, also compare against CFG parallelism and FSDP; SP is not automatically the best multi-GPU setting for every model.
## Choosing The Attention Exchange
<table style={{width: "100%", borderCollapse: "collapse", tableLayout: "fixed"}}>
<colgroup>
<col style={{width: "20%"}} />
<col style={{width: "30%"}} />
<col style={{width: "25%"}} />
<col style={{width: "25%"}} />
</colgroup>
<thead>
<tr style={{borderBottom: "2px solid #d55816"}}>
<th style={{textAlign: "left", padding: "10px 12px", fontWeight: 700}}>Mode</th>
<th style={{textAlign: "left", padding: "10px 12px", fontWeight: 700}}>Communication</th>
<th style={{textAlign: "left", padding: "10px 12px", fontWeight: 700}}>Memory</th>
<th style={{textAlign: "left", padding: "10px 12px", fontWeight: 700}}>Constraints</th>
</tr>
</thead>
<tbody>
<tr>
<td style={{padding: "9px 12px"}}><code>ulysses</code></td>
<td style={{padding: "9px 12px"}}>All-to-all before and after attention</td>
<td style={{padding: "9px 12px"}}>Full sequence with a shard of the attention heads during attention</td>
<td style={{padding: "9px 12px"}}>Attention head divisibility must match the Ulysses degree</td>
</tr>
<tr>
<td style={{padding: "9px 12px"}}><code>kv_gather</code></td>
<td style={{padding: "9px 12px"}}>All-gather K and V; Q and output remain sequence-sharded</td>
<td style={{padding: "9px 12px"}}>Replicates full K and V within the SP group</td>
<td style={{padding: "9px 12px"}}>Non-causal attention and <code>ring_degree=1</code>; no legacy varlen or video sparse attention</td>
</tr>
</tbody>
</table>
Neither exchange is universally faster. K/V gather avoids the reverse
all-to-all and can help when its local attention shape or collective is more
efficient, while Ulysses can use less attention activation memory. Benchmark
both on the target model, resolution, accelerator, and interconnect.
For SP degree `P`, the approximate per-rank network payload of K/V gather
relative to Ulysses is `P / 2`, excluding each rank's local shard. The payloads
are therefore similar at SP2, while K/V gather moves about 2x as much data at
SP4 and 4x at SP8. K/V gather may still be faster when all-gather and its local
attention layout are more efficient, especially at low SP degrees, but this
scaling makes the interconnect and input shape part of the selection policy.
## Recommended Commands
### Two-GPU Sequence Parallelism
### Ulysses Sequence Parallelism
The default mode needs only the total SP degree when ring parallelism is not
used:
```bash
sglang serve \
--model-path Qwen/Qwen-Image \
--num-gpus 4 \
--sp-degree 4 \
--port 8898
```
### K/V-Gather Sequence Parallelism
Use the same SP process-group layout and select the alternative attention
exchange explicitly:
```bash
sglang serve \
--model-path Qwen/Qwen-Image \
--num-gpus 4 \
--sp-degree 4 \
--sp-attention-mode kv_gather \
--port 8898
```
### Tensor Plus Sequence Parallelism
TP and SP use independent dimensions. With DP and CFG parallelism disabled,
the required GPU count is `tp_size * sp_degree`. This example creates two TP
groups across a two-rank SP dimension:
```bash
sglang serve \
--model-path Qwen/Qwen-Image \
--num-gpus 4 \
--tp-size 2 \
--sp-degree 2 \
--sp-attention-mode kv_gather \
--port 8898
```
Omit `--sp-attention-mode kv_gather` to use TP plus Ulysses with the same
`tp=2, sp=2` topology.
#### How TP Plus SP Works
TP and SP form orthogonal dimensions of the DiT process mesh. For `tp=2,
sp=2`, ranks `[0, 1]` and `[2, 3]` are TP groups, while ranks `[0, 2]` and
`[1, 3]` are SP groups. Each rank therefore belongs to one group of each type:
- TP shards supported attention and MLP projection weights and computation,
then communicates partial projection results inside the TP group.
- SP shards the latent sequence and attention activations, then uses Ulysses
or K/V gather inside the SP group.
Pure SP replicates the DiT weights on every SP rank. Adding TP reduces the
per-rank memory used by TP-sharded weights and keeps the sequence activation
sharding from SP, at the cost of adding TP communication to every applicable
DiT block. The exact memory reduction is model-dependent because not every
parameter or runtime buffer is TP-sharded.
TP plus SP should therefore be treated as a capacity and memory-latency Pareto
option, not as the default latency winner. On a single NVSwitch node, pure SP
often wins when the complete DiT weights fit on every GPU because it avoids
the repeated TP collectives. Try TP plus SP when pure SP does not fit, when
more memory headroom is required, or when its measured memory reduction is
worth a small latency increase.
The following representative eager results used eight H200 GPUs in one
NVSwitch node. Times are median scheduler-side end-to-end latency. They
illustrate the tradeoff rather than define a universal policy:
| Model and workload | Fastest tested topology | TP plus SP Pareto point | Tradeoff |
| --- | --- | --- | --- |
| Qwen-Image, 1536x1536 | CFG2xSP4 Ulysses: 972.6 ms, 62.8 GiB/GPU | CFG2xTP2xSP2 K/V: 1017.4 ms, 48.4 GiB/GPU | 4.6% slower, 22.9% less peak memory |
| Wan2.2-A14B, 832x480x81 | CFG2xSP4 K/V: 6573.5 ms, 61.9 GiB/GPU | CFG2xTP2xSP2 K/V: 7243.1 ms, 34.2 GiB/GPU | 10.2% slower, 44.7% less peak memory |
| LTX2.3, 768x512x241 | SP8 K/V: 7258.2 ms, 55.4 GiB/GPU | TP2xSP4 K/V: 10372.3 ms, 37.8 GiB/GPU | 42.9% slower, 31.8% less peak memory |
K/V gather can still improve TP plus SP at the same topology even when that
topology is not the global latency winner. In the same experiment it improved
TP2xSP4 by 4.2% for FLUX and 8.4% for LTX2.3, and improved
CFG2xTP2xSP2 by 6.0% for Qwen-Image and 2.7% for Wan2.2-A14B, relative to
Ulysses. Always compare the full candidate set, including pure SP, TP, CFG,
and their feasible combinations, rather than selecting the SP attention
backend first.
### FSDP Plus Sequence Parallelism
FSDP can shard DiT weights across the same workers that participate in SP.
Unlike TP times SP, the FSDP and SP degrees do not multiply the required GPU
count. This is useful when pure SP is fast enough but replicated DiT weights
or long-sequence activations leave too little memory headroom:
```bash
sglang serve \
--model-path Lightricks/LTX-2.3 \
--num-gpus 2 \
--use-fsdp-inference true \
--sp-degree 2 \
--sp-attention-mode kv_gather \
--port 8898
```
FSDP adds weight all-gather communication, so compare it with pure SP when
both fit. K/V gather has the same non-causal and `ring_degree=1` constraints
under FSDP.
### Ring Sequence Parallelism
This example uses two GPUs with `sp=2`, `ulysses=1`, and `ring=2`.