[Feature] Add AllGather-KV sequence-parallel attention str… (#4968)

Signed-off-by: zhuhh97 <zhuhonghan@huawei.com>
Signed-off-by: dph97 <nenbaying@163.com>
Signed-off-by: zhuhh97 <zhh@unknown>
Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com>
Co-authored-by: dengyunyang <584797741@qq.com>

tests/config/test_omni_config.py

@@ -431,6 +431,7 @@ def test_sub_config_fields_match_rfc_scopes():
         "sequence_parallel_size",
         "ulysses_degree",
         "ring_degree",
+        "allgather_degree",
         "ulysses_mode",
         "cfg_parallel_size",
         "vae_patch_parallel_size",
@@ -547,11 +548,40 @@ def test_from_pipeline_config_derives_sequence_parallel_size_from_degrees(tmp_pa
     assert stage.parallel_config.world_size == 6
 
 
+def test_from_pipeline_config_derives_sequence_parallel_size_from_allgather_degree(tmp_path):
+    deploy_path = tmp_path / "dreamzero_allgather_parallel.yaml"
+    deploy_path.write_text(
+        "\n".join(
+            [
+                "pipeline: dreamzero",
+                "async_chunk: false",
+                "stages:",
+                "  - stage_id: 0",
+                "    parallel_config:",
+                "      sequence_parallel_size: 99",
+                "      allgather_degree: 2",
+            ]
+        )
+    )
+
+    stage = _from_pipeline_key("dreamzero", deploy_config_path=str(deploy_path)).stage_by_id(0)
+
+    assert isinstance(stage, VllmOmniDiffusionStageConfig)
+    assert stage.parallel_config.allgather_degree == 2
+    assert stage.parallel_config.sequence_parallel_size == 2
+    assert stage.parallel_config.world_size == 2
+
+
 def test_diffusion_parallel_config_rejects_cfg_parallel_size_outside_current_bound():
     with pytest.raises(ValidationError):
         OmniStageDiffusionParallelConfig(cfg_parallel_size=4)
 
 
+def test_diffusion_parallel_config_rejects_allgather_with_ulysses_or_ring():
+    with pytest.raises(ValidationError):
+        OmniStageDiffusionParallelConfig(allgather_degree=2, ulysses_degree=2)
+
+
 def test_stage_realizations_use_stage_specific_parallel_config_types():
     qwen_config = _from_pipeline_key("qwen3_tts")
     ar_stage = qwen_config.stage_by_id(0)

tests/diffusion/attention/test_attention_sp.py

@@ -31,7 +31,11 @@ import tempfile
 import pytest
 import torch
 
+from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata, QueryRange
 from vllm_omni.diffusion.attention.layer import Attention
+from vllm_omni.diffusion.attention.parallel.allgather_kv import (
+    AllGatherKVParallelAttention,
+)
 from vllm_omni.diffusion.config import set_current_diffusion_config
 from vllm_omni.diffusion.data import (
     DiffusionParallelConfig,
@@ -155,14 +159,267 @@ class TestMultiLayerAttentionModel(torch.nn.Module):
         return hidden_states
 
 
+class _MockAllGatherSPGroup:
+    def __init__(self, *, rank: int, gather_chunks: list[list[torch.Tensor]]) -> None:
+        self.allgather_group = object()
+        self.allgather_world_size = len(gather_chunks[0])
+        self.allgather_rank = rank
+        self._gather_chunks = list(gather_chunks)
+        self.gathered_input_shapes: list[tuple[int, ...]] = []
+
+    def all_gather(self, input_: torch.Tensor, dim: int = 0, separate_tensors: bool = False, group=None):
+        assert not separate_tensors
+        assert group is self.allgather_group
+        self.gathered_input_shapes.append(tuple(input_.shape))
+        chunks = self._gather_chunks.pop(0)
+        return torch.cat(chunks, dim=dim)
+
+
+def test_allgather_kv_slices_full_dense_mask_to_local_query_rows():
+    rank = 1
+    joint_len = 1
+    img_seq_local = 2
+    img_seq_full = 4
+    key_chunks = [
+        torch.full((1, img_seq_local, 1, 1), 10.0),
+        torch.full((1, img_seq_local, 1, 1), 20.0),
+    ]
+    value_chunks = [
+        torch.full((1, img_seq_local, 1, 1), 30.0),
+        torch.full((1, img_seq_local, 1, 1), 40.0),
+    ]
+    strategy = AllGatherKVParallelAttention(
+        _MockAllGatherSPGroup(rank=rank, gather_chunks=[key_chunks, value_chunks]),
+    )
+
+    query = torch.zeros((1, img_seq_local, 1, 1))
+    key = key_chunks[rank]
+    value = value_chunks[rank]
+    joint = torch.ones((1, joint_len, 1, 1))
+    mask = torch.arange((joint_len + img_seq_full) * (joint_len + img_seq_full))
+    mask = (mask % 2 == 0).view(1, 1, joint_len + img_seq_full, joint_len + img_seq_full)
+    metadata = AttentionMetadata(
+        attn_mask=mask,
+        joint_query=joint,
+        joint_key=joint,
+        joint_value=joint,
+        joint_strategy="front",
+    )
+
+    q_local, k_full, v_full, metadata_local, _ = strategy.pre_attention(query, key, value, metadata)
+
+    expected_rows = torch.cat(
+        [
+            mask[..., :joint_len, :],
+            mask[
+                ...,
+                joint_len + rank * img_seq_local : joint_len + (rank + 1) * img_seq_local,
+                :,
+            ],
+        ],
+        dim=-2,
+    )
+    assert q_local.shape[1] == joint_len + img_seq_local
+    assert k_full.shape[1] == joint_len + img_seq_full
+    assert v_full.shape[1] == joint_len + img_seq_full
+    torch.testing.assert_close(k_full[:, joint_len:], torch.cat(key_chunks, dim=1))
+    torch.testing.assert_close(v_full[:, joint_len:], torch.cat(value_chunks, dim=1))
+    assert metadata_local is not None
+    assert torch.equal(metadata_local.attn_mask, expected_rows)
+
+
+def test_allgather_kv_slices_rear_joint_dense_mask():
+    rank = 1
+    joint_len = 1
+    img_seq_local = 2
+    img_seq_full = 4
+    key_chunks = [torch.zeros((1, img_seq_local, 1, 1)) for _ in range(2)]
+    value_chunks = [torch.zeros_like(chunk) for chunk in key_chunks]
+    strategy = AllGatherKVParallelAttention(
+        _MockAllGatherSPGroup(rank=rank, gather_chunks=[key_chunks, value_chunks]),
+    )
+    query = torch.zeros((1, img_seq_local, 1, 1))
+    joint = torch.ones((1, joint_len, 1, 1))
+    mask = torch.arange((img_seq_full + joint_len) ** 2)
+    mask = (mask % 2 == 0).view(1, 1, img_seq_full + joint_len, img_seq_full + joint_len)
+    metadata = AttentionMetadata(
+        attn_mask=mask,
+        joint_query=joint,
+        joint_key=joint,
+        joint_value=joint,
+        joint_strategy="rear",
+    )
+
+    q_local, _, _, metadata_local, _ = strategy.pre_attention(query, key_chunks[rank], value_chunks[rank], metadata)
+
+    img_start = rank * img_seq_local
+    expected_rows = torch.cat(
+        [
+            mask[..., img_start : img_start + img_seq_local, :],
+            mask[..., img_seq_full : img_seq_full + joint_len, :],
+        ],
+        dim=-2,
+    )
+    assert q_local.shape[1] == img_seq_local + joint_len
+    assert metadata_local is not None
+    assert torch.equal(metadata_local.attn_mask, expected_rows)
+
+
+def test_allgather_kv_rejects_invalid_joint_strategy():
+    chunks = [[torch.zeros((1, 2, 1, 1)) for _ in range(2)] for _ in range(2)]
+    strategy = AllGatherKVParallelAttention(_MockAllGatherSPGroup(rank=0, gather_chunks=chunks))
+    joint = torch.ones((1, 1, 1, 1))
+    metadata = AttentionMetadata(
+        joint_query=joint,
+        joint_key=joint,
+        joint_value=joint,
+        joint_strategy="back",
+    )
+
+    with pytest.raises(ValueError, match="Unsupported joint_strategy"):
+        strategy.pre_attention(torch.zeros((1, 2, 1, 1)), chunks[0][0], chunks[1][0], metadata)
+
+
+def test_allgather_kv_preserves_global_spans_and_sets_query_ranges():
+    strategy = AllGatherKVParallelAttention(
+        _MockAllGatherSPGroup(
+            rank=1,
+            gather_chunks=[
+                [torch.zeros((1, 2, 1, 1)), torch.zeros((1, 2, 1, 1))],
+                [torch.zeros((1, 2, 1, 1)), torch.zeros((1, 2, 1, 1))],
+            ],
+        ),
+    )
+    query = torch.zeros((1, 2, 1, 1))
+    key = torch.zeros((1, 2, 1, 1))
+    value = torch.zeros((1, 2, 1, 1))
+    joint = torch.ones((1, 1, 1, 1))
+    metadata = AttentionMetadata(
+        joint_query=joint,
+        joint_key=joint,
+        joint_value=joint,
+        full_attn_spans=[
+            [
+                (0, 1),  # joint span: kept on every rank.
+                (2, 5),  # image span crossing rank 0 and rank 1 image shards.
+            ]
+        ],
+    )
+
+    _, _, _, metadata_out, _ = strategy.pre_attention(query, key, value, metadata)
+
+    assert metadata_out is not None
+    assert metadata_out.full_attn_spans == [[(0, 1), (2, 5)]]
+    assert metadata_out.query_ranges == (
+        QueryRange(0, 1, 0),
+        QueryRange(1, 3, 3),
+    )
+
+
+def test_allgather_kv_query_ranges_include_reused_prefix_offset():
+    strategy = AllGatherKVParallelAttention(
+        _MockAllGatherSPGroup(
+            rank=1,
+            gather_chunks=[
+                [torch.zeros((1, 2, 1, 1)), torch.zeros((1, 2, 1, 1))],
+                [torch.zeros((1, 2, 1, 1)), torch.zeros((1, 2, 1, 1))],
+            ],
+        ),
+    )
+    query = torch.zeros((1, 2, 1, 1))
+    key = torch.zeros((1, 2, 1, 1))
+    value = torch.zeros((1, 2, 1, 1))
+    joint_query = torch.ones((1, 1, 1, 1))
+    # K/V retain two reused AR-prefix tokens that have already been removed
+    # from Q and from the attention mask's query rows.
+    joint_key = torch.ones((1, 3, 1, 1))
+    joint_value = torch.ones((1, 3, 1, 1))
+    mask = torch.arange(5 * 7).view(1, 1, 5, 7)
+    metadata = AttentionMetadata(
+        attn_mask=mask,
+        joint_query=joint_query,
+        joint_key=joint_key,
+        joint_value=joint_value,
+        full_attn_spans=[[(0, 7)]],
+    )
+
+    _, _, _, metadata_out, _ = strategy.pre_attention(query, key, value, metadata)
+
+    assert metadata_out is not None
+    assert metadata_out.query_ranges == (
+        QueryRange(0, 1, 2),
+        QueryRange(1, 3, 5),
+    )
+    expected_mask = torch.cat([mask[..., :1, :], mask[..., 3:5, :]], dim=-2)
+    assert torch.equal(metadata_out.attn_mask, expected_mask)
+
+
+def test_allgather_kv_allows_empty_full_attn_spans():
+    strategy = AllGatherKVParallelAttention(
+        _MockAllGatherSPGroup(
+            rank=0,
+            gather_chunks=[
+                [torch.zeros((1, 2, 1, 1)), torch.zeros((1, 2, 1, 1))],
+                [torch.zeros((1, 2, 1, 1)), torch.zeros((1, 2, 1, 1))],
+            ],
+        ),
+    )
+    query = torch.zeros((1, 2, 1, 1))
+    key = torch.zeros((1, 2, 1, 1))
+    value = torch.zeros((1, 2, 1, 1))
+    metadata = AttentionMetadata(full_attn_spans=[[]])
+
+    _, _, _, metadata_out, _ = strategy.pre_attention(query, key, value, metadata)
+
+    assert metadata_out is not None
+    assert metadata_out.full_attn_spans == [[]]
+    assert metadata_out.query_ranges == (QueryRange(0, 2, 0),)
+
+
+def test_allgather_kv_keeps_gathered_kv_compressed_for_gqa():
+    rank = 0
+    img_seq_local = 2
+    kv_heads = 2
+    repeat_num = 2
+    q_heads = kv_heads * repeat_num
+    key_chunks = [
+        torch.arange(1 * img_seq_local * kv_heads * 1, dtype=torch.float32).view(1, img_seq_local, kv_heads, 1),
+        torch.arange(100, 100 + 1 * img_seq_local * kv_heads * 1, dtype=torch.float32).view(
+            1, img_seq_local, kv_heads, 1
+        ),
+    ]
+    value_chunks = [chunk + 1000 for chunk in key_chunks]
+    sp_group = _MockAllGatherSPGroup(rank=rank, gather_chunks=[key_chunks, value_chunks])
+    strategy = AllGatherKVParallelAttention(sp_group)
+
+    query = torch.zeros((1, img_seq_local, q_heads, 1))
+    _, k_full, v_full, _, _ = strategy.pre_attention(query, key_chunks[rank], value_chunks[rank], AttentionMetadata())
+
+    assert sp_group.gathered_input_shapes == [
+        (1, img_seq_local, kv_heads, 1),
+        (1, img_seq_local, kv_heads, 1),
+    ]
+    assert k_full.shape == (1, img_seq_local * 2, kv_heads, 1)
+    assert v_full.shape == (1, img_seq_local * 2, kv_heads, 1)
+    expected_key = torch.cat(key_chunks, dim=1)
+    expected_value = torch.cat(value_chunks, dim=1)
+    torch.testing.assert_close(k_full, expected_key)
+    torch.testing.assert_close(v_full, expected_value)
+
+
 @pytest.mark.parametrize(
     "test_model_cls",
     [
         TestMultiLayerAttentionModel,
     ],
 )
-@pytest.mark.parametrize("ulysses_degree", [2])
-@pytest.mark.parametrize("ring_degree", [2])
+@pytest.mark.parametrize(
+    ("ulysses_degree", "ring_degree", "allgather_degree"),
+    [
+        pytest.param(2, 2, 1, id="ulysses-ring"),
+        pytest.param(1, 1, 2, id="allgather-kv"),
+    ],
+)
 @pytest.mark.parametrize("batch_size", [2])
 @pytest.mark.parametrize("seq_len", [16])
 @pytest.mark.parametrize("num_heads", [8])
@@ -175,6 +432,7 @@ class TestMultiLayerAttentionModel(torch.nn.Module):
 def test_sequence_parallel(
     ulysses_degree: int,
     ring_degree: int,
+    allgather_degree: int,
     test_model_cls: type[torch.nn.Module],
     dtype: torch.dtype,
     causal: bool,
@@ -186,8 +444,8 @@ def test_sequence_parallel(
     num_heads: int,
     head_size: int,
 ):
-    """Test Ulysses attention by comparing with and without SP enabled."""
-    sequence_parallel_size = ulysses_degree * ring_degree
+    """Compare Ulysses/Ring and AllGather-KV SP against a single-rank run."""
+    sequence_parallel_size = allgather_degree if allgather_degree > 1 else ulysses_degree * ring_degree
 
     # Skip if not enough GPUs available
     available_gpus = current_omni_platform.get_device_count()
@@ -223,6 +481,7 @@ def test_sequence_parallel(
                 use_compile,
                 1,  # ulysses_degree = 1
                 1,  # ring_degree = 1
+                1,  # allgather_degree = 1 (baseline: no AllGather-KV SP)
                 1,  # sequence_parallel_size = 1
                 baseline_output_file,
                 model_state_file,
@@ -233,7 +492,10 @@ def test_sequence_parallel(
         )
 
         # Step 2: Run with SP enabled
-        print(f"\n[SP Test] Running with SP (ulysses_degree={ulysses_degree}, ring_degree={ring_degree})...")
+        print(
+            "\n[SP Test] Running with SP "
+            f"(ulysses={ulysses_degree}, ring={ring_degree}, allgather={allgather_degree})..."
+        )
         torch.multiprocessing.spawn(
             ulysses_attention_on_test_model,
             args=(
@@ -250,6 +512,7 @@ def test_sequence_parallel(
                 use_compile,
                 ulysses_degree,
                 ring_degree,
+                allgather_degree,
                 sequence_parallel_size,
                 sp_output_file,
                 model_state_file,
@@ -345,6 +608,7 @@ def ulysses_attention_on_test_model(
     use_compile: bool,
     ulysses_degree: int,
     ring_degree: int,
+    allgather_degree: int,
     sequence_parallel_size: int,
     output_file: str,
     model_state_file: str,
@@ -356,7 +620,11 @@ def ulysses_attention_on_test_model(
     RANDOM_SEED = 42
     seed_everything(RANDOM_SEED)
 
-    mode_str = "Baseline (no SP)" if is_baseline else f"SP (ulysses={ulysses_degree}, ring={ring_degree})"
+    if allgather_degree > 1:
+        sp_kind = f"allgather={allgather_degree}"
+    else:
+        sp_kind = f"ulysses={ulysses_degree}, ring={ring_degree}"
+    mode_str = "Baseline (no SP)" if is_baseline else f"SP ({sp_kind})"
     print(f"\n[{mode_str}] Rank {local_rank}/{world_size} - Random seed set to {RANDOM_SEED}")
 
     device = torch.device(f"{current_omni_platform.device_type}:{local_rank}")
@@ -384,6 +652,7 @@ def ulysses_attention_on_test_model(
         sequence_parallel_size=sequence_parallel_size,
         ulysses_degree=ulysses_degree,
         ring_degree=ring_degree,
+        allgather_degree=allgather_degree,
         cfg_parallel_size=1,
     )
 
@@ -400,6 +669,7 @@ def ulysses_attention_on_test_model(
         sequence_parallel_size=sequence_parallel_size,
         ulysses_degree=ulysses_degree,
         ring_degree=ring_degree,
+        allgather_degree=allgather_degree,
         tensor_parallel_size=1,
         pipeline_parallel_size=1,
     )
@@ -519,7 +789,11 @@ def ulysses_attention_on_test_model(
             with open(output_file, "wb") as f:
                 pickle.dump(output_np, f)
 
-            mode_str = "baseline (no SP)" if is_baseline else f"SP (ulysses={ulysses_degree}, ring={ring_degree})"
+            if allgather_degree > 1:
+                sp_kind = f"allgather={allgather_degree}"
+            else:
+                sp_kind = f"ulysses={ulysses_degree}, ring={ring_degree}"
+            mode_str = "baseline (no SP)" if is_baseline else f"SP ({sp_kind})"
             print(
                 f"\n[{mode_str}] ✓ Saved output with shape {full_output.shape}:\n"
                 f"  - batch_size={batch_size}, seq_len={seq_len}\n"

tests/diffusion/attention/test_flash_attn.py

@@ -297,6 +297,36 @@ def test_flash_attn_func_preferred_over_varlen():
     print("✓ flash_attn_func forward works correctly!")
 
 
+def test_piecewise_flash_attn_uses_varlen_fallback(monkeypatch):
+    calls = []
+
+    def fake_varlen_func(q, k, v, **kwargs):
+        assert q.shape[1] == 4
+        assert k.shape[1] == 2
+        assert v.shape[1] == 2
+        calls.append(kwargs)
+        return q
+
+    monkeypatch.setattr(fa, "flash_attn_func", None)
+    monkeypatch.setattr(fa, "flash_attn_varlen_func", fake_varlen_func)
+    monkeypatch.setattr(fa, "HAS_FLASH_ATTN", True)
+
+    impl = FlashAttentionImpl(num_heads=4, num_kv_heads=2, head_size=4, softmax_scale=0.5, causal=False)
+    query = torch.randn(1, 3, 4, 4)
+    key = torch.randn(1, 3, 2, 4)
+    value = torch.randn(1, 3, 2, 4)
+    metadata = AttentionMetadata(full_attn_spans=[[(0, 3)]])
+
+    output = impl.forward_cuda(query, key, value, metadata)
+
+    assert output.shape == query.shape
+    assert len(calls) == 1
+    assert calls[0]["max_seqlen_q"] == 3
+    assert calls[0]["max_seqlen_k"] == 3
+    assert calls[0]["causal"] is False
+    assert calls[0]["softmax_scale"] == 0.5
+
+
 if __name__ == "__main__":
     print("Running FlashAttention Padding Tests...")
     print("=" * 60)

tests/diffusion/attention/test_piecewise_attn.py

@@ -19,7 +19,11 @@ import pytest
 import torch
 import torch.nn.functional as F
 
+import vllm_omni.diffusion.attention.backends.utils.piecewise_attn as piecewise_module
+from vllm_omni.diffusion.attention.backends.abstract import QueryRange
 from vllm_omni.diffusion.attention.backends.utils.piecewise_attn import (
+    Segment,
+    build_segments,
     piecewise_attn,
 )
 
@@ -71,6 +75,12 @@ BATCH_CASES = [
 ]
 
 
+def test_build_segments_full_span_retains_global_kv_end():
+    segments = build_segments([(2, 7)], query_offset=4, query_len=2)
+
+    assert segments == [Segment(q_start=4, q_end=6, kv_end=7, mode="full")]
+
+
 @pytest.mark.parametrize("global_spans", SPAN_CASES)
 @pytest.mark.parametrize("q_range", Q_RANGE_CASES)
 @pytest.mark.parametrize("batch_size", BATCH_CASES)
@@ -124,3 +134,80 @@ def test_piecewise_span_fully_before_qstart():
     )
     expected = _full_reference(query, key, value, global_spans, q_start, q_end, softmax_scale)
     torch.testing.assert_close(got, expected, atol=1e-5, rtol=1e-5)
+
+
+def test_piecewise_noncontiguous_query_ranges_match_full_mask():
+    torch.manual_seed(0)
+    B, H, D, Sk = 1, 2, 8, 8
+    key = torch.randn(B, Sk, H, D, device=DEVICE)
+    value = torch.randn(B, Sk, H, D, device=DEVICE)
+    query = torch.randn(B, 4, H, D, device=DEVICE)
+    spans = [(2, 7)]
+    query_ranges = (
+        QueryRange(0, 1, 0),
+        QueryRange(1, 4, 4),
+    )
+    softmax_scale = D**-0.5
+
+    got = piecewise_attn(
+        query,
+        key,
+        value,
+        full_attn_spans=[spans],
+        softmax_scale=softmax_scale,
+        attn_func=_sdpa_attn_func,
+        query_ranges=query_ranges,
+    )
+    expected = torch.cat(
+        [
+            _full_reference(query[:, :1], key, value, spans, 0, 1, softmax_scale),
+            _full_reference(query[:, 1:], key, value, spans, 4, 7, softmax_scale),
+        ],
+        dim=1,
+    )
+    torch.testing.assert_close(got, expected, atol=1e-5, rtol=1e-5)
+
+
+def test_piecewise_rejects_incomplete_query_ranges():
+    query = torch.zeros(1, 3, 1, 2)
+    key = torch.zeros(1, 3, 1, 2)
+    value = torch.zeros_like(key)
+
+    with pytest.raises(ValueError, match="cover local query contiguously"):
+        piecewise_attn(
+            query,
+            key,
+            value,
+            full_attn_spans=[[]],
+            softmax_scale=1.0,
+            attn_func=_sdpa_attn_func,
+            query_ranges=(QueryRange(1, 3, 1),),
+        )
+
+
+def test_piecewise_combines_noncontiguous_query_ranges_with_one_cat(monkeypatch):
+    query = torch.randn(1, 4, 2, 8)
+    key = torch.randn(1, 8, 2, 8)
+    value = torch.randn_like(key)
+    cat_calls = 0
+    original_cat = torch.cat
+
+    def counted_cat(tensors, dim=0):
+        nonlocal cat_calls
+        cat_calls += 1
+        return original_cat(tensors, dim=dim)
+
+    monkeypatch.setattr(piecewise_module.torch, "cat", counted_cat)
+
+    output = piecewise_attn(
+        query,
+        key,
+        value,
+        full_attn_spans=[[(0, 8)]],
+        softmax_scale=1.0,
+        attn_func=lambda q, k, v, **kwargs: q,
+        query_ranges=(QueryRange(0, 1, 0), QueryRange(1, 4, 4)),
+    )
+
+    assert cat_calls == 1
+    torch.testing.assert_close(output, query)

tests/diffusion/attention/test_sdpa.py

@@ -0,0 +1,57 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+
+import vllm_omni.diffusion.attention.backends.sdpa as sdpa_backend
+from vllm_omni.diffusion.attention.backends.sdpa import SDPAImpl
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_sdpa_expands_kv_when_native_gqa_kernel_is_unavailable(monkeypatch):
+    calls = []
+
+    def fake_sdpa(query, key, value, **kwargs):
+        calls.append((query.shape, key.shape, value.shape, kwargs))
+        return query
+
+    monkeypatch.setattr(sdpa_backend, "can_sdpa_use_fused_gqa", lambda *args: False)
+    monkeypatch.setattr(torch.nn.functional, "scaled_dot_product_attention", fake_sdpa)
+
+    impl = SDPAImpl(num_heads=4, num_kv_heads=2, head_size=8, softmax_scale=0.5)
+    output = impl.forward_cuda(
+        torch.randn(1, 3, 4, 8),
+        torch.randn(1, 3, 2, 8),
+        torch.randn(1, 3, 2, 8),
+    )
+
+    query_shape, key_shape, value_shape, kwargs = calls[0]
+    assert query_shape == key_shape == value_shape == (1, 4, 3, 8)
+    assert kwargs["enable_gqa"] is False
+    assert output.shape == (1, 3, 4, 8)
+
+
+def test_sdpa_keeps_compressed_kv_when_native_gqa_kernel_is_available(monkeypatch):
+    calls = []
+
+    def fake_sdpa(query, key, value, **kwargs):
+        calls.append((query.shape, key.shape, value.shape, kwargs))
+        return query
+
+    monkeypatch.setattr(sdpa_backend, "can_sdpa_use_fused_gqa", lambda *args: True)
+    monkeypatch.setattr(torch.nn.functional, "scaled_dot_product_attention", fake_sdpa)
+
+    impl = SDPAImpl(num_heads=4, num_kv_heads=2, head_size=8, softmax_scale=0.5)
+    output = impl.forward_cuda(
+        torch.randn(1, 3, 4, 8),
+        torch.randn(1, 3, 2, 8),
+        torch.randn(1, 3, 2, 8),
+    )
+
+    query_shape, key_shape, value_shape, kwargs = calls[0]
+    assert query_shape == (1, 4, 3, 8)
+    assert key_shape == value_shape == (1, 2, 3, 8)
+    assert kwargs["enable_gqa"] is True
+    assert output.shape == (1, 3, 4, 8)

tests/diffusion/distributed/test_parallel_state_sp_groups.py

@@ -40,7 +40,7 @@ def test_set_seq_parallel_pg_uses_explicit_sp_groups(rank, expected_ulysses, exp
     # tp=2, sp=2 -> SP groups are non-contiguous: [0,2] and [1,3]
     sp_group_ranks = RankGenerator(2, 2, 1, 1, 1, "tp-sp-pp-cfg-dp").get_ranks("sp")
 
-    ulysses_pg, ring_pg = set_seq_parallel_pg(
+    ulysses_pg, ring_pg, allgather_pg = set_seq_parallel_pg(
         sp_ulysses_degree=2,
         sp_ring_degree=1,
         rank=rank,
@@ -50,6 +50,43 @@ def test_set_seq_parallel_pg_uses_explicit_sp_groups(rank, expected_ulysses, exp
 
     assert ulysses_pg.ranks == expected_ulysses
     assert ring_pg.ranks == expected_ring
+    assert allgather_pg.ranks == [rank]
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize("rank", [0, 1, 2, 3])
+def test_set_seq_parallel_pg_builds_independent_allgather_group(rank, monkeypatch):
+    created_groups: list[SimpleNamespace] = []
+    monkeypatch.setattr(torch.distributed, "new_group", _fake_new_group_factory(created_groups))
+    sp_group_ranks = [[0, 2], [1, 3]]
+
+    ulysses_pg, ring_pg, allgather_pg = set_seq_parallel_pg(
+        sp_ulysses_degree=1,
+        sp_ring_degree=1,
+        sp_allgather_degree=2,
+        rank=rank,
+        world_size=4,
+        sp_group_ranks=sp_group_ranks,
+    )
+
+    assert ulysses_pg.ranks == [rank]
+    assert ring_pg.ranks == [rank]
+    assert allgather_pg.ranks == ([0, 2] if rank % 2 == 0 else [1, 3])
+
+
+@pytest.mark.cpu
+def test_set_seq_parallel_pg_rejects_allgather_with_ulysses_or_ring(monkeypatch):
+    created_groups: list[SimpleNamespace] = []
+    monkeypatch.setattr(torch.distributed, "new_group", _fake_new_group_factory(created_groups))
+
+    with pytest.raises(ValueError, match="mutually exclusive"):
+        set_seq_parallel_pg(
+            sp_ulysses_degree=2,
+            sp_ring_degree=1,
+            sp_allgather_degree=2,
+            rank=0,
+            world_size=4,
+        )
 
 
 @pytest.mark.cpu

tests/diffusion/models/hunyuan_image3/test_image_kv_cache_manager.py

@@ -45,6 +45,8 @@ def patched_mgr_env(sp_size=1):
     target = _TRANSFORMER_MODULE
     patches = [
         patch(f"{target}.get_sequence_parallel_world_size", return_value=sp_size),
+        patch(f"{target}.get_allgather_parallel_world_size", return_value=sp_size),
+        patch(f"{target}.get_ulysses_parallel_world_size", return_value=1, create=True),
         patch(f"{target}.get_sequence_parallel_rank", return_value=0),
         patch(f"{target}.Attention", MockAttention),
     ]

tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 import base64
 import copy
 import io
+import logging
 import os
 import tempfile
 from pathlib import Path
@@ -18,6 +19,8 @@ from tests.helpers.runtime import OmniServer
 
 pytestmark = [pytest.mark.full_model, pytest.mark.diffusion]
 
+logger = logging.getLogger(__name__)
+
 MODEL_NAME = "tencent/HunyuanImage-3.0-Instruct"
 SEED = 42
 NUM_INFERENCE_STEPS = 50
@@ -53,11 +56,12 @@ _DEPLOY_CONFIG = {
             "parallel_config": {
                 "pipeline_parallel_size": 1,
                 "data_parallel_size": 1,
-                "tensor_parallel_size": 4,
+                "tensor_parallel_size": 1,
                 "enable_expert_parallel": True,
                 "sequence_parallel_size": 1,
                 "ulysses_degree": 1,
                 "ring_degree": 1,
+                "allgather_degree": 4,
                 "cfg_parallel_size": 1,
                 "vae_patch_parallel_size": 1,
                 "use_hsdp": False,
@@ -83,9 +87,47 @@ def _devices() -> str:
 def _write_deploy_config(path: Path) -> None:
     config = copy.deepcopy(_DEPLOY_CONFIG)
     devices = _devices()
+    num_devices = len(devices.split(","))
     config["stages"][0]["devices"] = devices
-    config["stages"][0]["parallel_config"]["tensor_parallel_size"] = len(devices.split(","))
-    path.write_text(yaml.dump(config, default_flow_style=False, sort_keys=False))
+    parallel_config = config["stages"][0]["parallel_config"]
+    if parallel_config.get("allgather_degree", 1) > 1:
+        parallel_config["allgather_degree"] = num_devices
+    elif parallel_config.get("tensor_parallel_size", 1) > 1:
+        parallel_config["tensor_parallel_size"] = num_devices
+
+    sequence_parallel_size = (
+        parallel_config.get("allgather_degree", 1)
+        if parallel_config.get("allgather_degree", 1) > 1
+        else parallel_config.get("ulysses_degree", 1) * parallel_config.get("ring_degree", 1)
+    )
+    configured_world_size = (
+        parallel_config.get("pipeline_parallel_size", 1)
+        * parallel_config.get("data_parallel_size", 1)
+        * parallel_config.get("tensor_parallel_size", 1)
+        * sequence_parallel_size
+        * parallel_config.get("cfg_parallel_size", 1)
+    )
+    assert configured_world_size == num_devices, (
+        f"Configured parallel world size ({configured_world_size}) does not match "
+        f"the number of devices ({num_devices}): {parallel_config}"
+    )
+
+    config_text = yaml.dump(config, default_flow_style=False, sort_keys=False)
+    path.write_text(config_text)
+    logger.info(
+        "HunyuanImage3 pixel accuracy deploy config:\n"
+        "  path=%s\n"
+        "  CUDA_VISIBLE_DEVICES=%s\n"
+        "  devices=%s\n"
+        "  tensor_parallel_size=%s\n"
+        "  allgather_degree=%s\n%s",
+        path,
+        os.environ.get("CUDA_VISIBLE_DEVICES", "<unset>"),
+        devices,
+        parallel_config.get("tensor_parallel_size", 1),
+        parallel_config.get("allgather_degree", 1),
+        config_text,
+    )
 
 
 def _run_vllm_omni_hunyuan_image3_online(*, model: str, deploy_config: str, output_path: Path) -> Image.Image:

tests/test_config_factory.py

@@ -336,6 +336,31 @@ class TestStageConfig:
         assert "ring_degree" not in omega_config.engine_args
         assert "sequence_parallel_size" not in omega_config.engine_args
 
+    def test_to_omegaconf_diffusion_parallel_allgather_override_recomputes_sequence_parallel_size(self):
+        config = StageConfig(
+            stage_id=1,
+            model_stage="diffusion",
+            stage_type=StageType.DIFFUSION,
+            yaml_engine_args={
+                "parallel_config": {
+                    "sequence_parallel_size": 1,
+                    "ulysses_degree": 1,
+                    "ring_degree": 1,
+                    "allgather_degree": 1,
+                }
+            },
+            runtime_overrides={
+                "allgather_degree": 2,
+            },
+        )
+
+        omega_config = config.to_omegaconf()
+
+        assert omega_config.engine_args.parallel_config.allgather_degree == 2
+        assert omega_config.engine_args.parallel_config.sequence_parallel_size == 2
+        assert "allgather_degree" not in omega_config.engine_args
+        assert "sequence_parallel_size" not in omega_config.engine_args
+
     def test_to_omegaconf_diffusion_parallel_explicit_sequence_parallel_size_is_preserved(self):
         config = StageConfig(
             stage_id=1,

vllm_omni/config/omni_config.py

@@ -109,6 +109,7 @@ class _ParallelConfigEngineOverrides(TypedDict, total=False):
     sequence_parallel_size: int
     ulysses_degree: int
     ring_degree: int
+    allgather_degree: int
     ulysses_mode: str
     cfg_parallel_size: int
     vae_patch_parallel_size: int
@@ -360,6 +361,7 @@ class OmniStageDiffusionParallelConfig(OmniStageParallelConfig):
     sequence_parallel_size: int = Field(default=1, ge=1, init=False)
     ulysses_degree: int = Field(default=1, ge=1)
     ring_degree: int = Field(default=1, ge=1)
+    allgather_degree: int = Field(default=1, ge=1)
     ulysses_mode: str = "strict"
     cfg_parallel_size: int = Field(default=1, ge=1, le=3)
     vae_patch_parallel_size: int = Field(default=1, ge=1)
@@ -370,7 +372,11 @@ class OmniStageDiffusionParallelConfig(OmniStageParallelConfig):
     hsdp_replicate_size: int = Field(default=1, ge=1)
 
     def __post_init__(self) -> None:
-        self.sequence_parallel_size = self.ulysses_degree * self.ring_degree
+        self.sequence_parallel_size = (
+            self.allgather_degree if self.allgather_degree > 1 else self.ulysses_degree * self.ring_degree
+        )
+        if self.allgather_degree > 1 and (self.ulysses_degree > 1 or self.ring_degree > 1):
+            raise ValueError("allgather_degree > 1 is mutually exclusive with ulysses_degree/ring_degree > 1")
         if self.ulysses_mode not in {"strict", "advanced_uaa"}:
             raise ValueError("ulysses_mode must be 'strict' or 'advanced_uaa'")
         if self.vae_parallel_mode not in {"tile", "spatial_shard_height", "spatial_shard_width"}:

vllm_omni/config/stage_config.py

@@ -149,14 +149,17 @@ def _apply_diffusion_parallel_runtime_overrides(
             continue
         if parallel_config_dict is None:
             parallel_config_dict = {}
-        if key in ("ulysses_degree", "ring_degree"):
+        if key in ("ulysses_degree", "ring_degree", "allgather_degree"):
             degree_overridden = True
         parallel_config_dict[key] = runtime_overrides.pop(key)
 
     if parallel_config_dict is not None and degree_overridden and not sequence_parallel_explicit:
         ulysses_degree = parallel_config_dict.get("ulysses_degree") or 1
         ring_degree = parallel_config_dict.get("ring_degree") or 1
-        parallel_config_dict["sequence_parallel_size"] = ulysses_degree * ring_degree
+        allgather_degree = parallel_config_dict.get("allgather_degree") or 1
+        parallel_config_dict["sequence_parallel_size"] = (
+            allgather_degree if allgather_degree > 1 else ulysses_degree * ring_degree
+        )
 
     if parallel_config_dict is not None:
         engine_args["parallel_config"] = parallel_config_dict
@@ -366,6 +369,7 @@ class StageDeployConfig:
     ulysses_degree: int | None = None
     ulysses_mode: str | None = None
     ring_degree: int | None = None
+    allgather_degree: int | None = None
     sequence_parallel_size: int | None = None
     cfg_parallel_size: int | None = None
     vae_patch_parallel_size: int | None = None

vllm_omni/deploy/hunyuan_image3_dit.yaml

@@ -24,6 +24,7 @@ stages:
       sequence_parallel_size: 1
       ulysses_degree: 1
       ring_degree: 1
+      allgather_degree: 1
       cfg_parallel_size: 1
       vae_patch_parallel_size: 1
       use_hsdp: false

vllm_omni/diffusion/attention/backends/abstract.py

@@ -14,6 +14,7 @@ class AttentionBackend(ABC):
     """Abstract class for diffusion attention backends."""
 
     accept_output_buffer: bool = False
+    supports_piecewise_spans: bool = False
 
     @classmethod
     def supports_attention_mask(cls) -> bool:
@@ -51,6 +52,13 @@ class AttentionBackend(ABC):
         return (not supported_head_sizes) or head_size in supported_head_sizes
 
 
+@dataclass(frozen=True, slots=True)
+class QueryRange:
+    local_start: int
+    local_end: int
+    global_start: int
+
+
 @dataclass
 class AttentionMetadata:
     attn_mask: torch.Tensor | None = None
@@ -75,6 +83,7 @@ class AttentionMetadata:
     # Piecewise attention metadata (mixed causal/full masks).
     # full_attn_spans: per-sample [start, end) spans in global coordinates using full attention.
     full_attn_spans: list[list[tuple[int, int]]] | None = None
+    query_ranges: tuple[QueryRange, ...] | None = None
 
 
 T = TypeVar("T", bound=AttentionMetadata)

vllm_omni/diffusion/attention/backends/cudnn_attn.py

@@ -51,7 +51,6 @@ class CuDNNAttentionImpl(AttentionImpl):
     ) -> None:
         self.causal = causal
         self.softmax_scale = softmax_scale
-        self.requires_gqa = num_heads != num_kv_heads
 
     def forward_cuda(
         self,
@@ -64,6 +63,7 @@ class CuDNNAttentionImpl(AttentionImpl):
         if attn_metadata:
             attention_mask = _maybe_reshape_attn_mask(query, key, attn_metadata.attn_mask, mask_mode="broadcast_k")
 
+        enable_gqa = query.shape[2] != key.shape[2]
         query, key, value = (x.permute(0, 2, 1, 3) for x in (query, key, value))
         # Pin cuDNN exclusively. A priority list like [CUDNN, FLASH, MATH] hits a
         # PyTorch SDPA dispatch quirk: when FLASH rejects a non-None attn_mask,
@@ -85,7 +85,7 @@ class CuDNNAttentionImpl(AttentionImpl):
                     dropout_p=0.0,
                     is_causal=self.causal,
                     scale=self.softmax_scale,
-                    enable_gqa=self.requires_gqa,
+                    enable_gqa=enable_gqa,
                 )
         except RuntimeError as e:
             if "No available kernel" not in str(e):
@@ -102,6 +102,6 @@ class CuDNNAttentionImpl(AttentionImpl):
                 dropout_p=0.0,
                 is_causal=self.causal,
                 scale=self.softmax_scale,
-                enable_gqa=self.requires_gqa,
+                enable_gqa=enable_gqa,
             )
         return output.permute(0, 2, 1, 3)

vllm_omni/diffusion/attention/backends/flash_attn.py

@@ -21,6 +21,7 @@ logger = init_logger(__name__)
 
 class FlashAttentionBackend(AttentionBackend):
     accept_output_buffer: bool = True
+    supports_piecewise_spans: bool = True
 
     @classmethod
     def supports_attention_mask(cls) -> bool:
@@ -82,6 +83,28 @@ class FlashAttentionImpl(AttentionImpl):
     def _flash_wrapper(q, k, v, *, attn_func, **kwargs):
         return FlashAttentionImpl._unwrap_flash_output(attn_func(q, k, v, **kwargs))
 
+    @staticmethod
+    def _flash_varlen_wrapper(q, k, v, *, attn_func, causal, softmax_scale, **kwargs):
+        """Call a varlen-only FlashAttention backend for a dense segment."""
+        del kwargs
+        batch_size, q_len = q.shape[:2]
+        k_len = k.shape[1]
+        cu_seqlens_q = torch.arange(0, (batch_size + 1) * q_len, q_len, dtype=torch.int32, device=q.device)
+        cu_seqlens_k = torch.arange(0, (batch_size + 1) * k_len, k_len, dtype=torch.int32, device=q.device)
+        out = attn_func(
+            q=q.flatten(0, 1),
+            k=k.flatten(0, 1),
+            v=v.flatten(0, 1),
+            cu_seqlens_q=cu_seqlens_q,
+            cu_seqlens_k=cu_seqlens_k,
+            max_seqlen_q=q_len,
+            max_seqlen_k=k_len,
+            causal=causal,
+            softmax_scale=softmax_scale,
+        )
+        out = FlashAttentionImpl._unwrap_flash_output(out)
+        return out.reshape(batch_size, q_len, *out.shape[1:])
+
     def _forward_varlen_masked(
         self,
         query: torch.Tensor,
@@ -171,6 +194,7 @@ class FlashAttentionImpl(AttentionImpl):
         from vllm_omni.diffusion.attention.backends.utils.fa import (
             HAS_FLASH_ATTN,
             flash_attn_func,
+            flash_attn_varlen_func,
         )
 
         if not HAS_FLASH_ATTN:
@@ -186,10 +210,18 @@ class FlashAttentionImpl(AttentionImpl):
         # Try piecewise attention
         if full_attn_spans is not None:
             logger.debug("Using piecewise Flash Attention for mixed causal/full mask")
-            attn_func = partial(
-                FlashAttentionImpl._flash_wrapper,
-                attn_func=flash_attn_func,
-            )
+            if flash_attn_func is not None:
+                attn_func = partial(
+                    FlashAttentionImpl._flash_wrapper,
+                    attn_func=flash_attn_func,
+                )
+            elif flash_attn_varlen_func is not None:
+                attn_func = partial(
+                    FlashAttentionImpl._flash_varlen_wrapper,
+                    attn_func=flash_attn_varlen_func,
+                )
+            else:
+                raise ImportError("Piecewise FlashAttention requires a dense or varlen FlashAttention function")
 
             return piecewise_attn(
                 query,
@@ -198,6 +230,7 @@ class FlashAttentionImpl(AttentionImpl):
                 full_attn_spans,
                 self.softmax_scale,
                 attn_func,
+                query_ranges=attn_metadata.query_ranges,
             )
 
         if attention_mask is not None and torch.any(~attention_mask):

vllm_omni/diffusion/attention/backends/flash_attn_hub.py

@@ -12,9 +12,7 @@ from vllm_omni.diffusion.attention.backends.abstract import (
     AttentionImpl,
     AttentionMetadata,
 )
-from vllm_omni.diffusion.attention.backends.utils.piecewise_attn import (
-    piecewise_attn,
-)
+from vllm_omni.diffusion.attention.backends.utils.piecewise_attn import piecewise_attn
 
 logger = init_logger(__name__)
 
@@ -78,6 +76,7 @@ def _run_varlen_dense(
 
 class FlashAttentionHubBackend(AttentionBackend):
     accept_output_buffer: bool = True
+    supports_piecewise_spans: bool = True
 
     @classmethod
     def supports_attention_mask(cls) -> bool:
@@ -214,6 +213,7 @@ class FlashAttentionHubImpl(AttentionImpl):
                 full_attn_spans,
                 self.softmax_scale,
                 attn_func,
+                query_ranges=attn_metadata.query_ranges,
             )
 
         if attention_mask is not None and torch.any(~attention_mask):
@@ -243,6 +243,7 @@ class FlashAttentionHubImpl(AttentionImpl):
 
 class FlashAttention3HubBackend(AttentionBackend):
     accept_output_buffer: bool = True
+    supports_piecewise_spans: bool = True
 
     @classmethod
     def supports_attention_mask(cls) -> bool:
@@ -379,6 +380,7 @@ class FlashAttention3HubImpl(AttentionImpl):
                 full_attn_spans,
                 self.softmax_scale,
                 attn_func,
+                query_ranges=attn_metadata.query_ranges,
             )
 
         if attention_mask is not None and torch.any(~attention_mask):

vllm_omni/diffusion/attention/backends/sdpa.py

@@ -11,6 +11,7 @@ from vllm_omni.diffusion.attention.backends.abstract import (
     AttentionImpl,
     AttentionMetadata,
 )
+from vllm_omni.diffusion.attention.backends.utils.attn_runtime_selector import can_sdpa_use_fused_gqa
 
 logger = init_logger(__name__)
 
@@ -85,7 +86,6 @@ class SDPAImpl(AttentionImpl):
     ) -> None:
         self.causal = causal
         self.softmax_scale = softmax_scale
-        self.requires_gqa = num_heads != num_kv_heads
         if backend_kwargs:
             logger.warning("SDPAImpl ignoring backend_kwargs: %s", list(backend_kwargs.keys()))
 
@@ -103,7 +103,24 @@ class SDPAImpl(AttentionImpl):
         if attn_metadata:
             attention_mask = _maybe_reshape_attn_mask(query, key, attn_metadata.attn_mask, mask_mode=mask_mode)
 
+        enable_gqa = query.shape[2] != key.shape[2]
         query, key, value = (x.permute(0, 2, 1, 3) for x in (query, key, value))
+        # Only the PyTorch SDPA backend needs this dispatch check. If SDPA
+        # cannot select a fused GQA kernel for the runtime shape/mask, expand
+        # K/V locally so it can use the better-supported equal-head path.
+        if enable_gqa and not can_sdpa_use_fused_gqa(query, key, value, attention_mask, self.causal):
+            if query.shape[1] % key.shape[1] != 0:
+                raise ValueError(
+                    "GQA requires query heads to be a multiple of KV heads, "
+                    f"got q_heads={query.shape[1]} and kv_heads={key.shape[1]}."
+                )
+            repeat_num = query.shape[1] // key.shape[1]
+            key = key.repeat_interleave(repeat_num, dim=1)
+            value = value.repeat_interleave(repeat_num, dim=1)
+            enable_gqa = False
+            logger.debug(
+                "CUDA SDPA cannot use a fused native-GQA kernel for this shape; expanding K/V heads before SDPA."
+            )
         output = torch.nn.functional.scaled_dot_product_attention(
             query,
             key,
@@ -112,7 +129,7 @@ class SDPAImpl(AttentionImpl):
             dropout_p=0.0,
             is_causal=self.causal,
             scale=self.softmax_scale,
-            enable_gqa=self.requires_gqa,
+            enable_gqa=enable_gqa,
         )
         out = output.permute(0, 2, 1, 3)
         return out

vllm_omni/diffusion/attention/backends/utils/attn_runtime_selector.py

@@ -0,0 +1,40 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Runtime kernel-selection helpers for attention backends."""
+
+import torch
+
+
+def can_sdpa_use_fused_gqa(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attention_mask: torch.Tensor | None,
+    causal: bool,
+) -> bool:
+    """Return whether PyTorch SDPA can use a fused kernel for this GQA call.
+
+    This check is specific to the SDPA backend. Native GQA backends such as
+    FlashAttention should keep compressed K/V heads and do not use this helper.
+    """
+    if not query.is_cuda:
+        return False
+
+    try:
+        params = torch.backends.cuda.SDPAParams(
+            query,
+            key,
+            value,
+            attention_mask,
+            0.0,
+            causal,
+            True,
+        )
+        can_use_flash = torch.backends.cuda.can_use_flash_attention(params)
+        can_use_cudnn = getattr(torch.backends.cuda, "can_use_cudnn_attention", None)
+        return can_use_flash or (can_use_cudnn is not None and can_use_cudnn(params))
+    except (AttributeError, RuntimeError, TypeError):
+        # Older PyTorch versions may not expose GQA-aware SDPA capability checks.
+        # Returning False lets the SDPA backend expand K/V and use the mature
+        # equal-head fused-kernel path instead.
+        return False

vllm_omni/diffusion/attention/backends/utils/piecewise_attn.py

@@ -9,17 +9,24 @@ follows FlashAttention's bottom-right convention (``K[:e]`` is attended by
 
 Per segment:
   - causal segment ``[s, e)``: ``attn(Q[:, s:e], K[:, :e], V[:, :e], causal=True)``
-  - full-attn span ``[a, e)``: ``attn(Q[:, a:e], K[:, :e], V[:, :e], causal=False)``
+  - full-attn span ``[a, b)`` intersecting the query range at ``[s, e)``:
+    ``attn(Q[:, s:e], K[:, :b], V[:, :b], causal=False)``
 """
 
 from __future__ import annotations
 
-from typing import Literal, NamedTuple
+from typing import TYPE_CHECKING, Literal, NamedTuple
+
+import torch
+
+if TYPE_CHECKING:
+    from vllm_omni.diffusion.attention.backends.abstract import QueryRange
 
 
 class Segment(NamedTuple):
-    start: int
-    end: int
+    q_start: int
+    q_end: int
+    kv_end: int
     mode: Literal["causal", "full"]
 
 
@@ -30,30 +37,33 @@ def build_segments(full_attn_spans, query_offset, query_len):
     query_len: length of the query
 
     return:
-        List[Segment] in global coordinates, clipped to [query_offset, query_offset + query_len)
+        List[Segment] in global coordinates, clipped to
+        [query_offset, query_offset + query_len). Full-attention segments retain
+        the original span end as kv_end so a local query shard can attend past
+        its own boundary.
     """
     q_start = query_offset
     q_end = query_offset + query_len
 
-    segs: list[Segment] = []
+    segments: list[Segment] = []
     cur = q_start
 
-    for a, e in full_attn_spans:
+    for span_start, span_end in full_attn_spans:
         # clip span to query range
-        a_clipped = max(a, q_start)
-        e_clipped = min(e, q_end)
-        if a_clipped >= e_clipped:
+        overlap_start = max(span_start, q_start)
+        overlap_end = min(span_end, q_end)
+        if overlap_start >= overlap_end:
             continue
 
-        if cur < a_clipped:
-            segs.append(Segment(cur, a_clipped, "causal"))
-        segs.append(Segment(a_clipped, e_clipped, "full"))
-        cur = e_clipped
+        if cur < overlap_start:
+            segments.append(Segment(cur, overlap_start, overlap_start, "causal"))
+        segments.append(Segment(overlap_start, overlap_end, span_end, "full"))
+        cur = overlap_end
 
     if cur < q_end:
-        segs.append(Segment(cur, q_end, "causal"))
+        segments.append(Segment(cur, q_end, q_end, "causal"))
 
-    return segs
+    return segments
 
 
 def _check_homogeneous(
@@ -76,23 +86,40 @@ def piecewise_attn(
     full_attn_spans: list[list[tuple[int, int]]],
     softmax_scale: float,
     attn_func,
+    query_ranges: tuple[QueryRange, ...] | None = None,
 ):
-    B, Sq, H, D = query.shape
     _check_homogeneous(full_attn_spans)
-
-    query_offset = key.shape[1] - Sq
     spans = full_attn_spans[0]
-    out = query.new_zeros(B, Sq, H, D)
+    if query_ranges is None:
+        query_len = query.shape[1]
+        ranges = ((0, query_len, key.shape[1] - query_len),)
+    else:
+        ranges = tuple((r.local_start, r.local_end, r.global_start) for r in query_ranges)
 
-    for s, e, mode in build_segments(spans, query_offset, Sq):
-        q_s = s - query_offset
-        q_e = e - query_offset
-        out_seg = attn_func(
-            query[:, q_s:q_e],
-            key[:, :e],
-            value[:, :e],
-            causal=(mode == "causal"),
-            softmax_scale=softmax_scale,
-        )
-        out[:, q_s:q_e] = out_seg
-    return out
+    outputs = []
+    covered = 0
+    for local_start, local_end, global_start in ranges:
+        query_len = local_end - local_start
+        if local_start != covered or query_len < 0:
+            raise ValueError("query_ranges must cover local query contiguously")
+        for segment in build_segments(spans, global_start, query_len):
+            q_start = local_start + segment.q_start - global_start
+            q_end = local_start + segment.q_end - global_start
+            outputs.append(
+                attn_func(
+                    query[:, q_start:q_end],
+                    key[:, : segment.kv_end],
+                    value[:, : segment.kv_end],
+                    causal=(segment.mode == "causal"),
+                    softmax_scale=softmax_scale,
+                )
+            )
+        covered = local_end
+
+    if covered != query.shape[1]:
+        raise ValueError("query_ranges must cover the full local query")
+    if not outputs:
+        return torch.empty_like(query)
+    if len(outputs) == 1:
+        return outputs[0]
+    return torch.cat(outputs, dim=1)

vllm_omni/diffusion/attention/layer.py

@@ -140,6 +140,7 @@ class Attention(nn.Module):
             scatter_idx=scatter_idx,
             gather_idx=gather_idx,
             use_sync=use_sync,
+            causal=causal,
         )
         # Fallback strategy when SP is not active (outside sharded regions)
         self._no_parallel_strategy = NoParallelAttention()
@@ -282,6 +283,8 @@ class Attention(nn.Module):
         return out
 
     def _run_local_attention(self, query, key, value, attn_metadata):
+        self._assert_piecewise_compatible(attn_metadata)
+
         if query.dtype == torch.float32:
             logger.warning_once(
                 f"Only SDPA supports float32. Overriding user config {type(self.attention)} "
@@ -292,6 +295,20 @@ class Attention(nn.Module):
         # Fallback to standard attention
         return self.attention.forward(query, key, value, attn_metadata)
 
+    def _assert_piecewise_compatible(self, attn_metadata: AttentionMetadata | None) -> None:
+        if attn_metadata is None or attn_metadata.full_attn_spans is None:
+            return
+        if attn_metadata.attn_mask is not None and attn_metadata.attn_mask.ndim == 4:
+            return
+        backend_name = self.attn_backend.get_name()
+        if not self.attn_backend.supports_piecewise_spans:
+            raise ValueError(
+                f"Attention backend '{backend_name}' does not support "
+                f"piecewise attention (full_attn_spans without a 4D attn_mask). "
+                f"Use a Flash backend (FLASH_ATTN / FLASH_ATTN_HUB / FLASH_ATTN_3_HUB), "
+                f"or provide a 4D attn_mask that encodes the mixed causal/full pattern."
+            )
+
     def _run_ring_attention(self, query, key, value, attn_metadata):
         # Delegate to RingParallelAttention strategy if available
         if self.ring_runner is not None:

vllm_omni/diffusion/attention/parallel/allgather_kv.py

@@ -0,0 +1,182 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+from dataclasses import dataclass, replace
+from typing import TYPE_CHECKING
+
+import torch
+
+from vllm_omni.diffusion.attention.backends.abstract import QueryRange
+from vllm_omni.diffusion.attention.parallel.base import ParallelAttentionContext
+from vllm_omni.diffusion.distributed.group_coordinator import (
+    SequenceParallelGroupCoordinator,
+)
+
+if TYPE_CHECKING:
+    from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
+
+
+@dataclass(frozen=True, slots=True)
+class _AllGatherKVCtx(ParallelAttentionContext):
+    pass
+
+
+class AllGatherKVParallelAttention:
+    """Compute local Q against AllGathered K/V."""
+
+    def __init__(
+        self,
+        sp_group: SequenceParallelGroupCoordinator,
+    ) -> None:
+        self._sp_group = sp_group
+        self._allgather_group = sp_group.allgather_group
+        self._sp_size = sp_group.allgather_world_size
+        self._sp_rank = sp_group.allgather_rank
+
+    @property
+    def enabled(self) -> bool:
+        return True
+
+    @property
+    def name(self) -> str:
+        return "allgather_kv"
+
+    def pre_attention(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata | None,
+    ):
+        joint_q = joint_k = joint_v = None
+        joint_strategy = "front"
+        if attn_metadata is not None:
+            joint_q = attn_metadata.joint_query
+            joint_k = attn_metadata.joint_key
+            joint_v = attn_metadata.joint_value
+            joint_strategy = attn_metadata.joint_strategy or "front"
+        if joint_strategy not in {"front", "rear"}:
+            raise ValueError(f"Unsupported joint_strategy: {joint_strategy!r}")
+
+        k_img_full = self._sp_group.all_gather(key, dim=1, group=self._allgather_group)
+        v_img_full = self._sp_group.all_gather(value, dim=1, group=self._allgather_group)
+
+        if joint_k is not None:
+            if joint_k.shape[2] != key.shape[2]:
+                raise ValueError(
+                    "AllGather-KV SP expects joint_k to be GQA-compressed with the "
+                    f"same num_kv_heads as the image shard, got joint_k.heads="
+                    f"{joint_k.shape[2]} vs key.heads={key.shape[2]}."
+                )
+            if joint_strategy == "front":
+                k_full = torch.cat([joint_k, k_img_full], dim=1)
+                v_full = torch.cat([joint_v, v_img_full], dim=1)
+            else:
+                k_full = torch.cat([k_img_full, joint_k], dim=1)
+                v_full = torch.cat([v_img_full, joint_v], dim=1)
+        else:
+            k_full, v_full = k_img_full, v_img_full
+
+        joint_len = joint_q.shape[1] if joint_q is not None else 0
+        logical_q_full_len = joint_len + k_img_full.shape[1]
+        query_global_offset = k_full.shape[1] - logical_q_full_len
+        if query_global_offset < 0:
+            raise ValueError(
+                "AllGather-KV SP global K length is shorter than the logical Q length: "
+                f"k_len={k_full.shape[1]}, logical_q_len={logical_q_full_len}."
+            )
+
+        if joint_q is not None:
+            if joint_strategy == "front":
+                q_local = torch.cat([joint_q, query], dim=1)
+            else:
+                q_local = torch.cat([query, joint_q], dim=1)
+        else:
+            q_local = query
+
+        attn_metadata = self._slice_attn_metadata_for_local_query(
+            attn_metadata,
+            q_local_len=q_local.shape[1],
+            img_seq_local=query.shape[1],
+            img_seq_full=k_img_full.shape[1],
+            joint_len=joint_len,
+            joint_strategy=joint_strategy,
+            query_global_offset=query_global_offset,
+        )
+
+        return q_local, k_full, v_full, attn_metadata, _AllGatherKVCtx(name=self.name)
+
+    def _slice_attn_metadata_for_local_query(
+        self,
+        attn_metadata: AttentionMetadata | None,
+        *,
+        q_local_len: int,
+        img_seq_local: int,
+        img_seq_full: int,
+        joint_len: int,
+        joint_strategy: str,
+        query_global_offset: int,
+    ) -> AttentionMetadata | None:
+        """Slice global query metadata for the local query shard."""
+        if attn_metadata is None:
+            return None
+
+        img_start = self._sp_rank * img_seq_local
+        img_end = img_start + img_seq_local
+        if img_end > img_seq_full:
+            raise ValueError(
+                "AllGather-KV SP local image query range exceeds gathered image length: "
+                f"rank={self._sp_rank}, img_start={img_start}, img_end={img_end}, img_seq_full={img_seq_full}."
+            )
+
+        if joint_strategy == "front":
+            ranges = (
+                *((QueryRange(0, joint_len, query_global_offset),) if joint_len else ()),
+                QueryRange(joint_len, q_local_len, query_global_offset + joint_len + img_start),
+            )
+        else:
+            ranges = (
+                QueryRange(0, img_seq_local, query_global_offset + img_start),
+                *((QueryRange(img_seq_local, q_local_len, query_global_offset + img_seq_full),) if joint_len else ()),
+            )
+
+        if attn_metadata.attn_mask is None:
+            return replace(attn_metadata, query_ranges=ranges)
+
+        mask = attn_metadata.attn_mask
+        if mask.ndim != 4:
+            return replace(attn_metadata, query_ranges=ranges)
+
+        if mask.shape[-2] == q_local_len:
+            local_mask = mask
+        else:
+            q_full_len = joint_len + img_seq_full
+            if mask.shape[-2] != q_full_len:
+                raise ValueError(
+                    "AllGather-KV SP received an attention mask with incompatible Q length: "
+                    f"mask_q={mask.shape[-2]}, expected local_q={q_local_len} or full_q={q_full_len} "
+                    f"(joint_len={joint_len}, img_seq_local={img_seq_local}, img_seq_full={img_seq_full})."
+                )
+
+            parts = [
+                mask[..., local_start : local_start + (r.local_end - r.local_start), :]
+                for r in ranges
+                for local_start in (r.global_start - query_global_offset,)
+            ]
+            local_mask = torch.cat(parts, dim=-2)
+
+        if local_mask.shape[-2] != q_local_len:
+            raise ValueError(
+                "AllGather-KV SP produced an attention mask with incompatible local Q length: "
+                f"mask_q={local_mask.shape[-2]}, q_local={q_local_len}."
+            )
+        return replace(attn_metadata, attn_mask=local_mask.contiguous(), query_ranges=ranges)
+
+    def post_attention(
+        self,
+        attn_output: torch.Tensor,
+        ctx: ParallelAttentionContext | None,
+    ) -> torch.Tensor:
+        return attn_output

vllm_omni/diffusion/attention/parallel/factory.py

@@ -5,10 +5,19 @@ from __future__ import annotations
 
 from vllm.logger import init_logger
 
-from vllm_omni.diffusion.attention.parallel.base import NoParallelAttention, ParallelAttentionStrategy
+from vllm_omni.diffusion.attention.parallel.allgather_kv import (
+    AllGatherKVParallelAttention,
+)
+from vllm_omni.diffusion.attention.parallel.base import (
+    NoParallelAttention,
+    ParallelAttentionStrategy,
+)
 from vllm_omni.diffusion.attention.parallel.ring import RingParallelAttention
 from vllm_omni.diffusion.attention.parallel.ulysses import UlyssesParallelAttention
-from vllm_omni.diffusion.distributed.parallel_state import get_sequence_parallel_world_size, get_sp_group
+from vllm_omni.diffusion.distributed.parallel_state import (
+    get_sequence_parallel_world_size,
+    get_sp_group,
+)
 from vllm_omni.diffusion.forward_context import get_forward_context
 
 logger = init_logger(__name__)
@@ -19,6 +28,7 @@ def build_parallel_attention_strategy(
     scatter_idx: int,
     gather_idx: int,
     use_sync: bool,
+    causal: bool = False,
 ) -> ParallelAttentionStrategy:
     """Select a parallel attention strategy based on current diffusion config.
 
@@ -36,6 +46,7 @@ def build_parallel_attention_strategy(
 
     ulysses_degree = getattr(p, "ulysses_degree", 1)
     ring_degree = getattr(p, "ring_degree", 1)
+    allgather_degree = getattr(p, "allgather_degree", 1)
 
     try:
         sp_group = get_sp_group()
@@ -44,13 +55,26 @@ def build_parallel_attention_strategy(
             return NoParallelAttention()
     except Exception as e:
         # Log warning if SP is configured but group is not available
-        if ulysses_degree > 1 or ring_degree > 1:
+        if ulysses_degree > 1 or ring_degree > 1 or allgather_degree > 1:
             logger.warning(
-                f"SP configured (ulysses={ulysses_degree}, ring={ring_degree}) but SP group not available: {e}. "
+                f"SP configured (ulysses={ulysses_degree}, ring={ring_degree}, "
+                f"allgather={allgather_degree}) but SP group not available: {e}. "
                 f"Falling back to NoParallelAttention. This may cause incorrect results."
             )
         return NoParallelAttention()
 
+    if allgather_degree > 1:
+        if causal:
+            raise ValueError("AllGather-KV SP only supports non-causal attention.")
+        if ulysses_degree > 1 or ring_degree > 1:
+            raise ValueError(
+                f"AllGather-KV SP is mutually exclusive with Ulysses/Ring in v1 "
+                f"(got ulysses_degree={ulysses_degree}, ring_degree={ring_degree}, "
+                f"allgather_degree={allgather_degree})."
+            )
+        logger.debug(f"Using AllGatherKVParallelAttention (allgather_degree={allgather_degree})")
+        return AllGatherKVParallelAttention(sp_group=sp_group)
+
     # Ulysses (or Hybrid Ulysses+Ring)
     if ulysses_degree > 1:
         logger.debug(f"Using UlyssesParallelAttention (ulysses_degree={ulysses_degree})")

vllm_omni/diffusion/data.py

@@ -156,7 +156,9 @@ class DiffusionParallelConfig:
     """Enable expert parallelism for MoE layers (TP is still used for non-MoE layers)."""
 
     sequence_parallel_size: int | None = None
-    """Number of sequence parallel groups. sequence_parallel_size = ring_degree * ulysses_degree"""
+    """Number of sequence parallel groups.
+    sequence_parallel_size = ulysses_degree * ring_degree, or allgather_degree
+    when AllGather-KV is enabled."""
 
     ulysses_degree: int = 1
     """Number of GPUs used for ulysses sequence parallelism."""
@@ -164,6 +166,9 @@ class DiffusionParallelConfig:
     ring_degree: int = 1
     """Number of GPUs used for ring sequence parallelism."""
 
+    allgather_degree: int = 1
+    """Number of GPUs used for AllGather-KV sequence parallelism (causal=False only)."""
+
     ulysses_mode: str = "strict"
     """Ulysses sequence-parallel mode.
 
@@ -225,6 +230,7 @@ class DiffusionParallelConfig:
         assert self.sequence_parallel_size > 0, "Sequence parallel size must be > 0"
         assert self.ulysses_degree > 0, "Ulysses degree must be > 0"
         assert self.ring_degree > 0, "Ring degree must be > 0"
+        assert self.allgather_degree > 0, "AllGather degree must be > 0"
         assert self.cfg_parallel_size > 0, "CFG parallel size must be > 0"
         assert self.cfg_parallel_size in [1, 2, 3], (
             f"CFG parallel size must be 1, 2, or 3, but got {self.cfg_parallel_size}"
@@ -234,9 +240,17 @@ class DiffusionParallelConfig:
             "vae_parallel_mode must be one of {'tile', 'spatial_shard_height', 'spatial_shard_width'}, "
             f"but got {self.vae_parallel_mode!r}."
         )
-        assert self.sequence_parallel_size == self.ulysses_degree * self.ring_degree, (
-            "Sequence parallel size must be equal to the product of ulysses degree and ring degree,"
-            f" but got {self.sequence_parallel_size} != {self.ulysses_degree} * {self.ring_degree}"
+        if self.allgather_degree > 1:
+            assert self.ulysses_degree == 1 and self.ring_degree == 1, (
+                "AllGather-KV (allgather_degree>1) is mutually exclusive with Ulysses/Ring in v1. "
+                f"Got ulysses_degree={self.ulysses_degree}, ring_degree={self.ring_degree}, "
+                f"allgather_degree={self.allgather_degree}."
+            )
+        expected_sp_size = (
+            self.allgather_degree if self.allgather_degree > 1 else self.ulysses_degree * self.ring_degree
+        )
+        assert self.sequence_parallel_size == expected_sp_size, (
+            f"Sequence parallel size must be {expected_sp_size}, but got {self.sequence_parallel_size}"
         )
         assert self.ulysses_mode in {"strict", "advanced_uaa"}, (
             f"ulysses_mode must be one of {{'strict','advanced_uaa'}}, but got {self.ulysses_mode!r}."
@@ -250,15 +264,16 @@ class DiffusionParallelConfig:
 
     def __post_init__(self) -> None:
         if self.sequence_parallel_size is None:
-            self.sequence_parallel_size = self.ulysses_degree * self.ring_degree
+            self.sequence_parallel_size = (
+                self.allgather_degree if self.allgather_degree > 1 else self.ulysses_degree * self.ring_degree
+            )
 
         # Calculate world_size from other parallelism dimensions
         other_parallel_world_size = (
             self.pipeline_parallel_size
             * self.data_parallel_size
             * self.tensor_parallel_size
-            * self.ulysses_degree
-            * self.ring_degree
+            * self.sequence_parallel_size
             * self.cfg_parallel_size
         )
 

vllm_omni/diffusion/distributed/group_coordinator.py

@@ -199,9 +199,15 @@ class GroupCoordinator:
         return input_
 
     def all_gather(
-        self, input_: torch.Tensor, dim: int = 0, separate_tensors: bool = False
+        self,
+        input_: torch.Tensor,
+        dim: int = 0,
+        separate_tensors: bool = False,
+        group: ProcessGroup | None = None,
     ) -> torch.Tensor | list[torch.Tensor]:
-        world_size = self.world_size
+        if group is None:
+            group = self.device_group
+        world_size = torch.distributed.get_world_size(group)
         # Bypass the function if we are using only 1 GPU.
         if world_size == 1:
             return input_
@@ -214,7 +220,7 @@ class GroupCoordinator:
         input_size[0] *= world_size
         output_tensor = torch.empty(input_size, dtype=input_.dtype, device=input_.device)
         # All-gather.
-        torch.distributed.all_gather_into_tensor(output_tensor, input_.contiguous(), group=self.device_group)
+        torch.distributed.all_gather_into_tensor(output_tensor, input_.contiguous(), group=group)
         if dim != 0:
             input_size[0] //= world_size
             output_tensor = output_tensor.reshape(
@@ -1002,6 +1008,7 @@ class SequenceParallelGroupCoordinator(GroupCoordinator):
 
         ulysses_group = kwargs.get("ulysses_group", None)
         ring_group = kwargs.get("ring_group", None)
+        allgather_group = kwargs.get("allgather_group", None)
         if ulysses_group is None:
             raise RuntimeError(
                 "Please pass argument 'ulysses_group' when calling init func of SequenceParallelGroupCoordinator"
@@ -1010,10 +1017,17 @@ class SequenceParallelGroupCoordinator(GroupCoordinator):
             raise RuntimeError(
                 "Please pass argument 'ring_group' when calling init func of SequenceParallelGroupCoordinator"
             )
+        if allgather_group is None:
+            raise RuntimeError(
+                "Please pass argument 'allgather_group' when calling init func of SequenceParallelGroupCoordinator"
+            )
         self.ulysses_group = ulysses_group
         self.ring_group = ring_group
+        self.allgather_group = allgather_group
 
         self.ulysses_world_size = torch.distributed.get_world_size(self.ulysses_group)
         self.ulysses_rank = torch.distributed.get_rank(self.ulysses_group)
         self.ring_world_size = torch.distributed.get_world_size(self.ring_group)
         self.ring_rank = torch.distributed.get_rank(self.ring_group)
+        self.allgather_world_size = torch.distributed.get_world_size(self.allgather_group)
+        self.allgather_rank = torch.distributed.get_rank(self.allgather_group)

vllm_omni/diffusion/distributed/parallel_state.py

@@ -303,6 +303,14 @@ def get_ring_parallel_rank():
     return get_sp_group().ring_rank
 
 
+def get_allgather_parallel_world_size():
+    return get_sp_group().allgather_world_size
+
+
+def get_allgather_parallel_rank():
+    return get_sp_group().allgather_rank
+
+
 def get_expert_parallel_group_ranks() -> list[list[int]]:
     assert vllm_parallel_state._EP is not None, "expert parallel group is not initialized"
     assert _EXPERT_PARALLEL_GROUP_RANKS is not None, "expert parallel group ranks are not initialized"
@@ -548,13 +556,12 @@ def set_seq_parallel_pg(
     world_size: int,
     use_ulysses_low: bool = True,
     sp_group_ranks: list[list[int]] | None = None,
-) -> tuple[torch.distributed.ProcessGroup, torch.distributed.ProcessGroup]:
+    sp_allgather_degree: int = 1,
+) -> tuple[torch.distributed.ProcessGroup, torch.distributed.ProcessGroup, torch.distributed.ProcessGroup]:
     """
-    Initialize sequence-parallel Ulysses and Ring process groups.
+    Initialize Ulysses, Ring, and AllGather-KV process groups.
 
-    This builds sequence-parallel (SP) subgroups inside each data-parallel (DP)
-    slice. The SP group size is sp_ulysses_degree * sp_ring_degree, and
-    world_size must be divisible by that size.
+    AllGather-KV is mutually exclusive with Ulysses and Ring.
 
     Args:
         sp_ulysses_degree: Size of each Ulysses subgroup.
@@ -569,7 +576,7 @@ def set_seq_parallel_pg(
             ranges.
 
     Returns:
-        ulyssess_pg (torch.distributed.ProcessGroup): The Ulysses process group
+        ulysses_pg (torch.distributed.ProcessGroup): The Ulysses process group
             for this rank.
         ring_pg (torch.distributed.ProcessGroup): The Ring process group for
             this rank.
@@ -586,6 +593,31 @@ def set_seq_parallel_pg(
         - If sp_group_ranks is None, groups are auto-generated within each DP
           slice using offsets of size sp_size.
     """
+    if sp_allgather_degree > 1:
+        if sp_ulysses_degree > 1 or sp_ring_degree > 1:
+            raise ValueError("AllGather-KV is mutually exclusive with Ulysses and Ring")
+        sp_size = sp_allgather_degree
+        if sp_group_ranks is None:
+            sp_group_ranks = [list(range(offset, offset + sp_size)) for offset in range(0, world_size, sp_size)]
+        if len(sp_group_ranks) * sp_size != world_size or any(len(ranks) != sp_size for ranks in sp_group_ranks):
+            raise ValueError(f"Invalid sp_group_ranks: expected {world_size // sp_size} groups of size {sp_size}.")
+
+        ulysses_pg = ring_pg = allgather_pg = None
+        for group_ranks in sp_group_ranks:
+            group = torch.distributed.new_group(group_ranks)
+            if rank in group_ranks:
+                allgather_pg = group
+        for singleton_rank in range(world_size):
+            group = torch.distributed.new_group([singleton_rank])
+            if rank == singleton_rank:
+                ulysses_pg = group
+        for singleton_rank in range(world_size):
+            group = torch.distributed.new_group([singleton_rank])
+            if rank == singleton_rank:
+                ring_pg = group
+        assert ulysses_pg is not None and ring_pg is not None and allgather_pg is not None
+        return ulysses_pg, ring_pg, allgather_pg
+
     sp_size = sp_ring_degree * sp_ulysses_degree
     dp_size = world_size // sp_size
 
@@ -619,7 +651,7 @@ def set_seq_parallel_pg(
                     ulysses_ranks = group_ranks[i * sp_ulysses_degree : (i + 1) * sp_ulysses_degree]
                     group = torch.distributed.new_group(ulysses_ranks)
                     if rank in ulysses_ranks:
-                        ulyssess_pg = group
+                        ulysses_pg = group
                         local_ulysses = list(ulysses_ranks)
                 for i in range(num_ring_pgs):
                     ring_ranks = group_ranks[i::num_ring_pgs]
@@ -639,7 +671,7 @@ def set_seq_parallel_pg(
                     ulysses_ranks = group_ranks[i::num_ulysses_pgs]
                     group = torch.distributed.new_group(ulysses_ranks)
                     if rank in ulysses_ranks:
-                        ulyssess_pg = group
+                        ulysses_pg = group
                         local_ulysses = list(ulysses_ranks)
         if local_sp_group is not None:
             logger.info(
@@ -662,7 +694,7 @@ def set_seq_parallel_pg(
                     )
                     group = torch.distributed.new_group(ulysses_ranks)
                     if rank in ulysses_ranks:
-                        ulyssess_pg = group
+                        ulysses_pg = group
 
                 for i in range(num_ring_pgs):
                     ring_ranks = list(range(i + offset, sp_size + offset, num_ring_pgs))
@@ -683,9 +715,15 @@ def set_seq_parallel_pg(
                     ulysses_ranks = list(range(i + offset, sp_size + offset, num_ulysses_pgs))
                     group = torch.distributed.new_group(ulysses_ranks)
                     if rank in ulysses_ranks:
-                        ulyssess_pg = group
+                        ulysses_pg = group
 
-    return ulyssess_pg, ring_pg
+    allgather_pg = None
+    for singleton_rank in range(world_size):
+        group = torch.distributed.new_group([singleton_rank])
+        if rank == singleton_rank:
+            allgather_pg = group
+    assert allgather_pg is not None
+    return ulysses_pg, ring_pg, allgather_pg
 
 
 def initialize_model_parallel(
@@ -694,6 +732,7 @@ def initialize_model_parallel(
     sequence_parallel_size: int | None = None,
     ulysses_degree: int = 1,
     ring_degree: int = 1,
+    allgather_degree: int = 1,
     tensor_parallel_size: int = 1,
     pipeline_parallel_size: int = 1,
     fully_shard_degree: int = 1,
@@ -710,9 +749,12 @@ def initialize_model_parallel(
         data_parallel_size: number of data parallelism groups.
         cfg_parallel_size: number of GPUs used for Classifier Free Guidance (CFG) parallelism.
         sequence_parallel_size: number of GPUs used for sequence parallelism.
-            sequence_parallel_size = ulysses_degree * ring_degree
+            Uses allgather_degree when AllGather-KV is enabled, otherwise
+            ulysses_degree * ring_degree.
         ulysses_degree: number of GPUs used for ulysses sequence parallelism.
         ring_degree: number of GPUs used for ring sequence parallelism.
+        allgather_degree: number of GPUs used for AllGather-KV sequence parallelism
+            (causal=False only). Mutually exclusive with ulysses/ring in v1.
         tensor_parallel_size: number of GPUs used for tensor parallelism.
         pipeline_parallel_size: number of GPUs used for pipeline parallelism.
         fully_shard_degree: number of GPUs used for fully sharded data parallelism (HSDP shard dimension).
@@ -749,16 +791,22 @@ def initialize_model_parallel(
     world_size: int = torch.distributed.get_world_size()
     backend = backend or torch.distributed.get_backend(get_world_group().device_group)
 
-    if sequence_parallel_size is None:
-        sequence_parallel_size = ring_degree * ulysses_degree
-        logger.info(
-            f"sequence_parallel_size is not provided, using ring_degree * ulysses_degree = {sequence_parallel_size}"
+    if allgather_degree > 1:
+        assert ulysses_degree == 1 and ring_degree == 1, (
+            "AllGather-KV (allgather_degree>1) is mutually exclusive with Ulysses/Ring in v1. "
+            f"Got ulysses_degree={ulysses_degree}, ring_degree={ring_degree}, "
+            f"allgather_degree={allgather_degree}."
         )
 
-    if sequence_parallel_size != ring_degree * ulysses_degree:
+    expected_sequence_parallel_size = allgather_degree if allgather_degree > 1 else ring_degree * ulysses_degree
+    if sequence_parallel_size is None:
+        sequence_parallel_size = expected_sequence_parallel_size
+        logger.info("sequence_parallel_size is not provided, using %d", sequence_parallel_size)
+
+    if sequence_parallel_size != expected_sequence_parallel_size:
         raise ValueError(
-            "sequence_parallel_size is not equal to ring_degree * ulysses_degree,"
-            f" but got {sequence_parallel_size} != {ring_degree} * {ulysses_degree}"
+            f"sequence_parallel_size must be {expected_sequence_parallel_size} for the configured SP mode, "
+            f"but got {sequence_parallel_size}"
         )
 
     dit_parallel_size = (
@@ -835,9 +883,10 @@ def initialize_model_parallel(
 
     global _SP
     assert _SP is None, "sequence parallel group is already initialized"
-    ulysses_pg, ring_pg = set_seq_parallel_pg(
+    ulysses_pg, ring_pg, allgather_pg = set_seq_parallel_pg(
         sp_ulysses_degree=ulysses_degree,
         sp_ring_degree=ring_degree,
+        sp_allgather_degree=allgather_degree,
         rank=get_world_group().rank_in_group,
         world_size=dit_parallel_size,
         sp_group_ranks=sp_group_ranks,
@@ -849,6 +898,7 @@ def initialize_model_parallel(
         parallel_mode="sequence",
         ulysses_group=ulysses_pg,
         ring_group=ring_pg,
+        allgather_group=allgather_pg,
     )
     if use_moe_parallel_mapping:
         # Diffusion normally uses its own SP group. Map it to vLLM PCP only for

vllm_omni/diffusion/distributed/sp_plan.py

@@ -77,6 +77,7 @@ class SequenceParallelConfig:
 
     ulysses_degree: int = 1
     ring_degree: int = 1
+    allgather_degree: int = 1
     convert_to_fp32: bool = True
 
     # Internal state - populated by setup()
@@ -85,18 +86,18 @@ class SequenceParallelConfig:
     _device: torch.device | None = None
 
     def __post_init__(self) -> None:
-        if self.ulysses_degree < 1 or self.ring_degree < 1:
-            raise ValueError("`ulysses_degree` and `ring_degree` must be >= 1.")
+        if self.ulysses_degree < 1 or self.ring_degree < 1 or self.allgather_degree < 1:
+            raise ValueError("SP degrees must be >= 1.")
 
-        if self.ulysses_degree == 1 and self.ring_degree == 1:
-            raise ValueError(
-                "At least one of `ulysses_degree` or `ring_degree` must be > 1 to use sequence parallelism."
-            )
+        if self.allgather_degree > 1 and (self.ulysses_degree > 1 or self.ring_degree > 1):
+            raise ValueError("AllGather-KV is mutually exclusive with Ulysses and Ring.")
+        if self.ulysses_degree == self.ring_degree == self.allgather_degree == 1:
+            raise ValueError("At least one SP degree must be > 1.")
 
     @property
     def sequence_parallel_size(self) -> int:
         """Total sequence parallel world size."""
-        return self.ulysses_degree * self.ring_degree
+        return self.allgather_degree if self.allgather_degree > 1 else self.ulysses_degree * self.ring_degree
 
     def get_world_size(self) -> int:
         """Get the sequence parallel world size from parallel state.
@@ -178,13 +179,12 @@ class SequenceParallelConfig:
         self._world_size = world_size
         self._device = device
 
-        expected_sp_size = self.ulysses_degree * self.ring_degree
+        expected_sp_size = self.sequence_parallel_size
         actual_sp_size = self.get_world_size()
 
         if expected_sp_size != actual_sp_size:
             raise ValueError(
-                f"Configuration mismatch: ulysses_degree ({self.ulysses_degree}) * "
-                f"ring_degree ({self.ring_degree}) = {expected_sp_size}, but "
+                f"Configuration mismatch: expected SP size {expected_sp_size}, but "
                 f"actual sequence parallel world size is {actual_sp_size}."
             )
 

vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py

@@ -62,6 +62,7 @@ from vllm_omni.diffusion.attention.backends.abstract import (
 from vllm_omni.diffusion.attention.layer import Attention
 from vllm_omni.diffusion.cache.cachedit import CacheDiTAdapterConfig
 from vllm_omni.diffusion.distributed.parallel_state import (
+    get_allgather_parallel_world_size,
     get_cfg_group,
     get_classifier_free_guidance_rank,
     get_classifier_free_guidance_world_size,
@@ -78,6 +79,7 @@ from vllm_omni.diffusion.forward_context import set_forward_context_denoise_step
 from vllm_omni.diffusion.layers.norm import RMSNorm
 from vllm_omni.diffusion.layers.rope import RotaryEmbedding
 from vllm_omni.diffusion.models.hunyuan_image3.hunyuan_fused_moe import HunyuanFusedMoE
+from vllm_omni.diffusion.utils.kv_utils import repeat_kv
 from vllm_omni.model_executor.layers.timestep_embedding import timestep_embedding
 
 logger = logging.getLogger(__name__)
@@ -426,22 +428,6 @@ def apply_rotary_pos_emb(
     return q_embed, k_embed
 
 
-def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
-    """
-    Equivalent to torch.repeat_interleave(x, dim=2, repeats=n_rep).
-    Input:  (batch, seqlen, num_key_value_heads, head_dim)
-    Output: (batch, seqlen, num_attention_heads, head_dim)
-    """
-    batch, slen, num_key_value_heads, head_dim = hidden_states.shape
-
-    if n_rep == 1:
-        return hidden_states
-
-    hidden_states = hidden_states[:, :, :, None, :].expand(batch, slen, num_key_value_heads, n_rep, head_dim)
-
-    return hidden_states.reshape(batch, slen, num_key_value_heads * n_rep, head_dim)
-
-
 def default(value, default_value):
     return value if value is not None else default_value
 
@@ -957,6 +943,7 @@ class ImageKVCacheManager:
         self._injected_ar_kv: list[tuple[torch.Tensor, torch.Tensor]] | None = None
 
         self.sp_size = get_sequence_parallel_world_size()
+        self.allgather_size = get_allgather_parallel_world_size()
         self.sp_rank = get_sequence_parallel_rank()
         self.attn = Attention(
             num_heads=self.num_heads,
@@ -1130,6 +1117,7 @@ class ImageKVCacheManager:
         head_num_per_rank = query.shape[1]
         kv_head_num_per_rank = key.shape[1]
         repeat_num = head_num_per_rank // kv_head_num_per_rank
+        keep_kv_compressed = self.allgather_size > 1
         head_dim = query.shape[2]
 
         query = query.reshape(bs, q_len, head_num_per_rank, head_dim)
@@ -1171,11 +1159,12 @@ class ImageKVCacheManager:
                 joint_text_query = query[:, :0, :, :]
                 joint_text_key, joint_text_value = self._reuse_prompt_kv(key, value, seq_len, bs, shard_image_size)
 
-        key = repeat_kv(key, repeat_num)
-        value = repeat_kv(value, repeat_num)
-        if self.sp_size > 1:
-            joint_text_key = repeat_kv(joint_text_key, repeat_num)
-            joint_text_value = repeat_kv(joint_text_value, repeat_num)
+        if not keep_kv_compressed:
+            key = repeat_kv(key, repeat_num)
+            value = repeat_kv(value, repeat_num)
+            if self.sp_size > 1:
+                joint_text_key = repeat_kv(joint_text_key, repeat_num)
+                joint_text_value = repeat_kv(joint_text_value, repeat_num)
 
         attention_mask = attention_mask.contiguous()
 

vllm_omni/diffusion/registry.py

@@ -464,18 +464,25 @@ def _apply_sequence_parallel_if_enabled(model, od_config: OmniDiffusionConfig) -
             if plan is None:
                 continue
 
-            # Create SP config
-            sp_config = SequenceParallelConfig(
-                ulysses_degree=od_config.parallel_config.ulysses_degree,
-                ring_degree=od_config.parallel_config.ring_degree,
-            )
+            # AllGather-KV reuses the Ulysses sequence-sharding hooks.
+            allgather_degree = getattr(od_config.parallel_config, "allgather_degree", 1)
+            if allgather_degree > 1:
+                sp_config = SequenceParallelConfig(
+                    allgather_degree=allgather_degree,
+                )
+                mode = "allgather_kv"
+            else:
+                sp_config = SequenceParallelConfig(
+                    ulysses_degree=od_config.parallel_config.ulysses_degree,
+                    ring_degree=od_config.parallel_config.ring_degree,
+                )
+                # Apply hooks according to the plan
+                mode = (
+                    "hybrid"
+                    if sp_config.ulysses_degree > 1 and sp_config.ring_degree > 1
+                    else ("ulysses" if sp_config.ulysses_degree > 1 else "ring")
+                )
 
-            # Apply hooks according to the plan
-            mode = (
-                "hybrid"
-                if sp_config.ulysses_degree > 1 and sp_config.ring_degree > 1
-                else ("ulysses" if sp_config.ulysses_degree > 1 else "ring")
-            )
             logger.info(
                 f"Applying sequence parallelism to {transformer.__class__.__name__} ({attr}) "
                 f"(sp_size={sp_size}, mode={mode}, ulysses={sp_config.ulysses_degree}, ring={sp_config.ring_degree})"

vllm_omni/diffusion/utils/kv_utils.py

@@ -3,6 +3,19 @@
 import torch
 
 
+def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+    """Expand GQA KV from num_kv_heads to num_heads.
+
+    Input:  (batch, seqlen, num_key_value_heads, head_dim)
+    Output: (batch, seqlen, num_key_value_heads * n_rep, head_dim)
+    """
+    if n_rep == 1:
+        return hidden_states
+    batch, slen, num_key_value_heads, head_dim = hidden_states.shape
+    hidden_states = hidden_states[:, :, :, None, :].expand(batch, slen, num_key_value_heads, n_rep, head_dim)
+    return hidden_states.reshape(batch, slen, num_key_value_heads * n_rep, head_dim)
+
+
 def left_pad_stack(
     tensors: list[torch.Tensor],
 ) -> tuple[torch.Tensor, torch.Tensor | None]:

vllm_omni/diffusion/worker/diffusion_worker.py

@@ -315,6 +315,7 @@ class DiffusionWorker:
                 sequence_parallel_size=parallel_config.sequence_parallel_size,
                 ulysses_degree=parallel_config.ulysses_degree,
                 ring_degree=parallel_config.ring_degree,
+                allgather_degree=parallel_config.allgather_degree,
                 tensor_parallel_size=parallel_config.tensor_parallel_size,
                 pipeline_parallel_size=parallel_config.pipeline_parallel_size,
                 fully_shard_degree=parallel_config.hsdp_shard_size if parallel_config.use_hsdp else 1,

vllm_omni/engine/arg_utils.py

@@ -459,6 +459,7 @@ class OrchestratorArgs:
     ulysses_degree: int | None = None
     ulysses_mode: str = "strict"
     ring_degree: int | None = None
+    allgather_degree: int | None = None
     diffusion_quantization_config: str | None = None
     use_hsdp: bool = False
     hsdp_shard_size: int = -1

vllm_omni/engine/async_omni_engine.py

@@ -933,6 +933,7 @@ class AsyncOmniEngine:
         if parallel_config is None:
             ulysses_degree = normalized_kwargs.get("ulysses_degree") or 1
             ring_degree = normalized_kwargs.get("ring_degree") or 1
+            allgather_degree = normalized_kwargs.get("allgather_degree") or 1
             ulysses_mode = normalized_kwargs.get("ulysses_mode") or "strict"
             sequence_parallel_size = normalized_kwargs.get("sequence_parallel_size")
             pipeline_parallel_size = normalized_kwargs.get("pipeline_parallel_size") or 1
@@ -947,7 +948,7 @@ class AsyncOmniEngine:
             hsdp_shard_size = normalized_kwargs.get("hsdp_shard_size", -1)
             hsdp_replicate_size = normalized_kwargs.get("hsdp_replicate_size", 1)
             if sequence_parallel_size is None:
-                sequence_parallel_size = ulysses_degree * ring_degree
+                sequence_parallel_size = allgather_degree if allgather_degree > 1 else ulysses_degree * ring_degree
 
             parallel_config = DiffusionParallelConfig(
                 pipeline_parallel_size=pipeline_parallel_size,
@@ -957,6 +958,7 @@ class AsyncOmniEngine:
                 sequence_parallel_size=sequence_parallel_size,
                 ulysses_degree=ulysses_degree,
                 ring_degree=ring_degree,
+                allgather_degree=allgather_degree,
                 ulysses_mode=ulysses_mode,
                 cfg_parallel_size=cfg_parallel_size,
                 vae_patch_parallel_size=vae_patch_parallel_size,

vllm_omni/entrypoints/cli/serve.py

@@ -473,6 +473,14 @@ class OmniServeCommand(CLISubcommand):
             help="Ring Sequence Parallelism degree for diffusion models. "
             "Equivalent to setting DiffusionParallelConfig.ring_degree.",
         )
+        omni_config_group.add_argument(
+            "--allgather-degree",
+            dest="allgather_degree",
+            type=int,
+            default=None,
+            help="AllGather-KV Sequence Parallelism degree for non-causal diffusion attention. "
+            "Equivalent to setting DiffusionParallelConfig.allgather_degree.",
+        )
         omni_config_group.add_argument(
             "--diffusion-quantization-config",
             type=json.loads,