vllm_omni.worker.sampling_utils ¶

Sampling helpers shared by the GPU and NPU AR model runners.

logger `module-attribute` ¶

logger = init_logger(__name__)

apply_fixed_seed_to_sampling_params ¶

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.

build_model_sampler_extra_args ¶

build_model_sampler_extra_args(
    input_batch: Any, requests: Any
) -> list[dict | None]

Return per-request SamplingParams.extra_args in batch-row order.

call_model_sampler ¶

call_model_sampler(
    model: Any,
    model_sample: Any,
    logits: 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.

clamp_prompt_ids_to_penalty_padding ¶

clamp_prompt_ids_to_penalty_padding(
    prompt_token_ids: Tensor, logits_vocab: int
) -> 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.

get_tts_local_seed ¶

get_tts_local_seed(sampling_params: Any) -> int | None

Return the explicit per-request Talker MTP seed, if configured.

sanitize_min_tokens_stop_ids ¶

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.

sanitize_sampling_params_min_tokens_stop_ids ¶

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.