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