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

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

.buildkite/common/ci_source_file_dependencies.yml

@@ -16,6 +16,9 @@ x-qwen3-tts: &qwen3_tts
   - vllm_omni/worker/gpu_model_runner.py
   - vllm_omni/model_executor/models/qwen3_tts/
   - vllm_omni/model_executor/models/common/qwen3_code_predictor.py
+  - vllm_omni/model_executor/models/common/fused_sampling.py
+  - vllm_omni/model_executor/models/common/short_kv_attention.py
+  - vllm_omni/model_executor/models/common/talker_first_audio.py
   - vllm_omni/model_executor/models/common/snake_activation.py
   - vllm_omni/model_executor/models/common/alias_free_activation.py
   - vllm_omni/model_executor/stage_input_processors/qwen3_tts.py
@@ -85,11 +88,16 @@ x-qwen3-omni: &qwen3_omni
   - vllm_omni/entrypoints/openai/serving_chat.py
   - vllm_omni/model_executor/models/qwen3_omni/
   - vllm_omni/model_executor/models/common/qwen3_code_predictor.py
+  - vllm_omni/model_executor/models/common/fused_sampling.py
+  - vllm_omni/model_executor/models/common/short_kv_attention.py
+  - vllm_omni/model_executor/models/common/talker_first_audio.py
   - vllm_omni/model_executor/models/common/snake_activation.py
   - vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py
   - vllm_omni/model_executor/models/qwen3_tts/cuda_graph_decoder_wrapper.py
   - vllm_omni/model_executor/stage_input_processors/qwen3_omni.py
   - vllm_omni/deploy/qwen3_omni_moe.yaml
+  - vllm_omni/deploy/qwen3_omni_moe_mrv2.yaml
+  - vllm_omni/deploy/qwen3_omni_moe_mrv2_h200.yaml
   - vllm_omni/deploy/qwen3_omni_moe_thinking.yaml
 
 x-wan22: &wan22

docs/design/index.md

@@ -116,4 +116,5 @@ The pre-#5137 pages are preserved in the
 [legacy module archive](module/archive/README.md) for historical reference and
 are not active design contracts.
 
+- [Qwen3-Omni MRv2 performance](qwen3_omni_mrv2_performance.md)
 - [MiniCPM-o 4.5 turn-mode MRv2 performance](minicpm_o45_mrv2_performance.md)

docs/design/qwen3_omni_mrv2_performance.md

@@ -0,0 +1,22 @@
+# Qwen3-Omni MRv2 performance
+
+This model adaptation builds on #8184's existing stage runtime and output
+contracts. It adds Talker CPU metadata parsing, sampled-token embedding
+handoff, an optional first-frame decoder, length-grouped Code2Wav graph
+buckets, and opt-in incremental short-KV code prediction. The H200 profile
+also enables conservative fused sampling and bounded cuDNN algorithm search.
+The shared predictor keeps the upstream re-prefill helper and its default
+sampling behavior; short KV is explicitly selected by the deployment profile.
+
+Use `vllm_omni/deploy/qwen3_omni_moe_mrv2_h200.yaml` for the experimental H200
+profile. Preserve a control with `code_predictor_fused_sampling=false` and
+`codec_cudnn_benchmark=false` when C64 P99 is the primary requirement.
+cuDNN search restores process settings on exit and has startup cost. Kernel
+arithmetic is not bitwise-equivalent; the first-frame decoder uses additional
+model memory. Full-duplex performance is not established by turn-mode tests.
+
+First-frame delivery requires `talker_first_audio: true` in connector extra,
+a one-frame initial codec chunk, CUDA, a local single-rank Talker and prefix caching disabled.
+Other deployments retain the regular Code2Wav path.
+`codec_fused_snake` explicitly opts CUDA profiles into fused decoder blocks;
+the base V1 profile retains its original blocks and indexed text cache.

tests/config/test_environment_variables.py

@@ -169,7 +169,7 @@ def test_inventory_matches_reviewed_snapshot_counts():
         EnvironmentVariableCategory.PLATFORM_EXTERNAL: 29,
         EnvironmentVariableCategory.MODEL_SPECIFIC: 75,
         EnvironmentVariableCategory.BENCHMARK_TRANSITIONAL: 21,
-        EnvironmentVariableCategory.INTERNAL: 6,
+        EnvironmentVariableCategory.INTERNAL: 7,
     }
 
     disposition_counts = Counter(

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

tests/model_executor/stage_input_processors/test_qwen3_omni_chunk_ramp.py

@@ -0,0 +1,116 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+from collections import defaultdict
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+import vllm_omni.model_executor.stage_input_processors.qwen3_omni as q3
+from vllm_omni.model_executor.stage_input_processors.chunk_size_utils import ramp_decode_windows
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_Q = 2
+_RAMP = [1, 2, 4, 8, 16, 25]
+
+
+def _manager(**extra):
+    config = {"codec_chunk_frames": 25, "codec_left_context_frames": 25, "initial_codec_chunk_frames": 4, **extra}
+    return SimpleNamespace(
+        code_prompt_token_ids=defaultdict(list),
+        put_req_chunk=defaultdict(int),
+        ramp_chunk_count=defaultdict(int),
+        connector=SimpleNamespace(config={"extra": config}),
+    )
+
+
+def _request(request_id="req"):
+    return SimpleNamespace(external_req_id=request_id, sampling_params=SimpleNamespace(stop_token_ids=[]))
+
+
+def _frame(index: int) -> dict:
+    return {
+        "codes": {"audio": torch.tensor([[index, 1000 + index]], dtype=torch.long)},
+        "meta": {"codec_frame_valid": torch.ones(1, dtype=torch.int8)},
+    }
+
+
+def _stream(manager, num_frames, finish_after_last=True):
+    request = _request()
+    chunks = []
+
+    def emit(payload):
+        if payload is None:
+            return
+        codes = payload.codes.audio.reshape(_Q, -1)
+        chunks.append((codes[0].tolist(), int(payload.meta.left_context_size), bool(payload.meta.finished)))
+        # The transport counts every enqueued chunk.
+        manager.put_req_chunk[request.external_req_id] += 1
+        manager.ramp_chunk_count[request.external_req_id] += 1
+
+    for index in range(num_frames):
+        last = finish_after_last and index == num_frames - 1
+        emit(q3.talker2code2wav_async_chunk(manager, _frame(index), request, is_finished=last))
+    return chunks
+
+
+def test_ramp_decode_windows_include_left_context():
+    assert ramp_decode_windows(_RAMP, 25) == [1, 3, 7, 15, 31, 50]
+    assert ramp_decode_windows([1, 2], 0) == [1, 2]
+
+
+def test_ramp_chunks_grow_and_resend_left_context():
+    chunks = _stream(_manager(codec_chunk_ramp=_RAMP), 56 + 25, finish_after_last=False)
+
+    new_frames = [len(window) - left for window, left, _finished in chunks]
+    assert new_frames == [1, 2, 4, 8, 16, 25, 25]
+    assert [left for _window, left, _finished in chunks] == [0, 1, 3, 7, 15, 25, 25]
+    assert [len(window) for window, _left, _finished in chunks] == [1, 3, 7, 15, 31, 50, 50]
+    emitted = 0
+    for window, left, _finished in chunks:
+        # Each window ends at the newest frame and re-sends `left` earlier ones.
+        assert window == list(range(emitted - left, emitted + len(window) - left))
+        emitted += len(window) - left
+    # The first chunk is the stream's first frame alone (no context).
+    assert chunks[0][0] == [0]
+
+
+def test_ramp_flushes_partial_tail_with_left_context():
+    chunks = _stream(_manager(codec_chunk_ramp=_RAMP), 10)
+
+    assert [len(window) - left for window, left, _finished in chunks] == [1, 2, 4, 3]
+    window, left, finished = chunks[-1]
+    assert finished and left == 7 and window == list(range(0, 10))
+
+
+def test_ramp_finish_on_chunk_boundary_leaves_the_marker_to_the_transport():
+    manager = _manager(codec_chunk_ramp=_RAMP)
+    request = _request()
+    for index in range(3):
+        payload = q3.talker2code2wav_async_chunk(manager, _frame(index), request)
+        if payload is not None:
+            manager.ramp_chunk_count[request.external_req_id] += 1
+    # Frames 0 and 1-2 went out as chunks 0 and 1; the finish step has no frame.
+    finish = {"codes": {"audio": torch.zeros((1, _Q), dtype=torch.long)}, "meta": {"codec_frame_valid": torch.zeros(1)}}
+    assert q3.talker2code2wav_async_chunk(manager, finish, request, is_finished=True) is None
+
+
+def test_without_ramp_the_initial_chunk_setting_is_unchanged():
+    chunks = _stream(_manager(), 29 + 3)
+
+    assert [(len(window), left) for window, left, _finished in chunks] == [(4, 0), (29, 4), (28, 25)]
+
+
+def test_first_audio_marker_only_reaches_initial_codec_chunk():
+    manager = _manager(codec_chunk_ramp=[1, 1])
+    for index in range(2):
+        frame = _frame(index)
+        frame["meta"]["first_audio"] = torch.tensor([True])
+        payload = q3.talker2code2wav_async_chunk(manager, frame, _request())
+        assert bool(payload.meta.first_audio) is (index == 0)
+        manager.put_req_chunk["req"] += 1
+        manager.ramp_chunk_count["req"] += 1

tests/model_executor/stage_input_processors/test_qwen3_omni_first_token_handoff.py

@@ -0,0 +1,146 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+from collections import defaultdict
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+import vllm_omni.model_executor.stage_input_processors.qwen3_omni as q3
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+ACCEPT = 24
+HIDDEN = 8
+PROMPT = [151644, 872, 5, 6, 151645, 151644, 77091, 198]
+EOS = 151645
+
+
+class _Manager:
+    def __init__(self) -> None:
+        self.put_req_chunk: defaultdict[str, int] = defaultdict(int)
+        self.request_payload: dict[str, dict] = {}
+        self.config = SimpleNamespace(
+            hf_config=SimpleNamespace(talker_config=SimpleNamespace(accept_hidden_layer=ACCEPT))
+        )
+
+    def process(self, output, request, is_finished=False):
+        payload = q3.thinker2talker_async_chunk(self, output, request, is_finished=is_finished)
+        if payload is not None:
+            self.put_req_chunk[request.external_req_id] += 1
+        return payload
+
+
+def _request(output_token_ids, *, history=True, max_tokens=64):
+    # The native data plane keeps token history only while chunk 0 is pending.
+    return SimpleNamespace(
+        external_req_id="r",
+        prompt_token_ids=list(PROMPT) if history else [],
+        output_token_ids=list(output_token_ids) if history else [],
+        all_token_ids=list(PROMPT) + list(output_token_ids) if history else [],
+        output_token_count=len(output_token_ids),
+        last_output_token_id=output_token_ids[-1] if output_token_ids else None,
+        resumable=False,
+        additional_information=None,
+        sampling_params=SimpleNamespace(all_stop_token_ids={EOS}, max_tokens=max_tokens),
+    )
+
+
+def _embed(token_id: int) -> torch.Tensor:
+    return torch.full((1, HIDDEN), float(token_id), dtype=torch.bfloat16)
+
+
+def _output(token_ids, sampled=None):
+    layers = {
+        0: torch.cat([_embed(t) for t in token_ids]),
+        ACCEPT: torch.randn(len(token_ids), HIDDEN).to(torch.bfloat16),
+    }
+    embed = {"tts_bos": _embed(1), "tts_eos": _embed(2), "tts_pad": _embed(3)}
+    if sampled is not None:
+        embed["sampled"] = sampled
+    return {"hidden_states": {"layers": layers}, "embed": embed}
+
+
+def _run_turn(tokens, *, publish, max_tokens=64, prefill_chunks=(len(PROMPT),)):
+    manager = _Manager()
+    rows, chunks, start = [], 0, 0
+    for i, size in enumerate(prefill_chunks):
+        last = i == len(prefill_chunks) - 1
+        sampled = _embed(tokens[0]) if last else torch.empty(0, dtype=torch.bfloat16)
+        request = _request(tokens[:1] if last else [], max_tokens=max_tokens)
+        payload = manager.process(_output(PROMPT[start : start + size], sampled if publish else None), request)
+        start += size
+        if payload is not None:
+            chunks += 1
+            rows.extend(payload.embed.prefill[len(PROMPT) :])
+    for step in range(1, len(tokens)):
+        history = manager.put_req_chunk["r"] == 0
+        request = _request(tokens[: step + 1], history=history, max_tokens=max_tokens)
+        output = _output([tokens[step - 1]], _embed(tokens[step]) if publish else None)
+        payload = manager.process(output, request)
+        if payload is None:
+            continue
+        chunks += 1
+        rows.extend(payload.embed.prefill[len(PROMPT) :] if payload.embed.prefill is not None else payload.embed.decode)
+    return [int(row[0]) for row in rows], chunks, manager
+
+
+@pytest.mark.parametrize(
+    ("tokens", "max_tokens"),
+    [
+        ([11, 12, 13, EOS], 64),  # ends on a stop token
+        ([11, 12, 13], 3),  # ends on the max_tokens-th token
+        ([11, EOS], 64),
+    ],
+)
+def test_sampled_stream_matches_captured_stream(tokens, max_tokens):
+    late, _chunks, _manager = _run_turn(tokens, publish=False, max_tokens=max_tokens)
+    early, _chunks, manager = _run_turn(tokens, publish=True, max_tokens=max_tokens)
+    # The final token is never processed, so neither path feeds it to the Talker.
+    assert early == late == tokens[:-1]
+    assert q3._SAMPLED_TEXT_STREAM in manager.request_payload["r"]
+
+
+def test_sampled_stream_sends_each_token_one_step_earlier():
+    manager = _Manager()
+    chunk0 = manager.process(_output(PROMPT, sampled=_embed(11)), _request([11]))
+    assert chunk0 is not None  # at the prefill step, not after the next decode step
+    assert torch.equal(chunk0.embed.prefill, torch.cat([_embed(t) for t in [*PROMPT, 11]]))
+    assert chunk0.hidden_states.output.shape[0] == len(PROMPT)
+    assert chunk0.ids.all == [*PROMPT, 11]
+    # Decode step 1 processes 11 and samples 12: it carries 12.
+    chunk1 = manager.process(_output([11], sampled=_embed(12)), _request([11, 12], history=False))
+    assert torch.equal(chunk1.embed.decode, _embed(12))
+
+
+def test_chunked_prefill_uses_sample_of_final_chunk_only():
+    early, chunks, _manager = _run_turn([11, 12, EOS], publish=True, prefill_chunks=(5, 3))
+    assert early == [11, 12] and chunks == 2
+
+
+def test_terminal_first_token_keeps_captured_stream():
+    manager = _Manager()
+    assert manager.process(_output(PROMPT, sampled=_embed(EOS)), _request([EOS])) is None
+    assert q3._SAMPLED_TEXT_STREAM not in manager.request_payload["r"]
+
+
+@pytest.mark.parametrize("prefill_rows", [0, 1, 3])
+def test_missing_sample_on_sampled_stream(prefill_rows):
+    manager = _Manager()
+    manager.process(_output(PROMPT, sampled=_embed(11)), _request([11]))
+    chunk = manager.process(_output([11], sampled=_embed(12)), _request([11, 12], history=False))
+    assert torch.equal(chunk.embed.decode, _embed(12))
+    if not prefill_rows:
+        with pytest.raises(RuntimeError, match="sample embedding missing"):
+            manager.process(_output([12]), _request([11, 12], history=False))
+    else:
+        before = manager.put_req_chunk["r"]
+        for tokens in (PROMPT[:prefill_rows], [11]):
+            partial = _output(tokens, sampled=torch.empty(0))
+            assert manager.process(partial, _request([11, 12], history=False)) is None
+        assert manager.put_req_chunk["r"] == before
+        chunk = manager.process(_output([12], sampled=_embed(13)), _request([11, 12, 13], history=False))
+        assert torch.equal(chunk.embed.decode, _embed(13))

tests/model_executor/stage_input_processors/test_qwen3_omni_streaming_helpers.py

@@ -1047,3 +1047,80 @@ def test_qwen2_5_omni_thinker2talker_full_payload_noop() -> None:
 
     payload = thinker2talker_full_payload(None, {"any": "thing"}, None)
     assert payload is None
+
+
+def _codec_transfer_manager(request_id: str, *, initial_frames: int = 2, chunk_frames: int = 4):
+    return SimpleNamespace(
+        code_prompt_token_ids=defaultdict(list),
+        put_req_chunk=defaultdict(int, {request_id: 0}),
+        connector=SimpleNamespace(
+            config={
+                "extra": {
+                    "initial_codec_chunk_frames": initial_frames,
+                    "codec_chunk_frames": chunk_frames,
+                    "codec_left_context_frames": 4,
+                }
+            }
+        ),
+    )
+
+
+def _codec_request(request_id: str) -> SimpleNamespace:
+    return SimpleNamespace(
+        external_req_id=request_id,
+        sampling_params=SimpleNamespace(stop_token_ids=[2150], stop_token_id=None),
+    )
+
+
+def test_talker2code2wav_async_chunk_reads_last_valid_row_of_multi_row_span() -> None:
+    # MRv2 eager frames: a prefill span carries its first frame in its last row.
+    request_id = "eager_prefill"
+    manager = _codec_transfer_manager(request_id)
+    codes = torch.zeros((5, 3), dtype=torch.long)
+    codes[-1] = torch.tensor([11, 12, 13])
+    valid = torch.tensor([0, 0, 0, 0, 1], dtype=torch.int8)
+
+    payload = q3.talker2code2wav_async_chunk(
+        manager, {"codes": {"audio": codes}, "meta": {"codec_frame_valid": valid}}, _codec_request(request_id)
+    )
+
+    assert payload is None  # below the initial chunk size
+    assert [frame.tolist() for frame in manager.code_prompt_token_ids[request_id]] == [[[11, 12, 13]]]
+
+
+def test_talker2code2wav_async_chunk_skips_invalid_frame_and_flushes_on_finish() -> None:
+    request_id = "eager_eos"
+    manager = _codec_transfer_manager(request_id)
+    manager.code_prompt_token_ids[request_id] = [torch.tensor([[7, 8, 9]])]
+    # The step that sampled codec EOS carries an invalid (placeholder) frame.
+    output = {
+        "codes": {"audio": torch.tensor([[2150, 0, 0]])},
+        "meta": {"codec_frame_valid": torch.tensor([0], dtype=torch.int8)},
+    }
+
+    payload = q3.talker2code2wav_async_chunk(manager, output, _codec_request(request_id), is_finished=True)
+
+    assert len(manager.code_prompt_token_ids[request_id]) == 1
+    assert payload is not None
+    assert payload.codes.audio.tolist() == [7, 8, 9]
+    assert payload.meta.finished.item() is True
+
+
+def test_talker2code2wav_async_chunk_without_validity_keeps_one_row_contract() -> None:
+    # V1 / MRv2 deferred frames: one row per step, zero rows are prefill
+    # placeholders and a stop-token CB0 marks the post-EOS step.
+    request_id = "legacy_rows"
+    manager = _codec_transfer_manager(request_id)
+    request = _codec_request(request_id)
+
+    assert (
+        q3.talker2code2wav_async_chunk(manager, {"codes": {"audio": torch.zeros((4, 3), dtype=torch.long)}}, request)
+        is None
+    )
+    assert q3.talker2code2wav_async_chunk(manager, {"codes": {"audio": torch.tensor([[2150, 1, 1]])}}, request) is None
+    assert manager.code_prompt_token_ids[request_id] == []
+    q3.talker2code2wav_async_chunk(manager, {"codes": {"audio": torch.tensor([[5, 6, 7]])}}, request)
+    payload = q3.talker2code2wav_async_chunk(manager, {"codes": {"audio": torch.tensor([[8, 9, 10]])}}, request)
+
+    assert payload is not None  # two frames fill the initial chunk
+    assert payload.codes.audio.tolist() == [5, 8, 6, 9, 7, 10]

tests/worker_v2/test_omni_model_state.py

@@ -620,6 +620,12 @@ def test_publish_sampled_embeddings_for_rows_whose_sample_is_kept(monkeypatch) -
     assert [tuple(row.shape) for row in sampled] == [(1, 3), (1, 3), (0,), (1, 3)]
     assert [row.tolist() for row in sampled if row.numel()] == [[[v] * 3] for v in (11.0, 12.0, 14.0)]
 
+    batch.is_prefilling_np[:] = True
+    batch.num_computed_prefill_tokens_np[:] = 0
+    batch.num_scheduled_tokens = [1] * 4
+    extra, _done = state.publish_sampled_embeddings(batch, torch.tensor([[11], [12], [13], [14]]))
+    assert all(row.numel() == 0 for row in extra["embed"]["sampled"])
+
     # Speculative steps sample several tokens per row: not published.
     assert state.publish_sampled_embeddings(batch, torch.tensor([[11, 1], [12, 1], [13, 1], [14, 1]])) is None
 

vllm_omni/config/environment_variable_inventory.py

@@ -260,6 +260,7 @@ _BENCHMARK_TRANSITIONAL = (
 )
 
 _INTERNAL = (
+    "VLLM_OMNI_TALKER_FIRST_AUDIO",
     "VLLM_OMNI_CODEC_FIRST_CHUNK_EXPRESS",
     "VLLM_OMNI_CODEC_FIRST_CHUNK_EXPRESS_SLACK_S",
     "VLLM_OMNI_CONNECTOR_RECV_POLL_MS",

vllm_omni/deploy/qwen3_omni_moe_mrv2.yaml

@@ -0,0 +1,50 @@
+# Opt-in CUDA Model Runner V2 profile for Qwen3-Omni-MoE: the Thinker (stage 0),
+# Talker (stage 1) and Code2Wav (stage 2) run on MRv2 with its native data
+# plane. The base profile keeps every stage on V1.
+#
+# On MRv2 the Thinker hands each generated token's embedding to the Talker with
+# the step that sampled it (one Thinker step earlier than its next forward would
+# capture the same embedding) and replays decode in FULL CUDA graphs together
+# with the layers it captures for the Talker.
+#
+# When the stages share one GPU, run them under NVIDIA MPS (see
+# docs/configuration/stage_configs.md); startup warns when no MPS server is found.
+#
+# MRv2 downstream stages support turn-based sessions only: realtime (resumable)
+# requests and duplex sessions need the V1 chunk adapter.
+base_config: qwen3_omni_moe.yaml
+
+connectors:
+  connector_of_shared_memory:
+    extra:
+      # First Code2Wav chunk is the stream's first frame alone, then chunks
+      # double up to the steady codec_chunk_frames (replaces the 4-frame
+      # initial chunk). Code2Wav re-decodes up to codec_left_context_frames
+      # before each chunk, so windows are 1, 3, 7, 15, 31, then 50 frames.
+      # A one-frame first chunk is also what lets the Talker deliver it
+      # itself. Opt in explicitly; VLLM_OMNI_TALKER_FIRST_AUDIO=0 disables it.
+      talker_first_audio: true
+      codec_chunk_ramp: [1, 2, 4, 8, 16, 25]
+      # Code2Wav decodes a streaming step as a few length groups on multi-row
+      # CUDA graphs instead of one eager call padded to its longest window
+      # (with the ramp, over half of those padded frames were waste).
+      codec_graph_batch_sizes: [2, 3, 4, 6, 8, 12, 16]
+      codec_fused_snake: true
+
+stages:
+  - stage_id: 0
+    model_runner: v2
+  - stage_id: 1
+    model_runner: v2
+  - stage_id: 2
+    model_runner: v2
+
+platforms:
+  npu:
+    model_runner: v1
+  xpu:
+    model_runner: v1
+  rocm:
+    model_runner: v1
+  musa:
+    model_runner: v1

vllm_omni/deploy/qwen3_omni_moe_mrv2_h200.yaml

@@ -0,0 +1,27 @@
+# Opt-in single-H200 MRv2 profile, measured at client concurrency 32/64.
+# Run all stages under NVIDIA MPS and warm representative shapes first.
+# Retains Triton MoE and native nucleus reductions, fusing only the final
+# sampling mask/Gumbel selection. cuDNN searches convolution algorithms
+# during warmup; short-KV attention and convolution selection change rounding.
+# This is not a claim of bitwise audio or universal P99/throughput improvement.
+# Evidence: docs/design/qwen3_omni_mrv2_performance.md
+base_config: qwen3_omni_moe_mrv2.yaml
+
+connectors:
+  connector_of_shared_memory:
+    extra:
+      code_predictor_kv_cache: true
+      code_predictor_short_kv: true
+      code_predictor_fused_sampling: true
+      codec_cudnn_benchmark: true
+
+stages:
+  - stage_id: 0
+    devices: "0"
+    gpu_memory_utilization: 0.60
+  - stage_id: 1
+    devices: "0"
+    gpu_memory_utilization: 0.22
+  - stage_id: 2
+    devices: "0"
+    gpu_memory_utilization: 0.10

vllm_omni/model_executor/models/common/fused_sampling.py

@@ -0,0 +1,97 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Fused Gumbel selection after native nucleus probability reductions.
+
+Keep ATen's top-k threshold, sort, softmax and cumulative sum, including
+boundary ties. Random numbers remain indexed by original token id. Only the
+final mask, Gumbel transform and argmax are fused; no vocabulary scatter is
+needed. This avoids changing which tokens survive a nucleus boundary through
+different reduction orders.
+"""
+
+import torch
+from vllm.triton_utils import tl, triton
+
+
+@triton.jit
+def _sorted_nucleus_gumbel(
+    logits,
+    indices,
+    probabilities,
+    cumulative,
+    uniforms,
+    output,
+    vocab: tl.constexpr,
+    logits_stride: tl.constexpr,
+    token_stride: tl.constexpr,
+    uniform_stride: tl.constexpr,
+    top_p: tl.constexpr,
+    block: tl.constexpr,
+):
+    row = tl.program_id(0)
+    offsets = tl.arange(0, block)
+    valid = offsets < vocab
+    values = tl.load(logits + row * logits_stride + offsets * token_stride, valid, other=-float("inf")).to(tl.float32)
+    if top_p < 1.0:
+        ids = tl.load(indices + row * vocab + offsets, valid, other=0)
+        probs = tl.load(probabilities + row * vocab + offsets, valid, other=0)
+        cdf = tl.load(cumulative + row * vocab + offsets, valid, other=0)
+        values = tl.where((cdf - probs) >= top_p, -float("inf"), values)
+    else:
+        ids = offsets
+    noise = tl.load(uniforms + row * uniform_stride + ids, valid, other=0.5)
+    score = values - tl.extra.cuda.libdevice.log(-tl.extra.cuda.libdevice.log(noise))
+    # torch.argmax returns the first original token for a tie or NaN.
+    first_nan = tl.min(tl.where(valid & (score != score), ids, 2147483647), 0)
+    maximum = tl.max(tl.where(valid, score, -float("inf")), 0)
+    winner = tl.min(tl.where(valid & (score == maximum), ids, 2147483647), 0)
+    tl.store(output + row, tl.where(first_nan < 2147483647, first_nan, winner))
+
+
+def sample_top_k_top_p_gumbel(
+    logits: torch.Tensor, uniforms: torch.Tensor, *, top_k: int, top_p: float
+) -> torch.Tensor:
+    """Sample CUDA rows without scattering sorted logits back to the vocabulary.
+
+    Uniforms are caller-owned FP32 values with the same shape as logits; their
+    generation and generator advancement are deliberately outside this kernel.
+    Softmax/scan use the reference ATen operations. The Gumbel logarithms
+    still use libdevice, so this is not a universal bitwise-equivalence claim.
+    """
+    if logits.ndim != 2 or not logits.is_cuda or uniforms.device != logits.device or torch.version.hip is not None:
+        raise ValueError("fused sampling requires two-dimensional CUDA logits and uniforms")
+    if uniforms.shape != logits.shape or uniforms.dtype != torch.float32 or uniforms.stride(1) != 1:
+        raise ValueError("uniforms must be FP32 with the logits shape and contiguous rows")
+    rows, vocab = logits.shape
+    # The stored-mode caller disables top-k for both 0 and negative values
+    # (vLLM commonly uses -1). Keep that behavior when fusion is enabled.
+    top_k = max(top_k, 0)
+    if vocab == 0 or top_k > vocab or not 0 < top_p <= 1:
+        raise ValueError("invalid top-k/top-p sampling parameters")
+    if top_k:
+        threshold = logits.topk(top_k, dim=-1).values[:, -1:]
+        logits = logits.masked_fill(logits < threshold, -float("inf"))
+    indices = probabilities = cumulative = None
+    if top_p < 1:
+        logits, indices = logits.sort(dim=-1, descending=True)
+        probabilities = logits.softmax(dim=-1, dtype=torch.float32)
+        cumulative = probabilities.cumsum(-1)
+    result = torch.empty((rows, 1), dtype=torch.long, device=logits.device)
+    if rows:
+        _sorted_nucleus_gumbel[(rows,)](
+            logits,
+            indices,
+            probabilities,
+            cumulative,
+            uniforms,
+            result,
+            vocab,
+            logits.stride(0),
+            logits.stride(1),
+            uniforms.stride(0),
+            top_p,
+            triton.next_power_of_2(vocab),
+            num_warps=4,
+            enable_fp_fusion=False,
+        )
+    return result

vllm_omni/model_executor/models/common/qwen3_code_predictor.py

@@ -18,6 +18,7 @@ Shared by Qwen3-Omni and Qwen3-TTS talker models.
 from __future__ import annotations
 
 import dataclasses
+import time
 from collections.abc import Iterable, Sequence
 
 import torch
@@ -158,8 +159,8 @@ class CodePredictorAttention(nn.Module):
     """Multi-head self-attention for code predictor.
 
     Uses ``F.scaled_dot_product_attention`` with HF-compatible RoPE and RMSNorm.
-    No KV cache -- the code predictor always re-prefills the full (short)
-    sequence each AR step.
+    The default path re-prefills the short sequence; the opt-in cached path
+    appends only new positions.
 
     Input : [B, seq_len, hidden_size]
     Output: [B, seq_len, hidden_size]
@@ -316,6 +317,53 @@ class CodePredictorAttention(nn.Module):
         attn_out = attn_out.transpose(1, 2).reshape(bsz, seq_len, -1)
         return self.o_proj(attn_out)
 
+    def forward_cached(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
+        k_cache: torch.Tensor,
+        v_cache: torch.Tensor,
+        write_positions: torch.Tensor,
+        key_mask: torch.Tensor,
+    ) -> torch.Tensor:
+        """Attend only the new positions, reusing earlier positions' keys and values.
+
+        ``k_cache``/``v_cache`` are ``[B, num_kv_heads, max_seq, head_dim]``;
+        the new positions' keys and values are written at ``write_positions``
+        (``[n]``). ``key_mask`` (``[1, 1, n, max_seq]``, bool) keeps key
+        positions up to each query's, so slots not yet written in this call are
+        never read and every single-position step has the same shapes.
+        """
+        bsz, seq_len, _ = hidden_states.shape
+        qkv = self.qkv_proj(hidden_states)
+        q_raw, k_raw, v_raw = self._split_qkv(qkv)
+        q = self.q_norm(q_raw.view(bsz, seq_len, self.num_heads, self.head_dim)).transpose(1, 2)
+        k = self.k_norm(k_raw.view(bsz, seq_len, self.num_kv_heads, self.head_dim)).transpose(1, 2)
+        v = v_raw.view(bsz, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2)
+
+        cos, sin = position_embeddings
+        cos = cos.unsqueeze(1)
+        sin = sin.unsqueeze(1)
+        q = (q * cos) + (_rotate_half(q) * sin)
+        k = (k * cos) + (_rotate_half(k) * sin)
+        k_cache.index_copy_(2, write_positions, k.to(k_cache.dtype))
+        v_cache.index_copy_(2, write_positions, v.to(v_cache.dtype))
+        if getattr(self, "use_short_kv_attention", False):
+            from .short_kv_attention import short_kv_attention
+
+            attn_out = short_kv_attention(q, k_cache, v_cache, write_positions, self.scaling)
+        else:
+            attn_out = F.scaled_dot_product_attention(
+                q,
+                k_cache,
+                v_cache,
+                attn_mask=key_mask,
+                scale=self.scaling,
+                enable_gqa=self.is_gqa,
+            )
+        attn_out = attn_out.transpose(1, 2).reshape(bsz, seq_len, -1)
+        return self.o_proj(attn_out)
+
 
 # ===================================================================
 #  MLP
@@ -351,7 +399,7 @@ class CodePredictorMLP(nn.Module):
 
 
 class CodePredictorDecoderLayer(nn.Module):
-    """Transformer decoder layer (SDPA, no KV cache)."""
+    """Transformer decoder layer with default re-prefill and opt-in KV reuse."""
 
     def __init__(self, config, *, prefix: str = "") -> None:
         super().__init__()
@@ -383,9 +431,29 @@ class CodePredictorDecoderLayer(nn.Module):
         hidden_states = residual + hidden_states
         return hidden_states
 
+    def forward_cached(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
+        k_cache: torch.Tensor,
+        v_cache: torch.Tensor,
+        write_positions: torch.Tensor,
+        key_mask: torch.Tensor,
+    ) -> torch.Tensor:
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+        hidden_states = self.self_attn.forward_cached(
+            hidden_states, position_embeddings, k_cache, v_cache, write_positions, key_mask
+        )
+        hidden_states = residual + hidden_states
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        return residual + hidden_states
+
 
 # ===================================================================
-#  Base Transformer Model (re-prefill, no KV cache)
+#  Base Transformer Model (re-prefill, or incremental with a K/V buffer)
 # ===================================================================
 
 
@@ -454,6 +522,29 @@ class CodePredictorBaseModel(nn.Module):
             hidden_states = self.norm(hidden_states)
         return hidden_states.to(input_dtype)
 
+    def forward_cached(
+        self,
+        inputs_embeds: torch.Tensor,
+        position_ids: torch.Tensor,
+        kv_caches: Sequence[tuple[torch.Tensor, torch.Tensor]],
+        write_positions: torch.Tensor,
+        key_mask: torch.Tensor,
+    ) -> torch.Tensor:
+        """``forward`` for the new positions only (``position_ids`` ``[B, n]``), reusing earlier K/V."""
+        input_dtype = inputs_embeds.dtype
+        use_fp32 = input_dtype == torch.float16 and inputs_embeds.device.type not in ("cpu", "npu")
+        if use_fp32:
+            inputs_embeds = inputs_embeds.float()
+        hidden_states = inputs_embeds
+        with torch.amp.autocast(inputs_embeds.device.type, enabled=use_fp32, dtype=torch.float32):
+            position_embeddings = self.rotary_emb(hidden_states, position_ids)
+            for layer, (k_cache, v_cache) in zip(self.layers, kv_caches):
+                hidden_states = layer.forward_cached(
+                    hidden_states, position_embeddings, k_cache, v_cache, write_positions, key_mask
+                )
+            hidden_states = self.norm(hidden_states)
+        return hidden_states.to(input_dtype)
+
     # Mapping from HF checkpoint shard name to the (fused parameter,
     # slice index) pair used to assemble the fused tensor.  See
     # ``load_weights`` below.
@@ -626,8 +717,14 @@ class CodePredictorWrapper(nn.Module):
       4. torch.compile on inner transformer.
       5. Inline sampling (top-k + top-p) -- no custom op overhead.
       6. Optional manual CUDA graph capture per batch-size bucket.
+
+    ``code_predictor_kv_cache`` (CUDA, opt-in) switches to incremental K/V
+    decoding: at large batches the re-prefill FLOPs are not negligible.
     """
 
+    _kv_cache_enabled: bool = False
+    _fused_sampling_enabled: bool = False
+
     def __init__(
         self,
         *,
@@ -672,6 +769,7 @@ class CodePredictorWrapper(nn.Module):
         # Sampling defaults for "stored" mode
         self._top_k: int = 50
         self._top_p: float = 0.8
+        self._do_sample: bool = True
 
         # Lazily initialised state
         self._proj_buf: torch.Tensor | None = None
@@ -686,6 +784,7 @@ class CodePredictorWrapper(nn.Module):
         # ``None`` keeps the legacy power-of-two bucket derivation.
         self._execution_batch_buckets: list[int] | None = None
         prefix_graph_cfg = self._stage_connector_extra_config(vllm_config)
+        self._fused_sampling_enabled = self._parse_bool_config(prefix_graph_cfg.get("code_predictor_fused_sampling"))
         prefix_graphs_requested = self._parse_bool_config(prefix_graph_cfg.get("code_predictor_prefix_graphs"))
         is_npu = current_omni_platform.is_npu()
         # MRv2 captures the whole Talker MTP call. It can therefore capture
@@ -705,15 +804,28 @@ class CodePredictorWrapper(nn.Module):
             prefix_graph_cfg.get("code_predictor_prefix_graph_seq_lens")
         )
         self._prefix_reprefill_seq_lens = tuple(self._prefix_seq_lens(self._num_groups + 1))
+        # Incremental decoding (opt-in, CUDA): each AR step runs only its new
+        # position against a static K/V buffer instead of re-prefilling the
+        # whole prefix (at batch 64 the re-prefill runs 17x the positions).
+        self._kv_cache_enabled = (
+            self._parse_bool_config(prefix_graph_cfg.get("code_predictor_kv_cache")) and current_omni_platform.is_cuda()
+        )
+        self._short_kv_requested = self._parse_bool_config(prefix_graph_cfg.get("code_predictor_short_kv"))
+        self._short_kv_enabled = False
+        self._kv_buf: torch.Tensor | None = None
+        self._compiled_kv_fwd = None
+        # Per AR step: (new positions [n], key mask [1, 1, n, max_seq]).
+        self._kv_steps: dict[int, tuple[torch.Tensor, torch.Tensor]] = {}
 
     def get_input_embeddings(self) -> nn.ModuleList:
         return self.model.get_input_embeddings()
 
-    def set_sampling_params(self, top_k: int = 50, top_p: float = 0.8) -> None:
-        """Configure sampling parameters to maintain consistency with previous implementation."""
+    def set_sampling_params(self, top_k: int = 50, top_p: float = 0.8, do_sample: bool = True) -> None:
+        """Configure "stored"-mode sampling; ``do_sample=False`` selects greedy argmax."""
         self._top_k = top_k
         self._top_p = top_p
-        logger.debug("Sampling parameters updated: top_k=%d, top_p=%.2f", top_k, top_p)
+        self._do_sample = do_sample
+        logger.debug("Sampling parameters updated: top_k=%d, top_p=%.2f, do_sample=%s", top_k, top_p, do_sample)
 
     # ------------------------------------------------------------------
     #  Lazy-init helpers
@@ -730,6 +842,35 @@ class CodePredictorWrapper(nn.Module):
         ):
             return
         self._proj_buf = torch.zeros(bsz, max_seq, self._cp_hidden, dtype=dtype, device=device)
+        if self._kv_cache_enabled:
+            cfg = self.config
+            head_dim = int(getattr(cfg, "head_dim", None) or cfg.hidden_size // cfg.num_attention_heads)
+            self._kv_buf = torch.zeros(
+                int(cfg.num_hidden_layers),
+                2,
+                bsz,
+                int(cfg.num_key_value_heads),
+                max_seq,
+                head_dim,
+                dtype=dtype,
+                device=device,
+            )
+
+    def _kv_views(self, bsz: int) -> list[tuple[torch.Tensor, torch.Tensor]]:
+        kv_buf = self._kv_buf
+        assert kv_buf is not None
+        return [(kv_buf[i, 0, :bsz], kv_buf[i, 1, :bsz]) for i in range(kv_buf.shape[0])]
+
+    def _ensure_kv_steps(self, device: torch.device) -> None:
+        if self._kv_steps:
+            return
+        max_seq = self._num_groups + 1
+        keys = torch.arange(max_seq, device=device)
+        for step in range(1, self._num_groups):
+            # Step 1 runs positions 0 (talker hidden) and 1 (CB0); step s > 1 runs position s.
+            positions = torch.arange(0, 2, device=device) if step == 1 else torch.tensor([step], device=device)
+            key_mask = (keys[None, :] <= positions[:, None]).view(1, 1, positions.numel(), max_seq)
+            self._kv_steps[step] = (positions, key_mask)
 
     def _setup_compile(self) -> None:
         """Lazily set up torch.compile with optional device graph capture."""
@@ -742,6 +883,9 @@ class CodePredictorWrapper(nn.Module):
         self._model_dtype = next(self.model.parameters()).dtype
         self._lm_heads_list = list(self.lm_head)
         self._codec_embeds_list = list(self.model.codec_embedding)
+        if self._model_dtype == torch.float16 or not current_omni_platform.supports_torch_inductor():
+            # The fp16 body runs in fp32 under autocast; keep re-prefill there.
+            self._kv_cache_enabled = False
 
         # Torch 2.13 XPU Dynamo can double-register built-in handlers when
         # spawned workers compile this predictor. Keep this narrow path eager
@@ -768,10 +912,28 @@ class CodePredictorWrapper(nn.Module):
             dynamic=False,
             options={"epilogue_fusion": False},
         )
+        self._short_kv_enabled = (
+            self._kv_cache_enabled
+            and self._short_kv_requested
+            and current_omni_platform.is_cuda()
+            and self._model_dtype == torch.bfloat16
+        )
+        for layer in self.model.layers:
+            layer.self_attn.use_short_kv_attention = self._short_kv_enabled
+        if self._kv_cache_enabled:
+            self._compiled_kv_fwd = torch.compile(
+                self.model.forward_cached,
+                dynamic=self._short_kv_enabled,
+                options={"epilogue_fusion": False},
+            )
         with torch._dynamo.config.patch(cache_size_limit=self._compile_cache_size_limit()):
             self._warmup_buckets()
 
-        if self._wrapper_config.use_cuda_graphs:
+        if self._kv_cache_enabled:
+            # Fifteen short steps per call: they rely on the runner capturing
+            # the whole MTP call (MRv2) rather than on per-step inner graphs.
+            logger.info("code_predictor: torch.compile + incremental K/V decoding")
+        elif self._wrapper_config.use_cuda_graphs:
             self._capture_cuda_graphs()
             logger.info("code_predictor: torch.compile (no epilogue fusion) + CUDA graphs")
         else:
@@ -840,13 +1002,13 @@ class CodePredictorWrapper(nn.Module):
         return row_generators
 
     @classmethod
-    def _sample_codes_gumbel(
+    def _sampling_uniforms(
         cls,
         logits: torch.Tensor,
         generator: _GeneratorLike = None,
         uniforms: torch.Tensor | None = None,
     ) -> torch.Tensor:
-        """Sample ``logits`` via Gumbel-max with optional precomputed noise."""
+        """Draw one FP32 uniform per original token, preserving generator state."""
         if uniforms is None:
             row_generators = cls._normalize_generators(generator, int(logits.shape[0]))
             u = torch.empty_like(logits, dtype=torch.float32)
@@ -862,6 +1024,17 @@ class CodePredictorWrapper(nn.Module):
                     f"logits={tuple(logits.shape)}"
                 )
             u = uniforms
+        return u
+
+    @classmethod
+    def _sample_codes_gumbel(
+        cls,
+        logits: torch.Tensor,
+        generator: _GeneratorLike = None,
+        uniforms: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """Sample ``logits`` via Gumbel-max with optional precomputed noise."""
+        u = cls._sampling_uniforms(logits, generator, uniforms)
         return (logits.float() - torch.log(-torch.log(u))).argmax(dim=-1, keepdim=True)
 
     def _prefix_seq_lens(self, max_seq: int) -> list[int]:
@@ -924,6 +1097,8 @@ class CodePredictorWrapper(nn.Module):
             if needs_full_graph:
                 full_graph_entries += len(prefix_buckets)
             required_entries = full_graph_entries + len(prefix_buckets) * len(prefix_seq_lens)
+        if self._kv_cache_enabled:
+            required_entries = 4  # {first step, later steps} x {batch 1, dynamic batch}
         return max(torch._dynamo.config.cache_size_limit, required_entries)
 
     def _warmup_buckets(self) -> None:
@@ -938,7 +1113,29 @@ class CodePredictorWrapper(nn.Module):
         self._ensure_buffers(device, self._model_dtype, max(self._bucket_sizes))
         proj_buf = self._proj_buf
 
-        if self._prefix_reprefill_enabled:
+        if self._kv_cache_enabled:
+            self._ensure_kv_steps(device)
+            started = time.perf_counter()
+            for bsz in self._bucket_sizes:
+                caches = self._kv_views(bsz)
+                for step in (1, 2):
+                    positions, key_mask = self._kv_steps[step]
+                    lo = 0 if step == 1 else step
+                    hidden = proj_buf[:bsz, lo : step + 1, :]
+                    pos_ids = positions.unsqueeze(0).expand(bsz, -1)
+                    if bsz > 1:
+                        # One graph for every batch >= 2 per step shape (4 compiles, not 2 per bucket):
+                        # a cold per-bucket compile outlasted the stage handshake timeout.
+                        for tensor in (hidden, pos_ids, *(t for kv in caches for t in kv)):
+                            torch._dynamo.mark_dynamic(tensor, 0)
+                    for _ in range(2):
+                        self._compiled_kv_fwd(hidden, pos_ids, caches, positions, key_mask)
+            logger.info(
+                "code_predictor: incremental K/V decoding warmed for buckets %s in %.1f s",
+                self._bucket_sizes,
+                time.perf_counter() - started,
+            )
+        elif self._prefix_reprefill_enabled:
             prefix_seq_lens = self._prefix_reprefill_seq_lens
             needs_full_graph = set(prefix_seq_lens) != set(range(2, max_seq))
             for bsz in self._bucket_sizes:
@@ -1206,9 +1403,29 @@ class CodePredictorWrapper(nn.Module):
             all_codes = torch.empty(bsz, num_groups, dtype=torch.long, device=device)
             all_codes[:, 0] = layer0_code.reshape(bsz)
 
+        kv_caches = None
+        if self._kv_cache_enabled:
+            # Masked slots still enter the attention matmul with zero weight, so
+            # clear them: a non-finite value left by an earlier call would leak.
+            if not self._short_kv_enabled:
+                self._kv_buf[:, :, :padded_bsz].zero_()
+            kv_caches = self._kv_views(padded_bsz)
         # Autoregressive loop: predict layers 1..G-1
         for step in range(1, num_groups):
-            logits = self._predict_step_logits(proj_buf, bsz, padded_bsz, step, is_npu_capturing)
+            if kv_caches is not None:
+                # Step 1 runs positions 0-1; later steps append one position.
+                positions, key_mask = self._kv_steps[step]
+                lo = 0 if step == 1 else step
+                new_hidden = self._compiled_kv_fwd(
+                    proj_buf[:padded_bsz, lo : step + 1, :],
+                    positions.unsqueeze(0).expand(padded_bsz, -1),
+                    kv_caches,
+                    positions,
+                    key_mask,
+                )
+                logits = self._lm_heads_list[step - 1](new_hidden[:bsz, -1, :])
+            else:
+                logits = self._predict_step_logits(proj_buf, bsz, padded_bsz, step, is_npu_capturing)
 
             # Sample next code via Gumbel-max.
             #
@@ -1222,7 +1439,15 @@ class CodePredictorWrapper(nn.Module):
             # defensive than ``multinomial`` around fully-masked/NaN inputs),
             # and the helper below can honor either one batch generator or one
             # generator per seeded row.
-            if stored_mode:
+            if stored_mode and not self._do_sample:
+                code = logits.argmax(dim=-1, keepdim=True)
+            elif stored_mode and self._fused_sampling_enabled and logits.is_cuda and torch.version.hip is None:
+                from .fused_sampling import sample_top_k_top_p_gumbel
+
+                step_uniforms = sample_uniforms[:, step - 1, :] if sample_uniforms is not None else None
+                uniforms = self._sampling_uniforms(logits, sample_generator, step_uniforms)
+                code = sample_top_k_top_p_gumbel(logits, uniforms, top_k=s_top_k, top_p=s_top_p)
+            elif stored_mode:
                 # "stored" mode: top-k -> top-p -> Gumbel-max
                 if s_top_k > 0:
                     topk_vals, _ = logits.topk(s_top_k, dim=-1)

vllm_omni/model_executor/models/common/talker_first_audio.py

@@ -0,0 +1,29 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Shared eligibility; each model separately controls its first-audio opt-in."""
+
+from typing import Any
+
+import torch
+
+from vllm_omni.platforms import current_omni_platform
+
+
+def supports_talker_first_audio(vllm_config: Any) -> bool:
+    """Require CUDA MRv2 streaming, local single-rank execution and no prefix replay.
+
+    Other configurations retain the regular codec path. These conditions do
+    not enable the feature: model-specific options must also accept it.
+    """
+    model = vllm_config.model_config
+    if not (getattr(model, "use_v2_model_runner", False) and getattr(model, "async_chunk", False)):
+        return False
+    parallel = vllm_config.parallel_config
+    return (
+        current_omni_platform.is_cuda()
+        and torch.device(vllm_config.device_config.device).type == "cuda"
+        and parallel.tensor_parallel_size == 1
+        and parallel.pipeline_parallel_size == 1
+        and parallel.distributed_executor_backend in (None, "uni")
+        and not vllm_config.cache_config.enable_prefix_caching
+    )

vllm_omni/model_executor/models/qwen3_omni/first_frame_decoder.py

@@ -0,0 +1,99 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""First codec frame -> PCM inside the Qwen3-Omni Talker process.
+
+Code2Wav decodes each streaming chunk statelessly from its codes and left
+context. With a one-frame first chunk (``codec_chunk_ramp`` starting at 1),
+chunk 0 is a context-free decode of the stream's first frame, which the Talker
+can run as soon as eager MTP completes that frame, removing the connector hop
+and the Code2Wav step from time to first audio. The orchestrator then trims
+Code2Wav's copy of those samples.
+
+The decoder is a second ``Qwen3OmniMoeCode2Wav`` built and loaded exactly like
+the Code2Wav stage's, and captured with the same one-frame input shape as that
+stage's ``(batch=1, frames=1)`` graph. Optional cuDNN autotuning can select
+different algorithms in the two processes and change floating-point rounding.
+Each bucket gets its own graph pool because it replays on a side stream while
+the Talker's own graphs may still run.
+"""
+
+from __future__ import annotations
+
+import os
+from collections.abc import Iterable, Sequence
+from typing import Any
+
+import torch
+import torch.nn as nn
+from torch.cuda import CUDAGraph
+from vllm.logger import init_logger
+
+from vllm_omni.model_executor.models.common.qwen3_code_predictor import CodePredictorWrapper
+from vllm_omni.model_executor.models.common.talker_first_audio import supports_talker_first_audio
+from vllm_omni.model_executor.stage_input_processors.chunk_size_utils import parse_chunk_ramp
+
+logger = init_logger(__name__)
+
+
+def talker_first_audio_enabled(vllm_config: Any) -> bool:
+    """Opt in only when chunk 0 carries the first frame in the same step."""
+    extra = CodePredictorWrapper._stage_connector_extra_config(vllm_config)
+    ramp = parse_chunk_ramp(extra)
+    initial_frames = ramp[0] if ramp is not None else int(extra.get("initial_codec_chunk_frames") or 0)
+    return (
+        CodePredictorWrapper._parse_bool_config(extra.get("talker_first_audio"))
+        and initial_frames == 1
+        and os.environ.get("VLLM_OMNI_TALKER_FIRST_AUDIO", "1") == "1"
+        and supports_talker_first_audio(vllm_config)
+    )
+
+
+class Qwen3OmniFirstFrameDecoder(nn.Module):
+    def __init__(self, code2wav: nn.Module, sample_rate: int) -> None:
+        super().__init__()
+        self.code2wav = code2wav
+        self.num_quantizers = int(code2wav.config.num_quantizers)
+        self.sample_rate = int(sample_rate)
+        self._graphs: dict[int, tuple[CUDAGraph, torch.Tensor, torch.Tensor]] = {}
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        """Load the checkpoint's ``code2wav.*`` weights as the Code2Wav stage does."""
+        loaded = self.code2wav.load_weights(weights)
+        if hasattr(self.code2wav, "precompute_snake_caches"):
+            self.code2wav.precompute_snake_caches()
+        return loaded
+
+    @torch.inference_mode()
+    def capture(self, batch_sizes: Sequence[int] = (1, 2, 4, 8)) -> None:
+        device = next(self.code2wav.parameters()).device
+        if device.type != "cuda":
+            return
+        for batch_size in sorted(set(batch_sizes)):
+            static_input = torch.zeros(batch_size, self.num_quantizers, 1, dtype=torch.long, device=device)
+            self.code2wav(static_input)
+            torch.accelerator.synchronize(device)
+            graph = CUDAGraph()
+            with torch.cuda.graph(graph, pool=torch.cuda.graph_pool_handle()):
+                static_output = self.code2wav(static_input)
+            self._graphs[batch_size] = (graph, static_input, static_output)
+        logger.info("Captured Qwen3-Omni Talker first-frame decoder graphs for batch sizes %s", sorted(self._graphs))
+
+    @torch.inference_mode()
+    def decode(self, codes: torch.Tensor) -> torch.Tensor:
+        """``codes`` [n, num_quantizers] long on device -> float32 PCM [n, samples] (new tensor)."""
+        if not self._graphs:
+            return self.code2wav(codes.to(torch.long).unsqueeze(-1))[:, 0, :].float()
+        outputs = []
+        largest = max(self._graphs)
+        for start in range(0, int(codes.shape[0]), largest):
+            chunk = codes[start : start + largest]
+            rows = int(chunk.shape[0])
+            batch_size = min(size for size in self._graphs if size >= rows)
+            graph, static_input, static_output = self._graphs[batch_size]
+            static_input.zero_()
+            static_input[:rows, :, 0].copy_(chunk)
+            graph.replay()
+            # A replay overwrites graph-owned storage; .float() aliases FP32 outputs.
+            outputs.append(static_output[:rows, 0, :].to(dtype=torch.float32, copy=True))
+        return torch.cat(outputs, dim=0)

vllm_omni/model_executor/models/qwen3_omni/pipeline.py

@@ -45,6 +45,8 @@ QWEN3_OMNI_PIPELINE = PipelineConfig(
             engine_output_type="latent",
             custom_process_next_stage_input_func=(f"{_PROC}.thinker2talker_full_payload"),
             async_chunk_process_next_stage_input_func=(f"{_PROC}.thinker2talker_async_chunk"),
+            # Takes effect only when the deploy selects model_runner v2 for this stage.
+            supports_native_mrv2_data_plane=True,
             sampling_constraints={"detokenize": True},
         ),
         StagePipelineConfig(
@@ -57,6 +59,8 @@ QWEN3_OMNI_PIPELINE = PipelineConfig(
             sync_process_input_func=f"{_PROC}.thinker2talker_token_only",
             custom_process_next_stage_input_func=(f"{_PROC}.talker2code2wav_full_payload"),
             async_chunk_process_next_stage_input_func=(f"{_PROC}.talker2code2wav_async_chunk"),
+            # Each stage can use MRv2 when its deploy selects model_runner v2.
+            supports_native_mrv2_data_plane=True,
             sampling_constraints={
                 "detokenize": False,
                 "stop_token_ids": [2150],
@@ -72,6 +76,7 @@ QWEN3_OMNI_PIPELINE = PipelineConfig(
             final_output_type="audio",
             hf_config_name="thinker_config",
             engine_output_type="audio",
+            supports_native_mrv2_data_plane=True,
             sampling_constraints={"detokenize": True},
             requires_full_payload_input=True,
         ),

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py

@@ -4,7 +4,7 @@
 """Inference-only Qwen3-Omni-Moe unified model (thinker + talker + code2wav)."""
 
 import asyncio
-from collections.abc import AsyncGenerator, Iterable
+from collections.abc import AsyncGenerator, Iterable, Mapping
 from functools import cached_property
 from typing import Any
 
@@ -50,6 +50,7 @@ from vllm_omni.data_entry_keys import Embeddings, HiddenStates, Ids, OmniPayload
 from vllm_omni.metrics import definitions as defs
 from vllm_omni.model_executor.custom_process_mixin import CustomProcessMixin
 from vllm_omni.model_executor.models.output_templates import OmniOutput
+from vllm_omni.model_executor.models.qwen3_omni.first_frame_decoder import talker_first_audio_enabled
 from vllm_omni.model_executor.models.qwen3_omni.quantization import (
     apply_outer_quant_config_mapping,
 )
@@ -61,7 +62,9 @@ from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker import (
     Qwen3OmniMoeThinkerProcessingInfo,
 )
 from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights, safe_tensor_reshape
+from vllm_omni.model_executor.stage_input_processors.chunk_size_utils import parse_chunk_ramp, ramp_decode_windows
 from vllm_omni.platforms import current_omni_platform
+from vllm_omni.utils.device_copy import to_device_nonblocking
 
 # Special token IDs for Qwen3 Omni MoE
 # Reference: https://huggingface.co/Qwen/Qwen3-Omni-30B-A3B-Instruct/blob/main/tokenizer_config.json
@@ -88,6 +91,13 @@ TALKER_CODEC_THINK_EOS_ID = 4205  # Think mode end
 logger = init_logger(__name__)
 
 
+def _cpu_long(values: Any) -> torch.Tensor:
+    """Token ids as a host tensor (payload ids arrive as lists or CPU tensors)."""
+    if isinstance(values, torch.Tensor):
+        return values.detach().to("cpu", torch.long)
+    return torch.as_tensor(values, dtype=torch.long)
+
+
 @MULTIMODAL_REGISTRY.register_processor(
     Qwen3OmniMoeThinkerMultiModalProcessor,
     info=Qwen3OmniMoeThinkerProcessingInfo,
@@ -135,6 +145,11 @@ class Qwen3OmniMoeForConditionalGeneration(
 
     realtime_max_tokens = 64
 
+    # Set by the MRv2 model state once eager Talker-MTP frames are enabled.
+    eager_frames_active = False
+    # Postprocess only stores the last hidden for the deferred decode MTP.
+    eager_frames_need_postprocess = False
+
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         super().__init__()
         self.have_multimodal_outputs = True
@@ -180,9 +195,18 @@ class Qwen3OmniMoeForConditionalGeneration(
         # thinker layers and the forward must return what stock vLLM expects.
         self.is_staged_run = getattr(vllm_config.model_config, "model_stage", None) is not None
         self._returns_tuple = self.model_stage == "thinker" and self.is_staged_run
+        # The captured layers are a fixed-structure tensor dict on the token
+        # axis, so MRv2 may replay decode in FULL graphs and return them too.
+        self.supports_mrv2_full_graph_aux_outputs = self._returns_tuple
 
         if self.model_stage == "thinker":
             self.use_async_omni_output = True
+            # MRv2: the Talker's prompt ends with the first generated token's
+            # embedding; publishing it with the prefill output lets the first
+            # Thinker->Talker chunk leave one decode step earlier.
+            # Speculative sampling can change the token width between steps.
+            # Keep the existing capture-stream handoff for the entire request.
+            self.publishes_sampled_embeddings = self.is_staged_run and vllm_config.speculative_config is None
             # Initialize thinker model (multimodal processing + text generation)
             # Create a new vllm_config with thinker_config as the hf_config
             thinker_vllm_config = vllm_config.with_hf_config(
@@ -233,6 +257,19 @@ class Qwen3OmniMoeForConditionalGeneration(
             )
             self.model = self.talker
             self.code2wav = None
+            # Residual codebooks (MTP) output key, shared by the V1 runner
+            # (talker_mtp_*) and the MRv2 model state (mtp_*).
+            self.talker_mtp_output_key = ("codes", "audio")
+            # MRv2: the step that samples CB0 may complete the frame itself,
+            # since the residual codebooks depend only on CB0 and that step's
+            # last hidden, and the next input is their embedding sum plus the
+            # next text step (see ``talker_mtp``).
+            self.mtp_eager_frames = True
+            self._codec_codebook_size = int(talker_config.code_predictor_config.vocab_size)
+            self._apply_subtalker_sampling_params(getattr(vllm_config.model_config, "subtalker_sampling_params", None))
+
+            # Decode each stream's first frame here (see talker_first_audio_enabled).
+            self.first_frame_decoder = self._build_first_frame_decoder(vllm_config, code2wav_config, prefix)
 
             # for CI: Initialize special tokens embeddings early to avoid AttributeError when loading dummy weights
             self._init_special_tokens_embeddings()
@@ -247,6 +284,9 @@ class Qwen3OmniMoeForConditionalGeneration(
                 # talker MTP codec codes must stay on GPU to avoid a per-step D2H
                 # sync stall; build_mm_cpu handles the eventual D2H at payload time.
                 ("codes", "audio"),
+                # The growing thinker-embed cache is read and extended every
+                # decode step; on CPU it cost a blocking D2H per request per step.
+                ("embed", "cached_decode"),
             }
 
         elif self.model_stage == "code2wav":
@@ -285,6 +325,74 @@ class Qwen3OmniMoeForConditionalGeneration(
                     ["hidden_states", "residual", *incoming_captures], thinker_config.text_config.hidden_size
                 )
 
+    @staticmethod
+    def _build_first_frame_decoder(
+        vllm_config: VllmConfig,
+        code2wav_config: Qwen3OmniMoeCode2WavConfig,
+        prefix: str,
+    ) -> nn.Module | None:
+        """A Code2Wav copy for the Talker's first-frame audio, or None when disabled.
+
+        Only a streaming (async-chunk) Talker hands frames to Code2Wav chunk by
+        chunk, so only there is a stream's first chunk a context-free decode.
+        The copy is built like the Code2Wav stage's module (same class, config,
+        dtype and device context). Optional cuDNN autotuning is also shared,
+        although per-process algorithm selection can change PCM rounding.
+        """
+        if not talker_first_audio_enabled(vllm_config):
+            return None
+        from vllm_omni.model_executor.models.qwen3_omni.first_frame_decoder import Qwen3OmniFirstFrameDecoder
+
+        code2wav = init_vllm_registered_model(
+            vllm_config=vllm_config.with_hf_config(code2wav_config, architectures=["Qwen3OmniMoeCode2Wav"]),
+            prefix=maybe_prefix(prefix, "first_frame_decoder.code2wav"),
+            hf_config=code2wav_config,
+            architectures=["Qwen3OmniMoeCode2Wav"],
+        )
+        logger.info("Qwen3-Omni Talker decodes first-frame audio in-stage")
+        return Qwen3OmniFirstFrameDecoder(code2wav, defs.resolve_audio_sample_rate(code2wav_config))
+
+    def capture_first_frame_graphs(self) -> None:
+        """Called by the MRv2 runner after its own graph capture."""
+        decoder = getattr(self, "first_frame_decoder", None)
+        if decoder is not None:
+            decoder.capture()
+
+    def _apply_subtalker_sampling_params(self, params: Any) -> None:
+        """Apply the deploy's ``subtalker_sampling_params`` to the residual-codebook predictor.
+
+        The predictor samples in its "stored" mode (top-k 50, top-p 0.8 by
+        default) and ignores per-call sampling arguments; ``do_sample: false``
+        makes it greedy.
+        """
+        if not isinstance(params, Mapping) or not params:
+            return
+        predictor = self.talker.code_predictor
+        predictor.set_sampling_params(
+            top_k=int(params.get("top_k", predictor._top_k)),
+            top_p=float(params.get("top_p", predictor._top_p)),
+            do_sample=bool(params.get("do_sample", True)),
+        )
+
+    # MRv2 capability names read the V1 canonical values, so a platform patch
+    # of ``talker_mtp_graph_safe`` applies to both runners.
+    @property
+    def mtp(self):
+        """MRv2 name for ``talker_mtp``; only the Talker stage has an MTP."""
+        return self.talker_mtp if self.model_stage == "talker" else None
+
+    @property
+    def mtp_output_key(self) -> tuple[str, str]:
+        return getattr(self, "talker_mtp_output_key", ("codes", "audio"))
+
+    @property
+    def mtp_graph_safe(self) -> bool:
+        return bool(getattr(self, "talker_mtp_graph_safe", False))
+
+    def mtp_frame_valid(self, layer0: torch.Tensor) -> torch.Tensor:
+        """Rows whose CB0 is a codec id rather than codec EOS (or another special id)."""
+        return (layer0 >= 0) & (layer0 < self._codec_codebook_size)
+
     def _thinker_capture_layer_indices(self) -> list[int]:
         accept_layer = getattr(self.talker_config, "accept_hidden_layer", None)
         return [0, int(accept_layer)] if self.is_staged_run and accept_layer is not None else []
@@ -496,38 +604,47 @@ class Qwen3OmniMoeForConditionalGeneration(
         # ========== Stage 3: Code2Wav ==========
         elif self.model_stage == "code2wav":
             seq_token_counts: list[int] | None = kwargs.get("seq_token_counts")
+            num_quantizers = int(self.code2wav.config.num_quantizers)
 
             # Extract codec codes from input
-            if input_ids.shape[0] % 16 == 0:
+            if input_ids.shape[0] % num_quantizers == 0:
                 if seq_token_counts is not None:
-                    max_seq_len = max(seq_token_counts) // 16
+                    max_seq_len = max(seq_token_counts) // num_quantizers
                     batch_size = len(seq_token_counts)
                     split_codes = torch.split(input_ids, seq_token_counts, dim=0)
-                    codes = torch.zeros((batch_size, 16, max_seq_len), device=input_ids.device, dtype=input_ids.dtype)
+                    codes = torch.zeros(
+                        (batch_size, num_quantizers, max_seq_len), device=input_ids.device, dtype=input_ids.dtype
+                    )
                     for idx, code in enumerate(split_codes):
-                        seq_len = code.shape[0] // 16
-                        codes[idx, :, :seq_len] = code.reshape(16, seq_len)
+                        seq_len = code.shape[0] // num_quantizers
+                        codes[idx, :, :seq_len] = code.reshape(num_quantizers, seq_len)
                 else:
-                    codes = input_ids.reshape(1, 16, -1)
+                    codes = input_ids.reshape(1, num_quantizers, -1)
             else:
                 if seq_token_counts is None:
                     logger.debug(
-                        "Code2Wav warmup input length %s is not divisible by 16; padding with zeros.",
+                        "Code2Wav warmup input length %s is not divisible by %s; padding with zeros.",
                         input_ids.shape[0],
+                        num_quantizers,
                     )
                 else:
                     logger.warning_once(
-                        "Code2Wav input length is not divisible by 16; padding with zeros. "
-                        "This is expected only during cudagraph warmup."
+                        "Code2Wav input length is not divisible by %s; padding with zeros. "
+                        "This is expected only during cudagraph warmup.",
+                        num_quantizers,
                     )
                 input_ids_flatten = input_ids.reshape(-1)
                 input_ids_flatten = torch.cat(
                     [
                         input_ids_flatten,
-                        torch.zeros(16 - input_ids.shape[0] % 16, dtype=torch.long, device=input_ids.device),
+                        torch.zeros(
+                            num_quantizers - input_ids.shape[0] % num_quantizers,
+                            dtype=torch.long,
+                            device=input_ids.device,
+                        ),
                     ]
                 )
-                codes = input_ids_flatten.reshape(1, 16, -1)
+                codes = input_ids_flatten.reshape(1, num_quantizers, -1)
 
             # Generate audio from codec codes
             # Get every request's left_context_size from runtime_additional_information (passed via kwargs)
@@ -541,6 +658,38 @@ class Qwen3OmniMoeForConditionalGeneration(
                 logger.debug("No additional_information provided to code2wav stage.")
             audio_tensors = self.generate_audio(codes, left_context_size, seq_token_counts)
 
+            # Keep the skip decision attached to this forward's output, not
+            # mutable model state: asynchronous materialization can overlap steps.
+            from vllm_omni.data_entry_keys import FIRST_AUDIO_REQUIRED_KEY
+
+            flags = [
+                bool(info.get("meta", {}).get("first_audio", False)) for info in (runtime_additional_information or [])
+            ]
+            if flags and any(flags):
+                if len(flags) != len(audio_tensors):
+                    raise ValueError("First-audio flags must align with Code2Wav requests")
+                sample_rate = defs.resolve_audio_sample_rate(self.code2wav_config)
+                frame_counts = (
+                    [count // num_quantizers for count in seq_token_counts] if seq_token_counts else [codes.shape[-1]]
+                )
+                # The causal decoder withholds a right-edge tail. The first
+                # frame has fewer samples than total_upsample; trim its actual
+                # prefix length, including when chunk 0 contains several frames.
+                prefix_lengths = [
+                    max(0, audio.numel() - (frames - 1) * int(self.code2wav.total_upsample)) if skip else 0
+                    for audio, frames, skip in zip(audio_tensors, frame_counts, flags, strict=True)
+                ]
+                return OmniOutput(
+                    text_hidden_states=None,
+                    multimodal_outputs={
+                        "model_outputs": [
+                            audio.reshape(1, -1)[..., prefix:]
+                            for audio, prefix in zip(audio_tensors, prefix_lengths, strict=True)
+                        ],
+                        "sr": [torch.tensor(sample_rate, dtype=torch.int32) for _ in audio_tensors],
+                        FIRST_AUDIO_REQUIRED_KEY: [torch.tensor(skip) for skip in flags],
+                    },
+                )
             return audio_tensors
 
         # Fallback (shouldn't reach here)
@@ -605,7 +754,13 @@ class Qwen3OmniMoeForConditionalGeneration(
 
             if "runtime_additional_information" in kwargs and "model_intermediate_buffer" not in kwargs:
                 logger.warning_once("runtime_additional_information is deprecated, use model_intermediate_buffer")
+            # MRv2 passes each request's token span.
+            spans = kwargs.get("request_token_spans")
+            if spans is not None and self.eager_frames_active and self.vllm_config.model_config.async_chunk:
+                return self._make_eager_talker_output(talker_hidden, spans)
             code_predictor_codes = [info.get("codes", {}).get("audio") for info in info_dicts]
+            if spans is not None:
+                code_predictor_codes = self._align_codes_to_spans(code_predictor_codes, spans, talker_hidden.device)
             audio_codes = torch.cat(code_predictor_codes, dim=0)
             multimodal_outputs: OmniPayload = {"codes": {"audio": audio_codes}}
             span_len = audio_codes.shape[0]
@@ -627,6 +782,44 @@ class Qwen3OmniMoeForConditionalGeneration(
 
         return model_outputs
 
+    def _align_codes_to_spans(
+        self,
+        codes: list[Any],
+        spans: list[tuple[int, int]],
+        device: torch.device,
+    ) -> list[torch.Tensor]:
+        """One ``codes.audio`` block per request with exactly its span's rows.
+
+        MRv2 also runs requests that never went through preprocess (vLLM's
+        ``_warmup_`` requests), whose buffers carry no codes; they get zero
+        rows, which the downstream processor skips like prefill rows.
+        """
+        q = int(self.talker.num_code_groups)
+        aligned: list[torch.Tensor] = []
+        for block, (start, end) in zip(codes, spans, strict=True):
+            if isinstance(block, torch.Tensor) and block.ndim == 2 and block.shape[0] == end - start:
+                aligned.append(block)
+            else:
+                aligned.append(torch.zeros((end - start, q), dtype=torch.long, device=device))
+        return aligned
+
+    def _make_eager_talker_output(self, hidden: torch.Tensor, spans: list[tuple[int, int]]) -> OmniOutput:
+        """Async-chunk output when the MRv2 runner completes frames after sampling.
+
+        ``run_eager_mtp`` writes each sampled row's frame and validity into the
+        last token row of its request span; every other row carries no frame.
+        """
+        num_tokens = spans[-1][1] if spans else 0
+        q = int(self.talker.num_code_groups)
+        mm: OmniPayload = {
+            "codes": {"audio": torch.zeros((num_tokens, q), dtype=torch.long, device=hidden.device)},
+            "meta": {
+                "codec_frame_valid": torch.zeros((num_tokens,), dtype=torch.int8, device=hidden.device),
+                "first_audio": torch.zeros((num_tokens,), dtype=torch.int8, device=hidden.device),
+            },
+        }
+        return OmniOutput(text_hidden_states=hidden, multimodal_outputs=mm)
+
     # ==================== Audio Generation ====================
 
     def generate_audio(
@@ -849,7 +1042,7 @@ class Qwen3OmniMoeForConditionalGeneration(
 
         def _proj_from_thinker(x_opt: torch.Tensor | None) -> torch.Tensor:
             if isinstance(x_opt, torch.Tensor) and x_opt.numel() > 0:
-                xin = _ensure_1x1(x_opt).to(module_device)
+                xin = to_device_nonblocking(_ensure_1x1(x_opt), module_device)
             else:
                 xin = torch.zeros(
                     (1, thinker_embed.shape[-1]),
@@ -882,27 +1075,26 @@ class Qwen3OmniMoeForConditionalGeneration(
             voice_type = str(voice_type).lower().strip()
         start_index = meta.get("num_processed_tokens", 0)
         end_index = start_index + input_embeds.shape[0]
-        # Read thinker outputs for prefill
-        thinker_sequence_embeds = embed["prefill"].to(
-            device=self._module_device(self.talker), dtype=torch.bfloat16
+        # Read thinker outputs for prefill. Every copy here is H2D from the
+        # connector payload; pageable copies would stall the whole Talker batch.
+        talker_device = self._module_device(self.talker)
+        thinker_sequence_embeds = to_device_nonblocking(embed["prefill"], talker_device).to(
+            torch.bfloat16
         )  # Tensor [P,H]
-        thinker_hidden_states = hs["output"].to(
-            device=self._module_device(self.talker), dtype=torch.bfloat16
-        )  # Tensor [K,H]
+        thinker_hidden_states = to_device_nonblocking(hs["output"], talker_device).to(torch.bfloat16)  # Tensor [K,H]
+        # Token ids stay on CPU too: prefill parsing reads them on the host.
+        thinker_sequences_cpu = None if ids.get("all") is None else _cpu_long(ids["all"])
+        thinker_chatml_ids_cpu = None if ids.get("prompt") is None else _cpu_long(ids["prompt"])
         thinker_sequences = (
-            ids.get("all")
-            if ids.get("all") is None
-            else torch.as_tensor(ids["all"], device=self._module_device(self.talker))
+            None if thinker_sequences_cpu is None else to_device_nonblocking(thinker_sequences_cpu, talker_device)
         )
         thinker_chatml_ids = (
-            ids.get("prompt")
-            if ids.get("prompt") is None
-            else torch.as_tensor(ids["prompt"], device=self._module_device(self.talker))
+            None if thinker_chatml_ids_cpu is None else to_device_nonblocking(thinker_chatml_ids_cpu, talker_device)
         )
 
-        tts_bos_thinker = embed["tts_bos"].to(device=self._module_device(self.talker), dtype=torch.bfloat16)
-        tts_eos_thinker = embed["tts_eos"].to(device=self._module_device(self.talker), dtype=torch.bfloat16)
-        tts_pad_thinker = embed["tts_pad"].to(device=self._module_device(self.talker), dtype=torch.bfloat16)
+        tts_bos_thinker = to_device_nonblocking(embed["tts_bos"], talker_device).to(torch.bfloat16)
+        tts_eos_thinker = to_device_nonblocking(embed["tts_eos"], talker_device).to(torch.bfloat16)
+        tts_pad_thinker = to_device_nonblocking(embed["tts_pad"], talker_device).to(torch.bfloat16)
 
         if thinker_sequence_embeds is None or thinker_hidden_states is None:
             raise ValueError(
@@ -924,6 +1116,8 @@ class Qwen3OmniMoeForConditionalGeneration(
             )
             if ids_chatml.ndim == 1:
                 ids_chatml = ids_chatml.unsqueeze(0)
+            if thinker_chatml_ids_cpu is not None and thinker_chatml_ids_cpu.ndim == 1:
+                thinker_chatml_ids_cpu = thinker_chatml_ids_cpu.unsqueeze(0)
         else:
             # Fallback: create dummy ids if not provided
             ids_chatml = torch.zeros(
@@ -932,6 +1126,7 @@ class Qwen3OmniMoeForConditionalGeneration(
                 device=self._module_device(self.talker),
             )
             thinker_sequences = ids_chatml
+            thinker_chatml_ids_cpu = thinker_sequences_cpu = None
 
         speaker_id = self._get_text_spk_token_id(voice_type)
         req_input_ids, req_embeds, trailing_text_hidden = self._thinker_to_talker_prefill(
@@ -944,6 +1139,8 @@ class Qwen3OmniMoeForConditionalGeneration(
             tts_bos_thinker=tts_bos_thinker,
             tts_eos_thinker=tts_eos_thinker,
             tts_pad_thinker=tts_pad_thinker,
+            input_ids_cpu=thinker_chatml_ids_cpu,
+            thinker_result_ids_cpu=thinker_sequences_cpu,
         )
 
         # Queue trailing_text_hidden for decode (drop first for next steps),
@@ -993,12 +1190,11 @@ class Qwen3OmniMoeForConditionalGeneration(
             if cached_thinker_decode_embeds is None:
                 update_dict.setdefault("embed", {})["cached_decode"] = thinker_decode_embeds
             else:
-                cached_thinker_decode_embeds = cached_thinker_decode_embeds.to(
-                    device=self._module_device(self.talker), dtype=torch.bfloat16
-                )
-                thinker_decode_embeds = thinker_decode_embeds.to(
-                    device=self._module_device(self.talker), dtype=torch.bfloat16
+                talker_device = self._module_device(self.talker)
+                cached_thinker_decode_embeds = to_device_nonblocking(cached_thinker_decode_embeds, talker_device).to(
+                    torch.bfloat16
                 )
+                thinker_decode_embeds = to_device_nonblocking(thinker_decode_embeds, talker_device).to(torch.bfloat16)
                 update_dict.setdefault("embed", {})["cached_decode"] = torch.cat(
                     [cached_thinker_decode_embeds, thinker_decode_embeds], dim=0
                 )
@@ -1015,26 +1211,31 @@ class Qwen3OmniMoeForConditionalGeneration(
         tts_bos_thinker: torch.Tensor | None = None,
         tts_eos_thinker: torch.Tensor | None = None,
         tts_pad_thinker: torch.Tensor | None = None,
+        input_ids_cpu: torch.Tensor | None = None,
+        thinker_result_ids_cpu: torch.Tensor | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
         """
         Project thinker outputs to talker inputs during prefill stage.
 
+        ``input_ids_cpu`` / ``thinker_result_ids_cpu`` (host copies of the same
+        ids) move segment parsing off the GPU: each ``.item()`` / ``nonzero`` /
+        mask test on device ids is a host sync that stalls the Talker batch.
+
         Returns:
             (input_ids, input_embeds) for talker
         """
         target_len = thinker_result_ids.shape[-1]
-        im_start_indexes = torch.cat(
-            (
-                torch.nonzero(input_ids[0] == self.config.im_start_token_id).squeeze(-1),
-                torch.tensor([target_len], device=input_ids.device, dtype=input_ids.dtype),
-            ),
-            dim=-1,
-        )  # Shape [n_starts + 1]; Take batch 0 since batched inference is not supported here.
+        # Take batch 0 since batched inference is not supported here.
+        parse_ids = input_ids_cpu if input_ids_cpu is not None else input_ids
+        result_ids = thinker_result_ids_cpu if thinker_result_ids_cpu is not None else thinker_result_ids
+        im_start_indexes = torch.nonzero(parse_ids[0] == self.config.im_start_token_id).flatten().tolist()
+        im_start_indexes.append(target_len)
         multimodal_mask = (
-            (thinker_result_ids == self.thinker_config.audio_token_id) |
-            (thinker_result_ids == self.thinker_config.image_token_id) |
-            (thinker_result_ids == self.thinker_config.video_token_id)
-        ).to(input_ids.device)  # [t] # fmt: skip
+            (result_ids == self.thinker_config.audio_token_id) |
+            (result_ids == self.thinker_config.image_token_id) |
+            (result_ids == self.thinker_config.video_token_id)
+        )  # [t] # fmt: skip
+        role_tokens = parse_ids[0].tolist()
 
         tts_bos_embed, tts_eos_embed, tts_pad_embed = self._get_tts_embed(
             thinker_embed, tts_bos_thinker, tts_eos_thinker, tts_pad_thinker
@@ -1045,21 +1246,21 @@ class Qwen3OmniMoeForConditionalGeneration(
         trailing_text_hidden_all: torch.Tensor | None = None
         # For every chatml parts
         for i in range(len(im_start_indexes) - 1):
-            im_start_index = im_start_indexes[i].item()
-            segment_end_index = im_start_indexes[i + 1].item()
-            role_token = input_ids[0][im_start_index + 1]
+            im_start_index = im_start_indexes[i]
+            segment_end_index = im_start_indexes[i + 1]
+            role_token = role_tokens[im_start_index + 1]
             # Talker should ignore thinker system prompt
-            if (role_token == self.config.system_token_id).item():
+            if role_token == self.config.system_token_id:
                 continue
             # Talker takes word embeddings for tokens and hidden state from `accept_hidden_layer` for multimodal inputs
-            elif (role_token == self.config.user_token_id).item():
+            elif role_token == self.config.user_token_id:
                 talker_user_part = self._get_talker_user_parts(
                     im_start_index, segment_end_index, multimodal_mask, thinker_hidden, thinker_embed
                 )
                 talker_input_embeds.append(talker_user_part)
                 talker_input_ids.append(thinker_result_ids[im_start_index:segment_end_index])
             # Take assistant output (for now)
-            elif (role_token == self.config.assistant_token_id).item() and i == len(im_start_indexes) - 2:
+            elif role_token == self.config.assistant_token_id and i == len(im_start_indexes) - 2:
                 talker_assistant_embeds, talker_assistant_ids, trailing_text_hidden = self._get_talker_assistant_parts(
                     im_start_index,
                     segment_end_index,
@@ -1078,7 +1279,7 @@ class Qwen3OmniMoeForConditionalGeneration(
                 except Exception:
                     pass
             # History assistant output (ignore for now)
-            elif (role_token == self.config.assistant_token_id).item() and i != len(im_start_indexes) - 2:
+            elif role_token == self.config.assistant_token_id and i != len(im_start_indexes) - 2:
                 continue
             else:
                 raise AssertionError("Expect role id after <|im_start|> (assistant, user, system)")
@@ -1103,20 +1304,22 @@ class Qwen3OmniMoeForConditionalGeneration(
 
         cached_thinker_decode_embeds = embed.get("cached_decode", None)
         thinker_decode_embed = embed.get("decode", None)
+        if getattr(self.vllm_config.model_config, "use_v2_model_runner", False) and not meta.get("resumable", False):
+            return self._next_thinker_decode_text_step(
+                cached_thinker_decode_embeds, thinker_decode_embed, meta, device, update_dict
+            )
         start_index = meta.get("num_processed_tokens", 0)
 
         if cached_thinker_decode_embeds is not None and start_index < cached_thinker_decode_embeds.shape[0]:
-            cached_thinker_decode_embeds = cached_thinker_decode_embeds.to(device)
+            cached_thinker_decode_embeds = to_device_nonblocking(cached_thinker_decode_embeds, device)
             thinker_embed = cached_thinker_decode_embeds[start_index]
             if thinker_decode_embed is not None:
-                thinker_decode_embed = thinker_decode_embed.to(device)
+                thinker_decode_embed = to_device_nonblocking(thinker_decode_embed, device)
                 cached_thinker_decode_embeds = torch.cat([cached_thinker_decode_embeds, thinker_decode_embed], dim=0)
                 update_dict.setdefault("embed", {})["cached_decode"] = cached_thinker_decode_embeds
 
         elif thinker_decode_embed is not None:
-            thinker_embed = thinker_decode_embed
-            if thinker_embed.device != device:
-                thinker_embed = thinker_embed.to(device)
+            thinker_embed = to_device_nonblocking(thinker_decode_embed, device)
 
         else:
             # When the tokens output by the thinker are exhausted, an EOS token needs to be appended.
@@ -1129,6 +1332,41 @@ class Qwen3OmniMoeForConditionalGeneration(
         update_dict.setdefault("embed", {})["decode"] = None
         return self.talker.text_projection(thinker_embed).to(device)
 
+    def _next_thinker_decode_text_step(
+        self,
+        pending: torch.Tensor | None,
+        incoming: torch.Tensor | None,
+        meta: OmniPayloadMeta,
+        device: torch.device,
+        update_dict: OmniPayload,
+    ) -> torch.Tensor:
+        """Take the oldest unconsumed Thinker decode row as this step's text.
+
+        Rows arrive in order but not one per Talker step: the MRv2 receiver
+        hands over every row that arrived since the previous step (the V1
+        adapter delivers one chunk per step). New rows are appended to the
+        pending queue ``embed.cached_decode`` and each step consumes one. Once
+        the queue is empty the Thinker's text is exhausted: EOS once, then pad.
+        """
+        embed_update = update_dict.setdefault("embed", {})
+        embed_update["decode"] = None
+        if isinstance(incoming, torch.Tensor) and incoming.numel() > 0:
+            incoming = to_device_nonblocking(incoming, device).reshape(-1, incoming.shape[-1])
+            if isinstance(pending, torch.Tensor) and pending.shape[0] > 0:
+                pending = to_device_nonblocking(pending, device)
+                pending = torch.cat([pending, incoming.to(pending.dtype)], dim=0)
+            else:
+                pending = incoming
+        if not isinstance(pending, torch.Tensor) or pending.shape[0] == 0:
+            embed_update["cached_decode"] = None
+            if meta.get("eos_emitted", False):
+                return self.tts_pad_embed.to(device)
+            update_dict.setdefault("meta", {})["eos_emitted"] = True
+            return self.tts_eos_embed.to(device)
+        pending = to_device_nonblocking(pending, device)
+        embed_update["cached_decode"] = pending[1:] if pending.shape[0] > 1 else None
+        return self.talker.text_projection(pending[:1]).to(device)
+
     def talker_preprocess_decode(
         self, input_ids: torch.Tensor, input_embeds: torch.Tensor, update_dict: OmniPayload, payload: OmniPayload
     ):
@@ -1203,6 +1441,21 @@ class Qwen3OmniMoeForConditionalGeneration(
         )
 
         user_mm_mask = multimodal_mask[im_start_index:segment_end_index]
+        if user_mm_mask.device.type == "cpu":
+            # Host mask: no device sync for the test, and a text-only segment
+            # (the common TTS case) needs no masked gather at all.
+            if not bool(user_mm_mask.any()):
+                user_thinker_embed = thinker_embed[im_start_index:segment_end_index]
+                return self.talker.text_projection(user_thinker_embed).to(thinker_hidden.device, torch.bfloat16)
+            mm_index = to_device_nonblocking(torch.nonzero(user_mm_mask).flatten(), thinker_hidden.device)
+            text_index = to_device_nonblocking(torch.nonzero(~user_mm_mask).flatten(), thinker_hidden.device)
+            segment_hidden = thinker_hidden[im_start_index:segment_end_index]
+            segment_embed = thinker_embed[im_start_index:segment_end_index]
+            mm_hidden = self.talker.hidden_projection(segment_hidden.index_select(0, mm_index))
+            user_talker_part.index_copy_(0, mm_index, mm_hidden.to(user_talker_part.dtype))
+            text_hidden = self.talker.text_projection(segment_embed.index_select(0, text_index))
+            user_talker_part.index_copy_(0, text_index, text_hidden.to(user_talker_part.dtype))
+            return user_talker_part
         # Multimodal data exists
         if user_mm_mask.any():
             user_thinker_hidden_mm = thinker_hidden[im_start_index:segment_end_index][user_mm_mask]
@@ -1236,17 +1489,19 @@ class Qwen3OmniMoeForConditionalGeneration(
             ),
             dim=0,
         )
-        codec_special_tokens = torch.tensor(
-            [
-                self.config.talker_config.codec_nothink_id,
-                self.config.talker_config.codec_think_bos_id,
-                self.config.talker_config.codec_think_eos_id,
-                speaker_id,
-                self.config.talker_config.codec_pad_id,
-                self.config.talker_config.codec_bos_id,
-            ],
-            device=tts_pad_embed.device,
-            dtype=torch.long,
+        codec_special_tokens = to_device_nonblocking(
+            torch.tensor(
+                [
+                    self.config.talker_config.codec_nothink_id,
+                    self.config.talker_config.codec_think_bos_id,
+                    self.config.talker_config.codec_think_eos_id,
+                    speaker_id,
+                    self.config.talker_config.codec_pad_id,
+                    self.config.talker_config.codec_bos_id,
+                ],
+                dtype=torch.long,
+            ),
+            tts_pad_embed.device,
         )
         embed_input_ids = self.talker.embed_input_ids(codec_special_tokens).to(
             device=tts_pad_embed.device, dtype=torch.bfloat16
@@ -1381,18 +1636,31 @@ class Qwen3OmniMoeForConditionalGeneration(
 
     # ==================== Weight Loading ====================
 
-    def _get_codec_frame_config(self) -> tuple[int, int]:
-        """Extract codec_chunk_frames and codec_left_context_frames from stage connector config."""
+    def _get_codec_connector_extra(self) -> dict[str, Any]:
         model_cfg = getattr(self.vllm_config, "model_config", None)
         connector_cfg = getattr(model_cfg, "stage_connector_config", None)
         if isinstance(connector_cfg, dict):
             extra = connector_cfg.get("extra", {})
         else:
             extra = getattr(connector_cfg, "extra", None) or {}
+        return extra if isinstance(extra, Mapping) else {}
+
+    def _get_codec_frame_config(self) -> tuple[int, int]:
+        """Extract codec_chunk_frames and codec_left_context_frames from stage connector config."""
+        extra = self._get_codec_connector_extra()
         chunk_frames = int(extra.get("codec_chunk_frames", 0) or 0)
         left_frames = int(extra.get("codec_left_context_frames", 0) or 0)
         return chunk_frames, left_frames
 
+    def _get_codec_ramp_windows(self) -> list[int]:
+        """Decode windows of the connector's ``codec_chunk_ramp`` (empty without a ramp)."""
+        extra = self._get_codec_connector_extra()
+        chunk_frames = int(extra.get("codec_chunk_frames", 25) or 25)
+        ramp = parse_chunk_ramp(dict(extra), steady=chunk_frames)
+        if not ramp:
+            return []
+        return ramp_decode_windows(ramp, int(extra.get("codec_left_context_frames", 25) or 0))
+
     def _maybe_enable_code2wav_cudagraph(self) -> None:
         """Enable the inner Code2Wav CUDA graph unless this stage runs in eager mode."""
         if not self.code2wav or not hasattr(self.code2wav, "enable_cudagraph"):
@@ -1404,9 +1672,15 @@ class Qwen3OmniMoeForConditionalGeneration(
             return
 
         chunk_frames, left_frames = self._get_codec_frame_config()
+        streaming = bool(self.vllm_config.model_config.async_chunk)
+        # Opt-in multi-row graphs for streaming windows (connector extra
+        # ``codec_graph_batch_sizes``, e.g. [2, 3, 4, 6, 8, 12, 16]).
+        batch_sizes = self._get_codec_connector_extra().get("codec_graph_batch_sizes") or ()
         self.code2wav.enable_cudagraph(
             codec_chunk_frames=chunk_frames,
             codec_left_context_frames=left_frames,
+            extra_capture_sizes=self._get_codec_ramp_windows() if streaming else (),
+            streaming_batch_sizes=[int(b) for b in batch_sizes] if streaming else (),
         )
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
@@ -1439,6 +1713,12 @@ class Qwen3OmniMoeForConditionalGeneration(
             loaded_weights.update(talker_loaded)
             loaded_weights.update(self._init_special_tokens_embeddings())
 
+        # The Talker's first-frame decoder is a Code2Wav copy (same code2wav.* weights).
+        first_frame_decoder = getattr(self, "first_frame_decoder", None)
+        if first_frame_decoder is not None and code2wav_weights:
+            decoder_loaded = first_frame_decoder.load_weights(code2wav_weights)
+            loaded_weights.update(add_prefix_to_loaded_weights(decoder_loaded, "first_frame_decoder.code2wav"))
+
         # Load code2wav weights
         if self.code2wav and code2wav_weights:
             code2wav_loaded = self.code2wav.load_weights(code2wav_weights)

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py

@@ -5,7 +5,8 @@
 
 from __future__ import annotations
 
-from collections.abc import Iterable
+from collections.abc import Callable, Iterable, Sequence
+from contextlib import nullcontext
 
 import numpy as np
 import torch
@@ -19,6 +20,7 @@ from transformers.models.qwen3_omni_moe.modeling_qwen3_omni_moe import (
     Qwen3OmniMoeCode2WavDecoderBlock,
     Qwen3OmniMoeCode2WavTransformerModel,
     Qwen3OmniMoeConvNeXtBlock,
+    Qwen3OmniMoeSnakeBeta,
 )
 from vllm.config import VllmConfig  # type: ignore
 from vllm.logger import init_logger  # type: ignore
@@ -31,9 +33,93 @@ from vllm_omni.model_executor.models.common.snake_activation import SnakeBeta
 from vllm_omni.model_executor.models.qwen3_omni.quantization import (
     Qwen3OmniNestedSupportsQuant,
 )
+from vllm_omni.platforms import current_omni_platform
 
 logger = init_logger(__name__)
 
+# Fixed cost of one streaming decode call (graph replay or eager launch), in
+# decoded-frame equivalents: a call on a single frame costs about as much GPU
+# time as decoding this many extra frames inside a larger call.
+_DECODE_CALL_COST_FRAMES = 32
+
+
+def use_fused_snake(module: nn.Module) -> int:
+    """Replace HF ``Qwen3OmniMoeSnakeBeta`` activations under ``module`` with the fused one.
+
+    The HF decoder blocks recompute ``exp(alpha)`` / ``exp(beta)`` and launch
+    about seven broadcast kernels per call; the shared ``SnakeBeta`` computes the
+    same ``x + 1/b * sin^2(a*x)`` in one kernel from precomputed caches.
+    Parameter names and shapes match, so this runs before weight loading.
+    """
+    count = 0
+    for parent in list(module.modules()):
+        for name, child in list(parent.named_children()):
+            if isinstance(child, Qwen3OmniMoeSnakeBeta):
+                setattr(parent, name, SnakeBeta(child.in_features))
+                count += 1
+    return count
+
+
+def plan_decode_groups(
+    lengths: Sequence[int],
+    bucket_of: Callable[[int], int | None],
+    batch_sizes_for: Callable[[int], Sequence[int]],
+    call_cost_frames: int = _DECODE_CALL_COST_FRAMES,
+) -> list[tuple[list[int], int]]:
+    """Split a streaming batch into decode calls of rows with similar lengths.
+
+    Decoding every row at the batch's longest window wastes most of the work
+    when short ramp chunks (1-15 frames) share a step with steady 50-frame
+    windows. Rows are sorted by length and cut into contiguous groups of
+    graph buckets (``bucket_of``) so that the modelled cost, per call
+    ``call_cost_frames`` plus padded rows (to the next captured row count of
+    that bucket, ``batch_sizes_for``) times the bucket length, is smallest.
+    Returns ``(row indices, decode length)`` per call; the decode length is
+    the group's longest row.
+    """
+    count = len(lengths)
+    if count == 0:
+        return []
+    order = sorted(range(count), key=lambda row: lengths[row])
+    buckets = [bucket_of(int(lengths[row])) or int(lengths[row]) for row in order]
+    distinct = sorted(set(buckets))
+    members = [[row for row, bucket in zip(order, buckets) if bucket == size] for size in distinct]
+
+    def sizes_for(size: int) -> list[int]:
+        return sorted({int(b) for b in batch_sizes_for(size) if int(b) > 0} | {1})
+
+    def cost(rows: int, size: int) -> int:
+        batch_sizes = sizes_for(size)
+        calls, rest = divmod(rows, batch_sizes[-1])
+        padded = calls * batch_sizes[-1]
+        if rest:
+            calls += 1
+            padded += next(b for b in batch_sizes if b >= rest)
+        return calls * call_cost_frames + padded * size
+
+    best = [0] + [None] * len(distinct)
+    cut = [0] * (len(distinct) + 1)
+    for end in range(1, len(distinct) + 1):
+        rows = 0
+        for start in range(end - 1, -1, -1):
+            rows += len(members[start])
+            total = best[start] + cost(rows, distinct[end - 1])
+            if best[end] is None or total < best[end]:
+                best[end], cut[end] = total, start
+    groups: list[tuple[list[int], int]] = []
+    end = len(distinct)
+    while end > 0:
+        start = cut[end]
+        groups.append(([row for bucket_rows in members[start:end] for row in bucket_rows], distinct[end - 1]))
+        end = start
+    calls: list[tuple[list[int], int]] = []
+    for rows, size in reversed(groups):
+        max_batch = sizes_for(size)[-1]
+        for offset in range(0, len(rows), max_batch):
+            part = rows[offset : offset + max_batch]
+            calls.append((part, max(int(lengths[row]) for row in part)))
+    return calls
+
 
 class Qwen3OmniMoeCode2Wav(nn.Module, Qwen3OmniNestedSupportsQuant):
     """
@@ -74,6 +160,14 @@ class Qwen3OmniMoeCode2Wav(nn.Module, Qwen3OmniNestedSupportsQuant):
         super().__init__()
 
         self.config: Qwen3OmniMoeCode2WavConfig = vllm_config.model_config.hf_config
+        connector = getattr(vllm_config.model_config, "stage_connector_config", None)
+        extra = connector.get("extra", {}) if isinstance(connector, dict) else getattr(connector, "extra", {})
+        # Both the separate stage and the Talker's first-frame copy receive
+        # this connector config. Keep algorithm selection consistent between
+        # them, and leave other models' cuDNN settings unchanged.
+        self._cudnn_benchmark = (extra or {}).get("codec_cudnn_benchmark", False)
+        if not isinstance(self._cudnn_benchmark, bool):
+            raise ValueError("codec_cudnn_benchmark must be a boolean")
 
         # Calculate total upsampling factor
         self.total_upsample = np.prod(self.config.upsample_rates + self.config.upsampling_ratios)
@@ -122,10 +216,19 @@ class Qwen3OmniMoeCode2Wav(nn.Module, Qwen3OmniNestedSupportsQuant):
             Qwen3OmniMoeCausalConvNet(output_dim, 1, kernel_size=7),
         ]
         self.decoder = nn.ModuleList(decoder)
+        # Keep existing V1 and non-CUDA decoder blocks unchanged.
+        if (
+            (extra or {}).get("codec_fused_snake", False)
+            and current_omni_platform.is_cuda()
+            and torch.device(vllm_config.device_config.device).type == "cuda"
+        ):
+            use_fused_snake(self.decoder)
 
         # CUDA Graph support — reuses CUDAGraphDecoderWrapper from Qwen3-TTS
         self._cudagraph_enabled = False
         self._cudagraph_wrapper = None
+        # Captured row counts per graph size; empty keeps one padded decode per batch.
+        self._streaming_batch_sizes: dict[int, list[int]] = {}
 
     def precompute_snake_caches(self):
         """Precompute exp(alpha) and 1/(exp(beta)+eps) for all SnakeBeta modules."""
@@ -142,8 +245,19 @@ class Qwen3OmniMoeCode2Wav(nn.Module, Qwen3OmniNestedSupportsQuant):
         device: torch.device | None = None,
         codec_chunk_frames: int = 0,
         codec_left_context_frames: int = 0,
+        extra_capture_sizes: Iterable[int] = (),
+        streaming_batch_sizes: Iterable[int] = (),
     ):
-        """Enable CUDA graph acceleration (same pattern as Qwen3-TTS Code2Wav)."""
+        """Enable CUDA graph acceleration (same pattern as Qwen3-TTS Code2Wav).
+
+        ``extra_capture_sizes`` adds exact-size graphs (e.g. a chunk ramp's
+        decode windows) to the default buckets, so those chunks replay without
+        zero padding. ``streaming_batch_sizes`` (e.g. ``[2, 4, 8, 16]``) also
+        captures multi-row graphs for every bucket up to one streaming window
+        (chunk + left context); a streaming batch is then decoded as a few
+        length groups (``plan_decode_groups``) instead of one eager call
+        padded to its longest row.
+        """
         from vllm_omni.model_executor.models.qwen3_tts.cuda_graph_decoder_wrapper import (
             CUDAGraphDecoderWrapper,
         )
@@ -154,8 +268,30 @@ class Qwen3OmniMoeCode2Wav(nn.Module, Qwen3OmniNestedSupportsQuant):
             logger.warning("Cannot enable CUDA Graph: not on CUDA device (got %s)", device)
             return
 
+        extra = {int(size) for size in extra_capture_sizes if int(size) > 0}
+        batch_sizes = sorted({int(b) for b in streaming_batch_sizes if int(b) > 1})
+        capture_sizes = None
+        extra_shapes: list[tuple[int, int]] = []
+        if extra or batch_sizes:
+            capture_sizes = sorted(
+                set(
+                    CUDAGraphDecoderWrapper.compute_capture_sizes(
+                        codec_chunk_frames=codec_chunk_frames,
+                        codec_left_context_frames=codec_left_context_frames,
+                    )
+                )
+                | extra
+            )
+            window = codec_chunk_frames + codec_left_context_frames
+            streaming_sizes = [size for size in capture_sizes if window <= 0 or size <= window]
+            # A multi-row graph never holds more frames than the largest
+            # single-row one, so the shared graph pool does not grow.
+            largest = max(capture_sizes)
+            extra_shapes = [(b, size) for b in batch_sizes for size in streaming_sizes if b * size <= largest]
         wrapper = CUDAGraphDecoderWrapper(
             decoder=self,
+            capture_sizes=capture_sizes,
+            extra_capture_shapes=extra_shapes,
             num_quantizers=self.config.num_quantizers,
             enabled=True,
         )
@@ -172,10 +308,18 @@ class Qwen3OmniMoeCode2Wav(nn.Module, Qwen3OmniNestedSupportsQuant):
             raise
         self._cudagraph_wrapper = wrapper
         self._cudagraph_enabled = True
+        if extra_shapes and torch.cuda.is_available():
+            # Return the eager warm-up activations: this stage shares the GPU.
+            torch.accelerator.empty_cache()
+        self._streaming_batch_sizes = {
+            size: [1, *sorted(b for b, shape_size in extra_shapes if shape_size == size)]
+            for size in {shape_size for _b, shape_size in extra_shapes}
+        }
         logger.info(
-            "CUDA Graph enabled for Code2Wav: num_quantizers=%d, sizes=%s",
+            "CUDA Graph enabled for Code2Wav: num_quantizers=%d, sizes=%s, multi-row shapes=%s",
             self.config.num_quantizers,
             self._cudagraph_wrapper.capture_sizes,
+            extra_shapes,
         )
 
     def forward(self, codes: torch.Tensor) -> torch.Tensor:
@@ -200,6 +344,30 @@ class Qwen3OmniMoeCode2Wav(nn.Module, Qwen3OmniNestedSupportsQuant):
         hidden = self.pre_transformer(inputs_embeds=hidden).last_hidden_state
         # Shape: [batch, seq_len, hidden_size]
 
+        # Algorithm search happens during eager warmup before graph capture.
+        # Restore process settings afterward, including on a failed decode.
+        # Captured replays use the selected kernels without this Python path.
+        cudnn = torch.backends.cudnn
+        context = (
+            cudnn.flags(
+                enabled=cudnn.enabled,
+                benchmark=True,
+                benchmark_limit=10,
+                deterministic=cudnn.deterministic,
+                allow_tf32=cudnn.allow_tf32,
+            )
+            if (
+                self._cudnn_benchmark
+                and hidden.is_cuda
+                and torch.version.hip is None
+                and not torch.cuda.is_current_stream_capturing()
+            )
+            else nullcontext()
+        )
+        with context:
+            return self._decode_waveform(hidden)
+
+    def _decode_waveform(self, hidden: torch.Tensor) -> torch.Tensor:
         # Stage 3: Upsampling
         hidden = hidden.permute(0, 2, 1)  # [batch, hidden_size, seq_len]
         for blocks in self.upsample:
@@ -305,28 +473,38 @@ class Qwen3OmniMoeCode2Wav(nn.Module, Qwen3OmniNestedSupportsQuant):
                 "defaulting to left_context_size=zeros(len(codes)). This is expected during cudagraph warmup."
             )
             left_context_size = [0] * codes.shape[0]
-        # Decode chunk
-        wavs = []
-        if self._cudagraph_enabled and self._cudagraph_wrapper is not None:
-            batch_wav = self._cudagraph_wrapper.decode(codes)
-        else:
-            batch_wav = self(codes)
         if seq_token_counts is not None:
             code_seq_lens = [n // self.config.num_quantizers for n in seq_token_counts]
         else:
             # Fallback: assume all batch elements share the same sequence length.
             code_seq_lens = [codes.shape[-1]] * codes.shape[0]
-        # The decoder trims the same right-edge tail from every batch row.
-        # Infer it from the decoded window, not a shorter request's length:
-        # padding that request with codec zeros does not provide valid context.
-        tail = max(0, int(codes.shape[-1] * self.total_upsample) - batch_wav.shape[-1])
-        for idx, code_seq_len in enumerate(code_seq_lens):
-            # Refill the previous chunk's tail using the re-decoded left context,
-            # and exclude the current row's padded tail just as singleton decode does.
-            start = max(0, left_context_size[idx] * self.total_upsample - tail)
-            end = max(0, code_seq_len * self.total_upsample - tail)
-            wav_chunk = batch_wav[idx, :, start:end]
-            wavs.append(wav_chunk)
+        batch, width = int(codes.shape[0]), int(codes.shape[-1])
+        graphs = self._cudagraph_wrapper if self._cudagraph_enabled else None
+        batch_sizes = getattr(self, "_streaming_batch_sizes", None)
+        if graphs is not None and batch_sizes and batch > 1:
+            groups = plan_decode_groups(
+                code_seq_lens, graphs._get_padded_size, lambda size: batch_sizes.get(size, (1,))
+            )
+        else:
+            groups = [(list(range(batch)), width)]
+        wavs: list[torch.Tensor | None] = [None] * batch
+        for rows, length in groups:
+            if len(rows) == batch and length == width:
+                group_codes = codes
+            else:
+                # Basic indexing only: a device index tensor would need a host copy.
+                group_codes = torch.stack([codes[row, :, :length] for row in rows])
+            batch_wav = graphs.decode(group_codes) if graphs is not None else self(group_codes)
+            # The decoder trims the same right-edge tail from every batch row.
+            # Infer it from the decoded window, not a shorter request's length:
+            # padding that request with codec zeros does not provide valid context.
+            tail = max(0, int(length * self.total_upsample) - batch_wav.shape[-1])
+            for position, row in enumerate(rows):
+                # Refill the previous chunk's tail using the re-decoded left context,
+                # and exclude the current row's padded tail just as singleton decode does.
+                start = max(0, left_context_size[row] * self.total_upsample - tail)
+                end = max(0, code_seq_lens[row] * self.total_upsample - tail)
+                wavs[row] = batch_wav[position, :, start:end]
         return wavs
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

vllm_omni/model_executor/models/qwen3_tts/cached_code_predictor.py

@@ -18,8 +18,7 @@ from vllm_omni.model_executor.models.common.qwen3_code_predictor import (
     CodePredictorDecoderLayer,
     _rotate_half,
 )
-
-from .short_kv_attention import short_kv_attention
+from vllm_omni.model_executor.models.common.short_kv_attention import short_kv_attention
 
 
 def _attention_cached(

vllm_omni/model_executor/models/qwen3_tts/first_audio.py

@@ -4,9 +4,7 @@
 
 from typing import Any
 
-import torch
-
-from vllm_omni.platforms import current_omni_platform
+from vllm_omni.model_executor.models.common.talker_first_audio import supports_talker_first_audio
 
 
 def talker_first_audio_enabled(vllm_config: Any) -> bool:
@@ -18,16 +16,4 @@ def talker_first_audio_enabled(vllm_config: Any) -> bool:
     from .qwen3_tts_code_predictor_vllm import Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM as Predictor
 
     extra = Predictor._stage_connector_extra_config(vllm_config)
-    model = vllm_config.model_config
-    parallel = vllm_config.parallel_config
-    return (
-        Predictor._parse_bool_config(extra.get("talker_first_audio"))
-        and current_omni_platform.is_cuda()
-        and torch.device(vllm_config.device_config.device).type == "cuda"
-        and bool(getattr(model, "use_v2_model_runner", False))
-        and bool(getattr(model, "async_chunk", False))
-        and parallel.tensor_parallel_size == 1
-        and parallel.pipeline_parallel_size == 1
-        and parallel.distributed_executor_backend in (None, "uni")
-        and not vllm_config.cache_config.enable_prefix_caching
-    )
+    return Predictor._parse_bool_config(extra.get("talker_first_audio")) and supports_talker_first_audio(vllm_config)

vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code_predictor_vllm.py

@@ -55,6 +55,9 @@ class Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM(CodePredictorWrapp
             talker_hidden_size=int(talker_config.hidden_size),
             prefix=prefix,
         )
+        # TTS owns its frame-local cache and overrides _predict_step_logits.
+        # Do not also activate the shared wrapper's Omni incremental loop.
+        self._kv_cache_enabled = False
         # Store talker_config for backward compat (accessed by some callers)
         self.talker_config = talker_config
         self._vllm_config = vllm_config

vllm_omni/model_executor/stage_input_processors/chunk_size_utils.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 import logging
 from dataclasses import dataclass, field
@@ -99,6 +99,21 @@ def ramp_cumulative(index: int, ramp: list[int], steady: int) -> int:
     return sum(ramp) + (index + 1 - ramp_len) * steady
 
 
+def ramp_decode_windows(ramp: list[int], left_context: int) -> list[int]:
+    """Decode window (left context + new frames) of each ramp chunk for a stateless decoder.
+
+    Chunk ``i`` re-decodes up to ``left_context`` earlier frames with its
+    ``ramp[i]`` new ones, e.g. ``[1, 2, 4, 8, 16, 25]`` with 25 frames of
+    context -> ``[1, 3, 7, 15, 31, 50]``.
+    """
+    windows: list[int] = []
+    emitted = 0
+    for new_frames in ramp:
+        windows.append(new_frames + min(emitted, max(0, left_context)))
+        emitted += new_frames
+    return windows
+
+
 def compute_ramp_emit(
     length: int,
     chunk_index: int,

vllm_omni/model_executor/stage_input_processors/qwen3_omni.py

@@ -1,12 +1,12 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 # Copyright 2025 The Qwen team.
 """Stage input processor for Qwen3 Omni MoE: Thinker → Talker transition."""
 
 import logging
 from collections.abc import Mapping
 from dataclasses import dataclass, field
-from typing import Any
+from typing import Any, cast
 
 import torch
 from vllm.inputs import TextPrompt
@@ -23,6 +23,7 @@ from vllm_omni.data_entry_keys import (
 )
 from vllm_omni.engine import OmniEngineCoreRequest
 from vllm_omni.inputs.data import OmniTokensPrompt
+from vllm_omni.model_executor.stage_input_processors.chunk_size_utils import compute_ramp_emit, parse_chunk_ramp
 from vllm_omni.model_executor.stage_input_processors.tts_utils import (
     extract_language_from_prompt,
     extract_language_from_request,
@@ -76,6 +77,43 @@ def _get_accept_hidden_layer_index(transfer_manager: Any) -> int:
     return int(mc.hf_config.talker_config.accept_hidden_layer)
 
 
+# Set once chunk 0 carried the first generated token taken from the prefill
+# sample: from then on each chunk carries the token its step sampled.
+_SAMPLED_TEXT_STREAM = "_qwen3_omni_sampled_text_stream"
+
+
+def _is_final_token(request: Any) -> bool | None:
+    """Whether the token sampled last ends the Thinker's turn (``None``: unknown)."""
+    output_token_ids = _ensure_list(getattr(request, "output_token_ids", None) or [])
+    count = getattr(request, "output_token_count", None)
+    count = len(output_token_ids) if count is None else count
+    token = getattr(request, "last_output_token_id", None)
+    if token is None and output_token_ids:
+        token = output_token_ids[-1]
+    if token is None or count == 0:
+        return None
+    params = getattr(request, "sampling_params", None)
+    stop_ids = getattr(params, "all_stop_token_ids", None) or getattr(params, "stop_token_ids", None) or ()
+    max_tokens = getattr(params, "max_tokens", None)
+    return token in stop_ids or (max_tokens is not None and count >= max_tokens)
+
+
+def _sampled_token_embed(thinker_embed: Mapping[str, Any], request: Any) -> torch.Tensor | None:
+    """Embedding of the token this step sampled, when published and not final.
+
+    The Talker is fed the embeddings of the Thinker's generated tokens except
+    the final one (a stop token, or the last one ``max_tokens`` allows), which
+    no forward ever processes. For any other token the published embedding is
+    exactly what the next step's layer-0 capture holds, one step later.
+    """
+    sampled = thinker_embed.get("sampled")
+    if not isinstance(sampled, torch.Tensor) or sampled.dim() != 2 or sampled.shape[0] != 1:
+        return None
+    if getattr(request, "resumable", False) or _is_final_token(request) is not False:
+        return None
+    return sampled.detach().cpu()
+
+
 def _compute_talker_prompt_ids_length(info: OmniPayload, device: torch.device | str = "cuda") -> int:
     im_start_token_id = 151644
     system_token_id = 8948
@@ -298,6 +336,7 @@ def _get_qwen3_streaming_state(
     streaming_context: Any | None,
 ) -> _Qwen3OmniStreamingState:
     bridge_states = getattr(streaming_context, "bridge_states", None)
+    assert bridge_states is not None
     per_model_state = bridge_states.setdefault("qwen3_omni", {})
     state = per_model_state.get(request_id)
     if state is None:
@@ -387,7 +426,7 @@ def thinker2talker_async_chunk(
         return None
 
     thinker_hs = multimodal_output.get("hidden_states", {})
-    thinker_layers = thinker_hs.get("layers", {}) if isinstance(thinker_hs, dict) else {}
+    thinker_layers = cast(Mapping[int | str, Any], thinker_hs.get("layers", {}) if isinstance(thinker_hs, dict) else {})
     thinker_embed_raw = multimodal_output.get("embed", {})
     thinker_embed = thinker_embed_raw if isinstance(thinker_embed_raw, dict) else {}
 
@@ -441,27 +480,48 @@ def thinker2talker_async_chunk(
             speaker=speaker,
             language=language,
         )
-        if transfer_manager.request_payload.get(request_id) is None:
-            if not is_finished:
-                transfer_manager.request_payload[request_id] = to_dict(payload)
-                return None
-        else:
-            save_payload = transfer_manager.request_payload.pop(request_id)
+        assert payload.embed is not None and payload.embed.prefill is not None
+        assert payload.hidden_states is not None and payload.hidden_states.output is not None
+        save_payload = transfer_manager.request_payload.pop(request_id, None)
+        if save_payload is not None:
             payload.embed.prefill = torch.cat(
                 (save_payload.get("embed", {}).get("prefill"), payload.embed.prefill), dim=0
             )
             payload.hidden_states.output = torch.cat(
                 (save_payload.get("hidden_states", {}).get("output"), payload.hidden_states.output), dim=0
             )
-            prefill_shape = payload.embed.prefill.shape[0]
-            if not is_finished and prefill_shape <= len(prompt_token_ids):
-                transfer_manager.request_payload[request_id] = to_dict(payload)
-                return None
+        if not is_finished and payload.embed.prefill.shape[0] == len(prompt_token_ids):
+            # The Talker prompt ends with the first generated token's embedding,
+            # which the next decode step would capture; take it from the
+            # prefill step's sample instead when the producer published it.
+            first_token_embed = _sampled_token_embed(thinker_embed, request)
+            if first_token_embed is not None:
+                payload.embed.prefill = torch.cat((payload.embed.prefill, first_token_embed), dim=0)
+                transfer_manager.request_payload[request_id] = {_SAMPLED_TEXT_STREAM: True}
+        if not is_finished and payload.embed.prefill.shape[0] <= len(prompt_token_ids):
+            transfer_manager.request_payload[request_id] = to_dict(payload)
+            return None
     else:
         if request.resumable:
             return _construct_thinker2talker_streaming_input_async_chunk(
                 is_finished, request, thinker_emb, thinker_hid, transfer_manager
             )
+        stream = transfer_manager.request_payload.get(request_id)
+        if stream is not None and stream.get(_SAMPLED_TEXT_STREAM):
+            # This step's capture was already sent as the previous step's sample.
+            sampled = _sampled_token_embed(thinker_embed, request)
+            if sampled is None:
+                # A partial (including resumed) prefill has no accepted sample.
+                # Its captures replay text already sent before preemption.
+                raw_sample = thinker_embed.get("sampled")
+                if isinstance(raw_sample, torch.Tensor) and raw_sample.numel() == 0:
+                    return None
+                if _is_final_token(request) is not True and not is_finished:
+                    raise RuntimeError(f"Thinker sample embedding missing for request {request_id}")
+                # The final token is never part of the text; a finish goes out
+                # as the bare finish marker.
+                return None
+            thinker_emb = sampled
         if thinker_emb.shape[0] > 1:
             logger.warning(
                 "Unexpected multiple embeddings in thinker2talker_async_chunk for chunk_id %d: "
@@ -486,7 +546,7 @@ def thinker2talker_full_payload(
     transfer_manager: Any,
     pooling_output: dict[str, Any],
     request: OmniEngineCoreRequest,
-) -> dict[str, Any] | None:
+) -> OmniPayload | None:
     """Pack complete thinker output for the non-async connector path."""
     rid = getattr(request, "request_id", None)
     if not isinstance(pooling_output, Mapping):
@@ -612,7 +672,7 @@ def thinker2talker_token_only(
             )
         thinker_sequences = prompt_token_ids + output_ids
         thinker_input_ids = prompt_token_ids
-        info_for_len = {"ids": {"all": thinker_sequences, "prompt": thinker_input_ids}}
+        info_for_len: OmniPayload = {"ids": {"all": thinker_sequences, "prompt": thinker_input_ids}}
         prompt_len = _compute_talker_prompt_ids_length(info_for_len, device="cpu")
         # Keep this fallback until the connector reliably preserves voice metadata.
         additional_information = to_dict(
@@ -637,6 +697,14 @@ def thinker2talker_token_only(
 # =========================
 
 
+def _last_codec_row(codes: Any) -> torch.Tensor | None:
+    """The last ``[1, Q]`` row of token-major codec codes, or None if there is none."""
+    if not isinstance(codes, torch.Tensor) or codes.ndim != 2 or codes.shape[0] == 0:
+        return None
+    # Own the row: it outlives this step's (batch-wide) output tensor.
+    return codes[-1:].to(torch.long).clone()
+
+
 def talker2code2wav_async_chunk(
     transfer_manager: Any,
     multimodal_output: OmniPayload | dict[str, Any],
@@ -648,27 +716,47 @@ def talker2code2wav_async_chunk(
     """
     request_id = request.external_req_id
     code_predictor_codes = None
+    frame_valid = None
+    first_audio = False
     if isinstance(multimodal_output, Mapping):
         talker_codes = multimodal_output.get("codes", {})
         if isinstance(talker_codes, dict):
             code_predictor_codes = talker_codes.get("audio")
+        meta = multimodal_output.get("meta", {})
+        if isinstance(meta, Mapping):
+            frame_valid = meta.get("codec_frame_valid")
+            flag = meta.get("first_audio", False)
+            first_audio = (
+                bool(flag.numel() and flag.reshape(-1)[-1].item()) if isinstance(flag, torch.Tensor) else bool(flag)
+            )
 
-    sampling_params = getattr(request, "sampling_params", None)
-    stop_token_ids = set(getattr(sampling_params, "stop_token_ids", None) or [])
-    stop_token_id = getattr(sampling_params, "stop_token_id", None)
-    if stop_token_id is not None:
-        stop_token_ids.add(stop_token_id)
+    if isinstance(frame_valid, torch.Tensor) and frame_valid.numel() > 0:
+        # Token-major codes with explicit per-row validity (MRv2 eager
+        # frames): a step's frame, if any, is the last row of its span, and a
+        # multi-row prefill span carries its first frame there.
+        frame = _last_codec_row(code_predictor_codes)
+        append_codes = frame is not None and bool(frame_valid.reshape(-1)[-1].item())
+        code_predictor_codes = frame
+    else:
+        # One row per step (V1 and MRv2 deferred frames): all-zero rows are
+        # prefill placeholders and a stop-token CB0 marks the post-EOS step.
+        sampling_params = getattr(request, "sampling_params", None)
+        stop_token_ids = set(getattr(sampling_params, "stop_token_ids", None) or [])
+        stop_token_id = getattr(sampling_params, "stop_token_id", None)
+        if stop_token_id is not None:
+            stop_token_ids.add(stop_token_id)
 
-    append_codes = (
-        isinstance(code_predictor_codes, torch.Tensor)
-        and code_predictor_codes.numel() > 0
-        and bool(code_predictor_codes.any())
-    )
-    if append_codes:
-        first_codebook = int(code_predictor_codes[0, 0].item())
-        if first_codebook in stop_token_ids:
-            logger.debug("skip stop-token codec frame: first_codebook=%s", first_codebook)
-            append_codes = False
+        append_codes = (
+            isinstance(code_predictor_codes, torch.Tensor)
+            and code_predictor_codes.numel() > 0
+            and bool(code_predictor_codes.any())
+        )
+        if append_codes:
+            assert isinstance(code_predictor_codes, torch.Tensor)
+            first_codebook = int(code_predictor_codes[0, 0].item())
+            if first_codebook in stop_token_ids:
+                logger.debug("skip stop-token codec frame: first_codebook=%s", first_codebook)
+                append_codes = False
     if append_codes:
         transfer_manager.code_prompt_token_ids[request_id].append(code_predictor_codes)
     elif not is_finished:
@@ -680,12 +768,42 @@ def talker2code2wav_async_chunk(
     chunk_size_config = int(cfg.get("codec_chunk_frames", 25))
     left_context_size_config = int(cfg.get("codec_left_context_frames", 25))
     configured_initial_chunk_size = int(cfg.get("initial_codec_chunk_frames") or 0)
+    # The connector config is fixed for the adapter's lifetime: parse once.
+    if not hasattr(transfer_manager, "_qwen3_omni_chunk_ramp"):
+        transfer_manager._qwen3_omni_chunk_ramp = parse_chunk_ramp(cfg, steady=chunk_size_config)
+    ramp = transfer_manager._qwen3_omni_chunk_ramp
 
     chunk_id = transfer_manager.put_req_chunk[request_id]
     length = len(transfer_manager.code_prompt_token_ids[request_id])
     if length <= 0:
         return None
 
+    if ramp is not None:
+        # Ramp (replaces initial_codec_chunk_frames): chunk i carries ramp[i]
+        # new frames, then codec_chunk_frames. Code2Wav decodes statelessly, so
+        # every chunk re-sends up to codec_left_context_frames earlier frames.
+        # The counter restarts with each segment, like code_prompt_token_ids.
+        chunk_index = transfer_manager.ramp_chunk_count.get(request_id, 0)
+        emit, context_length = compute_ramp_emit(length, chunk_index, ramp, chunk_size_config, bool(is_finished))
+        if not emit or context_length <= 0:
+            # Finished on a chunk boundary: the transport sends the finish marker.
+            return None
+        left_context_size = min(length - context_length, left_context_size_config)
+        end_index = left_context_size + context_length
+        codes = (
+            torch.cat(transfer_manager.code_prompt_token_ids[request_id][-end_index:], dim=0)
+            .transpose(0, 1)
+            .reshape(-1)
+        )
+        return OmniPayloadStruct(
+            codes=CodesStruct(audio=codes),
+            meta=MetaStruct(
+                first_audio=torch.tensor(first_audio and chunk_id == 0),
+                left_context_size=left_context_size,
+                finished=torch.tensor(is_finished, dtype=torch.bool),
+            ),
+        )
+
     if configured_initial_chunk_size > 0:
         if chunk_id == 0:
             chunk_size_config = configured_initial_chunk_size
@@ -715,6 +833,7 @@ def talker2code2wav_async_chunk(
     return OmniPayloadStruct(
         codes=CodesStruct(audio=codes),
         meta=MetaStruct(
+            first_audio=torch.tensor(first_audio and chunk_id == 0),
             left_context_size=left_context_size,
             finished=torch.tensor(is_finished, dtype=torch.bool),
         ),

vllm_omni/worker_v2/model_states/omni_model_state.py

@@ -890,15 +890,15 @@ class OmniModelState(DefaultModelState):
         finishing = computed + scheduled >= input_batch.prefill_len_np[:num_reqs]
         kept = ~input_batch.is_prefilling_np[:num_reqs] | finishing
         rows = np.flatnonzero(kept).tolist()
-        if not rows:
-            return None
         device = sampled_token_ids.device
-        first = sampled_token_ids.reshape(num_reqs, -1)[:, 0].index_select(0, index_to_device(rows, device))
-        embeds = self.model.embed_input_ids(first.long())
-        empty = torch.empty(0, dtype=embeds.dtype)
-        sampled: list[torch.Tensor] = [empty] * num_reqs
-        for k, row in enumerate(rows):
-            sampled[row] = embeds[k : k + 1]
+        # Preserve the empty per-row prefill marker even when every request
+        # is prefilling; absence would look like a broken decode handoff.
+        sampled: list[torch.Tensor] = [torch.empty(0)] * num_reqs
+        if rows:
+            first = sampled_token_ids.reshape(num_reqs, -1)[:, 0].index_select(0, index_to_device(rows, device))
+            embeds = self.model.embed_input_ids(first.long())
+            for k, row in enumerate(rows):
+                sampled[row] = embeds[k : k + 1]
         done = torch.cuda.Event()
         done.record()
         extra: dict[str, Any] = {"embed": {"sampled": sampled}}