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
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
+63 -4
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,12 +1048,35 @@ 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
):
self.ulysses_degree = self.sp_degree
logger.info(
f"Automatically set ulysses_degree=sp_degree={self.ulysses_degree} for best performance"
)
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(
"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:
self.ulysses_degree = 1
@@ -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()