[Bugfix] Fix GLM-TTS DiT float32 crash with flash-attn backends (#7756)

Signed-off-by: Sworol <171231900+Sworol@users.noreply.github.com>
Co-authored-by: Sworol <171231900+Sworol@users.noreply.github.com>
Co-authored-by: Claude Code <noreply@anthropic.com>

tests/model_executor/models/test_glm_tts_dit_attention.py

@@ -0,0 +1,126 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Attention dtype routing for GLM-TTS's float32 DiT stage."""
+
+from __future__ import annotations
+
+import pytest
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from vllm_omni.model_executor.models.glm_tts.glm_tts_dit import DiTAttention
+
+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 _Spec:
+    def __init__(self) -> None:
+        self.name = "FLASH_ATTN"
+        self.explicit = False
+
+
+_SPEC = _Spec()
+
+
+class _FakeDiffusionAttention(nn.Module):
+    """Stands in for the project diffusion attention layer, recording calls."""
+
+    def __init__(self, **_kwargs) -> None:
+        super().__init__()
+        self.attn_backend = _BackendType(_SPEC.name)
+        self.backend_explicit = _SPEC.explicit
+        self.calls = 0
+
+    def forward(self, q, k, v, attn_metadata=None):
+        del k, v, attn_metadata
+        self.calls += 1
+        return torch.full_like(q, 7)
+
+
+@pytest.fixture
+def fake_attention(monkeypatch):
+    monkeypatch.setattr(
+        "vllm_omni.model_executor.models.glm_tts.glm_tts_dit.DiffusionAttention", _FakeDiffusionAttention
+    )
+
+
+def _build_attention(*, backend_name: str, explicit: bool) -> DiTAttention:
+    _SPEC.name = backend_name
+    _SPEC.explicit = explicit
+    attention = DiTAttention(dim=16, heads=2, dim_head=8, dropout=0.0).eval()
+    assert isinstance(attention.attn, _FakeDiffusionAttention)
+    return attention
+
+
+def _reference_forward(attention: DiTAttention, x: torch.Tensor) -> torch.Tensor:
+    """Independent SDPA reference through to_out, mirroring the forward path."""
+    query = attention.to_q(x)
+    key = attention.to_k(x)
+    value = attention.to_v(x)
+    batch_size, seq_len = x.shape[0], x.shape[1]
+    query = query.view(batch_size, seq_len, attention.heads, attention.dim_head)
+    key = key.view(batch_size, seq_len, attention.heads, attention.dim_head)
+    value = value.view(batch_size, seq_len, attention.heads, attention.dim_head)
+    out = F.scaled_dot_product_attention(
+        query.transpose(1, 2),
+        key.transpose(1, 2),
+        value.transpose(1, 2),
+        dropout_p=0.0,
+        is_causal=False,
+        scale=attention.scale,
+    ).transpose(1, 2)
+    out = out.view(batch_size, seq_len, attention.inner_dim)
+    return attention.to_out(out)
+
+
+@pytest.mark.parametrize("backend_name", _FLOAT32_INCOMPATIBLE_BACKENDS)
+def test_float32_uses_sdpa_for_automatic_incompatible_backend(backend_name, fake_attention):
+    attention = _build_attention(backend_name=backend_name, explicit=False)
+    x = torch.randn(2, 5, 16)
+
+    actual = attention(x)
+
+    torch.testing.assert_close(actual, _reference_forward(attention, x))
+    assert attention.attn.calls == 0
+
+
+@pytest.mark.parametrize("backend_name", _FLOAT32_INCOMPATIBLE_BACKENDS)
+def test_explicit_backend_remains_authoritative(backend_name, fake_attention):
+    attention = _build_attention(backend_name=backend_name, explicit=True)
+    x = torch.zeros(1, 3, 16)
+
+    actual = attention(x)
+
+    batch_size, seq_len = x.shape[0], x.shape[1]
+    expected_hidden = torch.full((batch_size, seq_len, attention.inner_dim), 7, dtype=x.dtype)
+    torch.testing.assert_close(actual, attention.to_out(expected_hidden))
+    assert attention.attn.calls == 1
+
+
+def test_float32_keeps_compatible_automatic_backend(fake_attention):
+    attention = _build_attention(backend_name="TORCH_SDPA", explicit=False)
+    x = torch.zeros(1, 3, 16)
+
+    actual = attention(x)
+
+    batch_size, seq_len = x.shape[0], x.shape[1]
+    expected_hidden = torch.full((batch_size, seq_len, attention.inner_dim), 7, dtype=x.dtype)
+    torch.testing.assert_close(actual, attention.to_out(expected_hidden))
+    assert attention.attn.calls == 1

vllm_omni/model_executor/models/glm_tts/glm_tts_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
 """GLM-TTS DiT (Diffusion Transformer) Model.
 
 Flow Matching model that converts speech tokens to mel-spectrogram.
@@ -305,6 +305,21 @@ class FeedForward(nn.Module):
         return self.ff(x)
 
 
+# These backends reject float32 Q/K/V. GLM-TTS runs its DiT stage in float32
+# (the wrapper keeps the Euler ODE state 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. Mirrors the MiniMax Music 3 fix (#7354).
+_FLOAT32_UNSUPPORTED_BACKENDS = {
+    "FLASH_ATTN",
+    "FLASH_ATTN_HUB",
+    "FLASH_ATTN_3_HUB",
+    "CUDNN_ATTN",
+    "FLASHINFER_ATTN",
+}
+
+
 class DiTAttention(nn.Module):
     """Attention module using diffusion infrastructure."""
 
@@ -327,6 +342,13 @@ class DiTAttention(nn.Module):
             softmax_scale=self.scale,
             causal=False,
         )
+        backend = getattr(self.attn, "attn_backend", None)
+        backend_name = backend.get_name() if backend is not None else None
+        self._auto_float32_sdpa = (
+            backend is not None
+            and not getattr(self.attn, "backend_explicit", False)
+            and backend_name in _FLOAT32_UNSUPPORTED_BACKENDS
+        )
 
     def forward(
         self,
@@ -369,6 +391,23 @@ class DiTAttention(nn.Module):
                 is_causal=False,
             )
             out = out.permute(0, 2, 1, 3)  # [B, T, H, D]
+        elif self._auto_float32_sdpa and query.dtype == torch.float32:
+            # The automatically selected backend rejects float32 Q/K/V, and
+            # GLM-TTS decodes in float32, so route through SDPA for this input.
+            # Explicit backend choices keep the shared layer's fail-fast
+            # behavior (matches MiniMax Music 3, #7354).
+            q = query.transpose(1, 2)  # [B, H, T, D]
+            k = key.transpose(1, 2)
+            v = value.transpose(1, 2)
+            out = F.scaled_dot_product_attention(
+                q,
+                k,
+                v,
+                dropout_p=0.0,
+                is_causal=False,
+                scale=self.scale,
+            )
+            out = out.transpose(1, 2)  # [B, T, H, D]
         else:
             out = self.attn(query, key, value, attn_metadata=None)
         out = out.view(batch_size, seq_len, self.inner_dim)