[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_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)