## inference. A common pattern in production are to [[thoughts/PD disaggregated serving#prefill/decode|disaggregate]] prefill and decode nodes such that each deployments are optimized for its own prefill/decode goals/target. ### sampling ```python title="vllm-project/vllm · sampler.py" # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """A layer that samples the next tokens from the model's outputs.""" import torch import torch.nn as nn from vllm.config.model import LogprobsMode from vllm.utils.platform_utils import is_pin_memory_available from vllm.v1.outputs import LogprobsTensors, SamplerOutput from vllm.v1.sample.metadata import SamplingMetadata from vllm.v1.sample.ops.bad_words import apply_bad_words from vllm.v1.sample.ops.logprobs import batched_count_greater_than from vllm.v1.sample.ops.penalties import apply_all_penalties from vllm.v1.sample.ops.topk_topp_sampler import TopKTopPSampler _SAMPLING_EPS = 1e-5 class Sampler(nn.Module): """ A layer that samples the next tokens from the model's outputs with the following steps in order: 1. If logprobs are requested: a) If `logprobs_mode` is `raw_logprobs`, compute logprobs as the final logprobs to return. b) If `logprobs_mode` is `raw_logits`, clone the logits as the final logprobs to return. 2. Convert logits to float32. 3. Apply allowed token ids whitelist. 4. Apply bad words exclusion. 5. Apply logit processors which are not argmax-invariant, i.e. that can impact greedy sampling. a) Min tokens processor b) Logit bias processor 6. Apply penalties a) Repetition penalty b) Frequency penalty c) Presence penalty 7. Sample the next tokens. `sample` method performs the following steps: a) If not `all_random`, perform greedy sampling. If `all_greedy`, return the greedily sampled tokens and final logprobs if requested. b) Apply temperature. c) Apply logit processors which are argmax-invariant, by default the min_p processor. d) Apply top_k and/or top_p. e) Sample the next tokens with the probability distribution. f) If `all_random` or temperature >= epsilon (1e-5), return the randomly sampled tokens and final logprobs if requested. Else, return the greedily sampled tokens and logprobs if requested. 8. Gather the logprobs of the top `max_num_logprobs` and sampled token (if requested). Note that if the sampled token is within the top `max_num_logprobs`, the logprob will be eventually merged in `LogprobsProcessor` during output processing. Therefore, the final output may contain either `max_num_logprobs + 1` or `max_num_logprobs` logprobs. 9. Return the final `SamplerOutput`. """ def __init__( self, logprobs_mode: LogprobsMode = "raw_logprobs", use_fp64_gumbel: bool = False, ): super().__init__() self.topk_topp_sampler = TopKTopPSampler(logprobs_mode, use_fp64_gumbel) self.pin_memory = is_pin_memory_available() self.logprobs_mode = logprobs_mode self.use_fp64_gumbel = use_fp64_gumbel def forward( self, logits: torch.Tensor, sampling_metadata: SamplingMetadata, predict_bonus_token: bool = False, logprobs_mode_override: LogprobsMode | None = None, ) -> SamplerOutput: logprobs_mode = logprobs_mode_override or self.logprobs_mode # NOTE(woosuk): Use the original logits (before any penalties or # temperature scaling) for the top-k logprobs. # This is different from the V0 sampler, which uses the logits that # is used for sampling (after penalties and temperature scaling). num_logprobs = sampling_metadata.max_num_logprobs raw_logprobs: torch.Tensor | None = None if num_logprobs is not None or sampling_metadata.logprob_token_ids: if logprobs_mode == "raw_logprobs": raw_logprobs = self.compute_logprobs(logits) elif logprobs_mode == "raw_logits": if logits.dtype == torch.float32: raw_logprobs = logits.clone() else: raw_logprobs = logits.to(torch.float32) # Use float32 for the logits. logits = logits.to(torch.float32) logits = self.apply_logits_processors( logits, sampling_metadata, predict_bonus_token ) # Sample the next token. sampled, processed_logprobs = self.sample(logits, sampling_metadata) if processed_logprobs is not None: raw_logprobs = processed_logprobs # Convert sampled token ids to int64 (long) type to ensure compatibility # with subsequent operations that may use these values as indices. # This conversion is necessary because FlashInfer sampling operations # return int32 (while PyTorch argmax and topk return int64). sampled = sampled.long() # Handle logprob_token_ids if specified (more efficient than full vocab) # This is used by generative_scoring API to get logprobs for specific tokens logprob_token_ids_tensors = None if sampling_metadata.logprob_token_ids: assert raw_logprobs is not None logprob_token_ids_tensors = self.gather_specific_token_logprobs( raw_logprobs, sampling_metadata.logprob_token_ids, sampled ) if num_logprobs is None: logprobs_tensors = logprob_token_ids_tensors elif num_logprobs == -1: # Return the full unsorted and unranked logprobs. logprobs_tensors = LogprobsTensors( torch.empty(0), raw_logprobs, torch.empty(0) ) else: # Gather the logprobs and ranks of the topk and sampled token. logprobs_tensors = self.gather_logprobs( raw_logprobs, num_logprobs, token_ids=sampled ) # If we have both num_logprobs and logprob_token_ids, prefer # logprob_token_ids as it's more specific if logprob_token_ids_tensors is not None and num_logprobs is not None: logprobs_tensors = logprob_token_ids_tensors # Use int32 to reduce the tensor size. sampled = sampled.to(torch.int32) # These are GPU tensors. sampler_output = SamplerOutput( # The sampled tokens are expanded to 2D tensor with shape # [num_requests, 1], where each row represents one generated # token per request. sampled_token_ids=sampled.unsqueeze(-1), logprobs_tensors=logprobs_tensors, ) return sampler_output def gather_specific_token_logprobs( self, logprobs: torch.Tensor, logprob_token_ids: dict[int, list[int]], sampled: torch.Tensor, ) -> LogprobsTensors | None: """Gather logprobs for specific token IDs requested per request. Used by the generative_scoring API to return logprobs for an explicit set of token ids rather than the top-k. Handles heterogeneous token id lists across requests by padding shorter lists to the max length. Args: logprobs: [batch_size, vocab_size] tensor of (raw) logprobs to gather from. logprob_token_ids: dict mapping req_index -> list of token IDs sampled: [batch_size] tensor of sampled token IDs Returns: LogprobsTensors with logprobs for the specified tokens, or None if no requests have logprob_token_ids. """ if not logprob_token_ids: return None batch_size = logprobs.shape[0] device = logprobs.device # Find max number of tokens across all requests max_num_tokens = max(len(tids) for tids in logprob_token_ids.values()) pin = self.pin_memory # Build the padded token_ids and valid_mask matrices on pinned CPU, # then upload non-blocking. token_ids_cpu = torch.zeros( batch_size, max_num_tokens + 1, dtype=torch.int64, pin_memory=pin ) # Create mask for valid positions (True = valid, False = padded) valid_mask_cpu = torch.zeros( batch_size, max_num_tokens + 1, dtype=torch.bool, pin_memory=pin ) valid_mask_cpu[:, 0] = True # Sampled token is always valid for req_idx, token_ids in logprob_token_ids.items(): num_tokens = len(token_ids) token_ids_cpu[req_idx, 1 : num_tokens + 1] = torch.as_tensor( token_ids, dtype=torch.int64 ) valid_mask_cpu[req_idx, 1 : num_tokens + 1] = True token_ids_tensor = token_ids_cpu.to(device, non_blocking=True) valid_mask = valid_mask_cpu.to(device, non_blocking=True) # Sampled token in column 0 — fill on-device from the sampled GPU # tensor so we don't need to D2H + re-upload. token_ids_tensor[:, 0] = sampled # Gather logprobs at the requested token ids. gathered_logprobs = logprobs.gather(-1, token_ids_tensor) # Mask invalid (padded) positions with -inf gathered_logprobs = gathered_logprobs.masked_fill(~valid_mask, float("-inf")) # Compute ranks for the sampled token. log_softmax is monotonic w.r.t. # the original logits, so ranks computed from logprobs are equivalent. sampled_logprobs = logprobs.gather(-1, sampled.unsqueeze(-1)) # Avoid 0/1 specialization recompile on the batch dimension of the # compiled batched_count_greater_than. See gather_logprobs for context. torch._dynamo.decorators.mark_unbacked(logprobs, 0) torch._dynamo.decorators.mark_unbacked(sampled_logprobs, 0) token_ranks = batched_count_greater_than(logprobs, sampled_logprobs) return LogprobsTensors( logprob_token_ids=token_ids_tensor.to(torch.int32), logprobs=gathered_logprobs, selected_token_ranks=token_ranks, ) @staticmethod def apply_temperature( logits: torch.Tensor, temp: torch.Tensor, all_random: bool, ) -> torch.Tensor: # Use in-place division to avoid creating a new tensor. # Avoid division by zero if there are greedy requests. if not all_random: temp = torch.where(temp < _SAMPLING_EPS, 1.0, temp) return logits.div_(temp.unsqueeze(dim=1)) @staticmethod def greedy_sample(logits: torch.Tensor) -> torch.Tensor: return logits.argmax(dim=-1).view(-1) def sample( self, logits: torch.Tensor, sampling_metadata: SamplingMetadata, logprobs_mode_override: LogprobsMode | None = None, ) -> tuple[torch.Tensor, torch.Tensor | None]: """Sample logits based on sampling metadata. The various logits processing functions called in this method may update the logits tensor in-place. """ 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 @staticmethod def compute_logprobs(logits: torch.Tensor) -> torch.Tensor: return logits.log_softmax(dim=-1, dtype=torch.float32) @staticmethod def gather_logprobs( logprobs: torch.Tensor, num_logprobs: int, token_ids: torch.Tensor, ) -> LogprobsTensors: """ Gather logprobs for topk and sampled/prompt token. Args: logprobs: (num tokens) x (vocab) tensor num_logprobs: maximum number of logprobs to retain per token token_ids: prompt tokens (if prompt logprobs) or sampled tokens (if sampled logprobs); 1D token ID tensor with (num tokens) elements Must be int64. Returns: Top-k int indices tensor, (num tokens) x (num_logprobs + 1) Top-k float logprobs tensor, (num tokens) x (num_logprobs + 1) Sampled token rank tensor, (num tokens) """ assert token_ids.dtype == torch.int64 # Find the topK values. topk_logprobs, topk_indices = torch.topk(logprobs, num_logprobs, dim=-1) # Get with the logprob of the prompt or sampled token. token_ids = token_ids.unsqueeze(-1) token_logprobs = logprobs.gather(-1, token_ids) # Compute the ranks of the actual token. # Avoid 0/1 specialization recompile on the batch dimension # of the compiled batched_count_greater_than. mark_unbacked makes # the size fully symbolic so dynamo doesn't specialize when # batch_size transitions from 1 to >=2. torch._dynamo.decorators.mark_unbacked(logprobs, 0) torch._dynamo.decorators.mark_unbacked(token_logprobs, 0) token_ranks = batched_count_greater_than(logprobs, token_logprobs) # Concatenate together with the topk. indices = torch.cat((token_ids, topk_indices), dim=1) logprobs = torch.cat((token_logprobs, topk_logprobs), dim=1) # Use int32 to reduce the tensor size. indices = indices.to(torch.int32) return LogprobsTensors(indices, logprobs, token_ranks) @staticmethod def _combine_outputs_with_spec_tokens( output_token_ids: list[list[int]], spec_token_ids: list[list[int]] | None = None, ) -> list[list[int]]: if spec_token_ids is None: return output_token_ids return [ [*out, *spec] if spec else out for out, spec in zip(output_token_ids, spec_token_ids) ] def apply_logits_processors( self, logits: torch.Tensor, sampling_metadata: SamplingMetadata, predict_bonus_token: bool, ) -> torch.Tensor: bad_words_token_ids = sampling_metadata.bad_words_token_ids any_penalties_or_bad_words = ( bool(bad_words_token_ids) or not sampling_metadata.no_penalties ) holder = sampling_metadata.thinking_budget_state_holder needs_thinking_combine = holder is not None and holder.has_tracked_requests() output_token_ids = sampling_metadata.output_token_ids if predict_bonus_token and ( any_penalties_or_bad_words or needs_thinking_combine ): # Combine base outputs with spec tokens when speculative decoding # is enabled. output_token_ids = self._combine_outputs_with_spec_tokens( output_token_ids, sampling_metadata.spec_token_ids, ) # Apply allowed token ids. if sampling_metadata.allowed_token_ids_mask is not None: logits.masked_fill_(sampling_metadata.allowed_token_ids_mask, float("-inf")) # Apply bad words exclusion. if bad_words_token_ids: apply_bad_words(logits, bad_words_token_ids, output_token_ids) # Apply logits processors which can impact greedy sampling. for processor in sampling_metadata.logitsprocs.non_argmax_invariant: logits = processor.apply(logits) # Apply penalties (e.g., freq_penalties). logits = self.apply_penalties(logits, sampling_metadata, output_token_ids) if holder is not None and holder.has_tracked_requests(): holder.update_state( output_token_ids, sampling_metadata.spec_token_ids, repeat_indices=None, ) logits = holder.apply_to_logits( logits, predict_bonus_token, sampling_metadata.spec_token_ids, ) return logits @staticmethod def apply_penalties( logits: torch.Tensor, sampling_metadata: SamplingMetadata, output_token_ids: list[list[int]], ) -> torch.Tensor: if sampling_metadata.no_penalties: return logits assert sampling_metadata.prompt_token_ids is not None return apply_all_penalties( logits, sampling_metadata.prompt_token_ids, sampling_metadata.presence_penalties, sampling_metadata.frequency_penalties, sampling_metadata.repetition_penalties, output_token_ids, ) ``` given logits $z_i$ and base distribution $p_i = \text{softmax}(z)_i$, sampling methods modify this distribution before drawing tokens. #### temperature rescales logits before softmax: $p_i^{(\tau)} = \text{softmax}(z/\tau)_i$ $$ \lim_{\tau \to 0^+} p_i^{(\tau)} = \mathbf{1}[i = \arg\max z], \quad \lim_{\tau \to \infty} p_i^{(\tau)} = 1/|\mathcal{V}| $$ entropy $H(p^{(\tau)})$ monotonically increases with $\tau$. ```python def apply_temperature( logits: torch.Tensor, temp: torch.Tensor ) -> torch.Tensor: temp = torch.where(temp < 1e-5, 1.0, temp) return logits.div_(temp.unsqueeze(dim=1)) ``` #### top-k retain $k$ highest-probability tokens, renormalize: $$ p_i^{(k)} = \frac{p_i \cdot \mathbf{1}[i \in \mathcal{V}_k]}{\sum_{j \in \mathcal{V}_k} p_j}, \quad \mathcal{V}_k = \{i : \text{rank}(p_i) \leq k\} $$ fixed $k$ ignores distribution shape—problematic when peaked (includes noise) or flat (truncates options). ```python def apply_top_k(logits: torch.Tensor, k: torch.Tensor) -> torch.Tensor: vocab_size = logits.shape[1] k = k.clamp(min=1, max=vocab_size) top_k_values = logits.topk(k.max(), dim=1).values threshold = top_k_values.gather(1, (k - 1).unsqueeze(1)) return logits.masked_fill_(logits < threshold, -float('inf')) ``` #### top-p (nucleus) \[@holtzman2020curiouscaseneuraltext\] dynamic vocabulary: smallest set with cumulative probability $\geq p$: $$ \mathcal{V}_p = \arg\min_{V \subseteq \mathcal{V}} |V| \quad \text{s.t.} \quad \sum_{i \in V} p_i \geq p $$ requires sorting: $O(|\mathcal{V}| \log |\mathcal{V}|)$. adapts to distribution—peaked → few tokens, flat → many. ```python def apply_top_p(logits: torch.Tensor, p: torch.Tensor) -> torch.Tensor: logits_sort, logits_idx = logits.sort(dim=-1, descending=False) probs_sort = logits_sort.softmax(dim=-1) probs_cumsum = probs_sort.cumsum(dim=-1) mask = probs_cumsum <= (1 - p.unsqueeze(1)) mask[:, -1] = False logits_sort.masked_fill_(mask, -float('inf')) return logits_sort.scatter(dim=-1, index=logits_idx, src=logits_sort) ``` #### min-p \[@nguyen2025turningheatminpsampling\] probability floor relative to maximum: $\mathcal{V}_{\min p} = \{i : p_i \geq \min_p \cdot \max_j p_j\}$ scales with model confidence—tighter threshold when confident, looser when uncertain. outperforms top-p at high temperatures ($\tau > 1$). [^minp-iclr] [^minp-iclr]: ICLR 2025 oral. addresses top-p failure mode: flat distribution includes low-quality tokens. ```python def apply_min_p(logits: torch.Tensor, min_p: torch.Tensor) -> torch.Tensor: probs = logits.softmax(dim=-1) max_probs = probs.amax(dim=-1, keepdim=True) threshold = max_probs * min_p.unsqueeze(1) return logits.masked_fill_(probs < threshold, -float('inf')) ``` #### pipeline order ```mermaid flowchart LR A["logits (B×V)"] --> B["temperature τ"] B --> C["min_p filter"] C --> D["top_k ∩ top_p"] D --> E["softmax → probs"] E --> F["multinomial sample"] F --> G["token_id"] style B fill:#f9d71c,stroke:#333 style C fill:#87ceeb,stroke:#333 style D fill:#98fb98,stroke:#333 ``` vLLM applies: **temperature → min\_p → top\_k/top\_p → sample** temperature before truncation means flatter distributions include more tokens in top-p. this order maximizes diversity control. ```python def sample(logits: torch.Tensor, params: SamplingParams) -> torch.Tensor: logits = apply_temperature(logits, params.temperature) logits = apply_min_p(logits, params.min_p) logits = apply_top_k(logits, params.top_k) logits = apply_top_p(logits, params.top_p) probs = logits.softmax(dim=-1) return torch.multinomial(probs, num_samples=1) ``` [[thoughts/Speculative decoding#von Neumann acceptance rejection|see also: speculative decoding]] ### KV The core “retrieval” bags that contains all previous stored key-value pair or newly added items. [[thoughts/PD disaggregated serving|Prefill disaggregation]] is pretty interesting in a sense that we can separate prefill stage to a separate nodes \[@qin2024mooncakekvcachecentricdisaggregatedarchitecture\] ![[thoughts/images/mooncake-pd.webp|KV-centric optimization]] ### next-token prediction. Sampling: we essentially look forward K-tokens, and then we sample from the distribution of the next token. ### multi-token prediction. @gloeckle2024betterfasterlarge, also used in [[thoughts/Speculative decoding]] for [[thoughts/DeepSeek|DeepSeek-V3 and DeepSeek-R1]] ![[thoughts/images/MTP-deepseek.webp|MTP implementation in DeepSeek, where they keep causal chain for prediction of each token at each depth]] tl/dr: predict $n$-tokens at once, via shared trunk and n dedicated attention heads [^attention-head] Note that during inference, we only employ _one attention head_ [^attention-head]: @gloeckle2024betterfasterlarge employs $n=4$. The order of the forward and backward in a n-token prediction model with $n=4$ heads of the shared trunk works as follow: ```python z = model.shared(x) d = z.detach() d.requires_grad = False for i in range(n): p = model.heads[i](d) loss(p, y[i]).backward() z.backward() ``` ## Byte-Latent Transformer idea: learn from raw-bytes and skip tokenizer/detokenizer protocol. ## Feynman-Kac Let $\mathcal{V}$ be the vocab of given transformers model, and $\mathcal{S} = \mathcal{V}^{*}$ the set of multi-token strings. Assume $\mathcal{V}$ contains token `EOS` and write $\mathcal{F} \subseteq \mathcal{S}$ for the set of `EOS`-terminated strings. > \[!definition\] Definition 1. _Feynman-Kac Transformer model_ > > is a tuple $(s_{0}, \{M_t\}_{t\ge 1}, \{G_t\}_{t\ge 1})$ where: > > - $s_{0} \in \mathcal{S}$ is an _initial state_, which will take as empty string $\epsilon$ > - $M_t(s_t \mid s_{t-1}, f_\theta)$ is a _Markov kernel_ from $s_{t-1} \in \mathcal{F}^c$ to $s_t \in \mathcal{S}$, parameterised by a transformer network $f_\theta: \mathcal{F}^c \to \mathbb{R}^{\mid \mathcal{V} \mid}$ mapping non-`EOS`-terminated strings to vectors of logits > - $G_t(s_{t-1}, s_t, f_\theta)$ is a _potential function_, mapping a pair $(s_{t-1}, s_t) \in \mathcal{F}^c \times \mathcal{S}$ to a real-valued non-negative score. Goal: generate from distribution $\mathbb{P}$ that reweights Markov chain $\mathbb{M}$ by potential functions $G_t$. We define ==_step-t filtering posteriors_==: $$ P_t(s_t) = \frac{\mathbb{E}_\mathbb{M} \left[ \prod_{i=1}^{t \wedge T} G_i(S_{i-1}, S_i, f_\theta) \cdot [S_t = s_t] \right]}{\mathbb{E}_\mathbb{M} \left[ \prod_{i=1}^{t \wedge T} G_i(S_{i-1}, S_i, f_\theta) \right]} $$ _Given that $T$ is mostly finite_ we can then define _overall posterior_ \[@lew2023sequentialmontecarlosteering\{see 2.2 for examples\}\] $$ \mathbb{P}(s) = \lim_{t \to \infty} \mathbb{P}_t(s) $$"How is LLaMa.cpp possible?"
— Andrej Karpathy (@karpathy) 15 août 2023
great post by @finbarrtimbers https://t.co/yF43inlY87
llama.cpp surprised many people (myself included) with how quickly you can run large LLMs on small computers, e.g. 7B runs @ ~16 tok/s on a MacBook. Wait don't you need supercomputers to work… pic.twitter.com/EIp9iPkZ6x
Algorithm 2 Sequential Monte Carlo Transformer Steering
Input: (# particles), (factor), Feynman-Kac Transformer model
Output: Weighted particle approximation of the posterior
Output: Unbiased estimate of the partition function
Initialize
Initialize for
Initialize
while for some do
for
for do
if then
Set
else
Generate for
Set for
end if
end for
Set normalized weights for and
Set
Set
Set and generate
for do
Set
if then
Set
Set
end if
end for
Set particles
Set particles
end while
return