[Bugfix] Route Music3 float32 attention to SDPA (#7354)

Signed-off-by: liuyao0322 <yaoliu548926@gmail.com>

tests/model_executor/models/test_minimax_music3_dit_attention.py

@@ -0,0 +1,117 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Attention dtype routing for MiniMax Music 3's float32 acoustic stage."""
+
+from __future__ import annotations
+
+import pytest
+import torch
+from torch import nn
+
+from vllm_omni.model_executor.models.minimax_music3 import dit
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_FLOAT32_INCOMPATIBLE_BACKENDS = [
+    "FLASH_ATTN",
+    "FLASH_ATTN_HUB",
+    "FLASH_ATTN_3_HUB",
+    "CUDNN_ATTN",
+    "FLASHINFER_ATTN",
+]
+
+
+class _BackendType:
+    def __init__(self, name: str) -> None:
+        self.name = name
+
+    def get_name(self) -> str:
+        return self.name
+
+
+class _Backend(nn.Module):
+    def __init__(self, *, name: str, explicit: bool) -> None:
+        super().__init__()
+        self.attn_backend = _BackendType(name)
+        self.backend_explicit = explicit
+        self.calls = 0
+
+    def forward(self, q, k, v):
+        del k, v
+        self.calls += 1
+        return torch.full_like(q, 7)
+
+
+def _attention(monkeypatch, *, backend_name: str, explicit: bool):
+    backend = _Backend(name=backend_name, explicit=explicit)
+    monkeypatch.setattr(dit, "_build_native_attention", lambda **_kwargs: backend)
+    return dit.Attention(8, head_dim=4), backend
+
+
+@pytest.mark.parametrize(
+    "backend_name",
+    _FLOAT32_INCOMPATIBLE_BACKENDS,
+)
+def test_float32_uses_sdpa_for_automatic_incompatible_backend(monkeypatch, backend_name):
+    attention, backend = _attention(monkeypatch, backend_name=backend_name, explicit=False)
+    q = torch.randn(2, 5, 2, 4)
+
+    actual = attention._attend(q, q, q)
+
+    expected = torch.nn.functional.scaled_dot_product_attention(
+        q.transpose(1, 2),
+        q.transpose(1, 2),
+        q.transpose(1, 2),
+        is_causal=False,
+        scale=attention.softmax_scale,
+    ).transpose(1, 2)
+    torch.testing.assert_close(actual, expected)
+    assert backend.calls == 0
+    assert attention.backend_name == f"TORCH_SDPA(float32)/{backend_name}(lower precision)"
+
+
+@pytest.mark.parametrize("backend_name", _FLOAT32_INCOMPATIBLE_BACKENDS)
+@pytest.mark.parametrize(
+    ("dtype", "explicit"),
+    [
+        (torch.float16, False),
+        (torch.bfloat16, False),
+        (torch.float32, True),
+    ],
+)
+def test_lower_precision_or_explicit_backend_remains_authoritative(monkeypatch, backend_name, dtype, explicit):
+    attention, backend = _attention(monkeypatch, backend_name=backend_name, explicit=explicit)
+    q = torch.zeros(1, 3, 2, 4, dtype=dtype)
+
+    actual = attention._attend(q, q, q)
+
+    assert backend.calls == 1
+    torch.testing.assert_close(actual, torch.full_like(q, 7))
+
+
+def test_float32_uses_sdpa_when_native_attention_is_unavailable(monkeypatch):
+    monkeypatch.setattr(dit, "_build_native_attention", lambda **_kwargs: None)
+    attention = dit.Attention(8, head_dim=4)
+    q = torch.randn(1, 3, 2, 4)
+
+    actual = attention._attend(q, q, q)
+
+    expected = torch.nn.functional.scaled_dot_product_attention(
+        q.transpose(1, 2),
+        q.transpose(1, 2),
+        q.transpose(1, 2),
+        is_causal=False,
+        scale=attention.softmax_scale,
+    ).transpose(1, 2)
+    torch.testing.assert_close(actual, expected)
+    assert attention.backend_name == "TORCH_SDPA"
+
+
+def test_float32_keeps_compatible_automatic_backend(monkeypatch):
+    attention, backend = _attention(monkeypatch, backend_name="SDPA", explicit=False)
+    q = torch.zeros(1, 3, 2, 4)
+
+    actual = attention._attend(q, q, q)
+
+    assert backend.calls == 1
+    torch.testing.assert_close(actual, torch.full_like(q, 7))

vllm_omni/model_executor/models/minimax_music3/dit.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Condition encoder and flow-matching transformer for MiniMax Music 3.
 
 Stage 1 turns the AR stage's conditioning frames into a DAC latent in two
@@ -16,13 +16,13 @@ a single ``to_qkv`` so each block runs one GEMM instead of three.
 :func:`remap_transformer_state` performs that fusion, and is the only place
 the checkpoint's layout is reinterpreted.
 
-Attention runs through the project's own diffusion attention layer, so this
-model picks up the same backend selection, sequence-parallel plumbing and
-KV-quant policy as every other transformer in the tree. That layer resolves
-its backend from the diffusion config, which an ``LLM_GENERATION`` stage such
-as this one does not set; it tolerates that and falls through to the platform
-default. If it cannot be built at all the block runs plain SDPA instead, a
-choice made once when the blocks are built and never inside ``forward``.
+Attention uses the project's diffusion attention layer when its selected
+backend supports the runtime dtype. This ``LLM_GENERATION`` stage does not set
+a diffusion backend and therefore gets the platform default. CUDA commonly
+resolves that default to FlashAttention, cuDNN, or FlashInfer, so the model
+routes its native float32 decode through SDPA while retaining the selected
+backend for supported lower-precision inputs. An explicitly selected backend
+remains authoritative.
 """
 
 from __future__ import annotations
@@ -75,6 +75,19 @@ _TRANSFORMER_IN_DIM = DIT_LATENT_CHANNELS * 2 + _CONDITION_DIM
 # reaches exactly zero, so the noise term never fully vanishes.
 _SIGMA_FLOOR = 1e-6
 
+# These backends reject float32 Q/K/V. MiniMax Music 3 decodes in float32, so
+# automatically selected incompatible backends must not receive its attention
+# tensors. This includes cuDNN and FlashInfer selected on Blackwell. Explicit
+# backend choices retain the shared layer's fail-fast contract instead of
+# being silently replaced here.
+_FLOAT32_UNSUPPORTED_BACKENDS = {
+    "FLASH_ATTN",
+    "FLASH_ATTN_HUB",
+    "FLASH_ATTN_3_HUB",
+    "CUDNN_ATTN",
+    "FLASHINFER_ATTN",
+}
+
 
 class FourierFeatures(nn.Module):
     """Random-Fourier timestep features: ``[B, 1] -> [B, out_features]``."""
@@ -149,10 +162,7 @@ def _build_native_attention(*, num_heads: int, head_size: int, softmax_scale: fl
     except Exception as exc:  # noqa: BLE001 - any failure means: use SDPA
         _NATIVE_ATTENTION_UNAVAILABLE = repr(exc)
         logger.warning(
-            "MiniMax Music 3 DiT could not build the diffusion attention layer "
-            "(%s); running torch SDPA instead. The stage decodes in float32, "
-            "for which the diffusion layer dispatches to SDPA anyway, so this "
-            "changes performance bookkeeping and not the audio.",
+            "MiniMax Music 3 DiT could not build the diffusion attention layer (%s); running torch SDPA instead.",
             exc,
         )
         return None
@@ -180,10 +190,19 @@ class Attention(nn.Module):
             softmax_scale=self.softmax_scale,
             prefix=prefix,
         )
-        self.backend_name = self.backend.attn_backend.get_name() if self.backend is not None else "TORCH_SDPA"
+        native_backend_name = self.backend.attn_backend.get_name() if self.backend is not None else None
+        self._auto_float32_sdpa = (
+            self.backend is not None
+            and not getattr(self.backend, "backend_explicit", False)
+            and native_backend_name in _FLOAT32_UNSUPPORTED_BACKENDS
+        )
+        if self._auto_float32_sdpa:
+            self.backend_name = f"TORCH_SDPA(float32)/{native_backend_name}(lower precision)"
+        else:
+            self.backend_name = native_backend_name or "TORCH_SDPA"
 
     def _attend(self, q: Tensor, k: Tensor, v: Tensor) -> Tensor:
-        if self.backend is not None:
+        if self.backend is not None and not (q.dtype == torch.float32 and self._auto_float32_sdpa):
             return self.backend(q, k, v)
         q, k, v = (t.transpose(1, 2) for t in (q, k, v))
         out = F.scaled_dot_product_attention(q, k, v, is_causal=False, scale=self.softmax_scale)