From ca6b4e21127ff3f36530438a42948fe17117d412 Mon Sep 17 00:00:00 2001 From: DanaerLee <102979255+MrlixiangWE@users.noreply.github.com> Date: Thu, 3 Sep 2026 03:54:55 +0800 Subject: [PATCH] [Perf][Bugfix][OmniVoice] Restore float16 serving, and fuse the generator hot loop (#6317) Signed-off-by: MrlixiangWE --- benchmarks/tts/model_configs.yaml | 11 + .../omnivoice/test_cuda_graph_generator.py | 35 +- .../omnivoice/test_fused_projection_load.py | 183 ++++++++++ .../models/omnivoice/test_fused_qkv_rope.py | 166 +++++++++ .../models/omnivoice/test_mask_dtype.py | 183 ++++++++++ .../models/omnivoice/test_triton_kernels.py | 11 +- .../models/omnivoice/fused_qkv_rope.py | 273 ++++++++++++++ .../models/omnivoice/omnivoice_generator.py | 337 +++++++++++------- 8 files changed, 1039 insertions(+), 160 deletions(-) create mode 100644 tests/model_executor/models/omnivoice/test_fused_projection_load.py create mode 100644 tests/model_executor/models/omnivoice/test_fused_qkv_rope.py create mode 100644 tests/model_executor/models/omnivoice/test_mask_dtype.py create mode 100644 vllm_omni/model_executor/models/omnivoice/fused_qkv_rope.py diff --git a/benchmarks/tts/model_configs.yaml b/benchmarks/tts/model_configs.yaml index 5450e8a7c..7cdf47eec 100644 --- a/benchmarks/tts/model_configs.yaml +++ b/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: {} diff --git a/tests/model_executor/models/omnivoice/test_cuda_graph_generator.py b/tests/model_executor/models/omnivoice/test_cuda_graph_generator.py index dc6126713..b55f8279c 100644 --- a/tests/model_executor/models/omnivoice/test_cuda_graph_generator.py +++ b/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) diff --git a/tests/model_executor/models/omnivoice/test_fused_projection_load.py b/tests/model_executor/models/omnivoice/test_fused_projection_load.py new file mode 100644 index 000000000..fc5db01c7 --- /dev/null +++ b/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) diff --git a/tests/model_executor/models/omnivoice/test_fused_qkv_rope.py b/tests/model_executor/models/omnivoice/test_fused_qkv_rope.py new file mode 100644 index 000000000..73a4bccc5 --- /dev/null +++ b/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) diff --git a/tests/model_executor/models/omnivoice/test_mask_dtype.py b/tests/model_executor/models/omnivoice/test_mask_dtype.py new file mode 100644 index 000000000..ff899e2a2 --- /dev/null +++ b/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 diff --git a/tests/model_executor/models/omnivoice/test_triton_kernels.py b/tests/model_executor/models/omnivoice/test_triton_kernels.py index 1598b0f74..b71c7c967 100644 --- a/tests/model_executor/models/omnivoice/test_triton_kernels.py +++ b/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 diff --git a/vllm_omni/model_executor/models/omnivoice/fused_qkv_rope.py b/vllm_omni/model_executor/models/omnivoice/fused_qkv_rope.py new file mode 100644 index 000000000..000a2f569 --- /dev/null +++ b/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"] diff --git a/vllm_omni/model_executor/models/omnivoice/omnivoice_generator.py b/vllm_omni/model_executor/models/omnivoice/omnivoice_generator.py index 087cba96f..a75c4e6f7 100644 --- a/vllm_omni/model_executor/models/omnivoice/omnivoice_generator.py +++ b/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.")