[Model] Add single-stage Qwen3-TTS with in-Talker streaming codec decode (#8259)

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

tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py

@@ -569,6 +569,27 @@ class TestCodePredictorPerRowGenerators:
             )
 
 
+@pytest.mark.parametrize("fused_requested", [False, True])
+@pytest.mark.parametrize("invalid_input", ["uniforms", "generators"])
+def test_sampling_contract_rejected_before_dispatch(mocker, loaded_target_classes, fused_requested, invalid_input):
+    predictor, talker_config = TestCodePredictorPerRowGenerators()._make_predictor(mocker, loaded_target_classes)
+    predictor._fused_requested = fused_requested
+    predictor._fused = mocker.Mock()
+    setup = mocker.patch.object(predictor, "_setup_compile")
+    uniforms = torch.full((2, 3, 1 if invalid_input == "uniforms" else 64), 0.5)
+    generators = [torch.Generator()] if invalid_input == "generators" else None
+    with pytest.raises(ValueError, match="sample_uniforms must have shape|one entry per row"):
+        predictor(
+            torch.zeros(2, dtype=torch.long),
+            torch.randn(2, talker_config.hidden_size),
+            torch.randn(2, talker_config.hidden_size),
+            sample_uniforms=uniforms,
+            generators=generators,
+        )
+    setup.assert_not_called()
+    predictor._fused.assert_not_called()
+
+
 class TestCodePredictorModelDtype:
     """Test the inner model forward with different dtypes."""
 
@@ -663,6 +684,8 @@ class TestCodePredictorGraphReplay:
 
         predictor = object.__new__(code_predictor_wrapper)
         torch.nn.Module.__init__(predictor)
+        predictor._fused_requested = False
+        predictor._fused = None
         predictor._num_groups = 3
         predictor._model_dtype = torch.float32
         predictor._setup_compile = mocker.Mock()

tests/model_executor/models/qwen3_tts/test_code_predictor_sampling.py

@@ -2,6 +2,8 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Supplied-uniform samples must match the existing ATen path exactly."""
 
+from types import SimpleNamespace
+
 import pytest
 import torch
 
@@ -12,6 +14,79 @@ pytestmark = [
 ]
 
 
+@pytest.mark.parametrize("top_k", [0, 1, 50, 2048])
+@torch.inference_mode()
+def test_fused_sample_kernel_preserves_signed_zero_cutoff_ties(top_k):
+    from vllm_omni.model_executor.models.qwen3_tts.fused_code_predictor import _cp_sample_kernel
+
+    vocab, hidden, groups = 2048, 64, 4
+    logits = torch.zeros(2, vocab, device="cuda", dtype=torch.bfloat16)
+    logits[0, 1] = -0.0
+    logits[1, ::2] = -0.0
+    uniforms = torch.full((2, groups - 1, vocab), 0.01, device="cuda")[:, 1, :]
+    uniforms[0, 1] = uniforms[1, 2] = 0.99
+    codes = torch.full((2, groups), -1, device="cuda", dtype=torch.int64)
+    table = torch.randn(vocab, hidden, device="cuda", dtype=torch.bfloat16)
+    next_input = torch.empty(2, hidden, device="cuda", dtype=table.dtype)
+
+    def sample():
+        _cp_sample_kernel[(2,)](
+            logits, uniforms, uniforms.stride(0), codes, 1, groups, table, next_input, 1.0,
+            V=vocab, TOPK=top_k, HID=hidden, HAS_NEXT=True, num_warps=4,
+        )  # fmt: skip
+
+    scaled = logits.float()
+    if top_k > 0:
+        scaled = scaled.masked_fill(scaled < scaled.topk(top_k).values[:, -1:], -torch.inf)
+    expected = (scaled - torch.log(-torch.log(uniforms))).argmax(-1)
+    assert expected.tolist() == [1, 2]
+    sample()
+    torch.testing.assert_close(codes[:, 1], expected)
+    torch.testing.assert_close(next_input, table[expected])
+    assert (codes[:, [0, 2, 3]] == -1).all()
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        sample()
+    codes.fill_(-1)
+    graph.replay()
+    torch.testing.assert_close(codes[:, 1], expected)
+    torch.testing.assert_close(next_input, table[expected])
+
+
+@torch.inference_mode()
+def test_fused_predictor_warmup_with_capacity_one():
+    from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import (
+        Qwen3TTSTalkerCodePredictorConfig,
+        Qwen3TTSTalkerConfig,
+    )
+    from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_code_predictor_vllm import (
+        Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM,
+    )
+
+    cp = Qwen3TTSTalkerCodePredictorConfig(
+        vocab_size=64, hidden_size=64, intermediate_size=128, num_hidden_layers=1,
+        num_attention_heads=4, num_key_value_heads=2, head_dim=16, num_code_groups=4,
+    )  # fmt: skip
+    talker = Qwen3TTSTalkerConfig(hidden_size=64, num_code_groups=4)
+    config = SimpleNamespace(
+        model_config=SimpleNamespace(stage_connector_config={}),
+        additional_config={"code_predictor_kv_cache": True, "code_predictor_fused": True},
+        scheduler_config=SimpleNamespace(max_num_seqs=1),
+    )
+    predictor = Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM(
+        vllm_config=config, config=cp, talker_config=talker
+    ).to(device="cuda", dtype=torch.bfloat16)
+    predictor._setup_compile()
+    assert predictor._fused is not None
+    hidden = torch.randn(1, 1, 64, device="cuda", dtype=torch.bfloat16)
+    result = predictor(
+        torch.zeros(1, 1, device="cuda", dtype=torch.long), hidden, hidden,
+        sample_uniforms=torch.full((1, 3, 64), 0.5, device="cuda"),
+    )  # fmt: skip
+    assert result.shape == (1, 4)
+    assert ((result >= 0) & (result < 64)).all()
+
+
 @pytest.mark.parametrize("batch", [1, 64])
 @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32])
 @pytest.mark.parametrize("top_k", [-1, 1, 50, 2048])
@@ -66,3 +141,47 @@ def test_topk_gumbel_accepts_strided_supplied_uniforms(top_k):
     expected = (scaled.float() - torch.log(-torch.log(uniforms))).argmax(-1, keepdim=True)
     actual = sample_code_topk_gumbel(logits, uniforms, top_k, 1 / 0.9)
     torch.testing.assert_close(actual, expected)
+
+
+@pytest.mark.parametrize("batch", [1, 4])
+@pytest.mark.parametrize("top_k", [1, 50])
+@torch.inference_mode()
+def test_fused_predictor_matches_full_reference_with_fixed_uniforms(mocker, batch, top_k):
+    from vllm.config import VllmConfig
+
+    from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import (
+        Qwen3TTSTalkerCodePredictorConfig,
+        Qwen3TTSTalkerConfig,
+    )
+    from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_code_predictor_vllm import (
+        Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM,
+    )
+
+    torch.manual_seed(42)
+    cp = Qwen3TTSTalkerCodePredictorConfig(
+        vocab_size=64, hidden_size=64, intermediate_size=128, num_hidden_layers=1,
+        num_attention_heads=4, num_key_value_heads=2, head_dim=16, num_code_groups=4,
+    )  # fmt: skip
+    config = mocker.Mock(
+        spec=VllmConfig,
+        model_config=mocker.Mock(stage_connector_config={}),
+        additional_config={"code_predictor_kv_cache": True, "code_predictor_fused": True},
+        scheduler_config=mocker.Mock(max_num_seqs=4),
+    )
+    predictor = Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM(
+        vllm_config=config, config=cp, talker_config=Qwen3TTSTalkerConfig(hidden_size=64, num_code_groups=4)
+    ).to(device="cuda", dtype=torch.bfloat16)
+    predictor._setup_compile()
+    assert predictor._fused is not None
+    ids = torch.randint(0, 64, (batch, 1), device="cuda")
+    embeds = torch.randn(batch, 1, 64, device="cuda", dtype=torch.bfloat16)
+    hidden = torch.randn_like(embeds)
+    uniforms = torch.rand(batch, 3, 64, device="cuda").clamp_(1e-6, 1 - 1e-6)
+    actual = predictor(ids, embeds, hidden, top_k=top_k, sample_uniforms=uniforms).clone()
+    # Use identical parameters through the full re-prefill and ATen sampling
+    # path, bypassing both the fused predictor and its frame-local KV cache.
+    predictor._fused_requested = False
+    predictor._fused = None
+    predictor._frame_cache = None
+    expected = predictor(ids, embeds, hidden, top_k=top_k, sample_uniforms=uniforms)
+    torch.testing.assert_close(actual, expected, rtol=0, atol=0)

tests/model_executor/models/qwen3_tts/test_first_audio_graphs.py

@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Prefix graphs must follow per-request first-audio delivery, not a capture hint."""
 
+import math
 from types import SimpleNamespace
 
 import pytest
@@ -18,6 +19,38 @@ pytestmark = [pytest.mark.core_model]
 DEVICE = torch.device("cuda:0")
 
 
+@pytest.mark.cpu
+@pytest.mark.parametrize("capacity,model_type", [(1, "base"), (128, "base"), (128, "custom_voice")])
+def test_stream_priming_graph_workspace_respects_token_budget(monkeypatch, capacity, model_type):
+    from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import Qwen3TTSTalkerForConditionalGeneration
+    from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz import streaming_decoder
+
+    captures = []
+
+    def capture(stream, sizes, frames=1):
+        captures.append((sizes, frames))
+        return SimpleNamespace(sizes=sizes, frames=frames)
+
+    monkeypatch.setattr(streaming_decoder, "StreamingDecodeGraphs", capture)
+    model = SimpleNamespace(
+        stream_decoder=SimpleNamespace(max_frames=33),
+        stream_graphs=None,
+        stream_prime_graphs=None,
+        stream_ref_context_frames=25,
+        stream_chunk_frames=25,
+        config=SimpleNamespace(tts_model_type=model_type),
+        vllm_config=SimpleNamespace(scheduler_config=SimpleNamespace(max_num_batched_tokens=512)),
+    )
+    sizes = [size for size in [1, 2, 4, 8, 16, 32, 64, 128] if size <= capacity]
+    Qwen3TTSTalkerForConditionalGeneration.capture_stream_decode_graphs(model, sizes)
+    expected = [(sizes, 1)]
+    if model_type == "base":
+        expected.append(([size for size in sizes if size * 25 <= 512], 25))
+    assert captures == expected
+    Qwen3TTSTalkerForConditionalGeneration.capture_stream_decode_graphs(model, sizes)
+    assert captures == expected
+
+
 @pytest.mark.cpu
 @pytest.mark.parametrize("skip_flags", [(False, False), (True, True), (False, True), (True, False)])
 @pytest.mark.parametrize("initial_frames", [1, 3])
@@ -124,3 +157,134 @@ def test_xvec_graphs_follow_request_first_audio_flags(initial_frames, time_major
             torch.testing.assert_close(got, want, atol=1e-4, rtol=1e-4)
             for key in ("ref_hidden", "ref_conv", "prefix_hidden", "suffix_quantized", "suffix_conv"):
                 torch.testing.assert_close(cache[key], expected_cache[key], atol=1e-5, rtol=1e-5)
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@torch.inference_mode()
+def test_stream_decoder_preemption_restores_all_state_into_a_different_slot():
+    from tests.model_executor.models.qwen3_tts.test_time_major_decoder import _make_decoder
+    from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.streaming_decoder import StreamingCodecDecoder
+
+    decoder = _make_decoder().to(device=DEVICE, dtype=torch.bfloat16)
+    for name, param in decoder.named_parameters():
+        if name.endswith("embedding_sum"):
+            param.normal_()
+    decoder.config.head_dim = decoder.config.hidden_size // decoder.config.num_attention_heads
+    stream = StreamingCodecDecoder(decoder, num_slots=2, dtype=torch.bfloat16)
+    codes = torch.randint(0, 32, (1, 8, 2), device=DEVICE)
+    slot = torch.tensor([0], device=DEVICE, dtype=torch.int32)
+    for pos in range(7):
+        stream(codes[:, pos : pos + 1], slot, torch.tensor([pos], device=DEVICE, dtype=torch.int32))
+    saved = stream.save_slot(0)
+    expected = stream(codes[:, 7:], slot, torch.tensor([7], device=DEVICE, dtype=torch.int32)).clone()
+    # Both the freed slot and the destination can hold an unrelated request.
+    other_codes = (codes[:, :3] + 1) % 32
+    for pos in range(3):
+        stream(
+            other_codes[:, pos : pos + 1].expand(2, -1, -1).contiguous(),
+            torch.tensor([0, 1], device=DEVICE, dtype=torch.int32),
+            torch.full((2,), pos, device=DEVICE, dtype=torch.int32),
+        )
+    without_restore = stream(
+        codes[:, 7:],
+        torch.tensor([1], device=DEVICE, dtype=torch.int32),
+        torch.tensor([7], device=DEVICE, dtype=torch.int32),
+    ).clone()
+    assert not torch.equal(without_restore, expected)
+    stream.restore_slot(1, saved)
+    actual = stream(
+        codes[:, 7:],
+        torch.tensor([1], device=DEVICE, dtype=torch.int32),
+        torch.tensor([7], device=DEVICE, dtype=torch.int32),
+    )
+    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("use_graph", [False, True])
+@torch.inference_mode()
+def test_batched_reference_priming_matches_serial_decoder_continuation(use_graph):
+    from tests.model_executor.models.qwen3_tts.test_time_major_decoder import _make_decoder
+    from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import Qwen3TTSTalkerForConditionalGeneration
+    from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.streaming_decoder import (
+        StreamingCodecDecoder,
+        StreamingDecodeGraphs,
+    )
+
+    decoder = _make_decoder().to(device=DEVICE, dtype=torch.bfloat16)
+    for name, param in decoder.named_parameters():
+        if name.endswith("embedding_sum"):
+            param.normal_()
+    decoder.config.head_dim = decoder.config.hidden_size // decoder.config.num_attention_heads
+    stream = StreamingCodecDecoder(decoder, num_slots=5, dtype=torch.bfloat16)
+    graphs = StreamingDecodeGraphs(stream, [2], frames=25) if use_graph else None
+    model = Qwen3TTSTalkerForConditionalGeneration.__new__(Qwen3TTSTalkerForConditionalGeneration)
+    torch.nn.Module.__init__(model)
+    model.stream_decoder = stream
+    model.stream_prime_graphs = graphs
+    model.stream_chunk_frames = 25
+    refs = [torch.randint(0, 32, (length, 2), device=DEVICE) for length in [25, 3, 51, 25]]
+    slots = [2, 0, 3, 4]
+    frame = torch.randint(0, 32, (4, 1, 2), device=DEVICE)
+    slot_tensor = torch.tensor(slots, device=DEVICE, dtype=torch.int32)
+    positions = torch.tensor([len(ref) for ref in refs], device=DEVICE, dtype=torch.int32)
+    for slot, ref in zip(slots, refs, strict=True):
+        for t0 in range(0, len(ref), 25):
+            stream(
+                ref[None, t0 : t0 + 25].to(torch.int32).contiguous(),
+                torch.tensor([slot], device=DEVICE, dtype=torch.int32),
+                torch.tensor([t0], device=DEVICE, dtype=torch.int32),
+            )
+    expected = stream(frame, slot_tensor, positions).clone()
+    untouched = stream.save_slot(1)
+    # Reuse dirty slots; pos=0 must reset every causal layer for each group.
+    model.prime_stream_decoder(list(zip(slots, refs, strict=True)))
+    actual = stream(frame, slot_tensor, positions)
+    torch.testing.assert_close(actual, expected, rtol=0, atol=3e-5)
+    for got, want in zip(stream.save_slot(1), untouched, strict=True):
+        torch.testing.assert_close(got, want, rtol=0, atol=0)
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("use_graph", [False, True])
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+@torch.inference_mode()
+def test_stream_frames_match_fp32_exact_decoder(use_graph, dtype):
+    from tests.model_executor.models.qwen3_tts.test_time_major_decoder import _make_decoder
+    from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.streaming_decoder import (
+        StreamingCodecDecoder,
+        StreamingDecodeGraphs,
+    )
+
+    decoder = _make_decoder().to(device=DEVICE, dtype=dtype)
+    # The tiny HF default initialization shrinks the deep conv stack to
+    # near-zero PCM. Keep the oracle signal large enough to catch lost state.
+    for module in decoder.modules():
+        if isinstance(module, (torch.nn.Conv1d, torch.nn.ConvTranspose1d, torch.nn.Linear)):
+            torch.nn.init.kaiming_uniform_(module.weight, a=math.sqrt(5))
+            if module.bias is not None:
+                module.bias.uniform_(-0.01, 0.01)
+    for name, param in decoder.named_parameters():
+        if name.endswith("embedding_sum"):
+            param.normal_()
+    decoder.config.head_dim = decoder.config.hidden_size // decoder.config.num_attention_heads
+    reference = _make_decoder().to(device=DEVICE, dtype=torch.float32)
+    reference.load_state_dict(decoder.state_dict())
+    reference.precompute_snake_caches()
+    # Cross both the attention window and KV ring wrap, checking every frame.
+    codes = torch.randint(0, 32, (2, 136, 2), device=DEVICE)
+    expected = reference._forward_exact(codes.transpose(1, 2))[:, 0]
+    assert expected.square().sum() > 1e-3
+    stream = StreamingCodecDecoder(decoder, num_slots=3, dtype=dtype)
+    decode = StreamingDecodeGraphs(stream, [2]) if use_graph else stream
+    slots = torch.tensor([2, 0], device=DEVICE, dtype=torch.int32)
+    spf = int(stream.spf)
+    for frame in range(codes.shape[1]):
+        positions = torch.full((2,), frame, device=DEVICE, dtype=torch.int32)
+        actual = decode(codes[:, frame : frame + 1].contiguous().to(torch.int32), slots, positions).reshape(2, -1)
+        exact_frame = expected[:, frame * spf : (frame + 1) * spf]
+        relative_rms = (actual - exact_frame).square().mean(1).sqrt() / exact_frame.square().mean(1).sqrt()
+        # BF16 kernels are approximate; compare every frame to the independent
+        # FP32 full decoder, not only one streaming implementation to itself.
+        tolerance = 0.05 if dtype == torch.bfloat16 else 5e-4
+        assert torch.all(relative_rms < tolerance), (frame, relative_rms)

tests/model_executor/models/qwen3_tts/test_streaming_col2im_cuda.py

@@ -0,0 +1,57 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Streaming col2im preserves PCM inputs and ping-pong slot state."""
+
+import pytest
+import torch
+from vllm.triton_utils import triton
+
+from tests.helpers.mark import hardware_test
+from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.streaming_decoder import _col2im_kernel
+
+pytestmark = [pytest.mark.core_model, pytest.mark.tts]
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("batch,frames,channels,rate", [(1, 1, 5, 2), (3, 3, 129, 3), (9, 1, 64, 4)])
+def test_col2im_partial_tiles_and_reordered_slot_state(batch: int, frames: int, channels: int, rate: int):
+    torch.manual_seed(23)
+    capacity = batch + 2
+    z = torch.randn(batch, frames, 2 * rate, channels).to(torch.bfloat16)
+    bias = torch.randn(channels).to(torch.bfloat16)
+    state = torch.randn(2, capacity, rate, channels).to(torch.bfloat16)
+    expected_state = state.clone()
+    slots = torch.arange(batch - 1, -1, -1, dtype=torch.int32)
+    positions = torch.arange(batch, dtype=torch.int32)
+    parity = torch.arange(capacity, dtype=torch.int32) % 2
+    expected = torch.empty(batch, frames, rate, channels, dtype=torch.bfloat16)
+    for b in range(batch):
+        slot = int(slots[b])
+        side = int(parity[slot])
+        for t in range(frames):
+            previous = z[b, t - 1, rate:] if t else state[side, slot] if positions[b] else torch.zeros(rate, channels)
+            expected[b, t] = (z[b, t, :rate].float() + previous.float() + bias.float()).to(torch.bfloat16)
+        expected_state[1 - side, slot] = z[b, -1, rate:]
+
+    z_gpu, bias_gpu, state_gpu = z.cuda(), bias.cuda(), state.cuda()
+    slots_gpu, positions_gpu, parity_gpu = slots.cuda(), positions.cuda(), parity.cuda()
+    actual = torch.empty_like(expected, device="cuda")
+    width = min(128, triton.next_power_of_2(channels))
+    rows = batch * frames
+    _col2im_kernel[(rows, rate, triton.cdiv(channels, width))](
+        z_gpu,
+        bias_gpu,
+        state_gpu,
+        slots_gpu,
+        positions_gpu,
+        parity_gpu,
+        actual,
+        frames,
+        capacity,
+        R=rate,
+        C=channels,
+        BC=width,
+        num_warps=4,
+    )
+    torch.testing.assert_close(actual.cpu(), expected, rtol=0, atol=0)
+    torch.testing.assert_close(state_gpu.cpu(), expected_state, rtol=0, atol=0)