[Perf][Bugfix][OmniVoice] Restore float16 serving, and fuse the generator hot loop (#6317)

Signed-off-by: MrlixiangWE <mrdanaer@gmail.com>

tests/model_executor/models/omnivoice/test_cuda_graph_generator.py

@@ -31,6 +31,7 @@ NUM_CB = 8
 VOCAB = 1025
 HIDDEN = 64
 CAPTURE_SIZES = [32, 64, 128]
+HEAD_DIM = 64
 
 
 # ---------------------------------------------------------------------------
@@ -53,22 +54,23 @@ class _SyntheticGenerator:
         self.config = self._Cfg()
         self.text_embedding = nn.Embedding(1000, HIDDEN).to(device).eval()
         self._linear = nn.Linear(HIDDEN, NUM_CB * VOCAB, bias=False).to(device).eval()
-        self._rope_cos: torch.Tensor | None = None
-        self._rope_sin: torch.Tensor | None = None
+        self._rope_table: torch.Tensor | None = None
 
-    def _ensure_rope(self, seq_len: int, device: torch.device) -> None:
-        if self._rope_cos is None or self._rope_cos.shape[0] < seq_len:
-            max_len = max(seq_len, 512)
-            self._rope_cos = torch.zeros(max_len, 64, device=device)
-            self._rope_sin = torch.zeros(max_len, 64, device=device)
+    @property
+    def model_dtype(self) -> torch.dtype:
+        """Part of the interface _OmniVoiceCUDAGraphForward expects of a generator."""
+        return self.text_embedding.weight.dtype
+
+    def _rope_table_for(self, seq_len: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
+        """Part of the interface _OmniVoiceCUDAGraphForward expects of a generator."""
+        return torch.zeros(seq_len, HEAD_DIM, device=device, dtype=dtype)
 
     def _step_forward(
         self,
         input_ids: torch.Tensor,
         audio_mask: torch.Tensor,
         attention_mask: torch.Tensor | None,
-        cos: torch.Tensor,
-        sin: torch.Tensor,
+        rope_table: torch.Tensor,
     ) -> torch.Tensor:
         two_b, _, S = input_ids.shape
         x = self.text_embedding(input_ids[:, 0, :].clamp(0, 999))  # [two_b, S, H]
@@ -82,11 +84,9 @@ class _SyntheticGenerator:
 
 
 def _eager(gen: _SyntheticGenerator, ids: torch.Tensor, mask: torch.Tensor, attn: torch.Tensor) -> torch.Tensor:
-    """Run _step_forward with zero RoPE tensors (synthetic gen ignores them)."""
-    S = ids.shape[-1]
-    cos = torch.zeros(S, 64, device=ids.device)
-    sin = torch.zeros(S, 64, device=ids.device)
-    return gen._step_forward(ids, mask, attn, cos, sin)
+    """Run _step_forward with a zero RoPE table (synthetic gen ignores it)."""
+    two_b, _, S = ids.shape
+    return gen._step_forward(ids, mask, attn, gen._rope_table_for(S, ids.device, torch.float32))
 
 
 def _make_inputs(seq_len: int, device: torch.device = DEVICE):
@@ -251,13 +251,10 @@ def test_cuda_graph_disabled_matches_eager_generator():
     mask = torch.ones(2, seq_len, dtype=torch.bool, device=DEVICE)
     attn = torch.ones(2, 1, seq_len, seq_len, dtype=torch.bool, device=DEVICE)
 
-    gen_eager._ensure_rope(seq_len, DEVICE)
-    dtype = gen_eager.text_embedding.weight.dtype
-    cos = gen_eager._rope_cos[:seq_len].to(device=DEVICE, dtype=dtype)
-    sin = gen_eager._rope_sin[:seq_len].to(device=DEVICE, dtype=dtype)
+    rope_table = gen_eager._rope_table_for(seq_len, DEVICE, gen_eager.model_dtype)
 
     with torch.no_grad():
-        eager_logits = gen_eager._step_forward(ids, mask, attn, cos, sin)
+        eager_logits = gen_eager._step_forward(ids, mask, attn, rope_table)
         graph_logits = gen_graph._cuda_graph_fwd(ids, mask, attn)
 
     torch.testing.assert_close(graph_logits, eager_logits, atol=0, rtol=0)

tests/model_executor/models/omnivoice/test_fused_projection_load.py

@@ -0,0 +1,183 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""The OmniVoice generator packs q/k/v and gate/up into fused projections.
+
+The HF checkpoint stores those five tensors separately, so packing them is a
+load-time step -- and the failure mode it introduces is silent. The generic
+name-based loader looks up ``q_proj``/``gate_proj`` as modules; once they no
+longer exist it finds nothing, logs a warning, and leaves the fused parameters
+at their random initialization. The model then serves requests and produces
+noise. These tests assert the property that failure mode violates: after a
+load, every fused parameter is fully written from the checkpoint, and anything
+that would leave one partly written raises instead.
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.models.omnivoice.omnivoice_generator import (
+    _FUSED_PROJECTIONS,
+    OmniVoiceGenerator,
+)
+from vllm_omni.transformers_utils.configs.omnivoice import OmniVoiceConfig
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+NUM_LAYERS = 2
+HIDDEN = 32
+HEAD_DIM = 8
+NUM_HEADS = 4
+NUM_KV_HEADS = 2
+INTERMEDIATE = 64
+
+
+def _config() -> OmniVoiceConfig:
+    return OmniVoiceConfig(
+        llm_config={
+            "hidden_size": HIDDEN,
+            "num_hidden_layers": NUM_LAYERS,
+            "num_attention_heads": NUM_HEADS,
+            "num_key_value_heads": NUM_KV_HEADS,
+            "head_dim": HEAD_DIM,
+            "intermediate_size": INTERMEDIATE,
+            "vocab_size": 64,
+            "max_position_embeddings": 128,
+        },
+        enable_cuda_graph=False,
+    )
+
+
+def _checkpoint_shards() -> dict[str, torch.Tensor]:
+    """The five per-layer tensors an OmniVoice checkpoint actually stores."""
+    torch.manual_seed(0)
+    shards: dict[str, torch.Tensor] = {}
+    for layer in range(NUM_LAYERS):
+        prefix = f"llm.layers.{layer}"
+        shards[f"{prefix}.self_attn.q_proj.weight"] = torch.randn(NUM_HEADS * HEAD_DIM, HIDDEN)
+        shards[f"{prefix}.self_attn.k_proj.weight"] = torch.randn(NUM_KV_HEADS * HEAD_DIM, HIDDEN)
+        shards[f"{prefix}.self_attn.v_proj.weight"] = torch.randn(NUM_KV_HEADS * HEAD_DIM, HIDDEN)
+        shards[f"{prefix}.mlp.gate_proj.weight"] = torch.randn(INTERMEDIATE, HIDDEN)
+        shards[f"{prefix}.mlp.up_proj.weight"] = torch.randn(INTERMEDIATE, HIDDEN)
+    return shards
+
+
+def test_every_fused_parameter_is_written() -> None:
+    generator = OmniVoiceGenerator(_config())
+    shards = _checkpoint_shards()
+
+    loaded = generator._load_fused_projections(shards)
+
+    assert loaded == set(shards), "some checkpoint shards were not consumed"
+    shards_per_layer = sum(len(paths) for paths in _FUSED_PROJECTIONS.values())
+    assert shards_per_layer == 5, "q,k,v and gate,up"
+    assert len(loaded) == NUM_LAYERS * shards_per_layer
+
+
+def test_packed_weight_matches_the_shards_it_came_from() -> None:
+    """Packing order is q,k,v and gate,up -- the split in forward() assumes it."""
+    generator = OmniVoiceGenerator(_config())
+    shards = _checkpoint_shards()
+    generator._load_fused_projections(shards)
+
+    for layer in range(NUM_LAYERS):
+        prefix = f"llm.layers.{layer}"
+        qkv = torch.cat(
+            [
+                shards[f"{prefix}.self_attn.q_proj.weight"],
+                shards[f"{prefix}.self_attn.k_proj.weight"],
+                shards[f"{prefix}.self_attn.v_proj.weight"],
+            ],
+            dim=0,
+        )
+        gate_up = torch.cat(
+            [shards[f"{prefix}.mlp.gate_proj.weight"], shards[f"{prefix}.mlp.up_proj.weight"]],
+            dim=0,
+        )
+        torch.testing.assert_close(generator.layers[layer].self_attn.qkv_proj.weight, qkv)
+        torch.testing.assert_close(generator.layers[layer].mlp.gate_up_proj.weight, gate_up)
+
+
+def test_no_fused_parameter_keeps_its_random_init() -> None:
+    """The corruption this guards against: a fused param never written at all."""
+    generator = OmniVoiceGenerator(_config())
+    before = {
+        name: param.detach().clone()
+        for name, param in generator.named_parameters()
+        if name.endswith(("qkv_proj.weight", "gate_up_proj.weight"))
+    }
+    assert len(before) == NUM_LAYERS * len(_FUSED_PROJECTIONS)
+
+    generator._load_fused_projections(_checkpoint_shards())
+
+    for name, initial in before.items():
+        current = dict(generator.named_parameters())[name]
+        assert not torch.equal(current, initial), f"{name} was left at its initialization"
+
+
+@pytest.mark.parametrize(
+    "dropped",
+    ["self_attn.k_proj", "self_attn.v_proj", "self_attn.q_proj", "mlp.gate_proj", "mlp.up_proj"],
+)
+def test_a_missing_shard_raises_instead_of_loading_partially(dropped: str) -> None:
+    generator = OmniVoiceGenerator(_config())
+    shards = _checkpoint_shards()
+    del shards[f"llm.layers.0.{dropped}.weight"]
+
+    with pytest.raises(ValueError, match="missing"):
+        generator._load_fused_projections(shards)
+
+
+def test_a_wrong_shaped_shard_raises() -> None:
+    generator = OmniVoiceGenerator(_config())
+    shards = _checkpoint_shards()
+    shards["llm.layers.0.self_attn.q_proj.weight"] = torch.randn(NUM_HEADS * HEAD_DIM + 1, HIDDEN)
+
+    with pytest.raises(ValueError, match="pack to"):
+        generator._load_fused_projections(shards)
+
+
+def test_a_checkpoint_missing_a_whole_layer_raises() -> None:
+    """Half-loaded is the dangerous state: it neither errors nor works."""
+    generator = OmniVoiceGenerator(_config())
+    shards = {k: v for k, v in _checkpoint_shards().items() if not k.startswith("llm.layers.1.")}
+
+    with pytest.raises(ValueError, match="fused projections"):
+        generator._load_fused_projections(shards)
+
+
+def test_a_checkpoint_with_no_fused_shards_is_left_alone() -> None:
+    """An unrelated state_dict must not trip the completeness check."""
+    generator = OmniVoiceGenerator(_config())
+    assert generator._load_fused_projections({"audio_heads.weight": torch.randn(4, 4)}) == set()
+
+
+def test_load_weights_wires_the_packing_in(tmp_path, monkeypatch: pytest.MonkeyPatch) -> None:
+    """The packing has to be reached from load_weights, not merely available.
+
+    Every test above calls _load_fused_projections directly, so deleting its call
+    site from load_weights would leave them all green while the model loaded
+    corrupted. This goes through the public entry point instead.
+    """
+    import safetensors.torch
+
+    generator = OmniVoiceGenerator(_config())
+    shards = _checkpoint_shards()
+    before = generator.layers[0].self_attn.qkv_proj.weight.detach().clone()
+
+    (tmp_path / "model.safetensors").write_bytes(b"")  # only its existence is checked
+    monkeypatch.setattr(safetensors.torch, "load_file", lambda *args, **kwargs: shards)
+    generator.load_weights(str(tmp_path), torch.device("cpu"))
+
+    expected = torch.cat(
+        [
+            shards["llm.layers.0.self_attn.q_proj.weight"],
+            shards["llm.layers.0.self_attn.k_proj.weight"],
+            shards["llm.layers.0.self_attn.v_proj.weight"],
+        ],
+        dim=0,
+    )
+    written = generator.layers[0].self_attn.qkv_proj.weight
+    assert not torch.equal(written, before), "load_weights left the fused projection at its init"
+    torch.testing.assert_close(written, expected)

tests/model_executor/models/omnivoice/test_fused_qkv_rope.py

@@ -0,0 +1,166 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""The fused attention prologue must equal the six steps it replaces.
+
+``fused_qkv_norm_rope`` collapses split, per-head RMSNorm of Q and K, RoPE on
+both, the GQA broadcast of K and V, and the transpose into SDPA's layout into
+one Triton pass. That is a lot of behaviour to fold into one kernel, so these
+check it against the unfused reference across the geometries and dtypes it
+claims to support, and check that unsupported geometries fall back rather than
+producing something wrong.
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+from vllm.triton_utils import HAS_TRITON
+
+from vllm_omni.model_executor.models.omnivoice.fused_qkv_rope import (
+    _eager_qkv_norm_rope,
+    fused_cuda_supported,
+    fused_qkv_norm_rope,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cuda]
+
+cuda_only = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+triton_only = pytest.mark.skipif(not HAS_TRITON, reason="Triton required")
+
+EPS = 1e-6
+TOL = {
+    torch.float32: dict(atol=2e-5, rtol=2e-5),
+    torch.float16: dict(atol=4e-3, rtol=4e-3),
+    torch.bfloat16: dict(atol=3e-2, rtol=3e-2),
+}
+
+
+def _inputs(batch, seq_len, num_heads, num_kv_heads, head_dim, dtype, device="cuda"):
+    torch.manual_seed(11)
+    total = num_heads + 2 * num_kv_heads
+    qkv = torch.randn(batch, seq_len, total, head_dim, device=device, dtype=dtype)
+    q_weight = torch.randn(head_dim, device=device, dtype=dtype)
+    k_weight = torch.randn(head_dim, device=device, dtype=dtype)
+    freqs = torch.randn(seq_len, head_dim // 2, device=device)
+    rope_table = torch.cat([freqs.cos(), freqs.sin()], dim=-1).to(dtype)
+    return qkv, q_weight, k_weight, rope_table
+
+
+@cuda_only
+@triton_only
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
+@pytest.mark.parametrize(
+    "batch,seq_len,num_heads,num_kv_heads,head_dim",
+    [
+        (2, 44, 16, 8, 128),  # OmniVoice: the CFG pair, 2:1 GQA
+        (2, 130, 16, 8, 128),
+        (1, 7, 8, 8, 128),  # no GQA broadcast
+        (3, 16, 16, 8, 64),
+        (2, 33, 32, 8, 128),  # 4:1 GQA
+    ],
+)
+def test_matches_the_unfused_reference(dtype, batch, seq_len, num_heads, num_kv_heads, head_dim):
+    qkv, q_weight, k_weight, rope_table = _inputs(batch, seq_len, num_heads, num_kv_heads, head_dim, dtype)
+    assert fused_cuda_supported(qkv, num_heads, num_kv_heads)
+
+    got = fused_qkv_norm_rope(qkv, q_weight, k_weight, rope_table, EPS, num_heads, num_kv_heads)
+    expected = _eager_qkv_norm_rope(qkv, q_weight, k_weight, rope_table, EPS, num_heads, num_kv_heads)
+
+    for name, a, b in zip(("q", "k", "v"), got, expected):
+        assert a.shape == (batch, num_heads, seq_len, head_dim), name
+        assert a.is_contiguous(), name
+        torch.testing.assert_close(a, b, msg=lambda m, n=name: f"{n}: {m}", **TOL[dtype])
+
+
+@cuda_only
+@triton_only
+def test_v_is_the_untouched_projection_broadcast_into_place():
+    """V gets no norm and no rotation, only the GQA broadcast and the transpose."""
+    batch, seq_len, num_heads, num_kv_heads, head_dim = 2, 44, 16, 8, 128
+    qkv, q_weight, k_weight, rope_table = _inputs(batch, seq_len, num_heads, num_kv_heads, head_dim, torch.float32)
+
+    _, _, v = fused_qkv_norm_rope(qkv, q_weight, k_weight, rope_table, EPS, num_heads, num_kv_heads)
+
+    raw_v = qkv[:, :, num_heads + num_kv_heads :]
+    expected = raw_v.repeat_interleave(num_heads // num_kv_heads, dim=2).permute(0, 2, 1, 3)
+    torch.testing.assert_close(v, expected.contiguous(), atol=0, rtol=0)
+
+
+@cuda_only
+@triton_only
+def test_kv_heads_are_broadcast_not_reordered():
+    """Query group g must see KV head g // repeat, not some transposed pairing."""
+    batch, seq_len, num_heads, num_kv_heads, head_dim = 1, 4, 16, 8, 128
+    qkv, q_weight, k_weight, rope_table = _inputs(batch, seq_len, num_heads, num_kv_heads, head_dim, torch.float32)
+
+    _, k, _ = fused_qkv_norm_rope(qkv, q_weight, k_weight, rope_table, EPS, num_heads, num_kv_heads)
+
+    repeat = num_heads // num_kv_heads
+    for head in range(num_heads):
+        torch.testing.assert_close(k[:, head], k[:, (head // repeat) * repeat], atol=0, rtol=0)
+
+
+@cuda_only
+@triton_only
+def test_position_zero_is_a_pure_rmsnorm():
+    """A rope table row of cos=1, sin=0 must leave the normalized value alone."""
+    batch, seq_len, num_heads, num_kv_heads, head_dim = 1, 4, 8, 8, 128
+    qkv, q_weight, k_weight, _ = _inputs(batch, seq_len, num_heads, num_kv_heads, head_dim, torch.float32)
+    half = head_dim // 2
+    rope_table = torch.cat(
+        [torch.ones(seq_len, half, device="cuda"), torch.zeros(seq_len, half, device="cuda")], dim=-1
+    )
+
+    q, _, _ = fused_qkv_norm_rope(qkv, q_weight, k_weight, rope_table, EPS, num_heads, num_kv_heads)
+
+    raw_q = qkv[:, :, :num_heads].to(torch.float32)
+    variance = raw_q.pow(2).mean(-1, keepdim=True)
+    expected = (raw_q * torch.rsqrt(variance + EPS) * q_weight).permute(0, 2, 1, 3)
+    torch.testing.assert_close(q, expected.contiguous(), **TOL[torch.float32])
+
+
+@cuda_only
+@pytest.mark.parametrize(
+    "num_heads,num_kv_heads,head_dim",
+    [(12, 4, 128), (16, 8, 96), (16, 6, 128)],
+)
+def test_unsupported_geometry_falls_back_and_still_matches(num_heads, num_kv_heads, head_dim):
+    """Head counts the tiling cannot split, or a head_dim that is not a power of two."""
+    batch, seq_len = 2, 8
+    qkv, q_weight, k_weight, rope_table = _inputs(batch, seq_len, num_heads, num_kv_heads, head_dim, torch.float32)
+    assert not fused_cuda_supported(qkv, num_heads, num_kv_heads)
+
+    got = fused_qkv_norm_rope(qkv, q_weight, k_weight, rope_table, EPS, num_heads, num_kv_heads)
+    expected = _eager_qkv_norm_rope(qkv, q_weight, k_weight, rope_table, EPS, num_heads, num_kv_heads)
+    for a, b in zip(got, expected):
+        torch.testing.assert_close(a, b, atol=0, rtol=0)
+
+
+@pytest.mark.cpu
+def test_cpu_tensors_take_the_eager_path():
+    batch, seq_len, num_heads, num_kv_heads, head_dim = 1, 5, 16, 8, 128
+    qkv, q_weight, k_weight, rope_table = _inputs(
+        batch, seq_len, num_heads, num_kv_heads, head_dim, torch.float32, device="cpu"
+    )
+    assert not fused_cuda_supported(qkv, num_heads, num_kv_heads)
+
+    q, k, v = fused_qkv_norm_rope(qkv, q_weight, k_weight, rope_table, EPS, num_heads, num_kv_heads)
+    assert q.shape == k.shape == v.shape == (batch, num_heads, seq_len, head_dim)
+
+
+@pytest.mark.cpu
+def test_a_mismatched_head_count_raises():
+    qkv = torch.randn(1, 4, 30, 128)  # 16 + 2*8 would be 32
+    w = torch.randn(128)
+    table = torch.zeros(4, 128)
+    with pytest.raises(ValueError, match="packed QKV"):
+        fused_qkv_norm_rope(qkv, w, w, table, EPS, 16, 8)
+
+
+@pytest.mark.cpu
+def test_a_short_rope_table_raises():
+    qkv = torch.randn(1, 10, 32, 128)
+    w = torch.randn(128)
+    table = torch.zeros(4, 128)
+    with pytest.raises(ValueError, match="rope_table"):
+        fused_qkv_norm_rope(qkv, w, w, table, EPS, 16, 8)

tests/model_executor/models/omnivoice/test_mask_dtype.py

@@ -0,0 +1,183 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""OmniVoice attention masks must follow the model dtype, not a hardcoded float32.
+
+SDPA requires the additive attention mask to carry the same dtype as the query.
+The generator used to materialize that mask as float32 unconditionally, in three
+places (the shared per-forward mask, the CUDA-graph capture buffer, and the
+replay-time normalization), so every half-precision deployment died on the first
+attention call with::
+
+    RuntimeError: invalid dtype for bias - should match query's dtype
+
+(the exact wording depends on which SDPA backend is selected for the shape; the
+math backend words the same rejection as ``attn_mask.dtype``.)
+
+That left OmniVoice servable only in float32, even though the upstream k2-fsa
+implementation runs it in float16.
+
+Note the split between the CPU and CUDA tests below: the CPU SDPA backend
+silently accepts the mismatched mask, so only the CUDA tests can pin the actual
+crash. The CPU tests cover the mask contract itself, which is what the fix
+changes and what a future regression would break first.
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.models.omnivoice.omnivoice_generator import (
+    OmniVoiceAttention,
+    OmniVoiceGenerator,
+    _additive_float_mask,
+)
+from vllm_omni.transformers_utils.configs.omnivoice import OmniVoiceConfig
+
+HALF_DTYPES = [torch.float16, torch.bfloat16]
+ALL_DTYPES = HALF_DTYPES + [torch.float32]
+
+cpu_test = pytest.mark.core_model
+cuda_test = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+
+
+def _tiny_config() -> OmniVoiceConfig:
+    """A real OmniVoiceConfig sized so the module fits in a unit test."""
+    return OmniVoiceConfig(
+        llm_config={
+            "hidden_size": 32,
+            "num_hidden_layers": 2,
+            "num_attention_heads": 4,
+            "num_key_value_heads": 2,
+            "head_dim": 8,
+            "intermediate_size": 64,
+            "vocab_size": 64,
+            "max_position_embeddings": 128,
+        },
+        enable_cuda_graph=False,
+    )
+
+
+def _inputs(dtype: torch.dtype, device: torch.device) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    hidden = torch.randn(2, 6, 32, device=device, dtype=dtype)
+    bool_mask = torch.ones(2, 1, 6, 6, dtype=torch.bool, device=device)
+    bool_mask[:, :, :, 4:] = False
+    rope_table = torch.zeros(6, 8, device=device, dtype=dtype)
+    rope_table[:, :4] = 1.0  # cos = 1, sin = 0: no rotation, so the test is about the mask
+    return hidden, bool_mask, rope_table
+
+
+# --------------------------------------------------------------------------
+# The mask contract (CPU)
+# --------------------------------------------------------------------------
+
+
+@cpu_test
+@pytest.mark.cpu
+@pytest.mark.parametrize("dtype", ALL_DTYPES)
+def test_additive_float_mask_uses_requested_dtype(dtype: torch.dtype) -> None:
+    out = _additive_float_mask(torch.tensor([[True, False]]), dtype)
+    assert out.dtype == dtype
+
+
+@cpu_test
+@pytest.mark.cpu
+def test_additive_float_mask_keeps_mask_semantics() -> None:
+    """Attend positions stay at 0.0 and masked ones at -inf, in half precision too."""
+    out = _additive_float_mask(torch.tensor([[True, False]]), torch.float16)
+    assert out[0, 0].item() == 0.0
+    assert out[0, 1].item() == float("-inf")
+
+
+@cpu_test
+@pytest.mark.cpu
+def test_additive_float_mask_requires_an_explicit_dtype() -> None:
+    """The float32 default is the bug; keeping the parameter required is the fix."""
+    with pytest.raises(TypeError):
+        _additive_float_mask(torch.tensor([[True]]))  # type: ignore[call-arg]
+
+
+@cpu_test
+@pytest.mark.cpu
+@pytest.mark.parametrize("dtype", ALL_DTYPES)
+def test_generator_model_dtype_tracks_its_weights(dtype: torch.dtype) -> None:
+    """The single source of truth the three former float32 literals now defer to."""
+    assert OmniVoiceGenerator(_tiny_config()).to(dtype).model_dtype == dtype
+
+
+@cpu_test
+@pytest.mark.cpu
+@pytest.mark.parametrize("dtype", HALF_DTYPES)
+def test_attention_accepts_a_model_dtype_mask(dtype: torch.dtype) -> None:
+    attn = OmniVoiceAttention(_tiny_config()).to(dtype).eval()
+    hidden, bool_mask, rope_table = _inputs(dtype, torch.device("cpu"))
+
+    with torch.inference_mode():
+        out = attn(hidden, rope_table, attention_mask=_additive_float_mask(bool_mask, dtype))
+
+    assert out.shape == hidden.shape
+    assert out.dtype == dtype
+    assert torch.isfinite(out).all()
+
+
+# --------------------------------------------------------------------------
+# The crash itself (CUDA only — the CPU backend does not reject the mismatch)
+# --------------------------------------------------------------------------
+
+
+@cuda_test
+@pytest.mark.parametrize("dtype", HALF_DTYPES)
+def test_float32_mask_is_what_broke_half_precision(dtype: torch.dtype) -> None:
+    """Pin the original failure so a float32 default cannot come back unnoticed."""
+    attn = OmniVoiceAttention(_tiny_config()).to(device="cuda:0", dtype=dtype).eval()
+    hidden, bool_mask, rope_table = _inputs(dtype, torch.device("cuda:0"))
+
+    with pytest.raises(RuntimeError, match=r"(invalid dtype for bias|attn_mask)"), torch.inference_mode():
+        attn(hidden, rope_table, attention_mask=_additive_float_mask(bool_mask, torch.float32))
+
+
+@cuda_test
+@pytest.mark.parametrize("dtype", HALF_DTYPES)
+def test_real_path_forward_runs_in_half_precision(dtype: torch.dtype) -> None:
+    """The regression, on the production path with the Triton norm kernels live."""
+    attn = OmniVoiceAttention(_tiny_config()).to(device="cuda:0", dtype=dtype).eval()
+    hidden, bool_mask, rope_table = _inputs(dtype, torch.device("cuda:0"))
+
+    with torch.inference_mode():
+        out = attn(hidden, rope_table, attention_mask=_additive_float_mask(bool_mask, dtype))
+
+    assert out.dtype == dtype
+    assert torch.isfinite(out).all()
+
+
+@cuda_test
+@pytest.mark.parametrize("dtype", HALF_DTYPES)
+def test_generator_forward_runs_in_half_precision(dtype: torch.dtype) -> None:
+    """End-to-end over the real iterative loop, which is where the float32 mask was built.
+
+    This is the test that fails on an unfixed tree: ``forward`` materialized the
+    shared SDPA mask as float32 regardless of the weights' dtype.
+    """
+    device = torch.device("cuda:0")
+    config = _tiny_config()
+    generator = OmniVoiceGenerator(config).to(device=device, dtype=dtype).eval()
+
+    seq_len, target_len = 12, 4
+    input_ids = torch.zeros(2, config.num_audio_codebook, seq_len, dtype=torch.long, device=device)
+    input_ids[:, 1:, :] = config.audio_mask_id
+    audio_mask = torch.zeros(2, seq_len, dtype=torch.bool, device=device)
+    audio_mask[:, seq_len - target_len :] = True
+    attention_mask = torch.ones(2, 1, seq_len, seq_len, dtype=torch.bool, device=device)
+
+    with torch.inference_mode():
+        tokens = generator(
+            input_ids,
+            audio_mask,
+            attention_mask,
+            target_lens=[target_len],
+            seed=0,
+            num_step=2,
+        )
+
+    assert tokens.shape == (1, config.num_audio_codebook, target_len)
+    assert tokens.dtype == torch.long

tests/model_executor/models/omnivoice/test_triton_kernels.py

@@ -15,6 +15,7 @@ import torch
 
 pytestmark = [
     pytest.mark.core_model,
+    pytest.mark.cuda,
     pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"),
 ]
 
@@ -48,7 +49,9 @@ def _ref_rms_norm(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Te
 
 
 def _ref_swiglu(gate: torch.Tensor, up: torch.Tensor) -> torch.Tensor:
-    # Kernel: gate → float32, silu, cast to up.dtype, mul up
+    # Kernel: gate → float32, silu, cast to up.dtype, mul up.
+    # It now reads both halves out of one packed [..., 2 * cols] activation, so
+    # the tests concatenate what they used to pass as two tensors.
     gate_f32 = gate.float()
     silu = gate_f32 * torch.sigmoid(gate_f32)
     return silu.to(up.dtype) * up
@@ -124,7 +127,7 @@ def test_swiglu_float32(rows, cols):
     up = torch.randn(rows, cols, device=DEVICE)
 
     ref = _ref_swiglu(gate, up)
-    out = triton_swiglu(gate, up)
+    out = triton_swiglu(torch.cat([gate, up], dim=-1))
 
     torch.testing.assert_close(out, ref, atol=1e-5, rtol=1e-5)
 
@@ -136,7 +139,7 @@ def test_swiglu_negative_gate():
     up = torch.ones(4, 64, device=DEVICE)
 
     ref = _ref_swiglu(gate, up)
-    out = triton_swiglu(gate, up)
+    out = triton_swiglu(torch.cat([gate, up], dim=-1))
 
     torch.testing.assert_close(out, ref, atol=1e-5, rtol=1e-5)
     assert not torch.isnan(out).any()
@@ -146,7 +149,7 @@ def test_swiglu_negative_gate():
 def test_swiglu_output_shape_preserved():
     gate = torch.randn(3, 8, 512, device=DEVICE)
     up = torch.randn(3, 8, 512, device=DEVICE)
-    out = triton_swiglu(gate, up)
+    out = triton_swiglu(torch.cat([gate, up], dim=-1))
     assert out.shape == gate.shape