177 lines
7.0 KiB
Python
177 lines
7.0 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""Sampling helpers shared by the GPU and NPU AR model runners."""
|
|
|
|
from typing import Any
|
|
|
|
import torch
|
|
from vllm.logger import init_logger
|
|
from vllm.sampling_params import SamplingParams
|
|
from vllm.v1.sample.logits_processor import LogitsProcessors, MinTokensLogitsProcessor
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
__all__ = [
|
|
"apply_fixed_seed_to_sampling_params",
|
|
"get_tts_local_seed",
|
|
"clamp_prompt_ids_to_penalty_padding",
|
|
"sanitize_min_tokens_stop_ids",
|
|
"sanitize_sampling_params_min_tokens_stop_ids",
|
|
"build_model_sampler_extra_args",
|
|
"call_model_sampler",
|
|
]
|
|
|
|
|
|
def get_tts_local_seed(sampling_params: Any) -> int | None:
|
|
"""Return the explicit per-request Talker MTP seed, if configured."""
|
|
extra_args = getattr(sampling_params, "extra_args", None) if sampling_params is not None else None
|
|
seed = extra_args.get("tts_local_seed") if isinstance(extra_args, dict) else None
|
|
return int(seed) if seed is not None else None
|
|
|
|
|
|
def apply_fixed_seed_to_sampling_params(
|
|
sampling_params: Any,
|
|
seed: int,
|
|
*,
|
|
seed_talker_mtp: bool,
|
|
) -> None:
|
|
"""Apply one request seed to a stage sampler and, for Talkers, its MTP RNG."""
|
|
seed = int(seed)
|
|
sampling_params.seed = seed
|
|
if not seed_talker_mtp:
|
|
return
|
|
extra_args = dict(getattr(sampling_params, "extra_args", None) or {})
|
|
extra_args["tts_local_seed"] = seed
|
|
sampling_params.extra_args = extra_args
|
|
|
|
|
|
def sanitize_sampling_params_min_tokens_stop_ids(
|
|
sampling_params: SamplingParams,
|
|
logits_vocab: int,
|
|
) -> None:
|
|
"""Remove unreachable ids before MRv2 builds persistent min-token state.
|
|
|
|
MRv2 copies ``SamplingParams.all_stop_token_ids`` into GPU-resident
|
|
``LogitBiasState`` when a request is registered. Narrow codec heads cannot
|
|
safely carry the text tokenizer EOS in that state. The engine-facing EOS
|
|
and explicit ``stop_token_ids`` fields remain unchanged; only the set used
|
|
by the sampler's min-token mask is normalized.
|
|
"""
|
|
if sampling_params.min_tokens <= 0:
|
|
return
|
|
stop_token_ids = sampling_params.all_stop_token_ids
|
|
unreachable = {token_id for token_id in stop_token_ids if token_id < 0 or token_id >= logits_vocab}
|
|
if not unreachable:
|
|
return
|
|
stop_token_ids.difference_update(unreachable)
|
|
logger.warning_once(
|
|
"min_tokens: dropped stop token ids %s that exceed the logits vocabulary (%d); "
|
|
"the model head cannot emit them.",
|
|
str(sorted(unreachable)),
|
|
logits_vocab,
|
|
)
|
|
|
|
|
|
def build_model_sampler_extra_args(input_batch: Any, requests: Any) -> list[dict | None]:
|
|
"""Return per-request ``SamplingParams.extra_args`` in batch-row order."""
|
|
request_states = requests or {}
|
|
per_req_extra_args: list[dict | None] = []
|
|
for req_id in getattr(input_batch, "req_ids", []):
|
|
state = request_states.get(req_id)
|
|
params = getattr(state, "sampling_params", None)
|
|
per_req_extra_args.append(getattr(params, "extra_args", None))
|
|
return per_req_extra_args
|
|
|
|
|
|
def call_model_sampler(
|
|
model: Any,
|
|
model_sample: Any,
|
|
logits: torch.Tensor,
|
|
sampling_metadata: Any,
|
|
*,
|
|
input_batch: Any,
|
|
requests: Any,
|
|
) -> Any:
|
|
"""Call an opted-in model sampler with per-request extra arguments.
|
|
|
|
The opt-in keeps the existing two-argument sampler contract unchanged for
|
|
every other model while allowing custom samplers to make row-local choices.
|
|
"""
|
|
if not getattr(model, "model_sampler_wants_extra_args", False):
|
|
return model_sample(logits, sampling_metadata)
|
|
return model_sample(
|
|
logits,
|
|
sampling_metadata,
|
|
per_req_extra_args=build_model_sampler_extra_args(input_batch, requests),
|
|
)
|
|
|
|
|
|
def clamp_prompt_ids_to_penalty_padding(prompt_token_ids: torch.Tensor, logits_vocab: int) -> torch.Tensor:
|
|
"""Clamp batch-level pad ids down to ``logits_vocab`` — upstream's
|
|
designed penalty padding value.
|
|
|
|
``max=logits_vocab`` (NOT ``logits_vocab - 1``) is deliberate: upstream
|
|
penalty computation allocates ``vocab_size + 1`` bins and drops the last
|
|
column, so ``vocab_size`` is the padding value that never affects
|
|
penalties (vllm/model_executor/layers/utils.py::
|
|
get_token_bin_counts_and_mask). Clamping one lower would count padding
|
|
as real occurrences of the last vocab token.
|
|
"""
|
|
return prompt_token_ids.clamp(max=logits_vocab)
|
|
|
|
|
|
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] = []
|
|
restore_reqs: list[int] = []
|
|
restore_tok_ids: list[int] = []
|
|
for index, (_, _, stop_tok_ids, uses_structured_output) in min_toks.items():
|
|
reqs.extend([index] * len(stop_tok_ids))
|
|
tok_ids.extend(stop_tok_ids)
|
|
if uses_structured_output:
|
|
restore_reqs.extend([index] * len(stop_tok_ids))
|
|
restore_tok_ids.extend(stop_tok_ids)
|
|
proc.logits_slice = (
|
|
proc._device_tensor(reqs, torch.int32),
|
|
proc._device_tensor(tok_ids, torch.int32),
|
|
)
|
|
proc.restore_logits_slice = (
|
|
proc._device_tensor(restore_reqs, torch.int32),
|
|
proc._device_tensor(restore_tok_ids, torch.int32),
|
|
)
|