diff --git a/tests/diffusion/attention/test_attention_config.py b/tests/diffusion/attention/test_attention_config.py index 656161589..4a053b1bf 100644 --- a/tests/diffusion/attention/test_attention_config.py +++ b/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): diff --git a/tests/diffusion/models/ltx2/test_ltx2_transformer.py b/tests/diffusion/models/ltx2/test_ltx2_transformer.py index f12fbc466..679b7ca37 100644 --- a/tests/diffusion/models/ltx2/test_ltx2_transformer.py +++ b/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")) diff --git a/tests/diffusion/models/mammoth_moda2/__init__.py b/tests/diffusion/models/mammoth_moda2/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/diffusion/models/mammoth_moda2/test_dit_attention.py b/tests/diffusion/models/mammoth_moda2/test_dit_attention.py new file mode 100644 index 000000000..bab2fee7d --- /dev/null +++ b/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) diff --git a/tests/diffusion/models/mammoth_moda2/test_dit_attention_cuda.py b/tests/diffusion/models/mammoth_moda2/test_dit_attention_cuda.py new file mode 100644 index 000000000..62415953a --- /dev/null +++ b/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() diff --git a/vllm_omni/diffusion/attention/layer.py b/vllm_omni/diffusion/attention/layer.py index 4b8ca8f49..4c588af0c 100644 --- a/vllm_omni/diffusion/attention/layer.py +++ b/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 diff --git a/vllm_omni/diffusion/models/mammoth_moda2/mammothmoda2_dit_model.py b/vllm_omni/diffusion/models/mammoth_moda2/mammothmoda2_dit_model.py index 000abe051..e2c086d4a 100644 --- a/vllm_omni/diffusion/models/mammoth_moda2/mammothmoda2_dit_model.py +++ b/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(