Sampler: the last mile from logits to token

源码版本v0.25.1

Responsibilities

Once the model forward finishes, GPUModelRunner has [num_tokens, vocab_size] logits. The remaining work is to turn logits into token ids: apply temperature, top_k, top_p, penalties; do greedy or multinomial sampling; and along the way collect top-k logprobs. Sampler(sampler.py:20) is the nn.Module that does this; each GPUModelRunner holds one instance(gpu_model_runner.py:532).

Sampler.forward is a fixed nine-step pipeline(sampler.py:22-59): first compute raw_logprobs / clone raw_logits as needed, convert logits to float32, apply the allowed_token_ids whitelist and bad_words blacklist, run the non-argmax-invariant logits processors(MinTokens, LogitBias), apply repetition / frequency / presence penalty, then enter sample() for greedy or random sampling. sample(sampler.py:243) further branches internally: greedy goes through argmax; random goes through apply_temperature -> argmax-invariant processor(default MinP) -> topk_topp_sampler -> torch.where(temp < eps, greedy, random).

The actual top-k/top-p implementation does not live in Sampler but in TopKTopPSampler(topk_topp_sampler.py:70), which on CUDA uses FlashInfer's top_p_sampling_from_probs / top_k_sampling_from_probs(topk_topp_sampler.py:471), on CPU/XPU/ROCm uses each platform's native implementation, and falls back to forward_native otherwise.

Design motivation

• argmax invariance classification: LogitsProcessor.is_argmax_invariant(interface.py:84-92) distinguishes whether the greedy result changes; MinP, MinTokens, and LogitBias fall on each side respectively. Greedy requests skip the whole random flow after the argmax-invariant processors(sampler.py:257-271).
• Temperature 0 is treated as greedy: apply_temperature rewrites positions where temp < eps to 1.0(sampler.py:233-237), and finally torch.where(temp < eps, greedy, random) merges once, avoiding a branch.
• In-place first: apply_temperature uses logits.div_, and MinP.apply uses masked_fill_(builtin.py:115), reducing peak memory.
• FlashInfer sampler is fast but limited: it does not return the filtered logits(topk_topp_sampler.py:86-95), so when logprobs_mode is processed_logits / processed_logprobs, forward_native is forced.
• logprobs collection is separate from sampling: gather_logprobs(sampler.py:308-356) runs after sampling, using the original (not temperature-scaled) logits to compute log_softmax, so penalties do not affect the logprob values returned to the user.
• Guard against batch-dim dynamic specialization: gather_logprobs uses torch._dynamo.decorators.mark_unbacked(sampler.py:345-346) so that dynamo does not recompile when batch_size=1 -> >=2.
• spec decode compatibility: apply_logits_processors, when predict_bonus_token is set, appends spec_token_ids after output_token_ids(sampler.py:384-393), so the penalty sees the draft tokens.

Key files

• Sampler class:20 — nn.Module, holds TopKTopPSampler, implements the 9-step pipeline.
• forward docstring:22-59 — the 9-step pipeline spec, annotating what each step does.
• Sampler.forward:72-149 — main entry, strings together apply_logits_processors + sample + gather_logprobs.
• Sampler.sample:243 — greedy / random branch, temperature application + top-k/top-p.
• apply_temperature:227-237 — in-place div_, rewrites zero temperature to 1.0 to avoid divide-by-zero.
• greedy_sample:239-241 — logits.argmax(dim=-1).view(-1).
• apply_logits_processors:371-420 — allowed_token_ids_mask + bad_words + non_argmax_invariant processors + penalties.
• apply_penalties:422-439 — delegates to the apply_all_penalties op.
• SamplingMetadata:14 — holds per-batch sampling metadata such as temperature/top_p/top_k/generators/penalties.
• TopKTopPSampler:70 — platform dispatch: CUDA->FlashInfer, CPU->native, XPU->kernel, ROCm->aiter.
• flashinfer_sample:471-508 — calls FlashInfer top_p_renorm_probs / top_k_sampling_from_probs; statistically equivalent but faster than rejection sampling.
• GPUModelRunner holds Sampler:532 — at construction Sampler(logprobs_mode, use_fp64_gumbel).

Data flow

The input at every step is [num_tokens, vocab_size] logits, where num_tokens is the total number of tokens to compute across the batch in the current step(prefill computes multiple prompt tokens; decode computes 1 per request). Below is the greedy vs. random branching inside sample:

# vllm/v1/sample/sampler.py L243-L302
def sample(
    self,
    logits: torch.Tensor,
    sampling_metadata: SamplingMetadata,
    logprobs_mode_override: LogprobsMode | None = None,
) -> tuple[torch.Tensor, torch.Tensor | None]:
    logprobs_mode = logprobs_mode_override or self.logprobs_mode
    assert not (sampling_metadata.all_greedy and sampling_metadata.all_random)
    if sampling_metadata.all_random:
        greedy_sampled = None
    else:
        greedy_sampled = self.greedy_sample(logits)
        if sampling_metadata.all_greedy:
            processed_logprobs = None
            if (
                sampling_metadata.max_num_logprobs is not None
                or sampling_metadata.logprob_token_ids
            ):
                if logprobs_mode == "processed_logits":
                    processed_logprobs = logits
                elif logprobs_mode == "processed_logprobs":
                    processed_logprobs = self.compute_logprobs(logits)
            return greedy_sampled, processed_logprobs

    assert sampling_metadata.temperature is not None

    # Apply temperature.
    logits = self.apply_temperature(
        logits, sampling_metadata.temperature, sampling_metadata.all_random
    )

    # Apply logits processors that only apply to random sampling
    # (argmax invariant)
    for processor in sampling_metadata.logitsprocs.argmax_invariant:
        logits = processor.apply(logits)

    # Apply top_k and/or top_p.
    random_sampled, processed_logprobs = self.topk_topp_sampler(
        logits,
        sampling_metadata.generators,
        sampling_metadata.top_k,
        sampling_metadata.top_p,
    )

    if greedy_sampled is None:
        return random_sampled, processed_logprobs

    sampled = torch.where(
        sampling_metadata.temperature < _SAMPLING_EPS,
        greedy_sampled,
        random_sampled,
        out=greedy_sampled,  # Reuse tensor
    )
    return sampled, processed_logprobs

all_greedy / all_random are batch-level flags pre-computed in SamplingMetadata so the whole batch takes the fast path. In a mixed batch both paths are computed, then torch.where selects per-request by temperature — this avoids the GPU slowing down due to branching.

Boundaries and failure modes

• all_greedy and all_random are mutually exclusive: sample opens with assert not (all_greedy and all_random)(sampler.py:256); configuring both raises AssertionError directly.
• Zero-temperature tokens do not go through random: torch.where(temp < eps, greedy, random)(sampler.py:296-301) guarantees greedy requests always get argmax, not polluted by random.
• logprobs use the original logits: compute_logprobs runs before apply_temperature(sampler.py:85-93); penalty / temperature do not affect the logprob values returned to the user. This is a difference from v0.
• logprob_token_ids takes a dedicated gather: gather_specific_token_logprobs(sampler.py:151-225) is used by the generative_scoring API; it skips top-k and directly gathers the specified tokens, with padding and valid_mask.
• spec decode disables logits processors: build_logitsprocs raises ValueError if it detects speculative_config and a custom processor is passed(__init__.py:201-209); only MinTokens is kept.
• FlashInfer outputs int32: at the end of forward, sampled = sampled.long()(sampler.py:109) unifies to int64, because FlashInfer sampler returns int32 while PyTorch argmax / topk return int64; subsequent index operations need to be compatible.

Summary

Sampler is vLLM's unified entry point for sampling: a nine-step pipeline + greedy/random branching + FlashInfer acceleration. It relies on the LogitsProcessor abstraction in /sampling/logits for parameterized transforms, and itself only cares about "how to turn logits into tokens". The returned SamplerOutput is passed up to /worker/gpu-model-runner, where OutputProcessor eventually assembles it into RequestOutput. For how the engine core schedules these steps, see /engine/engine-core.

See official docs: vLLM docs · README