diff --git a/tests/model_executor/models/omnivoice/__init__.py b/tests/model_executor/models/omnivoice/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/model_executor/models/omnivoice/test_cuda_graph_generator.py b/tests/model_executor/models/omnivoice/test_cuda_graph_generator.py new file mode 100644 index 000000000..dc6126713 --- /dev/null +++ b/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) diff --git a/tests/model_executor/models/omnivoice/test_decoder.py b/tests/model_executor/models/omnivoice/test_decoder.py new file mode 100644 index 000000000..db8e7328e --- /dev/null +++ b/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)" + ) diff --git a/tests/model_executor/models/omnivoice/test_triton_kernels.py b/tests/model_executor/models/omnivoice/test_triton_kernels.py new file mode 100644 index 000000000..1598b0f74 --- /dev/null +++ b/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 diff --git a/vllm_omni/model_executor/models/omnivoice/omnivoice_decoder.py b/vllm_omni/model_executor/models/omnivoice/omnivoice_decoder.py index 59f23d0d0..efb6d9dc7 100644 --- a/vllm_omni/model_executor/models/omnivoice/omnivoice_decoder.py +++ b/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) diff --git a/vllm_omni/model_executor/models/omnivoice/omnivoice_generator.py b/vllm_omni/model_executor/models/omnivoice/omnivoice_generator.py index eedb5828a..0de2a1f60 100644 --- a/vllm_omni/model_executor/models/omnivoice/omnivoice_generator.py +++ b/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) diff --git a/vllm_omni/transformers_utils/configs/omnivoice.py b/vllm_omni/transformers_utils/configs/omnivoice.py index d1bd3527f..1bd7207b6 100644 --- a/vllm_omni/transformers_utils/configs/omnivoice.py +++ b/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)