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