FlashSampling: Efficient Sampling in LLMs

• FlashSampling is an exact sampling primitive that integrates the Gumbel-Max trick into the LM-head matmul to enable efficient, on-chip token selection.
• It partitions the vocabulary into tiles to perform in-register Gumbel perturbations, reducing memory traffic and eliminating separate post-processing steps.
• The approach supports tensor-parallel and online contexts, achieving up to 2× kernel speedups and significant end-to-end latency improvements.

FlashSampling is an exact sampling primitive designed to accelerate large-vocabulary decoding in neural LLMs 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) (Ruiz et al., 16 Mar 2026). 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 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 (Ruiz et al., 16 Mar 2026)).

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 (Ruiz et al., 16 Mar 2026):

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 (Ruiz et al., 16 Mar 2026). 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.