diff --git a/tests/worker/test_min_tokens_stop_ids_guard.py b/tests/worker/test_min_tokens_stop_ids_guard.py new file mode 100644 index 000000000..5d7591c99 --- /dev/null +++ b/tests/worker/test_min_tokens_stop_ids_guard.py @@ -0,0 +1,115 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Guard against out-of-vocabulary stop ids in min-tokens masking (#4962). + +vLLM's input processor folds the stage tokenizer's EOS id into +``SamplingParams.all_stop_token_ids``. For AR stages whose lm_head is +narrower than the tokenizer vocabulary (codec talkers such as Qwen3-TTS: +3072 logits vs text EOS 151645), any ``min_tokens >= 1`` makes +``MinTokensLogitsProcessor`` write ``-inf`` at an out-of-range index and +the engine dies with a CUDA device-side assert. +""" + +import pytest +import torch +from vllm import SamplingParams +from vllm.v1.sample.logits_processor import ( + BatchUpdate, + LogitsProcessors, + MinTokensLogitsProcessor, +) + +from vllm_omni.worker.sampling_utils import sanitize_min_tokens_stop_ids + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +TALKER_VOCAB = 3072 +CODEC_EOS = 2150 +TEXT_EOS = 151645 + + +def _make_min_tokens_proc(stop_token_ids: list[int], eos_token_id: int | None) -> MinTokensLogitsProcessor: + """Build a MinTokensLogitsProcessor with one active request, mirroring + the production path (engine folds the tokenizer EOS via + ``update_from_generation_config``).""" + params = SamplingParams(min_tokens=2, stop_token_ids=stop_token_ids) + if eos_token_id is not None: + params.update_from_generation_config({}, eos_token_id) + proc = MinTokensLogitsProcessor(None, device=torch.device("cpu"), is_pin_memory=False) + proc.update_state( + BatchUpdate( + batch_size=1, + removed=[], + added=[(0, params, None, [])], + moved=[], + ) + ) + return proc + + +def test_oob_stop_id_crashes_without_guard(): + """Document the failure mode: the folded text EOS is out of range for a + narrow codec head and index_put_ raises (device-side assert on CUDA).""" + proc = _make_min_tokens_proc([CODEC_EOS], eos_token_id=TEXT_EOS) + logits = torch.zeros(1, TALKER_VOCAB) + with pytest.raises((IndexError, RuntimeError)): + proc.apply(logits) + + +def test_guard_filters_oob_and_keeps_in_range_mask(): + proc = _make_min_tokens_proc([CODEC_EOS], eos_token_id=TEXT_EOS) + sanitize_min_tokens_stop_ids(LogitsProcessors([proc]), TALKER_VOCAB) + + logits = torch.zeros(1, TALKER_VOCAB) + out = proc.apply(logits) + + assert out[0, CODEC_EOS] == -float("inf") + masked = (out == -float("inf")).nonzero() + assert masked.tolist() == [[0, CODEC_EOS]] + + +def test_guard_noop_when_all_stop_ids_in_range(): + """Full-width heads (e.g. CosyVoice3 talker) must be untouched: no + tensor rebuild, mask unchanged.""" + proc = _make_min_tokens_proc([CODEC_EOS], eos_token_id=None) + slice_before = proc.logits_slice + + sanitize_min_tokens_stop_ids(LogitsProcessors([proc]), TALKER_VOCAB) + + assert proc.logits_slice is slice_before + out = proc.apply(torch.zeros(1, TALKER_VOCAB)) + assert out[0, CODEC_EOS] == -float("inf") + + +def test_guard_sanitizes_params_only_once(): + """The stop-id set is shared with the request's SamplingParams, so the + second call finds nothing to drop (no repeated rebuilds/warnings).""" + proc = _make_min_tokens_proc([CODEC_EOS], eos_token_id=TEXT_EOS) + sanitize_min_tokens_stop_ids(LogitsProcessors([proc]), TALKER_VOCAB) + slice_after_first = proc.logits_slice + + sanitize_min_tokens_stop_ids(LogitsProcessors([proc]), TALKER_VOCAB) + + assert proc.logits_slice is slice_after_first + + +def test_guard_handles_multiple_requests(): + params_oob = SamplingParams(min_tokens=4, stop_token_ids=[CODEC_EOS]) + params_oob.update_from_generation_config({}, TEXT_EOS) + params_ok = SamplingParams(min_tokens=4, stop_token_ids=[7]) + proc = MinTokensLogitsProcessor(None, device=torch.device("cpu"), is_pin_memory=False) + proc.update_state( + BatchUpdate( + batch_size=2, + removed=[], + added=[(0, params_oob, None, []), (1, params_ok, None, [])], + moved=[], + ) + ) + + sanitize_min_tokens_stop_ids(LogitsProcessors([proc]), TALKER_VOCAB) + out = proc.apply(torch.zeros(2, TALKER_VOCAB)) + + assert out[0, CODEC_EOS] == -float("inf") + assert out[1, 7] == -float("inf") + assert int((out == -float("inf")).sum()) == 2 diff --git a/vllm_omni/deploy/qwen3_tts.yaml b/vllm_omni/deploy/qwen3_tts.yaml index 69085fc5a..34f0d6cdc 100644 --- a/vllm_omni/deploy/qwen3_tts.yaml +++ b/vllm_omni/deploy/qwen3_tts.yaml @@ -54,6 +54,10 @@ stages: temperature: 0.9 top_k: 50 max_tokens: 4096 + # Official qwen-tts generates the talker with min_new_tokens=2; without + # it, checkpoints that sample codec EOS on the first frame return empty + # audio (HTTP 400). See #4962. + min_tokens: 2 seed: 42 repetition_penalty: 1.05 subtalker_sampling_params: diff --git a/vllm_omni/deploy/qwen3_tts_high_concurrency.yaml b/vllm_omni/deploy/qwen3_tts_high_concurrency.yaml index 4fb6f9894..939396840 100644 --- a/vllm_omni/deploy/qwen3_tts_high_concurrency.yaml +++ b/vllm_omni/deploy/qwen3_tts_high_concurrency.yaml @@ -60,6 +60,10 @@ stages: temperature: 0.9 top_k: 50 max_tokens: 4096 + # Official qwen-tts generates the talker with min_new_tokens=2; without + # it, checkpoints that sample codec EOS on the first frame return empty + # audio (HTTP 400). See #4962. + min_tokens: 2 repetition_penalty: 1.05 subtalker_sampling_params: do_sample: true 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 a0abc7fa2..9a8a55709 100644 --- a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py +++ b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py @@ -46,6 +46,7 @@ from vllm_omni.distributed.omni_connectors.kv_transfer_manager import OmniKVTran from vllm_omni.outputs import OmniModelRunnerOutput from vllm_omni.platforms.npu.worker.npu_model_runner import OmniNPUModelRunner from vllm_omni.utils.mm_outputs import build_mm_cpu, partition_payload_list, to_payload_element +from vllm_omni.worker.sampling_utils import sanitize_min_tokens_stop_ids def _ensure_tensor_values(payload: dict[str, object]) -> dict[str, torch.Tensor]: @@ -922,6 +923,15 @@ class NPUARModelRunner(OmniNPUModelRunner): logits_vocab = logits.shape[-1] if self.input_batch.vocab_size > logits_vocab: smd.prompt_token_ids = smd.prompt_token_ids.clamp(max=logits_vocab) + + # Drop min-tokens stop ids the head cannot emit (e.g. the text + # tokenizer EOS folded into all_stop_token_ids on a narrow codec + # talker head); they would index_put_ out of bounds (#4962). + if logits is not None: + sanitize_min_tokens_stop_ids( + self.input_batch.sampling_metadata.logitsprocs, + logits.shape[-1], + ) # -------------------------------------- Omni-new ------------------------------------------------- diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index 2bf7a8af2..a5e4a9bd1 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -47,6 +47,7 @@ from vllm_omni.utils.mm_outputs import build_mm_cpu, partition_payload_list, to_ from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner from vllm_omni.worker.omni_connector_model_runner_mixin import OmniConnectorModelRunnerMixin from vllm_omni.worker.runner_assisted_metadata import RunnerAssistedFullAttentionMetadataRequest +from vllm_omni.worker.sampling_utils import sanitize_min_tokens_stop_ids logger = init_logger(__name__) @@ -1892,6 +1893,15 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): if self.input_batch.vocab_size > logits_vocab: smd.prompt_token_ids = smd.prompt_token_ids.clamp(max=logits_vocab) + # Drop min-tokens stop ids the head cannot emit (e.g. the text + # tokenizer EOS folded into all_stop_token_ids on a narrow codec + # talker head); they would index_put_ out of bounds (#4962). + if logits is not None: + sanitize_min_tokens_stop_ids( + self.input_batch.sampling_metadata.logitsprocs, + logits.shape[-1], + ) + with record_function_or_nullcontext("gpu_model_runner: sample"): sampler_output = self._sample(logits, spec_decode_metadata) diff --git a/vllm_omni/worker/sampling_utils.py b/vllm_omni/worker/sampling_utils.py new file mode 100644 index 000000000..c38a9bbdc --- /dev/null +++ b/vllm_omni/worker/sampling_utils.py @@ -0,0 +1,58 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Sampling-state guards shared by the GPU and NPU AR model runners.""" + +import torch +from vllm.logger import init_logger +from vllm.v1.sample.logits_processor import LogitsProcessors, MinTokensLogitsProcessor + +logger = init_logger(__name__) + +__all__ = ["sanitize_min_tokens_stop_ids"] + + +def sanitize_min_tokens_stop_ids(logitsprocs: LogitsProcessors, logits_vocab: int) -> None: + """Drop stop ids the model head cannot emit from min-tokens masking state. + + vLLM's input processor unconditionally folds the stage tokenizer's EOS id + into ``SamplingParams.all_stop_token_ids``. AR stages whose lm_head is + narrower than the tokenizer vocabulary (codec talkers such as Qwen3-TTS: + 3072 logits vs text EOS 151645) then crash on any ``min_tokens >= 1``: + ``MinTokensLogitsProcessor.apply`` writes ``-inf`` at an out-of-range + index and ``index_put_`` triggers a CUDA device-side assert (#4962). + + Out-of-range ids are unreachable for the head, so dropping them never + changes sampling or stopping behavior. The per-request stop-id set is + mutated in place (it is shared with the request's ``SamplingParams``), + so each request is sanitized at most once; the processor's device-side + mask slice is rebuilt only when an out-of-range id was found. + """ + for proc in logitsprocs.non_argmax_invariant: + if not isinstance(proc, MinTokensLogitsProcessor): + continue + min_toks = getattr(proc, "min_toks", None) + if not min_toks: + continue + needs_rebuild = False + for _, _, stop_tok_ids in min_toks.values(): + oob = [tok for tok in stop_tok_ids if tok >= logits_vocab] + if not oob: + continue + stop_tok_ids.difference_update(oob) + needs_rebuild = True + logger.warning_once( + "min_tokens: dropped stop token ids %s that exceed the logits vocabulary (%d); " + "the model head cannot emit them.", + str(sorted(oob)), + logits_vocab, + ) + if needs_rebuild: + reqs: list[int] = [] + tok_ids: list[int] = [] + for index, (_, _, stop_tok_ids) in min_toks.items(): + reqs.extend([index] * len(stop_tok_ids)) + tok_ids.extend(stop_tok_ids) + proc.logits_slice = ( + proc._device_tensor(reqs, torch.int32), + proc._device_tensor(tok_ids, torch.int32), + )