[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/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