diff --git a/tests/config/test_omni_config.py b/tests/config/test_omni_config.py index 3ee1d40a1..cee4a8214 100644 --- a/tests/config/test_omni_config.py +++ b/tests/config/test_omni_config.py @@ -551,6 +551,7 @@ def test_sub_config_fields_match_structured_scopes(): "enable_sleep_mode", "default_sampling_params", "subtalker_sampling_params", + "silence_ban_frames", "has_sampling_extra_args", "custom_voice_dir", "task_type", diff --git a/tests/engine/test_stage_engine_args.py b/tests/engine/test_stage_engine_args.py index e7c9b2311..9ec7e520a 100644 --- a/tests/engine/test_stage_engine_args.py +++ b/tests/engine/test_stage_engine_args.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from __future__ import annotations import copy @@ -87,6 +90,7 @@ _OMNI_ONLY_LLM_STAGE_ENGINE_FIELDS = frozenset( "num_weight_load_threads", "omni_kv_config", "parallel_config", + "silence_ban_frames", "subtalker_sampling_params", "task_type", "tokenizer_subdir", diff --git a/tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_silence_ban.py b/tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_silence_ban.py new file mode 100644 index 000000000..3552c5356 --- /dev/null +++ b/tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_silence_ban.py @@ -0,0 +1,322 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for the Qwen3-TTS leading-silence ban (#4966, #5048). + +``compute_logits`` suppresses the checkpoint-derived silence codec vocabulary +for the first ``silence_ban_frames`` decode frames of a request, and only for +requests running Base voice-clone in x-vector-only mode. ICL onsets are +legitimate audio and must not be masked. + +Two per-request gates guard the mask, and both index by batch position: + +* the decode step, from ``sampling_metadata.output_token_ids[i]`` -- the + request's generated-token history, so the window closes on its own and a + preempted request cannot re-enter it; +* the generation mode, recorded at prefill in ``preprocess_batch`` and looked + up through the batch req_id order the runner supplies via + ``set_batch_req_ids``. + +These tests drive the methods directly against a talker built with ``__new__``: +no config, no checkpoint, no device. Only the attributes the masking path reads +are populated, and ``logits_processor`` is stubbed to return a fixed tensor. +""" + +from types import SimpleNamespace + +import pytest +import torch + +from vllm_omni.model_executor.models.qwen3_tts.prompt_embeds_builder import resolve_x_vector_only +from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import ( + Qwen3TTSTalkerForConditionalGeneration, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +_VOCAB = 32 +_SILENCE_IDS = (3, 7, 11) + + +def _make_talker(*, ban_frames: int, logits: torch.Tensor, modes=None): + """Talker stub exposing only what ``compute_logits`` touches. + + ``modes`` is one entry per batch row: ``True`` for x-vector-only, ``False`` + for ICL, ``None`` for a request whose mode was never recorded. Defaults to + every row in scope, so a test that says nothing about mode exercises the + step gate alone. + """ + talker = Qwen3TTSTalkerForConditionalGeneration.__new__(Qwen3TTSTalkerForConditionalGeneration) + + silence_mask = torch.zeros((_VOCAB,), dtype=torch.bool) + for token in _SILENCE_IDS: + silence_mask[token] = True + + if modes is None: + modes = [True] * logits.shape[0] + req_ids = [f"req{i}" for i in range(len(modes))] + + object.__setattr__(talker, "_silence_ban_frames", ban_frames) + object.__setattr__(talker, "_silence_mask", silence_mask) + object.__setattr__(talker, "_codec_disallowed_mask", torch.zeros((_VOCAB,), dtype=torch.bool)) + object.__setattr__(talker, "_batch_req_ids", req_ids) + object.__setattr__( + talker, + "_req_x_vector_only", + {rid: mode for rid, mode in zip(req_ids, modes, strict=True) if mode is not None}, + ) + object.__setattr__(talker, "lm_head", None) + object.__setattr__(talker, "logits_processor", lambda _head, _hidden: logits.clone()) + return talker + + +def _metadata(*steps: int): + """sampling_metadata whose output_token_ids encode one step per request.""" + return SimpleNamespace(output_token_ids=[[0] * step for step in steps]) + + +def _banned(row: torch.Tensor) -> set[int]: + return {int(i) for i in torch.nonzero(torch.isinf(row) & (row < 0)).flatten()} + + +# -------------------- the step gate -------------------- + + +def test_ban_disabled_leaves_logits_unchanged() -> None: + logits = torch.zeros((1, _VOCAB)) + talker = _make_talker(ban_frames=0, logits=logits) + out = talker.compute_logits(torch.zeros((1, 4)), _metadata(0)) + assert _banned(out[0]) == set(), "ban_frames=0 must be a no-op" + + +@pytest.mark.parametrize("step", [0, 1, 2]) +def test_early_steps_mask_only_the_silence_vocabulary(step: int) -> None: + logits = torch.zeros((1, _VOCAB)) + talker = _make_talker(ban_frames=3, logits=logits) + out = talker.compute_logits(torch.zeros((1, 4)), _metadata(step)) + assert _banned(out[0]) == set(_SILENCE_IDS), f"step {step} must mask exactly the silence ids" + + +@pytest.mark.parametrize("step", [3, 4, 50]) +def test_steps_past_the_window_are_untouched(step: int) -> None: + logits = torch.zeros((1, _VOCAB)) + talker = _make_talker(ban_frames=3, logits=logits) + out = talker.compute_logits(torch.zeros((1, 4)), _metadata(step)) + assert _banned(out[0]) == set(), f"step {step} is past the window; nothing may be masked" + + +def test_mixed_step_batch_masks_only_the_early_request() -> None: + logits = torch.zeros((2, _VOCAB)) + talker = _make_talker(ban_frames=3, logits=logits) + out = talker.compute_logits(torch.zeros((2, 4)), _metadata(0, 3)) + assert _banned(out[0]) == set(_SILENCE_IDS), "the step-0 request must be masked" + assert _banned(out[1]) == set(), "the step-3 request must not be masked" + + +# -------------------- the mode gate -------------------- + + +def test_icl_request_is_never_masked() -> None: + logits = torch.zeros((1, _VOCAB)) + talker = _make_talker(ban_frames=3, logits=logits, modes=[False]) + out = talker.compute_logits(torch.zeros((1, 4)), _metadata(0)) + assert _banned(out[0]) == set(), "an ICL onset is real audio and must not be suppressed" + + +def test_mixed_mode_batch_masks_only_the_x_vector_request() -> None: + logits = torch.zeros((2, _VOCAB)) + talker = _make_talker(ban_frames=3, logits=logits, modes=[True, False]) + out = talker.compute_logits(torch.zeros((2, 4)), _metadata(0, 0)) + assert _banned(out[0]) == set(_SILENCE_IDS), "the x-vector-only request must be masked" + assert _banned(out[1]) == set(), "the ICL request in the same batch must not be" + + +def test_unrecorded_mode_is_left_alone() -> None: + logits = torch.zeros((2, _VOCAB)) + talker = _make_talker(ban_frames=3, logits=logits, modes=[True, None]) + out = talker.compute_logits(torch.zeros((2, 4)), _metadata(0, 0)) + assert _banned(out[0]) == set(_SILENCE_IDS) + assert _banned(out[1]) == set(), "a request that cannot be classified must not be masked" + + +def test_stale_batch_req_ids_skip_the_ban() -> None: + logits = torch.zeros((2, _VOCAB)) + talker = _make_talker(ban_frames=3, logits=logits) + # One req_id for a two-row batch: the recorded order no longer indexes the + # logits, so masking by position would hit the wrong request. + object.__setattr__(talker, "_batch_req_ids", ["req0"]) + out = talker.compute_logits(torch.zeros((2, 4)), _metadata(0, 0)) + assert _banned(out[0]) == set(), "a req_id/logits length mismatch must disable the ban" + assert _banned(out[1]) == set() + + +# -------------------- metadata plumbing -------------------- + + +def test_missing_sampling_metadata_does_not_crash() -> None: + logits = torch.zeros((1, _VOCAB)) + talker = _make_talker(ban_frames=3, logits=logits) + out = talker.compute_logits(torch.zeros((1, 4)), None) + assert _banned(out[0]) == set() + + +@pytest.mark.parametrize( + "metadata", + [SimpleNamespace(output_token_ids=None), SimpleNamespace(), SimpleNamespace(output_token_ids=[])], +) +def test_absent_output_token_ids_warns_without_masking(metadata) -> None: + logits = torch.zeros((1, _VOCAB)) + talker = _make_talker(ban_frames=3, logits=logits) + out = talker.compute_logits(torch.zeros((1, 4)), metadata) + assert _banned(out[0]) == set(), "the ban must skip, not crash, when history is unavailable" + + +def test_batch_length_mismatch_skips_the_ban() -> None: + logits = torch.zeros((2, _VOCAB)) + talker = _make_talker(ban_frames=3, logits=logits) + out = talker.compute_logits(torch.zeros((2, 4)), _metadata(0)) + assert _banned(out[0]) == set(), "a runner shape mismatch must disable the ban, not misindex it" + assert _banned(out[1]) == set() + + +# -------------------- mode recording and eviction -------------------- + + +def _mode_talker(ban_frames: int = 3): + talker = Qwen3TTSTalkerForConditionalGeneration.__new__(Qwen3TTSTalkerForConditionalGeneration) + object.__setattr__(talker, "_silence_ban_frames", ban_frames) + object.__setattr__(talker, "_req_x_vector_only", {}) + object.__setattr__(talker, "_batch_req_ids", []) + object.__setattr__(talker, "_prompt_builder", SimpleNamespace(preprocess_batch=lambda **_kwargs: None)) + return talker + + +def _info(*, task_type: str = "Base", xvec: bool = True, icl: bool | None = None) -> dict: + info: dict = {"task_type": [task_type], "x_vector_only_mode": [xvec]} + if icl is not None: + info["voice_clone_prompt"] = [{"icl_mode": icl}] + return info + + +def test_preprocess_batch_records_the_generation_mode() -> None: + talker = _mode_talker() + buffer = { + "req0": _info(xvec=True), + "req1": _info(xvec=False), + "req2": _info(task_type="CustomVoice"), + } + talker.preprocess_batch(req_ids=list(buffer), model_intermediate_buffer=buffer, device=torch.device("cpu")) + assert talker._req_x_vector_only == {"req0": True, "req1": False, "req2": None} + + +def test_recorded_mode_matches_the_prompt_builder_resolution() -> None: + """The recording path must agree with the resolution the prefill path uses; + a divergence here would mask requests the builder treats as in-context.""" + infos = [ + _info(xvec=True), + _info(xvec=False), + _info(xvec=True, icl=True), + _info(xvec=False, icl=False), + _info(task_type="VoiceDesign"), + ] + talker = _mode_talker() + buffer = {f"req{i}": info for i, info in enumerate(infos)} + talker.preprocess_batch(req_ids=list(buffer), model_intermediate_buffer=buffer, device=torch.device("cpu")) + for req_id, info in buffer.items(): + assert talker._req_x_vector_only[req_id] == resolve_x_vector_only(info) + + +def test_preprocess_batch_records_nothing_when_the_ban_is_off() -> None: + talker = _mode_talker(ban_frames=0) + buffer = {"req0": _info(xvec=True)} + talker.preprocess_batch(req_ids=["req0"], model_intermediate_buffer=buffer, device=torch.device("cpu")) + assert talker._req_x_vector_only == {} + + +def test_set_batch_req_ids_evicts_departed_requests() -> None: + talker = _mode_talker() + object.__setattr__(talker, "_req_x_vector_only", {"req0": True, "req1": False, "req2": True}) + talker.set_batch_req_ids(["req2", "req0"]) + assert talker._batch_req_ids == ["req2", "req0"], "batch order must be preserved, not sorted" + assert talker._req_x_vector_only == {"req0": True, "req2": True}, "req1 has left the batch" + + +# -------------------- _init_silence_mask sanity checks -------------------- +# +# The silence vocabulary is derived from the checkpoint at load time rather than +# hardcoded. A derivation that comes back empty, oversized, or out of range means +# the encode did not produce silence codes, so the ban disables itself instead of +# masking a guess. + +_TALKER_VOCAB = 2048 + + +def _make_init_talker(*, codes, ban_frames: int = 3, vocab: int = _TALKER_VOCAB): + """Talker stub exposing only what ``_init_silence_mask`` touches. + + ``codes`` is either the list ``_encode_ref_audio_batch`` should return or an + exception instance it should raise. + """ + talker = Qwen3TTSTalkerForConditionalGeneration.__new__(Qwen3TTSTalkerForConditionalGeneration) + + def _encode(_wavs, _sr, *, device): + if isinstance(codes, Exception): + raise codes + return codes + + def _parameters(): + return iter([torch.zeros(1)]) + + object.__setattr__(talker, "_silence_ban_frames", ban_frames) + object.__setattr__(talker, "_silence_mask", torch.zeros((vocab,), dtype=torch.bool)) + object.__setattr__(talker, "talker_config", SimpleNamespace(vocab_size=vocab)) + object.__setattr__(talker, "_encode_ref_audio_batch", _encode) + object.__setattr__(talker, "parameters", _parameters) + return talker + + +def _codes_from(ids) -> list[torch.Tensor]: + """One [T, Q] codec tensor whose codebook-0 column holds ``ids``.""" + cb0 = torch.tensor(list(ids), dtype=torch.long).unsqueeze(1) + return [torch.cat([cb0, torch.zeros_like(cb0)], dim=1)] + + +def test_plausible_derivation_populates_the_mask() -> None: + talker = _make_init_talker(codes=_codes_from([4, 4, 9, 17])) + talker._init_silence_mask() + assert talker._silence_ban_frames == 3 + assert set(torch.nonzero(talker._silence_mask).flatten().tolist()) == {4, 9, 17} + + +def test_empty_derivation_disables_the_ban() -> None: + talker = _make_init_talker(codes=[]) + talker._init_silence_mask() + assert talker._silence_ban_frames == 0, "an all-False mask must not stay enabled" + assert not talker._silence_mask.any() + + +def test_oversized_derivation_disables_the_ban() -> None: + limit = max(64, _TALKER_VOCAB // 20) + talker = _make_init_talker(codes=_codes_from(range(limit + 1))) + talker._init_silence_mask() + assert talker._silence_ban_frames == 0, "a speech-sized vocabulary must not be masked" + assert not talker._silence_mask.any() + + +def test_out_of_range_derivation_disables_the_ban() -> None: + talker = _make_init_talker(codes=_codes_from([4, 9, _TALKER_VOCAB + 1])) + talker._init_silence_mask() + assert talker._silence_ban_frames == 0, "a vocabulary-size disagreement must not be half-applied" + assert not talker._silence_mask.any() + + +def test_encoder_failure_disables_the_ban() -> None: + talker = _make_init_talker(codes=RuntimeError("encoder unavailable")) + talker._init_silence_mask() + assert talker._silence_ban_frames == 0 + assert not talker._silence_mask.any() + + +def test_disabled_ban_skips_derivation_entirely() -> None: + talker = _make_init_talker(codes=RuntimeError("must not be called"), ban_frames=0) + talker._init_silence_mask() + assert talker._silence_ban_frames == 0 diff --git a/vllm_omni/config/model.py b/vllm_omni/config/model.py index ef08dde8d..736a504bd 100644 --- a/vllm_omni/config/model.py +++ b/vllm_omni/config/model.py @@ -141,6 +141,7 @@ class OmniModelConfig(ModelConfig): } ) subtalker_sampling_params: dict[str, Any] | None = None + silence_ban_frames: int = 0 omni_kv_config: dict | None = None codec_frame_rate_hz: float | None = None task_type: str | None = None diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index 8dd9c3db6..c1e24ea4b 100644 --- a/vllm_omni/config/omni_config.py +++ b/vllm_omni/config/omni_config.py @@ -152,6 +152,7 @@ class _ModelEngineOverrides(TypedDict, total=False): active_stream_window: int enable_sleep_mode: bool subtalker_sampling_params: dict[str, Any] + silence_ban_frames: int has_sampling_extra_args: bool custom_voice_dir: str task_type: str @@ -387,6 +388,7 @@ class OmniStageModelConfig: enable_sleep_mode: bool = False default_sampling_params: dict[str, Any] | None = None subtalker_sampling_params: dict[str, Any] | None = None + silence_ban_frames: int = 0 has_sampling_extra_args: bool = False custom_voice_dir: str | None = None task_type: str | None = None diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index 0e4457733..2b1a9f479 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -301,6 +301,7 @@ class StageDeployConfig: default_sampling_params: dict[str, Any] | None = None default_pooling_params: dict[str, Any] | None = None subtalker_sampling_params: dict[str, Any] | None = None + silence_ban_frames: int = 0 # === Generic stage engine fields === # Parallelism, scheduler, and memory-capacity controls. diff --git a/vllm_omni/deploy/qwen3_tts.yaml b/vllm_omni/deploy/qwen3_tts.yaml index 08531869a..58b191f77 100644 --- a/vllm_omni/deploy/qwen3_tts.yaml +++ b/vllm_omni/deploy/qwen3_tts.yaml @@ -73,6 +73,10 @@ stages: # audio (HTTP 400). See #4962. min_tokens: 2 repetition_penalty: 1.05 + # Suppress codebook-0 silence tokens for the first N decode frames of a + # request (#4966). Applies to every request on this stage, so intended + # for x-vector-only (Base voice-clone) deploys. 0 = off. + silence_ban_frames: 0 subtalker_sampling_params: do_sample: true temperature: 0.9 diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index fcd11b173..152e5b3ee 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -184,6 +184,7 @@ class OmniEngineArgs(EngineArgs): custom_process_next_stage_input_func: str | None = None stage_connector_spec: dict[str, Any] = field(default_factory=dict) subtalker_sampling_params: dict[str, Any] | None = None + silence_ban_frames: int = 0 async_chunk: bool = False retains_state_across_chunks: bool = False # WS-A: Stage-1 active stream slots. 0 = legacy preempt-everything. @@ -407,6 +408,7 @@ class OmniEngineArgs(EngineArgs): custom_process_next_stage_input_func=self.custom_process_next_stage_input_func, stage_connector_config=stage_connector_config, subtalker_sampling_params=self.subtalker_sampling_params, + silence_ban_frames=self.silence_ban_frames, omni_kv_config=self.omni_kv_config, task_type=self.task_type, has_sampling_extra_args=self.has_sampling_extra_args, diff --git a/vllm_omni/model_executor/models/qwen3_tts/prompt_embeds_builder.py b/vllm_omni/model_executor/models/qwen3_tts/prompt_embeds_builder.py index 0d8f03294..4053ddae4 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/prompt_embeds_builder.py +++ b/vllm_omni/model_executor/models/qwen3_tts/prompt_embeds_builder.py @@ -85,6 +85,37 @@ def build_instruct_text(instruct: str) -> str: # --------------------------------------------------------------------------- +def resolve_x_vector_only(info_dict: dict) -> bool | None: + """Resolve whether a request runs Base voice-clone in x-vector-only mode. + + Mirrors the resolution in :meth:`Qwen3TTSPromptEmbedsBuilder.build_prompt_embeds` + for the ``task_type == "Base"`` branch: the ``x_vector_only_mode`` flag, then + the ``voice_clone_prompt.icl_mode`` override when present. + + Returns ``None`` when the mode does not apply (any non-Base task), so callers + can distinguish "in-context" from "not a voice-clone request at all". + """ + task_type = first_value(info_dict.get("task_type"), "CustomVoice") + if task_type != "Base": + return None + + xvec_only = bool((info_dict.get("x_vector_only_mode") or [False])[0]) + + raw = info_dict.get("voice_clone_prompt") + if isinstance(raw, list): + raw = raw[0] if raw else None + if isinstance(raw, list) and raw and isinstance(raw[0], dict): + raw = raw[0] + if isinstance(raw, dict) and "icl_mode" in raw: + icl_flag = raw.get("icl_mode") + if isinstance(icl_flag, list): + icl_flag = icl_flag[0] if icl_flag else None + if isinstance(icl_flag, bool): + xvec_only = not icl_flag + + return xvec_only + + def first_value(value: object, default: object = None) -> object: if isinstance(value, list): return value[0] if value else default diff --git a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py index 3daeb4cf2..535a447d3 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py +++ b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py @@ -30,7 +30,7 @@ from vllm_omni.utils.speaker_cache import ( ) from .configuration_qwen3_tts import Qwen3TTSConfig, Qwen3TTSSpeakerEncoderConfig, Qwen3TTSTalkerConfig -from .prompt_embeds_builder import PRECOMPUTED_TEXT_IDS_KEY, Qwen3TTSPromptEmbedsBuilder +from .prompt_embeds_builder import PRECOMPUTED_TEXT_IDS_KEY, Qwen3TTSPromptEmbedsBuilder, resolve_x_vector_only from .qwen3_tts_code_predictor_vllm import Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM from .tokenizer_12hz.configuration_qwen3_tts_tokenizer_v2 import Qwen3TTSTokenizerV2Config from .tokenizer_12hz.modeling_qwen3_tts_tokenizer_v2 import Qwen3TTSTokenizerV2Encoder @@ -421,6 +421,27 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): codec_mask[self._codec_eos_token_id] = True self.register_buffer("_codec_disallowed_mask", ~codec_mask, persistent=False) + # Silence-region codec tokens (#4966). The talker has real probability mass on + # these in the opening decode frames; how many it draws is a sampling outcome, + # which is the source of both the leading silence and its seed-to-seed variance. + # The vocabulary is checkpoint-specific, so it is derived at load time in + # :meth:`_init_silence_mask`; this is only the placeholder. + self._silence_ban_frames = max(0, int(getattr(vllm_config.model_config, "silence_ban_frames", 0) or 0)) + # vLLM only fills sampling_metadata.output_token_ids when penalties, bad + # words, or a logits processor need history. compute_logits reads it for + # the per-request decode step, so request it when the ban is on. + self.logitsprocs_need_output_token_ids = self._silence_ban_frames > 0 + self.register_buffer("_silence_mask", torch.zeros((vocab,), dtype=torch.bool), persistent=False) + + # Per-request generation mode for the silence ban (#4966). The ban is a + # Base voice-clone (x-vector-only) fix; ICL onsets are legitimate audio + # and must not be masked. Mode is only knowable at prefill, so it is + # recorded in preprocess_batch and read back at decode. The batch order + # comes from set_batch_req_ids, which the runner calls each step with + # the same req_id ordering that indexes sampling_metadata. + self._req_x_vector_only: dict[str, bool | None] = {} + self._batch_req_ids: list[str] = [] + # Keys that should stay on GPU in model_intermediate_buffer to avoid # CPU-to-GPU round-trips on every decode step. self.gpu_resident_buffer_keys: set[tuple[str, str]] = { @@ -560,6 +581,48 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): # Mask out invalid codec ids using the pre-built constant buffer. logits = logits.masked_fill(self._codec_disallowed_mask, float("-inf")) + # Suppress silence-region tokens for the first N decode frames (#4966). The + # per-request decode step is the length of its generated-token history. + ban_n = self._silence_ban_frames + if ban_n > 0 and sampling_metadata is not None: + output_token_ids = getattr(sampling_metadata, "output_token_ids", None) + if output_token_ids is None or len(output_token_ids) != logits.shape[0]: + logger.warning_once( + "Silence ban is enabled but output_token_ids (%s) does not match the " + "logits batch (%d); the ban will not take effect.", + "None" if output_token_ids is None else len(output_token_ids), + logits.shape[0], + ) + else: + # Both gates are tiny [B] checks, so keep them on CPU: once no + # request is both in scope and inside the window, skip the device + # work rather than paying a host-to-device copy and a [B, vocab] + # mask on every decode frame. + # + # Scope the ban to x-vector-only requests (#4966). A request whose + # mode was never recorded is left alone: masking a request that + # cannot be classified is the risk this gate exists to remove. + batch_req_ids = self._batch_req_ids + if len(batch_req_ids) == logits.shape[0]: + in_scope = [self._req_x_vector_only.get(rid) is True for rid in batch_req_ids] + else: + logger.warning_once( + "Silence ban is enabled but the recorded batch req_ids (%d) do not match " + "the logits batch (%d); the ban will not take effect.", + len(batch_req_ids), + logits.shape[0], + ) + in_scope = [False] * logits.shape[0] + + steps = [len(t) for t in output_token_ids] + if any(scoped and step < ban_n for scoped, step in zip(in_scope, steps, strict=True)): + early = torch.tensor( + [scoped and step < ban_n for scoped, step in zip(in_scope, steps, strict=True)], + device=logits.device, + dtype=torch.bool, + ).unsqueeze(1) + logits = logits.masked_fill(early & self._silence_mask, float("-inf")) + return logits # -------------------- Omni multimodal output plumbing -------------------- @@ -944,12 +1007,30 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): device: torch.device, ) -> None: """Delegate batched preprocess to :class:`Qwen3TTSPromptEmbedsBuilder`.""" + if self._silence_ban_frames > 0: + for req_id in req_ids: + info_dict = model_intermediate_buffer.get(req_id) + if isinstance(info_dict, dict): + self._req_x_vector_only[req_id] = resolve_x_vector_only(info_dict) self._prompt_builder.preprocess_batch( req_ids=req_ids, model_intermediate_buffer=model_intermediate_buffer, device=device, ) + def set_batch_req_ids(self, req_ids: Sequence[str]) -> None: + """Record the current batch's req_ids, in the order that indexes logits. + + Called by the runner after the forward and before ``compute_logits``, + with the same ``input_batch`` ordering that ``sampling_metadata`` + uses. Doubles as the eviction point: any request no longer in the batch + has finished or been aborted, so its recorded mode is dropped. + """ + self._batch_req_ids = list(req_ids) + if len(self._req_x_vector_only) > len(self._batch_req_ids): + live = set(self._batch_req_ids) + self._req_x_vector_only = {k: v for k, v in self._req_x_vector_only.items() if k in live} + def _encode_ref_audio_batch(self, wavs: list[np.ndarray], sr: int, *, device: torch.device) -> list[torch.Tensor]: fe = self._encoder_feature_extractor target_sr = int(fe.sampling_rate) @@ -1042,6 +1123,7 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): self.encoder.to(device=device, dtype=torch.bfloat16) self._init_runtime_buffers() + self._init_silence_mask() logger.info("Loaded %d weights for Qwen3TTSTalkerForConditionalGeneration", len(loaded)) self._build_stacked_codec_embed() @@ -1056,6 +1138,84 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): device=w.device, dtype=w.dtype ) + @torch.no_grad() + def _init_silence_mask(self) -> None: + """Derive the silence-region codec vocabulary from the checkpoint. + + Encodes silence at several durations and amplitudes and collects the + codebook-0 ids the tokenizer emits. These are the tokens the talker draws + when it produces leading silence (#4966). Derived rather than hardcoded + because the codec vocabulary is checkpoint-specific. + """ + if self._silence_ban_frames <= 0: + return + device = next(self.parameters()).device + sr = 24000 + try: + rng = np.random.RandomState(0) + wavs = [ + (amp * rng.randn(int(dur * sr))).astype(np.float32) + for amp in (0.0, 1e-5, 1e-4, 1e-3) + for dur in (0.3, 0.5, 1.0) + ] + codes = self._encode_ref_audio_batch(wavs, sr, device=device) + ids: set[int] = set() + for code in codes: + code = code.detach().cpu() + cb0 = code[:, 0] if code.ndim == 2 else code + ids.update(int(t) for t in cb0.tolist()) + except Exception as exc: + logger.warning("Could not derive the silence codec vocabulary; disabling the ban: %s", exc) + self._silence_ban_frames = 0 + return + vocab = int(self.talker_config.vocab_size) + in_range = sorted(t for t in ids if 0 <= t < vocab) + num_out_of_range = len(ids) - len(in_range) + if num_out_of_range: + # Every derived id should be a codebook-0 index for this checkpoint. + # Ids outside the talker vocabulary mean the encoder and the talker + # disagree on vocabulary size, so the rest of the derivation is not + # trustworthy either. + logger.warning( + "%d of %d derived silence codec tokens fall outside the talker vocabulary " + "(size %d); disabling the ban.", + num_out_of_range, + len(ids), + vocab, + ) + self._silence_ban_frames = 0 + return + if not in_range: + # An empty derivation would leave an all-False mask: the ban would + # suppress nothing while still paying the per-frame check. + logger.warning("Derived no silence codec tokens; disabling the ban.") + self._silence_ban_frames = 0 + return + # The silence region is a small corner of codebook 0 (12 tokens on + # Qwen3-TTS-12Hz-1.7B-Base). A derivation this large means the encode + # returned speech-like codes, and masking that much of the vocabulary + # would distort generation rather than trim the onset. + max_silence_tokens = max(64, vocab // 20) + if len(in_range) > max_silence_tokens: + logger.warning( + "Derived %d silence codec tokens, above the %d sanity limit for a vocabulary of %d; disabling the ban.", + len(in_range), + max_silence_tokens, + vocab, + ) + self._silence_ban_frames = 0 + return + mask = torch.zeros((vocab,), dtype=torch.bool) + for token in in_range: + mask[token] = True + self._silence_mask.copy_(mask.to(self._silence_mask.device)) + logger.info( + "Derived %d silence codec tokens, suppressed for the first %d decode frames: %s", + len(in_range), + self._silence_ban_frames, + in_range, + ) + @torch.no_grad() def _init_runtime_buffers(self) -> None: """Populate request-independent runtime buffers from frozen weights. diff --git a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py index 1829fafa2..c43a4f8d0 100644 --- a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py +++ b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py @@ -778,6 +778,15 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin, Duplex if callable(flush_pending_metadata): flush_pending_metadata(req_ids[:num_reqs]) + # [Omni] Hand the model the batch's req_ids in logits order, for + # models that gate logits per request. Mirrors gpu_ar_model_runner. + # Only valid without spec decode: there logits_indices carries several + # rows per request, so row i no longer corresponds to req_ids[i]. + if spec_decode_metadata is None: + set_batch_req_ids = getattr(self.model, "set_batch_req_ids", None) + if callable(set_batch_req_ids): + set_batch_req_ids(req_ids[:num_reqs]) + hidden_states, multimodal_outputs = self.extract_multimodal_outputs(hidden_states) if multimodal_outputs is not None: diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index c87d17cc2..1197f5b44 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -1248,6 +1248,15 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex flush_pending_metadata = getattr(self.model, "flush_pending_metadata", None) if callable(flush_pending_metadata): flush_pending_metadata(req_ids[:num_reqs]) + + # [Omni] Hand the model the batch's req_ids in logits order, for + # models that gate logits per request. Only valid without spec + # decode: there logits_indices carries several rows per request, + # so row i no longer corresponds to req_ids[i]. + if spec_decode_metadata is None: + set_batch_req_ids = getattr(self.model, "set_batch_req_ids", None) + if callable(set_batch_req_ids): + set_batch_req_ids(req_ids[:num_reqs]) finally: if runner_assisted_context_enabled: self._set_runner_assisted_full_attention_metadata_context(enabled=False)