Sampling, structured outputs, and speculative decoding

Active contributors: Cody Yu, Lily Liu, Wentao Ye, Andreas Karatzas, Cyrus Leung.

Purpose

The sampler turns logits into tokens. vLLM's sampler must run inside CUDA graphs, support per-request sampling parameters (temperature, top-p, top-k, penalties, beam search, logprobs), enforce structured-output constraints from a grammar engine, and verify draft tokens proposed by speculative decoding — all without falling off the fast path.

Directory layout

vllm/v1/sample/
├── sampler.py              # The Sampler module (~17 KB)
├── rejection_sampler.py    # Spec-decode verification (~35 KB)
├── thinking_budget_state.py # Reasoning-token budget enforcement (~22 KB)
├── metadata.py             # SamplingMetadata struct
├── ops/                    # Triton kernels: penalties, top-k/p, etc.
└── logits_processor/       # Per-request logit transforms

vllm/v1/structured_output/
├── __init__.py             # StructuredOutputManager
├── backend_xgrammar.py     # xgrammar (default)
├── backend_guidance.py     # Guidance
├── backend_outlines.py     # Outlines
├── backend_lm_format_enforcer.py
├── backend_types.py
├── request.py              # Per-request grammar state
└── utils.py

vllm/v1/spec_decode/
├── llm_base_proposer.py    # Big driver class (~82 KB)
├── eagle.py                # EAGLE 2 / EAGLE 3 wiring
├── medusa.py               # Medusa heads
├── dflash.py               # DFlash drafting
├── ngram_proposer.py / ngram_proposer_gpu.py
├── suffix_decoding.py
├── draft_model.py          # Standalone draft model
├── extract_hidden_states.py
├── metrics.py              # SpecDecodingStats
├── metadata.py             # SpecDecodeMetadata
└── utils.py

Sampling

graph TD
    L[logits from model.forward]
    LP[Per-request LogitsProcessors<br/>(penalties, structured-output mask, banlist, ...)]
    PT[Penalties / temperature]
    TK[Top-k / top-p]
    Smp[Categorical sample / argmax]
    LP_out[logprobs (optional)]
    Out[ModelRunnerOutput.sampled_token_ids]

    L --> LP --> PT --> TK --> Smp --> Out
    L --> LP_out

Implementation in vllm/v1/sample/sampler.py::Sampler.forward. SamplingMetadata is built from SamplingParams per request and packed into tensors so a single fused kernel handles the whole batch. Greedy and random paths share most of the pipeline; only the final step differs.

Per-request sampling parameters

vllm/sampling_params.py::SamplingParams (~945 lines) is the public API. Highlights:

• Decoding: temperature, top_p, top_k, min_p, seed, n (parallel sampling), best_of, use_beam_search.
• Penalties: presence_penalty, frequency_penalty, repetition_penalty.
• Stops: stop, stop_token_ids, min_tokens, max_tokens, max_completion_tokens, ignore_eos.
• Logprobs: logprobs, prompt_logprobs.
• Output kind: RequestOutputKind (cumulative / delta / final).
• Tools / structured output: structured_outputs: StructuredOutputsParams (json schema, regex, choice, grammar, json_object, structural_tag).
• Reasoning: min_reasoning_tokens, max_reasoning_tokens, reasoning_budget.
• Misc: truncate_prompt_tokens, prompt_logprobs, output_text_buffer_length.

vllm/sampling_params.py::StructuredOutputsParams is the cross-cutting config that gets routed into the structured-output backend.

Structured outputs

vllm/v1/structured_output/__init__.py::StructuredOutputManager wires together:

• A backend (xgrammar by default, fallbacks to guidance, outlines, lm-format-enforcer).
• Per-request grammar state held in vllm/v1/structured_output/request.py.
• A masking step inside the sampler that zeros out tokens forbidden by the grammar.

Backends:

Backend	File	Notes
xgrammar	backend_xgrammar.py	Default. CUDA-graph-friendly bitmask masking.
guidance	backend_guidance.py	LLGuidance integration
outlines	backend_outlines.py	Outlines compiler
lm-format-enforcer	backend_lm_format_enforcer.py	Vendor-specific enforcer

Selection happens in vllm/sampling_params.py::SamplingParams._validate_structured_output and is recorded in StructuredOutputsParams._backend.

Speculative decoding

V1's spec decode lives entirely under vllm/v1/spec_decode/. Proposers:

Proposer	File	Mechanism
n-gram	ngram_proposer.py / ngram_proposer_gpu.py	Match recent token n-gram against the prompt cache
suffix	suffix_decoding.py	Suffix automaton over the corpus / prompt
draft model	draft_model.py + llm_base_proposer.py	Run a smaller LM in parallel
EAGLE / EAGLE3	eagle.py	Hidden-state-based one-step draft heads
DFlash	dflash.py	Diff-flash speculative variant
Medusa	medusa.py	Multi-head Medusa drafts
MTP (Multi-Token Prediction)	model-side files (*_mtp.py)	Model emits N future tokens per step

Verification is done by RejectionSampler (vllm/v1/sample/rejection_sampler.py), which re-samples the target distribution under the modified Sampling kernel. Stats are surfaced in SpecDecodingStats (vllm/v1/spec_decode/metrics.py).

SpeculativeConfig (vllm/config/speculative.py, ~1,400 lines) is the user-facing knob: --speculative-config '{"method": "eagle", "model": "...", "num_speculative_tokens": 4, ...}'.

Reasoning budget

Modern models (DeepSeek-R1, Qwen-Reasoning, ...) emit reasoning content before the user-visible answer. thinking_budget_state.py enforces:

• min_reasoning_tokens / max_reasoning_tokens per request
• Stops that fire only after reasoning ends (reasoning_ended flag on EngineCoreRequest)

vllm/reasoning/ parsers strip the reasoning markers (e.g., <think>…</think>) from the streamed text.

Key source files

File	Purpose
vllm/sampling_params.py	SamplingParams, StructuredOutputsParams
vllm/v1/sample/sampler.py	The fused sampler
vllm/v1/sample/rejection_sampler.py	Spec-decode verification
vllm/v1/sample/thinking_budget_state.py	Reasoning budget
vllm/v1/structured_output/__init__.py	StructuredOutputManager
vllm/v1/spec_decode/llm_base_proposer.py	Generic LLM-based proposer
vllm/v1/spec_decode/eagle.py	EAGLE wiring
vllm/v1/spec_decode/ngram_proposer.py	n-gram proposer
vllm/config/speculative.py	SpeculativeConfig
vllm/config/structured_outputs.py	StructuredOutputsConfig

Entry points for modification

• New sampling op: add a Triton kernel to vllm/v1/sample/ops/, expose it from the sampler, and add a flag to SamplingParams if user-visible.
• New structured-output backend: subclass the backend interface, register in StructuredOutputManager, list in StructuredOutputsConfig.backend.
• New proposer: implement a class with propose(...) and add a branch in SpeculativeConfig.method.
• New stop condition: extend FinishReason and check_stop (vllm/v1/core/sched/utils.py).