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

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

benchmarks/tts/model_configs.yaml

@@ -132,3 +132,14 @@ models:
     endpoint: /v1/audio/speech
     task_extra_body:
       voice_design: {}
+
+  # OmniVoice serves plain text-to-speech over /v1/audio/speech; voice cloning
+  # and voice design are offline-only for this model, so default_voice is the
+  # only task the server can answer.
+  k2-fsa/OmniVoice:
+    supported_tasks: [default_voice]
+    backend: openai-audio-speech
+    endpoint: /v1/audio/speech
+    trust_remote_code: true
+    task_extra_body:
+      default_voice: {}

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
 
 

vllm_omni/model_executor/models/omnivoice/fused_qkv_rope.py

@@ -0,0 +1,273 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Fused attention prologue for models with a packed QKV projection.
+
+A Qwen3-style attention layer with GQA does six things between its projection
+and its SDPA call: split the packed QKV, RMSNorm Q and K per head, rotate both
+with RoPE, broadcast K and V across their query groups, and transpose all three
+into the ``[batch, heads, positions, head_dim]`` layout SDPA wants. Every one of
+those is cheap arithmetic over the whole activation, so together they cost
+several full round-trips of Q, K and V through memory.
+
+``fused_qkv_norm_rope`` does all of it in a single Triton pass. Each program
+owns one tile of heads of one token: it reads that tile out of the packed
+projection once, and writes the finished tensor straight into the attention
+layout, broadcasting K and V to their query groups from registers. RoPE's
+rotate-half partner is re-read from the same row rather than materialized as a
+concatenation, and the RMS is a per-head scalar, so the partner needs no second
+reduction.
+
+Measured on OmniVoice (28 layers x 32 unmasking steps, A800, float16), the
+unfused path launches 18 kernels between the packed projection and attention
+and this replaces them with 1, cutting the generator's forward by 25-36%
+depending on target length.
+
+Falls back to an eager reference for non-CUDA tensors, missing Triton, or a
+geometry the tiling cannot express.
+
+Nothing here is OmniVoice-specific beyond the calling convention, but OmniVoice
+is its only consumer today: every other model in this repo with a packed QKV
+projection runs on vLLM's own attention stack, whose backends already handle
+GQA and RoPE, so none of them performs the sequence this replaces. It therefore
+lives beside its consumer rather than in models/common/, the same placement
+fused_qk_norm_rope took under diffusion/. Lift it out when a second model wants
+it.
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+from vllm.triton_utils import HAS_TRITON, tl, triton
+
+# One tile of heads per program. Head counts must be a multiple of this so that
+# a tile never straddles the Q/K/V boundary, which keeps the branch uniform.
+_HEAD_TILE = 8
+
+
+if HAS_TRITON:
+
+    @triton.jit
+    def _qkv_norm_rope_kernel(
+        qkv_ptr,
+        q_weight_ptr,
+        k_weight_ptr,
+        rope_ptr,
+        q_out_ptr,
+        k_out_ptr,
+        v_out_ptr,
+        qkv_stride_token,
+        qkv_stride_head,
+        rope_stride_pos,
+        out_stride_batch,
+        out_stride_head,
+        out_stride_pos,
+        seq_len,
+        num_heads: tl.constexpr,
+        num_kv_heads: tl.constexpr,
+        head_dim: tl.constexpr,
+        rotary_half: tl.constexpr,
+        kv_repeat: tl.constexpr,
+        head_tile: tl.constexpr,
+        eps: tl.constexpr,
+    ):
+        token = tl.program_id(0)
+        tile_id = tl.program_id(1)
+        batch = token // seq_len
+        pos = token % seq_len
+
+        dims = tl.arange(0, head_dim)
+        heads = tile_id * head_tile + tl.arange(0, head_tile)
+        src = qkv_ptr + token * qkv_stride_token + heads[:, None] * qkv_stride_head + dims[None, :]
+
+        # rotate_half pairs dim d with d +/- rotary_half, negating the lower half.
+        partner = tl.where(dims < rotary_half, dims + rotary_half, dims - rotary_half)
+        sign = tl.where(dims < rotary_half, -1.0, 1.0)
+        freq = tl.where(dims < rotary_half, dims, dims - rotary_half)
+        cos = tl.load(rope_ptr + pos * rope_stride_pos + freq)
+        sin = tl.load(rope_ptr + pos * rope_stride_pos + rotary_half + freq)
+
+        q_tiles: tl.constexpr = num_heads // head_tile
+        kv_tiles: tl.constexpr = num_kv_heads // head_tile
+
+        if tile_id < q_tiles:
+            x = tl.load(src).to(tl.float32)
+            weight = tl.load(q_weight_ptr + dims).to(tl.float32)
+            inv_rms = tl.rsqrt(tl.sum(x * x, axis=1) / head_dim + eps)
+            normed = x * inv_rms[:, None] * weight[None, :]
+            pair_x = tl.load(
+                qkv_ptr + token * qkv_stride_token + heads[:, None] * qkv_stride_head + partner[None, :]
+            ).to(tl.float32)
+            pair_weight = tl.load(q_weight_ptr + partner).to(tl.float32)
+            pair_normed = pair_x * inv_rms[:, None] * pair_weight[None, :]
+            out = normed * cos[None, :].to(tl.float32) + sign[None, :] * pair_normed * sin[None, :].to(tl.float32)
+            dst = (
+                q_out_ptr
+                + batch * out_stride_batch
+                + heads[:, None] * out_stride_head
+                + pos * out_stride_pos
+                + dims[None, :]
+            )
+            tl.store(dst, out.to(q_out_ptr.dtype.element_ty))
+
+        elif tile_id < q_tiles + kv_tiles:
+            # K is broadcast across its query group here, from registers.
+            # Measured alternative: SDPA's enable_gqa does the broadcast with no
+            # copy at all, but on torch 2.13 no fused kernel accepts mismatched
+            # head counts, so it falls back to the math backend and runs 2.0-4.9x
+            # slower than materializing K and V. Hence the broadcast, but folded
+            # into a store that was happening anyway.
+            kv_head = heads - num_heads
+            x = tl.load(src).to(tl.float32)
+            weight = tl.load(k_weight_ptr + dims).to(tl.float32)
+            inv_rms = tl.rsqrt(tl.sum(x * x, axis=1) / head_dim + eps)
+            normed = x * inv_rms[:, None] * weight[None, :]
+            pair_x = tl.load(
+                qkv_ptr + token * qkv_stride_token + heads[:, None] * qkv_stride_head + partner[None, :]
+            ).to(tl.float32)
+            pair_weight = tl.load(k_weight_ptr + partner).to(tl.float32)
+            pair_normed = pair_x * inv_rms[:, None] * pair_weight[None, :]
+            out = normed * cos[None, :].to(tl.float32) + sign[None, :] * pair_normed * sin[None, :].to(tl.float32)
+            stored = out.to(k_out_ptr.dtype.element_ty)
+            # Broadcast to the query group straight from registers.
+            for rep in tl.static_range(kv_repeat):
+                dst = (
+                    k_out_ptr
+                    + batch * out_stride_batch
+                    + (kv_head[:, None] * kv_repeat + rep) * out_stride_head
+                    + pos * out_stride_pos
+                    + dims[None, :]
+                )
+                tl.store(dst, stored)
+
+        else:
+            kv_head = heads - num_heads - num_kv_heads
+            stored = tl.load(src)
+            for rep in tl.static_range(kv_repeat):
+                dst = (
+                    v_out_ptr
+                    + batch * out_stride_batch
+                    + (kv_head[:, None] * kv_repeat + rep) * out_stride_head
+                    + pos * out_stride_pos
+                    + dims[None, :]
+                )
+                tl.store(dst, stored)
+
+
+def _eager_qkv_norm_rope(
+    qkv: torch.Tensor,
+    q_weight: torch.Tensor,
+    k_weight: torch.Tensor,
+    rope_table: torch.Tensor,
+    eps: float,
+    num_heads: int,
+    num_kv_heads: int,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """The six unfused steps, kept as the reference the kernel is checked against."""
+    head_dim = qkv.shape[-1]
+    half = head_dim // 2
+    q, k, v = qkv.split([num_heads, num_kv_heads, num_kv_heads], dim=2)
+
+    cos = rope_table[: qkv.shape[1], :half]
+    sin = rope_table[: qkv.shape[1], half:]
+    cos = torch.cat([cos, cos], dim=-1).unsqueeze(0).unsqueeze(2)
+    sin = torch.cat([sin, sin], dim=-1).unsqueeze(0).unsqueeze(2)
+
+    def norm_rope(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
+        normed = F.rms_norm(x.to(torch.float32), (head_dim,), weight.to(torch.float32), eps)
+        rotated = torch.cat([-normed[..., half:], normed[..., :half]], dim=-1)
+        return (normed * cos + rotated * sin).to(x.dtype)
+
+    repeat = num_heads // num_kv_heads
+    q = norm_rope(q, q_weight)
+    k = norm_rope(k, k_weight).repeat_interleave(repeat, dim=2)
+    v = v.repeat_interleave(repeat, dim=2)
+    return (q.permute(0, 2, 1, 3).contiguous(), k.permute(0, 2, 1, 3).contiguous(), v.permute(0, 2, 1, 3).contiguous())
+
+
+def fused_cuda_supported(qkv: torch.Tensor, num_heads: int, num_kv_heads: int) -> bool:
+    """Whether the Triton path can express this geometry."""
+    head_dim = qkv.shape[-1]
+    return (
+        HAS_TRITON
+        and qkv.is_cuda
+        and qkv.is_contiguous()
+        and head_dim > 0
+        and head_dim & (head_dim - 1) == 0  # a single tl.arange tiles the head
+        and num_heads % _HEAD_TILE == 0
+        and num_kv_heads % _HEAD_TILE == 0  # tiles never straddle the Q/K/V boundary
+        and num_heads % num_kv_heads == 0
+    )
+
+
+def fused_qkv_norm_rope(
+    qkv: torch.Tensor,
+    q_weight: torch.Tensor,
+    k_weight: torch.Tensor,
+    rope_table: torch.Tensor,
+    eps: float,
+    num_heads: int,
+    num_kv_heads: int,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Split a packed QKV projection into attention-ready Q, K and V.
+
+    Args:
+        qkv: ``[batch, positions, num_heads + 2 * num_kv_heads, head_dim]``, the
+            output of one packed projection.
+        q_weight: ``[head_dim]`` per-head RMSNorm weight for Q.
+        k_weight: ``[head_dim]`` per-head RMSNorm weight for K.
+        rope_table: ``[positions, head_dim]``, the first half of each row holding
+            ``cos(theta)`` and the second ``sin(theta)``, each ``head_dim // 2``
+            wide. One row per position; the kernel repeats it across the batch.
+        eps: RMSNorm epsilon.
+        num_heads: number of query heads.
+        num_kv_heads: number of key/value heads; must divide ``num_heads``.
+
+    Returns:
+        ``(q, k, v)``, each ``[batch, num_heads, positions, head_dim]`` and
+        contiguous, with K and V already broadcast across their query groups.
+    """
+    batch, seq_len, total_heads, head_dim = qkv.shape
+    if total_heads != num_heads + 2 * num_kv_heads:
+        raise ValueError(f"packed QKV has {total_heads} heads, expected {num_heads} + 2 * {num_kv_heads}")
+    if rope_table.shape[0] < seq_len or rope_table.shape[-1] != head_dim:
+        raise ValueError(f"rope_table {tuple(rope_table.shape)} does not cover [{seq_len}, {head_dim}]")
+
+    if not fused_cuda_supported(qkv, num_heads, num_kv_heads):
+        return _eager_qkv_norm_rope(qkv, q_weight, k_weight, rope_table, eps, num_heads, num_kv_heads)
+
+    out_shape = (batch, num_heads, seq_len, head_dim)
+    q_out = torch.empty(out_shape, dtype=qkv.dtype, device=qkv.device)
+    k_out = torch.empty(out_shape, dtype=qkv.dtype, device=qkv.device)
+    v_out = torch.empty(out_shape, dtype=qkv.dtype, device=qkv.device)
+
+    rope = rope_table[:seq_len].contiguous()
+    grid = (batch * seq_len, total_heads // _HEAD_TILE)
+    _qkv_norm_rope_kernel[grid](
+        qkv,
+        q_weight,
+        k_weight,
+        rope,
+        q_out,
+        k_out,
+        v_out,
+        qkv.stride(1),
+        qkv.stride(2),
+        rope.stride(0),
+        q_out.stride(0),
+        q_out.stride(1),
+        q_out.stride(2),
+        seq_len,
+        num_heads=num_heads,
+        num_kv_heads=num_kv_heads,
+        head_dim=head_dim,
+        rotary_half=head_dim // 2,
+        kv_repeat=num_heads // num_kv_heads,
+        head_tile=_HEAD_TILE,
+        eps=eps,
+        num_warps=4,
+    )
+    return q_out, k_out, v_out
+
+
+__all__ = ["fused_qkv_norm_rope", "fused_cuda_supported"]

vllm_omni/model_executor/models/omnivoice/omnivoice_generator.py

@@ -23,6 +23,7 @@ import torch.nn as nn
 import torch.nn.functional as F
 from vllm.logger import init_logger
 
+from vllm_omni.model_executor.models.omnivoice.fused_qkv_rope import fused_qkv_norm_rope
 from vllm_omni.transformers_utils.configs.omnivoice import OmniVoiceConfig
 
 logger = init_logger(__name__)
@@ -100,44 +101,46 @@ try:
 
     @triton.jit
     def _swiglu_fwd_kernel(
-        gate_ptr,
-        up_ptr,
+        inp_ptr,
         out_ptr,
-        stride,
+        in_stride,
+        out_stride,
         n_cols: tl.constexpr,
         BLOCK_SIZE: tl.constexpr,  # noqa: N803
     ):
         pid = tl.program_id(0).to(tl.int64)
-        gate_ptr += pid * stride
-        up_ptr += pid * stride
-        out_ptr += pid * stride
+        inp_ptr += pid * in_stride
+        out_ptr += pid * out_stride
         cols = tl.arange(0, BLOCK_SIZE)
         mask = cols < n_cols
-        gate = tl.load(gate_ptr + cols, mask=mask, other=0).to(tl.float32)
-        up = tl.load(up_ptr + cols, mask=mask, other=0)
+        gate = tl.load(inp_ptr + cols, mask=mask, other=0).to(tl.float32)
+        up = tl.load(inp_ptr + n_cols + cols, mask=mask, other=0)
         silu_gate = gate * tl.sigmoid(gate)
         out = silu_gate.cast(up.dtype) * up
         tl.store(out_ptr + cols, out, mask=mask)
 
-    def triton_swiglu(gate: torch.Tensor, up: torch.Tensor) -> torch.Tensor:
-        gate = gate.contiguous()
-        up = up.contiguous()
-        shape = gate.shape
-        n_cols = shape[-1]
-        g2d = gate.view(-1, n_cols)
-        u2d = up.view(-1, n_cols)
-        out = torch.empty_like(g2d)
+    def triton_swiglu(gate_up: torch.Tensor) -> torch.Tensor:
+        """SwiGLU over a packed ``[..., 2 * intermediate]`` activation.
+
+        Reading both halves straight out of the fused projection's output keeps
+        the packing free: splitting it first would hand this kernel two strided
+        views and cost a full copy of each half per layer per step.
+        """
+        n_cols = gate_up.shape[-1] // 2
+        gate_up = gate_up.contiguous()
+        x2d = gate_up.view(-1, 2 * n_cols)
+        out = torch.empty(x2d.shape[0], n_cols, dtype=gate_up.dtype, device=gate_up.device)
         BLOCK_SIZE, num_warps = _calculate_settings(n_cols)
-        _swiglu_fwd_kernel[(g2d.shape[0],)](
-            g2d,
-            u2d,
+        _swiglu_fwd_kernel[(x2d.shape[0],)](
+            x2d,
             out,
+            x2d.stride(0),
             out.stride(0),
             n_cols=n_cols,
             BLOCK_SIZE=BLOCK_SIZE,
             num_warps=num_warps,
         )
-        return out.view(*shape)
+        return out.view(*gate_up.shape[:-1], n_cols)
 
     @triton.jit
     def _fused_add_rms_norm_fwd_kernel(
@@ -267,9 +270,11 @@ class OmniVoiceAttention(nn.Module):
         self.num_kv_heads = config.llm_num_key_value_heads
         self.head_dim = config.llm_head_dim
 
-        self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False)
-        self.k_proj = nn.Linear(self.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
-        self.v_proj = nn.Linear(self.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
+        # q/k/v are packed into one projection: the three are sibling GEMMs over
+        # the same activation, and at this model's shapes (2 x 44 rows) three
+        # small GEMMs cost noticeably more than one wide one.
+        self.num_qkv_heads = self.num_heads + 2 * self.num_kv_heads
+        self.qkv_proj = nn.Linear(self.hidden_size, self.num_qkv_heads * self.head_dim, bias=False)
         self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)
 
         # Qwen3 uses per-head QK norm
@@ -281,40 +286,25 @@ class OmniVoiceAttention(nn.Module):
     def forward(
         self,
         hidden_states: torch.Tensor,
+        rope_table: torch.Tensor,
         attention_mask: torch.Tensor | None = None,
-        cos: torch.Tensor | None = None,
-        sin: torch.Tensor | None = None,
     ) -> torch.Tensor:
         batch_size, seq_len, _ = hidden_states.shape
 
-        q = self.q_proj(hidden_states)
-        k = self.k_proj(hidden_states)
-        v = self.v_proj(hidden_states)
+        qkv = self.qkv_proj(hidden_states).view(batch_size, seq_len, self.num_qkv_heads, self.head_dim)
 
-        q = q.view(batch_size, seq_len, self.num_heads, self.head_dim)
-        k = k.view(batch_size, seq_len, self.num_kv_heads, self.head_dim)
-        v = v.view(batch_size, seq_len, self.num_kv_heads, self.head_dim)
-
-        # Per-head QK norm (Qwen3)
-        q = self.q_norm(q)
-        k = self.k_norm(k)
-
-        # Apply RoPE
-        if cos is not None and sin is not None:
-            q = _apply_rotary_pos_emb(q, cos, sin)
-            k = _apply_rotary_pos_emb(k, cos, sin)
-
-        # Expand KV heads for GQA (8 KV heads → 16 Q heads)
-        if self.num_kv_heads != self.num_heads:
-            repeat_factor = self.num_heads // self.num_kv_heads
-            k = k.repeat_interleave(repeat_factor, dim=2)
-            v = v.repeat_interleave(repeat_factor, dim=2)
-
-        # Full bidirectional attention via SDPA with proper mask support
-        # Permute to (batch, heads, seq, head_dim) for SDPA
-        q = q.permute(0, 2, 1, 3)
-        k = k.permute(0, 2, 1, 3)
-        v = v.permute(0, 2, 1, 3)
+        # One kernel for the whole prologue: split the packed projection, RMSNorm
+        # Q and K per head, rotate both, broadcast K and V across their query
+        # groups, and emit SDPA's [batch, heads, positions, head_dim] layout.
+        q, k, v = fused_qkv_norm_rope(
+            qkv,
+            self.q_norm.weight,
+            self.k_norm.weight,
+            rope_table,
+            self.q_norm.eps,
+            self.num_heads,
+            self.num_kv_heads,
+        )
 
         # Caller passes a float mask; materialize float form if a bool slips through.
         sdpa_mask = attention_mask
@@ -340,14 +330,24 @@ class OmniVoiceMLP(nn.Module):
 
     def __init__(self, config: OmniVoiceConfig):
         super().__init__()
-        self.gate_proj = nn.Linear(config.llm_hidden_size, config.llm_intermediate_size, bias=False)
-        self.up_proj = nn.Linear(config.llm_hidden_size, config.llm_intermediate_size, bias=False)
-        self.down_proj = nn.Linear(config.llm_intermediate_size, config.llm_hidden_size, bias=False)
+        self.intermediate_size = config.llm_intermediate_size
+        self.gate_up_proj = nn.Linear(config.llm_hidden_size, 2 * self.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(self.intermediate_size, config.llm_hidden_size, bias=False)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
+        gate_up = self.gate_up_proj(x)
         if _TRITON_AVAILABLE:
-            return self.down_proj(triton_swiglu(self.gate_proj(x), self.up_proj(x)))
-        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
+            return self.down_proj(triton_swiglu(gate_up))
+        gate, up = gate_up.chunk(2, dim=-1)
+        return self.down_proj(F.silu(gate) * up)
+
+
+# Fused parameter -> the HF checkpoint shards it absorbs, in packing order.
+# The checkpoint keeps q/k/v and gate/up separate; load_weights packs them.
+_FUSED_PROJECTIONS: dict[str, tuple[str, ...]] = {
+    "self_attn.qkv_proj": ("self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj"),
+    "mlp.gate_up_proj": ("mlp.gate_proj", "mlp.up_proj"),
+}
 
 
 class OmniVoiceTransformerBlock(nn.Module):
@@ -364,12 +364,34 @@ class OmniVoiceTransformerBlock(nn.Module):
         self,
         hidden_states: torch.Tensor,
         attention_mask: torch.Tensor | None = None,
-        cos: torch.Tensor | None = None,
-        sin: torch.Tensor | None = None,
-    ) -> torch.Tensor:
-        residual = hidden_states
-        hidden_states = self.input_layernorm(hidden_states)
-        hidden_states = self.self_attn(hidden_states, attention_mask=attention_mask, cos=cos, sin=sin)
+        rope_table: torch.Tensor | None = None,
+        residual: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Returns ``(normed_hidden_states, residual)``.
+
+        The residual is threaded across the block boundary rather than being
+        added at the end. The MLP's residual add and the next block's input
+        RMSNorm are the same read of the same tensor, so handing the pending
+        residual on lets both happen in one fused kernel instead of a bare add
+        followed by a separate norm. Only the first block, which has no pending
+        residual, still pays for a standalone norm.
+        """
+        if residual is None:
+            residual = hidden_states
+            hidden_states = self.input_layernorm(hidden_states)
+        elif _TRITON_AVAILABLE:
+            hidden_states, residual = triton_fused_add_rms_norm(
+                hidden_states,
+                residual,
+                self.input_layernorm.weight,
+                self.input_layernorm.eps,
+            )
+        else:
+            hidden_states = residual + hidden_states
+            residual = hidden_states
+            hidden_states = self.input_layernorm(hidden_states)
+
+        hidden_states = self.self_attn(hidden_states, rope_table, attention_mask=attention_mask)
 
         if _TRITON_AVAILABLE:
             # Fused: (attn_out + residual) + RMSNorm in one kernel
@@ -384,9 +406,7 @@ class OmniVoiceTransformerBlock(nn.Module):
             residual = hidden_states
             hidden_states = self.post_attention_layernorm(hidden_states)
 
-        hidden_states = self.mlp(hidden_states)
-        hidden_states = residual + hidden_states
-        return hidden_states
+        return self.mlp(hidden_states), residual
 
 
 # ---------------------------------------------------------------------------
@@ -394,30 +414,22 @@ class OmniVoiceTransformerBlock(nn.Module):
 # ---------------------------------------------------------------------------
 
 
-def _precompute_rope(
+def _precompute_rope_table(
     head_dim: int,
     max_seq_len: int,
     theta: float = 1000000.0,
     device: torch.device | None = None,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    """Precompute RoPE cos/sin pre-doubled to full head_dim so the hot path skips a per-call cat."""
+) -> torch.Tensor:
+    """Precompute the packed ``[max_seq_len, head_dim]`` RoPE table.
+
+    Layout is what ``fused_qkv_norm_rope`` expects: the first half of each row
+    holds ``cos(theta)`` and the second half ``sin(theta)``, each of width
+    ``head_dim // 2``. Precomputing it keeps the hot path free of any cat.
+    """
     inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2, device=device, dtype=torch.float32) / head_dim))
     t = torch.arange(max_seq_len, device=device, dtype=torch.float32)
     freqs = torch.outer(t, inv_freq)
-    cos = torch.cat([freqs.cos(), freqs.cos()], dim=-1)
-    sin = torch.cat([freqs.sin(), freqs.sin()], dim=-1)
-    return cos, sin
-
-
-def _apply_rotary_pos_emb(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
-    """Apply rotary position embedding. x shape: (B, S, H, D). cos/sin are pre-doubled to width D."""
-    seq_len = x.shape[1]
-    cos = cos[:seq_len].unsqueeze(0).unsqueeze(2)  # (1, S, 1, D)
-    sin = sin[:seq_len].unsqueeze(0).unsqueeze(2)
-    x1 = x[..., : x.shape[-1] // 2]
-    x2 = x[..., x.shape[-1] // 2 :]
-    rotated = torch.cat([-x2, x1], dim=-1)
-    return x * cos + rotated * sin
+    return torch.cat([freqs.cos(), freqs.sin()], dim=-1)
 
 
 # ---------------------------------------------------------------------------
@@ -449,13 +461,17 @@ def _maybe_enable_tf32() -> None:
 # ---------------------------------------------------------------------------
 
 
-def _additive_float_mask(mask: torch.Tensor, dtype: torch.dtype = torch.float32) -> torch.Tensor:
+def _additive_float_mask(mask: torch.Tensor, dtype: torch.dtype) -> torch.Tensor:
     """Convert a boolean attention mask to its additive float form.
 
     ``True`` (attend) maps to ``0.0`` and ``False`` (masked) to ``-inf``. A bool
     mask must never be copied straight into a float buffer: the implicit cast
     maps True/False to 1.0/0.0, which leaves masked positions at 0.0 and so
     silently *unmasks* them.
+
+    ``dtype`` is required rather than defaulting to float32: SDPA rejects an
+    additive mask whose dtype differs from the query, so the mask has to follow
+    the model dtype and a default would just hide that coupling.
     """
     if mask.dtype != torch.bool:
         return mask
@@ -536,10 +552,7 @@ class _OmniVoiceCUDAGraphForward:
         _, bucket = key
         device = input_ids.device
 
-        self._gen._ensure_rope(bucket, device)
-        model_dtype = self._gen.text_embedding.weight.dtype
-        static_cos = self._gen._rope_cos[:bucket].to(device=device, dtype=model_dtype).contiguous()
-        static_sin = self._gen._rope_sin[:bucket].to(device=device, dtype=model_dtype).contiguous()
+        static_rope_table = self._gen._rope_table_for(bucket, device, self._gen.model_dtype)
 
         static_input_ids = input_ids.clone()
         static_audio_mask = audio_mask.clone()
@@ -550,8 +563,7 @@ class _OmniVoiceCUDAGraphForward:
                 static_input_ids,
                 static_audio_mask,
                 static_attn_mask,
-                static_cos,
-                static_sin,
+                static_rope_table,
             )
         torch.accelerator.synchronize(device)
 
@@ -569,8 +581,7 @@ class _OmniVoiceCUDAGraphForward:
                     static_input_ids,
                     static_audio_mask,
                     static_attn_mask,
-                    static_cos,
-                    static_sin,
+                    static_rope_table,
                 )
 
         entry = {
@@ -578,8 +589,7 @@ class _OmniVoiceCUDAGraphForward:
             "static_input_ids": static_input_ids,
             "static_audio_mask": static_audio_mask,
             "static_attn_mask": static_attn_mask,
-            "static_cos": static_cos,
-            "static_sin": static_sin,
+            "static_rope_table": static_rope_table,
             "static_output": static_output,
         }
         logger.info("OmniVoice CUDA Graph captured for key %s", key)
@@ -600,8 +610,9 @@ class _OmniVoiceCUDAGraphForward:
             key = (two_b, bucket)
             dummy_ids = torch.zeros(two_b, num_cb, bucket, dtype=torch.long, device=device)
             dummy_mask = torch.zeros(two_b, bucket, dtype=torch.bool, device=device)
-            # Capture with a float mask to match what forward() feeds at replay time.
-            dummy_attn = torch.zeros(two_b, 1, bucket, bucket, dtype=torch.float32, device=device)
+            # Capture with a float mask to match what forward() feeds at replay time,
+            # in the model dtype so replay can copy_ into it without a cast.
+            dummy_attn = torch.zeros(two_b, 1, bucket, bucket, dtype=self._gen.model_dtype, device=device)
             self._graphs[key] = self._capture_for_key(key, dummy_ids, dummy_mask, dummy_attn)
         logger.info("OmniVoice CUDA Graph warmup complete (%d graphs)", len(self._graphs))
 
@@ -612,12 +623,8 @@ class _OmniVoiceCUDAGraphForward:
         attention_mask: torch.Tensor | None,
     ) -> torch.Tensor:
         if torch.cuda.is_current_stream_capturing():
-            seq_len = input_ids.shape[-1]
-            self._gen._ensure_rope(seq_len, input_ids.device)
-            dtype = self._gen.text_embedding.weight.dtype
-            cos = self._gen._rope_cos[:seq_len].to(device=input_ids.device, dtype=dtype)
-            sin = self._gen._rope_sin[:seq_len].to(device=input_ids.device, dtype=dtype)
-            return self._gen._step_forward(input_ids, audio_mask, attention_mask, cos, sin)
+            rope_table = self._gen._rope_table_for(input_ids.shape[-1], input_ids.device, self._gen.model_dtype)
+            return self._gen._step_forward(input_ids, audio_mask, attention_mask, rope_table)
 
         seq_len = input_ids.shape[-1]
         two_b = input_ids.shape[0]
@@ -626,7 +633,7 @@ class _OmniVoiceCUDAGraphForward:
         # Graphs are captured with (and their static buffers hold) the additive
         # float mask, so normalize here, before padding or any copy_ into them.
         if attention_mask is not None:
-            attention_mask = _additive_float_mask(attention_mask)
+            attention_mask = _additive_float_mask(attention_mask, self._gen.model_dtype)
 
         if bucket is None:
             # Lazy capture: oversized sequence or non-unit batch (no pre-warmed bucket).
@@ -726,25 +733,38 @@ class OmniVoiceGenerator(nn.Module):
         )
 
         # Precompute RoPE
-        self._rope_cos = None
-        self._rope_sin = None
+        self._rope_table = None
 
         # CUDA Graph (bucket-size pre-capture; lazy fallback for oversized shapes)
         self._cuda_graph_fwd: _OmniVoiceCUDAGraphForward | None = (
             _OmniVoiceCUDAGraphForward(self, config.cuda_graph_capture_sizes) if config.enable_cuda_graph else None
         )
 
+    @property
+    def model_dtype(self) -> torch.dtype:
+        """The dtype every activation and mask in the generator has to match."""
+        return self.text_embedding.weight.dtype
+
     def _ensure_rope(self, seq_len: int, device: torch.device) -> None:
-        """Lazily compute RoPE cos/sin if needed."""
-        if self._rope_cos is None or self._rope_cos.shape[0] < seq_len:
+        """Lazily compute the packed RoPE table if needed."""
+        if self._rope_table is None or self._rope_table.shape[0] < seq_len:
             max_len = max(seq_len, 4096)
-            self._rope_cos, self._rope_sin = _precompute_rope(
+            self._rope_table = _precompute_rope_table(
                 self.config.llm_head_dim,
                 max_len,
                 theta=self.config.llm_rope_theta,
                 device=device,
             )
 
+    def _rope_table_for(self, seq_len: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
+        """The packed [seq_len, head_dim] table the fused prologue indexes.
+
+        It depends only on the bucket and dtype, so callers build it once per
+        request or per captured graph, never per layer.
+        """
+        self._ensure_rope(seq_len, device)
+        return self._rope_table[:seq_len].to(device=device, dtype=dtype).contiguous()
+
     def _prepare_embeddings(
         self,
         input_ids: torch.Tensor,
@@ -781,27 +801,21 @@ class OmniVoiceGenerator(nn.Module):
         self,
         inputs_embeds: torch.Tensor,
         attention_mask: torch.Tensor | None = None,
-        cos: torch.Tensor | None = None,
-        sin: torch.Tensor | None = None,
+        rope_table: torch.Tensor | None = None,
     ) -> torch.Tensor:
         """Run through transformer layers.
 
         Args:
             inputs_embeds: [B, S, hidden_size]
             attention_mask: [B, 1, S, S] or None
-            cos, sin: optional precomputed RoPE tables in target dtype
+            rope_table: optional precomputed [B * S, head_dim] RoPE table
 
         Returns:
             hidden_states: [B, S, hidden_size]
         """
-        device = inputs_embeds.device
-        seq_len = inputs_embeds.shape[1]
-        self._ensure_rope(seq_len, device)
-
         hidden_states = inputs_embeds
-        if cos is None or sin is None:
-            cos = self._rope_cos.to(device=device, dtype=hidden_states.dtype)
-            sin = self._rope_sin.to(device=device, dtype=hidden_states.dtype)
+        if rope_table is None:
+            rope_table = self._rope_table_for(inputs_embeds.shape[1], inputs_embeds.device, hidden_states.dtype)
 
         # Safety: convert bool mask if caller hasn't (e.g. external paths beyond forward()).
         if attention_mask is not None and attention_mask.dtype == torch.bool:
@@ -809,15 +823,16 @@ class OmniVoiceGenerator(nn.Module):
                 ~attention_mask, float("-inf")
             )
 
+        residual = None
         for layer in self.layers:
-            hidden_states = layer(
+            hidden_states, residual = layer(
                 hidden_states,
                 attention_mask=attention_mask,
-                cos=cos,
-                sin=sin,
+                rope_table=rope_table,
+                residual=residual,
             )
 
-        return self.norm(hidden_states)
+        return self.norm(hidden_states + residual)
 
     def _get_logits(self, hidden_states: torch.Tensor) -> torch.Tensor:
         """Project hidden states to per-codebook logits.
@@ -842,14 +857,16 @@ class OmniVoiceGenerator(nn.Module):
         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:
-        """Single unmasking-step forward using pre-cast RoPE tensors (CUDA graph safe)."""
+        """Single unmasking-step forward using a pre-cast RoPE table (CUDA graph safe)."""
         hidden_states = self._prepare_embeddings(input_ids, audio_mask)
+        residual = None
         for layer in self.layers:
-            hidden_states = layer(hidden_states, attention_mask=attention_mask, cos=cos, sin=sin)
-        return self._get_logits(self.norm(hidden_states))
+            hidden_states, residual = layer(
+                hidden_states, attention_mask=attention_mask, rope_table=rope_table, residual=residual
+            )
+        return self._get_logits(self.norm(hidden_states + residual))
 
     @torch.inference_mode()
     def forward(
@@ -926,19 +943,14 @@ class OmniVoiceGenerator(nn.Module):
         c_lens = attention_mask[:B, 0, 0].sum(dim=-1).tolist()
 
         # Materialize the SDPA float mask once so the captured graph (and eager path) skip per-layer conversion.
-        sdpa_attn_mask = torch.zeros_like(attention_mask, dtype=torch.float32).masked_fill_(
-            ~attention_mask, float("-inf")
-        )
+        sdpa_attn_mask = _additive_float_mask(attention_mask, self.model_dtype)
 
         use_cuda_graph = self._cuda_graph_fwd is not None and input_ids.is_cuda
         if not use_cuda_graph:
             # Eager-path-only constants (the cuda-graph captures its own).
             text_embeds_cached = self.text_embedding(input_ids[:, 0, :])
             audio_mask_3d = audio_mask.unsqueeze(-1)
-            self._ensure_rope(input_ids.shape[-1], device)
-            target_dtype = text_embeds_cached.dtype
-            cos = self._rope_cos.to(device=device, dtype=target_dtype)
-            sin = self._rope_sin.to(device=device, dtype=target_dtype)
+            rope_table = self._rope_table_for(input_ids.shape[-1], device, text_embeds_cached.dtype)
 
         # Main iterative loop
         for step in range(num_step):
@@ -946,11 +958,11 @@ class OmniVoiceGenerator(nn.Module):
                 # Float mask skips per-layer conversion; fp32 cast deferred to the per-item slices below.
                 batch_logits = self._cuda_graph_fwd(input_ids, audio_mask, sdpa_attn_mask)
             else:
-                # Eager fallback reuses hoisted constants (text embeds, sdpa mask, cos/sin).
+                # Eager fallback reuses hoisted constants (text embeds, sdpa mask, rope table).
                 inputs_embeds = self._prepare_embeddings(
                     input_ids, audio_mask, text_embeds=text_embeds_cached, audio_mask_3d=audio_mask_3d
                 )
-                hidden_states = self._transformer_forward(inputs_embeds, sdpa_attn_mask, cos=cos, sin=sin)
+                hidden_states = self._transformer_forward(inputs_embeds, sdpa_attn_mask, rope_table=rope_table)
                 # fp32 cast deferred to the per-item slices below.
                 batch_logits = self._get_logits(hidden_states)
             # batch_logits: [2*B, 8, S, 1025]
@@ -1014,6 +1026,52 @@ class OmniVoiceGenerator(nn.Module):
 
         return tokens
 
+    def _load_fused_projections(self, state_dict: dict[str, torch.Tensor]) -> set[str]:
+        """Pack the checkpoint's separate q/k/v and gate/up shards into the fused params.
+
+        This has to happen explicitly. The generic per-tensor path below looks
+        the destination up by name, and ``q_proj``/``gate_proj`` no longer exist
+        as modules -- so it would find nothing, log a warning nobody reads, and
+        leave the fused parameters at their random initialization. A corrupted
+        model that still answers requests is worse than a failed load, so a
+        missing or wrong-shaped shard raises here.
+        """
+        loaded: set[str] = set()
+        packed_params = 0
+        for idx, layer in enumerate(self.layers):
+            for fused_path, shard_paths in _FUSED_PROJECTIONS.items():
+                keys = [f"llm.layers.{idx}.{path}.weight" for path in shard_paths]
+                missing = [k for k in keys if k not in state_dict]
+                if len(missing) == len(keys):
+                    continue
+                if missing:
+                    raise ValueError(
+                        f"OmniVoice checkpoint is missing {missing} needed to build "
+                        f"layers.{idx}.{fused_path}; refusing to load a partially "
+                        f"initialized fused projection."
+                    )
+                module = layer
+                for part in fused_path.split("."):
+                    module = getattr(module, part)
+                packed = torch.cat([state_dict[k] for k in keys], dim=0)
+                if packed.shape != module.weight.shape:
+                    raise ValueError(
+                        f"OmniVoice checkpoint shards {keys} pack to {tuple(packed.shape)} "
+                        f"but layers.{idx}.{fused_path}.weight is {tuple(module.weight.shape)}."
+                    )
+                module.weight.data.copy_(packed)
+                loaded.update(keys)
+                packed_params += 1
+
+        expected = len(self.layers) * len(_FUSED_PROJECTIONS)
+        if loaded and packed_params != expected:
+            raise ValueError(
+                f"OmniVoice checkpoint filled {packed_params}/{expected} fused projections; "
+                f"the remaining ones would stay randomly initialized."
+            )
+        logger.info("Generator: packed %d/%d fused projections", packed_params, expected)
+        return loaded
+
     def load_weights(self, model_dir: str, device: torch.device) -> None:
         """Load weights from HuggingFace OmniVoice model.safetensors.
 
@@ -1053,8 +1111,13 @@ class OmniVoiceGenerator(nn.Module):
                 self.audio_heads.weight.data.copy_(state_dict[key])
                 loaded_keys.add(key)
 
-        # 4. Transformer layers: llm.layers.N.* -> layers.N.*
+        # 4a. Fused projections, packed from their separate checkpoint shards.
+        loaded_keys |= self._load_fused_projections(state_dict)
+
+        # 4b. Remaining transformer weights: llm.layers.N.* -> layers.N.*
         for key, value in state_dict.items():
+            if key in loaded_keys:
+                continue
             if key.startswith("llm.layers."):
                 # llm.layers.0.self_attn.q_proj.weight -> layers.0.self_attn.q_proj.weight
                 our_key = key.replace("llm.layers.", "layers.")