---
title: 'FlashSampling: Efficient Sampling in LLMs'
url: https://www.emergentmind.com/topics/flashsampling
type: topic
---

# FlashSampling: Efficient Sampling in LLMs

FlashSampling is an exact sampling primitive designed to accelerate large-vocabulary decoding in neural language models by fusing the categorical sampling step directly into the language-model head (LM-head) matrix multiplication (matmul), thereby eliminating the overhead associated with separate post-processing kernels and the materialization of logits tensors in high-bandwidth memory (HBM) [2603.15854]. The approach exploits the mathematical decomposability of the Gumbel-Max trick and implements a highly memory- and bandwidth-efficient sampling strategy suitable for both standard and tensor-parallel transformer decoders.

## 1. High-Level Algorithmic Strategy

In conventional large language model (LLM) decoding, the output token distribution is sampled via a two-stage process: (a) a General Matrix Multiply (GEMM) computes the logits matrix $Y = H W^\top$ of shape $[B, V]$ (with $B$ the batch size and $V$ the vocabulary size), which is then materialized in HBM; (b) a series of additional kernels apply temperature scaling, bias, masking, softmax (or cumulative distribution function), and finally, stochastic sampling.

FlashSampling eliminates extraneous memory traffic and kernel launches by fusion: sampling is implemented as an epilogue to the matmul using an on-chip tile-by-tile approach. The vocabulary is partitioned into $T$ disjoint tiles $\{V_1, \dots, V_T\}$ of size $K_{\rm tile}$. For each tile, partial dot products are computed in-register, noise is injected via the Gumbel-Max trick, and only the per-tile maximizer (score and index) is stored in HBM. A final lightweight reduction kernel gathers these $T$ candidates per row, returning the global sample via $\arg\max$. This construction guarantees exactness (not an approximation), as $\arg\max$ is associative across a partition.

## 2. Mathematical Foundations and Sampling Decomposition

FlashSampling employs the Gumbel-Max trick for exact categorical sampling:

\[
i^\star = \arg\max_{i \in [V]} (\tilde\ell_i + g_i), \quad g_i \overset{\mathrm{i.i.d.}}{\sim} \mathrm{Gumbel}(0,1)
\]
\[
\mathbb{P}(i^\star = i) = \frac{e^{\tilde\ell_i}}{\sum_j e^{\tilde\ell_j}}
\]

Here, $\tilde\ell_i$ denotes the transformed logit for index $i$ (i.e., post temperature/bias/mask). Partitioning $[V]$ into tiles, the global maximizer is computed as the maximum among the tile-local maxima:

\[
M_t = \max_{i \in V_t} (\tilde\ell_i + g_i), \quad i_t^* = \arg\max_{i \in V_t} (\tilde\ell_i + g_i)
\]
\[
i^* = i_{t^*}^*, \quad t^* = \arg\max_t M_t
\]

By the max-stability property of Gumbel distributions, this two-stage procedure is mathematically equivalent to global Gumbel-Max on $[V]$ (see Lemma 4.6 in [2603.15854]).

Hierarchical (grouped) decompositions allow further flexibility. If the vocabulary is split into $m$ groups $\{\mathcal{G}_k\}$, one defines

\[
L_k = \log \sum_{i \in \mathcal{G}_k} e^{\tilde\ell_i}
\]
\[
M_k = \max_{i \in \mathcal{G}_k} (\tilde\ell_i + g_i) \sim \mathrm{Gumbel}(L_k,1)
\]

Sampling can proceed by first selecting a group via $K^* = \arg\max_k M_k$, followed by a within-group argmax. Online merging facilitates $O(1)$-memory sampling over streamed groups, and tensor-parallel schemes treat each device's vocabulary shard as a group, needing only an $O(n)$-sized all-gather.

## 3. Implementation Specifics

To avoid full logits materialization, FlashSampling employs a custom Triton or CUDA kernel:

- Each $[B, K_{\rm tile}]$ subblock is computed in registers/SRAM, perturbed with Gumbel samples, and immediately reduced to a $[B, 1]$ identifier (best score and index).
- Per-row, an array of $T$ tile-local maxima is accumulated. A minimal $\arg\max$ reduction over these yields the sample token for each batch row.
- Deterministic transformations (temperature, bias, or masking by setting logits to $-\infty$) are applied in-tile.
- A counter-based RNG (Philox) indexed by $(b,i)$ generates reproducible Gumbel variates, ensuring thread safety.
- All dot-products and score perturbations use FP32 accumulators for numerical stability, supporting BF16 activations/weights if present.
- The kernel is designed to substitute the standard matmul + “sample_from_logits” pipeline in transformer decoders.

## 4. Performance and Latency Benchmarks

FlashSampling achieves considerable kernel speedup and latency reduction in multiple workflows [2603.15854]:

**Kernel-level speedups (for $D=4096$, $V=151{,}936$):**

| Baseline                         | H100 | H200 | B200 | B300 |
|-----------------------------------|------|------|------|------|
| Multinomial (cuBLAS + softmax + multinomial) | 1.60× | 1.60× | 1.84× | 1.81× |
| FlashInfer top-k/top-p           |      |      |      | 2.52× |
| FlashInfer Gumbel-Max            |      |      | 1.39× | 1.37× |

- At batch size $B=1$, speedups of 1.25× (H100), 1.28× (H200), and 1.46× (B200/B300) versus Multinomial.
- In the latency-sensitive decode regime ($B \le 64$), baseline sampling overhead is 12–29% of kernel time; FlashSampling reduces this to 2–6%.
- End-to-end vLLM experiments (B200 GPU): up to 19% reduction in time-per-output-token (TPOT) for Qwen3-1.7B, 3–7% for Qwen3-8B, 1–2% for Qwen3-32B, and 0.3–2.4% for GPT-OSS-120B. The gain is largest for smaller LMs where LM-head matmul is a larger runtime fraction.

## 5. Memory, Bandwidth, and Architectural Insights

Baseline decode traffic incurs at least $2BV$ floats (write/read of logits), with $4BV$ bytes in BF16. Relative to mandatory LM-head weight fetch ($2VD$ bytes), fusion offers a memory bandwidth saving of $2B/D$ ($\sim$3.1% for $D=4096$, $B=64$).

Empirically, performance gains exceed what is explainable by traffic reduction alone: measured end-to-end speedups are 50–80% over naive sampling kernels. The predominant advantages derive from the elimination of kernel launches, inter-kernel synchronization, and full-HBM passes.

In memory-bound regimes ($B \le 64$), FlashSampling achieves 85–95% of peak HBM throughput, compared to 65–75% for baseline approaches. The algorithm thus tracks closer to the roofline bandwidth bound due to full fusion of post-processing.

**Limitations:** For $B > 128$, GEMM efficiency is dominated by vendor-optimized BLAS (e.g., cuBLAS), and the Triton-based kernel can be outperformed by the baseline pipeline. Nucleus (top-$p$) sampling, which requires a non-tileable global softmax/cumsum, typically applies top-$k$ with large $k$ as a bottleneck mitigation, followed by a tiny softmax/top-$p$ post-filtering. This procedure is orthogonal to the core FlashSampling method.

## 6. Integration in Online and Tensor-Parallel Contexts

FlashSampling's hierarchical reduction approach facilitates online group merging and compatibility with tensor-parallel sharding: distributed settings treat each worker’s vocabulary chunk as a group. Each rank computes local maxima and log-masses; a subsequent all-gather and Gumbel-Max over $n$ scalars yields the correct global sample without $O(V)$ communication. Grouped or hierarchical variants (Group-Gumbel-Max) are mathematically exact by independent max-stability of Gumbels, preserving the categorical distribution over the full vocabulary.

## 7. Summary and Practical Impact

FlashSampling demonstrates that mathematically exact categorical sampling is achievable as an integral lightweight epilogue to the LM-head matmul, removing explicit post-processing bottlenecks and on-HBM tensor materialization [2603.15854]. By exploiting the associativity of $\arg\max$ and hierarchical Gumbel-Max decompositions, it ensures equivalence to traditional full-vocabulary sampling, while enabling up to 2× kernel speedups and up to 19% end-to-end token latency reduction in vLLM workloads. The method provides a robust, general-purpose primitive for efficient sampling on modern GPU architectures, with strong empirical and theoretical guarantees of correctness and performance.

Source: https://www.emergentmind.com/topics/flashsampling