[Model] Optimize Qwen3-TTS streaming throughput and first audio on CUDA (#8163)

Signed-off-by: Sy03 <1370724210@qq.com>

tests/model_executor/models/qwen3_tts/test_cached_code_predictor.py

@@ -0,0 +1,50 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Cache isolation and graph replay for the opt-in TTS predictor."""
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.models.common.qwen3_code_predictor import CodePredictorBaseModel
+from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import Qwen3TTSTalkerCodePredictorConfig
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cuda]
+
+
+@pytest.mark.skipif(not torch.cuda.is_available() or torch.version.hip is not None, reason="requires CUDA")
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+def test_cached_frames_match_causal_prefill_after_graph_replay(dtype):
+    from vllm_omni.model_executor.models.qwen3_tts.cached_code_predictor import FrameLocalKVCache
+
+    config = Qwen3TTSTalkerCodePredictorConfig(
+        vocab_size=64,
+        hidden_size=64,
+        intermediate_size=128,
+        num_hidden_layers=2,
+        num_attention_heads=4,
+        num_key_value_heads=2,
+        head_dim=16,
+        num_code_groups=4,
+    )
+    torch.manual_seed(171)
+    model = CodePredictorBaseModel(config).to(device="cuda", dtype=dtype).eval()
+    inputs = torch.randn(3, 5, 64, device="cuda", dtype=dtype)
+    positions = torch.arange(5, device="cuda").expand(3, -1)
+    cache = FrameLocalKVCache(model, max_batch=3)
+    with torch.inference_mode():
+        for batch in (1, 3):
+            for _ in range(2):
+                for step in (1, 2, 3):
+                    cache(inputs, batch, step)
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            output = torch.cat([cache(inputs, 3, step) for step in (1, 2, 3)], dim=1)
+        for _ in range(3):
+            inputs.normal_()
+            # No previous frame (including a smaller batch) may leak through.
+            cache(inputs, 1, 1)
+            cache.buffer.fill_(torch.nan)
+            graph.replay()
+            expected = model(inputs, positions)[:, :4]
+            tolerance = 1e-5 if dtype == torch.float32 else 0.04
+            torch.testing.assert_close(output, expected, atol=tolerance, rtol=tolerance)

tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py

@@ -681,6 +681,7 @@ class TestCodePredictorGraphReplay:
             sampling_mode="per_call",
             return_proj_buf=False,
         )
+        predictor._frame_cache = None
         predictor._prefix_graphs_enabled = False
         predictor._prefix_reprefill_enabled = False
         predictor._bucket_pos_ids = {1: torch.arange(4).unsqueeze(0)}
@@ -1398,3 +1399,36 @@ class TestMTPExecutionBuckets:
             )
             wrapper2._bucket_sizes = [1, 2, 4]  # stand-in for warmed-up state
             wrapper2.configure_mtp_execution_buckets([3])
+
+
+@pytest.mark.parametrize("enabled", [False, True])
+def test_tts_cache_falls_back_on_cpu(mocker, loaded_target_classes, enabled):
+    cp_config, talker_config = _make_tiny_config(loaded_target_classes)
+    cfg = _make_vllm_config(mocker)
+    cfg.model_config.stage_connector_config = {"extra": {"code_predictor_kv_cache": enabled}}
+    predictor = loaded_target_classes[2](vllm_config=cfg, config=cp_config, talker_config=talker_config)
+    common = sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"]
+    setup = mocker.patch.object(common.CodePredictorWrapper, "_setup_compile")
+    predictor._setup_compile()
+    setup.assert_called_once_with()
+    assert predictor._frame_cache is None
+
+
+def test_shared_predictor_ignores_tts_fast_path_options(mocker, loaded_target_classes):
+    cp_config, _ = _make_tiny_config(loaded_target_classes)
+    cfg = _make_vllm_config(mocker)
+    cfg.model_config.stage_connector_config = {
+        "extra": {"code_predictor_kv_cache": True, "code_predictor_fused_sampling": True}
+    }
+    common = sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"]
+    predictor = common.CodePredictorWrapper(
+        vllm_config=cfg,
+        cp_config=cp_config,
+        wrapper_config=common.CodePredictorWrapperConfig(),
+    )
+    assert not hasattr(predictor, "_frame_cache")
+    logits = torch.randn(3, cp_config.vocab_size)
+    uniforms = torch.rand_like(logits)
+    expected = (logits.float() - torch.log(-torch.log(uniforms))).argmax(-1, keepdim=True)
+    actual = predictor._sample_per_call(logits, 1.0, 0, None, uniforms)
+    torch.testing.assert_close(actual, expected)

tests/model_executor/models/qwen3_tts/test_code_predictor_sampling.py

@@ -0,0 +1,68 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Supplied-uniform samples must match the existing ATen path exactly."""
+
+import pytest
+import torch
+
+pytestmark = [
+    pytest.mark.core_model,
+    pytest.mark.cuda,
+    pytest.mark.skipif(not torch.cuda.is_available() or torch.version.hip is not None, reason="requires CUDA"),
+]
+
+
+@pytest.mark.parametrize("batch", [1, 64])
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32])
+@pytest.mark.parametrize("top_k", [-1, 1, 50, 2048])
+def test_topk_gumbel_exact_samples_and_graph(batch, dtype, top_k):
+    from vllm_omni.model_executor.models.qwen3_tts.code_predictor_sampling import sample_code_topk_gumbel
+
+    with torch.inference_mode():
+        torch.manual_seed(42)
+        logits = torch.randn(batch, 2048, device="cuda", dtype=dtype) * 5
+        uniforms = torch.rand(batch, 15, 2048, device="cuda")[:, 7, :]
+        uniforms.clamp_(1e-6, 1 - 1e-6)
+        for case in ("random", "ties", "all_masked", "part_masked", "uniform_zero", "uniform_one", "nan"):
+            if case == "ties":
+                logits.zero_()
+                uniforms.fill_(0.5)
+            elif case == "all_masked":
+                logits.fill_(-torch.inf)
+            elif case == "part_masked":
+                logits[:, 4:9] = 1
+            elif case == "uniform_zero":
+                logits.normal_()
+                uniforms.zero_()
+            elif case == "uniform_one":
+                uniforms.fill_(1)
+            elif case == "nan":
+                uniforms.fill_(0.5)
+                logits[:, 7] = torch.nan
+            scaled = logits * (1 / 0.9)
+            if top_k > 0:
+                threshold = scaled.topk(top_k, dim=-1).values[:, -1:]
+                scaled = scaled.masked_fill(scaled < threshold, -torch.inf)
+            expected = (scaled.float() - torch.log(-torch.log(uniforms))).argmax(-1, keepdim=True)
+            actual = sample_code_topk_gumbel(logits, uniforms, top_k, 1 / 0.9)
+            torch.testing.assert_close(actual, expected, atol=0, rtol=0)
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            captured = sample_code_topk_gumbel(logits, uniforms, top_k, 1 / 0.9)
+        graph.replay()
+        torch.testing.assert_close(captured, expected, atol=0, rtol=0)
+
+
+@pytest.mark.parametrize("top_k", [0, 1, 50, 97])
+def test_topk_gumbel_accepts_strided_supplied_uniforms(top_k):
+    from vllm_omni.model_executor.models.qwen3_tts.code_predictor_sampling import sample_code_topk_gumbel
+
+    torch.manual_seed(81)
+    logits = torch.randn(3, 194, device="cuda", dtype=torch.bfloat16)[:, ::2]
+    uniforms = torch.rand(3, 15, 194, device="cuda")[:, 3, ::2].clamp_(1e-6, 1 - 1e-6)
+    scaled = logits * (1 / 0.9)
+    if top_k > 0:
+        scaled = scaled.masked_fill(scaled < scaled.topk(top_k).values[:, -1:], -torch.inf)
+    expected = (scaled.float() - torch.log(-torch.log(uniforms))).argmax(-1, keepdim=True)
+    actual = sample_code_topk_gumbel(logits, uniforms, top_k, 1 / 0.9)
+    torch.testing.assert_close(actual, expected)

tests/model_executor/models/qwen3_tts/test_first_audio_graphs.py

@@ -0,0 +1,126 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Prefix graphs must follow per-request first-audio delivery, not a capture hint."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from tests.helpers.mark import hardware_test
+from tests.model_executor.models.qwen3_tts.test_qwen3_tts_incremental_decode import _decoder_stub
+from vllm_omni.model_executor.models.qwen3_tts.segmented_graph_wrapper import CUDAGraphDecoderWrapper
+from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.modeling_qwen3_tts_tokenizer_v2 import (
+    Qwen3TTSTokenizerV2Decoder,
+)
+
+pytestmark = [pytest.mark.core_model]
+DEVICE = torch.device("cuda:0")
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize("skip_flags", [(False, False), (True, True), (False, True), (True, False)])
+@pytest.mark.parametrize("initial_frames", [1, 3])
+@pytest.mark.parametrize("state_only_available", [False, True])
+def test_xvec_first_chunk_replays_prefix_graph(monkeypatch, skip_flags, initial_frames, state_only_available):
+    monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False)
+    wrapper = CUDAGraphDecoderWrapper.__new__(CUDAGraphDecoderWrapper)
+    wrapper.prefix_length = 72
+    wrapper.initial_chunk_frames = initial_frames
+    wrapper.codec_chunk_frames = 25
+    wrapper.capture_batch_sizes = [1, 2]
+    wrapper._icl_previous_frames_by_target = {26: 1}
+    wrapper._xvec_previous_frames_by_target = {initial_frames + 25: initial_frames}
+    replayed = []
+    wrapper.xvec_prefix_states = {
+        2: {
+            "graph": SimpleNamespace(replay=lambda: replayed.append("full")),
+            "input": {"codes": torch.zeros(2, 2, initial_frames)},
+            "output": torch.arange(4 * initial_frames, dtype=torch.float32).view(2, 1, 2 * initial_frames),
+            "cache": {
+                "ref_hidden": torch.zeros(2, 2, 0),
+                "ref_conv": torch.zeros(2, 0, 3),
+                "prefix_hidden": torch.zeros(2, 0, 3),
+                "ref_upsample": torch.zeros(2, 3, 0),
+                "ref_wav": torch.zeros(2, 1, 0),
+                "suffix_quantized": torch.ones(2, 2, initial_frames),
+                "suffix_conv": torch.ones(2, initial_frames, 3),
+                "past_key_values": SimpleNamespace(layers=[]),
+            },
+        }
+    }
+    wrapper.xvec_prefix_state_only_states = {}
+    if state_only_available and initial_frames == 1:
+        wrapper.xvec_prefix_state_only_states[2] = {
+            **wrapper.xvec_prefix_states[2],
+            "graph": SimpleNamespace(replay=lambda: replayed.append("state")),
+            "state_only": True,
+        }
+    wrapper._record_graph_hit = lambda *_args: None
+    wrapper._record_graph_fallback = lambda *_args: None
+    wrapper._ensure_suffix_buffers = lambda cache: None
+    wrapper.decoder = _decoder_stub(
+        capture_first_audio_state_only=True,
+        total_upsample=2,
+        _is_suffix_cache_rolling=lambda previous, cached: False,
+        _decode_xvec_first_chunk=lambda *_args: pytest.fail("unexpected eager fallback"),
+        _slice_dynamic_cache=Qwen3TTSTokenizerV2Decoder._slice_dynamic_cache,
+    )
+
+    caches = [{"prefix_frames": 0, "skip_first_audio": skip} for skip in skip_flags]
+    outputs = wrapper._batched_request_decode(
+        [torch.full((1, 2, initial_frames), 3), torch.full((1, 2, initial_frames), 4)],
+        caches,
+    )
+
+    assert replayed == ["state" if all(skip_flags) and state_only_available and initial_frames == 1 else "full"]
+    torch.testing.assert_close(wrapper.xvec_prefix_states[2]["input"]["codes"][0], torch.full((2, initial_frames), 3.0))
+    torch.testing.assert_close(wrapper.xvec_prefix_states[2]["input"]["codes"][1], torch.full((2, initial_frames), 4.0))
+    assert [cache["decoder_prefix_frames"] for cache in caches] == [0, 0]
+    assert [cache["suffix_frames"] for cache in caches] == [initial_frames, initial_frames]
+    assert len(outputs) == 2
+    for row, skip in enumerate(skip_flags):
+        expected = wrapper.xvec_prefix_states[2]["output"][row : row + 1, :, 2 if skip else 0 :]
+        torch.testing.assert_close(outputs[row], expected)
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("initial_frames", [1, 3])
+@pytest.mark.parametrize("time_major", [False, True])
+@torch.inference_mode()
+def test_xvec_graphs_follow_request_first_audio_flags(initial_frames, time_major):
+    from tests.model_executor.models.qwen3_tts.test_time_major_decoder import _make_decoder
+    from vllm_omni.model_executor.models.qwen3_tts.segmented_graph_wrapper import (
+        CUDAGraphDecoderWrapper as SegmentedWrapper,
+    )
+
+    decoder = _make_decoder().to(DEVICE)
+    decoder.precompute_snake_caches()
+    if time_major:
+        decoder.enable_time_major_conv()
+    decoder.capture_first_audio_state_only = True
+    wrapper = SegmentedWrapper(
+        decoder,
+        capture_modes=("xvec",),
+        capture_batch_sizes=[1, 2],
+        num_quantizers=2,
+        initial_chunk_frames=initial_frames,
+    )
+    # This regression covers prefix graphs. Avoid compiling unrelated suffix
+    # shapes, especially the time-major kernels for long continuation chunks.
+    wrapper._xvec_previous_frames_by_target = {}
+    wrapper.warmup(DEVICE)
+    assert set(wrapper.xvec_prefix_states) == {1, 2}
+    assert set(wrapper.xvec_prefix_state_only_states) == ({1, 2} if initial_frames == 1 else set())
+
+    for skip_flags in ((False, False), (True, True), (False, True), (True, False)):
+        codes = [torch.randint(0, 32, (1, 2, initial_frames), device=DEVICE) for _ in skip_flags]
+        caches = [{"prefix_frames": 0, "skip_first_audio": skip} for skip in skip_flags]
+        expected_caches = [{"prefix_frames": 0, "skip_first_audio": skip} for skip in skip_flags]
+        expected = [decoder._decode_stream_first_chunk(code, cache) for code, cache in zip(codes, expected_caches)]
+        actual = wrapper._decode_xvec_prefix_batch(codes, caches)
+        assert actual is not None  # Eligible batches must not silently become per-request eager calls.
+        for got, want, cache, expected_cache in zip(actual, expected, caches, expected_caches):
+            torch.testing.assert_close(got, want, atol=1e-4, rtol=1e-4)
+            for key in ("ref_hidden", "ref_conv", "prefix_hidden", "suffix_quantized", "suffix_conv"):
+                torch.testing.assert_close(cache[key], expected_cache[key], atol=1e-5, rtol=1e-5)

tests/model_executor/models/qwen3_tts/test_qwen3_tts_code2wav.py

@@ -752,3 +752,37 @@ def test_load_weights_uses_model_dtype_before_precomputing_caches(dtype):
     assert _load_weights_noop(model) == {"decoder.fake_weight"}
     assert model.decoder.weight.dtype is dtype
     assert cache_dtypes == [dtype]
+
+
+@pytest.mark.parametrize("capture_fails", [False, True])
+def test_decode_autotune_restores_process_flags(mocker, monkeypatch, capture_fails):
+    monkeypatch.setattr(torch.backends.cudnn, "benchmark", False)
+    monkeypatch.setattr(torch.backends.cudnn, "benchmark_limit", 5)
+    original = (torch.backends.cudnn.enabled, torch.backends.cudnn.deterministic, torch.backends.cudnn.allow_tf32)
+    model = _make_model(
+        async_chunk=True,
+        device=torch.device("cuda"),
+        stage_connector_config={"extra": {"decode_cudnn_benchmark": True}},
+    )
+
+    observed = []
+
+    def capture(**kwargs):
+        observed.append(
+            (
+                torch.backends.cudnn.benchmark,
+                torch.backends.cudnn.benchmark_limit,
+                torch.backends.cudnn.enabled,
+                torch.backends.cudnn.deterministic,
+                torch.backends.cudnn.allow_tf32,
+            )
+        )
+        if capture_fails:
+            raise RuntimeError("capture failed")
+
+    capture_mock = mocker.patch.object(model, "_maybe_enable_decoder_cudagraph", side_effect=capture)
+    _load_weights_noop(model)
+    capture_mock.assert_called_once()
+    assert observed == [(True, 10, *original)]
+    assert torch.backends.cudnn.benchmark is False
+    assert torch.backends.cudnn.benchmark_limit == 5

tests/model_executor/models/qwen3_tts/test_short_kv_attention.py

@@ -0,0 +1,44 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Causal/GQA correctness and graph replay for the frame-local KV kernel."""
+
+import pytest
+import torch
+
+pytestmark = [
+    pytest.mark.core_model,
+    pytest.mark.cuda,
+    pytest.mark.skipif(not torch.cuda.is_available() or torch.version.hip is not None, reason="requires CUDA"),
+]
+
+
+@pytest.mark.parametrize("batch", [1, 3, 64])
+@pytest.mark.parametrize("heads,kv_heads,dim", [(16, 8, 128), (8, 8, 64)])
+@pytest.mark.parametrize("positions", [[0, 1], [2], [8], [15]])
+def test_short_attention_causal_gqa_and_graph(batch, heads, kv_heads, dim, positions):
+    from vllm_omni.model_executor.models.qwen3_tts.short_kv_attention import short_kv_attention
+
+    with torch.inference_mode():
+        torch.manual_seed(42)
+        pos = torch.tensor(positions, device="cuda")
+        q = torch.randn(batch, len(positions), heads, dim, device="cuda", dtype=torch.bfloat16).transpose(1, 2)
+        cache = torch.randn(2, batch + 3, kv_heads, 17, dim, device="cuda", dtype=torch.bfloat16)
+        k, v = cache[0, :batch], cache[1, :batch]
+        mask = torch.arange(17, device="cuda")[None, :] <= pos[:, None]
+        kr = k.float().repeat_interleave(heads // kv_heads, dim=1)
+        vr = v.float().repeat_interleave(heads // kv_heads, dim=1)
+        scores = (q.float() @ kr.transpose(-1, -2)) * dim**-0.5
+        expected = scores.masked_fill(~mask, -torch.inf).softmax(-1) @ vr
+        actual = short_kv_attention(q, k, v, pos, dim**-0.5)
+        torch.testing.assert_close(actual.float(), expected, atol=0.016, rtol=0.008)
+
+        # Unwritten cache slots must not affect a new frame, even if poisoned.
+        k[:, :, max(positions) + 1 :] = torch.nan
+        v[:, :, max(positions) + 1 :] = torch.nan
+        poisoned = short_kv_attention(q, k, v, pos, dim**-0.5)
+        torch.testing.assert_close(poisoned, actual, atol=0, rtol=0)
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            captured = short_kv_attention(q, k, v, pos, dim**-0.5)
+        graph.replay()
+        torch.testing.assert_close(captured, actual, atol=0, rtol=0)

tests/model_executor/models/qwen3_tts/test_skip_delivered_first.py

@@ -0,0 +1,65 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Code2Wav drops only frames explicitly marked as delivered by the Talker."""
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.modeling_qwen3_tts_tokenizer_v2 import (
+    Qwen3TTSTokenizerV2Decoder,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+UPSAMPLE = 4
+
+
+def _decoder(skip: bool):
+    decoder = Qwen3TTSTokenizerV2Decoder.__new__(Qwen3TTSTokenizerV2Decoder)
+    decoder.__dict__["capture_first_audio_state_only"] = skip
+    decoder.__dict__["total_upsample"] = UPSAMPLE
+    calls = []
+
+    def full(codes, caches):
+        calls.append("full")
+        frames = int(codes.shape[-1])
+        return torch.arange(frames * UPSAMPLE, dtype=torch.float32).view(1, 1, -1)
+
+    def state_only(codes, caches):
+        calls.append("state")
+        return torch.zeros(1)
+
+    decoder.__dict__["_decode_xvec_first_chunk"] = full
+    decoder.__dict__["_decode_xvec_first_chunk_state_only"] = state_only
+    return decoder, calls
+
+
+def test_one_frame_first_chunk_only_builds_state():
+    decoder, calls = _decoder(skip=True)
+    out = decoder._decode_stream_first_chunk(torch.zeros(1, 16, 1, dtype=torch.long), {"skip_first_audio": True})
+    assert out.shape[-1] == 0 and calls == ["state"]
+
+
+def test_longer_first_chunk_keeps_all_but_the_delivered_frame():
+    decoder, calls = _decoder(skip=True)
+    out = decoder._decode_stream_first_chunk(
+        torch.zeros(1, 16, 3, dtype=torch.long),
+        {"skip_first_audio": True},
+    )
+    assert calls == ["full"]
+    assert torch.equal(out, torch.arange(UPSAMPLE, 3 * UPSAMPLE, dtype=torch.float32).view(1, 1, -1))
+
+
+def test_without_skip_the_whole_first_chunk_is_decoded():
+    decoder, calls = _decoder(skip=False)
+    out = decoder._decode_stream_first_chunk(
+        torch.zeros(1, 16, 3, dtype=torch.long),
+        {"skip_first_audio": False},
+    )
+    assert calls == ["full"] and out.shape[-1] == 3 * UPSAMPLE
+
+
+def test_request_without_first_audio_uses_regular_decoder_even_when_graph_skip_is_enabled():
+    decoder, calls = _decoder(skip=True)
+    out = decoder._decode_stream_first_chunk(torch.zeros(1, 16, 1, dtype=torch.long), {})
+    assert calls == ["full"] and out.shape[-1] == UPSAMPLE

tests/model_executor/models/qwen3_tts/test_time_major_conv.py

@@ -0,0 +1,112 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz import time_major_conv as tmc
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _snake_nct(x, exp_alpha, inv_beta):
+    return x + inv_beta[None, :, None] * torch.sin(x * exp_alpha[None, :, None]).square()
+
+
+@pytest.mark.parametrize(("kernel_size", "dilation"), [(1, 1), (7, 1), (7, 3), (3, 9)])
+@pytest.mark.parametrize("with_bias_and_snake", [False, True])
+def test_im2col_gemm_matches_causal_conv1d(kernel_size, dilation, with_bias_and_snake):
+    torch.manual_seed(0)
+    bsz, t, c_in, c_out = 2, 11, 6, 5
+    conv = torch.nn.Conv1d(c_in, c_out, kernel_size, dilation=dilation)
+    x = torch.randn(bsz, c_in, t)
+    in_bias = torch.randn(c_in) if with_bias_and_snake else None
+    snake = (torch.rand(c_in) + 0.5, torch.rand(c_in) + 0.5) if with_bias_and_snake else None
+
+    ref_in = x if in_bias is None else x + in_bias[None, :, None]
+    if snake is not None:
+        ref_in = _snake_nct(ref_in, *snake)
+    expected = conv(F.pad(ref_in, ((kernel_size - 1) * dilation, 0)))
+
+    cols = tmc.snake_im2col(x.transpose(1, 2).contiguous(), kernel_size, dilation, in_bias, snake)
+    actual = torch.addmm(conv.bias, cols, tmc.conv_gemm_weight(conv.weight)).view(bsz, t, c_out)
+    torch.testing.assert_close(actual.transpose(1, 2), expected, rtol=1e-5, atol=1e-5)
+
+
+@pytest.mark.parametrize(("stride", "kernel_size"), [(2, 2), (3, 6), (8, 16)])
+def test_transposed_overlap_add_matches_trimmed_conv_transpose(stride, kernel_size):
+    torch.manual_seed(0)
+    bsz, t, c_in, c_out = 2, 5, 4, 3
+    conv = torch.nn.ConvTranspose1d(c_in, c_out, kernel_size, stride=stride)
+    x = torch.randn(bsz, c_in, t)
+    expected = conv(x)[..., : t * stride]  # causal: drop the kernel - stride right tail
+
+    z = (x.transpose(1, 2).reshape(bsz * t, c_in) @ tmc.transposed_gemm_weight(conv.weight)).view(
+        bsz, t, kernel_size, c_out
+    )
+    actual = tmc.transposed_overlap_add(z, stride, conv.bias)
+    torch.testing.assert_close(actual.transpose(1, 2), expected, rtol=1e-5, atol=1e-5)
+
+
+def test_transposed_overlap_add_rejects_other_kernels():
+    with pytest.raises(ValueError, match="kernel stride or 2"):
+        tmc.transposed_overlap_add(torch.zeros(1, 2, 5, 3), 2, torch.zeros(3))
+
+
+@pytest.mark.parametrize(("kernel_size", "dilation"), [(1, 1), (7, 3)])
+@pytest.mark.parametrize("with_residual", [False, True])
+def test_causal_conv_matches_conv1d(kernel_size, dilation, with_residual):
+    torch.manual_seed(0)
+    bsz, t, c_in, c_out = 2, 13, 6, 5
+    conv = torch.nn.Conv1d(c_in, c_out, kernel_size, dilation=dilation)
+    x = torch.randn(bsz, c_in, t)
+    residual = torch.randn(bsz, t, c_out) if with_residual else None
+    expected = conv(F.pad(x, ((kernel_size - 1) * dilation, 0))).transpose(1, 2)
+    if residual is not None:
+        expected = expected + residual
+    actual = tmc.causal_conv(
+        x.transpose(1, 2).contiguous(), tmc.conv_gemm_weight(conv.weight), conv.bias, kernel_size, dilation, residual
+    )
+    torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-5)
+
+
+def _gpu_tolerances(dtype):
+    # bf16 results may differ by one ulp (sin rounding, fp32 vs bf16 accumulation order).
+    return {"rtol": 1.6e-2, "atol": 1e-2} if dtype == torch.bfloat16 else {"rtol": 1e-4, "atol": 1e-4}
+
+
+@pytest.mark.skipif(not torch.cuda.is_available() or not tmc.HAS_TRITON, reason="needs CUDA and Triton")
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+def test_triton_kernels_match_reference(dtype):
+    torch.manual_seed(0)
+    tol = _gpu_tolerances(dtype)
+    x = torch.randn(3, 37, 96, device="cuda", dtype=dtype)
+    bias = torch.randn(96, device="cuda", dtype=dtype)
+    snake = (torch.rand(96, device="cuda", dtype=dtype) + 0.5, torch.rand(96, device="cuda", dtype=dtype) + 0.5)
+    for kernel_size, dilation in ((1, 1), (7, 9)):
+        torch.testing.assert_close(
+            tmc.snake_im2col(x, kernel_size, dilation, bias, snake),
+            tmc.snake_im2col_reference(x, kernel_size, dilation, bias, snake),
+            **tol,
+        )
+    z = torch.randn(3, 37, 10, 48, device="cuda", dtype=dtype)
+    torch.testing.assert_close(
+        tmc.transposed_overlap_add(z, 5, bias[:48]), tmc.transposed_overlap_add_reference(z, 5, bias[:48]), **tol
+    )
+
+
+@pytest.mark.skipif(not torch.cuda.is_available() or not tmc.HAS_TRITON, reason="needs CUDA and Triton")
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+@pytest.mark.parametrize(("c_in", "c_out", "kernel_size", "dilation"), [(96, 96, 7, 9), (192, 96, 1, 1), (96, 1, 7, 1)])
+def test_triton_causal_conv_matches_reference(dtype, c_in, c_out, kernel_size, dilation):
+    torch.manual_seed(0)
+    x = torch.randn(3, 101, c_in, device="cuda", dtype=dtype)
+    weight = torch.randn(kernel_size * c_in, c_out, device="cuda", dtype=dtype) / (kernel_size * c_in) ** 0.5
+    bias = torch.randn(c_out, device="cuda", dtype=dtype)
+    residual = torch.randn(3, 101, c_out, device="cuda", dtype=dtype)
+    expected = tmc.causal_conv_reference(
+        x.float(), weight.float(), bias.float(), kernel_size, dilation, residual.float()
+    )
+    actual = tmc.causal_conv(x, weight, bias, kernel_size, dilation, residual)
+    torch.testing.assert_close(actual.float(), expected, **_gpu_tolerances(dtype))

tests/model_executor/models/qwen3_tts/test_time_major_decoder.py

@@ -0,0 +1,85 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.configuration_qwen3_tts_tokenizer_v2 import (
+    Qwen3TTSTokenizerV2DecoderConfig,
+)
+from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.modeling_qwen3_tts_tokenizer_v2 import (
+    Qwen3TTSTokenizerV2Decoder,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _make_decoder() -> Qwen3TTSTokenizerV2Decoder:
+    config = Qwen3TTSTokenizerV2DecoderConfig(
+        codebook_size=32,
+        hidden_size=16,
+        latent_dim=16,
+        codebook_dim=16,
+        num_attention_heads=2,
+        num_key_value_heads=2,
+        intermediate_size=32,
+        num_hidden_layers=1,
+        num_quantizers=2,
+        decoder_dim=32,
+        upsample_rates=(8, 5, 4, 3),
+        upsampling_ratios=(2, 2),
+        sliding_window=72,
+    )
+    torch.manual_seed(0)
+    decoder = Qwen3TTSTokenizerV2Decoder(config).eval()
+    with torch.no_grad():
+        for name, param in decoder.named_parameters():
+            # Non-trivial SnakeBeta and ConvNeXt scales so every term matters.
+            if name.endswith(("alpha", "beta")):
+                param.uniform_(-0.5, 0.5)
+            elif name.endswith("gamma"):
+                param.uniform_(0.5, 1.0)
+    decoder.precompute_snake_caches()
+    return decoder
+
+
+@torch.inference_mode()
+def test_time_major_conv_decode_matches_nct_stack() -> None:
+    decoder = _make_decoder()
+    hidden = torch.randn(2, 5, decoder.config.latent_dim)
+    expected = decoder._conv_decode(hidden)
+
+    decoder.enable_time_major_conv()
+    actual = decoder._conv_decode(hidden)
+
+    assert actual.shape == expected.shape == (2, 1, 5 * int(decoder.total_upsample))
+    scale = expected.abs().max()
+    torch.testing.assert_close(actual / scale, expected / scale, rtol=0, atol=1e-5)
+    # The stack holds weight copies outside the module tree.
+    assert not any("time_major" in key for key in decoder.state_dict())
+
+
+@torch.inference_mode()
+def test_time_major_conv_decode_takes_sliced_hidden() -> None:
+    decoder = _make_decoder()
+    hidden = torch.randn(1, 9, decoder.config.latent_dim)
+    expected = decoder._conv_decode(hidden[:, -4:, :])
+    decoder.enable_time_major_conv()
+    actual = decoder._conv_decode(hidden[:, -4:, :])  # non-contiguous, like the suffix decode's slice
+    scale = expected.abs().max()
+    torch.testing.assert_close(actual / scale, expected / scale, rtol=0, atol=1e-5)
+
+
+@torch.inference_mode()
+def test_first_frame_single_key_attention_matches_codec_without_graphs():
+    from vllm_omni.model_executor.models.qwen3_tts.first_frame_decoder import Qwen3TTSFirstFrameDecoder
+
+    decoder = _make_decoder()
+    codes = torch.randint(0, 32, (3, 2, 1))
+    expected = decoder._decode_xvec_first_chunk(codes, {})[:, 0, :]
+    first = Qwen3TTSFirstFrameDecoder.__new__(Qwen3TTSFirstFrameDecoder)
+    torch.nn.Module.__init__(first)
+    first.decoder = decoder
+    first._graphs = {}
+    actual = first.decode(codes[:, :, 0])
+    torch.testing.assert_close(actual, expected, rtol=1e-4, atol=1e-5)