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`.
python/sglang/multimodal_gen/runtime/distributed/sp_shard_utils.py
+5
View File
@@ -218,6 +218,11 @@ def plan_text_strategy(txt_len: int) -> str:
sp_size = get_sp_world_size()
if sp_size <= 1:
return "replicate"
local_len = (txt_len + sp_size - 1) // sp_size
num_pad = local_len * sp_size - txt_len
# padding must fit in the final shard to remain one global-tail block
if num_pad > local_len:
return "replicate"
if txt_len % sp_size != 0 and get_ring_parallel_world_size() > 1:
return "replicate"
if txt_len < _TEXT_SHARD_MIN:
python/sglang/multimodal_gen/runtime/entrypoints/vla/protocol.py
+1
View File
@@ -172,6 +172,7 @@ def action_metadata(server_args: ServerArgs) -> dict[str, Any]:
"sp_degree": server_args.sp_degree,
"ulysses_degree": server_args.ulysses_degree,
"ring_degree": server_args.ring_degree,
"kv_gather_degree": server_args.kv_gather_degree,
"prefix_strategy": pipeline_config.prefix_parallel_strategy,
"action_strategy": pipeline_config.action_parallel_strategy,
"layout_version": pipeline_config.parallel_layout_version,
python/sglang/multimodal_gen/runtime/layers/attention/layer.py
+349 -11
View File
@@ -75,6 +75,57 @@ _PYTORCH_DEFAULT_CUDA_SDP_BACKENDS = [
_VARLEN_FA_ENABLED = os.environ.get("SGLANG_VARLEN_FA", "1") != "0"
def _resolve_sp_attention_mode(
*, causal: bool, sparse_backend: bool
) -> tuple[str, bool]:
"""Resolve one layer's SP exchange; returns (mode, is_auto).
``kv_gather_degree > 1`` selects the gather exchange for the SP rows. When
the degree was auto-assigned, layers the gather path cannot serve fall
back to Ulysses; an explicit degree fails closed instead of degrading.
"""
from sglang.multimodal_gen.runtime.server_args import get_global_server_args
args = get_global_server_args()
if args.kv_gather_degree <= 1:
return "ulysses", False
if causal or sparse_backend:
if args.sp_split_auto:
return "ulysses", True
if causal:
raise ValueError("K/V-gather SP does not support causal attention.")
raise NotImplementedError(
"K/V-gather SP does not support sparse attention backends."
)
return "kv_gather", args.sp_split_auto
def _kv_gather_unsupported_reason(
*,
qkv_pre_all_to_all: bool,
replicated_mode_count: int,
attn_mask: torch.Tensor | None,
num_replicated_kv_prefix: int,
) -> str | None:
"""Call shapes the gather path does not take; explicit mode fails closed
on these, auto falls back to the Ulysses exchange for the call."""
if qkv_pre_all_to_all:
return (
"K/V-gather SP expects sequence-sharded Q/K/V; "
"caller-side pre-all-to-all is Ulysses-only."
)
if replicated_mode_count > 1:
return "K/V-gather SP supports at most one replicated-token mode per call."
if attn_mask is not None:
if num_replicated_kv_prefix:
return "K/V-gather SP masked attention does not support a KV-only prefix."
if attn_mask.dim() != 2:
return "K/V-gather SP masked attention expects a [B, S_local] mask."
if torch.is_floating_point(attn_mask):
return "K/V-gather SP supports boolean or integer padding masks."
return None
def build_varlen_mask_meta(
key_mask: torch.Tensor,
) -> dict:
@@ -265,6 +316,45 @@ class UlyssesAttention(nn.Module):
self.num_kv_heads = num_kv_heads
self.backend = attn_backend.get_enum()
self.dtype = dtype
self.causal = causal
self.sp_attention_mode, self.sp_attention_mode_is_auto = (
_resolve_sp_attention_mode(
causal=causal, sparse_backend=self.backend.is_sparse
)
)
def _forward_with_kv_gather(
self,
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
ctx_attn_metadata,
replicated_q: torch.Tensor | None,
replicated_k: torch.Tensor | None,
replicated_v: torch.Tensor | None,
seq_lens: list[int] | None,
) -> tuple[torch.Tensor, torch.Tensor | None]:
if seq_lens is not None:
raise NotImplementedError(
"K/V-gather SP does not support varlen UlyssesAttention."
)
if any(x is not None for x in (replicated_q, replicated_k, replicated_v)):
if any(x is None for x in (replicated_q, replicated_k, replicated_v)):
raise ValueError("Replicated Q, K, and V must be provided together.")
k = sequence_model_parallel_all_gather(k, dim=1)
v = sequence_model_parallel_all_gather(v, dim=1)
local_query_len = q.shape[1]
if replicated_q is not None:
q = torch.cat([q, replicated_q], dim=1)
k = torch.cat([k, replicated_k], dim=1)
v = torch.cat([v, replicated_v], dim=1)
output = self.attn_impl.forward(q, k, v, ctx_attn_metadata)
if replicated_q is None:
return output, None
return output[:, :local_query_len], output[:, local_query_len:]
def forward(
self,
@@ -294,12 +384,26 @@ class UlyssesAttention(nn.Module):
# Check input shapes
assert q.dim() == 4 and k.dim() == 4 and v.dim() == 4, "Expected 4D tensors"
batch_size, seq_len, num_heads, head_dim = q.shape
local_rank = get_sp_parallel_rank()
world_size = get_sp_world_size()
forward_context: ForwardContext = get_forward_context()
ctx_attn_metadata = forward_context.attn_metadata
if self.sp_attention_mode == "kv_gather" and not (
self.sp_attention_mode_is_auto and seq_lens is not None
):
return self._forward_with_kv_gather(
q,
k,
v,
ctx_attn_metadata,
replicated_q,
replicated_k,
replicated_v,
seq_lens,
)
local_rank = get_sp_parallel_rank()
world_size = get_sp_world_size()
if seq_lens is not None:
assert (
replicated_q is None and replicated_k is None and replicated_v is None
@@ -384,6 +488,10 @@ class UlyssesAttention_VSA(UlyssesAttention):
- o (torch.Tensor): Output tensor after attention for the main sequence
- replicated_o (Optional[torch.Tensor]): Output tensor for replicated tokens, if provided
"""
if self.sp_attention_mode == "kv_gather":
raise NotImplementedError(
"K/V-gather SP does not support video sparse attention."
)
# Check text tokens are not supported for VSA now
assert (
replicated_q is None and replicated_k is None and replicated_v is None
@@ -535,11 +643,12 @@ class LocalAttention(nn.Module):
class USPAttention(nn.Module):
"""
Ulysses Sequence Parallelism with Ring Attention.
Sequence-parallel attention with Ulysses, K/V gather, and Ring Attention.
This class implements the USP algorithm, which is a combination of
Ulysses-style all-to-all communication for sequence-head dimension sharding
and Ring Attention for fine-grained sequence parallelism within subgroups.
The default path implements USP, which combines Ulysses-style all-to-all
communication for sequence-head dimension sharding with Ring Attention
inside subgroups. The K/V-gather path keeps queries sequence-sharded and
gathers keys and values within the SP group.
"""
_usp_a2a_stream = None
@@ -612,6 +721,11 @@ class USPAttention(nn.Module):
self.skip_sequence_parallel = skip_sequence_parallel
self.enable_packed_qkv_input_a2a = bool(enable_packed_qkv_input_a2a)
self.sp_attention_mode, self.sp_attention_mode_is_auto = (
_resolve_sp_attention_mode(
causal=causal, sparse_backend=self.backend.is_sparse
)
)
def _get_usp_a2a_stream(self):
if USPAttention._usp_a2a_stream is None:
@@ -681,6 +795,40 @@ class USPAttention(nn.Module):
and not effective_skip_sp
and get_sequence_parallel_world_size() > 1
)
replicated_mode_count = sum(
value > 0
for value in (
num_replicated_prefix,
num_replicated_suffix,
num_replicated_kv_prefix,
)
)
if (
self.sp_attention_mode == "kv_gather"
and not effective_skip_sp
and get_sequence_parallel_world_size() > 1
):
unsupported = _kv_gather_unsupported_reason(
qkv_pre_all_to_all=qkv_pre_all_to_all,
replicated_mode_count=replicated_mode_count,
attn_mask=attn_mask,
num_replicated_kv_prefix=num_replicated_kv_prefix,
)
if unsupported is None:
return self._forward_with_kv_gather(
q,
k,
v,
ctx_attn_metadata,
attn_mask,
attn_mask_meta,
num_replicated_prefix,
num_replicated_suffix,
num_replicated_kv_prefix,
)
if not self.sp_attention_mode_is_auto:
raise NotImplementedError(unsupported)
if attn_mask is not None or meta_only_pad:
def _prepare_sdpa_mask(
@@ -944,11 +1092,7 @@ class USPAttention(nn.Module):
return out
sp_size = get_ulysses_parallel_world_size()
if (
(num_replicated_prefix > 0 and num_replicated_suffix > 0)
or (num_replicated_prefix > 0 and num_replicated_kv_prefix > 0)
or (num_replicated_suffix > 0 and num_replicated_kv_prefix > 0)
):
if replicated_mode_count > 1:
raise ValueError(
"USPAttention supports at most one replicated-token mode per call."
)
@@ -1003,6 +1147,193 @@ class USPAttention(nn.Module):
return out
@staticmethod
def _gather_sharded_sequence(
tensor: torch.Tensor,
num_replicated_prefix: int = 0,
num_replicated_suffix: int = 0,
) -> torch.Tensor:
if num_replicated_prefix and num_replicated_suffix:
raise ValueError(
"Replicated prefix and suffix cannot be used at the same time."
)
if num_replicated_prefix:
replicated = tensor[:, :num_replicated_prefix]
sharded = tensor[:, num_replicated_prefix:]
gathered = sequence_model_parallel_all_gather(sharded, dim=1)
return torch.cat([replicated, gathered], dim=1)
if num_replicated_suffix:
replicated = tensor[:, -num_replicated_suffix:]
sharded = tensor[:, :-num_replicated_suffix]
gathered = sequence_model_parallel_all_gather(sharded, dim=1)
return torch.cat([gathered, replicated], dim=1)
return sequence_model_parallel_all_gather(tensor, dim=1)
def _forward_with_kv_gather(
self,
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
ctx_attn_metadata,
attn_mask: torch.Tensor | None,
attn_mask_meta: dict | None,
num_replicated_prefix: int,
num_replicated_suffix: int,
num_replicated_kv_prefix: int,
) -> torch.Tensor:
if attn_mask is not None and num_replicated_kv_prefix:
raise NotImplementedError(
"K/V-gather SP masked attention does not support a KV-only prefix."
)
kv_prefix = num_replicated_prefix or num_replicated_kv_prefix
k = self._gather_sharded_sequence(
k,
num_replicated_prefix=kv_prefix,
num_replicated_suffix=num_replicated_suffix,
)
v = self._gather_sharded_sequence(
v,
num_replicated_prefix=kv_prefix,
num_replicated_suffix=num_replicated_suffix,
)
if attn_mask is None and attn_mask_meta is None:
return self.attn_impl.forward(q, k, v, ctx_attn_metadata)
explicit_mask = attn_mask is not None
if attn_mask is None:
local_pad = int(attn_mask_meta.get("local_pad", 0))
attn_mask = torch.ones(q.shape[:2], dtype=torch.bool, device=q.device)
if local_pad:
attn_mask[:, -local_pad:] = False
elif attn_mask.dim() != 2:
raise NotImplementedError(
"K/V-gather SP masked attention expects a [B, S_local] mask."
)
else:
if attn_mask.dtype not in (
torch.bool,
torch.uint8,
torch.int32,
torch.int64,
):
raise NotImplementedError(
"K/V-gather SP supports boolean or integer padding masks."
)
attn_mask = attn_mask.to(dtype=torch.bool)
cache_key = (
tuple(attn_mask.shape),
num_replicated_prefix,
num_replicated_suffix,
explicit_mask,
)
mask_cache = None
if attn_mask_meta is not None:
mask_cache = attn_mask_meta.get("_kv_gather_cache")
cached = mask_cache.get(cache_key) if mask_cache is not None else None
if cached is None:
key_mask = self._gather_sharded_sequence(
attn_mask,
num_replicated_prefix=num_replicated_prefix,
num_replicated_suffix=num_replicated_suffix,
)
cached = {"key_mask": key_mask}
if attn_mask_meta is not None:
if mask_cache is None:
mask_cache = {}
attn_mask_meta["_kv_gather_cache"] = mask_cache
mask_cache[cache_key] = cached
else:
key_mask = cached["key_mask"]
if (
_VARLEN_FA_ENABLED
and self.backend == AttentionBackendEnum.FA
and q.device.type == "cuda"
and q.dtype in (torch.float16, torch.bfloat16)
):
if (
explicit_mask
and attn_mask_meta is not None
and all(
key in attn_mask_meta
for key in (
"indices",
"cu_seqlens",
"max_seqlen",
"inv_indices",
)
)
):
query_meta = attn_mask_meta
else:
query_meta = cached.get("query_meta")
if query_meta is None:
query_meta = build_varlen_mask_meta(attn_mask)
cached["query_meta"] = query_meta
key_meta = cached.get("key_meta")
if key_meta is None:
key_meta = build_varlen_mask_meta(key_mask)
cached["key_meta"] = key_meta
batch_size, query_len = q.shape[:2]
q_unpad = q.reshape(-1, *q.shape[2:]).index_select(0, query_meta["indices"])
k_unpad = k.reshape(-1, *k.shape[2:]).index_select(0, key_meta["indices"])
v_unpad = v.reshape(-1, *v.shape[2:]).index_select(0, key_meta["indices"])
out_unpad = flash_attn_varlen_func(
q=q_unpad,
k=k_unpad,
v=v_unpad,
cu_seqlens_q=query_meta["cu_seqlens"],
cu_seqlens_k=key_meta["cu_seqlens"],
max_seqlen_q=query_meta["max_seqlen"],
max_seqlen_k=key_meta["max_seqlen"],
softmax_scale=self.softmax_scale,
causal=False,
ver=_fa_backend.fa_ver,
)
return fused_scatter_to_padded(
out_unpad,
query_meta["inv_indices"],
batch_size,
query_len,
)
q_ = q.transpose(1, 2)
k_ = k.transpose(1, 2)
v_ = v.transpose(1, 2)
if q_.shape[1] != k_.shape[1]:
if q_.shape[1] % k_.shape[1] != 0:
raise ValueError(
f"Query heads ({q_.shape[1]}) must be divisible by "
f"KV heads ({k_.shape[1]})."
)
repeat_factor = q_.shape[1] // k_.shape[1]
k_ = k_.repeat_interleave(repeat_factor, dim=1)
v_ = v_.repeat_interleave(repeat_factor, dim=1)
sdpa_context = (
sdpa_kernel(_PYTORCH_DEFAULT_CUDA_SDP_BACKENDS)
if self.allow_cudnn_sdp and q_.device.type == "cuda"
else nullcontext()
)
with sdpa_context:
out = torch.nn.functional.scaled_dot_product_attention(
q_,
k_,
v_,
attn_mask=key_mask[:, None, None, :],
dropout_p=0.0,
is_causal=False,
scale=self.softmax_scale,
).transpose(1, 2)
return out * attn_mask[:, :, None, None]
def _forward_with_replicated_prefix(
self,
q: torch.Tensor,
@@ -1087,6 +1418,13 @@ class USPAttention(nn.Module):
v = torch.cat([v_prefix, v_suffix], dim=1)
return self.attn_impl.forward(q, k, v, ctx_attn_metadata)
if self.sp_attention_mode == "kv_gather":
k_suffix = sequence_model_parallel_all_gather(k_suffix, dim=1)
v_suffix = sequence_model_parallel_all_gather(v_suffix, dim=1)
k = torch.cat([k_prefix, k_suffix], dim=1)
v = torch.cat([v_prefix, v_suffix], dim=1)
return self.attn_impl.forward(q, k, v, ctx_attn_metadata)
if get_ulysses_parallel_world_size() == 1:
k = torch.cat([k_prefix, k_suffix], dim=1)
v = torch.cat([v_prefix, v_suffix], dim=1)
python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py
+3 -1
View File
@@ -768,7 +768,9 @@ class QwenImageCrossAttention(nn.Module):
# Joint order [text, image]; join_seqs relocates any SP text tail-pad
# behind the image (see sp_shard.join_seqs for why).
seg_qkv = None
if sp_text_sharded:
# The segmented pre-all-to-all emits Ulysses layout; K/V-gather takes
# the join_seqs path and exchanges inside the attention instead.
if sp_text_sharded and self.attn.sp_attention_mode == "ulysses":
from sglang.multimodal_gen.runtime.layers.usp import (
_ipc_input_a2a_qkv_segmented,
)
python/sglang/multimodal_gen/runtime/server_args/server_args.py
+60 -1
View File
@@ -219,6 +219,12 @@ class ServerArgs(DisaggServerArgsMixin):
# sequence parallelism
ulysses_degree: Optional[int] = None
ring_degree: Optional[int] = None
# rows split inside attention, exchanged with one K/V all-gather instead
# of Ulysses a2a or ring rotation; auto-assigned at sp_degree=2 when no SP
# degree is set explicitly
kv_gather_degree: Optional[int] = None
# whether the SP split was auto-assigned (lets layers fall back per call)
sp_split_auto: bool = False
# data parallelism
# number of data parallelism groups
dp_size: int = 1
@@ -1042,11 +1048,34 @@ class ServerArgs(DisaggServerArgsMixin):
if (
self.ulysses_degree is None
and self.ring_degree is None
and self.kv_gather_degree is None
and self.sp_degree != 1
):
if self.sp_degree == 2:
# measured-win zone for the K/V-gather exchange; layers whose
# calls the gather path cannot take fall back to Ulysses
self.kv_gather_degree = 2
self.sp_split_auto = True
logger.info(
"Automatically set kv_gather_degree=sp_degree=2; set "
"--ulysses-degree explicitly to keep the Ulysses exchange"
)
else:
self.ulysses_degree = self.sp_degree
logger.info(
f"Automatically set ulysses_degree=sp_degree={self.ulysses_degree} for best performance"
"Automatically set ulysses_degree=sp_degree=%d for the "
"sequence-parallel process-group layout",
self.ulysses_degree,
)
if self.kv_gather_degree is None:
self.kv_gather_degree = 1
if self.kv_gather_degree > 1:
if (self.ulysses_degree or 1) != 1 or (self.ring_degree or 1) != 1:
raise ValueError(
"kv_gather_degree does not compose with ulysses_degree or "
"ring_degree yet; set exactly one of them above 1"
)
if self.ulysses_degree is None:
@@ -1059,6 +1088,13 @@ class ServerArgs(DisaggServerArgsMixin):
self.ring_degree = 1
logger.debug(f"Ring degree not set, using default value {self.ring_degree}")
if self.kv_gather_degree > 1:
# K/V-gather rows occupy the contiguous inner SP dimension; the
# process groups are built from ulysses_degree, so alias it until
# gather gets a first-class dimension (needed only once it
# composes with Ulysses).
self.ulysses_degree = self.kv_gather_degree
def _model_default_uses_cfg(self) -> bool:
"""
Check whether the model uses classifier-free guidance by default.
@@ -1451,6 +1487,20 @@ class ServerArgs(DisaggServerArgsMixin):
"`replicate` disables both. The default is `auto`."
),
)
parser.add_argument(
"--kv-gather-degree",
type=int,
default=ServerArgs.kv_gather_degree,
help=(
"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."
),
)
parser.add_argument(
"--enable-cfg-parallel",
action=StoreBoolean,
@@ -2398,6 +2448,15 @@ class ServerArgs(DisaggServerArgsMixin):
)
def _validate_parallelism(self):
if self.kv_gather_degree < 1:
raise ValueError("kv_gather_degree must be >= 1")
if self.kv_gather_degree > 1 and self.sp_degree != self.kv_gather_degree:
raise ValueError(
f"kv_gather_degree ({self.kv_gather_degree}) must equal "
f"sp_degree ({self.sp_degree}); check how many GPUs remain for "
"sequence parallelism after dp/tp/cfg"
)
if self.sp_degree > self.num_gpus or self.num_gpus % self.sp_degree != 0:
raise ValueError(
f"num_gpus ({self.num_gpus}) must be >= and divisible by sp_degree ({self.sp_degree})"
python/sglang/multimodal_gen/test/unit/conftest.py
+2
View File
@@ -41,6 +41,8 @@ def _make_unit_server_args():
enable_breakable_cuda_graph=False,
enable_layerwise_nvtx_marker=False,
enable_torch_compile=False,
kv_gather_degree=1,
sp_split_auto=False,
model_loaded={},
model_paths={},
pipeline_config=pipeline_config,
python/sglang/multimodal_gen/test/unit/test_ideogram4.py
+56 -9
View File
@@ -182,6 +182,8 @@ def _fake_server_args(cfg=None):
disable_autocast=False,
enable_cfg_parallel=False,
attention_backend_config=None,
kv_gather_degree=1,
sp_split_auto=False,
)
@@ -357,7 +359,12 @@ class TestIdeogram4(unittest.TestCase):
prev_args = server_args_module._global_server_args
try:
set_global_server_args(
SimpleNamespace(attention_backend="torch_sdpa", comfyui_mode=False)
SimpleNamespace(
attention_backend="torch_sdpa",
comfyui_mode=False,
kv_gather_degree=1,
sp_split_auto=False,
)
)
torch.manual_seed(0)
batch_size, seq_len, num_heads, head_dim = 2, 5, 2, 8
@@ -769,7 +776,12 @@ class TestIdeogram4(unittest.TestCase):
prev_args = server_args_module._global_server_args
try:
set_global_server_args(
SimpleNamespace(attention_backend="torch_sdpa", comfyui_mode=False)
SimpleNamespace(
attention_backend="torch_sdpa",
comfyui_mode=False,
kv_gather_degree=1,
sp_split_auto=False,
)
)
with patch(
"sglang.multimodal_gen.runtime.layers.attention.layer.get_ring_parallel_world_size",
@@ -794,7 +806,12 @@ class TestIdeogram4(unittest.TestCase):
prev_args = server_args_module._global_server_args
try:
set_global_server_args(
SimpleNamespace(attention_backend="torch_sdpa", comfyui_mode=False)
SimpleNamespace(
attention_backend="torch_sdpa",
comfyui_mode=False,
kv_gather_degree=1,
sp_split_auto=False,
)
)
with patch(
"sglang.multimodal_gen.runtime.layers.attention.layer.get_ring_parallel_world_size",
@@ -842,7 +859,12 @@ class TestIdeogram4(unittest.TestCase):
prev_args = server_args_module._global_server_args
try:
set_global_server_args(
SimpleNamespace(attention_backend="torch_sdpa", comfyui_mode=False)
SimpleNamespace(
attention_backend="torch_sdpa",
comfyui_mode=False,
kv_gather_degree=1,
sp_split_auto=False,
)
)
with (
patch(
@@ -888,7 +910,12 @@ class TestIdeogram4(unittest.TestCase):
prev_args = server_args_module._global_server_args
try:
set_global_server_args(
SimpleNamespace(attention_backend="torch_sdpa", comfyui_mode=False)
SimpleNamespace(
attention_backend="torch_sdpa",
comfyui_mode=False,
kv_gather_degree=1,
sp_split_auto=False,
)
)
with (
patch(
@@ -941,7 +968,12 @@ class TestIdeogram4(unittest.TestCase):
prev_args = server_args_module._global_server_args
try:
set_global_server_args(
SimpleNamespace(attention_backend="torch_sdpa", comfyui_mode=False)
SimpleNamespace(
attention_backend="torch_sdpa",
comfyui_mode=False,
kv_gather_degree=1,
sp_split_auto=False,
)
)
with patch(
"sglang.multimodal_gen.runtime.layers.attention.layer.get_ring_parallel_world_size",
@@ -997,7 +1029,12 @@ class TestIdeogram4(unittest.TestCase):
prev_args = server_args_module._global_server_args
try:
set_global_server_args(
SimpleNamespace(attention_backend="torch_sdpa", comfyui_mode=False)
SimpleNamespace(
attention_backend="torch_sdpa",
comfyui_mode=False,
kv_gather_degree=1,
sp_split_auto=False,
)
)
with (
patch(
@@ -1210,7 +1247,12 @@ class TestIdeogram4(unittest.TestCase):
prev_args = server_args_module._global_server_args
try:
set_global_server_args(
SimpleNamespace(attention_backend="torch_sdpa", comfyui_mode=False)
SimpleNamespace(
attention_backend="torch_sdpa",
comfyui_mode=False,
kv_gather_degree=1,
sp_split_auto=False,
)
)
with (
patch.dict(os.environ, {W8A8_FP8_GEMM_ENV: "1"}),
@@ -1250,7 +1292,12 @@ class TestIdeogram4(unittest.TestCase):
prev_args = server_args_module._global_server_args
try:
set_global_server_args(
SimpleNamespace(attention_backend="torch_sdpa", comfyui_mode=False)
SimpleNamespace(
attention_backend="torch_sdpa",
comfyui_mode=False,
kv_gather_degree=1,
sp_split_auto=False,
)
)
with (
patch(
python/sglang/multimodal_gen/test/unit/test_pi05_action_api.py
+2
View File
@@ -33,6 +33,7 @@ def _server_args(config: Pi05PipelineConfig | None = None) -> SimpleNamespace:
sp_degree=1,
ulysses_degree=1,
ring_degree=1,
kv_gather_degree=1,
pipeline_config=config or Pi05PipelineConfig(),
)
@@ -174,6 +175,7 @@ def test_action_metadata_reports_policy_shape_and_capabilities():
assert metadata["runtime"]["materialize_dtype"] == "bf16"
assert metadata["runtime"]["enable_autocast"] is True
assert metadata["runtime"]["parallelism"]["num_gpus"] == 1
assert metadata["runtime"]["parallelism"]["kv_gather_degree"] == 1
assert metadata["runtime"]["parallelism"]["prefix_strategy"] == "tp"
assert metadata["runtime"]["parallelism"]["action_strategy"] == "sp"
assert metadata["defaults"]["prefix_cache"] is False
python/sglang/multimodal_gen/test/unit/test_server_args.py
+101
View File
@@ -1760,6 +1760,107 @@ class TestOffloadDefaults(unittest.TestCase):
self.assertEqual(server_args.ltx2_two_stage_device_mode, "original")
class TestKVGatherDegree(unittest.TestCase):
def test_sp2_defaults_to_kv_gather(self):
args = _from_dict_without_model_resolution(
{
"model_path": "/fake",
"num_gpus": 2,
"performance_mode": "manual",
}
)
self.assertEqual(args.kv_gather_degree, 2)
self.assertTrue(args.sp_split_auto)
# gather rows occupy the contiguous inner SP dimension
self.assertEqual(args.ulysses_degree, 2)
self.assertEqual(args.sp_degree, 2)
def test_higher_sp_defaults_to_ulysses(self):
args = _from_dict_without_model_resolution(
{
"model_path": "/fake",
"num_gpus": 4,
"performance_mode": "manual",
}
)
self.assertEqual(args.kv_gather_degree, 1)
self.assertFalse(args.sp_split_auto)
self.assertEqual(args.ulysses_degree, 4)
def test_explicit_ulysses_is_not_overridden(self):
args = _from_dict_without_model_resolution(
{
"model_path": "/fake",
"num_gpus": 2,
"ulysses_degree": 2,
"performance_mode": "manual",
}
)
self.assertEqual(args.kv_gather_degree, 1)
self.assertEqual(args.ulysses_degree, 2)
def test_explicit_degree_is_not_auto(self):
args = _from_dict_without_model_resolution(
{
"model_path": "/fake",
"num_gpus": 2,
"kv_gather_degree": 2,
"performance_mode": "manual",
}
)
self.assertEqual(args.kv_gather_degree, 2)
self.assertFalse(args.sp_split_auto)
def test_kv_gather_supports_tp(self):
args = _from_dict_without_model_resolution(
{
"model_path": "/fake",
"num_gpus": 4,
"tp_size": 2,
"sp_degree": 2,
"kv_gather_degree": 2,
"performance_mode": "manual",
}
)
self.assertEqual(args.tp_size, 2)
self.assertEqual(args.sp_degree, 2)
self.assertEqual(args.kv_gather_degree, 2)
def test_kv_gather_supports_fsdp(self):
args = _from_dict_without_model_resolution(
{
"model_path": "/fake",
"num_gpus": 2,
"sp_degree": 2,
"kv_gather_degree": 2,
"use_fsdp_inference": True,
"performance_mode": "manual",
}
)
self.assertTrue(args.use_fsdp_inference)
self.assertEqual(args.kv_gather_degree, 2)
def test_kv_gather_does_not_compose_yet(self):
for extra in ({"ulysses_degree": 2}, {"ring_degree": 2}):
with self.assertRaisesRegex(ValueError, "does not compose"):
_from_dict_without_model_resolution(
{
"model_path": "/fake",
"num_gpus": 4,
"sp_degree": 4,
"kv_gather_degree": 2,
"performance_mode": "manual",
**extra,
}
)
class TestFSDPShardConditions(unittest.TestCase):
def test_helpers_match_only_direct_block_entries(self):
self.assertTrue(
python/sglang/multimodal_gen/test/unit/test_sp_shard.py
+10
View File
@@ -140,6 +140,16 @@ def test_strategy_shard_when_legal(monkeypatch):
assert sps.plan_text_strategy(16) == "shard"
def test_strategy_replicates_when_padding_spans_multiple_shards(monkeypatch):
_fake_sp(monkeypatch, 8)
assert sps.plan_text_strategy(1) == "replicate"
assert sps.plan_text_strategy(6) == "replicate"
assert sps.plan_text_strategy(7) == "shard"
assert sps.plan_text_strategy(9) == "replicate"
assert sps.plan_text_strategy(13) == "replicate"
assert sps.plan_text_strategy(14) == "shard"
def test_strategy_ring_blocks_padded_shard(monkeypatch):
_fake_sp(monkeypatch, 2, ring=2)
assert sps.plan_text_strategy(15) == "replicate" # padded shard needs mask
python/sglang/multimodal_gen/test/unit/test_usp_attention_kv_gather.py
+348
View File
@@ -0,0 +1,348 @@
import unittest
from types import SimpleNamespace
from unittest.mock import patch
import torch
import torch.nn.functional as F
from sglang.multimodal_gen.runtime.layers.attention.layer import (
UlyssesAttention,
UlyssesAttention_VSA,
USPAttention,
_kv_gather_unsupported_reason,
_resolve_sp_attention_mode,
)
from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
_LAYER = "sglang.multimodal_gen.runtime.layers.attention.layer"
class _SdpaAttention:
def __init__(self, scale: float):
self.scale = scale
def forward(self, q, k, v, _ctx):
return F.scaled_dot_product_attention(
q.transpose(1, 2),
k.transpose(1, 2),
v.transpose(1, 2),
dropout_p=0.0,
is_causal=False,
scale=self.scale,
).transpose(1, 2)
def _make_attention(head_dim: int) -> USPAttention:
obj = USPAttention.__new__(USPAttention)
obj.causal = False
obj.backend = AttentionBackendEnum.TORCH_SDPA
obj.softmax_scale = head_dim**-0.5
obj.attn_impl = _SdpaAttention(obj.softmax_scale)
obj.allow_cudnn_sdp = False
obj.skip_sequence_parallel = False
obj.sp_attention_mode = "kv_gather"
obj.sp_attention_mode_is_auto = False
return obj
def _reference_attention(q, k, v, scale, key_mask=None, query_mask=None):
out = F.scaled_dot_product_attention(
q.transpose(1, 2),
k.transpose(1, 2),
v.transpose(1, 2),
attn_mask=None if key_mask is None else key_mask[:, None, None, :],
dropout_p=0.0,
is_causal=False,
scale=scale,
).transpose(1, 2)
if query_mask is not None:
out = out * query_mask[:, :, None, None]
return out
class TestUSPAttentionKVGather(unittest.TestCase):
def setUp(self):
torch.manual_seed(0)
self.heads = 3
self.head_dim = 4
self.attn = _make_attention(self.head_dim)
def _run(self, q, k, v, gathered, **kwargs):
with (
patch(f"{_LAYER}.get_ring_parallel_world_size", return_value=1),
patch(
f"{_LAYER}.sequence_model_parallel_all_gather",
side_effect=gathered,
),
):
return self.attn._forward_with_kv_gather(
q,
k,
v,
None,
kwargs.pop("attn_mask", None),
kwargs.pop("attn_mask_meta", None),
kwargs.pop("num_replicated_prefix", 0),
kwargs.pop("num_replicated_suffix", 0),
kwargs.pop("num_replicated_kv_prefix", 0),
)
def test_local_queries_attend_gathered_kv(self):
q = torch.randn(1, 3, self.heads, self.head_dim)
k = torch.randn(1, 3, self.heads, self.head_dim)
v = torch.randn(1, 3, self.heads, self.head_dim)
full_k = torch.randn(1, 6, self.heads, self.head_dim)
full_v = torch.randn(1, 6, self.heads, self.head_dim)
out = self._run(q, k, v, [full_k, full_v])
expected = _reference_attention(q, full_k, full_v, self.attn.softmax_scale)
torch.testing.assert_close(out, expected)
def test_replicated_prefix_is_not_duplicated(self):
prefix = 2
q = torch.randn(1, 4, self.heads, self.head_dim)
k = torch.randn(1, 4, self.heads, self.head_dim)
v = torch.randn(1, 4, self.heads, self.head_dim)
gathered_k_suffix = torch.randn(1, 4, self.heads, self.head_dim)
gathered_v_suffix = torch.randn(1, 4, self.heads, self.head_dim)
full_k = torch.cat([k[:, :prefix], gathered_k_suffix], dim=1)
full_v = torch.cat([v[:, :prefix], gathered_v_suffix], dim=1)
out = self._run(
q,
k,
v,
[gathered_k_suffix, gathered_v_suffix],
num_replicated_prefix=prefix,
)
expected = _reference_attention(q, full_k, full_v, self.attn.softmax_scale)
torch.testing.assert_close(out, expected)
def test_padding_mask_uses_local_queries_and_global_keys(self):
q = torch.randn(2, 3, self.heads, self.head_dim)
k = torch.randn(2, 3, self.heads, self.head_dim)
v = torch.randn(2, 3, self.heads, self.head_dim)
full_k = torch.randn(2, 6, self.heads, self.head_dim)
full_v = torch.randn(2, 6, self.heads, self.head_dim)
query_mask = torch.tensor([[True, True, False], [True, True, True]])
key_mask = torch.tensor(
[
[True, True, False, True, False, False],
[True, True, True, True, True, False],
]
)
out = self._run(
q,
k,
v,
[full_k, full_v, key_mask],
attn_mask=query_mask,
attn_mask_meta={},
)
expected = _reference_attention(
q,
full_k,
full_v,
self.attn.softmax_scale,
key_mask=key_mask,
query_mask=query_mask,
)
torch.testing.assert_close(out, expected)
def test_separate_replicated_kv_prefix_gathers_only_suffix(self):
q = torch.randn(1, 3, self.heads, self.head_dim)
k_prefix = torch.randn(1, 2, self.heads, self.head_dim)
v_prefix = torch.randn(1, 2, self.heads, self.head_dim)
k_suffix = torch.randn(1, 3, self.heads, self.head_dim)
v_suffix = torch.randn(1, 3, self.heads, self.head_dim)
gathered_k_suffix = torch.randn(1, 6, self.heads, self.head_dim)
gathered_v_suffix = torch.randn(1, 6, self.heads, self.head_dim)
full_k = torch.cat([k_prefix, gathered_k_suffix], dim=1)
full_v = torch.cat([v_prefix, gathered_v_suffix], dim=1)
with (
patch(
f"{_LAYER}.get_forward_context",
return_value=SimpleNamespace(attn_metadata=None),
),
patch(f"{_LAYER}.get_sequence_parallel_world_size", return_value=2),
patch(f"{_LAYER}.get_ring_parallel_world_size", return_value=1),
patch(
f"{_LAYER}.sequence_model_parallel_all_gather",
side_effect=[gathered_k_suffix, gathered_v_suffix],
),
):
out = self.attn.forward_with_replicated_kv_prefix(
q, k_prefix, v_prefix, k_suffix, v_suffix
)
expected = _reference_attention(q, full_k, full_v, self.attn.softmax_scale)
torch.testing.assert_close(out, expected)
class TestUlyssesAttentionKVGather(unittest.TestCase):
def setUp(self):
torch.manual_seed(1)
self.heads = 3
self.head_dim = 4
self.attn = UlyssesAttention.__new__(UlyssesAttention)
self.attn.causal = False
self.attn.backend = AttentionBackendEnum.TORCH_SDPA
self.attn.softmax_scale = self.head_dim**-0.5
self.attn.attn_impl = _SdpaAttention(self.attn.softmax_scale)
self.attn.sp_attention_mode = "kv_gather"
self.attn.sp_attention_mode_is_auto = False
def _run(self, q, k, v, gathered, **kwargs):
with (
patch(
f"{_LAYER}.get_forward_context",
return_value=SimpleNamespace(attn_metadata=None),
),
patch(f"{_LAYER}.get_ring_parallel_world_size", return_value=1),
patch(
f"{_LAYER}.sequence_model_parallel_all_gather",
side_effect=gathered,
),
):
return self.attn.forward(
q,
k,
v,
kwargs.get("replicated_q"),
kwargs.get("replicated_k"),
kwargs.get("replicated_v"),
kwargs.get("seq_lens"),
)
def test_local_queries_attend_gathered_kv(self):
q = torch.randn(1, 3, self.heads, self.head_dim)
k = torch.randn(1, 3, self.heads, self.head_dim)
v = torch.randn(1, 3, self.heads, self.head_dim)
full_k = torch.randn(1, 6, self.heads, self.head_dim)
full_v = torch.randn(1, 6, self.heads, self.head_dim)
out, replicated_out = self._run(q, k, v, [full_k, full_v])
expected = _reference_attention(q, full_k, full_v, self.attn.softmax_scale)
torch.testing.assert_close(out, expected)
self.assertIsNone(replicated_out)
def test_replicated_suffix_is_computed_without_head_sharding(self):
q = torch.randn(1, 3, self.heads, self.head_dim)
k = torch.randn(1, 3, self.heads, self.head_dim)
v = torch.randn(1, 3, self.heads, self.head_dim)
replicated_q = torch.randn(1, 2, self.heads, self.head_dim)
replicated_k = torch.randn(1, 2, self.heads, self.head_dim)
replicated_v = torch.randn(1, 2, self.heads, self.head_dim)
full_k = torch.randn(1, 6, self.heads, self.head_dim)
full_v = torch.randn(1, 6, self.heads, self.head_dim)
out, replicated_out = self._run(
q,
k,
v,
[full_k, full_v],
replicated_q=replicated_q,
replicated_k=replicated_k,
replicated_v=replicated_v,
)
full_q = torch.cat([q, replicated_q], dim=1)
full_k = torch.cat([full_k, replicated_k], dim=1)
full_v = torch.cat([full_v, replicated_v], dim=1)
expected = _reference_attention(full_q, full_k, full_v, self.attn.softmax_scale)
torch.testing.assert_close(out, expected[:, : q.shape[1]])
torch.testing.assert_close(replicated_out, expected[:, q.shape[1] :])
def test_varlen_is_rejected(self):
q = torch.randn(1, 3, self.heads, self.head_dim)
with self.assertRaisesRegex(NotImplementedError, "varlen"):
self._run(q, q, q, [], seq_lens=[3, 3])
def test_video_sparse_attention_is_rejected(self):
attn = UlyssesAttention_VSA.__new__(UlyssesAttention_VSA)
attn.sp_attention_mode = "kv_gather"
q = torch.randn(1, 3, self.heads, self.head_dim)
with self.assertRaisesRegex(NotImplementedError, "video sparse"):
attn.forward(q, q, q, gate_compress=q)
class TestSpAttentionModeResolution(unittest.TestCase):
def _resolve(self, *, degree=2, auto=True, causal=False, sparse=False):
stub = SimpleNamespace(kv_gather_degree=degree, sp_split_auto=auto)
with patch(
"sglang.multimodal_gen.runtime.server_args.get_global_server_args",
return_value=stub,
):
return _resolve_sp_attention_mode(causal=causal, sparse_backend=sparse)
def test_gather_degree_selects_the_gather_exchange(self):
self.assertEqual(self._resolve(), ("kv_gather", True))
self.assertEqual(self._resolve(auto=False), ("kv_gather", False))
def test_degree_one_is_plain_ulysses(self):
self.assertEqual(self._resolve(degree=1), ("ulysses", False))
self.assertEqual(self._resolve(degree=1, causal=True), ("ulysses", False))
def test_auto_degree_falls_back_for_unsupported_layers(self):
self.assertEqual(self._resolve(causal=True), ("ulysses", True))
self.assertEqual(self._resolve(sparse=True), ("ulysses", True))
def test_explicit_degree_fails_closed(self):
with self.assertRaises(ValueError):
self._resolve(auto=False, causal=True)
with self.assertRaises(NotImplementedError):
self._resolve(auto=False, sparse=True)
class TestKVGatherCallSupport(unittest.TestCase):
def _reason(self, **overrides):
kwargs = dict(
qkv_pre_all_to_all=False,
replicated_mode_count=0,
attn_mask=None,
num_replicated_kv_prefix=0,
)
kwargs.update(overrides)
return _kv_gather_unsupported_reason(**kwargs)
def test_plain_and_masked_calls_are_supported(self):
self.assertIsNone(self._reason())
self.assertIsNone(self._reason(attn_mask=torch.ones(1, 4, dtype=torch.bool)))
def test_unsupported_shapes_are_reported(self):
self.assertIn("pre-all-to-all", self._reason(qkv_pre_all_to_all=True))
self.assertIn("replicated-token", self._reason(replicated_mode_count=2))
self.assertIn(
"KV-only prefix",
self._reason(
attn_mask=torch.ones(1, 4, dtype=torch.bool),
num_replicated_kv_prefix=2,
),
)
self.assertIn("[B, S_local]", self._reason(attn_mask=torch.ones(1, 1, 4)))
self.assertIn("integer padding", self._reason(attn_mask=torch.ones(1, 4)))
def test_explicit_mode_raises_and_auto_falls_back_at_dispatch(self):
attn = _make_attention(4)
attn.skip_sequence_parallel = False
q = torch.randn(1, 4, 3, 4)
with (
patch(
f"{_LAYER}.get_forward_context",
return_value=SimpleNamespace(attn_metadata=None),
),
patch(f"{_LAYER}.get_sequence_parallel_world_size", return_value=2),
):
attn.sp_attention_mode_is_auto = False
with self.assertRaisesRegex(NotImplementedError, "pre-all-to-all"):
attn.forward(q, q, q, qkv_pre_all_to_all=True)
if __name__ == "__main__":
unittest.main()