[Perf][Diffusion] Run MammothModa2 DiT attention through the shared attention layer (#7094)

Signed-off-by: MrlixiangWE <mrdanaer@gmail.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

tests/diffusion/attention/test_attention_config.py

@@ -10,6 +10,7 @@ Tests cover:
 - AttentionMetadata.extra field
 """
 
+from dataclasses import replace
 from types import SimpleNamespace
 from unittest.mock import Mock
 
@@ -17,7 +18,12 @@ import pytest
 import torch
 
 import vllm_omni.diffusion.attention.layer as layer_mod
-from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
+import vllm_omni.diffusion.attention.selector as selector_mod
+from vllm_omni.diffusion.attention.backends.abstract import (
+    AttentionMetadata,
+    PackedPaddingMetadata,
+    QueryRange,
+)
 from vllm_omni.diffusion.attention.backends.flash_attn import FlashAttentionBackend
 from vllm_omni.diffusion.attention.layer import Attention
 from vllm_omni.diffusion.config import (
@@ -477,6 +483,7 @@ class TestAttentionInitUsesCurrentDiffusionConfig:
     @pytest.mark.parametrize("mindiesd_available", [False, True])
     def test_paged_npu_memory_profile_uses_sdpa_only_without_mindiesd(self, monkeypatch, mindiesd_available):
         attention = Attention.__new__(Attention)
+        attention.allow_fp32_fallback = False
         attention._scheduler_paged_kv = True
         attention.paged_kv_cache_role = "primary"
         attention.backend_pref = "FLASH_ATTN"
@@ -519,6 +526,7 @@ class TestAttentionInitUsesCurrentDiffusionConfig:
     @pytest.mark.parametrize("fa4_available", [False, True])
     def test_paged_cuda_memory_profile_uses_sdpa_without_dense_fa4(self, monkeypatch, fa4_available):
         attention = Attention.__new__(Attention)
+        attention.allow_fp32_fallback = False
         attention._scheduler_paged_kv = True
         attention.paged_kv_cache_role = "primary"
         attention.backend_pref = "FLASH_ATTN"
@@ -561,6 +569,7 @@ class TestAttentionInitUsesCurrentDiffusionConfig:
 
     def test_dense_flash_import_error_is_not_reclassified(self, monkeypatch):
         attention = Attention.__new__(Attention)
+        attention.allow_fp32_fallback = False
         attention._scheduler_paged_kv = False
         attention.paged_kv_cache_role = "primary"
         attention.backend_pref = "FLASH_ATTN"
@@ -900,6 +909,7 @@ class TestAttentionInitUsesCurrentDiffusionConfig:
             return sentinel
 
         fake_attention = SimpleNamespace(
+            allow_fp32_fallback=False,
             attention=SimpleNamespace(
                 forward=_selected_forward,
             ),
@@ -938,6 +948,164 @@ class TestAttentionInitUsesCurrentDiffusionConfig:
             Attention._run_ring_attention(fake_attention, query, query, query, metadata)
 
 
+class TestOptInFloat32Fallback:
+    @staticmethod
+    def _make_attention(monkeypatch, *, backend=None, allow=True, kv_cache_dtype=None):
+        events, state = [], {}
+        context = object()
+
+        class _SelectedImpl:
+            def __init__(self, **kwargs):
+                self.kwargs = kwargs
+
+            @staticmethod
+            def supports_kv_cache_dtype(*args):
+                return True
+
+            def forward(self, query, key, value, metadata):
+                events.append("selected")
+                state["kernel_call"] = (query, key, value, metadata)
+                state["kernel_output"] = query + 4
+                return state["kernel_output"]
+
+        class _FallbackImpl(_SelectedImpl):
+            def forward(self, query, key, value, metadata):
+                output = super().forward(query, key, value, metadata)
+                events[-1] = "fallback"
+                return output
+
+        class _FlashBackend(FlashAttentionBackend):
+            @staticmethod
+            def get_impl_cls():
+                return _SelectedImpl
+
+        class _SDPABackend(layer_mod.SDPABackend):
+            @staticmethod
+            def get_impl_cls():
+                return _SelectedImpl
+
+        class _Strategy:
+            def pre_attention(self, query, key, value, metadata):
+                events.append("pre")
+                state["prepared"] = (
+                    query + 1,
+                    key + 2,
+                    value + 3,
+                    replace(metadata) if metadata is not None else None,
+                )
+                return (*state["prepared"], context)
+
+            def post_attention(self, output, ctx):
+                events.append("post")
+                assert output is state["kernel_output"]
+                assert ctx is context
+                return output + 7
+
+        # Keep real role/config resolution and construction; replace hardware
+        # lookup and kernels. These are CPU dispatch tests, not SP collectives.
+        monkeypatch.setattr(
+            selector_mod,
+            "_cached_get_backend_cls",
+            lambda name, *args, **kwargs: _SDPABackend if name == "TORCH_SDPA" else _FlashBackend,
+        )
+        monkeypatch.setattr(layer_mod.SDPABackend, "get_impl_cls", staticmethod(lambda: _FallbackImpl))
+        monkeypatch.setattr(layer_mod, "build_parallel_attention_strategy", lambda **kwargs: _Strategy())
+        monkeypatch.setattr(layer_mod, "is_forward_context_available", lambda: False)
+        config = OmniDiffusionConfig(
+            diffusion_attention_config={"default": backend} if backend else {},
+            diffusion_kv_cache_dtype=kv_cache_dtype,
+        )
+        with set_current_diffusion_config(config):
+            attention = Attention(
+                num_heads=4,
+                num_kv_heads=2,
+                head_size=8,
+                causal=False,
+                softmax_scale=0.25,
+                **({"allow_fp32_fallback": True} if allow else {}),
+            )
+        return attention, events, state
+
+    @pytest.mark.parametrize(
+        "backend,allow,dtype,is_cuda,expected",
+        [
+            pytest.param(None, True, torch.float32, True, "fallback", id="auto-fp32"),
+            pytest.param(None, False, torch.float32, True, "selected", id="default-opt-out"),
+            pytest.param("FLASH_ATTN", True, torch.float32, True, "selected", id="explicit-flash"),
+            pytest.param("TORCH_SDPA", True, torch.float32, True, "selected", id="explicit-sdpa"),
+            pytest.param(None, True, torch.bfloat16, True, "selected", id="bf16"),
+            pytest.param(None, True, torch.float32, False, "selected", id="cpu-fp32"),
+        ],
+    )
+    def test_public_forward_preserves_strategy_and_backend_choice(
+        self, monkeypatch, backend, allow, dtype, is_cuda, expected
+    ):
+        attention, events, state = self._make_attention(monkeypatch, backend=backend, allow=allow)
+        # Only the dispatch predicate is simulated; every tensor stays on CPU.
+        monkeypatch.setattr(torch.Tensor, "is_cuda", property(lambda self: is_cuda))
+        query = torch.ones(1, 2, 4, 8, dtype=dtype)
+        key = torch.full((1, 2, 2, 8), 2, dtype=dtype)
+        value = torch.full_like(key, 3)
+        metadata = AttentionMetadata(attn_mask=torch.tensor([[True, False]]))
+
+        output = attention(query, key, value, metadata)
+
+        assert events == ["pre", expected, "post"]
+        assert all(got is prepared for got, prepared in zip(state["kernel_call"], state["prepared"]))
+        assert state["kernel_call"][3] is not metadata
+        assert state["kernel_call"][3].attn_mask is metadata.attn_mask
+        torch.testing.assert_close(output, query + 12)
+        assert attention.backend_explicit is (backend is not None)
+        impl = attention.sdpa_fallback if expected == "fallback" else attention.attention
+        assert impl.kwargs["num_kv_heads"] == 2
+        assert impl.kwargs["softmax_scale"] == 0.25
+
+    def test_explicit_backend_failure_is_not_retried_with_sdpa(self, monkeypatch):
+        attention, events, _ = self._make_attention(monkeypatch, backend="FLASH_ATTN")
+        monkeypatch.setattr(torch.Tensor, "is_cuda", property(lambda self: True))
+        selected = Mock(side_effect=RuntimeError("selected backend rejects this dtype"))
+        monkeypatch.setattr(attention.attention, "forward", selected)
+        query = torch.ones(1, 2, 4, 8)
+
+        with pytest.raises(RuntimeError, match="selected backend rejects this dtype"):
+            attention(query, query, query)
+
+        selected.assert_called_once()
+        assert events == ["pre"]
+
+    @pytest.mark.parametrize("contract", ["packed", "packed-padding", "piecewise", "quantized-kv"])
+    def test_backend_metadata_is_not_silently_sent_to_dense_sdpa(self, monkeypatch, contract):
+        quant_dtype = "fp8" if contract == "quantized-kv" else None
+        attention, events, state = self._make_attention(monkeypatch, kv_cache_dtype=quant_dtype)
+        monkeypatch.setattr(torch.Tensor, "is_cuda", property(lambda self: True))
+        lengths = torch.tensor([0, 1, 2], dtype=torch.int32)
+        metadata = AttentionMetadata()
+        if contract == "packed":
+            metadata.extra = {
+                "cu_seqlens_q": lengths,
+                "cu_seqlens_k": lengths,
+                "max_seqlen_q": 1,
+                "max_seqlen_k": 1,
+            }
+        elif contract == "packed-padding":
+            metadata.packed_padding = PackedPaddingMetadata(1, 1, lengths[:2], lengths[:2])
+        elif contract == "piecewise":
+            metadata.full_attn_spans = [[(0, 1)]]
+            metadata.query_ranges = (QueryRange(0, 2, 0),)
+        query = torch.ones(1, 2, 4, 8)
+
+        attention(query, query, query, metadata)
+
+        assert events == ["pre", "selected", "post"]
+        forwarded = state["kernel_call"][3]
+        if quant_dtype:
+            # Quantization is injected by the public forward from config.
+            assert forwarded.extra == {"kv_cache_dtype": "fp8"}
+            assert metadata.extra == {}
+        else:
+            assert forwarded is state["prepared"][3]
+
+
 class TestDiffusionKvCacheQuantization:
     @staticmethod
     def _install_attention_init_stubs(monkeypatch, *, supports_kv_cache_dtype=True):

tests/diffusion/models/ltx2/test_ltx2_transformer.py

@@ -233,6 +233,7 @@ def _fake_ltx_parallel_attention():
 
     layer = object.__new__(ltx2_transformer._LTX2ParallelAttention)
     nn.Module.__init__(layer)
+    layer.allow_fp32_fallback = False
     layer.attn_backend = Backend
     layer.attention = SimpleNamespace(forward=run("native"))
     layer.sdpa_fallback = SimpleNamespace(forward=run("sdpa"))

tests/diffusion/models/mammoth_moda2/test_dit_attention.py

@@ -0,0 +1,189 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""MammothModa2 DiT self-attention runs through the shared Omni attention layer.
+
+``_reference_attention`` is the arithmetic the processor used before it was
+wired to the shared layer: SDPA over K/V replicated to the query head count,
+padded rows zeroed. Both sides read the same module weights, so any drift in
+the projections, the QK RMSNorm, the interleaved real RoPE, the softmax scale
+or the padding mask shows up as a mismatch here.
+"""
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from vllm_omni.diffusion.attention.backends.registry import DiffusionAttentionBackendEnum
+from vllm_omni.diffusion.config import set_current_diffusion_config
+from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.diffusion.models.mammoth_moda2.mammothmoda2_dit_model import (
+    TransformerBlock,
+    apply_real_rotary_emb,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+# Small stand-in for the real 2520 / 21 / 7 (head_dim 120) layout.
+DIM, HEADS = 48, 6
+SEQ, BATCH = 37, 2
+
+# The shared layer resolves its backend from the diffusion config at construction;
+# without one it takes the platform default, which on a CUDA host is FLASH_ATTN
+# and cannot run CPU tensors. Pin SDPA the way a deployment would.
+_SDPA_CONFIG = OmniDiffusionConfig(diffusion_attention_config={"default": {"backend": "TORCH_SDPA"}})
+
+
+def _block(kv_heads: int, modulation: bool = True) -> TransformerBlock:
+    torch.manual_seed(0)
+    with set_current_diffusion_config(_SDPA_CONFIG):
+        block = TransformerBlock(
+            DIM,
+            HEADS,
+            kv_heads,
+            multiple_of=8,
+            ffn_dim_multiplier=1.0,
+            norm_eps=1e-5,
+            modulation=modulation,
+        )
+    return block.eval()
+
+
+def _inputs(head_dim: int):
+    torch.manual_seed(1)
+    hidden = torch.randn(BATCH, SEQ, DIM)
+    mask = torch.ones(BATCH, SEQ, dtype=torch.bool)
+    mask[0, SEQ - 9 :] = False
+    mask[1, SEQ - 2 :] = False
+    angles = torch.rand(1, SEQ, head_dim) * 6.283
+    rotary = (angles.cos(), angles.sin())
+    temb = torch.randn(BATCH, min(DIM, 1024))
+    return hidden, mask, rotary, temb
+
+
+def _reference_attention(attn, hidden, mask, rotary):
+    """The pre-change processor, kept verbatim as the oracle."""
+    batch, seq, _ = hidden.shape
+    query, key, value = attn.to_q(hidden), attn.to_k(hidden), attn.to_v(hidden)
+    head_dim = query.shape[-1] // attn.heads
+    kv_heads = key.shape[-1] // head_dim
+    query = attn.norm_q(query.view(batch, seq, attn.heads, head_dim))
+    key = attn.norm_k(key.view(batch, seq, kv_heads, head_dim))
+    value = value.view(batch, seq, kv_heads, head_dim)
+    query = apply_real_rotary_emb(query, rotary[0], rotary[1])
+    key = apply_real_rotary_emb(key, rotary[0], rotary[1])
+    query, key, value = query.transpose(1, 2), key.transpose(1, 2), value.transpose(1, 2)
+    if kv_heads < attn.heads:
+        key = key.repeat_interleave(attn.heads // kv_heads, dim=1)
+        value = value.repeat_interleave(attn.heads // kv_heads, dim=1)
+    out = F.scaled_dot_product_attention(query, key, value, attn_mask=mask.view(batch, 1, 1, seq), scale=attn.scale)
+    out = (out * mask[:, None, :, None]).transpose(1, 2).reshape(batch, seq, attn.heads * head_dim)
+    return attn.to_out[1](attn.to_out[0](out))
+
+
+@pytest.mark.parametrize("kv_heads", [2, HEADS], ids=["gqa_3to1", "mha"])
+def test_processor_matches_previous_arithmetic(kv_heads):
+    block = _block(kv_heads)
+    hidden, mask, rotary, _ = _inputs(block.head_dim)
+    with torch.no_grad():
+        got = block.attn(
+            hidden_states=hidden, encoder_hidden_states=hidden, attention_mask=mask, image_rotary_emb=rotary
+        )
+        want = _reference_attention(block.attn, hidden, mask, rotary)
+    torch.testing.assert_close(got, want, atol=1e-5, rtol=1e-5)
+    assert torch.count_nonzero(got[~mask]) == 0, "padded rows must stay zero"
+
+
+def test_no_mask_equals_all_valid_mask():
+    block = _block(2)
+    hidden, _, rotary, _ = _inputs(block.head_dim)
+    full = torch.ones(BATCH, SEQ, dtype=torch.bool)
+    with torch.no_grad():
+        unmasked = block.attn(
+            hidden_states=hidden, encoder_hidden_states=hidden, attention_mask=None, image_rotary_emb=rotary
+        )
+        masked = block.attn(
+            hidden_states=hidden, encoder_hidden_states=hidden, attention_mask=full, image_rotary_emb=rotary
+        )
+    torch.testing.assert_close(unmasked, masked, atol=1e-6, rtol=1e-6)
+
+
+@pytest.mark.parametrize("modulation", [True, False], ids=["noise_block", "context_refiner"])
+def test_block_forward_hands_native_kv_heads_to_shared_layer(monkeypatch, modulation):
+    """Through the block's real forward, on both of its branches (the modulated
+    noise/refiner blocks and the unmodulated context refiner): the shared layer
+    is what runs, and it receives K/V with their native head count -- no
+    replication before the call."""
+    block = _block(2, modulation=modulation)
+    hidden, mask, rotary, temb = _inputs(block.head_dim)
+    if not modulation:
+        temb = None
+    seen = []
+    original = block.attn.omni_attn.forward
+
+    def spy(query, key, value, attn_metadata=None):
+        seen.append((tuple(query.shape), tuple(key.shape), tuple(value.shape), attn_metadata))
+        return original(query, key, value, attn_metadata)
+
+    monkeypatch.setattr(block.attn.omni_attn, "forward", spy)
+    with torch.no_grad():
+        out = block(hidden, mask, rotary, temb)
+    assert out.shape == hidden.shape and torch.isfinite(out).all()
+    assert len(seen) == 1
+    q_shape, k_shape, v_shape, metadata = seen[0]
+    assert q_shape == (BATCH, SEQ, HEADS, block.head_dim)
+    assert k_shape == v_shape == (BATCH, SEQ, 2, block.head_dim)
+    assert metadata is not None and metadata.attn_mask.dtype == torch.bool
+    assert torch.equal(metadata.attn_mask, mask)
+
+
+def test_backend_selection_reaches_the_dit_layer():
+    """The diffusion config selects the backend when the DiT layer is constructed."""
+    assert _block(2).attn.omni_attn.attn_backend is DiffusionAttentionBackendEnum.TORCH_SDPA.get_class()
+    with set_current_diffusion_config(
+        OmniDiffusionConfig(diffusion_attention_config={"default": {"backend": "FLASH_ATTN"}})
+    ):
+        block = TransformerBlock(DIM, HEADS, 2, multiple_of=8, ffn_dim_multiplier=1.0, norm_eps=1e-5)
+    assert block.attn.omni_attn.attn_backend is DiffusionAttentionBackendEnum.FLASH_ATTN.get_class()
+
+
+def test_shared_layer_owns_no_parameters_and_keeps_checkpoint_keys():
+    block = _block(2)
+    assert sum(p.numel() for p in block.attn.omni_attn.parameters()) == 0
+    assert not list(block.attn.omni_attn.buffers())
+    keys = set(block.state_dict())
+    assert not any("omni_attn" in key for key in keys)
+    # The projections and QK norms the checkpoint provides are still where they were.
+    for expected in (
+        "attn.to_q.weight",
+        "attn.to_k.weight",
+        "attn.to_v.weight",
+        "attn.to_out.0.weight",
+        "attn.norm_q.weight",
+        "attn.norm_k.weight",
+    ):
+        assert expected in keys, expected
+
+
+@pytest.mark.parametrize("with_mask", [True, False], ids=["empty_mask", "no_mask"])
+def test_empty_text_stream_skips_the_kernel(monkeypatch, with_mask):
+    """CFG's unconditional branch carries zero text tokens (the pipeline's default
+    negative_prompt_embeds has no rows), so the context refiner attends over an
+    empty sequence. The flash-attention varlen fallback cannot take that
+    (arange step 0), so the processor must not reach the shared layer at all."""
+    block = _block(kv_heads=2, modulation=False)
+    attn = block.attn
+
+    def _never(*args, **kwargs):
+        raise AssertionError("shared layer called with an empty sequence")
+
+    monkeypatch.setattr(attn.omni_attn, "forward", _never)
+    hidden = torch.randn(BATCH, 0, DIM)
+    mask = torch.ones(BATCH, 0, dtype=torch.bool) if with_mask else None
+    angles = torch.rand(1, 0, block.head_dim)
+    out = attn(
+        hidden_states=hidden,
+        encoder_hidden_states=hidden,
+        attention_mask=mask,
+        image_rotary_emb=(angles.cos(), angles.sin()),
+    )
+    assert out.shape == (BATCH, 0, DIM)

tests/diffusion/models/mammoth_moda2/test_dit_attention_cuda.py

@@ -0,0 +1,151 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Real-shape check of the MammothModa2 DiT attention on CUDA: the shared
+backend (FLASH_ATTN by default, head_dim 120) against the pre-change arithmetic
+in bf16, with the padding mask the joint text+image sequence carries."""
+
+import pytest
+import torch
+
+import vllm_omni.diffusion.attention.backends.sdpa as sdpa_backend
+from tests.helpers.mark import hardware_marks
+from vllm_omni.diffusion.models.mammoth_moda2.mammothmoda2_dit_model import TransformerBlock
+
+from .test_dit_attention import _reference_attention
+
+pytestmark = [
+    pytest.mark.core_model,
+    pytest.mark.advanced_model,
+    pytest.mark.cuda,
+    *hardware_marks(res={"cuda": "L4"}, num_cards=1),
+]
+
+# MammothModa2-Preview gen_dit_config: hidden 2520, 21 heads, 7 KV heads, head_dim 120.
+DIM, HEADS, KV_HEADS = 2520, 21, 7
+
+
+@pytest.mark.parametrize("seq", [77 + 4096, 512], ids=["t2i_1024", "short"])
+def test_real_shape_matches_previous_arithmetic_bf16(seq):
+    torch.manual_seed(0)
+    block = (
+        TransformerBlock(DIM, HEADS, KV_HEADS, multiple_of=256, ffn_dim_multiplier=1.0, norm_eps=1e-5)
+        .cuda()
+        .to(torch.bfloat16)
+        .eval()
+    )
+    hidden = torch.randn(2, seq, DIM, device="cuda", dtype=torch.bfloat16)
+    mask = torch.ones(2, seq, dtype=torch.bool, device="cuda")
+    mask[0, seq - 300 :] = False
+    mask[1, seq - 17 :] = False
+    angles = torch.rand(1, seq, block.head_dim, device="cuda") * 6.283
+    rotary = (angles.cos().to(torch.bfloat16), angles.sin().to(torch.bfloat16))
+    with torch.no_grad():
+        got = block.attn(
+            hidden_states=hidden, encoder_hidden_states=hidden, attention_mask=mask, image_rotary_emb=rotary
+        )
+        want = _reference_attention(block.attn, hidden, mask, rotary)
+    assert torch.isfinite(got).all()
+    assert torch.count_nonzero(got[~mask]) == 0
+    diff = (got.float() - want.float()).abs()[mask]
+    # bf16 kernels differ in accumulation order; measured on A800: max 1.0e-3, mean 5e-5.
+    assert diff.max().item() < 2e-2, diff.max().item()
+    assert diff.mean().item() < 1e-3, diff.mean().item()
+
+
+def test_empty_text_stream_on_the_default_backend():
+    """The recipe's text-to-image request with text_guidance_scale > 1 runs the
+    context refiner on a zero-token unconditional prompt. Before the guard this
+    raised ``RuntimeError: step must be nonzero`` from the FA varlen fallback."""
+    torch.manual_seed(0)
+    block = (
+        TransformerBlock(DIM, HEADS, KV_HEADS, multiple_of=256, ffn_dim_multiplier=1.0, norm_eps=1e-5, modulation=False)
+        .cuda()
+        .to(torch.bfloat16)
+        .eval()
+    )
+    hidden = torch.randn(1, 0, DIM, device="cuda", dtype=torch.bfloat16)
+    mask = torch.ones(1, 0, dtype=torch.bool, device="cuda")
+    angles = torch.rand(1, 0, block.head_dim, device="cuda")
+    with torch.no_grad():
+        out = block.attn(
+            hidden_states=hidden,
+            encoder_hidden_states=hidden,
+            attention_mask=mask,
+            image_rotary_emb=(angles.cos().to(torch.bfloat16), angles.sin().to(torch.bfloat16)),
+        )
+    assert out.shape == (1, 0, DIM)
+
+
+@pytest.mark.parametrize(
+    ("seq", "force_gqa_fallback"),
+    [(512, False), (77 + 4096, False), (32, True)],
+    ids=["short", "t2i_1024", "unsupported_native_gqa"],
+)
+def test_fp32_falls_back_to_sdpa_and_matches_reference(monkeypatch, seq, force_gqa_fallback):
+    """FP32 keeps its dtype and reaches shared SDPA's GQA compatibility check.
+
+    Cover the long masked shape that otherwise falls back to quadratic math
+    attention, and force the unsupported-GQA branch independently of hardware.
+    """
+    torch.manual_seed(0)
+    dt = torch.float32
+    block = (
+        TransformerBlock(DIM, HEADS, KV_HEADS, multiple_of=256, ffn_dim_multiplier=1.0, norm_eps=1e-5)
+        .cuda()
+        .to(dt)
+        .eval()
+    )
+    hidden = torch.randn(1, seq, DIM, device="cuda", dtype=dt)
+    mask = torch.ones(1, seq, dtype=torch.bool, device="cuda")
+    mask[:, -13:] = False
+    angles = torch.rand(1, seq, block.head_dim, device="cuda")
+    rotary = (angles.cos().to(dt), angles.sin().to(dt))
+    capabilities, calls, phases = [], [], []
+    can_use_fused_gqa = sdpa_backend.can_sdpa_use_fused_gqa
+    sdpa = torch.nn.functional.scaled_dot_product_attention
+    strategy = block.attn.omni_attn._get_active_parallel_strategy()
+    pre_attention, post_attention = strategy.pre_attention, strategy.post_attention
+
+    def prepare(*args):
+        phases.append("pre")
+        return pre_attention(*args)
+
+    def restore(*args):
+        phases.append("post")
+        return post_attention(*args)
+
+    def check_gqa(*args):
+        supported = False if force_gqa_fallback else can_use_fused_gqa(*args)
+        capabilities.append(supported)
+        return supported
+
+    def record_sdpa(query, key, value, **kwargs):
+        phases.append("sdpa")
+        calls.append((query.shape, key.shape, value.shape, query.dtype, kwargs))
+        return sdpa(query, key, value, **kwargs)
+
+    with torch.no_grad(), monkeypatch.context() as patch:
+        patch.setattr(sdpa_backend, "can_sdpa_use_fused_gqa", check_gqa)
+        patch.setattr(torch.nn.functional, "scaled_dot_product_attention", record_sdpa)
+        patch.setattr(strategy, "pre_attention", prepare)
+        patch.setattr(strategy, "post_attention", restore)
+        got = block.attn(hidden, hidden, attention_mask=mask, image_rotary_emb=rotary)
+
+    assert phases == ["pre", "sdpa", "post"], "FP32 must preserve the shared parallel dispatch"
+    assert len(capabilities) == len(calls) == 1, "FP32 must use shared SDPA's runtime GQA check"
+    q_shape, k_shape, v_shape, dtype, kwargs = calls[0]
+    expected_kv_heads = KV_HEADS if capabilities[0] else HEADS
+    assert q_shape == (1, HEADS, seq, block.head_dim)
+    assert k_shape == v_shape == (1, expected_kv_heads, seq, block.head_dim)
+    assert dtype == got.dtype == torch.float32
+    assert kwargs["enable_gqa"] == capabilities[0]
+    assert kwargs["scale"] == block.attn.scale
+    assert kwargs["is_causal"] is False
+    assert torch.equal(kwargs["attn_mask"], mask[:, None, None, :])
+
+    with torch.no_grad():
+        want = _reference_attention(block.attn, hidden, mask, rotary)
+    assert torch.isfinite(got).all()
+    assert torch.count_nonzero(got[~mask]) == 0
+    diff = (got.float() - want.float()).abs()[mask]
+    assert diff.max().item() < 1e-3, diff.max().item()

vllm_omni/diffusion/attention/layer.py

@@ -81,6 +81,8 @@ class Attention(nn.Module):
         # varlen attention with learned sink logits). The shared Attention
         # layer still owns parallel dispatch and compile boundaries.
         custom_attention: nn.Module | None = None,
+        # Preserve dense FP32 inference for models opting into CUDA auto fallback.
+        allow_fp32_fallback: bool = False,
     ):
         super().__init__()
 
@@ -193,8 +195,8 @@ class Attention(nn.Module):
                 role=role,
                 backend_explicit=self.backend_explicit,
             )
-            # Some model-specific paths call this helper directly. The main backend
-            # dispatch below never switches to it implicitly.
+            # Compatibility kernels run inside shared dispatch, between the
+            # parallel strategy's input preparation and output restoration.
             self.sdpa_fallback = SDPABackend.get_impl_cls()(
                 num_heads=num_heads,
                 head_size=head_size,
@@ -221,6 +223,7 @@ class Attention(nn.Module):
         self.use_sync = use_sync
         self.causal = causal
         self.skip_sequence_parallel = skip_sequence_parallel
+        self.allow_fp32_fallback = allow_fp32_fallback
 
         self.use_ring = False
         self.ring_pg = None
@@ -497,13 +500,33 @@ class Attention(nn.Module):
 
         self._assert_metadata_compatible(attn_metadata)
 
+        if (
+            self.allow_fp32_fallback
+            and query.is_cuda
+            and query.dtype == torch.float32
+            and query.ndim == 4
+            and not self.backend_explicit
+            and self.attn_backend.get_name() == "FLASH_ATTN"
+            and (
+                attn_metadata is None
+                or (
+                    attn_metadata.full_attn_spans is None
+                    and attn_metadata.query_ranges is None
+                    and attn_metadata.video_layout is None
+                    and attn_metadata.packed_padding is None
+                    and not attn_metadata.extra
+                )
+            )
+        ):
+            logger.warning_once("Using SDPA for this layer's FP32 input with automatic CUDA FlashAttention selection.")
+            return self.sdpa_fallback.forward(query, key, value, attn_metadata)
+
         in_kv_memory_profile = is_forward_context_available() and get_forward_context().in_diffusion_kv_memory_profile
         # The startup KV-capacity profile needs tensor shapes, not a paged
         # attention result. If dense FLASH_ATTN deps are absent (NPU MindIE-SD
         # or CUDA CuTe FA4), SDPA provides that profile forward. Formal paged
         # requests never use this branch because their Worker adapter is active.
-        # Do not silently remap float32 (or other dtypes) to SDPA: a
-        # user-explicit backend must run or raise.
+        # A user-explicit backend must run or raise.
         if (
             self._scheduler_paged_kv
             and self.paged_kv_cache_role is not None

vllm_omni/diffusion/models/mammoth_moda2/mammothmoda2_dit_model.py

@@ -1,4 +1,5 @@
-import math
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 import torch
 import torch.nn.functional as F
@@ -6,27 +7,14 @@ from diffusers.configuration_utils import ConfigMixin, register_to_config
 from diffusers.models.attention_processor import Attention
 from diffusers.models.embeddings import TimestepEmbedding, Timesteps
 from diffusers.models.modeling_utils import ModelMixin
-from einops import rearrange, repeat
+from einops import rearrange
 from torch import nn
 from transformers.models.qwen2.modeling_qwen2 import Qwen2RMSNorm
 
-from .rope_real import RotaryPosEmbedReal
+from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
+from vllm_omni.diffusion.attention.layer import Attention as OmniAttention
 
-try:
-    from transformers.modeling_flash_attention_utils import (  # type: ignore
-        flash_attn_varlen_func,  # pyright: ignore[reportAttributeAccessIssue]
-        is_flash_attn_available,
-    )
-except Exception:  # pragma: no cover - best-effort compatibility
-    flash_attn_varlen_func = None  # type: ignore[assignment]
-
-    def is_flash_attn_available() -> bool:  # type: ignore[override]
-        return False
-
-
-from .rope_real import apply_real_rotary_emb
-
-_HAS_FLASH_ATTN_VARLEN = bool(is_flash_attn_available()) and flash_attn_varlen_func is not None
+from .rope_real import RotaryPosEmbedReal, apply_real_rotary_emb
 
 
 class LuminaRMSNormZero(nn.Module):
@@ -278,9 +266,14 @@ class SimpleQFormerImageRefiner(nn.Module):
 
 
 class AttnProcessor:
-    def __init__(self) -> None:
-        if not hasattr(F, "scaled_dot_product_attention"):
-            raise ImportError("AttnProcessor requires PyTorch 2.0+ (F.scaled_dot_product_attention).")
+    """Self-attention for the MammothModa2 DiT through the shared Omni attention layer.
+
+    The model-specific parts stay here: the projections, the per-head QK RMSNorm
+    and the interleaved real RoPE. The kernel, the padding mask and the
+    grouped-query heads are the shared backend's job, so K/V are handed over
+    with their native head count instead of being replicated to the query
+    head count first.
+    """
 
     def __call__(
         self,
@@ -289,118 +282,55 @@ class AttnProcessor:
         encoder_hidden_states: torch.Tensor,
         attention_mask: torch.Tensor | None = None,
         image_rotary_emb: torch.Tensor | None = None,
-        base_sequence_length: int | None = None,
     ) -> torch.Tensor:
         batch_size, sequence_length, _ = hidden_states.shape
 
-        # Get Query-Key-Value Pair
         query = attn.to_q(hidden_states)
         key = attn.to_k(encoder_hidden_states)
         value = attn.to_v(encoder_hidden_states)
 
-        query_dim = query.shape[-1]
-        inner_dim = key.shape[-1]
-        head_dim = query_dim // attn.heads
+        head_dim = query.shape[-1] // attn.heads
+        kv_heads = key.shape[-1] // head_dim
         dtype = query.dtype
 
-        # Get key-value heads
-        kv_heads = inner_dim // head_dim
+        if sequence_length == 0 or key.shape[1] == 0:
+            # The pipeline's default unconditional branch is a text stream with
+            # zero tokens (negative_prompt_embeds has no rows), so under CFG the
+            # context refiner attends over an empty sequence. SDPA returns an
+            # empty tensor for that; the flash-attention varlen fallback builds
+            # cu_seqlens with arange(step=seq_len) and rejects step 0. Nothing
+            # to attend to either way, so hand the projection an empty input.
+            empty = query.new_zeros(batch_size, sequence_length, attn.heads * head_dim)
+            return attn.to_out[1](attn.to_out[0](empty))
 
-        # Reshape tensors for attention computation
         query = query.view(batch_size, -1, attn.heads, head_dim)
         key = key.view(batch_size, -1, kv_heads, head_dim)
         value = value.view(batch_size, -1, kv_heads, head_dim)
 
-        # Apply Query-Key normalization
         if attn.norm_q is not None:
             query = attn.norm_q(query)
         if attn.norm_k is not None:
             key = attn.norm_k(key)
 
-        # Apply Rotary Position Embeddings
         if image_rotary_emb is not None:
             query = apply_real_rotary_emb(query, image_rotary_emb[0], image_rotary_emb[1])
             key = apply_real_rotary_emb(key, image_rotary_emb[0], image_rotary_emb[1])
 
         query, key = query.to(dtype), key.to(dtype)
 
-        # Calculate attention scale
-        if base_sequence_length is not None:
-            softmax_scale = math.sqrt(math.log(sequence_length, base_sequence_length)) * attn.scale
-        else:
-            softmax_scale = attn.scale
-
-        if _HAS_FLASH_ATTN_VARLEN and attention_mask is not None and hidden_states.is_cuda:
-            # Flash-Attn varlen expects packed tokens + cu_seqlens. Here we only need
-            # the self-attention case (q/k/v share the same padding mask).
+        if attention_mask is not None:
             attention_mask = attention_mask.to(torch.bool)
-            seqlens = attention_mask.sum(dim=-1, dtype=torch.int32)
-            indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
-            max_seqlen = int(seqlens.max().item())
-            cu_seqlens = F.pad(torch.cumsum(seqlens, dim=0, dtype=torch.int32), (1, 0))
 
-            query_states = query.reshape(batch_size * sequence_length, attn.heads, head_dim)[indices]
-            key_states = key.reshape(batch_size * sequence_length, kv_heads, head_dim)[indices]
-            value_states = value.reshape(batch_size * sequence_length, kv_heads, head_dim)[indices]
+        attn_metadata = AttentionMetadata(attn_mask=attention_mask) if attention_mask is not None else None
+        hidden_states = attn.omni_attn(query, key, value, attn_metadata)
 
-            if kv_heads < attn.heads:
-                key_states = repeat(key_states, "l h c -> l (h k) c", k=attn.heads // kv_heads)
-                value_states = repeat(value_states, "l h c -> l (h k) c", k=attn.heads // kv_heads)
+        if attention_mask is not None:
+            # Padded rows carry nothing downstream; keep them at zero as the
+            # previous varlen path did, since neither backend does it for us.
+            hidden_states = hidden_states * attention_mask[:, :, None, None]
 
-            attn_output_unpad = flash_attn_varlen_func(
-                query_states,
-                key_states,
-                value_states,
-                cu_seqlens_q=cu_seqlens,
-                cu_seqlens_k=cu_seqlens,
-                max_seqlen_q=max_seqlen,
-                max_seqlen_k=max_seqlen,
-                dropout_p=0.0,
-                causal=False,
-                softmax_scale=softmax_scale,
-            )
+        hidden_states = hidden_states.reshape(batch_size, sequence_length, attn.heads * head_dim).type_as(query)
 
-            out = torch.zeros(
-                (batch_size * sequence_length, attn.heads, head_dim),
-                device=attn_output_unpad.device,
-                dtype=attn_output_unpad.dtype,
-            )
-            out[indices] = attn_output_unpad
-            hidden_states = out.view(batch_size, sequence_length, attn.heads, head_dim).flatten(-2)
-            hidden_states = hidden_states.type_as(query)
-        else:
-            # PyTorch SDPA path.
-            attn_mask = None
-            if attention_mask is not None:
-                attention_mask = attention_mask.to(torch.bool)
-                attn_mask = attention_mask.view(batch_size, 1, 1, -1)
-
-            query = query.transpose(1, 2)  # [B, H, S, D]
-            key = key.transpose(1, 2)  # [B, H_kv, S, D]
-            value = value.transpose(1, 2)
-
-            if kv_heads < attn.heads:
-                key = key.repeat_interleave(attn.heads // kv_heads, dim=1)
-                value = value.repeat_interleave(attn.heads // kv_heads, dim=1)
-
-            hidden_states = F.scaled_dot_product_attention(
-                query,
-                key,
-                value,
-                attn_mask=attn_mask,
-                dropout_p=0.0,
-                is_causal=False,
-                scale=softmax_scale,
-            )
-
-            if attention_mask is not None:
-                # Keep padding tokens consistent with the flash-varlen path (zero output).
-                hidden_states = hidden_states * attention_mask[:, None, :, None]
-
-            hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
-            hidden_states = hidden_states.type_as(query)
-
-        # Apply output projection
         hidden_states = attn.to_out[0](hidden_states)
         hidden_states = attn.to_out[1](hidden_states)
 
@@ -441,6 +371,18 @@ class TransformerBlock(nn.Module):
         # 显式使用 transformers 的 Qwen2RMSNorm，避免依赖 diffusers 内部创建的 `RMSNorm` 再做递归替换。
         self.attn.norm_q = Qwen2RMSNorm(self.head_dim, eps=1e-5)
         self.attn.norm_k = Qwen2RMSNorm(self.head_dim, eps=1e-5)
+        # The kernel itself runs through the shared Omni attention layer: backend
+        # selection, padding-mask handling and native grouped-query heads live
+        # there. It owns no parameters, so checkpoint keys are unchanged. The
+        # diffusers Attention above is kept for its projections and QK norms.
+        self.attn.omni_attn = OmniAttention(
+            num_heads=num_attention_heads,
+            head_size=self.head_dim,
+            causal=False,
+            softmax_scale=self.attn.scale,
+            num_kv_heads=num_kv_heads,
+            allow_fp32_fallback=True,
+        )
 
         # Initialize feed-forward network
         self.feed_forward = LuminaFeedForward(