Merge current Omni main and retain conditioning regression coverage

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>

tests/model_executor/models/cosyvoice3/test_cosyvoice3_model_helpers.py

@@ -928,4 +928,51 @@ def test_talker_output_contract(monkeypatch, mrv2, packed):
         model._sampling_eps = 1e-6
         model.mrv2_custom_sampler(SimpleNamespace(penalties_state=None))
     output = model.forward(torch.ones(2, dtype=torch.long), torch.arange(2), inputs_embeds=hidden)
-    assert output is (hidden if mrv2 or packed else sentinel)
+    assert output is hidden
+
+
+def test_request_conditioning_normalizes_singletons_without_changing_schema():
+    from vllm_omni.data_entry_keys import to_struct
+    from vllm_omni.model_executor.models.cosyvoice3.cosyvoice3 import _normalize_request_conditioning
+
+    token = torch.tensor([[7, 8]], dtype=torch.int32)
+    feat = torch.randn(1, 4, 80)
+    speaker = torch.randn(1, 192)
+    length = torch.tensor([[2]], dtype=torch.int32)
+    raw = {"embed": {"speech_token": token, "speech_feat": feat, "embedding": speaker, "speech_token_len": [length]}}
+    result = to_struct(_normalize_request_conditioning(raw))
+    assert result.embed.speech_token is token
+    assert result.embed.speech_feat is feat
+    assert result.embed.embedding is speaker
+    assert result.embed.speech_token_len is length
+    assert raw["embed"]["speech_token_len"] == [length]
+
+
+def test_request_conditioning_rejects_unsplit_voices():
+    from vllm_omni.model_executor.models.cosyvoice3.cosyvoice3 import _normalize_request_conditioning
+
+    with pytest.raises(ValueError, match="unsplit batch"):
+        _normalize_request_conditioning({"embed": {"embedding": [torch.ones(1, 192), torch.zeros(1, 192)]}})
+
+
+def test_request_conditioning_reaches_code2wav_with_true_prompt_length():
+    model = _make_code2wav_model()
+    output = model.forward(
+        torch.tensor([1, 2]),
+        torch.arange(2),
+        model_intermediate_buffer=[
+            {
+                "embed": {
+                    "speech_token": torch.tensor([[3, 2, 0, 0]], dtype=torch.int32),
+                    "speech_feat": torch.randn(1, 8, 80),
+                    "embedding": torch.randn(1, 192),
+                    "speech_token_len": [torch.tensor([[2]], dtype=torch.int32)],
+                },
+                "meta": {"req_id": ["voice"]},
+            }
+        ],
+    )
+    call = model.code2wav.forward_calls[0]
+    assert call["prompt_token"].shape == (1, 2)
+    assert call["prompt_feat"].shape == (1, 4, 80)
+    assert output.multimodal_outputs

tests/model_executor/models/qwen3_tts/test_qwen3_tts_prompt_text_batch.py

@@ -0,0 +1,140 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from typing import Any
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import (
+    Qwen3TTSConfig,
+    Qwen3TTSTalkerConfig,
+)
+from vllm_omni.model_executor.models.qwen3_tts.prompt_embeds_builder import (
+    PRECOMPUTED_TEXT_IDS_KEY,
+    Qwen3TTSPromptEmbedsBuilder,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_batch_preprocess_projects_new_non_streaming_texts_once():
+    torch.manual_seed(0)
+    text_embedding = torch.nn.Embedding(64, 4)
+    projection = torch.nn.Linear(4, 4)
+    projected_shapes: list[tuple[int, ...]] = []
+
+    def text_projection(x: torch.Tensor) -> torch.Tensor:
+        projected_shapes.append(tuple(x.shape))
+        return projection(x)
+
+    builder = Qwen3TTSPromptEmbedsBuilder.__new__(Qwen3TTSPromptEmbedsBuilder)
+    builder._text_embedding = text_embedding
+    builder._text_projection = text_projection
+    builder._batched_text_embeds = {}
+
+    ids_a = list(range(1, 13))
+    ids_b = list(range(20, 31))
+    # Serving stores the assistant-template ids wrapped in a one-element list.
+    serving_ids = [ids_a]
+    buf: dict[str, dict[str, Any]] = {
+        "a": {"req_id": "a", "task_type": ["CustomVoice"], PRECOMPUTED_TEXT_IDS_KEY: [ids_a]},
+        "b": {"req_id": "b", "task_type": ["VoiceDesign"], PRECOMPUTED_TEXT_IDS_KEY: [torch.tensor(ids_b)]},
+        "streaming": {
+            "req_id": "streaming",
+            "task_type": ["CustomVoice"],
+            "non_streaming_mode": [False],
+            PRECOMPUTED_TEXT_IDS_KEY: serving_ids,
+        },
+        "base": {"req_id": "base", "task_type": ["Base"], PRECOMPUTED_TEXT_IDS_KEY: serving_ids},
+        "built": {
+            "req_id": "built",
+            "task_type": ["CustomVoice"],
+            "embed": {"prefill": torch.zeros(1)},
+            PRECOMPUTED_TEXT_IDS_KEY: serving_ids,
+        },
+    }
+
+    builder.preprocess_infos_batch(req_infos=list(buf.values()), device=torch.device("cpu"))
+
+    # One embedding + projection over both requests' text tokens (template stripped).
+    assert projected_shapes == [(1, (len(ids_a) - 8) + (len(ids_b) - 8), 4)]
+    assert set(builder._batched_text_embeds) == {"a", "b"}
+    with torch.no_grad():
+        for req_id, ids in (("a", ids_a), ("b", ids_b)):
+            expected = projection(text_embedding(torch.tensor([ids[3:-5]])))
+            torch.testing.assert_close(builder._batched_text_embeds[req_id], expected)
+            assert buf[req_id][PRECOMPUTED_TEXT_IDS_KEY].tolist() == [ids]
+    for skipped in ("streaming", "base", "built"):
+        assert buf[skipped][PRECOMPUTED_TEXT_IDS_KEY] is serving_ids
+
+
+class _Tokenizer:
+    def __call__(self, text: str, **kwargs: Any) -> dict[str, torch.Tensor]:
+        return {"input_ids": torch.tensor([[2, 3, 4, 5]])}
+
+
+def _prompt_builder() -> Qwen3TTSPromptEmbedsBuilder:
+    """Use the real constructor and prompt assembly with small embedding tables."""
+    torch.manual_seed(42)
+    config = Qwen3TTSConfig(tts_bos_token_id=50, tts_eos_token_id=51, tts_pad_token_id=52)
+    talker_config = Qwen3TTSTalkerConfig(
+        codec_nothink_id=10,
+        codec_think_id=11,
+        codec_think_bos_id=12,
+        codec_think_eos_id=13,
+        codec_pad_id=14,
+        codec_bos_id=15,
+        codec_language_id={"english": 16},
+        spk_id={"vivian": 17, "serena": 18},
+    )
+    builder = Qwen3TTSPromptEmbedsBuilder(
+        config=config,
+        talker_config=talker_config,
+        model_path="",
+        text_embedding=torch.nn.Embedding(64, 4),
+        text_projection=torch.nn.Linear(4, 4),
+        codec_embed=torch.nn.Embedding(64, 4),
+        residual_code_embeddings=lambda: [],
+        speaker_encoder=torch.nn.Identity(),
+        tts_pad_embed=torch.zeros(1, 4),
+        encode_ref_audio_batch=lambda *args, **kwargs: [],
+    )
+    builder._text_tokenizer = _Tokenizer()
+    return builder
+
+
+@pytest.mark.parametrize("task_type", ["CustomVoice", "VoiceDesign"])
+@pytest.mark.parametrize("non_streaming", [False, True])
+def test_batched_prompt_consumption_matches_serial_and_clears_cache(task_type: str, non_streaming: bool):
+    batched, serial = _prompt_builder(), _prompt_builder()
+
+    def request(req_id: str, ids: list[int], speaker: str) -> dict[str, Any]:
+        return {
+            "req_id": req_id,
+            "text": ["hello"],
+            "task_type": [task_type],
+            "language": ["English"],
+            "speaker": [speaker],
+            "instruct": ["Speak calmly"],
+            "non_streaming_mode": [non_streaming],
+            PRECOMPUTED_TEXT_IDS_KEY: [ids],
+        }
+
+    with torch.inference_mode():
+        # The same request id can be used again after its earlier prompt was consumed.
+        for ids_a, ids_b in [(list(range(12)), list(range(20, 31))), (list(range(10, 24)), list(range(30, 46)))]:
+            pending = [request("a", ids_a, "Vivian"), request("b", ids_b, "Serena")]
+            expected = [request("a", ids_a, "Vivian"), request("b", ids_b, "Serena")]
+            batched.preprocess_infos_batch(req_infos=pending, device=torch.device("cpu"))
+            assert set(batched._batched_text_embeds) == ({"a", "b"} if non_streaming else set())
+            for actual_info, expected_info in zip(pending, expected, strict=True):
+                actual = batched.build_prompt_embeds(task_type=task_type, info_dict=actual_info)
+                reference = serial.build_prompt_embeds(task_type=task_type, info_dict=expected_info)
+                torch.testing.assert_close(actual[0], reference[0])
+                torch.testing.assert_close(actual[1], reference[1])
+                assert actual[2:] == reference[2:]
+                assert PRECOMPUTED_TEXT_IDS_KEY not in actual_info
+            assert not batched._batched_text_embeds
+        batched.preprocess_infos_batch(req_infos=[], device=torch.device("cpu"))
+        assert not batched._batched_text_embeds

tests/model_executor/stage_input_processors/test_cosyvoice3_stage_input_processors.py

@@ -13,6 +13,7 @@ from vllm_omni.model_executor.stage_input_processors.cosyvoice3 import (
     text2flow_full_payload,
     text2flow_token_only,
 )
+from vllm_omni.outputs.mm_outputs import MultimodalPayload
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -121,6 +122,24 @@ def test_text2flow_full_payload_does_not_send_codec_ids():
     assert torch.equal(payload["embed"]["speech_token"], torch.tensor([[1, 2]], dtype=torch.long))
 
 
+def test_text2flow_full_payload_accepts_structured_multimodal_payload():
+    speech_token = torch.tensor([[1, 2]], dtype=torch.long)
+    payload = text2flow_full_payload(
+        None,
+        MultimodalPayload(
+            tensors={
+                "embed.speech_token": speech_token,
+                "codes.audio": torch.tensor([7, 8, 9], dtype=torch.long),
+            }
+        ),
+        SimpleNamespace(external_req_id="req-structured"),
+    )
+
+    assert payload is not None
+    assert "codes" not in payload
+    assert torch.equal(payload["embed"]["speech_token"], speech_token)
+
+
 def test_talker2code2wav_async_chunk_final_payload_uses_absolute_token_offset():
     transfer_manager = _transfer_manager()
     request = SimpleNamespace(
@@ -355,3 +374,36 @@ def test_talker2code2wav_async_chunk_emits_terminal_eof_without_duplicate_audio(
     assert payload_final is not None
     assert payload_final.meta.finished.item() is True
     assert payload_final.codes.audio.tolist() == []
+
+
+def test_full_payload_retains_true_length_for_padded_conditioning():
+    from vllm_omni.data_entry_keys import to_struct
+
+    token = torch.tensor([[7, 8, 0, 0]], dtype=torch.int32)
+    feat = torch.randn(1, 8, 80)
+    speaker = torch.randn(1, 192)
+    length = torch.tensor([[2]], dtype=torch.int32)
+    result = text2flow_full_payload(
+        None,
+        {
+            "embed.speech_token": token,
+            "embed.speech_feat": feat,
+            "embed.embedding": speaker,
+            "embed.speech_token_len": length,
+        },
+        SimpleNamespace(request_id="padded"),
+    )
+    payload = to_struct(result)
+    assert payload.embed.speech_token_len is length
+
+
+def test_reference_features_are_snapshots_under_output_accumulation():
+    from vllm_omni.outputs.mm_outputs import MultimodalPayload
+    from vllm_omni.outputs.output_modality import OutputModality
+
+    first = torch.ones(1, 174, 80)
+    latest = torch.zeros(1, 1500, 80)
+    payload = MultimodalPayload.from_dict({"embed.speech_feat": first})
+    payload = payload.merged_with(MultimodalPayload.from_dict({"embed.speech_feat": latest}))
+    payload.consolidate_tensors(OutputModality.LATENT)
+    assert payload["embed.speech_feat"] is latest

tests/model_executor/test_seeded_exponential_platform.py

@@ -0,0 +1,116 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from unittest.mock import Mock
+
+import pytest
+import torch
+
+from tests.helpers.mark import hardware_test
+from vllm_omni.utils import seeded_exponential
+
+pytestmark = [pytest.mark.core_model]
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize("is_nvidia", [False, True], ids=["rocm", "cuda"])
+def test_cuda_tensor_requires_nvidia_distribution_contract(monkeypatch, is_nvidia: bool) -> None:
+    monkeypatch.setattr(seeded_exponential.current_platform, "is_cuda", lambda: is_nvidia)
+    # HIP tensors expose the same Torch device predicate as NVIDIA tensors.
+    q = Mock(spec=torch.Tensor)
+    q.is_cuda = True
+    q.dtype = torch.float32
+    q.dim.return_value = 2
+    q.is_contiguous.return_value = True
+    generators = {0: object(), 1: object()}
+
+    assert seeded_exponential.batched_seeded_exponential_supported(q, generators) is is_nvidia
+
+
+@pytest.mark.cpu
+def test_shared_generator_keeps_ordered_torch_draws(monkeypatch) -> None:
+    monkeypatch.setattr(seeded_exponential.current_platform, "is_cuda", lambda: True)
+    q = Mock(spec=torch.Tensor)
+    q.is_cuda = True
+    q.dtype = torch.float32
+    q.dim.return_value = 2
+    q.is_contiguous.return_value = True
+    shared = object()
+
+    assert not seeded_exponential.batched_seeded_exponential_supported(q, {0: shared, 1: shared})
+
+
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
+@pytest.mark.parametrize("seeded_rows", [(0, 1, 2), (0, 2)], ids=["all_seeded", "mixed"])
+@pytest.mark.parametrize("use_fp64_gumbel", [False, True])
+def test_torch_fallback_preserves_samples_and_generator_state(monkeypatch, seeded_rows, use_fp64_gumbel) -> None:
+    from vllm.v1.sample.ops import topk_topp_sampler as sampler_ops
+
+    import vllm_omni.patch  # noqa: F401
+
+    original = sampler_ops.random_sample.__wrapped__
+    monkeypatch.setattr(seeded_exponential.current_platform, "is_cuda", lambda: False)
+
+    def reject_kernel(*args, **kwargs):
+        raise AssertionError("the CUDA distribution kernel must not run on the fallback path")
+
+    monkeypatch.setattr(seeded_exponential, "fill_exponential_rows", reject_kernel)
+    probs = torch.arange(1, 258, device="cuda", dtype=torch.float32).expand(3, -1).contiguous()
+    probs /= probs.sum(dim=-1, keepdim=True)
+    expected_gens = {row: torch.Generator(device="cuda").manual_seed(42 + row) for row in seeded_rows}
+    actual_gens = {row: torch.Generator(device="cuda").manual_seed(42 + row) for row in seeded_rows}
+    expected_default = torch.cuda.get_rng_state()
+    actual_default = expected_default.clone()
+    try:
+        for _ in range(5):
+            torch.cuda.set_rng_state(expected_default)
+            # FP32 sampling divides probabilities in place. Give both paths
+            # the same untouched input on every step.
+            expected = original(probs.clone(), expected_gens, use_fp64_gumbel)
+            expected_default = torch.cuda.get_rng_state()
+            torch.cuda.set_rng_state(actual_default)
+            actual = sampler_ops.random_sample(probs.clone(), actual_gens, use_fp64_gumbel)
+            actual_default = torch.cuda.get_rng_state()
+            assert torch.equal(actual, expected)
+            assert torch.equal(actual_default, expected_default)
+            for row in seeded_rows:
+                assert torch.equal(actual_gens[row].get_state(), expected_gens[row].get_state())
+    finally:
+        torch.cuda.set_rng_state(expected_default)
+
+
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
+@pytest.mark.parametrize("rows", [None, [0, 2, 4]], ids=["grouped_codebooks", "selected_rows"])
+def test_direct_fill_fallback_preserves_ordered_draws(monkeypatch, rows) -> None:
+    monkeypatch.setattr(seeded_exponential.current_platform, "is_cuda", lambda: False)
+
+    def reject_kernel(*args, **kwargs):
+        raise AssertionError("the NVIDIA kernel must not run for direct fallback callers")
+
+    monkeypatch.setattr(seeded_exponential, "torch_exponential_policy", reject_kernel)
+    expected = torch.ones(6, 257, device="cuda")
+    actual = expected.clone()
+    expected_gen = torch.Generator(device="cuda").manual_seed(71)
+    actual_gen = torch.Generator(device="cuda").manual_seed(71)
+    # A shared seeded generator and an unseeded request make loop order and
+    # default-state advancement observable, including grouped codebooks.
+    expected_gens = [expected_gen, None, expected_gen]
+    actual_gens = [actual_gen, None, actual_gen]
+    original_default = torch.cuda.get_rng_state()
+    expected_default = original_default.clone()
+    actual_default = original_default.clone()
+    try:
+        for _ in range(5):
+            torch.cuda.set_rng_state(expected_default)
+            targets = expected.reshape(3, -1) if rows is None else [expected[row] for row in rows]
+            for target, generator in zip(targets, expected_gens):
+                target.exponential_(generator=generator)
+            expected_default = torch.cuda.get_rng_state()
+            torch.cuda.set_rng_state(actual_default)
+            assert seeded_exponential.fill_exponential_rows(actual, actual_gens, rows) is actual
+            actual_default = torch.cuda.get_rng_state()
+            assert torch.equal(actual, expected)
+            assert torch.equal(actual_gen.get_state(), expected_gen.get_state())
+            assert torch.equal(actual_default, expected_default)
+    finally:
+        torch.cuda.set_rng_state(original_default)