Sampler: die letzte Meile von logits zum Token

源码版本v0.25.1

Verantwortung

Wenn der Modell-Forward abgeschlossen ist, hält GPUModelRunner die logits der Form [num_tokens, vocab_size]. Was übrig bleibt, ist die Umwandlung der logits in Token-IDs: Temperature, top_k, top_p und Penalties anwenden, greedy oder multinomial samplen und nebenbei die Top-k-Logprobs einsammeln. Sampler (sampler.py:20) ist genau das — ein nn.Module, das jeder GPUModelRunner als Instanz hält (gpu_model_runner.py:532).

Sampler.forward ist eine feste neunstufige Pipeline (sampler.py:22-59): bei Bedarf raw_logprobs / clone raw_logits berechnen, logits nach float32 wandeln, allowed_token_ids-Whitelist und bad_words-Blacklist anwenden, die nicht-argmax-invarianten Logits-Prozessoren (MinTokens, LogitBias) laufen lassen, repetition / frequency / presence penalty anwenden und schließlich in sample() greedy oder random samplen. sample (sampler.py:243) teilt sich intern weiter auf: greedy geht über argmax, random über apply_temperature → argmax-invariant processor (standardmäßig MinP) → topk_topp_sampler → torch.where(temp < eps, greedy, random).

Die konkrete top-k/top-p-Implementierung liegt nicht in Sampler, sondern in TopKTopPSampler (topk_topp_sampler.py:70); auf CUDA läuft sie über FlashInfers top_p_sampling_from_probs / top_k_sampling_from_probs (topk_topp_sampler.py:471), auf CPU/XPU/ROCm über die jeweiligen nativen Implementierungen, und wenn nichts davon zutrifft, fällt sie auf forward_native zurück.

Entwurfsmotivation

Klassifikation der argmax-Invarianz: LogitsProcessor.is_argmax_invariant (interface.py:84-92) unterscheidet, ob das greedy-Ergebnis verändert wird; MinP, MinTokens und LogitBias fallen jeweils auf unterschiedliche Seiten. Greedy-Requests überspringen die gesamte random-Pipeline nach den argmax-invarianten Prozessoren (sampler.py:257-271).
Temperatur 0 als greedy betrachtet: apply_temperature setzt Positionen mit temp < eps auf 1.0 (sampler.py:233-237); am Ende erfolgt eine einzige Zusammenführung über torch.where(temp < eps, greedy, random), um Verzweigungen zu vermeiden.
In-place bevorzugt: apply_temperature nutzt logits.div_, MinP.apply nutzt masked_fill_ (builtin.py:115), um den Speicher-Peak zu senken.
FlashInfer-Sampler ist schnell, aber eingeschränkt: Er gibt die gefilterten logits nicht zurück (topk_topp_sampler.py:86-95); wenn logprobs_mode processed_logits oder processed_logprobs ist, wird erzwungen auf forward_native zurückgefallen.
Logprobs-Erfassung und Sampling getrennt: gather_logprobs (sampler.py:308-356) läuft nach dem Sampling und berechnet log_softmax auf den ursprünglichen (nicht temperatur-skalierten) logits, damit Penalty die an den Nutzer zurückgegebenen Logprob-Werte nicht beeinflusst.
Schutz gegen dynamische Batch-Spezialisierung: gather_logprobs nutzt torch._dynamo.decorators.mark_unbacked (sampler.py:345-346), damit sich dynamo bei batch_size=1 → ≥2 nicht neu kompiliert.
Kompatibilität mit Spec-Decode: apply_logits_processors hängt beim predict_bonus_token die spec_token_ids an output_token_ids an (sampler.py:384-393), damit die Penalty die Draft-Token sieht.

Schlüsseldateien

Sampler Klasse:20 — nn.Module, hält TopKTopPSampler und implementiert die 9-stufige Pipeline.
forward docstring:22-59 — Spezifikation der 9-stufigen Pipeline; jeder Schritt ist kommentiert.
Sampler.forward:72-149 — Haupteinstieg; verbindet apply_logits_processors + sample + gather_logprobs.
Sampler.sample:243 — Aufteilung greedy / random, Anwendung von Temperatur + top-k/top-p.
apply_temperature:227-237 — in-place div_; setzt Temperatur 0 auf 1.0, um Division durch null zu vermeiden.
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 — Delegiert an das Op apply_all_penalties.
SamplingMetadata:14 — Hält temperature/top_p/top_k/generators/penalties und weitere pro-Batch-Sampling-Metadaten.
TopKTopPSampler:70 — Plattform-Dispatch: CUDA→FlashInfer, CPU→native, XPU→kernel, ROCm→aiter.
flashinfer_sample:471-508 — Ruft FlashInfer top_p_renorm_probs / top_k_sampling_from_probs auf; statistisch äquivalent, aber schneller als Rejection.
GPUModelRunner hält Sampler:532 — Im Konstruktor Sampler(logprobs_mode, use_fp64_gumbel).

Datenfluss

Die Eingabe jedes Schritts sind logits der Form [num_tokens, vocab_size], wobei num_tokens die Gesamtzahl der Token ist, die in diesem Schritt für die gesamte Charge an Requests berechnet werden (bei Prefill mehrere Prompt-Token, bei Decode pro Request einer). Im Folgenden die Aufteilung in greedy und random innerhalb der sample-Methode:

python

# 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 sind vorab berechnete Batch-Flags in SamplingMetadata, die der gesamten Charge erlauben, direkt den fast path zu nehmen. Bei gemischten Batches werden beide Pfade berechnet und am Ende über torch.where anhand der per-Request-Temperatur ausgewählt — so wird die GPU nicht durch Verzweigungen ausgebremst.

Grenzen und Fehler

all_greedy und all_random sind exklusiv: Am Anfang von sample steht assert not (all_greedy and all_random) (sampler.py:256); wenn beide gleichzeitig konfiguriert sind, wird direkt ein AssertionError geworfen.
Token mit Temperatur 0 geht nicht über random: torch.where(temp < eps, greedy, random) (sampler.py:296-301) stellt sicher, dass greedy-Requests immer argmax erhalten und nicht durch random verunreinigt werden.
Logprobs verwenden die ursprünglichen logits: compute_logprobs wird vor apply_temperature berechnet (sampler.py:85-93); Penalty und Temperatur beeinflussen die an den Nutzer zurückgegebenen Logprob-Werte nicht — das ist ein Unterschied zu v0.
logprob_token_ids mit speziellem gather: gather_specific_token_logprobs (sampler.py:151-225) wird in der generative_scoring-API verwendet; es überspringt Top-k und gather direkt die angegebenen Token, mit Padding und valid_mask.
Bei Spec-Decode sind Logits-Prozessoren deaktiviert: build_logitsprocs wirft ValueError, wenn speculative_config gesetzt und ein benutzerdefinierter Processor übergeben wird (__init__.py:201-209); nur MinTokens bleibt erhalten.
FlashInfer gibt int32 aus: Am Ende von forward wird sampled = sampled.long() (sampler.py:109) ausgeführt, um auf int64 zu vereinheitlichen, da der FlashInfer-Sampler int32 zurückgibt, während PyTorch argmax / topk int64 liefert; nachfolgende Index-Operationen müssen kompatibel sein.

Zusammenfassung

Sampler ist der einheitliche Eingang für das Sampling in vLLM: eine neunstufige Pipeline + Aufteilung in greedy/random + Beschleunigung über FlashInfer. Es verlässt sich auf die LogitsProcessor-Abstraktion aus /sampling/logits für parametrisierte Transformationen und kümmert sich selbst nur um die eine Sache: „wie werden logits zu Token". Der zurückgegebene SamplerOutput wird nach oben an /worker/gpu-model-runner weitergegeben und schließlich über den OutputProcessor zum RequestOutput zusammengefügt. Wie der Engine-Core diese Schritte schedult, siehe /engine/engine-core.

Siehe offizielle Dokumentation: vLLM-Dokumentation · README.

Sampler: die letzte Meile von logits zum Token ​

Verantwortung ​

Entwurfsmotivation ​

Schlüsseldateien ​

Datenfluss ​

Grenzen und Fehler ​

Zusammenfassung ​