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

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

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

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
+
+<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

@@ -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

@@ -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

@@ -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
+    The default path implements USP, which combines Ulysses-style all-to-all
-    Ulysses-style all-to-all communication for sequence-head dimension sharding
+    communication for sequence-head dimension sharding with Ring Attention
-    and Ring Attention for fine-grained sequence parallelism within subgroups.
+    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 (
+        if replicated_mode_count > 1:
-            (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)
-        ):
             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

@@ -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

@@ -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
+            if self.sp_degree == 2:
-            logger.info(
+                # measured-win zone for the K/V-gather exchange; layers whose
-                f"Automatically set ulysses_degree=sp_degree={self.ulysses_degree} for best performance"
+                # 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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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()