[Model] Add MiniMax H3 diffusion support (#5691)

tests/diffusion/attention/test_flash_attn.py

@@ -20,7 +20,6 @@ from vllm_omni.platforms import current_omni_platform
 
 pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
 
-
 is_gpu = current_omni_platform.is_cuda_alike() or current_omni_platform.is_xpu()
 HAS_FLASH_ATTN = fa.HAS_FLASH_ATTN
 flash_attn_func = fa.flash_attn_func  # noqa: N813
@@ -330,6 +329,48 @@ def test_piecewise_flash_attn_uses_varlen_fallback(monkeypatch):
     assert calls[0]["softmax_scale"] == 0.5
 
 
+def test_packed_varlen_metadata_bypasses_mask_unpadding(monkeypatch):
+    calls = []
+
+    def fake_varlen_func(q, k, v, **kwargs):
+        calls.append((q.shape, k.shape, v.shape, kwargs))
+        return q
+
+    monkeypatch.setattr(fa, "HAS_FLASH_ATTN", True)
+    monkeypatch.setattr(fa, "flash_attn_varlen_func", fake_varlen_func)
+
+    impl = FlashAttentionImpl(num_heads=2, head_size=4, softmax_scale=0.5, causal=False)
+    query = torch.randn(1, 8, 2, 4)
+    cu_seqlens = torch.tensor([0, 6, 8], dtype=torch.int32)
+    metadata = AttentionMetadata(
+        attn_mask=torch.tensor([[True] * 6 + [False] * 2]),
+        extra={
+            "cu_seqlens_q": cu_seqlens,
+            "cu_seqlens_k": cu_seqlens,
+            "max_seqlen_q": 6,
+            "max_seqlen_k": 6,
+        },
+    )
+
+    output = impl.forward_cuda(query, query, query, metadata)
+
+    assert torch.equal(output, query)
+    assert len(calls) == 1
+    assert calls[0][0] == torch.Size([8, 2, 4])
+    assert calls[0][3]["cu_seqlens_q"] is cu_seqlens
+    assert calls[0][3]["max_seqlen_q"] == 6
+
+
+def test_packed_varlen_metadata_must_be_complete(monkeypatch):
+    monkeypatch.setattr(fa, "HAS_FLASH_ATTN", True)
+    impl = FlashAttentionImpl(num_heads=2, head_size=4, softmax_scale=0.5, causal=False)
+    query = torch.randn(1, 8, 2, 4)
+    metadata = AttentionMetadata(extra={"cu_seqlens_q": torch.tensor([0, 8], dtype=torch.int32)})
+
+    with pytest.raises(ValueError, match="Incomplete packed FlashAttention metadata"):
+        impl.forward_cuda(query, query, query, metadata)
+
+
 if __name__ == "__main__":
     print("Running FlashAttention Padding Tests...")
     print("=" * 60)

tests/diffusion/attention/test_ring_fa4.py

@@ -0,0 +1,74 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.attention.backends.ring import ring_kernels
+from vllm_omni.diffusion.attention.backends.ring.ring_selector import AttnType
+from vllm_omni.diffusion.attention.parallel import ring as ring_parallel
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+def test_fa4_wrapper_returns_lse_and_normalizes_unbounded_window(monkeypatch):
+    expected_out = object()
+    expected_lse = object()
+    captured = {}
+
+    def fake_fa4(q, k, v, **kwargs):
+        captured.update(q=q, k=k, v=v, **kwargs)
+        return expected_out, expected_lse
+
+    monkeypatch.setattr(ring_kernels, "HAS_FA4", True)
+    monkeypatch.setattr(ring_kernels, "fa4_attn_func", fake_fa4)
+
+    out, lse = ring_kernels.flash_attn4_func_forward(
+        "q",
+        "k",
+        "v",
+        softmax_scale=0.125,
+        causal=True,
+        window_size=(-1, -1),
+    )
+
+    assert (out, lse) == (expected_out, expected_lse)
+    assert captured == {
+        "q": "q",
+        "k": "k",
+        "v": "v",
+        "softmax_scale": 0.125,
+        "causal": True,
+        "window_size": (None, None),
+        "softcap": 0.0,
+        "return_lse": True,
+    }
+
+
+def test_ring_prefers_fa4_on_blackwell(monkeypatch):
+    captured = {}
+
+    def fake_ring_flash_attn_func(*args, **kwargs):
+        captured.update(kwargs)
+        return "output"
+
+    monkeypatch.setattr(ring_parallel, "HAS_FA4", True)
+    monkeypatch.setattr(ring_parallel, "HAS_FA3", True)
+    monkeypatch.setattr(torch.cuda, "get_device_capability", lambda _device: (10, 3))
+
+    from vllm_omni.diffusion.attention.backends import ring_flash_attn
+
+    monkeypatch.setattr(ring_flash_attn, "ring_flash_attn_func", fake_ring_flash_attn_func)
+
+    query = SimpleNamespace(
+        dtype=torch.bfloat16,
+        shape=(1, 16, 4, 128),
+        is_cuda=True,
+        device=torch.device("cuda:0"),
+    )
+    strategy = ring_parallel.RingParallelAttention(SimpleNamespace(ring_group=object()))
+
+    assert strategy.run_attention(query, "key", "value", None) == "output"
+    assert captured["attn_type"] is AttnType.FA4