[Perf] [OmniVoice] Triton kernel fusion + CUDA Graph acceleration (#3336)

Signed-off-by: univa-HARRY <harry@univa.co.kr>
Signed-off-by: Harry <harry@univa.co.kr>
Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>

tests/model_executor/models/omnivoice/test_cuda_graph_generator.py

@@ -0,0 +1,263 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Tests for OmniVoice CUDA Graph generator wrapper numerical equivalence.
+
+Verifies that _OmniVoiceCUDAGraphForward produces results equivalent to
+eager mode across three scenarios:
+  - Exact-size inputs (no padding) → bit-identical
+  - Padded inputs (padded to nearest bucket) → correct slicing, exact match
+    for position-independent models
+  - Oversized inputs (fallback to eager) → bit-identical
+
+Uses a lightweight synthetic generator to keep warmup fast and avoid
+loading actual model weights.
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+import torch.nn as nn
+
+pytestmark = [
+    pytest.mark.core_model,
+    pytest.mark.cuda,
+    pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"),
+]
+
+DEVICE = torch.device("cuda:0")
+NUM_CB = 8
+VOCAB = 1025
+HIDDEN = 64
+CAPTURE_SIZES = [32, 64, 128]
+
+
+# ---------------------------------------------------------------------------
+# Synthetic generator: matches _OmniVoiceCUDAGraphForward's interface
+# ---------------------------------------------------------------------------
+
+
+class _SyntheticGenerator:
+    """Position-independent Embedding+Linear replacing OmniVoiceGenerator's 28-layer transformer.
+
+    Each token position is processed independently (no attention), so padded
+    runs produce bit-identical results to non-padded runs on the same positions.
+    This makes the padding/slicing logic easy to verify exactly.
+    """
+
+    class _Cfg:
+        num_audio_codebook = NUM_CB
+
+    def __init__(self, device: torch.device):
+        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
+
+    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)
+
+    def _step_forward(
+        self,
+        input_ids: torch.Tensor,
+        audio_mask: torch.Tensor,
+        attention_mask: torch.Tensor | None,
+        cos: torch.Tensor,
+        sin: torch.Tensor,
+    ) -> torch.Tensor:
+        two_b, _, S = input_ids.shape
+        x = self.text_embedding(input_ids[:, 0, :].clamp(0, 999))  # [two_b, S, H]
+        logits = self._linear(x)  # [two_b, S, 8*1025]
+        return logits.view(two_b, S, NUM_CB, VOCAB).permute(0, 2, 1, 3)  # [two_b, 8, S, 1025]
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+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)
+
+
+def _make_inputs(seq_len: int, device: torch.device = DEVICE):
+    two_b = 2
+    ids = torch.randint(0, 100, (two_b, NUM_CB, seq_len), dtype=torch.long, device=device)
+    mask = torch.ones(two_b, seq_len, dtype=torch.bool, device=device)
+    attn = torch.ones(two_b, 1, seq_len, seq_len, dtype=torch.bool, device=device)
+    return ids, mask, attn
+
+
+# ---------------------------------------------------------------------------
+# Fixtures
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture(scope="module")
+def gen():
+    torch.manual_seed(42)
+    return _SyntheticGenerator(DEVICE)
+
+
+@pytest.fixture(scope="module")
+def wrapper(gen):
+    from vllm_omni.model_executor.models.omnivoice.omnivoice_generator import (
+        _OmniVoiceCUDAGraphForward,
+    )
+
+    w = _OmniVoiceCUDAGraphForward(gen, capture_sizes=CAPTURE_SIZES)
+    w.warmup(DEVICE)
+    return w
+
+
+# ---------------------------------------------------------------------------
+# 1. Exact-size inputs → bit-identical
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("seq_len", CAPTURE_SIZES)
+def test_exact_size_bit_identical(gen, wrapper, seq_len):
+    """When input exactly matches a captured bucket, output must be bit-identical to eager."""
+    ids, mask, attn = _make_inputs(seq_len)
+    with torch.no_grad():
+        eager_out = _eager(gen, ids, mask, attn)
+        graph_out = wrapper(ids, mask, attn)
+    torch.testing.assert_close(graph_out, eager_out, atol=0, rtol=0)
+
+
+# ---------------------------------------------------------------------------
+# 2. Padded inputs → correct shape, values match at actual positions
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("seq_len", [1, 15, 33, 60, 100])
+def test_padded_output_shape(gen, wrapper, seq_len):
+    """Graph output must be sliced back to actual seq_len, not the bucket size."""
+    ids, mask, attn = _make_inputs(seq_len)
+    with torch.no_grad():
+        graph_out = wrapper(ids, mask, attn)
+    assert graph_out.shape == (2, NUM_CB, seq_len, VOCAB)
+
+
+@pytest.mark.parametrize("seq_len", [15, 33, 60, 100])
+def test_padded_output_matches_eager(gen, wrapper, seq_len):
+    """Padded graph output must equal eager output at actual positions.
+
+    The synthetic model has no attention across positions, so zero-padding
+    does not affect non-padded positions — exact match is expected.
+    """
+    ids, mask, attn = _make_inputs(seq_len)
+    with torch.no_grad():
+        eager_out = _eager(gen, ids, mask, attn)
+        graph_out = wrapper(ids, mask, attn)
+    torch.testing.assert_close(graph_out, eager_out, atol=0, rtol=0)
+
+
+# ---------------------------------------------------------------------------
+# 3. Oversized inputs → fallback to eager (lazy capture), bit-identical
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("seq_len", [129, 200, 256])
+def test_fallback_eager_bit_identical(gen, wrapper, seq_len):
+    """Sequences exceeding the largest bucket fall back to lazy capture → bit-identical."""
+    ids, mask, attn = _make_inputs(seq_len)
+    with torch.no_grad():
+        eager_out = _eager(gen, ids, mask, attn)
+        graph_out = wrapper(ids, mask, attn)
+    torch.testing.assert_close(graph_out, eager_out, atol=0, rtol=0)
+
+
+# ---------------------------------------------------------------------------
+# 4. Determinism
+# ---------------------------------------------------------------------------
+
+
+def test_deterministic_across_calls(wrapper):
+    """Same input must produce identical output on repeated CUDA graph replays."""
+    ids, mask, attn = _make_inputs(32)
+    with torch.no_grad():
+        out1 = wrapper(ids, mask, attn)
+        out2 = wrapper(ids, mask, attn)
+    torch.testing.assert_close(out1, out2, atol=0, rtol=0)
+
+
+# ---------------------------------------------------------------------------
+# 5. _find_bucket logic (CPU, no CUDA graph)
+# ---------------------------------------------------------------------------
+
+
+def test_find_bucket_returns_nearest_bucket():
+    """_find_bucket must return the smallest bucket >= seq_len, or None if all are smaller."""
+    from vllm_omni.model_executor.models.omnivoice.omnivoice_generator import (
+        _OmniVoiceCUDAGraphForward,
+    )
+
+    w = _OmniVoiceCUDAGraphForward.__new__(_OmniVoiceCUDAGraphForward)
+    w._capture_sizes = [32, 64, 128]
+    w._graphs = {}
+
+    assert w._find_bucket(1) == 32
+    assert w._find_bucket(32) == 32
+    assert w._find_bucket(33) == 64
+    assert w._find_bucket(64) == 64
+    assert w._find_bucket(128) == 128
+    assert w._find_bucket(129) is None
+    assert w._find_bucket(1000) is None
+
+
+# ---------------------------------------------------------------------------
+# 6. enable_cuda_graph=False produces same tokens as enable_cuda_graph=True
+# ---------------------------------------------------------------------------
+
+
+def test_cuda_graph_disabled_matches_eager_generator():
+    """OmniVoiceGenerator with enable_cuda_graph=False must produce the same
+    _step_forward output as one with enable_cuda_graph=True (after warmup).
+
+    Uses a 2-layer config so warmup completes quickly without loading weights.
+    """
+    from vllm_omni.model_executor.models.omnivoice.omnivoice_generator import OmniVoiceGenerator
+    from vllm_omni.transformers_utils.configs.omnivoice import OmniVoiceConfig
+
+    cfg = OmniVoiceConfig()
+    cfg.llm_num_hidden_layers = 2
+    cfg.num_hidden_layers = 2
+
+    torch.manual_seed(0)
+    gen_eager = OmniVoiceGenerator(cfg).to(DEVICE).eval()
+    gen_eager.config.enable_cuda_graph = False
+    gen_eager._cuda_graph_fwd = None
+
+    torch.manual_seed(0)
+    gen_graph = OmniVoiceGenerator(cfg).to(DEVICE).eval()
+    gen_graph.config.cuda_graph_capture_sizes = [64]
+    from vllm_omni.model_executor.models.omnivoice.omnivoice_generator import _OmniVoiceCUDAGraphForward
+
+    gen_graph._cuda_graph_fwd = _OmniVoiceCUDAGraphForward(gen_graph, capture_sizes=[64])
+    gen_graph._cuda_graph_fwd.warmup(DEVICE)
+
+    seq_len = 64
+    ids = torch.zeros(2, cfg.num_audio_codebook, seq_len, dtype=torch.long, device=DEVICE)
+    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)
+
+    with torch.no_grad():
+        eager_logits = gen_eager._step_forward(ids, mask, attn, cos, sin)
+        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_decoder.py

@@ -0,0 +1,214 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Correctness and regression tests for OmniVoiceDecoder.
+
+Covers three areas:
+1. NaN regression — verifies the fp16 overflow fix:
+     acoustic decoder must operate in float32 even when fc2 weights are fp16.
+     Demonstrates that the pre-fix code path (no .float() after fc2) produces
+     NaN/Inf with large-scale weights, while the fixed path does not.
+2. Output shape / value sanity — correct [B, 1, T*upsample] shape, finite values.
+3. Acoustic decoder float32 invariant — weights must be float32 after build.
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+import torch.nn as nn
+
+from vllm_omni.model_executor.models.omnivoice.omnivoice_decoder import (
+    HiggsAudioRVQ,
+    OmniVoiceDecoder,
+)
+from vllm_omni.transformers_utils.configs.omnivoice import OmniVoiceConfig
+
+pytestmark = [
+    pytest.mark.core_model,
+    pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"),
+]
+
+DEVICE = torch.device("cuda:0")
+UPSAMPLE = 4  # synthetic acoustic decoder stride (real DAC uses 960)
+T_FRAMES = 25  # number of audio token frames to decode
+
+
+# ---------------------------------------------------------------------------
+# Tiny acoustic decoder: ConvTranspose1d upsample by UPSAMPLE
+# ---------------------------------------------------------------------------
+
+
+class _TinyAcousticDecoder(nn.Module):
+    def __init__(self, in_channels: int = 256, upsample: int = UPSAMPLE):
+        super().__init__()
+        self.conv = nn.ConvTranspose1d(in_channels, 1, kernel_size=upsample, stride=upsample)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.conv(x)  # [B, 1, T*upsample]
+
+
+# ---------------------------------------------------------------------------
+# Builder
+# ---------------------------------------------------------------------------
+
+
+def _build_decoder(
+    *,
+    fc2_dtype: torch.dtype = torch.float32,
+    acoustic_dtype: torch.dtype = torch.float32,
+    weight_scale: float = 1.0,
+) -> OmniVoiceDecoder:
+    """Build an OmniVoiceDecoder with synthetic components (no model checkpoint required).
+
+    Args:
+        fc2_dtype: dtype for the fc2 linear layer weights.
+        acoustic_dtype: dtype for the acoustic decoder weights.
+        weight_scale: constant weight value for overflow testing.
+    """
+    config = OmniVoiceConfig()
+    decoder = OmniVoiceDecoder(config)
+
+    decoder.quantizer = HiggsAudioRVQ(num_quantizers=8, codebook_size=1024, codebook_dim=64, hidden_size=1024).to(
+        DEVICE
+    )
+
+    decoder.fc2 = nn.Linear(1024, 256).to(DEVICE).to(fc2_dtype)
+
+    acoustic = _TinyAcousticDecoder(in_channels=256, upsample=UPSAMPLE).to(DEVICE).to(acoustic_dtype)
+    if weight_scale != 1.0:
+        with torch.no_grad():
+            nn.init.constant_(acoustic.conv.weight, weight_scale)
+            nn.init.constant_(acoustic.conv.bias, weight_scale)
+    decoder.acoustic_decoder = acoustic
+
+    decoder._loaded = True
+    return decoder
+
+
+# ---------------------------------------------------------------------------
+# 1. NaN regression
+# ---------------------------------------------------------------------------
+
+
+def test_no_nan_fp16_fc2_float32_acoustic():
+    """fp16 fc2 + float32 acoustic decoder (the fixed path) must not produce NaN or Inf."""
+    decoder = _build_decoder(fc2_dtype=torch.float16, acoustic_dtype=torch.float32)
+    tokens = torch.randint(0, 1024, (1, 8, T_FRAMES), device=DEVICE)
+
+    audio = decoder(tokens)
+
+    assert not torch.isnan(audio).any(), "NaN in output — fp16 overflow may have reached acoustic decoder"
+    assert not torch.isinf(audio).any(), "Inf in output"
+
+
+def test_fp16_acoustic_decoder_produces_nan_without_fix():
+    """Pre-fix scenario: fp16 acoustic decoder with large weights overflows to NaN/Inf.
+
+    This test documents the original bug and will fail (raising AssertionError)
+    if the underlying hardware / Triton version somehow avoids overflow, but
+    with weights scaled to 300.0 the ConvTranspose1d intermediate activations
+    reliably exceed fp16 max (~65504) in practice.
+    """
+    decoder = _build_decoder(fc2_dtype=torch.float16, acoustic_dtype=torch.float16, weight_scale=300.0)
+    with torch.no_grad():
+        nn.init.constant_(decoder.fc2.weight, 300.0)
+        nn.init.constant_(decoder.fc2.bias, 300.0)
+        for q in decoder.quantizer.quantizers:
+            nn.init.constant_(q.codebook.weight, 300.0)
+
+    tokens = torch.randint(0, 1024, (1, 8, T_FRAMES), device=DEVICE)
+
+    # Simulate pre-fix forward: no .float() after fc2
+    codes = tokens.transpose(0, 1).long()
+    quantized = decoder.quantizer.decode(codes)
+    quantized_fp16 = decoder.fc2(quantized.transpose(1, 2).to(decoder.fc2.weight.dtype)).transpose(
+        1, 2
+    )  # stays fp16, no .float()
+    audio_broken = decoder.acoustic_decoder(quantized_fp16)
+
+    assert torch.isnan(audio_broken).any() or torch.isinf(audio_broken).any(), (
+        "Expected NaN/Inf in pre-fix fp16 path with large weights — test setup may need a higher weight_scale"
+    )
+
+
+def test_acoustic_decoder_receives_float32_input():
+    """The .float() cast after fc2 must ensure acoustic decoder always gets float32 input,
+    regardless of fc2 weight dtype."""
+    decoder = _build_decoder(fc2_dtype=torch.float16, acoustic_dtype=torch.float32)
+
+    captured_dtypes: list[torch.dtype] = []
+    original_fwd = decoder.acoustic_decoder.forward
+
+    def _hook(x: torch.Tensor) -> torch.Tensor:
+        captured_dtypes.append(x.dtype)
+        return original_fwd(x)
+
+    decoder.acoustic_decoder.forward = _hook  # type: ignore[method-assign]
+    tokens = torch.randint(0, 1024, (1, 8, T_FRAMES), device=DEVICE)
+    decoder(tokens)
+
+    assert captured_dtypes, "acoustic_decoder.forward was never called"
+    assert captured_dtypes[0] == torch.float32, (
+        f"Acoustic decoder received {captured_dtypes[0]}, expected torch.float32 — "
+        "the .float() upcast after fc2 may be missing"
+    )
+
+
+# ---------------------------------------------------------------------------
+# 2. Output shape and value sanity
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("batch,t_frames", [(1, 10), (1, 50), (2, 25)])
+def test_output_shape(batch, t_frames):
+    """Decoder must return [B, 1, T*upsample] waveform tensor."""
+    decoder = _build_decoder()
+    tokens = torch.randint(0, 1024, (batch, 8, t_frames), device=DEVICE)
+
+    audio = decoder(tokens)
+
+    assert audio.shape == (batch, 1, t_frames * UPSAMPLE), (
+        f"Expected ({batch}, 1, {t_frames * UPSAMPLE}), got {audio.shape}"
+    )
+
+
+def test_no_nan_float32_baseline():
+    """Float32 fc2 + float32 acoustic decoder must always produce finite output."""
+    decoder = _build_decoder(fc2_dtype=torch.float32, acoustic_dtype=torch.float32)
+    tokens = torch.randint(0, 1024, (1, 8, T_FRAMES), device=DEVICE)
+
+    audio = decoder(tokens)
+
+    assert not torch.isnan(audio).any()
+    assert not torch.isinf(audio).any()
+
+
+def test_output_is_3d():
+    """Output must be 3-D: [B, 1, samples]. OmniVoiceDecoder.forward unsqueezes if needed."""
+    decoder = _build_decoder()
+    tokens = torch.randint(0, 1024, (1, 8, 10), device=DEVICE)
+
+    audio = decoder(tokens)
+
+    assert audio.dim() == 3, f"Expected 3-D output, got {audio.dim()}-D"
+    assert audio.shape[1] == 1, f"Expected channel dim=1, got {audio.shape[1]}"
+
+
+# ---------------------------------------------------------------------------
+# 3. Acoustic decoder float32 invariant
+# ---------------------------------------------------------------------------
+
+
+def test_acoustic_decoder_weights_are_float32():
+    """Acoustic decoder weights must be float32 — any fp16 parameter is a regression indicator.
+
+    In production, load_weights() calls self.acoustic_decoder.float() to enforce this.
+    The synthetic builder mirrors that invariant.
+    """
+    decoder = _build_decoder(fc2_dtype=torch.float16, acoustic_dtype=torch.float32)
+
+    for name, param in decoder.acoustic_decoder.named_parameters():
+        assert param.dtype == torch.float32, (
+            f"acoustic_decoder.{name} is {param.dtype} — must be float32 (see load_weights fix)"
+        )

tests/model_executor/models/omnivoice/test_triton_kernels.py

@@ -0,0 +1,200 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Correctness tests for OmniVoice Triton kernels.
+
+Verifies that RMSNorm, SwiGLU, and fused-add+RMSNorm Triton kernels
+produce results numerically equivalent to the PyTorch reference.
+All computations follow the kernel's internal float32 promotion rules.
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+pytestmark = [
+    pytest.mark.core_model,
+    pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"),
+]
+
+DEVICE = torch.device("cuda:0")
+EPS = 1e-6
+
+try:
+    from vllm_omni.model_executor.models.omnivoice.omnivoice_generator import (
+        _TRITON_AVAILABLE,
+        triton_fused_add_rms_norm,
+        triton_rms_norm,
+        triton_swiglu,
+    )
+except ImportError:
+    _TRITON_AVAILABLE = False
+
+triton_available = pytest.mark.skipif(not _TRITON_AVAILABLE, reason="Triton not available on this platform")
+
+
+# ---------------------------------------------------------------------------
+# PyTorch reference implementations (mirror the kernel's internal arithmetic)
+# ---------------------------------------------------------------------------
+
+
+def _ref_rms_norm(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor:
+    # Kernel promotes x to float32, computes RMS, casts back to x.dtype, then mul weight
+    x_f32 = x.float()
+    ms = x_f32.pow(2).mean(-1, keepdim=True)
+    normed = (x_f32 * torch.rsqrt(ms + eps)).to(x.dtype)
+    return normed * weight.to(normed.dtype)
+
+
+def _ref_swiglu(gate: torch.Tensor, up: torch.Tensor) -> torch.Tensor:
+    # Kernel: gate → float32, silu, cast to up.dtype, mul up
+    gate_f32 = gate.float()
+    silu = gate_f32 * torch.sigmoid(gate_f32)
+    return silu.to(up.dtype) * up
+
+
+def _ref_fused_add_rms_norm(
+    x: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor, eps: float
+) -> tuple[torch.Tensor, torch.Tensor]:
+    # Kernel: s_f32 = x+r in float32, s stored as x.dtype; y = norm(s_f32)*weight
+    s_f32 = x.float() + residual.float()
+    s = s_f32.to(x.dtype)
+    ms = s_f32.pow(2).mean(-1, keepdim=True)
+    normed = (s_f32 * torch.rsqrt(ms + eps)).to(x.dtype)
+    y = normed * weight.to(normed.dtype)
+    return y, s
+
+
+# ---------------------------------------------------------------------------
+# RMSNorm
+# ---------------------------------------------------------------------------
+
+
+@triton_available
+@pytest.mark.parametrize("rows,cols", [(1, 64), (4, 512), (16, 1024), (32, 3072)])
+def test_rms_norm_float32(rows, cols):
+    torch.manual_seed(0)
+    x = torch.randn(rows, cols, device=DEVICE)
+    w = torch.ones(cols, device=DEVICE)
+
+    ref = _ref_rms_norm(x, w, EPS)
+    out = triton_rms_norm(x, w, EPS)
+
+    torch.testing.assert_close(out, ref, atol=1e-5, rtol=1e-5)
+
+
+@triton_available
+@pytest.mark.parametrize("shape", [(2, 16, 64), (4, 32, 512)])
+def test_rms_norm_3d_input(shape):
+    """Triton wrapper must handle 3-D [B, S, H] inputs by flattening and restoring shape."""
+    torch.manual_seed(1)
+    hidden = shape[-1]
+    x = torch.randn(*shape, device=DEVICE)
+    w = torch.ones(hidden, device=DEVICE)
+
+    ref = _ref_rms_norm(x, w, EPS)
+    out = triton_rms_norm(x, w, EPS)
+
+    torch.testing.assert_close(out, ref, atol=1e-5, rtol=1e-5)
+    assert out.shape == x.shape
+
+
+@triton_available
+def test_rms_norm_no_nan_on_zero_input():
+    """RMSNorm of the zero vector must not produce NaN (eps prevents divide-by-zero)."""
+    x = torch.zeros(4, 256, device=DEVICE)
+    w = torch.ones(256, device=DEVICE)
+
+    out = triton_rms_norm(x, w, EPS)
+
+    assert not torch.isnan(out).any()
+
+
+# ---------------------------------------------------------------------------
+# SwiGLU
+# ---------------------------------------------------------------------------
+
+
+@triton_available
+@pytest.mark.parametrize("rows,cols", [(1, 64), (4, 512), (16, 3072)])
+def test_swiglu_float32(rows, cols):
+    torch.manual_seed(2)
+    gate = torch.randn(rows, cols, device=DEVICE)
+    up = torch.randn(rows, cols, device=DEVICE)
+
+    ref = _ref_swiglu(gate, up)
+    out = triton_swiglu(gate, up)
+
+    torch.testing.assert_close(out, ref, atol=1e-5, rtol=1e-5)
+
+
+@triton_available
+def test_swiglu_negative_gate():
+    """silu(x) = x*sigmoid(x) is defined for negative x; result must not be NaN."""
+    gate = torch.full((4, 64), -10.0, device=DEVICE)
+    up = torch.ones(4, 64, device=DEVICE)
+
+    ref = _ref_swiglu(gate, up)
+    out = triton_swiglu(gate, up)
+
+    torch.testing.assert_close(out, ref, atol=1e-5, rtol=1e-5)
+    assert not torch.isnan(out).any()
+
+
+@triton_available
+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)
+    assert out.shape == gate.shape
+
+
+# ---------------------------------------------------------------------------
+# Fused Add + RMSNorm
+# ---------------------------------------------------------------------------
+
+
+@triton_available
+@pytest.mark.parametrize("rows,cols", [(1, 64), (4, 512), (16, 1024)])
+def test_fused_add_rms_norm_float32(rows, cols):
+    torch.manual_seed(3)
+    x = torch.randn(rows, cols, device=DEVICE)
+    residual = torch.randn(rows, cols, device=DEVICE)
+    weight = torch.ones(cols, device=DEVICE)
+
+    ref_y, ref_s = _ref_fused_add_rms_norm(x, residual, weight, EPS)
+    out_y, out_s = triton_fused_add_rms_norm(x, residual, weight, EPS)
+
+    torch.testing.assert_close(out_y, ref_y, atol=1e-5, rtol=1e-5)
+    # Sum output s = x + r must be bit-exact (no numerics involved, just addition)
+    torch.testing.assert_close(out_s, ref_s, atol=0, rtol=0)
+
+
+@triton_available
+def test_fused_add_rms_norm_sum_is_bitexact():
+    """The residual sum s = x + r must be bit-exact between Triton and PyTorch."""
+    torch.manual_seed(4)
+    x = torch.randn(8, 512, device=DEVICE)
+    r = torch.randn(8, 512, device=DEVICE)
+    w = torch.ones(512, device=DEVICE)
+
+    _, ref_s = _ref_fused_add_rms_norm(x, r, w, EPS)
+    _, out_s = triton_fused_add_rms_norm(x, r, w, EPS)
+
+    torch.testing.assert_close(out_s, ref_s, atol=0, rtol=0)
+
+
+@triton_available
+def test_fused_add_rms_norm_returns_two_tensors():
+    """Wrapper must return exactly two tensors: (normed_output, sum)."""
+    x = torch.randn(4, 256, device=DEVICE)
+    r = torch.randn(4, 256, device=DEVICE)
+    w = torch.ones(256, device=DEVICE)
+
+    result = triton_fused_add_rms_norm(x, r, w, EPS)
+
+    assert len(result) == 2
+    y, s = result
+    assert y.shape == x.shape
+    assert s.shape == x.shape

vllm_omni/model_executor/models/omnivoice/omnivoice_decoder.py

@@ -108,7 +108,10 @@ class OmniVoiceDecoder(nn.Module):
         quantized = self.quantizer.decode(codes)
 
         # Project: [B, 1024, T] → fc2 → [B, 256, T]
-        quantized = self.fc2(quantized.transpose(1, 2)).transpose(1, 2)
+        # Cast to fc2 weight dtype (may be fp16 when checkpoint stores weights as fp16),
+        # then upcast back to float32 — acoustic decoder ConvTranspose1d upsampling
+        # produces intermediate values that exceed the fp16 range (~65504), causing NaN.
+        quantized = self.fc2(quantized.transpose(1, 2).to(self.fc2.weight.dtype)).transpose(1, 2).float()
 
         # Acoustic decoder: [B, 256, T] → [B, 1, T*960]
         audio = self.acoustic_decoder(quantized)
@@ -192,6 +195,11 @@ class OmniVoiceDecoder(nn.Module):
                 param.data.copy_(state_dict[higgs_name])
                 loaded += 1
 
+        # Checkpoint weights may be fp16; force float32 so ConvTranspose1d
+        # upsampling doesn't overflow (fp16 max ~65504 is too small for
+        # intermediate activations in the DAC upsampling chain).
+        self.acoustic_decoder.float()
+
         # Apply HiggsAudioV2 output padding adjustment
         self._adjust_output_padding(self.acoustic_decoder)
 

vllm_omni/model_executor/models/omnivoice/omnivoice_generator.py

@@ -13,6 +13,8 @@ via FlashAttention/SageAttention/SDPA backends.
 from __future__ import annotations
 
 import math
+import threading
+from collections import OrderedDict
 
 import torch
 import torch.nn as nn
@@ -24,6 +26,184 @@ from vllm_omni.transformers_utils.configs.omnivoice import OmniVoiceConfig
 logger = init_logger(__name__)
 
 
+# ---------------------------------------------------------------------------
+# Triton kernels (inference-only; graceful fallback when triton is absent)
+# ---------------------------------------------------------------------------
+
+_TRITON_AVAILABLE = False
+try:
+    import triton
+    import triton.language as tl
+
+    try:
+        from triton.language.extra.libdevice import rsqrt  # noqa: F401
+    except ModuleNotFoundError:
+        from triton.language.extra.cuda.libdevice import rsqrt  # noqa: F401
+
+    def _calculate_settings(n: int) -> tuple[int, int]:
+        MAX_FUSED_SIZE = 65536
+        BLOCK_SIZE = triton.next_power_of_2(n)
+        if BLOCK_SIZE > MAX_FUSED_SIZE:
+            raise RuntimeError(f"n={n} exceeds max Triton block size {MAX_FUSED_SIZE}")
+        num_warps = 4
+        if BLOCK_SIZE >= 32768:
+            num_warps = 32
+        elif BLOCK_SIZE >= 8192:
+            num_warps = 16
+        elif BLOCK_SIZE >= 2048:
+            num_warps = 8
+        return BLOCK_SIZE, num_warps
+
+    @triton.jit
+    def _rms_norm_fwd_kernel(
+        Y_ptr,  # noqa: N803
+        Y_stride,  # noqa: N803
+        X_ptr,  # noqa: N803
+        X_stride,  # noqa: N803
+        W_ptr,  # noqa: N803
+        n_cols,
+        eps,
+        BLOCK_SIZE: tl.constexpr,  # noqa: N803
+    ):
+        row = tl.program_id(0).to(tl.int64)
+        cols = tl.arange(0, BLOCK_SIZE)
+        mask = cols < n_cols
+        x = tl.load(X_ptr + row * X_stride + cols, mask=mask, other=0.0)
+        w = tl.load(W_ptr + cols, mask=mask, other=0.0)
+        x_f32 = x.to(tl.float32)
+        ms = tl.sum(x_f32 * x_f32, axis=0) / n_cols
+        rstd = rsqrt(ms + eps)
+        y = (x_f32 * rstd).to(x.dtype) * w
+        tl.store(Y_ptr + row * Y_stride + cols, y, mask=mask)
+
+    def triton_rms_norm(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor:
+        shape = x.shape
+        n_cols = shape[-1]
+        x2d = x.contiguous().view(-1, n_cols)
+        w = weight.contiguous()
+        BLOCK_SIZE, num_warps = _calculate_settings(n_cols)
+        y = torch.empty_like(x2d)
+        _rms_norm_fwd_kernel[(x2d.shape[0],)](
+            y,
+            y.stride(0),
+            x2d,
+            x2d.stride(0),
+            w,
+            n_cols,
+            eps,
+            BLOCK_SIZE=BLOCK_SIZE,
+            num_warps=num_warps,
+        )
+        return y.view(*shape)
+
+    @triton.jit
+    def _swiglu_fwd_kernel(
+        gate_ptr,
+        up_ptr,
+        out_ptr,
+        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
+        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)
+        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)
+        BLOCK_SIZE, num_warps = _calculate_settings(n_cols)
+        _swiglu_fwd_kernel[(g2d.shape[0],)](
+            g2d,
+            u2d,
+            out,
+            out.stride(0),
+            n_cols=n_cols,
+            BLOCK_SIZE=BLOCK_SIZE,
+            num_warps=num_warps,
+        )
+        return out.view(*shape)
+
+    @triton.jit
+    def _fused_add_rms_norm_fwd_kernel(
+        Y_ptr,  # noqa: N803
+        Y_stride,  # noqa: N803
+        S_ptr,  # noqa: N803
+        S_stride,  # noqa: N803
+        X_ptr,  # noqa: N803
+        X_stride,  # noqa: N803
+        R_ptr,  # noqa: N803
+        R_stride,  # noqa: N803
+        W_ptr,  # noqa: N803
+        n_cols,
+        eps,
+        BLOCK_SIZE: tl.constexpr,  # noqa: N803
+    ):
+        row = tl.program_id(0).to(tl.int64)
+        cols = tl.arange(0, BLOCK_SIZE)
+        mask = cols < n_cols
+        x = tl.load(X_ptr + row * X_stride + cols, mask=mask, other=0.0)
+        r = tl.load(R_ptr + row * R_stride + cols, mask=mask, other=0.0)
+        dtype = x.dtype
+        s_f32 = x.to(tl.float32) + r.to(tl.float32)
+        s = s_f32.to(dtype)
+        tl.store(S_ptr + row * S_stride + cols, s, mask=mask)
+        w = tl.load(W_ptr + cols, mask=mask, other=0.0)
+        ms = tl.sum(s_f32 * s_f32, axis=0) / n_cols
+        rstd = rsqrt(ms + eps)
+        y = (s_f32 * rstd).to(dtype) * w
+        tl.store(Y_ptr + row * Y_stride + cols, y, mask=mask)
+
+    def triton_fused_add_rms_norm(
+        x: torch.Tensor,
+        residual: torch.Tensor,
+        weight: torch.Tensor,
+        eps: float,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        shape = x.shape
+        n_cols = shape[-1]
+        x2d = x.contiguous().view(-1, n_cols)
+        r2d = residual.contiguous().view(-1, n_cols)
+        w = weight.contiguous()
+        BLOCK_SIZE, num_warps = _calculate_settings(n_cols)
+        y = torch.empty_like(x2d)
+        s = torch.empty_like(x2d)
+        _fused_add_rms_norm_fwd_kernel[(x2d.shape[0],)](
+            y,
+            y.stride(0),
+            s,
+            s.stride(0),
+            x2d,
+            x2d.stride(0),
+            r2d,
+            r2d.stride(0),
+            w,
+            n_cols,
+            eps,
+            BLOCK_SIZE=BLOCK_SIZE,
+            num_warps=num_warps,
+        )
+        return y.view(*shape), s.view(*shape)
+
+    _TRITON_AVAILABLE = True
+    logger.debug("OmniVoice Triton kernels loaded")
+except Exception:
+    logger.debug("Triton not available; using PyTorch fallback for OmniVoice kernels")
+
+
 # ---------------------------------------------------------------------------
 # Unmasking schedule helpers
 # ---------------------------------------------------------------------------
@@ -63,6 +243,8 @@ class OmniVoiceRMSNorm(nn.Module):
         self.eps = eps
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if _TRITON_AVAILABLE:
+            return triton_rms_norm(x, self.weight, self.eps)
         variance = x.to(torch.float32).pow(2).mean(-1, keepdim=True)
         x = x * torch.rsqrt(variance + self.eps)
         return self.weight * x.to(self.weight.dtype)
@@ -158,6 +340,8 @@ class OmniVoiceMLP(nn.Module):
         self.down_proj = nn.Linear(config.llm_intermediate_size, config.llm_hidden_size, bias=False)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
+        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))
 
 
@@ -181,13 +365,22 @@ class OmniVoiceTransformerBlock(nn.Module):
         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)
-        hidden_states = residual + hidden_states
 
-        residual = hidden_states
-        hidden_states = self.post_attention_layernorm(hidden_states)
+        if _TRITON_AVAILABLE:
+            # Fused: (attn_out + residual) + RMSNorm in one kernel
+            hidden_states, residual = triton_fused_add_rms_norm(
+                hidden_states,
+                residual,
+                self.post_attention_layernorm.weight,
+                self.post_attention_layernorm.eps,
+            )
+        else:
+            hidden_states = residual + hidden_states
+            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
 
 
@@ -222,6 +415,214 @@ def _apply_rotary_pos_emb(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor)
     return x * torch.cat([cos, cos], dim=-1) + rotated * torch.cat([sin, sin], dim=-1)
 
 
+# ---------------------------------------------------------------------------
+# CUDA Graph wrapper
+# ---------------------------------------------------------------------------
+
+
+class _OmniVoiceCUDAGraphForward:
+    """Pre-captures CUDA graphs for predefined sequence-length buckets.
+
+    Memory layout: all graphs share a single per-instance pool handle, which
+    isolates OmniVoice CUDA memory from other vllm modules.  Sequential replay
+    (one step at a time) means pool sharing is safe.
+    """
+
+    # Default bucket count is 10; 16 gives modest headroom for edge cases
+    # (seq_len > max bucket or non-CFG batch) without unbounded GPU growth.
+    _MAX_LAZY_GRAPHS: int = 16
+
+    def __init__(self, generator: OmniVoiceGenerator, capture_sizes: list[int]) -> None:
+        self._gen = generator
+        self._capture_sizes = sorted(capture_sizes)
+        # Pre-warmed graphs keyed by (two_b, bucket); fixed set, never evicted.
+        self._graphs: dict[tuple[int, int], dict] = {}
+        # Lazy-captured graphs for oversized / non-CFG shapes; capped via LRU.
+        self._lazy_graphs: OrderedDict[tuple[int, int], dict] = OrderedDict()
+        self._lock = threading.Lock()
+        # Per-instance pool handle: isolates OmniVoice CUDA memory from other
+        # vllm modules while still allowing safe re-use across sequential replays.
+        self._pool_handle: int | None = None
+
+    def _find_bucket(self, seq_len: int) -> int | None:
+        for bucket in self._capture_sizes:
+            if bucket >= seq_len:
+                return bucket
+        return None
+
+    def _pad_inputs(
+        self,
+        input_ids: torch.Tensor,
+        audio_mask: torch.Tensor,
+        attention_mask: torch.Tensor | None,
+        bucket: int,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
+        S = input_ids.shape[-1]
+        if S == bucket:
+            return input_ids, audio_mask, attention_mask
+
+        two_b = input_ids.shape[0]
+        num_cb = input_ids.shape[1]
+
+        ids_padded = torch.zeros(two_b, num_cb, bucket, dtype=input_ids.dtype, device=input_ids.device)
+        ids_padded[:, :, :S] = input_ids
+
+        mask_padded = torch.zeros(two_b, bucket, dtype=torch.bool, device=audio_mask.device)
+        mask_padded[:, :S] = audio_mask
+
+        if attention_mask is not None:
+            attn_padded = torch.zeros(
+                two_b,
+                1,
+                bucket,
+                bucket,
+                dtype=attention_mask.dtype,
+                device=attention_mask.device,
+            )
+            attn_padded[:, :, :S, :S] = attention_mask
+        else:
+            attn_padded = None
+
+        return ids_padded, mask_padded, attn_padded
+
+    def _capture_for_key(
+        self,
+        key: tuple[int, int],
+        input_ids: torch.Tensor,
+        audio_mask: torch.Tensor,
+        attention_mask: torch.Tensor | None,
+    ) -> dict:
+        _, 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_input_ids = input_ids.clone()
+        static_audio_mask = audio_mask.clone()
+        static_attn_mask = attention_mask.clone() if attention_mask is not None else None
+
+        with torch.no_grad():
+            _ = self._gen._step_forward(
+                static_input_ids,
+                static_audio_mask,
+                static_attn_mask,
+                static_cos,
+                static_sin,
+            )
+        torch.accelerator.synchronize(device)
+
+        # Lazy-init per-instance pool handle: isolates OmniVoice CUDA Graph
+        # memory from other vllm modules (unlike get_global_graph_pool which
+        # shares a single pool across all captured graphs and can cause memory
+        # aliasing when two graphs replay concurrently).
+        if self._pool_handle is None:
+            self._pool_handle = torch.cuda.graph_pool_handle()
+
+        graph = torch.cuda.CUDAGraph()
+        with torch.no_grad():
+            with torch.cuda.graph(graph, pool=self._pool_handle):
+                static_output = self._gen._step_forward(
+                    static_input_ids,
+                    static_audio_mask,
+                    static_attn_mask,
+                    static_cos,
+                    static_sin,
+                )
+
+        entry = {
+            "graph": graph,
+            "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_output": static_output,
+        }
+        logger.info("OmniVoice CUDA Graph captured for key %s", key)
+        return entry
+
+    def warmup(self, device: torch.device) -> None:
+        """Pre-capture graphs for all bucket sizes with B=1 (two_b=2 for CFG)."""
+        if not torch.cuda.is_available():
+            return
+        logger.info(
+            "OmniVoice CUDA Graph warmup: capturing %d bucket sizes %s",
+            len(self._capture_sizes),
+            self._capture_sizes,
+        )
+        two_b = 2
+        num_cb = self._gen.config.num_audio_codebook
+        for bucket in self._capture_sizes:
+            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)
+            dummy_attn = torch.ones(two_b, 1, bucket, bucket, dtype=torch.bool, 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))
+
+    def __call__(
+        self,
+        input_ids: torch.Tensor,
+        audio_mask: torch.Tensor,
+        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)
+
+        seq_len = input_ids.shape[-1]
+        two_b = input_ids.shape[0]
+        bucket = self._find_bucket(seq_len) if two_b == 2 else None
+
+        if bucket is None:
+            # Lazy capture: oversized sequence or non-unit batch (no pre-warmed bucket).
+            # Lock prevents concurrent threads from double-capturing the same key.
+            # _lazy_graphs is capped at _MAX_LAZY_GRAPHS with LRU eviction to
+            # prevent unbounded GPU memory growth when seq_len varies widely.
+            key = (two_b, seq_len)
+            ids_in, mask_in, attn_in = input_ids, audio_mask, attention_mask
+            with self._lock:
+                entry = self._lazy_graphs.get(key)
+                if entry is None:
+                    entry = self._capture_for_key(key, ids_in, mask_in, attn_in)
+                    if len(self._lazy_graphs) >= self._MAX_LAZY_GRAPHS:
+                        evicted_key, _ = self._lazy_graphs.popitem(last=False)
+                        logger.warning("OmniVoice CUDA Graph lazy cache full; evicted key %s", evicted_key)
+                    self._lazy_graphs[key] = entry
+        else:
+            key = (two_b, bucket)
+            ids_in, mask_in, attn_in = self._pad_inputs(input_ids, audio_mask, attention_mask, bucket)
+            with self._lock:
+                entry = self._graphs.get(key)
+                if entry is None:
+                    entry = self._capture_for_key(key, ids_in, mask_in, attn_in)
+                    self._graphs[key] = entry
+
+        entry["static_input_ids"].copy_(ids_in)
+        entry["static_audio_mask"].copy_(mask_in)
+        if attn_in is not None and entry["static_attn_mask"] is not None:
+            entry["static_attn_mask"].copy_(attn_in)
+
+        entry["graph"].replay()
+
+        output = entry["static_output"]
+        if bucket is not None and bucket != seq_len:
+            output = output[:, :, :seq_len, :]
+        return output
+
+    def clear(self) -> None:
+        with self._lock:
+            self._graphs.clear()
+            self._lazy_graphs.clear()
+
+
 # ---------------------------------------------------------------------------
 # Generator model
 # ---------------------------------------------------------------------------
@@ -278,6 +679,11 @@ class OmniVoiceGenerator(nn.Module):
         self._rope_cos = None
         self._rope_sin = 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
+        )
+
     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:
@@ -364,6 +770,20 @@ class OmniVoiceGenerator(nn.Module):
             self.config.audio_vocab_size,
         ).permute(0, 2, 1, 3)  # [B, 8, S, 1025]
 
+    def _step_forward(
+        self,
+        input_ids: torch.Tensor,
+        audio_mask: torch.Tensor,
+        attention_mask: torch.Tensor | None,
+        cos: torch.Tensor,
+        sin: torch.Tensor,
+    ) -> torch.Tensor:
+        """Single unmasking-step forward using pre-cast RoPE tensors (CUDA graph safe)."""
+        hidden_states = self._prepare_embeddings(input_ids, audio_mask)
+        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))
+
     @torch.inference_mode()
     def forward(
         self,
@@ -440,10 +860,12 @@ class OmniVoiceGenerator(nn.Module):
 
         # Main iterative loop
         for step in range(num_step):
-            # Prepare embeddings and run transformer
-            inputs_embeds = self._prepare_embeddings(input_ids, audio_mask)
-            hidden_states = self._transformer_forward(inputs_embeds, attention_mask)
-            batch_logits = self._get_logits(hidden_states).to(torch.float32)
+            if self._cuda_graph_fwd is not None and input_ids.is_cuda:
+                batch_logits = self._cuda_graph_fwd(input_ids, audio_mask, attention_mask).to(torch.float32)
+            else:
+                inputs_embeds = self._prepare_embeddings(input_ids, audio_mask)
+                hidden_states = self._transformer_forward(inputs_embeds, attention_mask)
+                batch_logits = self._get_logits(hidden_states).to(torch.float32)
             # batch_logits: [2*B, 8, S, 1025]
 
             for i in range(B):
@@ -586,3 +1008,6 @@ class OmniVoiceGenerator(nn.Module):
             )
         else:
             logger.info("Generator: all %d weights loaded", len(loaded_keys))
+
+        if self._cuda_graph_fwd is not None:
+            self._cuda_graph_fwd.warmup(device)

vllm_omni/transformers_utils/configs/omnivoice.py

@@ -2,6 +2,8 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """OmniVoice configuration for vLLM-Omni two-stage pipeline."""
 
+import os
+
 from transformers import AutoConfig
 from transformers.configuration_utils import PretrainedConfig
 
@@ -80,6 +82,14 @@ class OmniVoiceConfig(PretrainedConfig):
 
         # Serving
         self.speculative_config = None
+        # Precedence: config.json > OMNIVOICE_CUDA_GRAPH env var > default (enabled).
+        # Read once at init; changing OMNIVOICE_CUDA_GRAPH at runtime has no effect.
+        self.enable_cuda_graph = getattr(self, "enable_cuda_graph", os.environ.get("OMNIVOICE_CUDA_GRAPH", "1") != "0")
+        self.cuda_graph_capture_sizes = getattr(
+            self,
+            "cuda_graph_capture_sizes",
+            [128, 192, 256, 320, 384, 448, 512, 640, 768, 1024],
+        )
 
 
 AutoConfig.register("omnivoice", OmniVoiceConfig)