FlashSampling: Fast and Memory-Efficient Exact Sampling

Abstract: Sampling from a categorical distribution is mathematically simple, but in large-vocabulary decoding, it often triggers extra memory traffic and extra kernels after the LM head. We present FlashSampling, an exact sampling primitive that fuses sampling into the LM-head matmul and never materializes the logits tensor in HBM. The method is simple: compute logits tile-by-tile on chip, add Gumbel noise, keep only one maximizer per row and per vocabulary tile, and finish with a small reduction over tiles. The fused tiled kernel is exact because $\argmax$ decomposes over a partition; grouped variants for online and tensor-parallel settings are exact by hierarchical factorization of the categorical distribution. Across H100, H200, B200, and B300 GPUs, FlashSampling speeds up kernel-level decode workloads, and in end-to-end vLLM experiments, it reduces time per output token by up to $19%$ on the models we test. These results show that exact sampling, with no approximation, can be integrated into the matmul itself, turning a bandwidth-bound postprocessing step into a lightweight epilogue. Project Page: https://github.com/FlashSampling/FlashSampling.

• The paper introduces a fused matmul-epilogue technique that integrates exact Gumbel-max sampling directly into LLM inference, eliminating full-logits materialization.
• It employs a hierarchical group-Gumbel-max decomposition to enable exact sampling in tensor-parallel and streaming regimes while ensuring strict distributional correctness.
• Empirical results on NVIDIA GPUs demonstrate up to 2.52× speedup and significant reductions in Time Per Output Token, especially for low-batch decoding.

FlashSampling: Fused Matmul-Epilogue Exact Sampling for Large-Vocabulary LLM Decoding

Motivation and Problem Formulation

Autoregressive LLM decoding repeatedly samples categorical distributions over large vocabularies, incurring significant memory and kernel overhead. Conventional pipelines materialize the $[B,V]$ logits tensor in HBM, followed by normalization and sampling steps that detract from hardware efficiency and increase latency. While the mathematical operation—drawing a sample from $\mathrm{Cat}(softmax(\ell))$—is trivial, the systems bottleneck arises from unnecessary memory round-trips and kernel launches at low batch sizes, especially in the decode regime. Large-vocabulary models ($V \gg 10^5$) further exacerbate HBM and synchronization costs.

Figure 1: Conventional multinomial sampling writes full logits to HBM and launches a separate sampling kernel, while FlashSampling fuses sampling into the matmul epilogue, avoiding logits materialization and reducing HBM traffic.

FlashSampling Algorithmic Design

FlashSampling deploys two core innovations:

1. Fused Matmul-Epilogue Sampling: Sampling is fused directly into the LM-head matmul kernel. Logits are computed tile-by-tile in on-chip memory, perturbed via i.i.d. Gumbel noise, and only the tile-local maximizer per row is retained. The epilogue avoids writing the $[B,V]$ logits tensor, instead performing a lightweight reduction over vocabulary tiles to select the global maximizer.
2. Hierarchical Group-Gumbel-Max Decomposition: Exact sampling remains feasible in tensor-parallel and streaming regimes through log-mass factorization. Each shard/group computes a local sample and log-mass, and a secondary reduction samples over these summarized states.

Algorithmically, FlashSampling bypasses explicit probability formation and normalization. The fused kernel’s exactness is pathwise via $\argmax$ decomposition; distributed and online variants achieve exact distributional correctness via hierarchical composition of categorical distributions and binary merge rules. Random number generation is deterministic via counter-based indexing for reproducibility.

Theoretical Guarantees

FlashSampling delivers strict exactness with respect to the target $\mathrm{Cat}(softmax(\ell))$ distribution. The fusion’s correctness arises from max-stability: $$\max_{i\in[V]} x_i = \max_t \max_{i\in\mathcal{T}_t} x_i$$ for partitioned vocabulary tiles. For online and tensor-parallel variants, the hierarchical log-mass factorization ensures that structured sampling and binary merge regimes remain exact. No approximation is introduced: the method always returns the $\argmax$ of perturbed logits according to the classical Gumbel-Max trick.

Practical Integration: Systems and Kernel Fusion

FlashSampling leverages existing trends in IO-aware kernel fusion, as seen in FlashAttention and similar efforts applied to cross-entropy. By retaining all expensive computation within the matmul epilogue, the method eliminates additional memory-bound passes for normalization and sampling. HBM savings due to avoiding logits writes/reads are minimal in isolation, but the holistic removal of kernel launches and synchronizations yields substantial real-world speedups. Triton implementation allows platform-agnostic deployment, albeit with some tradeoffs versus cuBLAS GEMM efficiency at high batch sizes.

Empirical Evaluation and Results

FlashSampling was benchmarked across NVIDIA H100, H200, B200, and B300 GPUs. Microbenchmarks and end-to-end vLLM integration were performed on models such as Qwen3-1.7B, Qwen3-8B, Qwen3-32B, and gpt-oss-120b. Relative performance is summarized below.

Figure 2: FlashSampling outperforms Multinomial Sampling and FlashInfer baselines across all batch sizes on B300, achieving fastest decode in the low-batch regime.

FlashSampling exhibits up to 1.84$\times$ speedup vs. Multinomial Sampling and up to 2.52$\times$ speedup vs. FlashInfer FI1 at $B \le 64$; speedups narrow as batch grows due to baseline GEMM efficiency improvements. Sampling runtime grows steeply with batch size for baselines, while FlashSampling absorbs sampling into the matmul with negligible overhead.

Figure 3: Sampling and matmul runtimes in $\mu$s vs. batch size show FlashSampling’s minimal epilogue overhead and baseline kernel launch costs.

Roofline analysis demonstrates that FlashSampling achieves higher HBM bandwidth utilization, tracking the memory-bound slope and outperforming baselines in decode-dominated workloads.

Figure 4: FlashSampling sits above baselines on the roofline plot and achieves superior bandwidth utilization on H100 for small batch sizes.

End-to-end vLLM evaluation reveals reductions in Time Per Output Token (TPOT) of up to 19\% on Qwen3-1.7B. Speedups scale proportionally with the fraction of decode time spent on the LM head, diminishing for larger models where attention and FFN dominate.

Figure 5: TPOT reductions vs. concurrency for four models show FlashSampling’s strongest gains on small architectures; larger models are dominated by attention and FFN.

Correctness verification at the kernel and inference levels establishes strict statistical fidelity to reference sampling implementations; GSM8K accuracy is statistically indistinguishable between FlashSampling and conventional methods.

Extension to Top-$k$, Nucleus, and Masking

The tile-based structure allows efficient integration of top-$k$ sampling and localized soft/truncate operations; Gumbel-Top-$k$ extensions are immediate. Nucleus sampling (top-$p$) is less amenable to tile-local reduction but can be combined sequentially after top-$k$. Masking is supported by setting logits to $-\infty$ pre-perturbation.

Related Work

FlashSampling builds on statistical and systems foundations in the Gumbel-Max trick, IO-aware kernel fusion, and large-vocabulary LLM inference. Prior works such as FlashInfer, Qrita, and SIMPLE address sampling via pre-materialized logits, introducing various levels of approximation or hardware offloading. FlashSampling’s contribution is to achieve strict exactness with zero approximation and maximal memory locality via kernel fusion.

Implications and Future Directions

FlashSampling’s stricter separation of model computation from memory-bound post-processing yields demonstrable efficiency improvements in practical LLM inference, with particular benefit to low-batch, memory-bound workloads. The method’s hierarchical factorization allows for scalable, distributed, and online variants suitable for both CPU-GPU heterogeneous inference and large-cluster serving. As Transformer vocabulary sizes and serving workloads continue to scale, systems-level kernel fusion for sampling will become a critical optimization. Future work may incorporate approximate softmax/normalizers as optional features, extend to more complex sampling strategies (beam search, stochastic sequence sampling), or explore fully integrated whole-model epilogues.

Conclusion

FlashSampling provides a fused, tile-wise matmul epilogue for exact categorical sampling, eliminating unnecessary memory-bound post-processing and kernel overhead. The method’s simplicity and exactness are theoretically validated via max-decomposition and log-mass factorization. Empirical results across multiple hardware platforms and LLM architectures demonstrate consistent kernel and inference time reductions in the decode regime. FlashSampling sets a new standard for efficient large-vocabulary sampling in neural LLM inference, with substantial implications for future AI deployment at scale.