diff --git a/docs/docs/sglang-diffusion/api/cli.mdx b/docs/docs/sglang-diffusion/api/cli.mdx index 4a96e2d3b..2a80d595f 100644 --- a/docs/docs/sglang-diffusion/api/cli.mdx +++ b/docs/docs/sglang-diffusion/api/cli.mdx @@ -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 diff --git a/docs/docs/sglang-diffusion/ring_sp_performance.mdx b/docs/docs/sglang-diffusion/ring_sp_performance.mdx index 7844b86ec..0d337c0cb 100644 --- a/docs/docs/sglang-diffusion/ring_sp_performance.mdx +++ b/docs/docs/sglang-diffusion/ring_sp_performance.mdx @@ -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 + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
ModeCommunicationMemoryConstraints
ulyssesAll-to-all before and after attentionFull sequence with a shard of the attention heads during attentionAttention head divisibility must match the Ulysses degree
kv_gatherAll-gather K and V; Q and output remain sequence-shardedReplicates full K and V within the SP groupNon-causal attention and ring_degree=1; no legacy varlen or video sparse attention
+ +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`. diff --git a/python/sglang/multimodal_gen/runtime/distributed/sp_shard_utils.py b/python/sglang/multimodal_gen/runtime/distributed/sp_shard_utils.py index b589798ee..23b62e586 100644 --- a/python/sglang/multimodal_gen/runtime/distributed/sp_shard_utils.py +++ b/python/sglang/multimodal_gen/runtime/distributed/sp_shard_utils.py @@ -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: diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/vla/protocol.py b/python/sglang/multimodal_gen/runtime/entrypoints/vla/protocol.py index 16e8d2e8a..c8ef4ddba 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/vla/protocol.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/vla/protocol.py @@ -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, diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/layer.py b/python/sglang/multimodal_gen/runtime/layers/attention/layer.py index 63966182b..aa23827b1 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/layer.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/layer.py @@ -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) diff --git a/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py b/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py index b8d62b305..cbb2ea280 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py @@ -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, ) diff --git a/python/sglang/multimodal_gen/runtime/server_args/server_args.py b/python/sglang/multimodal_gen/runtime/server_args/server_args.py index 5ce970671..b8f3f4dd1 100644 --- a/python/sglang/multimodal_gen/runtime/server_args/server_args.py +++ b/python/sglang/multimodal_gen/runtime/server_args/server_args.py @@ -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})" diff --git a/python/sglang/multimodal_gen/test/unit/conftest.py b/python/sglang/multimodal_gen/test/unit/conftest.py index 5dd663c05..25bccd9fb 100644 --- a/python/sglang/multimodal_gen/test/unit/conftest.py +++ b/python/sglang/multimodal_gen/test/unit/conftest.py @@ -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, diff --git a/python/sglang/multimodal_gen/test/unit/test_ideogram4.py b/python/sglang/multimodal_gen/test/unit/test_ideogram4.py index ad569cd67..3abe8ea31 100644 --- a/python/sglang/multimodal_gen/test/unit/test_ideogram4.py +++ b/python/sglang/multimodal_gen/test/unit/test_ideogram4.py @@ -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( diff --git a/python/sglang/multimodal_gen/test/unit/test_pi05_action_api.py b/python/sglang/multimodal_gen/test/unit/test_pi05_action_api.py index c5482db05..720e12086 100644 --- a/python/sglang/multimodal_gen/test/unit/test_pi05_action_api.py +++ b/python/sglang/multimodal_gen/test/unit/test_pi05_action_api.py @@ -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 diff --git a/python/sglang/multimodal_gen/test/unit/test_server_args.py b/python/sglang/multimodal_gen/test/unit/test_server_args.py index 7b033bbfd..5d94d9b54 100644 --- a/python/sglang/multimodal_gen/test/unit/test_server_args.py +++ b/python/sglang/multimodal_gen/test/unit/test_server_args.py @@ -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( diff --git a/python/sglang/multimodal_gen/test/unit/test_sp_shard.py b/python/sglang/multimodal_gen/test/unit/test_sp_shard.py index 4f7e8261b..9d7b29fc4 100644 --- a/python/sglang/multimodal_gen/test/unit/test_sp_shard.py +++ b/python/sglang/multimodal_gen/test/unit/test_sp_shard.py @@ -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 diff --git a/python/sglang/multimodal_gen/test/unit/test_usp_attention_kv_gather.py b/python/sglang/multimodal_gen/test/unit/test_usp_attention_kv_gather.py new file mode 100644 index 000000000..a04bb5b8a --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/test_usp_attention_kv_gather.py @@ -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()