From 45a3e77351b9ef2b36612159bb0d5a3df812466b Mon Sep 17 00:00:00 2001 From: leo <92263211+liuyao0322@users.noreply.github.com> Date: Wed, 16 Sep 2026 09:08:06 +0800 Subject: [PATCH] [Bugfix] Route Music3 float32 attention to SDPA (#7354) Signed-off-by: liuyao0322 --- .../test_minimax_music3_dit_attention.py | 117 ++++++++++++++++++ .../models/minimax_music3/dit.py | 47 ++++--- 2 files changed, 150 insertions(+), 14 deletions(-) create mode 100644 tests/model_executor/models/test_minimax_music3_dit_attention.py diff --git a/tests/model_executor/models/test_minimax_music3_dit_attention.py b/tests/model_executor/models/test_minimax_music3_dit_attention.py new file mode 100644 index 000000000..cdbb57796 --- /dev/null +++ b/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)) diff --git a/vllm_omni/model_executor/models/minimax_music3/dit.py b/vllm_omni/model_executor/models/minimax_music3/dit.py index cf180ac52..1e7af63fc 100644 --- a/vllm_omni/model_executor/models/minimax_music3/dit.py +++ b/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)