[Core][Model] Enable Qwen3-TTS MRV2 and optimize the TTS pipeline (#7781)

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

tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py

@@ -682,6 +682,7 @@ class TestCodePredictorGraphReplay:
             return_proj_buf=False,
         )
         predictor._prefix_graphs_enabled = False
+        predictor._prefix_reprefill_enabled = False
         predictor._bucket_pos_ids = {1: torch.arange(4).unsqueeze(0)}
         graph = mocker.Mock()
         graph_output = torch.zeros(1, 4, 4)
@@ -1362,3 +1363,38 @@ class TestCodePredictorFusedProjections:
 
         with pytest.raises(RuntimeError, match="missing fused parameters"):
             model.load_weights(weights)
+
+
+class TestMTPExecutionBuckets:
+    """Outer MRv2 MTP graph buckets must align with the predictor bucket set."""
+
+    def test_configure_mtp_execution_buckets(self, mocker: MockerFixture, loaded_target_classes) -> None:
+        _ = 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_npu", return_value=False)
+
+        cp_config, _ = _make_tiny_config(loaded_target_classes)
+        vllm_config = _make_vllm_config(mocker, max_num_seqs=4)
+        wrapper = common_mod.CodePredictorWrapper(
+            vllm_config=vllm_config,
+            cp_config=cp_config,
+            wrapper_config=common_mod.CodePredictorWrapperConfig(use_cuda_graphs=False),
+            talker_hidden_size=cp_config.hidden_size,
+        )
+
+        # Default keeps the legacy power-of-two derivation.
+        assert wrapper._execution_batch_buckets is None
+        assert wrapper._batch_bucket_sizes() == [1, 2, 4]
+
+        wrapper.configure_mtp_execution_buckets([3, 1, 3, 0, 9])
+        assert wrapper._batch_bucket_sizes() == [1, 3, 4]  # deduplicate/filter declaration and retain max
+
+        with pytest.raises(RuntimeError, match="before the first warmup"):
+            wrapper2 = common_mod.CodePredictorWrapper(
+                vllm_config=vllm_config,
+                cp_config=cp_config,
+                wrapper_config=common_mod.CodePredictorWrapperConfig(use_cuda_graphs=False),
+                talker_hidden_size=cp_config.hidden_size,
+            )
+            wrapper2._bucket_sizes = [1, 2, 4]  # stand-in for warmed-up state
+            wrapper2.configure_mtp_execution_buckets([3])

tests/model_executor/models/qwen3_tts/test_cuda_graph_decoder.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
 """
 Tests for CUDA Graph decoder wrapper numerical equivalence.
 
@@ -142,6 +142,28 @@ def test_padded_output_shape_and_length(decoder, wrapper, seq_len):
     assert graph_out.shape[-1] == expected_len
 
 
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+def test_padded_batch_size_uses_larger_capture_bucket(decoder):
+    """Non-exact batch sizes should replay the next captured batch bucket."""
+    w = CUDAGraphDecoderWrapper(
+        decoder=decoder,
+        capture_sizes=[25],
+        capture_batch_sizes=[1, 4],
+        num_quantizers=NUM_QUANTIZERS,
+        enabled=True,
+    )
+    w.warmup(DEVICE)
+    codes = _random_codes(25, batch_size=3)
+
+    with torch.no_grad():
+        eager_out = decoder(codes)
+        graph_out = w.decode(codes)
+
+    assert (4, 25) in w.graphs
+    assert graph_out.shape == eager_out.shape
+    torch.testing.assert_close(graph_out, eager_out, atol=0, rtol=0)
+
+
 @hardware_test(res={"cuda": "L4"}, num_cards=1)
 @pytest.mark.parametrize("seq_len", [10, 30, 47, 73, 99])
 def test_padded_interior_positions_close(decoder, wrapper, seq_len):

tests/model_executor/models/qwen3_tts/test_qwen3_tts_output_async_cuda.py

@@ -0,0 +1,55 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Output metadata must not wait for the preceding model forward on CUDA."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from tests.helpers.mark import hardware_test
+from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import (
+    Qwen3TTSTalkerForConditionalGeneration,
+)
+from vllm_omni.worker.gpu_ar_model_runner import _snapshot_tensor_payload_to_cpu_async
+
+pytestmark = [pytest.mark.core_model]
+
+
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+@pytest.mark.parametrize("validity_device", ["cpu", "cuda"])
+@torch.inference_mode()
+def test_output_validity_preserves_values_without_waiting_for_forward(validity_device):
+    talker = Qwen3TTSTalkerForConditionalGeneration.__new__(Qwen3TTSTalkerForConditionalGeneration)
+    talker.vllm_config = SimpleNamespace(model_config=SimpleNamespace(async_chunk=True))
+    hidden = torch.zeros((3, 8), device="cuda")
+    infos = [
+        {
+            "codes": {"audio": torch.zeros((frames, 16), device="cuda", dtype=torch.long)},
+            "meta": {"codec_frame_valid": torch.tensor(valid, device=validity_device)},
+        }
+        for frames, valid in [(2, True), (1, False)]
+    ]
+    # Warm allocation paths before checking the stream-ordering contract.
+    talker.make_omni_output(hidden, model_intermediate_buffer=infos)
+    copy_stream = torch.cuda.Stream()
+    torch.accelerator.synchronize()
+
+    preceding_forward = torch.cuda.Event()
+    torch.cuda._sleep(500_000_000)
+    preceding_forward.record()
+    output = talker.make_omni_output(hidden, model_intermediate_buffer=infos)
+    waited_for_forward = preceding_forward.query()
+    validity = output.multimodal_outputs["meta"]["codec_frame_valid"]
+    snapshot = _snapshot_tensor_payload_to_cpu_async(
+        output.multimodal_outputs, copy_stream=copy_stream, pin_memory=True
+    )
+    snapshot.wait()
+
+    assert not waited_for_forward, "output construction synchronized the preceding forward"
+    assert validity.device.type == "cuda"
+    assert validity.dtype == torch.int8
+    cpu_validity = snapshot.payload["meta"]["codec_frame_valid"]
+    assert cpu_validity.device.type == "cpu"
+    assert cpu_validity.tolist() == [1, 1, 0]
+    assert all(info["meta"]["codec_frame_valid"].device.type == validity_device for info in infos)

tests/model_executor/models/qwen3_tts/test_qwen3_tts_stateless_capture.py

@@ -0,0 +1,100 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Exercise actual stateless decoder capture, not a mocked capture indicator."""
+
+import pytest
+import torch
+
+from tests.helpers.mark import hardware_marks
+from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.configuration_qwen3_tts_tokenizer_v2 import (
+    Qwen3TTSTokenizerV2DecoderConfig,
+)
+from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.modeling_qwen3_tts_tokenizer_v2 import (
+    Qwen3TTSTokenizerV2DecoderTransformerModel,
+)
+
+pytestmark = [pytest.mark.core_model, *hardware_marks(res={"cuda": "L4"}, num_cards=1)]
+
+
+@pytest.mark.parametrize("batch_size", [1, 2])
+@pytest.mark.parametrize("input_mode", ["implicit", "positions"])
+@torch.inference_mode()
+def test_stateless_decoder_capture_replays_new_inputs(batch_size, input_mode):
+    sequence_length = 8
+    torch.manual_seed(42)
+    config = Qwen3TTSTokenizerV2DecoderConfig(
+        hidden_size=32,
+        latent_dim=16,
+        max_position_embeddings=512,
+        num_attention_heads=4,
+        num_key_value_heads=4,
+        intermediate_size=64,
+        num_hidden_layers=1,
+        num_quantizers=2,
+        sliding_window=16,
+        decoder_dim=32,
+        upsample_rates=(2,),
+        upsampling_ratios=(2,),
+    )
+    config._attn_implementation = "sdpa"
+    model = Qwen3TTSTokenizerV2DecoderTransformerModel._from_config(config).cuda().eval()
+    inputs = torch.randn(batch_size, sequence_length, config.latent_dim, device="cuda")
+    positions = {}
+    if input_mode == "positions":
+        cache_position = torch.arange(sequence_length, device="cuda")
+        positions = {"cache_position": cache_position, "position_ids": cache_position.unsqueeze(0)}
+    stream = torch.cuda.Stream()
+    stream.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(stream):
+        for _ in range(3):
+            model(inputs_embeds=inputs, use_cache=False, **positions)
+    torch.cuda.current_stream().wait_stream(stream)
+    torch.accelerator.synchronize()
+
+    # No wrapper or eager fallback: an illegal capture must fail this test.
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        captured_output = model(inputs_embeds=inputs, use_cache=False, **positions).last_hidden_state
+
+    outputs = []
+    for _ in range(3):
+        new_inputs = torch.randn_like(inputs)
+        expected = model(inputs_embeds=new_inputs, use_cache=False, **positions).last_hidden_state
+        if input_mode == "implicit":
+            # Compare the cached mask against the independent live-mask path,
+            # not only against eager execution using the same cached mask.
+            reference = model(
+                inputs_embeds=new_inputs,
+                use_cache=False,
+                position_ids=torch.arange(sequence_length, device="cuda").unsqueeze(0),
+            ).last_hidden_state
+            torch.testing.assert_close(expected, reference, atol=1e-5, rtol=1e-4)
+        inputs.copy_(new_inputs)
+        graph.replay()
+        torch.accelerator.synchronize()
+        torch.testing.assert_close(captured_output, expected, atol=1e-5, rtol=1e-4)
+        outputs.append(captured_output.clone())
+    assert not torch.equal(outputs[0], outputs[1]), "replay must consume the new input"
+
+
+@pytest.mark.parametrize("batch_size", [1, 2])
+@torch.inference_mode()
+def test_stateless_codec_wrapper_captures_and_replays_waveform(batch_size):
+    from tests.model_executor.models.qwen3_tts.test_qwen3_tts_incremental_decode import _make_small_decoder
+    from vllm_omni.model_executor.models.qwen3_tts.segmented_graph_wrapper import CUDAGraphDecoderWrapper
+
+    torch.manual_seed(42)
+    decoder = _make_small_decoder().cuda().eval()
+    wrapper = CUDAGraphDecoderWrapper(decoder, num_quantizers=2, async_chunk=False)
+    # Exercise the production capture method, including quantization and
+    # waveform convolutions, without accepting warmup's eager fallback.
+    wrapper._capture_stateless(batch_size, 8, torch.device("cuda"), torch.long)
+    state = wrapper.stateless_states[(batch_size, 8)]
+    for _ in range(3):
+        codes = torch.randint(0, 32, (batch_size, 2, 8), device="cuda")
+        expected = decoder._forward_exact(codes)
+        state["input"]["codes"].copy_(codes)
+        state["graph"].replay()
+        torch.accelerator.synchronize()
+        torch.testing.assert_close(state["output"], expected, atol=1e-5, rtol=1e-4)
+        assert torch.isfinite(state["output"]).all()

tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_preprocess.py

@@ -1,5 +1,9 @@
-from collections import OrderedDict
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from collections import OrderedDict, UserDict
 from types import SimpleNamespace
+from typing import Any
 
 import numpy as np
 import pytest
@@ -36,6 +40,7 @@ def _make_minimal_talker(
     """
     model = Qwen3TTSTalkerForConditionalGeneration.__new__(Qwen3TTSTalkerForConditionalGeneration)
     model.talker_config = SimpleNamespace(codec_pad_id=7, num_code_groups=16)
+    model._codebook_vocab_size = 2048
     model._embedding_dtype = torch.bfloat16
     if tts_pad_embed is None:
         tts_pad_embed = torch.zeros((1, 4), dtype=torch.bfloat16)
@@ -51,6 +56,22 @@ def _make_minimal_talker(
     return model
 
 
+def test_postprocess_batch_gathers_each_request_tail():
+    from vllm_omni.model_executor.models.output_templates import OwnedBatchTensor
+
+    model = _make_minimal_talker()
+    hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0], [7.0, 8.0]])
+
+    key, values = model.postprocess_batch_mrv2(
+        hidden_states=hidden,
+        last_token_indices=torch.tensor([1, 3]),
+    )
+
+    assert key == ("hidden_states", "last")
+    assert isinstance(values, OwnedBatchTensor)  # freshly allocated rows transfer ownership
+    assert torch.equal(values.tensor, torch.tensor([[3.0, 4.0], [7.0, 8.0]]))
+
+
 def _make_minimal_builder(
     *,
     config: SimpleNamespace | None = None,
@@ -95,6 +116,7 @@ def _make_minimal_builder(
     )
     builder._speaker_cache = None
     builder._text_tokenizer = None
+    builder._projected_token_cache = {}
     builder._embedding_dtype = torch.bfloat16
     builder._ref_audio_artifact_cache_max_entries = 256
     builder._ref_audio_artifact_cache = OrderedDict()
@@ -259,6 +281,30 @@ def test_decode_compacts_long_trailing_text_after_large_offset():
     assert torch.equal(update["hidden_states"]["trailing_text"], trailing_text[65:])
 
 
+def test_decode_marks_codec_frame_validity_from_processed_input_token():
+    model = _make_minimal_talker()
+    model.embed_input_ids = lambda input_ids: input_ids.to(torch.float32).reshape(1, 1, 1).expand(1, 1, 4)
+    common = {
+        "input_embeds": None,
+        "text": ["hello"],
+        "task_type": ["CustomVoice"],
+        "hidden_states": {
+            "trailing_text": torch.ones((2, 4)),
+            "last": torch.ones(4),
+        },
+        "meta": {"talker_text_offset": 0},
+        "_omni_is_prefill": False,
+        "_omni_num_computed_tokens": 2,
+        "_omni_prompt_len": 2,
+    }
+
+    _, _, valid_update = model.preprocess(input_ids=torch.tensor([123]), **common)
+    _, _, eos_update = model.preprocess(input_ids=torch.tensor([4198]), **common)
+
+    assert valid_update["meta"]["codec_frame_valid"].item() is True
+    assert eos_update["meta"]["codec_frame_valid"].item() is False
+
+
 def test_decode_replay_span_embeds_all_tokens_without_mutating_decode_state():
     model = _make_minimal_talker()
 
@@ -286,14 +332,17 @@ def test_decode_replay_span_embeds_all_tokens_without_mutating_decode_state():
         out_embeds.cpu(),
         torch.tensor([[101.0] * 4, [202.0] * 4, [303.0] * 4], dtype=torch.bfloat16),
     )
-    assert update == {"meta": {"codec_streaming": True}}
+    assert update["meta"]["codec_streaming"] is True
+    assert update["meta"]["codec_frame_valid"].item() is False
 
 
-def test_decode_batch_preprocess_matches_decode_state_updates():
+@pytest.mark.parametrize("mrv2", [False, True])
+def test_decode_batch_preprocess_matches_decode_state_updates(mrv2):
     tts_pad = torch.full((1, 4), -1.0, dtype=torch.bfloat16)
     model = _make_minimal_talker(tts_pad_embed=tts_pad)
 
     def fake_embed_input_ids(input_ids):
+        assert not mrv2, "MRV2 must reuse the prepared token embeddings"
         return input_ids.to(torch.float32).reshape(-1, 1, 1).expand(-1, 1, 4)
 
     model.embed_input_ids = fake_embed_input_ids
@@ -302,7 +351,10 @@ def test_decode_batch_preprocess_matches_decode_state_updates():
     last_a = torch.full((4,), 2.0, dtype=torch.float32)
     last_b = torch.full((4,), 3.0, dtype=torch.float32)
 
-    out_ids, out_embeds, past_hidden, text_step, updates = model.preprocess_decode_batch(
+    preprocess = model.preprocess_decode_batch_mrv2 if mrv2 else model.preprocess_decode_batch
+    prepared = torch.tensor([[101.0] * 4, [202.0] * 4], dtype=torch.bfloat16)
+    out_ids, out_embeds, past_hidden, text_step, updates = preprocess(
+        **({"input_embeds": prepared} if mrv2 else {}),
         input_ids=torch.tensor([101, 202], dtype=torch.long),
         req_infos=[
             {
@@ -321,15 +373,19 @@ def test_decode_batch_preprocess_matches_decode_state_updates():
     )
 
     assert out_ids.tolist() == [101, 202]
+    if mrv2:
+        assert out_embeds.data_ptr() == prepared.data_ptr()
     assert torch.equal(out_embeds.cpu(), torch.tensor([[101.0] * 4, [202.0] * 4], dtype=torch.bfloat16))
     assert torch.equal(past_hidden.cpu(), torch.stack([last_a, last_b]).to(torch.bfloat16))
     assert torch.equal(text_step[0].cpu(), trailing_a[1].to(torch.bfloat16))
     assert torch.equal(text_step[1].cpu(), tts_pad.reshape(-1).to(torch.bfloat16))
     assert updates[0]["meta"]["talker_text_offset"] == 2
     assert updates[0]["meta"]["codec_streaming"] is True
+    assert updates[0]["meta"]["codec_frame_valid"].item() is True
     assert "hidden_states" not in updates[0]
     assert updates[1]["meta"]["talker_text_offset"] == 0
     assert updates[1]["meta"]["codec_streaming"] is False
+    assert updates[1]["meta"]["codec_frame_valid"].item() is True
     assert updates[1]["hidden_states"]["trailing_text"].numel() == 0
 
 
@@ -449,11 +505,11 @@ def test_base_voice_clone_batch_preprocess_encodes_ref_code_by_sample_rate():
 
         def __call__(self, texts, *, padding=False):
             self.calls.append((texts, padding))
-            return {"input_ids": [[idx + 1, idx + 2, idx + 3] for idx, _ in enumerate(texts)]}
+            return UserDict({"input_ids": [[1, 2, 3], [2, 3, 4, 5]]})
 
     text_tok = FakeTextTokenizer()
     builder._text_tokenizer = text_tok
-    buf = {
+    buf: dict[str, dict[str, Any]] = {
         "r1": {
             "task_type": ["Base"],
             "text": ["one"],
@@ -488,9 +544,17 @@ def test_base_voice_clone_batch_preprocess_encodes_ref_code_by_sample_rate():
     assert buf["r2"][NORMALIZED_REF_AUDIO_KEY][0] is wav2
     assert len(text_tok.calls) == 2
     assert torch.equal(buf["r1"][PRECOMPUTED_TEXT_IDS_KEY], torch.tensor([1, 2, 3]))
-    assert torch.equal(buf["r2"][PRECOMPUTED_TEXT_IDS_KEY], torch.tensor([2, 3, 4]))
+    assert torch.equal(buf["r2"][PRECOMPUTED_TEXT_IDS_KEY], torch.tensor([2, 3, 4, 5]))
     assert torch.equal(buf["r1"][PRECOMPUTED_REF_IDS_KEY], torch.tensor([1, 2, 3]))
-    assert torch.equal(buf["r2"][PRECOMPUTED_REF_IDS_KEY], torch.tensor([2, 3, 4]))
+    assert torch.equal(buf["r2"][PRECOMPUTED_REF_IDS_KEY], torch.tensor([2, 3, 4, 5]))
+    assert (
+        buf["r1"][PRECOMPUTED_TEXT_IDS_KEY].untyped_storage().data_ptr()
+        == buf["r2"][PRECOMPUTED_TEXT_IDS_KEY].untyped_storage().data_ptr()
+    )
+    assert (
+        buf["r1"][PRECOMPUTED_REF_IDS_KEY].untyped_storage().data_ptr()
+        == buf["r2"][PRECOMPUTED_REF_IDS_KEY].untyped_storage().data_ptr()
+    )
 
 
 def test_base_voice_clone_batch_preprocess_reuses_singleton_normalized_audio_without_speech_tokenizer():
@@ -506,7 +570,7 @@ def test_base_voice_clone_batch_preprocess_reuses_singleton_normalized_audio_wit
             return {"input_ids": [[7, 8, 9] for _ in texts]}
 
     builder._text_tokenizer = FakeTextTokenizer()
-    buf = {
+    buf: dict[str, dict[str, Any]] = {
         "r1": {
             "task_type": ["Base"],
             "text": ["one"],
@@ -539,7 +603,7 @@ def test_base_voice_clone_batch_preprocess_skips_after_initial_prefill_state_exi
     builder._encode_ref_audio_batch_fn = lambda *a, **kw: (_ for _ in ()).throw(
         AssertionError("speech tokenizer not expected")
     )
-    buf = {
+    buf: dict[str, dict[str, Any]] = {
         "r1": {
             "task_type": ["Base"],
             "text": ["one"],
@@ -574,7 +638,7 @@ def test_base_voice_clone_batch_preprocess_uses_serving_artifact_cache_key_witho
             return {"input_ids": [[7, 8, 9] for _ in texts]}
 
     builder._text_tokenizer = FakeTextTokenizer()
-    buf = {
+    buf: dict[str, dict[str, Any]] = {
         "r1": {
             "task_type": ["Base"],
             "text": ["one"],

tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_ref_codes.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
 """Regression tests for cross-request ``codes.ref`` leakage (#4370).
 
 ``make_omni_output`` used to collapse every request's reference codec
@@ -11,11 +11,14 @@ vocoder chunks then decoded with another request's reference voice as
 context (audible onset timbre deformation at any concurrency > 1).
 """
 
+from types import SimpleNamespace
+
 import pytest
 import torch
 
 from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import (
     Qwen3TTSTalkerForConditionalGeneration,
+    _qwen3_tts_gpu_resident_buffer_keys,
 )
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -23,6 +26,18 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 _NUM_CODE_GROUPS = 16
 
 
+def test_large_tts_prefill_artifacts_stay_gpu_resident_only_on_mrv2() -> None:
+    v1_keys = _qwen3_tts_gpu_resident_buffer_keys(False)
+    v2_keys = _qwen3_tts_gpu_resident_buffer_keys(True)
+
+    assert ("embed", "prefill") not in v1_keys
+    assert ("codes", "ref") not in v1_keys
+    assert ("meta", "codec_frame_valid") not in v1_keys
+    assert ("embed", "prefill") in v2_keys
+    assert ("codes", "ref") in v2_keys
+    assert ("meta", "codec_frame_valid") in v2_keys
+
+
 def _make_talker() -> Qwen3TTSTalkerForConditionalGeneration:
     # make_omni_output reads no instance state, so skip __init__ (no real
     # config / checkpoint needed), same as the preprocess tests.
@@ -89,3 +104,42 @@ def test_make_omni_output_keeps_hidden_states_for_replay_spans_without_audio_cod
 
     assert torch.equal(out.text_hidden_states, hidden)
     assert out.multimodal_outputs["codes"]["audio"].shape == (1, _NUM_CODE_GROUPS)
+
+
+def test_make_omni_output_preserves_per_request_codec_frame_validity() -> None:
+    talker = _make_talker()
+    talker.vllm_config = SimpleNamespace(model_config=SimpleNamespace(async_chunk=True))
+    valid = _info(1, None)
+    invalid = _info(1, None)
+    valid["meta"]["codec_frame_valid"] = True
+    invalid["meta"]["codec_frame_valid"] = False
+
+    out = talker.make_omni_output(
+        torch.zeros((2, 8)),
+        model_intermediate_buffer=[valid, invalid],
+    )
+
+    assert out.multimodal_outputs["meta"]["codec_frame_valid"].tolist() == [1, 0]
+
+
+def test_make_omni_output_publishes_ref_once_for_kv_resumed_first_decode() -> None:
+    talker = _make_talker()
+    talker.vllm_config = SimpleNamespace(
+        model_config=SimpleNamespace(async_chunk=True),
+    )
+    ref_code = torch.ones((5, _NUM_CODE_GROUPS), dtype=torch.long)
+    info = _info(1, ref_code)
+    info["_omni_num_computed_tokens"] = 9
+    info["_omni_prompt_len"] = 5
+
+    first = talker.make_omni_output(
+        torch.zeros((1, 8)),
+        model_intermediate_buffer=[info],
+    )
+    second = talker.make_omni_output(
+        torch.zeros((1, 8)),
+        model_intermediate_buffer=[info],
+    )
+
+    assert first.multimodal_outputs["codes"]["ref"][0] is ref_code
+    assert "ref" not in second.multimodal_outputs["codes"]

tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_silence_ban.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
 """Unit tests for the Qwen3-TTS leading-silence ban (#4966, #5048).
 
 ``compute_logits`` suppresses the checkpoint-derived silence codec vocabulary
@@ -59,6 +59,7 @@ def _make_talker(*, ban_frames: int, logits: torch.Tensor, modes=None):
     object.__setattr__(talker, "_silence_mask", silence_mask)
     object.__setattr__(talker, "_codec_disallowed_mask", torch.zeros((_VOCAB,), dtype=torch.bool))
     object.__setattr__(talker, "_batch_req_ids", req_ids)
+    object.__setattr__(talker, "_mrv2_silence_ban_mask", None)
     object.__setattr__(
         talker,
         "_req_x_vector_only",
@@ -78,7 +79,26 @@ def _banned(row: torch.Tensor) -> set[int]:
     return {int(i) for i in torch.nonzero(torch.isinf(row) & (row < 0)).flatten()}
 
 
-# -------------------- the step gate --------------------
+def test_mrv2_silence_ban_tracks_batch_order_scope_and_decode_window() -> None:
+    talker = _make_talker(ban_frames=3, logits=torch.zeros((3, _VOCAB)), modes=[True, False, True])
+    with talker.mrv2_sampling_context(
+        req_ids=["req2", "req1", "req0"],
+        num_output_tokens=torch.tensor([3, 0, 1]),
+    ):
+        logits = talker.compute_logits(torch.zeros((3, 4)))
+
+    assert _banned(logits[0]) == set()
+    assert _banned(logits[1]) == set()
+    assert _banned(logits[2]) == set(_SILENCE_IDS)
+
+    # Prompt logprobs are computed after sampling and must remain unmasked.
+    assert not torch.isneginf(talker.compute_logits(torch.zeros((3, 4)))).any()
+
+    with talker.mrv2_sampling_context(
+        req_ids=["req2", "req1", "req0"],
+        num_output_tokens=torch.tensor([4, 1, 3]),
+    ):
+        assert not torch.isneginf(talker.compute_logits(torch.zeros((3, 4)))).any()
 
 
 def test_ban_disabled_leaves_logits_unchanged() -> None: