[Model] Optimize Qwen3-Omni MRv2 handoff, code prediction and audio decoding (#8223)

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

tests/model_executor/models/common/test_fused_sampling.py

@@ -0,0 +1,90 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+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"),
+]
+
+
+def _reference(logits, uniforms, top_k, top_p):
+    if top_k > 0:
+        threshold = logits.topk(top_k, dim=-1).values[:, -1:]
+        logits = logits.masked_fill(logits < threshold, -torch.inf)
+    if top_p < 1:
+        values, indices = logits.sort(dim=-1, descending=True)
+        probs = values.softmax(dim=-1, dtype=torch.float32)
+        values[(probs.cumsum(-1) - probs) >= top_p] = -torch.inf
+        logits = values.scatter(1, indices, values)
+    return (logits.float() - torch.log(-torch.log(uniforms))).argmax(-1, keepdim=True)
+
+
+@pytest.mark.parametrize("batch,vocab", [(1, 128), (7, 2048), (64, 2048)])
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
+@pytest.mark.parametrize("top_k,top_p", [(-1, 0.8), (0, 1.0), (50, 0.8), (50, 0.01)])
+def test_fixed_uniform_sampling_and_graph(batch, vocab, dtype, top_k, top_p):
+    from vllm_omni.model_executor.models.common.fused_sampling import sample_top_k_top_p_gumbel
+
+    torch.manual_seed(173)
+    logits = torch.randn(batch, vocab, device="cuda", dtype=dtype)
+    # The MTP caller supplies a noncontiguous row slice of [B, 15, vocab].
+    uniforms = torch.rand(batch, 15, vocab, device="cuda")[:, 3].clamp_min_(1e-7)
+    expected = _reference(logits, uniforms, top_k, top_p)
+    actual = sample_top_k_top_p_gumbel(logits, uniforms, top_k=top_k, top_p=top_p)
+    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        captured = sample_top_k_top_p_gumbel(logits, uniforms, top_k=top_k, top_p=top_p)
+    logits.add_(torch.randn_like(logits) * 0.1)
+    uniforms.uniform_(1e-7, 1 - 1e-7)
+    graph.replay()
+    torch.testing.assert_close(captured, _reference(logits, uniforms, top_k, top_p), rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("top_k,top_p", [(3, 0.5), (3, 1.0), (0, 0.8)])
+def test_threshold_ties_and_degenerate_rows(top_k, top_p):
+    from vllm_omni.model_executor.models.common.fused_sampling import sample_top_k_top_p_gumbel
+
+    logits = torch.tensor(
+        [[2, 2, 2, 2, 2, 1, 0, -1], [1] * 8, [-torch.inf] * 8, [torch.nan, 1, 2, 3, 4, 5, 6, 7]],
+        device="cuda",
+        dtype=torch.float32,
+    )
+    uniforms = torch.tensor([[0.4, 0.7, 0.8, 0.9, 0.2, 0.3, 0.6, 0.5]], device="cuda").expand(4, -1)
+    actual = sample_top_k_top_p_gumbel(logits, uniforms, top_k=top_k, top_p=top_p)
+    torch.testing.assert_close(actual, _reference(logits, uniforms, top_k, top_p), rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
+def test_native_nucleus_boundaries(dtype):
+    from vllm_omni.model_executor.models.common.fused_sampling import sample_top_k_top_p_gumbel
+
+    torch.manual_seed(1923)
+    logits = torch.randn(1, 2048, device="cuda", dtype=dtype)
+    values, indices = logits.sort(descending=True)
+    probabilities = values.softmax(-1, dtype=torch.float32)
+    before = probabilities.cumsum(-1) - probabilities
+    for rank in [1, 19, 49, 1000, 2000]:
+        boundary = before[0, rank]
+        uniforms = torch.full_like(logits, 0.001, dtype=torch.float32)
+        uniforms[0, indices[0, rank]] = 1 - 1e-7
+        for direction in [-torch.inf, torch.inf]:
+            top_p = torch.nextafter(boundary, boundary.new_tensor(direction)).item()
+            actual = sample_top_k_top_p_gumbel(logits, uniforms, top_k=0, top_p=top_p)
+            torch.testing.assert_close(actual, _reference(logits, uniforms, 0, top_p), rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("top_p", [0.8, 1.0])
+def test_strided_logits_and_non_power_of_two_vocabulary(top_p):
+    from vllm_omni.model_executor.models.common.fused_sampling import sample_top_k_top_p_gumbel
+
+    for seed in range(16):
+        torch.manual_seed(seed)
+        logits = torch.randn(128, 2050, device="cuda")[:, ::2]
+        uniforms = torch.rand(128, 3, 1025, device="cuda")[:, 1].clamp_min_(1e-7)
+        actual = sample_top_k_top_p_gumbel(logits, uniforms, top_k=-1, top_p=top_p)
+        torch.testing.assert_close(actual, _reference(logits, uniforms, -1, top_p), rtol=0, atol=0)

tests/model_executor/models/common/test_short_kv_attention.py

@@ -16,7 +16,7 @@ pytestmark = [
 @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
+    from vllm_omni.model_executor.models.common.short_kv_attention import short_kv_attention
 
     with torch.inference_mode():
         torch.manual_seed(42)

tests/model_executor/models/qwen3_omni/test_first_frame_cuda.py

@@ -0,0 +1,32 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.models.qwen3_omni.first_frame_decoder import Qwen3OmniFirstFrameDecoder
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cuda]
+
+
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+@pytest.mark.parametrize("rows", [16, 19])
+@torch.inference_mode()
+def test_first_frame_graph_outputs_own_storage(dtype, rows):
+    class Decoder(torch.nn.Linear):
+        config = SimpleNamespace(num_quantizers=1)
+
+        def forward(self, codes):
+            return codes.to(dtype) * self.weight
+
+    model = Decoder(1, 1, bias=False, device="cuda", dtype=dtype)
+    model.weight.fill_(1)
+    decoder = Qwen3OmniFirstFrameDecoder(model, sample_rate=24000)
+    decoder.capture()
+    codes = torch.arange(rows, device="cuda").reshape(-1, 1)
+    first = decoder.decode(codes)
+    second = decoder.decode(codes + 32)
+    assert first.dtype == second.dtype == torch.float32
+    torch.testing.assert_close(first, codes.float(), rtol=0, atol=0)
+    torch.testing.assert_close(second, (codes + 32).float(), rtol=0, atol=0)

tests/model_executor/models/qwen3_omni/test_qwen3_omni_code2wav_cudnn.py

@@ -0,0 +1,87 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import Qwen3OmniMoeCode2WavConfig
+
+from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_code2wav import Qwen3OmniMoeCode2Wav
+
+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"),
+]
+
+
+def _codec(enabled):
+    config = Qwen3OmniMoeCode2WavConfig(
+        hidden_size=32,
+        intermediate_size=64,
+        num_hidden_layers=1,
+        num_attention_heads=4,
+        num_key_value_heads=4,
+        head_dim=8,
+        decoder_dim=32,
+        upsample_rates=[2, 2],
+        upsampling_ratios=[2],
+        codebook_size=16,
+        num_quantizers=2,
+    )
+    config._attn_implementation = "sdpa"
+    vllm_config = SimpleNamespace(
+        model_config=SimpleNamespace(
+            hf_config=config, stage_connector_config={"extra": {"codec_cudnn_benchmark": enabled}}
+        )
+    )
+    model = Qwen3OmniMoeCode2Wav(vllm_config=vllm_config).cuda().eval()
+    model.precompute_snake_caches()
+    return model
+
+
+def _flags():
+    c = torch.backends.cudnn
+    return c.enabled, c.benchmark, c.benchmark_limit, c.deterministic, c.allow_tf32
+
+
+@torch.inference_mode()
+def test_autotuned_codec_graph_replay_restores_flags():
+    torch.manual_seed(21)
+    model = _codec(True)
+    codes = torch.randint(0, 16, (2, 2, 8), device="cuda")
+    original = _flags()
+    for _ in range(3):
+        model(codes)
+    assert _flags() == original
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        captured = model(codes)
+    assert _flags() == original
+    codes.random_(0, 16)
+    graph.replay()
+    expected = model(codes)
+    torch.testing.assert_close(captured, expected)
+    assert _flags() == original
+
+
+@pytest.mark.parametrize("enabled", [False, True])
+@torch.inference_mode()
+def test_codec_failure_restores_cudnn_settings(mocker, enabled):
+    model = _codec(enabled)
+    codes = torch.randint(0, 16, (1, 2, 8), device="cuda")
+    original = _flags()
+
+    def fail(hidden):
+        if enabled:
+            assert torch.backends.cudnn.benchmark
+            assert torch.backends.cudnn.benchmark_limit == 10
+        else:
+            assert _flags() == original
+        raise RuntimeError("decode failed")
+
+    mocker.patch.object(model, "_decode_waveform", side_effect=fail)
+    with pytest.raises(RuntimeError, match="decode failed"):
+        model(codes)
+    assert _flags() == original

tests/model_executor/models/qwen3_omni/test_qwen3_omni_code2wav_snake.py

@@ -0,0 +1,88 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import copy
+from types import SimpleNamespace
+
+import pytest
+import torch
+from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import Qwen3OmniMoeCode2WavConfig
+from transformers.models.qwen3_omni_moe.modeling_qwen3_omni_moe import (
+    Qwen3OmniMoeCode2WavDecoderBlock,
+    Qwen3OmniMoeSnakeBeta,
+)
+
+from vllm_omni.model_executor.models.common.snake_activation import SnakeBeta
+from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_code2wav import Qwen3OmniMoeCode2Wav, use_fused_snake
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _blocks():
+    config = Qwen3OmniMoeCode2WavConfig(decoder_dim=32, upsample_rates=[2, 2])
+    torch.manual_seed(0)
+    blocks = torch.nn.ModuleList(Qwen3OmniMoeCode2WavDecoderBlock(config, i) for i in range(2)).eval()
+    with torch.no_grad():
+        for module in blocks.modules():
+            if isinstance(module, Qwen3OmniMoeSnakeBeta):
+                module.alpha.uniform_(-0.5, 0.5)
+                module.beta.uniform_(-0.5, 0.5)
+    return blocks
+
+
+def test_decoder_blocks_swap_every_hf_snake_and_keep_weight_names():
+    blocks = _blocks()
+    hf_count = sum(isinstance(m, Qwen3OmniMoeSnakeBeta) for m in blocks.modules())
+    keys = set(blocks.state_dict())
+
+    assert use_fused_snake(blocks) == hf_count > 0
+    assert not any(isinstance(m, Qwen3OmniMoeSnakeBeta) for m in blocks.modules())
+    assert sum(isinstance(m, SnakeBeta) for m in blocks.modules()) == hf_count
+    assert set(blocks.state_dict()) == keys  # checkpoint weights load unchanged
+
+
+def test_fused_snake_blocks_match_hf_blocks():
+    reference = _blocks()
+    fused = copy.deepcopy(reference)
+    use_fused_snake(fused)
+    fused.load_state_dict(reference.state_dict())
+    x = torch.randn(2, 32, 20)
+    with torch.no_grad():
+        expected = x
+        actual = x
+        for ref_block, fused_block in zip(reference, fused):
+            expected = ref_block(expected)
+            actual = fused_block(actual)
+    torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-5)
+
+
+@pytest.mark.parametrize(
+    ("option", "device", "cuda", "fused"),
+    [
+        (False, "cuda", True, False),
+        (True, "cuda", True, True),
+        (True, "cpu", True, False),
+        (True, "cuda", False, False),
+    ],
+)
+def test_decoder_block_fusion_requires_cuda_opt_in(monkeypatch, option, device, cuda, fused):
+    from vllm_omni.platforms import current_omni_platform
+
+    monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: cuda)
+    config = Qwen3OmniMoeCode2WavConfig(
+        hidden_size=32,
+        intermediate_size=64,
+        num_hidden_layers=1,
+        num_attention_heads=2,
+        num_key_value_heads=2,
+        decoder_dim=32,
+    )
+    model = Qwen3OmniMoeCode2Wav(
+        vllm_config=SimpleNamespace(
+            model_config=SimpleNamespace(
+                hf_config=config, stage_connector_config={"extra": {"codec_fused_snake": option}}
+            ),
+            device_config=SimpleNamespace(device=device),
+        )
+    )
+    assert any(isinstance(m, Qwen3OmniMoeSnakeBeta) for m in model.decoder.modules()) is (not fused)

tests/model_executor/models/qwen3_omni/test_qwen3_omni_first_audio.py

@@ -0,0 +1,208 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import pytest
+import torch
+from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import Qwen3OmniMoeCode2WavConfig
+
+from vllm_omni.model_executor.models.qwen3_omni.first_frame_decoder import Qwen3OmniFirstFrameDecoder
+from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni import Qwen3OmniMoeForConditionalGeneration
+from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_code2wav import (
+    Qwen3OmniMoeCode2Wav,
+    plan_decode_groups,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_Q = 4
+_CODEBOOK = 16
+
+
+def _code2wav(quantizers=_Q) -> Qwen3OmniMoeCode2Wav:
+    torch.manual_seed(0)
+    config = Qwen3OmniMoeCode2WavConfig(
+        codebook_size=_CODEBOOK,
+        num_quantizers=quantizers,
+        hidden_size=32,
+        intermediate_size=64,
+        num_hidden_layers=2,
+        num_attention_heads=2,
+        num_key_value_heads=2,
+        decoder_dim=32,
+        sliding_window=8,
+    )
+    model = Qwen3OmniMoeCode2Wav(vllm_config=SimpleNamespace(model_config=SimpleNamespace(hf_config=config)))
+    return model.eval()
+
+
+def test_first_frame_pcm_equals_code2wav_streaming_chunk0():
+    code2wav = _code2wav()
+    decoder = Qwen3OmniFirstFrameDecoder(code2wav, sample_rate=24000)
+    codes = torch.randint(0, _CODEBOOK, (3, _Q))
+    with torch.inference_mode():
+        for row in range(codes.shape[0]):
+            pcm = decoder.decode(codes[row : row + 1])
+            chunk0 = code2wav.chunked_decode_streaming(
+                codes[row].reshape(1, _Q, 1), left_context_size=[0], seq_token_counts=[_Q]
+            )[0]
+            assert pcm.dtype == torch.float32
+            assert torch.equal(pcm[0], chunk0.reshape(-1).float())
+            assert 0 < pcm.shape[-1] < int(code2wav.total_upsample)
+        batched = decoder.decode(codes)
+    assert batched.shape == (3, pcm.shape[-1])
+
+
+def test_first_frame_is_a_prefix_of_a_longer_first_chunk():
+    code2wav = _code2wav()
+    decoder = Qwen3OmniFirstFrameDecoder(code2wav, sample_rate=24000)
+    codes = torch.randint(0, _CODEBOOK, (_Q, 4))
+    with torch.inference_mode():
+        first = decoder.decode(codes[:, :1].T)[0]
+        chunk0 = code2wav.chunked_decode_streaming(
+            codes.reshape(1, _Q, 4), left_context_size=[0], seq_token_counts=[4 * _Q]
+        )[0]
+    torch.testing.assert_close(chunk0.reshape(-1)[: first.numel()].float(), first)
+
+
+def test_chunk_ramp_adds_exact_code2wav_graph_sizes(monkeypatch):
+    import vllm_omni.model_executor.models.qwen3_tts.cuda_graph_decoder_wrapper as wrapper_module
+
+    real = wrapper_module.CUDAGraphDecoderWrapper
+    wrapper = MagicMock(compute_capture_sizes=real.compute_capture_sizes)
+    monkeypatch.setattr(wrapper_module, "CUDAGraphDecoderWrapper", wrapper)
+    extra = {"codec_chunk_frames": 25, "codec_left_context_frames": 25, "codec_chunk_ramp": [1, 2, 4, 8, 16, 25]}
+    code2wav = SimpleNamespace(config=SimpleNamespace(num_quantizers=_Q))
+    code2wav.enable_cudagraph = lambda **kw: Qwen3OmniMoeCode2Wav.enable_cudagraph(
+        code2wav, device=torch.device("cuda"), **kw
+    )
+    model = object.__new__(Qwen3OmniMoeForConditionalGeneration)
+    torch.nn.Module.__init__(model)
+    model.code2wav = code2wav
+    model.vllm_config = SimpleNamespace(
+        model_config=SimpleNamespace(async_chunk=True, enforce_eager=False, stage_connector_config={"extra": extra})
+    )
+    for batches in ([], [2, 4, 8]):
+        extra["codec_graph_batch_sizes"] = batches
+        Qwen3OmniMoeForConditionalGeneration._maybe_enable_code2wav_cudagraph(model)
+        args = wrapper.call_args.kwargs
+        sizes = args["capture_sizes"]
+        assert {1, 3, 7, 15, 31, 50} <= set(sizes)
+        assert set(real.compute_capture_sizes(codec_chunk_frames=25, codec_left_context_frames=25)) <= set(sizes)
+        assert set(args["extra_capture_shapes"]) == {
+            (b, size) for b in batches for size in sizes if size <= 50 and b * size <= 325
+        }
+    assert code2wav._streaming_batch_sizes[50] == [1, 2, 4]
+    extra.pop("codec_chunk_ramp")
+    extra.pop("codec_graph_batch_sizes")
+    Qwen3OmniMoeForConditionalGeneration._maybe_enable_code2wav_cudagraph(model)
+    assert wrapper.call_args.kwargs["capture_sizes"] is None
+
+
+_SIZES = [1, 2, 3, 4, 7, 8, 15, 16, 25, 31, 32, 50, 64]
+
+
+def _bucket(length):
+    return next((size for size in _SIZES if size >= length), None)
+
+
+def test_decode_groups_keep_a_uniform_batch_whole_and_split_by_the_largest_graph():
+    assert plan_decode_groups([50] * 6, _bucket, lambda size: [1, 2, 4, 8]) == [(list(range(6)), 50)]
+    calls = plan_decode_groups([25] * 20, _bucket, lambda size: [1, 2, 4, 8, 16])
+    assert [len(rows) for rows, _length in calls] == [16, 4]
+    calls = plan_decode_groups([1] * 8 + [50] * 8, _bucket, lambda size: [1, 2, 4, 8] if size < 50 else [1, 2, 4])
+    assert [(len(rows), length) for rows, length in calls] == [(8, 1), (4, 50), (4, 50)]
+    assert plan_decode_groups([], _bucket, lambda size: [1, 2]) == []
+
+
+@pytest.mark.parametrize("frames", [65, 129])
+def test_decode_groups_beyond_largest_graph_use_eager(monkeypatch, frames):
+    from vllm_omni.model_executor.models.qwen3_tts.cuda_graph_decoder_wrapper import CUDAGraphDecoderWrapper
+
+    wrapper = CUDAGraphDecoderWrapper.__new__(CUDAGraphDecoderWrapper)
+    wrapper.enabled = wrapper._warmed_up = True
+    wrapper._bucket_sizes = [8, 32, 64]
+    wrapper._compiled_shapes = {(1, 64), (8, 64)}
+    wrapper._compiled_graphs = wrapper.graphs = dict.fromkeys(wrapper._compiled_shapes, MagicMock())
+    wrapper.decoder = MagicMock(side_effect=lambda codes: codes.float() * 2)
+    monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False)
+    groups = plan_decode_groups([frames, frames], wrapper._get_padded_size, lambda _: (1,))
+    assert groups == [([0], frames), ([1], frames)]
+    codes = torch.arange(2 * frames).reshape(2, 1, frames)
+    for rows, width in groups:
+        actual = wrapper.decode(codes[rows, :, :width])
+        torch.testing.assert_close(actual, codes[rows].float() * 2, rtol=0, atol=0)
+    assert wrapper.decoder.call_count == 2
+    for graph in (*wrapper.graphs.values(), *wrapper._compiled_graphs.values()):
+        graph.replay.assert_not_called()
+
+
+class _PaddingGraphs:
+    def __init__(self, model, batch_sizes):
+        self.model = model
+        self.batch_sizes = sorted(batch_sizes)
+        self.calls = []
+
+    def _get_padded_size(self, length):
+        return _bucket(length)
+
+    def decode(self, codes):
+        rows, _q, frames = codes.shape
+        size = _bucket(frames)
+        batch = next(b for b in self.batch_sizes if b >= rows)
+        self.calls.append((rows, frames))
+        padded = torch.zeros(batch, codes.shape[1], size, dtype=codes.dtype)
+        padded[:rows, :, :frames] = codes
+        out = self.model(padded)
+        return out[:rows, :, : out.shape[-1] - (size - frames) * int(self.model.total_upsample)].clone()
+
+
+def test_grouped_streaming_decode_matches_one_padded_batch():
+    code2wav = _code2wav()
+    lengths = [1, 3, 50, 7, 31]
+    lefts = [0, 1, 25, 3, 15]
+    codes = torch.zeros(len(lengths), _Q, max(lengths), dtype=torch.long)
+    for row, length in enumerate(lengths):
+        codes[row, :, :length] = torch.randint(0, _CODEBOOK, (_Q, length))
+    counts = [length * _Q for length in lengths]
+    with torch.inference_mode():
+        reference = code2wav.chunked_decode_streaming(codes, left_context_size=lefts, seq_token_counts=counts)
+        graphs = _PaddingGraphs(code2wav, [1, 2, 4])
+        code2wav._cudagraph_wrapper, code2wav._cudagraph_enabled = graphs, True
+        code2wav._streaming_batch_sizes = {size: [1, 2, 4] for size in _SIZES}
+        grouped = code2wav.chunked_decode_streaming(codes, left_context_size=lefts, seq_token_counts=counts)
+
+    for want, got in zip(reference, grouped, strict=True):
+        assert want.shape == got.shape
+        torch.testing.assert_close(got, want, rtol=1e-4, atol=1e-5)
+
+
+@pytest.mark.parametrize(("frames", "quantizers"), [(1, 16), (4, 16), (4, 4)])
+def test_delivered_first_frame_is_removed_from_code2wav_output(frames, quantizers):
+    from vllm_omni.data_entry_keys import FIRST_AUDIO_REQUIRED_KEY
+
+    codec = _code2wav(quantizers)
+    model = object.__new__(Qwen3OmniMoeForConditionalGeneration)
+    torch.nn.Module.__init__(model)
+    model.model_stage, model.code2wav, model.code2wav_config = "code2wav", codec, codec.config
+    model.generate_audio = lambda codes, left, counts: codec.chunked_decode_streaming(
+        codes, left_context_size=left, seq_token_counts=counts
+    )
+    codes = torch.randint(0, _CODEBOOK, (2, quantizers, frames))
+    with torch.inference_mode():
+        reference = model.generate_audio(codes, [0, 0], [quantizers * frames] * 2)
+        first = Qwen3OmniFirstFrameDecoder(codec, sample_rate=24000).decode(codes[0, :, :1].T)[0]
+        out = model.forward(
+            input_ids=codes.flatten(),
+            positions=None,
+            seq_token_counts=[quantizers * frames] * 2,
+            runtime_additional_information=[
+                {"meta": {"left_context_size": 0, "first_audio": True}},
+                {"meta": {"left_context_size": 0, "first_audio": False}},
+            ],
+        ).multimodal_outputs
+    torch.testing.assert_close(torch.cat([first, out["model_outputs"][0].flatten()]), reference[0].flatten())
+    torch.testing.assert_close(out["model_outputs"][1], reference[1].reshape(1, -1))
+    assert [bool(flag) for flag in out[FIRST_AUDIO_REQUIRED_KEY]] == [True, False]

tests/model_executor/models/qwen3_omni/test_qwen3_omni_mrv2.py

@@ -0,0 +1,116 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+from torch import nn
+
+from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni import Qwen3OmniMoeForConditionalGeneration
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_H = 4
+_Q = 3
+
+
+class _Talker(nn.Module):
+    num_code_groups = _Q
+
+    def __init__(self):
+        super().__init__()
+        self.text_projection = nn.Identity()
+        self.code_predictor = SimpleNamespace(
+            _top_k=50,
+            _top_p=0.8,
+            set_sampling_params=lambda **kwargs: setattr(self, "sampling", kwargs),
+        )
+
+
+def _talker(*, async_chunk: bool = True) -> Qwen3OmniMoeForConditionalGeneration:
+    model = object.__new__(Qwen3OmniMoeForConditionalGeneration)
+    nn.Module.__init__(model)
+    model.model_stage = "talker"
+    model.talker = _Talker()
+    model.vllm_config = SimpleNamespace(model_config=SimpleNamespace(async_chunk=async_chunk, use_v2_model_runner=True))
+    model.tts_eos_embed = torch.full((_H,), -1.0)
+    model.tts_pad_embed = torch.full((_H,), -2.0)
+    model._codec_codebook_size = 2048
+    return model
+
+
+def _rows(*values: float) -> torch.Tensor:
+    return torch.tensor([[value] * _H for value in values])
+
+
+def _step(model, payload: dict) -> torch.Tensor:
+    update: dict = {}
+    text_step = model._thinker_decode_to_talker_decode(payload, torch.device("cpu"), update)
+    for key, value in update.items():
+        payload.setdefault(key, {}).update(value)
+    return text_step.reshape(-1)
+
+
+def test_mtp_frame_valid_rejects_codec_eos_and_special_ids():
+    layer0 = torch.tensor([0, 2047, 2048, 2150, -1])
+    assert _talker().mtp_frame_valid(layer0).tolist() == [True, True, False, False, False]
+
+
+def test_decode_rows_are_consumed_one_per_step_in_arrival_order():
+    model = _talker()
+    # Two rows were already pending (cached at prefill); two more arrive at once.
+    payload = {"embed": {"cached_decode": _rows(1, 2), "decode": _rows(3, 4)}, "meta": {}}
+    steps = [_step(model, payload)[0].item()]
+    for arrival in (None, _rows(5), None, None, None, None):
+        payload["embed"]["decode"] = arrival
+        steps.append(_step(model, payload)[0].item())
+    assert steps == [1, 2, 3, 4, 5, -1, -2]
+    assert payload["embed"]["cached_decode"] is None
+
+
+@pytest.mark.parametrize(("v2", "resumable"), [(True, True), (False, False), (False, True)])
+def test_v1_and_resumable_requests_keep_the_indexed_cache_path(v2, resumable):
+    model = _talker()
+    model.vllm_config.model_config.use_v2_model_runner = v2
+    payload = {"embed": {"cached_decode": _rows(1, 2, 3), "decode": None}, "meta": {"resumable": resumable}}
+    payload["meta"]["num_processed_tokens"] = 2
+    assert _step(model, payload)[0].item() == 3
+
+
+def test_make_omni_output_aligns_codes_to_request_spans():
+    model = _talker()
+    hidden = torch.zeros((8, _H))  # 5 real tokens + graph padding
+    decode_codes = torch.tensor([[5, 6, 7]])
+    buffers = [
+        {"req_id": "_warmup_0_"},  # no preprocess, no codes
+        {"req_id": "a", "codes": {"audio": torch.zeros((3, _Q), dtype=torch.long)}},
+        {"req_id": "b", "codes": {"audio": decode_codes}},
+    ]
+    out = model.make_omni_output(
+        hidden, model_intermediate_buffer=buffers, request_token_spans=[(0, 1), (1, 4), (4, 5)]
+    )
+    codes = out.multimodal_outputs["codes"]["audio"]
+    assert codes.shape == (5, _Q)
+    assert codes[4].tolist() == [5, 6, 7]
+    assert out.text_hidden_states.shape[0] == 5
+
+
+def test_make_omni_output_without_spans_keeps_v1_concat():
+    model = _talker()
+    buffers = [{"codes": {"audio": torch.tensor([[1, 2, 3]])}}, {"codes": {"audio": torch.tensor([[4, 5, 6]])}}]
+    out = model.make_omni_output(torch.zeros((4, _H)), model_intermediate_buffer=buffers)
+    assert out.multimodal_outputs["codes"]["audio"].tolist() == [[1, 2, 3], [4, 5, 6]]
+    assert out.text_hidden_states.shape[0] == 2
+
+
+def test_eager_output_is_token_major_with_invalid_rows():
+    model = _talker()
+    model.eager_frames_active = True
+    hidden = torch.zeros((6, _H))
+    buffers = [{"req_id": "a"}, {"req_id": "b"}]
+    out = model.make_omni_output(hidden, model_intermediate_buffer=buffers, request_token_spans=[(0, 3), (3, 4)])
+    mm = out.multimodal_outputs
+    assert mm["codes"]["audio"].shape == (4, _Q)
+    assert mm["meta"]["codec_frame_valid"].tolist() == [0, 0, 0, 0]
+    assert out.text_hidden_states is hidden

tests/model_executor/models/qwen3_omni/test_qwen3_omni_pp_capture.py

@@ -104,6 +104,7 @@ def test_factory_declares_only_captures_from_earlier_ranks(monkeypatch, start, e
     if staged:
         model_config.model_stage = "thinker"
     vllm_config = SimpleNamespace(model_config=model_config, with_hf_config=lambda *args, **kwargs: None)
+    vllm_config.speculative_config = None
     thinker = nn.Module()
     thinker.weight = nn.Parameter(torch.ones(1))
     thinker.language_model = SimpleNamespace(model=SimpleNamespace(start_layer=start))

tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py

@@ -1432,3 +1432,24 @@ def test_shared_predictor_ignores_tts_fast_path_options(mocker, loaded_target_cl
     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)
+
+
+def test_tts_cached_predictor_keeps_its_own_step_dispatch(mocker: MockerFixture, loaded_target_classes) -> None:
+    """The TTS cache must not enter Omni's uninitialized incremental loop."""
+    _, _, wrapper, _, _ = loaded_target_classes
+    common_mod = sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"]
+    mocker.patch.object(common_mod.current_omni_platform, "is_cuda", return_value=True)
+    mocker.patch.object(common_mod.current_omni_platform, "is_npu", return_value=False)
+    mocker.patch.object(wrapper, "_stage_connector_extra_config", return_value={"code_predictor_kv_cache": True})
+    cp_config, talker_config = _make_tiny_config(loaded_target_classes)
+    predictor = wrapper(vllm_config=_make_vllm_config(mocker), config=cp_config, talker_config=talker_config)
+    assert predictor._kv_requested
+    predictor._model_dtype = torch.float32
+    predictor._codec_embeds_list = list(predictor.model.codec_embedding)
+    predictor._bucket_sizes = [1]
+    mocker.patch.object(predictor, "_setup_compile")
+    step = mocker.patch.object(predictor, "_predict_step_logits", return_value=torch.ones(1, cp_config.vocab_size))
+    codes = predictor(torch.zeros(1, dtype=torch.long), torch.zeros(1, 32), torch.zeros(1, 32), do_sample=False)
+    assert step.call_count == cp_config.num_code_groups - 1
+    assert codes.shape == (1, cp_config.num_code_groups)
+    assert predictor._kv_buf is None

tests/model_executor/models/test_talker_first_audio.py

@@ -0,0 +1,70 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from types import SimpleNamespace
+
+import pytest
+
+from vllm_omni.model_executor.models.qwen3_omni.first_frame_decoder import talker_first_audio_enabled
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.mark.parametrize(
+    ("env", "extra", "async_chunk", "v2", "expected"),
+    [
+        (None, {}, True, True, False),
+        ("1", {}, True, True, False),
+        (None, {"talker_first_audio": True, "codec_chunk_ramp": [1, 2]}, True, True, True),
+        ("0", {"talker_first_audio": True, "codec_chunk_ramp": [1, 2]}, True, True, False),
+        (None, {"talker_first_audio": True, "codec_chunk_ramp": [4, 8]}, True, True, False),
+        (None, {"talker_first_audio": True, "initial_codec_chunk_frames": 1}, True, True, True),
+        (None, {"talker_first_audio": True, "initial_codec_chunk_frames": 4}, True, True, False),
+        (None, {"talker_first_audio": True, "codec_chunk_ramp": [1]}, True, True, False),
+        (None, {"talker_first_audio": True, "codec_chunk_ramp": [1, 2]}, False, True, False),
+        (None, {"talker_first_audio": True, "codec_chunk_ramp": [1, 2]}, True, False, False),
+    ],
+)
+def test_talker_first_audio_default_and_conditions(monkeypatch, env, extra, async_chunk, v2, expected):
+    if env is None:
+        monkeypatch.delenv("VLLM_OMNI_TALKER_FIRST_AUDIO", raising=False)
+    else:
+        monkeypatch.setenv("VLLM_OMNI_TALKER_FIRST_AUDIO", env)
+    from vllm_omni.platforms import current_omni_platform
+
+    monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: True)
+    config = SimpleNamespace(
+        model_config=SimpleNamespace(
+            async_chunk=async_chunk, use_v2_model_runner=v2, stage_connector_config={"extra": extra}
+        ),
+        device_config=SimpleNamespace(device="cuda"),
+        parallel_config=SimpleNamespace(
+            tensor_parallel_size=1, pipeline_parallel_size=1, distributed_executor_backend=None
+        ),
+        cache_config=SimpleNamespace(enable_prefix_caching=False),
+    )
+    assert talker_first_audio_enabled(config) is expected
+    for section, field, value in [
+        ("device_config", "device", "cpu"),
+        ("parallel_config", "tensor_parallel_size", 2),
+        ("parallel_config", "pipeline_parallel_size", 2),
+        ("parallel_config", "distributed_executor_backend", "mp"),
+        ("cache_config", "enable_prefix_caching", True),
+    ]:
+        with monkeypatch.context() as patch:
+            patch.setattr(getattr(config, section), field, value)
+            assert not talker_first_audio_enabled(config)
+    monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: False)
+    assert not talker_first_audio_enabled(config)
+
+
+@pytest.mark.parametrize("option", [False, True])
+def test_tts_requires_connector_opt_in(monkeypatch, option):
+    from vllm_omni.model_executor.models.qwen3_tts import first_audio
+    from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_code_predictor_vllm import (
+        Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM as Predictor,
+    )
+
+    monkeypatch.setattr(Predictor, "_stage_connector_extra_config", lambda _: {"talker_first_audio": option})
+    monkeypatch.setattr(first_audio, "supports_talker_first_audio", lambda _: True)
+    assert first_audio.talker_first_audio_enabled(None) is option