[Maintenance] Refresh MOSS-TTS optimization on main

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

tests/model_executor/models/audio8_tts/test_audio8_tts_codec_decoder.py

@@ -0,0 +1,141 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import functools
+
+import pytest
+import torch
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+NUM_CODEBOOKS = 10
+#: Samples the fake codec returns per frame, so trims are easy to assert.
+SAMPLES_PER_FRAME = 100
+
+
+@functools.lru_cache(maxsize=1)
+def _decoder_cls():
+    from vllm_omni.model_executor.models.audio8_tts.audio8_tts_codec_decoder import Audio8TTSCodecDecoder
+
+    return Audio8TTSCodecDecoder
+
+
+class _FakeCodec:
+    """Returns ``SAMPLES_PER_FRAME`` ramp samples per frame."""
+
+    def __init__(self):
+        self.calls: list[tuple[int, ...]] = []
+
+    def decode(self, codes_bqf: torch.Tensor) -> torch.Tensor:
+        self.calls.append(tuple(codes_bqf.shape))
+        frames = int(codes_bqf.shape[-1])
+        total = frames * SAMPLES_PER_FRAME
+        return torch.arange(total, dtype=torch.float32).view(1, 1, total)
+
+
+def _make_decoder(codec: _FakeCodec, num_codebooks: int = NUM_CODEBOOKS, splits=None):
+    decoder = object.__new__(_decoder_cls())
+    torch.nn.Module.__init__(decoder)
+    decoder._codec = codec
+    decoder._codec_device = torch.device("cpu")
+    decoder._codec_dtype = torch.float32
+    decoder._num_codebooks = num_codebooks
+    decoder._sample_rate = 44100
+    decoder._frame_size = 2048
+    decoder._logged_stats = True
+    decoder._ensure_codec_loaded = lambda: None
+    if splits is not None:
+        decoder._split_request_ids = lambda ids, seq_token_counts=None: splits
+    return decoder
+
+
+def test_left_context_frames_are_trimmed_from_the_emitted_chunk():
+    """The decoder must emit *delta* audio: the left-context prefix exists only
+    so the causal codec has history and must never reach the client twice."""
+    codec = _FakeCodec()
+    decoder = _make_decoder(codec)
+    codes = torch.arange(NUM_CODEBOOKS * 4, dtype=torch.long).reshape(NUM_CODEBOOKS, 4)
+
+    out = decoder.forward(
+        input_ids=torch.tensor([0], dtype=torch.long),
+        runtime_additional_information=[{"codes": {"audio": codes}, "meta": {"left_context_size": 3}}],
+    )
+
+    audios = out.multimodal_outputs["model_outputs"]
+    assert len(audios) == 1
+    # 4 frames decoded, 3 of context trimmed => 1 frame of new audio.
+    assert audios[0].shape[0] == SAMPLES_PER_FRAME
+    assert codec.calls == [(1, NUM_CODEBOOKS, 4)]
+
+
+def test_tensor_codes_from_runtime_info_win_over_input_ids():
+    codec = _FakeCodec()
+    decoder = _make_decoder(codec, num_codebooks=2)
+    out = decoder.forward(
+        input_ids=torch.zeros(64, dtype=torch.long),
+        runtime_additional_information=[{"codes": {"audio": torch.tensor([[1, 2], [3, 4]], dtype=torch.long)}}],
+    )
+    assert codec.calls == [(1, 2, 2)]
+    assert out.multimodal_outputs["model_outputs"][0].shape[0] == 2 * SAMPLES_PER_FRAME
+
+
+def test_flat_input_ids_are_reshaped_codebook_major():
+    codec = _FakeCodec()
+    decoder = _make_decoder(codec, num_codebooks=2)
+    out = decoder.forward(input_ids=torch.arange(6, dtype=torch.long), runtime_additional_information=[{}])
+    assert codec.calls == [(1, 2, 3)]
+    assert out.multimodal_outputs["model_outputs"][0].shape[0] == 3 * SAMPLES_PER_FRAME
+
+
+def test_every_return_path_sets_model_outputs_and_sr():
+    """A branch that omits ``model_outputs`` makes the serving layer drop the
+    chunk silently, so assert the key exists on all the early-exit paths."""
+    codec = _FakeCodec()
+
+    empty = _make_decoder(codec).forward(input_ids=torch.zeros(0, dtype=torch.long))
+    assert empty.multimodal_outputs["model_outputs"][0].numel() == 0
+    assert empty.multimodal_outputs["sr"]
+
+    ragged = _make_decoder(codec, num_codebooks=4).forward(
+        input_ids=torch.arange(6, dtype=torch.long),  # 6 % 4 != 0
+        runtime_additional_information=[{}],
+    )
+    assert ragged.multimodal_outputs["model_outputs"][0].numel() == 0
+
+    codes = torch.zeros((NUM_CODEBOOKS, 2), dtype=torch.long)
+    no_new = _make_decoder(codec).forward(
+        input_ids=torch.tensor([0], dtype=torch.long),
+        runtime_additional_information=[{"codes": {"audio": codes}, "meta": {"left_context_size": 2}}],
+    )
+    assert no_new.multimodal_outputs["model_outputs"][0].numel() == 0
+    assert codec.calls == [], "a context-only chunk must not be decoded"
+
+
+def test_mixed_batch_keeps_per_request_alignment():
+    """An empty request must not shift the following request's audio into its
+    slot: the output list index is the request index."""
+    codec = _FakeCodec()
+    splits = [
+        torch.empty((0,), dtype=torch.long),
+        torch.arange(NUM_CODEBOOKS * 2, dtype=torch.long),
+    ]
+    decoder = _make_decoder(codec, splits=splits)
+    out = decoder.forward(
+        input_ids=torch.arange(NUM_CODEBOOKS * 2, dtype=torch.long),
+        runtime_additional_information=[{}, {"meta": {"left_context_size": 1}}],
+    )
+    audios = out.multimodal_outputs["model_outputs"]
+    assert len(audios) == 2
+    assert audios[0].numel() == 0
+    assert audios[1].shape[0] == SAMPLES_PER_FRAME
+
+
+def test_codes_with_wrong_codebook_count_fall_back_to_input_ids():
+    codec = _FakeCodec()
+    decoder = _make_decoder(codec, num_codebooks=2)
+    out = decoder.forward(
+        input_ids=torch.arange(4, dtype=torch.long),
+        runtime_additional_information=[{"codes": {"audio": torch.zeros((5, 3), dtype=torch.long)}}],
+    )
+    assert codec.calls == [(1, 2, 2)]
+    assert out.multimodal_outputs["model_outputs"][0].shape[0] == 2 * SAMPLES_PER_FRAME

tests/model_executor/models/audio8_tts/test_audio8_tts_prompt_utils.py

@@ -0,0 +1,112 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import pytest
+
+from vllm_omni.model_executor.models.audio8_tts.prompt_utils import (
+    build_text_only_prompt_ids,
+    build_voice_clone_prompt_ids,
+    estimate_voice_clone_prompt_len,
+    normalize_text,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _SegmentTokenizer:
+    """Tokenizer stub that maps each encoded string to one sentinel id.
+
+    Lets the tests assert the *segment sequence* the prompt builder emits, which
+    is what must stay byte-identical to the reference ``ArkttsProcessor``.
+    """
+
+    def __init__(self):
+        self.vocab: dict[str, int] = {}
+
+    def encode(self, text: str, add_special_tokens: bool = True):
+        assert add_special_tokens is False, "Audio8 prompts must not add extra special tokens"
+        return [self.vocab.setdefault(text, 1000 + len(self.vocab))]
+
+    def segment_of(self, token_id: int) -> str:
+        return next(text for text, tid in self.vocab.items() if tid == token_id)
+
+
+def _segments(tokenizer: _SegmentTokenizer, ids: list[int]) -> list[str]:
+    return [tokenizer.segment_of(i) for i in ids]
+
+
+def test_text_only_prompt_matches_reference_segment_order():
+    tokenizer = _SegmentTokenizer()
+    ids, normalized = build_text_only_prompt_ids(tokenizer, "  Hello   world  ")
+
+    assert normalized == "Hello world"
+    assert _segments(tokenizer, ids) == [
+        "<|im_start|>system\n",
+        "convert the provided text to speech",
+        "<|im_end|>\n",
+        "<|im_start|>user\n",
+        "Hello world",
+        "<|im_end|>\n",
+        "<|im_start|>assistant\n<|voice|>",
+    ]
+
+
+def test_voice_clone_prompt_places_reference_codes_between_prefix_and_suffix():
+    tokenizer = _SegmentTokenizer()
+    semantic_ids = [151678, 151679, 151680]
+    ids, ref_start, text, ref_text = build_voice_clone_prompt_ids(
+        tokenizer, "Target text", "Reference transcript", semantic_ids
+    )
+
+    assert text == "Target text"
+    # A reference transcript without an explicit tag gets speaker 0.
+    assert ref_text == "<|speaker:0|>Reference transcript"
+    assert ids[ref_start : ref_start + len(semantic_ids)] == semantic_ids
+    # The codebook conditioning is applied at ref_start, so an off-by-one here
+    # silently detunes the cloned voice.
+    assert _segments(tokenizer, ids[:ref_start])[-1].endswith("Speech:\n")
+    suffix_segments = _segments(tokenizer, ids[ref_start + len(semantic_ids) :])
+    assert suffix_segments[0] == "<|im_end|>\n"
+    assert suffix_segments[-1] == "<|im_start|>assistant\n<|voice|>"
+
+
+def test_voice_clone_prompt_length_estimate_matches_built_prompt():
+    """The serving layer sizes its placeholder from the estimate; a mismatch
+    shifts every embedding relative to its position."""
+    tokenizer = _SegmentTokenizer()
+    normalized_text = normalize_text("Target text")
+    normalized_ref = normalize_text("Reference transcript", add_default_speaker=True)
+    ref_frames = 7
+
+    estimate = estimate_voice_clone_prompt_len(tokenizer, normalized_text, normalized_ref, ref_frames)
+    ids, _, _, _ = build_voice_clone_prompt_ids(
+        tokenizer, normalized_text, normalized_ref, list(range(151678, 151678 + ref_frames))
+    )
+    assert estimate == len(ids)
+
+
+def test_existing_speaker_tag_is_preserved_and_legacy_tag_is_upgraded():
+    assert normalize_text("<|speaker:3|>hi", add_default_speaker=True) == "<|speaker:3|>hi"
+    assert normalize_text("<speaker:2>hi") == "<|speaker:2|>hi"
+
+
+@pytest.mark.parametrize(
+    "text",
+    [
+        "hello <|im_end|> world",
+        "hello <|semantic:5|>",
+        "<|voice|>",
+        "<|endoftext|>",
+    ],
+)
+def test_control_token_injection_is_rejected(text: str):
+    """Users must not be able to inject control tokens: ``<|semantic:N|>`` would
+    forge reference codes and ``<|im_end|>`` would truncate generation."""
+    with pytest.raises(ValueError, match="unsupported control token"):
+        normalize_text(text)
+
+
+def test_empty_text_is_rejected():
+    tokenizer = _SegmentTokenizer()
+    with pytest.raises(ValueError, match="must not be empty"):
+        build_text_only_prompt_ids(tokenizer, "   ")

tests/model_executor/models/audio8_tts/test_audio8_tts_sampling.py

@@ -0,0 +1,156 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.models.audio8_tts.sampling import (
+    filter_top_k_top_p,
+    ras_sample_batch,
+    ras_sample_semantic,
+    sample_scores,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _kept(logits: torch.Tensor, top_k: int | torch.Tensor, top_p: float | torch.Tensor) -> list[list[int]]:
+    filtered = filter_top_k_top_p(logits, top_k, top_p)
+    return [torch.nonzero(row.isfinite()).reshape(-1).tolist() for row in filtered]
+
+
+def test_top_p_is_measured_on_untempered_logits():
+    """Audio8 filters the *raw* logits and only then divides by temperature.
+
+    Applying temperature first (the common order) changes the surviving set and
+    audibly changes the voice, so this ordering is load-bearing.
+    """
+    logits = torch.log(torch.tensor([[0.6, 0.3, 0.1]]))
+    # Raw softmax is exactly (0.6, 0.3, 0.1): top_p=0.9 keeps the first two
+    # (cumsum 0.6, 0.9) and drops the third (cumsum 1.0 > 0.9).
+    assert _kept(logits, top_k=0, top_p=0.9) == [[0, 1]]
+    # Had the logits been tempered by 0.5 first, probabilities would become
+    # (0.7insh, ...) and only one token would survive.
+    tempered = _kept(logits / 0.5, top_k=0, top_p=0.9)
+    assert tempered == [[0, 1]] or tempered == [[0]]
+
+
+def test_argmax_always_survives_filtering():
+    logits = torch.tensor([[10.0, -1.0, -2.0]])
+    # top_p=0.0 would otherwise remove every candidate and produce NaNs.
+    assert _kept(logits, top_k=1, top_p=0.0) == [[0]]
+
+
+def test_per_row_top_k_and_top_p_are_honoured():
+    logits = torch.log(torch.tensor([[0.5, 0.3, 0.15, 0.05], [0.5, 0.3, 0.15, 0.05]]))
+    kept = _kept(logits, top_k=torch.tensor([1, 3]), top_p=torch.tensor([1.0, 1.0]))
+    assert kept == [[0], [0, 1, 2]]
+
+
+def test_zero_temperature_row_falls_back_to_argmax():
+    logits = torch.tensor([[1.0, 5.0, 2.0], [1.0, 5.0, 2.0]])
+    out = sample_scores(
+        logits,
+        temperature=torch.tensor([0.0, 1.0]),
+        top_k=0,
+        top_p=1.0,
+        generator=torch.Generator().manual_seed(0),
+    )
+    assert int(out[0]) == 1
+
+
+def test_ras_substitutes_only_repeated_semantic_ids():
+    """RAS re-rolls a repeat from a flatter distribution; the EOS slot and
+    non-repeats must pass through untouched, or utterances get truncated."""
+    num_semantic = 2
+    # Row 0 will draw id 0 (a repeat) -> substituted.
+    # Row 1 will draw id 2 (the EOS slot) -> never substituted even if recent.
+    logits = torch.tensor(
+        [
+            [10.0, -20.0, -20.0],
+            [-20.0, -20.0, 10.0],
+        ]
+    )
+    recent = torch.tensor([[0, -1], [2, -1]])
+    out = ras_sample_semantic(
+        logits,
+        recent,
+        temperature=1e-6,  # deterministic: argmax within each draw
+        top_k=0,
+        top_p=1.0,
+        ras_temperature=1e-6,
+        ras_top_p=1.0,
+        num_semantic_ids=num_semantic,
+    )
+    # Both draws are argmax, so substitution is observable only through which
+    # branch was taken; assert the EOS row is preserved verbatim.
+    assert int(out[1]) == 2
+    assert int(out[0]) == 0
+
+
+def test_ras_without_history_returns_the_plain_draw():
+    logits = torch.tensor([[10.0, 0.0, -5.0]])
+    out = ras_sample_semantic(
+        logits,
+        None,
+        temperature=1e-6,
+        top_k=0,
+        top_p=1.0,
+        ras_temperature=1.0,
+        ras_top_p=0.9,
+        num_semantic_ids=2,
+    )
+    assert int(out[0]) == 0
+
+
+def test_sampling_is_reproducible_for_a_seeded_generator():
+    logits = torch.randn(4, 32, generator=torch.Generator().manual_seed(7))
+    first = sample_scores(logits, temperature=0.7, top_k=8, top_p=0.9, generator=torch.Generator().manual_seed(1234))
+    second = sample_scores(logits, temperature=0.7, top_k=8, top_p=0.9, generator=torch.Generator().manual_seed(1234))
+    assert torch.equal(first, second)
+
+
+def test_batch_routes_each_row_to_its_own_generator():
+    """Regression for the shared slot-0 generator bug.
+
+    A shared generator over ``[num_reqs, vocab]`` made every row draw from slot
+    0's RNG stream, so a request's audio depended on its batch neighbours and
+    per-request ``seed`` was ignored for slots > 0. Flat logits make the draw
+    purely noise-driven so the wrong generator changes the token.
+    """
+    # Flat logits: the sampled id is determined only by the generator's noise,
+    # so routing the generator to the wrong slot changes the result (coincidence ~1/256).
+    logits = torch.zeros(2, 256)
+    kw = dict(temperature=1.0, top_k=0, top_p=1.0, ras_temperature=1.0, ras_top_p=0.9, num_semantic_ids=200)
+
+    # (1) Each batched row equals sampling that row alone with its own generator.
+    batched = ras_sample_batch(
+        logits,
+        None,
+        generators={0: torch.Generator().manual_seed(11), 1: torch.Generator().manual_seed(22)},
+        num_reqs=2,
+        **kw,
+    )
+    row0 = ras_sample_semantic(logits[0:1], None, generator=torch.Generator().manual_seed(11), **kw)
+    row1 = ras_sample_semantic(logits[1:2], None, generator=torch.Generator().manual_seed(22), **kw)
+    assert int(batched[0]) == int(row0)
+    assert int(batched[1]) == int(row1)
+
+    # (2) Row 1 must not change when only slot 0's seed changes.
+    other = ras_sample_batch(
+        logits,
+        None,
+        generators={0: torch.Generator().manual_seed(9999), 1: torch.Generator().manual_seed(22)},
+        num_reqs=2,
+        **kw,
+    )
+    assert int(other[1]) == int(batched[1])
+
+
+def test_single_request_batch_keeps_the_fast_path():
+    """``num_reqs == 1`` must match a direct single-row draw (no per-row loop)."""
+    logits = torch.randn(1, 48, generator=torch.Generator().manual_seed(5))
+    kw = dict(temperature=0.8, top_k=16, top_p=0.9, ras_temperature=1.0, ras_top_p=0.9, num_semantic_ids=40)
+    batched = ras_sample_batch(logits, None, generators={0: torch.Generator().manual_seed(99)}, num_reqs=1, **kw)
+    direct = ras_sample_semantic(logits, None, generator=torch.Generator().manual_seed(99), **kw)
+    assert torch.equal(batched, direct)

tests/model_executor/models/fish_speech/test_dac_modules.py

@@ -0,0 +1,110 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""CPU tests for the vendored DAC codec modules (dac_modules).
+
+Guards the checkpoint contract: the module tree built by build_dac_codec()
+must keep the exact parameter/buffer names and shapes that codec.pth was
+saved with (verified against the upstream fish-speech 0.1.0 package), and
+the encode/decode surfaces used by dac_encoder.py / fish_speech_dac_decoder.py
+must keep their shapes.
+"""
+
+import pytest
+import torch
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+# Spot checks covering every submodule family and both weight-norm styles.
+# Derived from the upstream fish-speech 0.1.0 state dict (541 entries total);
+# legacy weight_norm (weight_g/weight_v) comes from the descript quantizer
+# projections, parametrizations.weight.original0/1 from the causal convs.
+EXPECTED_KEY_SAMPLES = [
+    # encoder: first conv + deepest encoder-block transformer
+    "encoder.block.0.conv.parametrizations.weight.original0",
+    "encoder.block.4.block.5.layers.3.attention.wqkv.weight",
+    "encoder.block.4.block.5.freqs_cis",
+    "encoder.block.1.block.0.block.0.alpha",
+    # quantizer: semantic + residual VQ (legacy weight_norm keys)
+    "quantizer.semantic_quantizer.quantizers.0.codebook.weight",
+    "quantizer.quantizer.quantizers.8.in_proj.weight_g",
+    "quantizer.quantizer.quantizers.8.in_proj.weight_v",
+    # quantizer pre/post transformers + down/upsample convnext
+    "quantizer.pre_module.layers.7.feed_forward.w1.weight",
+    "quantizer.post_module.norm.weight",
+    "quantizer.downsample.1.1.dwconv.conv.weight",
+    "quantizer.upsample.0.0.conv.weight",
+    # decoder: transposed conv + final conv
+    "decoder.model.1.block.1.conv.parametrizations.weight.original1",
+    "decoder.model.6.conv.parametrizations.weight.original0",
+]
+
+NUM_STATE_DICT_ENTRIES = 541
+DAC_NUM_CODEBOOKS = 10  # 1 semantic + 9 residual
+
+
+@pytest.fixture(scope="module")
+def codec():
+    from vllm_omni.model_executor.models.fish_speech.dac_utils import build_dac_codec
+
+    model = build_dac_codec()
+    model.eval()
+    return model
+
+
+def test_state_dict_contract(codec):
+    sd = codec.state_dict()
+    assert len(sd) == NUM_STATE_DICT_ENTRIES
+    for key in EXPECTED_KEY_SAMPLES:
+        assert key in sd, f"checkpoint-contract key missing: {key}"
+
+
+def test_no_fish_speech_import_required(codec):
+    import sys
+
+    assert "fish_speech" not in sys.modules, "vendored codec must not import the fish-speech package"
+
+
+@torch.no_grad()
+def test_encode_shapes(codec):
+    frame_length = codec.frame_length  # 2048 samples per code frame
+    wav = torch.randn(1, 1, frame_length * 3 + 100)
+    lengths = torch.tensor([wav.shape[-1]], dtype=torch.long)
+    codes, code_lengths = codec.encode(wav, lengths)
+    assert codes.shape[0] == 1
+    assert codes.shape[1] == DAC_NUM_CODEBOOKS
+    assert codes.shape[2] == 4  # ceil(3.x frames)
+    assert code_lengths.tolist() == [4]
+    assert codes.dtype in (torch.int64, torch.int32)
+
+
+@torch.no_grad()
+def test_decode_shapes(codec):
+    frames = 8
+    codes = torch.randint(0, 1024, (1, DAC_NUM_CODEBOOKS, frames), dtype=torch.long)
+    feature_lengths = torch.tensor([frames], dtype=torch.long)
+    wav, audio_lengths = codec.decode(codes, feature_lengths)
+    assert wav.shape[0] == 1 and wav.shape[1] == 1
+    assert audio_lengths.tolist() == [frames * codec.frame_length]
+    assert wav.shape[-1] >= audio_lengths[0] - codec.frame_length
+
+
+@torch.no_grad()
+def test_decoder_runner_mutations(codec):
+    """The stage-1 runner prunes encode-only parts and bakes weight norm;
+    decode must still work afterwards (mirrors fish_speech_dac_decoder)."""
+    from torch.nn.utils.parametrize import remove_parametrizations
+
+    codec.encoder = None
+    codec.quantizer.pre_module = None
+    codec.quantizer.downsample = None
+    for module in codec.modules():
+        parametrizations = getattr(module, "parametrizations", None)
+        if not parametrizations:
+            continue
+        for name in list(parametrizations.keys()):
+            remove_parametrizations(module, name, leave_parametrized=True)
+
+    codes = torch.randint(0, 1024, (2, DAC_NUM_CODEBOOKS, 5), dtype=torch.long)
+    wav, audio_lengths = codec.decode(codes, torch.tensor([5, 3], dtype=torch.long))
+    assert wav.shape[0] == 2
+    assert audio_lengths.tolist() == [5 * codec.frame_length, 3 * codec.frame_length]

tests/model_executor/models/higgs_audio_v3/test_higgs_audio_v3.py

@@ -11,13 +11,15 @@ import asyncio
 import json
 import os
 import time
-from collections import OrderedDict, defaultdict
+from collections import defaultdict
 from types import SimpleNamespace
 
 import pytest
 import torch
 
 from tests.helpers.stage_config import get_deploy_config_path
+from vllm_omni.entrypoints.openai.tts_adapters.base import SpeechServingContext
+from vllm_omni.entrypoints.openai.tts_adapters.higgs_audio_v3 import HiggsAudioV3Adapter
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -1683,31 +1685,21 @@ class TestVoiceCloneReferenceCache:
         assert tok._resolve_audio_tokenizer_dir() is None
 
     def test_higgs_v3_ref_code_cache_returns_clone(self):
-        from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech
-
-        serving = object.__new__(OmniOpenAIServingSpeech)
-        serving._higgs_audio_v3_ref_code_cache = OrderedDict()
-        serving._higgs_audio_v3_ref_code_cache_bytes = 0
-        serving._higgs_audio_v3_ref_code_inflight = {}
+        adapter = HiggsAudioV3Adapter(SpeechServingContext(server=SimpleNamespace()))
 
         codes = torch.arange(16, dtype=torch.long).reshape(2, 8)
-        serving._put_higgs_audio_v3_ref_codes("ref-a", codes)
+        adapter._put_higgs_audio_v3_ref_codes("ref-a", codes)
 
-        cached = serving._get_higgs_audio_v3_ref_codes("ref-a")
+        cached = adapter._get_higgs_audio_v3_ref_codes("ref-a")
         assert cached is not None
         cached.fill_(0)
 
-        cached_again = serving._get_higgs_audio_v3_ref_codes("ref-a")
+        cached_again = adapter._get_higgs_audio_v3_ref_codes("ref-a")
         assert cached_again is not None
         assert torch.equal(cached_again, codes)
 
     def test_higgs_v3_ref_code_inflight_deduplicates_concurrent_encode(self):
-        from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech
-
-        serving = object.__new__(OmniOpenAIServingSpeech)
-        serving._higgs_audio_v3_ref_code_cache = OrderedDict()
-        serving._higgs_audio_v3_ref_code_cache_bytes = 0
-        serving._higgs_audio_v3_ref_code_inflight = {}
+        adapter = HiggsAudioV3Adapter(SpeechServingContext(server=SimpleNamespace()))
         calls = 0
 
         def encode_reference_audio(_wav, _sr):
@@ -1722,7 +1714,7 @@ class TestVoiceCloneReferenceCache:
         async def run():
             return await asyncio.gather(
                 *[
-                    serving._resolve_higgs_audio_v3_ref_codes(
+                    adapter._resolve_higgs_audio_v3_ref_codes(
                         "ref-a",
                         object(),
                         24000,
@@ -1741,7 +1733,7 @@ class TestVoiceCloneReferenceCache:
             assert torch.equal(codes, torch.arange(16, dtype=torch.long).reshape(2, 8) + 1)
 
         cached, cache_hit, inflight_wait = asyncio.run(
-            serving._resolve_higgs_audio_v3_ref_codes(
+            adapter._resolve_higgs_audio_v3_ref_codes(
                 "ref-a",
                 object(),
                 24000,

tests/model_executor/models/mimo_audio/test_mimo_audio_code2wav_batch_decode.py

@@ -91,7 +91,11 @@ def _minimal_model(mocker: MockerFixture):
         decoder=decoder,
         config=SimpleNamespace(avg_pooler=2, stride_size=2, hop_length=240),
     )
-    model._tokenizer_service = SimpleNamespace(audio_tokenizer=audio_tok, cuda_graph_wrapper=None)
+    model._tokenizer_service = SimpleNamespace(
+        audio_tokenizer=audio_tok,
+        cuda_graph_wrapper=None,
+        device="cpu",
+    )
     return model, audio_tok
 
 

tests/model_executor/models/mimo_audio/test_mimo_audio_tokenizer_worker_cache.py

@@ -0,0 +1,115 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Regression tests for the MiMo-Audio tokenizer-worker cache key.
+
+The worker cache key used to keep only the device *type*, so an explicit
+``MIMO_AUDIO_TOKENIZER_DEVICE=cuda:1`` silently loaded the tokenizer on the
+process's current CUDA device and shared one cache entry with ``cuda:0``.
+The index must survive both the cache key and the worker constructor.
+
+CPU-only and weight-free: the worker class is replaced by a stub so no
+tokenizer weights are loaded.
+"""
+
+from __future__ import annotations
+
+import functools
+
+import pytest
+
+torch = pytest.importorskip("torch")
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_TOK_PATH = "/models/mimo-audio-tokenizer"
+
+
+@functools.lru_cache(maxsize=1)
+def _code2wav_module():
+    """Defer the import (pulls vLLM model_executor) until first use."""
+    from vllm_omni.model_executor.models.mimo_audio import mimo_audio_code2wav
+
+    return mimo_audio_code2wav
+
+
+class _StubWorker:
+    def __init__(self, device_str: str, config_path: str | None, audio_tokenizer_path: str):
+        self.device = device_str
+        self.config_path = config_path
+        self.audio_tokenizer_path = audio_tokenizer_path
+
+
+@pytest.fixture
+def isolated_worker_cache(monkeypatch):
+    """Empty worker cache with the real worker replaced by a stub."""
+    mod = _code2wav_module()
+    monkeypatch.setattr(mod, "MiMoAudioTokenizerWorker", _StubWorker)
+    monkeypatch.setattr(mod, "_TOKENIZER_WORKER_CACHE", {})
+    return mod
+
+
+@pytest.mark.parametrize("device", [torch.device("cuda:1"), "cuda:1"])
+def test_cache_key_keeps_cuda_index(device):
+    """Both torch.device and string spellings retain the index."""
+    mod = _code2wav_module()
+    key = mod._normalize_tokenizer_worker_cache_key(device, None, _TOK_PATH)
+    assert key[0] == "cuda:1"
+
+
+def test_cache_key_cpu_unchanged():
+    mod = _code2wav_module()
+    key = mod._normalize_tokenizer_worker_cache_key(torch.device("cpu"), None, _TOK_PATH)
+    assert key[0] == "cpu"
+
+
+def test_cache_key_bare_cuda_resolves_to_current_device(monkeypatch):
+    """A bare "cuda" shares the entry with its equivalent indexed spelling."""
+    mod = _code2wav_module()
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
+    monkeypatch.setattr(torch.accelerator, "current_device_index", lambda: 3)
+
+    bare = mod._normalize_tokenizer_worker_cache_key(torch.device("cuda"), None, _TOK_PATH)
+    indexed = mod._normalize_tokenizer_worker_cache_key(torch.device("cuda:3"), None, _TOK_PATH)
+    assert bare[0] == "cuda:3"
+    assert bare == indexed
+
+
+def test_worker_constructed_on_requested_index(isolated_worker_cache):
+    mod = isolated_worker_cache
+    worker = mod.get_tokenizer_worker(
+        device=torch.device("cuda:1"),
+        config_path=None,
+        audio_tokenizer_path=_TOK_PATH,
+    )
+    assert worker.device == "cuda:1"
+
+
+def test_distinct_cuda_indices_get_distinct_workers(isolated_worker_cache):
+    mod = isolated_worker_cache
+    first = mod.get_tokenizer_worker(
+        device=torch.device("cuda:0"),
+        config_path=None,
+        audio_tokenizer_path=_TOK_PATH,
+    )
+    second = mod.get_tokenizer_worker(
+        device=torch.device("cuda:1"),
+        config_path=None,
+        audio_tokenizer_path=_TOK_PATH,
+    )
+    assert first is not second
+    assert (first.device, second.device) == ("cuda:0", "cuda:1")
+
+
+def test_same_device_reuses_cached_worker(isolated_worker_cache):
+    mod = isolated_worker_cache
+    first = mod.get_tokenizer_worker(
+        device=torch.device("cuda:1"),
+        config_path=None,
+        audio_tokenizer_path=_TOK_PATH,
+    )
+    second = mod.get_tokenizer_worker(
+        device="cuda:1",
+        config_path=None,
+        audio_tokenizer_path=_TOK_PATH,
+    )
+    assert first is second

tests/model_executor/models/ming_tts/test_ming_shared_modules.py

@@ -1,12 +1,16 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from types import SimpleNamespace
 
 import pytest
 import torch
 
-from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech
+from vllm_omni.entrypoints.openai.tts_adapters.base import SpeechServingContext
+from vllm_omni.entrypoints.openai.tts_adapters.ming_tts import (
+    MingTTSAdapter,
+    parse_ming_instruction_fields,
+)
 from vllm_omni.model_executor.models.common.ming.aggregator import Aggregator
 from vllm_omni.model_executor.models.common.ming.audio_vae import AudioVAEConfig
 from vllm_omni.model_executor.models.common.ming.cfm_cudagraph import CFMGraphExecutor, CFMSampler
@@ -174,19 +178,19 @@ def test_ming_dense_validation_rejects_semantic_audio_vae_config():
 
 def test_ming_instruction_parser_preserves_dense_and_flash_defaults():
     """Ming dense and Ming flash keep distinct instruction defaults."""
-    serving = object.__new__(OmniOpenAIServingSpeech)
-    serving.uploaded_speakers = {"uploaded": {}}
+    serving = SimpleNamespace(uploaded_speakers={"uploaded": {}})
+    adapter = MingTTSAdapter(SpeechServingContext(server=serving))
 
-    dense_plain = serving._parse_ming_instruction(SimpleNamespace(instructions="calm", language=None, voice=None))
+    dense_plain = adapter._parse_ming_instruction(SimpleNamespace(instructions="calm", language=None, voice=None))
     assert dense_plain == "calm"
 
-    dense_with_fields = serving._parse_ming_instruction(
+    dense_with_fields = adapter._parse_ming_instruction(
         SimpleNamespace(instructions="calm", language="Auto", voice="灵小甄")
     )
     assert dense_with_fields == {"IP": "灵小甄", "风格": "calm"}
 
-    flash_fields = serving._parse_ming_instruction_fields(
-        SimpleNamespace(instructions="calm", language="粤语", voice="灵小甄")
+    flash_fields = parse_ming_instruction_fields(
+        serving, SimpleNamespace(instructions="calm", language="粤语", voice="灵小甄")
     )
     assert flash_fields == {"风格": "calm"}
 

tests/model_executor/models/minicpmo_4_5/test_audio_embeds_input.py

@@ -0,0 +1,138 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Regression tests for the MiniCPM-o 4.5 ``audio_embeds`` input path.
+
+Two defects made the OpenAI ``audio_embeds`` content part unusable with this model:
+
+1. ``MiniCPMOAudioEmbeddingItems`` registered itself under ``modality="image"``
+   (copy-pasted from ``MiniCPMVImageEmbeddingItems``), so every request failed with
+   ``Modality 'image' not found. Available modalities: {'audio'}``.
+2. ``_get_prompt_updates`` sized the placeholder from
+   ``sum(map(len, single_audio_embeds))``. ``MiniCPMOAudioEmbeddingInputs`` declares
+   ``TensorShape("bn", "s", "h")``, so one item is ``(s, h)`` and that expression
+   iterates rows and sums the hidden size -- ``s * h`` instead of ``s``. A one-second
+   slice of shape ``(10, 4096)`` asked for 40960 placeholder tokens instead of 10,
+   pushing every real request past ``max_model_len``.
+"""
+
+from __future__ import annotations
+
+import math
+from types import SimpleNamespace
+
+import pytest
+import torch
+from vllm.multimodal.parse import MultiModalDataItems
+
+from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_llm import (
+    MiniCPMO45OmniLLMMultiModalProcessor,
+    MiniCPMO45OmniLLMProcessingInfo,
+    MiniCPMOAudioEmbeddingItems,
+    _minicpmo_field_config,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+POOL_STEP = 5
+HIDDEN_SIZE = 4096
+#: Audio embedding vectors produced per second: 100 mel -> 50 cnn -> (50-5)//5+1.
+EMBEDS_PER_SECOND = 10
+
+
+class _FakeProcessor:
+    """Minimal processor surface used by ``get_audio_placeholder``."""
+
+    def __init__(self) -> None:
+        self.pool_step = POOL_STEP
+        self.audio_processor = SimpleNamespace(hop_length=160)
+        self.image_processor = SimpleNamespace(mean=[0.0], std=[1.0])
+
+    def get_audio_placeholder(
+        self,
+        audio_lens: int,
+        chunk_input: bool = True,
+        chunk_length: int = 1,
+    ) -> str:
+        del chunk_input, chunk_length
+        feature_lens = math.ceil(audio_lens / self.audio_processor.hop_length)
+        cnn_feature_lens = (feature_lens - 1) // 2 + 1
+        output_lens = (cnn_feature_lens - self.pool_step) // self.pool_step + 1
+        return "<unk>" * output_lens
+
+
+class _IdentityTokenizer:
+    """``_get_prompt_updates`` round-trips the image/video patterns through the
+    tokenizer to detect encode/decode drift; an identity pair keeps that a no-op."""
+
+    def encode(self, text: str, add_special_tokens: bool = True) -> str:
+        del add_special_tokens
+        return text
+
+    def decode(self, tokens: str) -> str:
+        return tokens
+
+
+def _processing_info() -> MiniCPMO45OmniLLMProcessingInfo:
+    processor = _FakeProcessor()
+    tokenizer = _IdentityTokenizer()
+    info = object.__new__(MiniCPMO45OmniLLMProcessingInfo)
+    info.ctx = SimpleNamespace(
+        tokenizer=tokenizer,
+        get_tokenizer=lambda: tokenizer,
+        get_hf_config=lambda: SimpleNamespace(audio_pool_step=POOL_STEP),
+        get_hf_processor=lambda **kwargs: processor,
+    )
+    return info
+
+
+def _audio_placeholder_for(single_audio_embeds: torch.Tensor) -> str:
+    """Run the real prompt-update path and return the audio placeholder text.
+
+    Goes through ``_get_prompt_updates`` rather than recomputing the length, so the
+    ``sum(map(len, ...))`` regression is actually exercised.
+    """
+    info = _processing_info()
+    proc = object.__new__(MiniCPMO45OmniLLMMultiModalProcessor)
+    proc.info = info
+
+    items = MiniCPMOAudioEmbeddingItems(
+        {"audio_embeds": single_audio_embeds.unsqueeze(0)},
+        fields_factory=_minicpmo_field_config,
+    )
+    mm_items = MultiModalDataItems({"audio": items})
+
+    updates = proc._get_prompt_updates(mm_items, {}, {})
+    audio_update = next(u for u in updates if u.modality == "audio")
+    return audio_update.content(0).full
+
+
+def test_audio_embedding_items_use_the_audio_modality() -> None:
+    """Registering under ``image`` makes every ``audio_embeds`` request fail."""
+    items = MiniCPMOAudioEmbeddingItems(
+        {"audio_embeds": torch.zeros(1, EMBEDS_PER_SECOND, HIDDEN_SIZE)},
+        fields_factory=_minicpmo_field_config,
+    )
+
+    assert items.modality == "audio"
+
+
+@pytest.mark.parametrize("num_seconds", [1, 3, 30])
+def test_placeholder_length_matches_embedding_count(num_seconds: int) -> None:
+    """An ``(s, h)`` item must yield exactly ``s`` placeholder tokens."""
+    num_embeds = num_seconds * EMBEDS_PER_SECOND
+
+    placeholder = _audio_placeholder_for(torch.zeros(num_embeds, HIDDEN_SIZE))
+
+    assert placeholder.count("<unk>") == num_embeds
+
+
+def test_placeholder_does_not_scale_with_hidden_size() -> None:
+    """The count must depend on ``s`` alone, never on ``h``.
+
+    ``sum(map(len, ...))`` returned ``s * h``, so the placeholder grew with the model's
+    hidden size. Two items with the same ``s`` but different ``h`` must agree.
+    """
+    narrow = _audio_placeholder_for(torch.zeros(EMBEDS_PER_SECOND, 8))
+    wide = _audio_placeholder_for(torch.zeros(EMBEDS_PER_SECOND, HIDDEN_SIZE))
+
+    assert narrow.count("<unk>") == wide.count("<unk>") == EMBEDS_PER_SECOND

tests/model_executor/models/minicpmo_4_5/test_code2wav_model_dir.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 MiniCPM-o 4.5 code2wav model-dir resolution (#5442).
 
 In hub/CI deployments ``model_config.model`` is a repo id rather than a local
@@ -27,7 +27,7 @@ def _no_hub(monkeypatch):
     def _fail(*args, **kwargs):  # pragma: no cover - must not be reached
         raise AssertionError("snapshot_download must not be called here")
 
-    monkeypatch.setattr(huggingface_hub, "snapshot_download", _fail)
+    monkeypatch.setattr(huggingface_hub.HfApi, "snapshot_download", _fail)
 
 
 def test_local_directory_is_returned_unchanged(tmp_path, monkeypatch):
@@ -38,13 +38,13 @@ def test_local_directory_is_returned_unchanged(tmp_path, monkeypatch):
 def test_repo_id_resolves_via_snapshot_download(tmp_path, monkeypatch):
     calls = {}
 
-    def _fake_snapshot_download(model_ref, revision=None, allow_patterns=None):
+    def _fake_snapshot_download(self, model_ref, revision=None, allow_patterns=None):
         calls["model_ref"] = model_ref
         calls["revision"] = revision
         calls["allow_patterns"] = allow_patterns
         return str(tmp_path / "snapshot")
 
-    monkeypatch.setattr(huggingface_hub, "snapshot_download", _fake_snapshot_download)
+    monkeypatch.setattr(huggingface_hub.HfApi, "snapshot_download", _fake_snapshot_download)
     resolved = _resolve_model_dir("openbmb/MiniCPM-o-4_5", revision="abc123")
     assert resolved == str(tmp_path / "snapshot")
     assert calls["model_ref"] == "openbmb/MiniCPM-o-4_5"
@@ -56,7 +56,7 @@ def test_snapshot_download_failure_propagates(monkeypatch):
     def _raise(*args, **kwargs):
         raise FileNotFoundError("offline and not cached")
 
-    monkeypatch.setattr(huggingface_hub, "snapshot_download", _raise)
+    monkeypatch.setattr(huggingface_hub.HfApi, "snapshot_download", _raise)
     with pytest.raises(FileNotFoundError):
         _resolve_model_dir("openbmb/MiniCPM-o-4_5")
 

tests/model_executor/models/minicpmo_4_5/test_long_audio_processing.py

@@ -10,6 +10,10 @@ Covers:
   - ``process_audios``: >30s audio unpads per chunk without ``TypeError``
   - multiple >30s audios unpad in chunk order
   - <=30s audio unpadding byte-identical to the pre-fix behavior
+  - ``MiniCPMOAudioFeatureInputs``: a batch mixing a >30s and a <=30s audio
+    (different slice counts) validates instead of raising, whether from a
+    single request with multiple audios of different lengths or from
+    multiple requests sharing an encoder batch
 """
 
 from __future__ import annotations
@@ -26,6 +30,7 @@ from vllm.multimodal.inputs import (
 
 from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_llm import (
     MiniCPMO45OmniLLMMultiModalProcessor,
+    MiniCPMOAudioFeatureInputs,
     MiniCPMOMultiModalDataParser,
     _minicpmo_field_config,
 )
@@ -152,3 +157,43 @@ class TestProcessAudiosUnpadding:
         assert len(features) == 1
         assert features[0].shape == (_FEAT_DIM, 980)
         assert torch.equal(features[0], fake_outputs["audio_features"][0][:, :980])
+
+
+class TestMixedChunkCountBatchValidation:
+    def test_batch_mixes_different_slice_counts(self) -> None:
+        # A >30s audio (3 slices) and a <=30s audio (2 slices) scheduled in
+        # the same encoder batch: audio_feature_lens is a list of per-audio
+        # tensors with different lengths (3 vs 2). This is the shape produced
+        # both by a single request with two audios of different lengths and
+        # by two separate requests batched together. Before the fix this
+        # raised ValueError: "audio_feature_lens contains inconsistent shapes".
+        audio_feature_lens = [
+            torch.tensor([_CHUNK_FRAMES, _CHUNK_FRAMES, 1]),
+            torch.tensor([_CHUNK_FRAMES, 1]),
+        ]
+        audio_features = [torch.zeros(_FEAT_DIM, _CHUNK_FRAMES) for _ in range(5)]
+
+        inputs = MiniCPMOAudioFeatureInputs(
+            type="audio_features",
+            audio_features=audio_features,
+            audio_feature_lens=audio_feature_lens,
+        )
+
+        assert inputs.audio_feature_lens is audio_feature_lens
+
+    def test_single_slice_count_batch_still_validated(self) -> None:
+        # Same slice count per audio (the common case) must still validate,
+        # unaffected by marking "s" dynamic.
+        audio_feature_lens = [
+            torch.tensor([_CHUNK_FRAMES, 1]),
+            torch.tensor([_CHUNK_FRAMES, 1]),
+        ]
+        audio_features = [torch.zeros(_FEAT_DIM, _CHUNK_FRAMES) for _ in range(4)]
+
+        inputs = MiniCPMOAudioFeatureInputs(
+            type="audio_features",
+            audio_features=audio_features,
+            audio_feature_lens=audio_feature_lens,
+        )
+
+        assert inputs.audio_feature_lens is audio_feature_lens

tests/model_executor/models/omnivoice/test_cuda_graph_generator.py

@@ -6,19 +6,27 @@ Tests for OmniVoice CUDA Graph generator wrapper numerical equivalence.
 Verifies that _OmniVoiceCUDAGraphForward produces results equivalent to
 eager mode across three scenarios:
   - Exact-size inputs (no padding) → bit-identical
-  - Padded inputs (padded to nearest bucket) → correct slicing, exact match
-    for position-independent models
-  - Oversized inputs (fallback to eager) → bit-identical
+  - Padded inputs (padded to nearest bucket) → correct slicing and exact match
+  - Oversized inputs (128-aligned lazy capture) → bit-identical
 
-Uses a lightweight synthetic generator to keep warmup fast and avoid
-loading actual model weights.
+Uses a small randomly initialized OmniVoiceGenerator so every test exercises
+the production embedding, transformer, varlen-attention, and logits paths
+without loading checkpoint weights.
 """
 
 from __future__ import annotations
 
+from types import SimpleNamespace
+
 import pytest
 import torch
-import torch.nn as nn
+from vllm.utils.math_utils import round_up
+
+from vllm_omni.model_executor.models.omnivoice.omnivoice_generator import (
+    OmniVoiceGenerator,
+    _OmniVoiceCUDAGraphForward,
+)
+from vllm_omni.transformers_utils.configs.omnivoice import OmniVoiceConfig
 
 pytestmark = [
     pytest.mark.core_model,
@@ -29,53 +37,7 @@ pytestmark = [
 DEVICE = torch.device("cuda:0")
 NUM_CB = 8
 VOCAB = 1025
-HIDDEN = 64
 CAPTURE_SIZES = [32, 64, 128]
-HEAD_DIM = 64
-
-
-# ---------------------------------------------------------------------------
-# Synthetic generator: matches _OmniVoiceCUDAGraphForward's interface
-# ---------------------------------------------------------------------------
-
-
-class _SyntheticGenerator:
-    """Position-independent Embedding+Linear replacing OmniVoiceGenerator's 28-layer transformer.
-
-    Each token position is processed independently (no attention), so padded
-    runs produce bit-identical results to non-padded runs on the same positions.
-    This makes the padding/slicing logic easy to verify exactly.
-    """
-
-    class _Cfg:
-        num_audio_codebook = NUM_CB
-
-    def __init__(self, device: torch.device):
-        self.config = self._Cfg()
-        self.text_embedding = nn.Embedding(1000, HIDDEN).to(device).eval()
-        self._linear = nn.Linear(HIDDEN, NUM_CB * VOCAB, bias=False).to(device).eval()
-        self._rope_table: torch.Tensor | None = None
-
-    @property
-    def model_dtype(self) -> torch.dtype:
-        """Part of the interface _OmniVoiceCUDAGraphForward expects of a generator."""
-        return self.text_embedding.weight.dtype
-
-    def _rope_table_for(self, seq_len: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
-        """Part of the interface _OmniVoiceCUDAGraphForward expects of a generator."""
-        return torch.zeros(seq_len, HEAD_DIM, device=device, dtype=dtype)
-
-    def _step_forward(
-        self,
-        input_ids: torch.Tensor,
-        audio_mask: torch.Tensor,
-        attention_mask: torch.Tensor | None,
-        rope_table: torch.Tensor,
-    ) -> torch.Tensor:
-        two_b, _, S = input_ids.shape
-        x = self.text_embedding(input_ids[:, 0, :].clamp(0, 999))  # [two_b, S, H]
-        logits = self._linear(x)  # [two_b, S, 8*1025]
-        return logits.view(two_b, S, NUM_CB, VOCAB).permute(0, 2, 1, 3)  # [two_b, 8, S, 1025]
 
 
 # ---------------------------------------------------------------------------
@@ -83,18 +45,37 @@ class _SyntheticGenerator:
 # ---------------------------------------------------------------------------
 
 
-def _eager(gen: _SyntheticGenerator, ids: torch.Tensor, mask: torch.Tensor, attn: torch.Tensor) -> torch.Tensor:
-    """Run _step_forward with a zero RoPE table (synthetic gen ignores it)."""
-    two_b, _, S = ids.shape
-    return gen._step_forward(ids, mask, attn, gen._rope_table_for(S, ids.device, torch.float32))
+def _eager(gen: OmniVoiceGenerator, ids: torch.Tensor, mask: torch.Tensor, cu_seqs) -> torch.Tensor:
+    """Run the production varlen step outside CUDA Graph."""
+    seq_len = ids.shape[0]
+    rope_table = gen._rope_table_for(seq_len, ids.device, gen.model_dtype)
+    return gen._step_forward(ids, mask, cu_seqs, rope_table)
+
+
+def _bucketed_eager(gen, wrapper, ids, mask, cu_seqs, batch_size) -> torch.Tensor:
+    """Run eager with the exact padded shape and metadata used by Graph."""
+    seq_len = ids.shape[0]
+    bucket = wrapper._find_bucket(batch_size, seq_len)
+    if bucket is None:
+        bucket = round_up(seq_len, wrapper._LAZY_CAPTURE_ALIGNMENT)
+    padded_ids, padded_mask = wrapper._pad_inputs(ids, mask, bucket)
+    padded_cu_seqs = cu_seqs.clone()
+    padded_cu_seqs[-1] = bucket
+    return _eager(gen, padded_ids, padded_mask, padded_cu_seqs)[:, :seq_len, :]
 
 
 def _make_inputs(seq_len: int, device: torch.device = DEVICE):
-    two_b = 2
-    ids = torch.randint(0, 100, (two_b, NUM_CB, seq_len), dtype=torch.long, device=device)
-    mask = torch.ones(two_b, seq_len, dtype=torch.bool, device=device)
-    attn = torch.ones(two_b, 1, seq_len, seq_len, dtype=torch.bool, device=device)
-    return ids, mask, attn
+    ids = torch.randint(0, 100, (seq_len, NUM_CB), dtype=torch.long, device=device)
+    mask = torch.ones(seq_len, dtype=torch.bool, device=device)
+    cond_len = (seq_len + 1) // 2
+    uncond_len = seq_len - cond_len
+
+    cu_seqs = torch.tensor(
+        [0, cond_len, cond_len + uncond_len, cond_len + uncond_len],
+        dtype=torch.int32,
+        device=device,
+    )
+    return ids, mask, cu_seqs
 
 
 # ---------------------------------------------------------------------------
@@ -105,15 +86,26 @@ def _make_inputs(seq_len: int, device: torch.device = DEVICE):
 @pytest.fixture(scope="module")
 def gen():
     torch.manual_seed(42)
-    return _SyntheticGenerator(DEVICE)
+    config = OmniVoiceConfig(
+        audio_vocab_size=VOCAB,
+        num_audio_codebook=NUM_CB,
+        llm_config={
+            "hidden_size": 64,
+            "num_hidden_layers": 2,
+            "num_attention_heads": 4,
+            "num_key_value_heads": 2,
+            "intermediate_size": 128,
+            "vocab_size": 128,
+            "max_position_embeddings": 4096,
+            "head_dim": 16,
+        },
+        enable_cuda_graph=False,
+    )
+    return OmniVoiceGenerator(config, SimpleNamespace(max_num_seqs=2)).to(DEVICE).eval()
 
 
 @pytest.fixture(scope="module")
 def wrapper(gen):
-    from vllm_omni.model_executor.models.omnivoice.omnivoice_generator import (
-        _OmniVoiceCUDAGraphForward,
-    )
-
     w = _OmniVoiceCUDAGraphForward(gen, capture_sizes=CAPTURE_SIZES)
     w.warmup(DEVICE)
     return w
@@ -127,10 +119,10 @@ def wrapper(gen):
 @pytest.mark.parametrize("seq_len", CAPTURE_SIZES)
 def test_exact_size_bit_identical(gen, wrapper, seq_len):
     """When input exactly matches a captured bucket, output must be bit-identical to eager."""
-    ids, mask, attn = _make_inputs(seq_len)
+    ids, mask, cu_seqs = _make_inputs(seq_len)
     with torch.no_grad():
-        eager_out = _eager(gen, ids, mask, attn)
-        graph_out = wrapper(ids, mask, attn)
+        eager_out = _eager(gen, ids, mask, cu_seqs)
+        graph_out = wrapper(ids, mask, cu_seqs, 1)
     torch.testing.assert_close(graph_out, eager_out, atol=0, rtol=0)
 
 
@@ -140,40 +132,36 @@ def test_exact_size_bit_identical(gen, wrapper, seq_len):
 
 
 @pytest.mark.parametrize("seq_len", [1, 15, 33, 60, 100])
-def test_padded_output_shape(gen, wrapper, seq_len):
+def test_padded_output_shape(wrapper, seq_len):
     """Graph output must be sliced back to actual seq_len, not the bucket size."""
-    ids, mask, attn = _make_inputs(seq_len)
+    ids, mask, cu_seqs = _make_inputs(seq_len)
     with torch.no_grad():
-        graph_out = wrapper(ids, mask, attn)
-    assert graph_out.shape == (2, NUM_CB, seq_len, VOCAB)
+        graph_out = wrapper(ids, mask, cu_seqs, 1)
+    assert graph_out.shape == (NUM_CB, seq_len, VOCAB)
 
 
 @pytest.mark.parametrize("seq_len", [15, 33, 60, 100])
 def test_padded_output_matches_eager(gen, wrapper, seq_len):
-    """Padded graph output must equal eager output at actual positions.
-
-    The synthetic model has no attention across positions, so zero-padding
-    does not affect non-padded positions — exact match is expected.
-    """
-    ids, mask, attn = _make_inputs(seq_len)
+    """Padded graph output must equal eager output at actual positions."""
+    ids, mask, cu_seqs = _make_inputs(seq_len)
     with torch.no_grad():
-        eager_out = _eager(gen, ids, mask, attn)
-        graph_out = wrapper(ids, mask, attn)
+        eager_out = _bucketed_eager(gen, wrapper, ids, mask, cu_seqs, 1)
+        graph_out = wrapper(ids, mask, cu_seqs, 1)
     torch.testing.assert_close(graph_out, eager_out, atol=0, rtol=0)
 
 
 # ---------------------------------------------------------------------------
-# 3. Oversized inputs → fallback to eager (lazy capture), bit-identical
+# 3. Oversized inputs → aligned lazy capture, bit-identical
 # ---------------------------------------------------------------------------
 
 
 @pytest.mark.parametrize("seq_len", [129, 200, 256])
-def test_fallback_eager_bit_identical(gen, wrapper, seq_len):
-    """Sequences exceeding the largest bucket fall back to lazy capture → bit-identical."""
-    ids, mask, attn = _make_inputs(seq_len)
+def test_aligned_lazy_capture_bit_identical(gen, wrapper, seq_len):
+    """Sequences beyond the static plan use aligned lazy capture and remain bit-identical."""
+    ids, mask, cu_seqs = _make_inputs(seq_len)
     with torch.no_grad():
-        eager_out = _eager(gen, ids, mask, attn)
-        graph_out = wrapper(ids, mask, attn)
+        eager_out = _bucketed_eager(gen, wrapper, ids, mask, cu_seqs, 1)
+        graph_out = wrapper(ids, mask, cu_seqs, 1)
     torch.testing.assert_close(graph_out, eager_out, atol=0, rtol=0)
 
 
@@ -184,13 +172,73 @@ def test_fallback_eager_bit_identical(gen, wrapper, seq_len):
 
 def test_deterministic_across_calls(wrapper):
     """Same input must produce identical output on repeated CUDA graph replays."""
-    ids, mask, attn = _make_inputs(32)
+    ids, mask, cu_seqs = _make_inputs(32)
     with torch.no_grad():
-        out1 = wrapper(ids, mask, attn)
-        out2 = wrapper(ids, mask, attn)
+        out1 = wrapper(ids, mask, cu_seqs, 1).clone()
+        out2 = wrapper(ids, mask, cu_seqs, 1).clone()
     torch.testing.assert_close(out1, out2, atol=0, rtol=0)
 
 
+def test_replay_uses_updated_cu_seqs(gen, wrapper):
+    """The same graph key must honor new sequence boundaries on replay."""
+    ids, mask, _ = _make_inputs(48)
+    cu_seqs_a = torch.tensor([0, 32, 48, 48], dtype=torch.int32, device=DEVICE)
+    cu_seqs_b = torch.tensor([0, 17, 48, 48], dtype=torch.int32, device=DEVICE)
+
+    with torch.no_grad():
+        eager_a = _bucketed_eager(gen, wrapper, ids, mask, cu_seqs_a, 1)
+        graph_a = wrapper(ids, mask, cu_seqs_a, 1).clone()
+        eager_b = _bucketed_eager(gen, wrapper, ids, mask, cu_seqs_b, 1)
+        graph_b = wrapper(ids, mask, cu_seqs_b, 1).clone()
+
+    torch.testing.assert_close(graph_a, eager_a, atol=0, rtol=0)
+    torch.testing.assert_close(graph_b, eager_b, atol=0, rtol=0)
+    assert not torch.equal(graph_a, graph_b)
+
+
+def test_batch_two_cu_seqs_matches_eager(gen, wrapper):
+    """B=2 uses four real sequence segments plus the fixed tail slot."""
+    ids, mask, _ = _make_inputs(48)
+    cu_seqs = torch.tensor([0, 10, 18, 32, 48, 48], dtype=torch.int32, device=DEVICE)
+
+    assert cu_seqs.numel() == 2 * 2 + 2
+    with torch.no_grad():
+        eager_out = _bucketed_eager(gen, wrapper, ids, mask, cu_seqs, 2)
+        graph_out = wrapper(ids, mask, cu_seqs, 2).clone()
+
+    torch.testing.assert_close(graph_out, eager_out, atol=0, rtol=0)
+
+
+def test_lazy_graph_cache_uses_lru_eviction(wrapper):
+    """A lazy-cache hit must protect that graph from the next eviction."""
+    original_limit = wrapper._MAX_LAZY_GRAPHS
+    wrapper._lazy_graphs.clear()
+    wrapper._MAX_LAZY_GRAPHS = 2
+    try:
+        # Static B=1 coverage ends at 128. These lengths round to three
+        # distinct 128-aligned lazy keys: 256, 384, and 512.
+        for seq_len in (129, 257):
+            ids, mask, cu_seqs = _make_inputs(seq_len)
+            with torch.no_grad():
+                wrapper(ids, mask, cu_seqs, 1)
+        assert list(wrapper._lazy_graphs) == [(1, 256), (1, 384)]
+
+        # Refresh key 256, making key 384 the least recently used entry.
+        ids, mask, cu_seqs = _make_inputs(129)
+        with torch.no_grad():
+            wrapper(ids, mask, cu_seqs, 1)
+        assert list(wrapper._lazy_graphs) == [(1, 384), (1, 256)]
+
+        # Inserting key 512 must evict key 384, not the recently hit key 256.
+        ids, mask, cu_seqs = _make_inputs(385)
+        with torch.no_grad():
+            wrapper(ids, mask, cu_seqs, 1)
+        assert list(wrapper._lazy_graphs) == [(1, 256), (1, 512)]
+    finally:
+        wrapper._MAX_LAZY_GRAPHS = original_limit
+        wrapper._lazy_graphs.clear()
+
+
 # ---------------------------------------------------------------------------
 # 5. _find_bucket logic (CPU, no CUDA graph)
 # ---------------------------------------------------------------------------
@@ -203,58 +251,17 @@ def test_find_bucket_returns_nearest_bucket():
     )
 
     w = _OmniVoiceCUDAGraphForward.__new__(_OmniVoiceCUDAGraphForward)
-    w._capture_sizes = [32, 64, 128]
+    w.capture_bucket_sizes_by_batch = {
+        1: [32, 64, 128],
+        2: [64, 128, 256],
+    }
     w._graphs = {}
 
-    assert w._find_bucket(1) == 32
-    assert w._find_bucket(32) == 32
-    assert w._find_bucket(33) == 64
-    assert w._find_bucket(64) == 64
-    assert w._find_bucket(128) == 128
-    assert w._find_bucket(129) is None
-    assert w._find_bucket(1000) is None
-
-
-# ---------------------------------------------------------------------------
-# 6. enable_cuda_graph=False produces same tokens as enable_cuda_graph=True
-# ---------------------------------------------------------------------------
-
-
-def test_cuda_graph_disabled_matches_eager_generator():
-    """OmniVoiceGenerator with enable_cuda_graph=False must produce the same
-    _step_forward output as one with enable_cuda_graph=True (after warmup).
-
-    Uses a 2-layer config so warmup completes quickly without loading weights.
-    """
-    from vllm_omni.model_executor.models.omnivoice.omnivoice_generator import OmniVoiceGenerator
-    from vllm_omni.transformers_utils.configs.omnivoice import OmniVoiceConfig
-
-    cfg = OmniVoiceConfig()
-    cfg.llm_num_hidden_layers = 2
-    cfg.num_hidden_layers = 2
-
-    torch.manual_seed(0)
-    gen_eager = OmniVoiceGenerator(cfg).to(DEVICE).eval()
-    gen_eager.config.enable_cuda_graph = False
-    gen_eager._cuda_graph_fwd = None
-
-    torch.manual_seed(0)
-    gen_graph = OmniVoiceGenerator(cfg).to(DEVICE).eval()
-    gen_graph.config.cuda_graph_capture_sizes = [64]
-    from vllm_omni.model_executor.models.omnivoice.omnivoice_generator import _OmniVoiceCUDAGraphForward
-
-    gen_graph._cuda_graph_fwd = _OmniVoiceCUDAGraphForward(gen_graph, capture_sizes=[64])
-    gen_graph._cuda_graph_fwd.warmup(DEVICE)
-
-    seq_len = 64
-    ids = torch.zeros(2, cfg.num_audio_codebook, seq_len, dtype=torch.long, device=DEVICE)
-    mask = torch.ones(2, seq_len, dtype=torch.bool, device=DEVICE)
-    attn = torch.ones(2, 1, seq_len, seq_len, dtype=torch.bool, device=DEVICE)
-
-    rope_table = gen_eager._rope_table_for(seq_len, DEVICE, gen_eager.model_dtype)
-
-    with torch.no_grad():
-        eager_logits = gen_eager._step_forward(ids, mask, attn, rope_table)
-        graph_logits = gen_graph._cuda_graph_fwd(ids, mask, attn)
-
-    torch.testing.assert_close(graph_logits, eager_logits, atol=0, rtol=0)
+    assert w._find_bucket(1, 1) == 32
+    assert w._find_bucket(1, 32) == 32
+    assert w._find_bucket(1, 33) == 64
+    assert w._find_bucket(1, 128) == 128
+    assert w._find_bucket(1, 129) is None
+    assert w._find_bucket(2, 33) == 64
+    assert w._find_bucket(2, 129) == 256
+    assert w._find_bucket(3, 64) is None

tests/model_executor/models/omnivoice/test_fused_projection_load.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
 """The OmniVoice generator packs q/k/v and gate/up into fused projections.
 
 The HF checkpoint stores those five tensors separately, so packing them is a
@@ -14,6 +14,8 @@ that would leave one partly written raises instead.
 
 from __future__ import annotations
 
+from types import SimpleNamespace
+
 import pytest
 import torch
 
@@ -49,6 +51,13 @@ def _config() -> OmniVoiceConfig:
     )
 
 
+def _generator() -> OmniVoiceGenerator:
+    return OmniVoiceGenerator(
+        _config(),
+        SimpleNamespace(max_num_seqs=1),
+    )
+
+
 def _checkpoint_shards() -> dict[str, torch.Tensor]:
     """The five per-layer tensors an OmniVoice checkpoint actually stores."""
     torch.manual_seed(0)
@@ -64,7 +73,7 @@ def _checkpoint_shards() -> dict[str, torch.Tensor]:
 
 
 def test_every_fused_parameter_is_written() -> None:
-    generator = OmniVoiceGenerator(_config())
+    generator = _generator()
     shards = _checkpoint_shards()
 
     loaded = generator._load_fused_projections(shards)
@@ -77,7 +86,7 @@ def test_every_fused_parameter_is_written() -> None:
 
 def test_packed_weight_matches_the_shards_it_came_from() -> None:
     """Packing order is q,k,v and gate,up -- the split in forward() assumes it."""
-    generator = OmniVoiceGenerator(_config())
+    generator = _generator()
     shards = _checkpoint_shards()
     generator._load_fused_projections(shards)
 
@@ -101,7 +110,7 @@ def test_packed_weight_matches_the_shards_it_came_from() -> None:
 
 def test_no_fused_parameter_keeps_its_random_init() -> None:
     """The corruption this guards against: a fused param never written at all."""
-    generator = OmniVoiceGenerator(_config())
+    generator = _generator()
     before = {
         name: param.detach().clone()
         for name, param in generator.named_parameters()
@@ -121,7 +130,7 @@ def test_no_fused_parameter_keeps_its_random_init() -> None:
     ["self_attn.k_proj", "self_attn.v_proj", "self_attn.q_proj", "mlp.gate_proj", "mlp.up_proj"],
 )
 def test_a_missing_shard_raises_instead_of_loading_partially(dropped: str) -> None:
-    generator = OmniVoiceGenerator(_config())
+    generator = _generator()
     shards = _checkpoint_shards()
     del shards[f"llm.layers.0.{dropped}.weight"]
 
@@ -130,7 +139,7 @@ def test_a_missing_shard_raises_instead_of_loading_partially(dropped: str) -> No
 
 
 def test_a_wrong_shaped_shard_raises() -> None:
-    generator = OmniVoiceGenerator(_config())
+    generator = _generator()
     shards = _checkpoint_shards()
     shards["llm.layers.0.self_attn.q_proj.weight"] = torch.randn(NUM_HEADS * HEAD_DIM + 1, HIDDEN)
 
@@ -140,7 +149,7 @@ def test_a_wrong_shaped_shard_raises() -> None:
 
 def test_a_checkpoint_missing_a_whole_layer_raises() -> None:
     """Half-loaded is the dangerous state: it neither errors nor works."""
-    generator = OmniVoiceGenerator(_config())
+    generator = _generator()
     shards = {k: v for k, v in _checkpoint_shards().items() if not k.startswith("llm.layers.1.")}
 
     with pytest.raises(ValueError, match="fused projections"):
@@ -149,7 +158,7 @@ def test_a_checkpoint_missing_a_whole_layer_raises() -> None:
 
 def test_a_checkpoint_with_no_fused_shards_is_left_alone() -> None:
     """An unrelated state_dict must not trip the completeness check."""
-    generator = OmniVoiceGenerator(_config())
+    generator = _generator()
     assert generator._load_fused_projections({"audio_heads.weight": torch.randn(4, 4)}) == set()
 
 
@@ -162,7 +171,7 @@ def test_load_weights_wires_the_packing_in(tmp_path, monkeypatch: pytest.MonkeyP
     """
     import safetensors.torch
 
-    generator = OmniVoiceGenerator(_config())
+    generator = _generator()
     shards = _checkpoint_shards()
     before = generator.layers[0].self_attn.qkv_proj.weight.detach().clone()
 

tests/model_executor/models/omnivoice/test_mask_dtype.py

@@ -1,183 +1,79 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""OmniVoice attention masks must follow the model dtype, not a hardcoded float32.
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""OmniVoice SDPA fallback masks must not depend on the model dtype.
 
-SDPA requires the additive attention mask to carry the same dtype as the query.
-The generator used to materialize that mask as float32 unconditionally, in three
-places (the shared per-forward mask, the CUDA-graph capture buffer, and the
-replay-time normalization), so every half-precision deployment died on the first
-attention call with::
+The previous padded attention path used an additive floating-point mask. SDPA
+requires such a mask to have the same dtype as the query, and constructing it
+as float32 caused half-precision OmniVoice inference to fail on CUDA.
 
-    RuntimeError: invalid dtype for bias - should match query's dtype
-
-(the exact wording depends on which SDPA backend is selected for the shape; the
-math backend words the same rejection as ``attn_mask.dtype``.)
-
-That left OmniVoice servable only in float32, even though the upstream k2-fsa
-implementation runs it in float16.
-
-Note the split between the CPU and CUDA tests below: the CPU SDPA backend
-silently accepts the mismatched mask, so only the CUDA tests can pin the actual
-crash. The CPU tests cover the mask contract itself, which is what the fix
-changes and what a future regression would break first.
+The packed-varlen path no longer needs an additive mask. When SDPA fallback is
+required, ``_attention_metadata_from_cu_seqs`` builds a boolean block mask.
+Boolean SDPA masks have no dtype coupling with the query, so no mask coercion
+is required for float16 or bfloat16 inference.
 """
 
 from __future__ import annotations
 
 import pytest
 import torch
+import torch.nn.functional as F
 
 from vllm_omni.model_executor.models.omnivoice.omnivoice_generator import (
-    OmniVoiceAttention,
-    OmniVoiceGenerator,
-    _additive_float_mask,
+    _attention_metadata_from_cu_seqs,
 )
-from vllm_omni.transformers_utils.configs.omnivoice import OmniVoiceConfig
 
 HALF_DTYPES = [torch.float16, torch.bfloat16]
-ALL_DTYPES = HALF_DTYPES + [torch.float32]
 
 cpu_test = pytest.mark.core_model
 cuda_test = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
 
 
-def _tiny_config() -> OmniVoiceConfig:
-    """A real OmniVoiceConfig sized so the module fits in a unit test."""
-    return OmniVoiceConfig(
-        llm_config={
-            "hidden_size": 32,
-            "num_hidden_layers": 2,
-            "num_attention_heads": 4,
-            "num_key_value_heads": 2,
-            "head_dim": 8,
-            "intermediate_size": 64,
-            "vocab_size": 64,
-            "max_position_embeddings": 128,
-        },
-        enable_cuda_graph=False,
+@cpu_test
+@pytest.mark.cpu
+def test_sdpa_fallback_mask_is_dtype_independent() -> None:
+    """The fallback mask is boolean rather than an additive model-dtype mask."""
+    cu_seqs = torch.tensor([0, 2, 3, 6, 6], dtype=torch.int32)
+
+    metadata = _attention_metadata_from_cu_seqs(
+        cu_seqs,
+        6,
+        needs_sdpa_mask=True,
     )
 
-
-def _inputs(dtype: torch.dtype, device: torch.device) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-    hidden = torch.randn(2, 6, 32, device=device, dtype=dtype)
-    bool_mask = torch.ones(2, 1, 6, 6, dtype=torch.bool, device=device)
-    bool_mask[:, :, :, 4:] = False
-    rope_table = torch.zeros(6, 8, device=device, dtype=dtype)
-    rope_table[:, :4] = 1.0  # cos = 1, sin = 0: no rotation, so the test is about the mask
-    return hidden, bool_mask, rope_table
-
-
-# --------------------------------------------------------------------------
-# The mask contract (CPU)
-# --------------------------------------------------------------------------
-
-
-@cpu_test
-@pytest.mark.cpu
-@pytest.mark.parametrize("dtype", ALL_DTYPES)
-def test_additive_float_mask_uses_requested_dtype(dtype: torch.dtype) -> None:
-    out = _additive_float_mask(torch.tensor([[True, False]]), dtype)
-    assert out.dtype == dtype
-
-
-@cpu_test
-@pytest.mark.cpu
-def test_additive_float_mask_keeps_mask_semantics() -> None:
-    """Attend positions stay at 0.0 and masked ones at -inf, in half precision too."""
-    out = _additive_float_mask(torch.tensor([[True, False]]), torch.float16)
-    assert out[0, 0].item() == 0.0
-    assert out[0, 1].item() == float("-inf")
-
-
-@cpu_test
-@pytest.mark.cpu
-def test_additive_float_mask_requires_an_explicit_dtype() -> None:
-    """The float32 default is the bug; keeping the parameter required is the fix."""
-    with pytest.raises(TypeError):
-        _additive_float_mask(torch.tensor([[True]]))  # type: ignore[call-arg]
-
-
-@cpu_test
-@pytest.mark.cpu
-@pytest.mark.parametrize("dtype", ALL_DTYPES)
-def test_generator_model_dtype_tracks_its_weights(dtype: torch.dtype) -> None:
-    """The single source of truth the three former float32 literals now defer to."""
-    assert OmniVoiceGenerator(_tiny_config()).to(dtype).model_dtype == dtype
-
-
-@cpu_test
-@pytest.mark.cpu
-@pytest.mark.parametrize("dtype", HALF_DTYPES)
-def test_attention_accepts_a_model_dtype_mask(dtype: torch.dtype) -> None:
-    attn = OmniVoiceAttention(_tiny_config()).to(dtype).eval()
-    hidden, bool_mask, rope_table = _inputs(dtype, torch.device("cpu"))
-
-    with torch.inference_mode():
-        out = attn(hidden, rope_table, attention_mask=_additive_float_mask(bool_mask, dtype))
-
-    assert out.shape == hidden.shape
-    assert out.dtype == dtype
-    assert torch.isfinite(out).all()
-
-
-# --------------------------------------------------------------------------
-# The crash itself (CUDA only — the CPU backend does not reject the mismatch)
-# --------------------------------------------------------------------------
+    assert metadata.attn_mask is not None
+    assert metadata.attn_mask.dtype == torch.bool
 
 
 @cuda_test
 @pytest.mark.parametrize("dtype", HALF_DTYPES)
-def test_float32_mask_is_what_broke_half_precision(dtype: torch.dtype) -> None:
-    """Pin the original failure so a float32 default cannot come back unnoticed."""
-    attn = OmniVoiceAttention(_tiny_config()).to(device="cuda:0", dtype=dtype).eval()
-    hidden, bool_mask, rope_table = _inputs(dtype, torch.device("cuda:0"))
-
-    with pytest.raises(RuntimeError, match=r"(invalid dtype for bias|attn_mask)"), torch.inference_mode():
-        attn(hidden, rope_table, attention_mask=_additive_float_mask(bool_mask, torch.float32))
-
-
-@cuda_test
-@pytest.mark.parametrize("dtype", HALF_DTYPES)
-def test_real_path_forward_runs_in_half_precision(dtype: torch.dtype) -> None:
-    """The regression, on the production path with the Triton norm kernels live."""
-    attn = OmniVoiceAttention(_tiny_config()).to(device="cuda:0", dtype=dtype).eval()
-    hidden, bool_mask, rope_table = _inputs(dtype, torch.device("cuda:0"))
-
-    with torch.inference_mode():
-        out = attn(hidden, rope_table, attention_mask=_additive_float_mask(bool_mask, dtype))
-
-    assert out.dtype == dtype
-    assert torch.isfinite(out).all()
-
-
-@cuda_test
-@pytest.mark.parametrize("dtype", HALF_DTYPES)
-def test_generator_forward_runs_in_half_precision(dtype: torch.dtype) -> None:
-    """End-to-end over the real iterative loop, which is where the float32 mask was built.
-
-    This is the test that fails on an unfixed tree: ``forward`` materialized the
-    shared SDPA mask as float32 regardless of the weights' dtype.
-    """
+def test_sdpa_fallback_mask_needs_no_half_precision_coercion(
+    dtype: torch.dtype,
+) -> None:
+    """Half-precision SDPA accepts the production bool mask without casting."""
     device = torch.device("cuda:0")
-    config = _tiny_config()
-    generator = OmniVoiceGenerator(config).to(device=device, dtype=dtype).eval()
+    seq_len = 6
+    cu_seqs = torch.tensor([0, 2, 3, 6, 6], dtype=torch.int32, device=device)
 
-    seq_len, target_len = 12, 4
-    input_ids = torch.zeros(2, config.num_audio_codebook, seq_len, dtype=torch.long, device=device)
-    input_ids[:, 1:, :] = config.audio_mask_id
-    audio_mask = torch.zeros(2, seq_len, dtype=torch.bool, device=device)
-    audio_mask[:, seq_len - target_len :] = True
-    attention_mask = torch.ones(2, 1, seq_len, seq_len, dtype=torch.bool, device=device)
+    metadata = _attention_metadata_from_cu_seqs(
+        cu_seqs,
+        seq_len,
+        needs_sdpa_mask=True,
+    )
+
+    assert metadata.attn_mask is not None
+    assert metadata.attn_mask.dtype == torch.bool
+
+    query = torch.randn(1, 2, seq_len, 8, device=device, dtype=dtype)
+    key = torch.randn_like(query)
+    value = torch.randn_like(query)
 
     with torch.inference_mode():
-        tokens = generator(
-            input_ids,
-            audio_mask,
-            attention_mask,
-            target_lens=[target_len],
-            seed=0,
-            num_step=2,
+        output = F.scaled_dot_product_attention(
+            query,
+            key,
+            value,
+            attn_mask=metadata.attn_mask,
         )
 
-    assert tokens.shape == (1, config.num_audio_codebook, target_len)
-    assert tokens.dtype == torch.long
+    assert output.dtype == dtype
+    assert torch.isfinite(output).all()

tests/model_executor/models/omnivoice/test_nonverbal_tags_tokenize.py

@@ -1,12 +1,15 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 from __future__ import annotations
 
 import pytest
-from huggingface_hub import hf_hub_download
 from tokenizers import Tokenizer as HFTokenizer
 
 from vllm_omni.diffusion.models.omnivoice.pipeline_omnivoice import (
     _tokenize_with_nonverbal_tags,
 )
+from vllm_omni.transformers_utils.repo_utils import hf_api
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -18,7 +21,7 @@ def _ids(tok: HFTokenizer, s: str) -> list[int]:
 class TestNonVerbalTags:
     @classmethod
     def setup_class(cls):
-        tokenizer_path = hf_hub_download(repo_id="k2-fsa/OmniVoice", filename="tokenizer.json")
+        tokenizer_path = hf_api().hf_hub_download(repo_id="k2-fsa/OmniVoice", filename="tokenizer.json")
         cls.tokenizer = HFTokenizer.from_file(tokenizer_path)
 
     def test_plain_text_fallback(self):

tests/model_executor/models/omnivoice/test_pipeline_batching.py

@@ -0,0 +1,228 @@
+# 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.diffusion.data import DiffusionOutput
+from vllm_omni.diffusion.models.omnivoice.pipeline_omnivoice import (
+    OmniVoicePipeline,
+    _PreparedOmniVoiceRequest,
+)
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.model_executor.models.omnivoice.omnivoice_generator import (
+    _attention_metadata_from_cu_seqs,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_sdpa_fallback_mask_preserves_packed_sequence_boundaries() -> None:
+    cu_seqs = torch.tensor([0, 2, 3, 6, 6], dtype=torch.int32)
+
+    metadata = _attention_metadata_from_cu_seqs(cu_seqs, 6, needs_sdpa_mask=True)
+
+    assert metadata.attn_mask is not None
+    expected = torch.tensor(
+        [
+            [1, 1, 0, 0, 0, 0],
+            [1, 1, 0, 0, 0, 0],
+            [0, 0, 1, 0, 0, 0],
+            [0, 0, 0, 1, 1, 1],
+            [0, 0, 0, 1, 1, 1],
+            [0, 0, 0, 1, 1, 1],
+        ],
+        dtype=torch.bool,
+    )
+    torch.testing.assert_close(metadata.attn_mask[0, 0], expected)
+
+
+def test_eager_attention_metadata_honors_exact_max_seqlen() -> None:
+    cu_seqs = torch.tensor([0, 2, 3, 6, 6], dtype=torch.int32)
+
+    metadata = _attention_metadata_from_cu_seqs(
+        cu_seqs,
+        6,
+        needs_sdpa_mask=False,
+        max_seqlen=3,
+    )
+
+    assert metadata.extra["max_seqlen_q"] == 3
+    assert metadata.extra["max_seqlen_k"] == 3
+
+
+def test_request_batch_prepare_error_does_not_skip_valid_request() -> None:
+    """A malformed request must not prevent another request from completing."""
+    prepared = _PreparedOmniVoiceRequest(
+        input_ids=torch.zeros(5, 8, dtype=torch.long),
+        audio_mask=torch.ones(5, dtype=torch.bool),
+        cond_len=3,
+        target_len=2,
+        seed=42,
+    )
+
+    def prepare_request(prompt, extra):
+        del extra
+        if prompt == "bad":
+            return DiffusionOutput(error="invalid prompt")
+        return prepared
+
+    def collate_requests(requests):
+        assert requests == [prepared]
+        return prepared.input_ids, prepared.audio_mask, [prepared.cond_len]
+
+    generator_calls = []
+
+    def generate_tokens(**kwargs):
+        generator_calls.append(kwargs)
+        assert kwargs["target_lens"] == [prepared.target_len]
+        return torch.ones(1, 8, prepared.target_len, dtype=torch.long)
+
+    pipeline = SimpleNamespace(
+        _prepare_request_input=prepare_request,
+        _collate_request_inputs=collate_requests,
+        generator=generate_tokens,
+        decoder=lambda tokens: tokens.float(),
+        num_step=32,
+        guidance_scale=7.0,
+        t_shift=1.0,
+        layer_penalty_factor=0.0,
+        position_temperature=0.0,
+        class_temperature=0.0,
+    )
+    sampling = OmniDiffusionSamplingParams(num_inference_steps=32, guidance_scale=7.0)
+    batch = DiffusionRequestBatch(
+        requests=[
+            OmniDiffusionRequest(prompt="bad", sampling_params=sampling, request_id="bad"),
+            OmniDiffusionRequest(
+                prompt="good",
+                sampling_params=OmniDiffusionSamplingParams(num_inference_steps=32, guidance_scale=7.0),
+                request_id="good",
+            ),
+        ]
+    )
+
+    outputs = OmniVoicePipeline.forward(pipeline, batch)
+
+    assert len(outputs) == 2
+    assert outputs[0].error == "invalid prompt"
+    assert outputs[1].error is None
+    torch.testing.assert_close(outputs[1].output, torch.ones(1, 8, 2))
+    assert generator_calls[0]["guidance_scale"] == 7.0
+
+
+def test_request_batch_honors_explicit_sampling_overrides() -> None:
+    """Request steps and guidance must override the OmniVoice defaults."""
+    prepared = _PreparedOmniVoiceRequest(
+        input_ids=torch.zeros(5, 8, dtype=torch.long),
+        audio_mask=torch.ones(5, dtype=torch.bool),
+        cond_len=3,
+        target_len=2,
+        seed=42,
+    )
+    captured_sampling = []
+
+    def generate_tokens(**kwargs):
+        captured_sampling.append((kwargs["num_step"], kwargs["guidance_scale"]))
+        return torch.ones(1, 8, prepared.target_len, dtype=torch.long)
+
+    pipeline = SimpleNamespace(
+        _prepare_request_input=lambda prompt, extra: prepared,
+        _collate_request_inputs=lambda requests: (prepared.input_ids, prepared.audio_mask, [prepared.cond_len]),
+        generator=generate_tokens,
+        decoder=lambda tokens: tokens.float(),
+        num_step=32,
+        guidance_scale=2.0,
+        t_shift=1.0,
+        layer_penalty_factor=0.0,
+        position_temperature=0.0,
+        class_temperature=0.0,
+    )
+    batch = DiffusionRequestBatch(
+        requests=[
+            OmniDiffusionRequest(
+                prompt="hello",
+                sampling_params=OmniDiffusionSamplingParams(num_inference_steps=12, guidance_scale=6.5),
+                request_id="request",
+            )
+        ]
+    )
+
+    outputs = OmniVoicePipeline.forward(pipeline, batch)
+
+    assert outputs[0].error is None
+    assert captured_sampling == [(12, 6.5)]
+
+
+def test_request_batch_prepare_error_preserves_surrounding_output_indices() -> None:
+    """A middle prepare error must not shift valid outputs into the wrong slots."""
+    prepared_before = _PreparedOmniVoiceRequest(
+        input_ids=torch.zeros(3, 8, dtype=torch.long),
+        audio_mask=torch.ones(3, dtype=torch.bool),
+        cond_len=2,
+        target_len=1,
+        seed=1,
+    )
+    prepared_after = _PreparedOmniVoiceRequest(
+        input_ids=torch.zeros(5, 8, dtype=torch.long),
+        audio_mask=torch.ones(5, dtype=torch.bool),
+        cond_len=3,
+        target_len=2,
+        seed=2,
+    )
+
+    def prepare_request(prompt, extra):
+        del extra
+        if prompt == "before":
+            return prepared_before
+        if prompt == "after":
+            return prepared_after
+        return DiffusionOutput(error="invalid middle request")
+
+    def collate_requests(requests):
+        assert requests == [prepared_before, prepared_after]
+        return (
+            torch.cat([prepared_before.input_ids, prepared_after.input_ids]),
+            torch.cat([prepared_before.audio_mask, prepared_after.audio_mask]),
+            [prepared_before.cond_len, prepared_after.cond_len],
+        )
+
+    def generate_tokens(**kwargs):
+        assert kwargs["target_lens"] == [1, 2]
+        before = torch.full((1, 8, 1), 11, dtype=torch.long)
+        after = torch.full((1, 8, 2), 22, dtype=torch.long)
+        return torch.cat([before, after], dim=-1)
+
+    pipeline = SimpleNamespace(
+        _prepare_request_input=prepare_request,
+        _collate_request_inputs=collate_requests,
+        generator=generate_tokens,
+        decoder=lambda tokens: tokens.float(),
+        num_step=32,
+        guidance_scale=2.0,
+        t_shift=1.0,
+        layer_penalty_factor=0.0,
+        position_temperature=0.0,
+        class_temperature=0.0,
+    )
+    batch = DiffusionRequestBatch(
+        requests=[
+            OmniDiffusionRequest(
+                prompt=prompt,
+                sampling_params=OmniDiffusionSamplingParams(num_inference_steps=32),
+                request_id=prompt,
+            )
+            for prompt in ("before", "bad", "after")
+        ]
+    )
+
+    outputs = OmniVoicePipeline.forward(pipeline, batch)
+
+    assert len(outputs) == 3
+    torch.testing.assert_close(outputs[0].output, torch.full((1, 8, 1), 11.0))
+    assert outputs[1].error == "invalid middle request"
+    torch.testing.assert_close(outputs[2].output, torch.full((1, 8, 2), 22.0))

tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.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 code predictor dtype alignment (fix for #2385).
 
@@ -20,6 +20,9 @@ import pytest
 import torch
 import torch.nn.functional as F
 from pytest_mock import MockerFixture
+from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import (
+    Qwen3OmniMoeTalkerCodePredictorConfig,
+)
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -41,6 +44,7 @@ _COMMON = os.path.join(_MODELS, "common")
 def _load_module(name: str, filename: str):
     path = os.path.abspath(os.path.join(_BASE, filename))
     spec = importlib.util.spec_from_file_location(name, path)
+    assert spec is not None and spec.loader is not None
     mod = importlib.util.module_from_spec(spec)
     sys.modules[name] = mod  # register before exec (needed for dataclasses etc.)
     spec.loader.exec_module(mod)
@@ -83,7 +87,7 @@ def _build_mock_modules(mocker: MockerFixture) -> dict[str, object]:
     vllm_parallel_mock = mocker.MagicMock()
     vllm_parallel_mock.VocabParallelEmbedding = torch.nn.Embedding
     custom_op_mock = types.ModuleType("vllm_omni.diffusion.layers.custom_op")
-    custom_op_mock.CustomOp = NativeCustomOp
+    setattr(custom_op_mock, "CustomOp", NativeCustomOp)
 
     return {
         "vllm_omni": mocker.MagicMock(),
@@ -119,6 +123,7 @@ def _load_target_classes(mocker: MockerFixture):
     common_spec = importlib.util.spec_from_file_location(
         "vllm_omni.model_executor.models.common.qwen3_code_predictor", common_cp_path
     )
+    assert common_spec is not None and common_spec.loader is not None
     common_cp_mod = importlib.util.module_from_spec(common_spec)
     sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"] = common_cp_mod
     common_spec.loader.exec_module(common_cp_mod)
@@ -209,6 +214,80 @@ def test_npu_custom_ops_use_fused_norm_and_cached_rope(mocker: MockerFixture, lo
     torch.testing.assert_close(sin, rotary.sin_cached[position_ids].to(torch.float16))
 
 
+@pytest.mark.parametrize(
+    ("rope_kwargs", "expected_theta"),
+    [
+        pytest.param(
+            {"rope_theta": 1_000_000.0},
+            1_000_000.0,
+            id="serialized_checkpoint",
+        ),
+        pytest.param(
+            {
+                "rope_theta": 10_000.0,
+                "rope_parameters": {"rope_type": "default", "rope_theta": 1_000_000.0},
+            },
+            1_000_000.0,
+            id="nested_precedence",
+        ),
+        pytest.param({}, 10_000.0, id="default"),
+    ],
+)
+def test_code_predictor_rotary_uses_qwen3_omni_rope_parameters(
+    loaded_target_classes,
+    rope_kwargs,
+    expected_theta,
+) -> None:
+    """Follow Transformers 5 deserialization and precedence for Qwen3-Omni."""
+    _ = loaded_target_classes
+    config = Qwen3OmniMoeTalkerCodePredictorConfig(
+        hidden_size=16,
+        num_attention_heads=1,
+        head_dim=16,
+        num_code_groups=16,
+        **rope_kwargs,
+    )
+    assert not hasattr(config, "rope_theta")
+    assert config.rope_parameters["rope_theta"] == expected_theta
+    common_mod = sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"]
+
+    rotary = common_mod._RotaryEmbedding(config)
+
+    expected = 1.0 / (expected_theta ** (torch.arange(0, 16, 2, dtype=torch.float32) / 16))
+    torch.testing.assert_close(rotary.inv_freq, expected, rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("theta", [10_000.0, 1_000_000.0])
+@pytest.mark.parametrize(
+    "legacy_top_level_only",
+    [False, True],
+    ids=["normalized", "legacy_top_level"],
+)
+def test_code_predictor_rotary_preserves_qwen3_tts_rope_theta(
+    loaded_target_classes,
+    theta,
+    legacy_top_level_only,
+) -> None:
+    """Keep the shared predictor compatible with current and legacy Qwen3-TTS."""
+    config_class = loaded_target_classes[0]
+    config = config_class(
+        hidden_size=16,
+        num_attention_heads=1,
+        head_dim=16,
+        num_code_groups=16,
+        rope_theta=theta,
+    )
+    assert config.rope_parameters["rope_theta"] == theta
+    if legacy_top_level_only:
+        delattr(config, "rope_parameters")
+    common_mod = sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"]
+
+    rotary = common_mod._RotaryEmbedding(config)
+
+    expected = 1.0 / (theta ** (torch.arange(0, 16, 2, dtype=torch.float32) / 16))
+    torch.testing.assert_close(rotary.inv_freq, expected, rtol=0, atol=0)
+
+
 class TestCodePredictorDtypeAlignment:
     """Test that code predictor buffers match model parameter dtype."""