[Bugfix] Drop out-of-vocabulary stop ids from min-tokens masking (qwen3-tts min_tokens engine crash) (#4971)

Signed-off-by: Yueqian Lin <linyueqian@outlook.com>

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

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:

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

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 -------------------------------------------------
 
 

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)
 

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),
+            )