Sampling on the GPU

Run Llama-3-8B on an H100 at batch 256. The final lm_head projection emits logits of shape [256, 128256]. In fp32 — and both engines upcast to fp32 before sampling — that is 131.3 MB. Now suppose you did the obvious thing and sampled on the CPU: copy the logits down, apply each request's temperature and top-p in Python, pick a token, copy the ids back.

The copy alone is fatal. 131.3 MB over PCIe Gen5 x16 at a generous 64 GB/s is 2.05 ms. §0.4 put the bandwidth floor of a Llama-3-8B bf16 decode step at 4.48 ms — the time to stream 15.01 GB of weights through 3.35 TB/s of HBM. So a one-way device-to-host copy of the logits is 46% of the entire step's floor, before the CPU has sorted a single row and before the token ids go back up. Worse, the copy is a synchronisation point: the GPU must finish, the host must work, the GPU must be re-fed. That destroys the overlap the scheduler depends on (§1.3) and makes the step un-capturable as a CUDA graph (§8.1).

A full FP32 read/write sweep has an ideal HBM floor of about 78 microseconds, roughly 1.75% of the cited weight-streaming floor. These are lower-bound traffic models, not measured step fractions. Sorting, reductions, RNG, transcendental instructions, occupancy and launches also matter; reducing passes is one important design lever.

Think of the logits as a very wide, very short matrix: $B$ rows of $V = 128{,}256$ float32s. Every sampling operator is a row-wise reduction or mask over that matrix. Temperature divides each row by a scalar. Penalties subtract from selected columns. Top-k and top-p delete all but a prefix of the row's sorted order. The multinomial draw reduces each row to one index.

A temperature sweep has low arithmetic intensity: roughly one scalar operation per eight logical bytes for FP32 read/write. Sorting, reductions and random draws have different instruction and synchronization costs. Bytes divided by peak bandwidth is a lower bound under the assumed traffic, not exact runtime; scalar instructions should not be compared only to a tensor-core peak. Fusion, selection algorithms, occupancy and launch reduction all matter.

Let $z \in \mathbb{R}^{V}$ be one row of logits, $V = 128{,}256$ for Llama-3. Define $p = \mathrm{softmax}(z)$. The operators, in the order both engines compose them:

Temperature. Divide logits by positive T. Increasing T flattens the distribution. As T approaches zero, mass becomes uniform over exactly tied maxima; only a unique maximum yields a point mass. The separate greedy path may have a different tie rule.

Penalties. These are the only operators that read the request's history, not just the current step. Let $c_v$ be the count of token $v$ in the generated output so far, and $m^{\text{prompt}}_v, m^{\text{out}}_v \in \{0,1\}$ indicate presence in prompt and output.

$\rho$ is the repetition penalty (multiplicative, sign-aware, neutral at 1.0), $\alpha_f$ the frequency penalty (scales with count), $\alpha_p$ the presence penalty (fires once). vLLM implements exactly this, and the sign-aware branch is the reason repetition penalty is a custom kernel rather than a multiply:

    num_seqs, vocab_size = logits.shape
    _, prompt_mask = get_token_bin_counts_and_mask(
        prompt_tokens_tensor, vocab_size, num_seqs
    )
    output_bin_counts, output_mask = get_token_bin_counts_and_mask(
        output_tokens_tensor, vocab_size, num_seqs
    )

    # Apply repetition penalties as a custom op
    from vllm._custom_ops import apply_repetition_penalties

    apply_repetition_penalties(logits, prompt_mask, output_mask, repetition_penalties)

    # We follow the definition in OpenAI API.
    # Refer to https://platform.openai.com/docs/api-reference/parameter-details
    logits -= frequency_penalties.unsqueeze(dim=1) * output_bin_counts
    logits -= presence_penalties.unsqueeze(dim=1) * output_mask
    return logits

Note the cost hiding in get_token_bin_counts_and_mask: it allocates a [num_seqs, vocab_size + 1] int64 buffer and scatter-adds the whole history into it. At batch 256 that is $256 \times 128{,}257 \times 8 = 262.7$ MB per call, twice (prompt and output) — and the source says so itself, at vllm/v1/sample/ops/penalties.py:L27-L28: "The penalties implementation is currently quite inefficient and will be reworked anyhow." The rework exists at this SHA in the V2 model runner's own sampler stack, whose Triton penalty kernel is quoted in §6.1.4.

Top-k. Keep the $k$ largest logits, mask the rest to $-\infty$. Purely order-based; no probability arithmetic needed.

Top-p (nucleus). Sort descending, keep the shortest prefix whose cumulative probability reaches $p$. Formally, with $\pi$ the descending permutation, keep ranks $r \le r^\ast$ where $r^\ast = \min\{ r : \sum_{j \le r} p_{\pi(j)} \ge p \}$. Unlike top-k, the number kept is data-dependent: a confident step may keep one token, an ambiguous one may keep hundreds.

Min-p. Keep tokens whose probability is at least a fixed fraction of the most likely token's: keep $v$ where $p_v \ge \mu \cdot \max_j p_j$. Because it is relative to the peak, it adapts automatically — sharp distributions keep few tokens, flat ones keep many. The V2 runner's vLLM kernel does it in log space, which avoids the softmax entirely: threshold = max_val + tl.log(min_p) at vllm/v1/worker/gpu/sample/min_p.py:L34.

The draw. Given the masked row, sample $v \sim \mathrm{softmax}(z)$. vLLM avoids torch.multinomial in the hot path because it synchronises with the host. The standard trick is the Gumbel-max / exponential-race identity: if $q_v \sim \mathrm{Exp}(1)$ independently, then $\arg\max_v p_v / q_v$ is distributed exactly as a draw from $p$. vLLM's random_sample says so in as many words:

def random_sample(
    probs: torch.Tensor,
    generators: dict[int, torch.Generator],
    use_fp64_gumbel: bool = False,
) -> torch.Tensor:
    """Randomly sample from the probabilities.

    We use this function instead of torch.multinomial because torch.multinomial
    causes CPU-GPU synchronization.
    """
    q = empty_exponential_noise_like(probs, use_fp64_gumbel)
    # NOTE(woosuk): To batch-process the requests without their own seeds,
    # which is the common case, we first assume that every request does
    # not have its own seed. Then, we overwrite the values for the requests
    # that have their own seeds.
    if len(generators) != probs.shape[0]:
        q.exponential_()
    if generators:
        # TODO(woosuk): This can be slow because we handle each request
        # one by one. Optimize this.
        for i, generator in generators.items():
            q[i].exponential_(generator=generator)

SGLang only half-agrees. Its FlashInfer backend never reaches a PyTorch multinomial, and its seeded path uses the hashed-Gumbel argmax below — but its unseeded pytorch backend calls torch.multinomial outright, both in the no-filter case (sampling_from_probs_torch) and after the top-k/top-p sort at python/sglang/srt/layers/sampler.py:L597-L598. So "nobody calls torch.multinomial" is a vLLM property, not a universal one.

Order matters, and the engines do not fully agree

vLLM's Sampler docstring is the canonical statement of the order — masks, then penalties, then temperature, then min-p, then top-k/top-p, then the draw (vllm/v1/sample/sampler.py:L21-L59). SGLang's PyTorch backend applies temperature, then top-k, then top-p, then min-p, all on one descending sort:

    probs_sort, probs_idx = probs.sort(dim=-1, descending=True)
    probs_sum = torch.cumsum(probs_sort, dim=-1)
    probs_sort[
        torch.arange(0, probs.shape[-1], device=probs.device).view(1, -1)
        >= top_ks.view(-1, 1)
    ] = 0.0
    probs_sort[(probs_sum - probs_sort) > top_ps.view(-1, 1)] = 0.0

Two consequences fall out of reading these side by side. First, the top-p rule itself agrees. vLLM sorts ascending and masks where the inclusive cumulative mass from the bottom is $\le 1 - p$; SGLang sorts descending and masks where the exclusive prefix mass exceeds $p$. For a token at descending rank $r$, vLLM keeps it iff $1 - \text{excl}(r) > 1 - p$, i.e. $\text{excl}(r) < p$; SGLang keeps it iff $\text{excl}(r) \le p$. The rules differ only on exact equality, and both force the top token to survive — vLLM with an explicit top_p_mask[:, -1] = False, SGLang because $\text{excl}(1) = 0$.

Second, when a request sets both min_p and top_p the two engines can legitimately keep different sets. vLLM masks with min-p first, so the softmax inside its top-p step renormalises over the survivors and reaches cumulative $p$ sooner — keeping fewer tokens. SGLang computes top-p against the full, unfiltered distribution and applies min-p afterwards. This is an arithmetic consequence of the two code paths, not a measurement; it is exactly the kind of divergence that makes "same params, same seed, same output" a false expectation across engines.

A production batch is not one workload. Row 0 is a code-completion request at $T=0$. Row 1 is creative writing at $T=1.1$, $p=0.95$. Row 2 is a JSON extraction at $T=0.2$, $k=40$, with a seed. Row 3 wants none of it. The kernel cannot branch on request identity — there is no request identity inside a CUDA block, only blockIdx.x. Every parameter must arrive as a tensor indexed by row.

vLLM maintains those tensors as a persistent batch: pinned CPU staging arrays of length max_num_seqs written by index when a request is admitted, and matching device tensors copied from them. Admission writes the row:

        if sampling_params := request.sampling_params:
            if sampling_params.sampling_type == SamplingType.GREEDY:
                # Should avoid division by zero later when apply_temperature.
                self.temperature_cpu[req_index] = 0.0
                self.greedy_reqs.add(req_id)
            else:
                self.temperature_cpu[req_index] = sampling_params.temperature
                self.random_reqs.add(req_id)

            self.top_p_cpu[req_index] = sampling_params.top_p
            if sampling_params.top_p < 1:
                self.top_p_reqs.add(req_id)
            top_k = sampling_params.top_k
            if 0 < top_k < self.vocab_size:
                self.top_k_reqs.add(req_id)
            else:
                top_k = self.vocab_size
            self.top_k_cpu[req_index] = top_k

Two design decisions are doing all the work here. Neutral encoding: a request that does not want top-k is stored as $k = V$, which is a no-op the kernel can execute unconditionally. Same for $p = 1.0$, $\rho = 1.0$, $\alpha = 0$. There is no "disabled" branch — only a value that happens to do nothing. Set-based gating: alongside the array, vLLM keeps a set of request ids that actually use each feature. no_top_p and no_top_k are just len(...) == 0 (vllm/v1/worker/gpu_input_batch.py:L1126-L1132), and the metadata build skips both the H2D copy and the tensor entirely when the set is empty:

    def _make_sampling_metadata(self) -> SamplingMetadata:
        num_reqs = self.num_reqs
        if not self.all_greedy:
            temperature = copy_slice(
                self.temperature_cpu_tensor, self.temperature, num_reqs
            )
        else:
            temperature = None
        if not self.no_top_p:
            copy_slice(self.top_p_cpu_tensor, self.top_p, num_reqs)
        if not self.no_top_k:
            copy_slice(self.top_k_cpu_tensor, self.top_k, num_reqs)

The result is a SamplingMetadata whose top_p and top_k are either a [num_reqs] tensor or None (vllm/v1/sample/metadata.py:L20-L21) — and None means the entire kernel is skipped. This is how a batch where nobody uses top-k avoids paying for it. It is not how a batch where one row uses top-k avoids penalising the other 255: that row forces the tensor to exist, and the kernel runs over all rows with $k = V$ for the rest.

The V2 model runner — the default for a dense model, per the callout in §6.1.3 — closes that gap by pushing the per-row check into the kernel and returning early before touching the logits at all:

    token_idx = tl.program_id(0).to(tl.int64)
    req_state_idx = tl.load(expanded_idx_mapping_ptr + token_idx)
    rep_penalty = tl.load(repetition_penalty_ptr + req_state_idx)
    freq_penalty = tl.load(frequency_penalty_ptr + req_state_idx)
    pres_penalty = tl.load(presence_penalty_ptr + req_state_idx)

    use_rep_penalty = rep_penalty != 1.0
    use_freq_penalty = freq_penalty != 0.0
    use_pres_penalty = pres_penalty != 0.0
    use_penalty = use_rep_penalty or use_freq_penalty or use_pres_penalty
    if not use_penalty:
        # Early return to avoid loading logits.
        return

That is the whole trick, and it is worth stating plainly: the branch is per row, not per request, and it is taken before the loads. A block whose row has no penalty exits having read three floats. The same pattern appears in the temperature kernel (if temperature == 0.0 or temperature == 1.0: return, vllm/v1/worker/gpu/sample/gumbel.py:L29-L31) and the min-p kernel. Because the logits row is 513 KB and the parameter is 4 bytes, an early return is five orders of magnitude cheaper than the work it skips.

SGLang builds the same tensors but from a list comprehension over the batch each time the batch composition changes, and keeps the per-feature gates as booleans on the struct:

            sampling_seed=sampling_seed,
            is_all_greedy=all(r.sampling_params.top_k <= 1 for r in reqs),
            is_any_greedy=any(r.sampling_params.top_k <= 1 for r in reqs),
            need_top_p_sampling=any(r.sampling_params.top_p != 1.0 for r in reqs),
            need_top_k_sampling=any(r.sampling_params.top_k != TOP_K_ALL for r in reqs),
            need_min_p_sampling=any(r.sampling_params.min_p > 0 for r in reqs),

Under continuous batching the row set changes constantly, so both engines need index surgery. vLLM swaps and condenses rows in place (swap_states, condense) and reports the moves to logits processors via a BatchUpdate record carrying removed, added, moved lists (vllm/v1/sample/logits_processor/interface.py:L36-L57). SGLang instead gathers with an index tensor — filter_batch does setattr(self, item, value[keep_indices_device]) for each parameter tensor (python/sglang/srt/sampling/sampling_batch_info.py:L307-L317) and merge_batch concatenates. The tradeoff is familiar: vLLM's in-place slots keep tensor addresses stable, which matters for CUDA graph capture; SGLang's gather is simpler and reallocates.

Top-p needs an ordering; top-k needs only a threshold. That asymmetry drives everything.

A full torch.sort over the vocabulary moves 4-byte keys plus 8-byte indices: 98.5 MB of payload at batch 64, or about 197 MB for a full read plus write. Actual radix traffic varies by implementation. vLLM's PyTorch fallback sorts and its docstring is candid about it: "If a top-p is used, this function will sort the logits tensor, which can be slow for large batches" (vllm/v1/sample/ops/topk_topp_sampler.py:L374-L378).

For top-k alone you can skip the sort:

def apply_top_k_only(logits: torch.Tensor, k: torch.Tensor) -> torch.Tensor:
    """
    Apply top-k mask to the logits.

    This implementation doesn't involve sorting the entire vocab.
    Note however that it involves a GPU->CPU sync which can be detrimental for
    async scheduling performance.

    The logits tensor may be updated in-place.
    """
    no_top_k_mask = k == logits.shape[1]
    # Set non-top-k rows to 1 so that we can gather.
    k = k.masked_fill(no_top_k_mask, 1)
    max_top_k = k.max()
    # topk.values tensor has shape [batch_size, max_top_k].
    # Convert top k to 0-based index in range [0, max_top_k).
    k_index = k.sub_(1).unsqueeze(1)
    top_k_mask = logits.topk(max_top_k, dim=1).values.gather(1, k_index.long())
    # Handle non-topk rows.
    top_k_mask.masked_fill_(no_top_k_mask.unsqueeze(1), -float("inf"))

This is heterogeneity handled by taking the max: run one topk at $k_{\max} = \max_i k_i$, then gather each row's own cutoff value out of the result. The whole batch pays for the greediest row. And k.max() is a device-to-host read, which is why this path is gated behind allow_cpu_sync and only used on CPU at this SHA.

The path that actually runs on an H100 is neither. vLLM's default is a Triton kernel that does pivot-based truncation: estimate mean and standard deviation from one 8192-wide sample tile, use a lookup table to convert $k/V$ into a Gaussian $\sigma$ cutoff, gather the outliers above that pivot into a per-program scratch buffer, then ternary-search the pivot until exactly $k$ survive (vllm/v1/sample/ops/topk_topp_triton.py:L127-L200). Nothing is sorted; the row is swept in tiles. The launch config is one program per SM, capped by batch:

    # Smaller tiles compile and run faster on CPU; GPU benefits from larger tiles.
    # On XPU, large BLOCK_SIZE causes precision loss in the single-pass pivot
    # approximation; use smaller tiles for accurate top-p results.
    launch_kwargs = {}
    if logits.device.type == "cpu":
        block_size, block_size_trunc = 256, 128
    elif logits.device.type == "xpu":
        block_size, block_size_trunc = 4096, 2048
    else:
        block_size, block_size_trunc = 8192, 4096
        # Each program serially sweeps the vocab row in BLOCK_SIZE tiles, so
        # per-tile latency bounds kernel latency, and Triton's default of 4
        # warps leaves an 8192-wide tile at 16 elements per lane. 8 warps is
        # faster on every arch measured (SM90, SM100, SM120, gfx950); 16 is not.
        launch_kwargs["num_warps"] = 8

The dispatcher picks it only when the batch is wide enough to fill the machine — if HAS_TRITON and logits.shape[0] >= 8: return apply_top_k_top_p_triton(...), otherwise the PyTorch sort (vllm/v1/sample/ops/topk_topp_sampler.py:L355-L364). Below 8 rows there are not enough programs to hide the serial sweep; above it the sort's bandwidth dominates. The third option, FlashInfer, avoids the sort with a rejection loop over a shrinking pivot and is used when available and no per-request generator is present.

The prefill/decode kernel split, concretely

csrc/libtorch_stable/sampler.cu contains two separately-instantiated top-k-per-row kernels with different shapes of input. The prefill one takes explicit per-row [rowStarts, rowEnds]:

template <int kNumThreadsPerBlock, bool useRadixSort>
static __global__ __launch_bounds__(kNumThreadsPerBlock) void topKPerRowPrefill(
    const float* logits, const int* rowStarts, const int* rowEnds,
    int* outIndices, int stride0, int stride1, const int topK,
    const int offsetIndex) {
  // The number of bins in the histogram.
  static constexpr int kNumBins = 2048;

  // The row computed by this block.
  int rowIdx = blockIdx.x + offsetIndex;

  // The range of logits within the row.
  int rowStart = rowStarts[rowIdx];
  int rowEnd = rowEnds[rowIdx];

The decode one instead derives the row's live length from a sequence-length tensor and the speculative fan-out next_n:

static __global__ __launch_bounds__(kNumThreadsPerBlock) void topKPerRowDecode(
    const float* logits, const int* seqLens, int* outIndices, int stride0,
    int stride1, const int topK, int next_n, int seqLensIs2D = 0,
    float* outLogits = nullptr, const int numBlocksToMerge = 0,
    const int* indices = nullptr) {
  // The number of bins in the histogram.
  static constexpr int kNumBins = 2048;

  // The row computed by this block.
  int rowIdx = blockIdx.x;

  // The range of logits within the row.
  int rowStart = 0;
  int batch_idx = rowIdx / next_n;
  int next_n_idx = rowIdx % next_n;

The launch configs diverge more than the kernels do. Decode branches on the row width: insertion sort below 12,288 columns, radix sort between there and 200,000, and a split-and-merge two-kernel scheme above that (csrc/libtorch_stable/sampler.cu:L717-L751). Prefill branches on the row count instead, running the first 12,288 blocks with the insertion-sort instantiation and any excess rows with the radix one:

  constexpr int kSortingAlgorithmThreshold = 12288;
  constexpr int kNumThreadsPerBlock = 512;
  const torch::stable::accelerator::DeviceGuard device_guard(
      logits.get_device_index());
  const cudaStream_t stream = get_current_cuda_stream();

  int numInsertionBlocks =
      std::min(static_cast<int>(numRows), kSortingAlgorithmThreshold);
  vllm::topKPerRowPrefill<kNumThreadsPerBlock, false>
      <<<numInsertionBlocks, kNumThreadsPerBlock, topK * sizeof(int32_t),
         stream>>>(logits.const_data_ptr<float>(),
                   rowStarts.const_data_ptr<int>(),
                   rowEnds.const_data_ptr<int>(),
                   indices.mutable_data_ptr<int>(), static_cast<int>(stride0),
                   static_cast<int>(stride1), static_cast<int>(topK), 0);

That is the shape argument in miniature. Prefill produces many short rows — one per token in the chunk, thousands of them, each covering only that token's causal window, so the grid is huge and per-row work is small. Decode produces few very long rows — one per sequence, spanning the whole context, so the grid is tiny and each block must be able to split its row across multiple blocks and merge. One kernel with runtime branches would carry both register footprints and get the worse occupancy of the two; two instantiations of a shared topKPerRowJob device function get each shape its own launch bounds and shared-memory budget.

Greedy is a separate path, not $T \to 0$

Both engines refuse to divide by zero, but they encode greedy differently. vLLM makes it a sampling type: SamplingParams.sampling_type returns GREEDY when temperature < 1e-5 (vllm/sampling_params.py:L745-L750), and __post_init__ then wipes the other knobs so they cannot interact:

        if self.temperature < _SAMPLING_EPS:
            # Zero temperature means greedy sampling.
            self.top_p = 1.0
            self.top_k = 0
            self.min_p = 0.0
            self._verify_greedy_sampling()

The temperature tensor still gets a 0.0 in that row, and apply_temperature guards it: temp = torch.where(temp < _SAMPLING_EPS, 1.0, temp) (vllm/v1/sample/sampler.py:L235-L238). Then sample() computes both answers and selects per row with torch.where(temperature < _SAMPLING_EPS, greedy_sampled, random_sampled) (vllm/v1/sample/sampler.py:L298-L303) — unless the batch is homogeneous, in which case the all_greedy / all_random flags let it skip one branch entirely.

SGLang normalises greedy away at the parameter level instead:

        # Process some special cases
        if 0 <= self.temperature < _SAMPLING_EPS:
            # top_k = 1 means greedy sampling
            self.temperature = 1.0
            self.top_k = 1
        if self.top_k == -1:
            self.top_k = TOP_K_ALL  # whole vocabulary

TOP_K_ALL is 1 << 30 — the same neutral-value trick, with a sentinel larger than any vocabulary. Greedy becomes $k=1$, and is_all_greedy is all(top_k <= 1), which takes a plain torch.argmax. The encodings are semantically equivalent but not observationally identical: in SGLang a request that explicitly sets top_k=1 with temperature=0.7 counts as greedy for that predicate, so a batch in which every request does this takes the greedy branch, where the returned logprobs come from log_softmax(logits) with no temperature division (python/sglang/srt/layers/sampler.py:L143-L145). Note the gate is a whole-batch flag, not a per-row one — the same request in a batch containing one random row takes the sampling branch instead, and gets its temperature applied.

A shared global RNG is the wrong primitive

If two seeded requests share a batch, drawing from one global generator makes each request's output depend on who else was in the batch — which under continuous batching is nondeterministic by construction. Reproducibility requires the noise for row $i$ to be a pure function of $(\text{seed}_i, \text{position}_i, \text{token id})$ and nothing else.

vLLM's V1 path holds a dict[int, torch.Generator] keyed by batch index (vllm/v1/sample/metadata.py:L23), created once per request from its seed (vllm/v1/worker/gpu_model_runner.py:L1331-L1338), and fills the exponential noise row-by-row in a Python loop — the TODO(woosuk) in the quote above. It is correct and it is slow, and it is also why FlashInfer gets disabled the moment any seeded request appears: "FlashInfer 0.2.3+ does not support per-request generators. Falling back to PyTorch-native implementation." (vllm/v1/sample/ops/topk_topp_sampler.py:L164-L170).

Both projects have converged on the same fix: derive the noise from a hash instead of a stream. SGLang's multinomial_with_seed hashes (seed, position, column index) with murmur3 and turns the result into Gumbel noise in float64:

    n, m = logprobs.shape
    seed = seed.to(torch.uint64)
    col_indices = torch.arange(m, device=logprobs.device)
    hashed = murmur_hash32(seed, positions, col_indices)

    # NOTE (sehoon): it is critical to keep gumbel noise calculation in float64 to avoid numerical instability.
    # keeping logprobs in float64 is less critical, but we found it's still safer to keep it in float64.
    x = hashed.to(torch.float64) / torch.iinfo(torch.uint32).max

    # x is a uniform sample in [0, 1]. get gumbel noise from it.
    # which is equivalent to -log(-log(x))
    # keep everything in in-place operations to avoid unnecessary memory allocations.
    # clamp both ends: x == 1 gives gumbel +inf (NaN at -inf logprobs); the cap is
    # the hash spacing so that bucket matches its neighbor instead of dominating
    x.log_().clamp_(min=torch.finfo(x.dtype).min, max=-(2.0**-32)).neg_()
    x.log_().neg_()  # -log(-log(x)) == gumbel noise

    # add gumbel noise to logprobs
    x.add_(logprobs.to(torch.float64))

    return torch.argmax(x, dim=1, keepdim=True)

vLLM's V2 runner does the same thing with Triton's counter-based RNG, and gives every request a seed — random if the user did not supply one — so there is only one code path (vllm/v1/worker/gpu/sample/states.py:L50-L54). The comment on gumbel_noised_argmax names the second payoff: "`keys` indexes the noise, so the same token draws the same noise wherever it appears; `pos` and `seed` place the draw in the request's stream, which is what lets a draft and its verification agree" (vllm/v1/worker/gpu/sample/gumbel.py:L97-L99) — that is the property §6.2's rejection sampler needs.

Seeding the noise does not make the whole step reproducible: the logits themselves shift with batch composition because reductions reassociate. That is the subject of §10.4; the sampler is one of the two sources it has to pin down. SGLang's --enable-deterministic-inference forces sampling_backend to pytorch for exactly this reason (python/sglang/srt/arg_groups/overrides.py:L2031-L2038).

Batch of four, Llama-3-8B, one decode step. Logits arrive as [4, 128256] bf16. This is the V1 runner's vllm/v1/sample/sampler.py, traced because it is the one whose stages map onto the operator algebra above one-for-one; on a default Llama-3-8B server the V2 runner runs instead (§6.1.3) and fuses several of these steps into the Triton kernels quoted in §6.1.4.

1. GPUModelRunner calls Sampler.forward(logits, sampling_metadata).
2. forward computes raw_logprobs = self.compute_logprobs(logits) first, if any row asked for logprobs — deliberately before penalties and temperature, so returned logprobs describe the model, not the sampler (vllm/v1/sample/sampler.py:L86-L94).
3. logits = logits.to(torch.float32) — 16.4 MB in, 32.8 MB out at batch 64; 2.1 MB at batch 4.
4. apply_logits_processors: allowed-token-ids mask (masked_fill_ with a [max_num_reqs, vocab] bool tensor), apply_bad_words, then every non-argmax-invariant logits processor, then apply_penalties. Each is skipped if its gate says no row needs it.
5. sample(): all_greedy false and all_random false here, so it computes greedy_sampled = logits.argmax(dim=-1) for the whole batch — cheap, one read pass — and keeps it.
6. apply_temperature divides in place, with the torch.where guard on the greedy rows.
7. Argmax-invariant processors run: by default just MinPLogitsProcessor, which returns immediately if self.min_p_count == 0 (vllm/v1/sample/logits_processor/builtin.py:L102-L104).
8. TopKTopPSampler.__call__ dispatches. On CUDA with FlashInfer available and no generators, forward_cuda → flashinfer_sample. Otherwise forward_native → apply_top_k_top_p (Triton at batch ≥ 8, PyTorch sort below) → softmax → random_sample.
9. torch.where(temperature < 1e-5, greedy_sampled, random_sampled) merges the two answers, writing into greedy_sampled to reuse the buffer.
10. If logprobs were requested, gather_logprobs runs torch.topk(logprobs, num_logprobs), gathers the sampled token's own logprob, and computes its rank with batched_count_greater_than — a full (x >= values).sum(-1) over the vocab (vllm/v1/sample/ops/logprobs.py:L25-L27).
11. Return SamplerOutput(sampled_token_ids=sampled.unsqueeze(-1), logprobs_tensors=...), all still on device.

SGLang's path is shorter because more of it is fused into SamplingBatchInfo: ModelRunner.sample → ModelRunner._preprocess_logits (grammar mask, logit bias, accumulated penalties; python/sglang/srt/model_executor/model_runner.py:L1745-L1760) → Sampler.forward, which runs its own _preprocess_logits (custom logit processors, then the env-gated NaN sanitisation above) and then either torch.argmax, or logits.div_(temperatures) then logits[:] = torch.softmax(logits, dim=-1) assigned back into the buffer, then _sample_from_probs (python/sglang/srt/layers/sampler.py:L207-L215). The destination is reused, but the right-hand-side softmax ordinarily allocates a result first. This syntax alone does not prove removal of a full temporary allocation.

Beam search: still present in vLLM, absent from SGLang

Beam search maintains $w$ candidate sequences per request and expands all of them each step. That requires a per-request tree of hypotheses whose KV state forks — and forking KV fights everything Part 2 built. Under paged allocation (§2.2) a fork means block-level copy-on-write; under continuous batching it means $w$ rows appearing and disappearing together; under prefix caching it means the tree's shared prefix is a cache entry with $w$ live readers. It also multiplies the decode batch by $w$ for a single user's throughput. Meanwhile instruction tuning largely removed the payoff: for open-ended generation from an aligned model, the highest-likelihood continuation is usually not the best one.

At a556f3f, vLLM has not removed beam search — it has moved it out of the engine. There is no beam logic in vllm/v1/ at all. Instead vllm/entrypoints/generate/beam_search/ implements it as a client of the ordinary generate API, one token at a time:

        # generate 2 * beam_width candidates at each step
        # following the huggingface transformers implementation
        # at https://github.com/huggingface/transformers/blob/e15687fffe5c9d20598a19aeab721ae0a7580f8a/src/transformers/generation/beam_search.py#L534 # noqa
        base_sampling_params = SamplingParams(
            logprobs=2 * beam_width,
            max_tokens=1,
            temperature=temperature,
            detokenize=False,
            skip_clone=True,  # Internal beam search, safe to skip clone
        )

Every step is a fresh one-token request carrying the full candidate prefix, scored by top-$2w$ logprobs, ranked on the host, and re-submitted. The engine never learns what a beam is; prefix caching does the deduplication that KV forking would have done. BeamSearchParams is a separate struct from SamplingParams (vllm/sampling_params.py:L1246-L1261) precisely because it never reaches the sampler.

SGLang does not implement it at all. A case-insensitive grep for beam across python/sglang/ at 7d89325 returns exactly one hit, and it is a comment in a multimodal generation pipeline telling the caller not to set num_beams.

Logprobs are the most expensive optional feature in the sampler

Returning top-$n$ logprobs costs a log_softmax over the full vocab, a topk over the full vocab, and a rank computation that is another full-vocab reduction — three extra sweeps of the 131 MB tensor at batch 256, roughly 0.16 ms derived, on top of a 4.48 ms step. Worse, it does not fit the "gate on the set being empty" pattern, because both engines take the maximum $n$ across the batch and run one topk at that width: max_k = max(top_logprobs_nums); values, indices = logprobs.topk(max_k, dim=-1) (python/sglang/srt/layers/logprob_processor.py:L93-L94), and topk_logprobs, topk_indices = torch.topk(logprobs, num_logprobs, dim=-1) in vLLM (vllm/v1/sample/sampler.py:L335). One request asking for 20 logprobs makes the whole batch run a top-20. vLLM caps it at max_logprobs: int = Field(default=20, ge=-1) (vllm/config/model.py:L250) for that reason. Both engines also offer an explicit token-id list as a cheaper alternative — a gather instead of a topk.

A mask can change the normalization and therefore the smallest top-p nucleus even when it removes no token from the old nucleus. Here p=0.8 retains the first two tokens before masking; removing the third raises the first token above the threshold. A constant log-normalizer, in contrast, cannot alter Gumbel argmax. This check isolates probability semantics, not a GPU sampler's RNG sequence.

import numpy as np

def nucleus(probs, threshold):
    order = np.argsort(-probs, kind="stable")
    count = np.searchsorted(np.cumsum(probs[order]), threshold) + 1
    return set(order[:count])

p = np.array([0.70, 0.15, 0.15])
masked = p * np.array([True, True, False])
masked /= masked.sum()
assert nucleus(p, 0.8) == {0, 1}
assert nucleus(masked, 0.8) == {0}
noise = np.array([0.2, -0.1, 0.5])
assert np.argmax(np.log(p) + noise) == np.argmax(np.log(p) - 7.0 + noise)
tied = np.array([2.0, 2.0, 0.0])
cold = np.exp((tied - tied.max()) / 0.001)
cold /= cold.sum()
np.testing.assert_allclose(cold, [0.5, 0.5, 0.0])
print("Nucleus, normalization and tied-temperature contracts pass.")

Price the sampler by toggling its backend, holding everything else fixed.

VLLM_USE_FLASHINFER_SAMPLER=0 vllm serve meta-llama/Meta-Llama-3-8B-Instruct --max-num-seqs 256
# then, in another shell
vllm bench serve --model meta-llama/Meta-Llama-3-8B-Instruct \
  --dataset-name random --random-input-len 512 --random-output-len 256 \
  --num-prompts 500 --max-concurrency 256

Compare greedy requests, top-p requests, and identical top-p requests with explicit seeds. Record runner V1/V2 and actual kernel names. Disabling FlashInfer does not force a PyTorch sort; Triton selection and counter-based V2 sampling are alternatives described above. A seed does not universally force a Python generator loop. Keep request/length distributions fixed and inspect the trace before attributing a TPOT difference to one sampler operation.

python -m sglang.launch_server --model-path meta-llama/Meta-Llama-3-8B-Instruct --sampling-backend flashinfer
python -m sglang.launch_server --model-path meta-llama/Meta-Llama-3-8B-Instruct --sampling-backend pytorch
python -m sglang.launch_server --model-path meta-llama/Meta-Llama-3-8B-Instruct --enable-deterministic-inference

Run these long-lived servers one at a time, stopping each before reusing the device/port. The deterministic-inference mode can change model kernels as well as sampling, so its latency versus the PyTorch sampler alone does not isolate hash cost. Record actual dispatch and profile with the Lab 10 methodology; neither a fixed sampler percentage nor one expected bottleneck is established by these commands.

• Holtzman et al., "The Curious Case of Neural Text Degeneration" (2019) — arXiv:1904.09751. The nucleus-sampling paper; the motivation for top-p over top-k is the whole first half.
• Park et al., "Qrita: High-performance Top-k and Top-p Algorithm for GPUs using Pivot-based Truncation and Selection" — arXiv:2602.01518, cited by name in the header of vllm/v1/sample/ops/topk_topp_triton.py as the algorithm that kernel implements.
• vLLM PR #26987 — github.com/vllm-project/vllm/pull/26987, referenced in topk_topp_sampler.py:L106-L107 for the PowerPC argmax-under-torch.compile fallback. A short, concrete example of how platform-specific sampler bugs get fixed.
• FlashInfer sampling kernels — github.com/flashinfer-ai/flashinfer. The rejection-based top-k/top-p sampler both engines call into; read top_k_top_p_sampling_from_probs for the loop that avoids the sort.
• OpenAI parameter details — platform.openai.com, the definition vLLM's penalty code links to directly and matches.
• Thinking Machines, "Defeating Nondeterminism in LLM Inference" (2025) — thinkingmachines.ai, the batch-invariance argument that §10.4 develops and that motivates SGLang's --enable-deterministic-inference.