LogitsProcessor: pre-sampling logits transforms

源码版本v0.25.1

Responsibilities

After Sampler gets the raw logits, it still runs a pile of transforms before sampling can start: MinP filters low-probability tokens, LogitBias adds biases to specific tokens, MinTokens forbids EOS until enough tokens are generated, repetition/frequency/presence penalties penalize repetition, structured output masks tokens outside the grammar, and thinking budget adds a budget for reasoning models. The unified abstraction for these transforms is LogitsProcessor(interface.py:60). Each processor holds its own batch state and is called in sequence by Sampler.apply_logits_processors(sampler.py:371) before every forward.

The LogitsProcessor interface is thin, just four methods: __init__, apply(logits) -> logits, is_argmax_invariant() -> bool(interface.py:84), and update_state(batch_update)(interface.py:94). is_argmax_invariant decides whether it lands in LogitsProcessors.argmax_invariant or non_argmax_invariant(state.py:148-160). The former only runs under random sampling (MinP); the latter also runs under greedy sampling (MinTokens, LogitBias). update_state is called whenever the batch composition changes, giving the processor a chance to refresh the per-request tensors it holds.

Built-in processors live in BUILTIN_LOGITS_PROCESSORS(__init__.py:49-53), in order [MinTokensLogitsProcessor, LogitBiasLogitsProcessor, MinPLogitsProcessor]. Beyond that, entry-point plugins (the vllm.logits_processors group) and FQCN strings are also supported(__init__.py:86-155). build_logitsprocs(__init__.py:184) merges all three categories and instantiates them.

Design motivation

• Argmax-invariance classification: MinPLogitsProcessor.is_argmax_invariant() = True(builtin.py:47-49), because it only filters extremely low-probability tokens and does not affect argmax. MinTokens / LogitBias return False(builtin.py:183-186), because they can change the greedy result (masking EOS or adding a bias may change argmax).
• Incremental batch state update: update_state takes a BatchUpdate(removed/added/moved)(interface.py:36-57). The processor only recomputes its tensors when the batch composition actually changes (MinP uses min_p_count to skip empty batches)(builtin.py:54-100).
• Dual CPU/GPU tensors: MinPLogitsProcessor holds min_p_cpu(pinned) + min_p_device(builtin.py:30-44). Values are updated on CPU, then H2D-copied once before apply to avoid cross-device sync on every apply.
• MinP operates on softmax, not logits: apply first computes softmax(logits) to find the max prob, then filters by the max_prob * min_p threshold(builtin.py:102-116), because min_p is a probability-space threshold, not a logit-space one.
• LogitBias uses sparse indexing: LogitBiasLogitsProcessor.apply does logits[req_idx, tok_id] += bias(builtin.py:159-162), only modifying positions that have a bias, instead of adding across the entire vocab.
• MinTokens masks EOS: MinTokensLogitsProcessor masks stop_token_ids until the generated count reaches min_tokens(builtin.py:165-186), forcing the model to generate long enough.
• Spec decode disables custom processors: build_logitsprocs keeps only MinTokensLogitsProcessor when speculative_config is non-empty(__init__.py:201-209). MinP / LogitBias are both turned off because they are semantically incompatible with rejection sampling.

Key files

• LogitsProcessor ABC:60 — four-method abstract base class: __init__, apply, is_argmax_invariant, update_state.
• BatchUpdate:36 — frozen dataclass holding the three sequences removed/added/moved.
• LogitsProcessors:148 — container, split into two paths by argmax_invariant.
• LogitsProcessors.all:162-165 — chains the two paths into one iterator.
• MinPLogitsProcessor:23 — probability-space threshold filter, argmax-invariant.
• LogitBiasLogitsProcessor:119 — sparse bias addition, non-argmax-invariant.
• MinTokensLogitsProcessor:165 — masks stop tokens until enough tokens are generated, non-argmax-invariant.
• BUILTIN_LOGITS_PROCESSORS:49-53 — default three-piece order: MinTokens -> LogitBias -> MinP.
• build_logitsprocs:184 — factory entry point, handles pooling / spec-decode / custom short-circuits.
• _load_logitsprocs_by_fqcns:86-155 — lazy-loads <module>:<Qualname> strings + issubclass check.
• ThinkingBudgetStateHolder:33 — reasoning budget processor; apply_to_logits is called by Sampler during forward.

Data flow

Before every forward, Scheduler or GPUModelRunner builds a BatchUpdate telling the processor which requests were added, removed, or moved in this batch. The processor uses this chance to refresh its per-request tensors. MinP's update_state is a typical example:

# vllm/v1/sample/logits_processor/builtin.py L54-L100
def update_state(self, batch_update: BatchUpdate | None):
    if not batch_update:
        return

    needs_update = False
    # Process added requests.
    for index, params, _, _ in batch_update.added:
        min_p = params.min_p
        min_p_before = self.min_p_cpu[index]
        if min_p_before != min_p:
            needs_update = True
            self.min_p_cpu[index] = min_p
            if min_p and not min_p_before:
                self.min_p_count += 1
            elif not min_p and min_p_before:
                self.min_p_count -= 1

    if self.min_p_count:
        # Process removed requests.
        if batch_update.removed:
            needs_update = True
            for index in batch_update.removed:
                if self.min_p_cpu[index]:
                    self.min_p_cpu[index] = 0
                    self.min_p_count -= 1

        # Process moved requests, unidirectional (a->b) and swap (a<->b).
        for adx, bdx, direct in batch_update.moved:
            min_p_a, min_p_b = self.min_p_cpu[adx], self.min_p_cpu[bdx]
            if min_p_a != min_p_b:
                needs_update = True
                self.min_p_cpu[bdx] = min_p_a
                if direct == MoveDirectionality.SWAP:
                    self.min_p_cpu[adx] = min_p_b
            if direct == MoveDirectionality.UNIDIRECTIONAL:
                if min_p_a:
                    self.min_p_cpu[adx] = 0
                if min_p_b:
                    self.min_p_count -= 1

    # Update tensors if needed.
    size = batch_update.batch_size
    if self.min_p_count and (needs_update or self.min_p.shape[0] != size):
        self.min_p = self.min_p_device[:size]
        if self.use_double_tensor:
            self.min_p.copy_(self.min_p_cpu_tensor[:size], non_blocking=True)
        self.min_p.unsqueeze_(1)

min_p_count is an escape counter: as long as no request in the batch has min_p enabled, apply returns immediately without any computation. The apply body is also short:

# vllm/v1/sample/logits_processor/builtin.py L102-L116
def apply(self, logits: torch.Tensor) -> torch.Tensor:
    if not self.min_p_count:
        return logits

    # Convert logits to probability distribution
    probability_values = torch.nn.functional.softmax(logits, dim=-1)
    # Calculate maximum probabilities per sequence
    max_probabilities = torch.amax(probability_values, dim=-1, keepdim=True)
    # Adjust min_p
    adjusted_min_p = max_probabilities.mul_(self.min_p)
    # Identify valid tokens using threshold comparison
    invalid_token_mask = probability_values < adjusted_min_p
    # Apply mask using boolean indexing
    logits.masked_fill_(invalid_token_mask, -float("inf"))
    return logits

Sampler.apply_logits_processors calls non-argmax-invariant and argmax-invariant processors in two phases: the former also runs on the greedy path, the latter only on the random path(sampler.py:403-405). Penalties go through the standalone apply_all_penalties op, not through the LogitsProcessor layer.

Boundaries and failure modes

• Pooling models do not support processors: build_logitsprocs returns an empty LogitsProcessors when is_pooling_model=True, and raises STR_POOLING_REJECTS_LOGITSPROCS if any custom ones were passed in(__init__.py:191-198).
• Spec decode disables custom processors: STR_SPEC_DEC_REJECTS_LOGITSPROCS(__init__.py:201-209). Only MinTokensLogitsProcessor is kept, because MinP / LogitBias are incompatible with rejection sampling.
• TPU does not support custom processors: _load_custom_logitsprocs returns an empty list when is_tpu()(__init__.py:176-179). v1 has not yet wired custom processors on TPU.
• Plugin load failures raise hard: _load_logitsprocs_plugins directly raise RuntimeError when an entry point fails to load(__init__.py:78-82). It does not silently skip, so users do not discover at runtime that some processor never took effect.
• FQCN must be a LogitsProcessor subclass: _load_logitsprocs_by_fqcns splits on <module>:<Qualname>(__init__.py:128) and raises ValueError if the target is not a subclass, preventing plain functions from being passed in by mistake.
• batch_update reference semantics: The output_tok_ids in each tuple of BatchUpdate.added is a reference to the request's running tokens list(interface.py:44-50). The processor sees the latest generated tokens through this reference, no extra updates needed.

Summary

LogitsProcessor is the pluggable transform layer above Sampler and below sampling: argmax-invariance decides whether it runs on the greedy path or only the random path, and update_state lets it track batch changes incrementally. The built-in trio MinP / LogitBias / MinTokens is extended via plugins and FQCN strings. For the actual sampling process see /sampling/sampler. Penalties and bad_words are called directly inside Sampler via apply_all_penalties / apply_bad_words, not through this abstraction.

See official docs: vLLM docs · README