[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/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),
     ]