# FlashSampling: Fast and Memory-Efficient Exact Sampling

arXiv: https://arxiv.org/abs/2603.15854 · HTML: https://mlsanity.com/paper/2603.15854/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 29 of its 43 references that are in the map.
- First posted on arXiv 16 Mar 2026.

## Builds on (most cited first, 29 of 29)

- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [Training Verifiers to Solve Math Word Problems](https://mlsanity.com/paper/2110.14168/) (2021)
- [Qwen3 Technical Report](https://mlsanity.com/paper/2505.09388/) (2025)
- [Efficient Memory Management for Large Language Model Serving with PagedAttention](https://mlsanity.com/paper/2309.06180/) (2023)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [The Curious Case of Neural Text Degeneration](https://mlsanity.com/paper/1904.09751/) (2019)
- [Categorical Reparameterization with Gumbel-Softmax](https://mlsanity.com/paper/1611.01144/) (2016)
- [FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning](https://mlsanity.com/paper/2307.08691/) (2023)
- [Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism](https://mlsanity.com/paper/1909.08053/) (2019)
- [FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision](https://mlsanity.com/paper/2407.08608/) (2024)
- [FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving](https://mlsanity.com/paper/2501.01005/) (2025)
- [Stochastic Beams and Where to Find Them: The Gumbel-Top-k Trick for Sampling Sequences Without Replacement](https://mlsanity.com/paper/1903.06059/) (2019)
- [Online normalizer calculation for softmax](https://mlsanity.com/paper/1805.02867/) (2018)
- [A* Sampling](https://mlsanity.com/paper/1411.0030/) (2014)
- [Turning Up the Heat: Min-p Sampling for Creative and Coherent LLM Outputs](https://mlsanity.com/paper/2407.01082/) (2024)
- [Liger Kernel: Efficient Triton Kernels for LLM Training](https://mlsanity.com/paper/2410.10989/) (2024)
- [A Review of the Gumbel-max Trick and its Extensions for Discrete Stochasticity in Machine Learning](https://mlsanity.com/paper/2110.01515/) (2021)
- [Cut Your Losses in Large-Vocabulary Language Models](https://mlsanity.com/paper/2411.09009/) (2024)
- [Sampled Softmax with Random Fourier Features](https://mlsanity.com/paper/1907.10747/) (2019)
- [FlashRNN: I/O-Aware Optimization of Traditional RNNs on modern hardware](https://mlsanity.com/paper/2412.07752/) (2024)
- [Approximate Top-k for Increased Parallelism](https://mlsanity.com/paper/2412.04358/) (2024)
- [Fast Generating A Large Number of Gumbel-Max Variables](https://mlsanity.com/paper/2002.00413/) (2020)
- [FlashFormer: Whole-Model Kernels for Efficient Low-Batch Inference](https://mlsanity.com/paper/2505.22758/) (2025)
- [From Projection to Prediction: Beyond Logits for Scalable Language Models](https://mlsanity.com/paper/2511.17599/) (2025)
- [SIMPLE: Disaggregating Sampling from GPU Inference into a Decision Plane for Faster Distributed LLM Serving](https://mlsanity.com/paper/2512.00719/) (2025)
- [Entropy-Aligned Decoding of LMs for Better Writing and Reasoning](https://mlsanity.com/paper/2601.01714/) (2026)
- [Qrita: High-performance Top-k and Top-p using Pivot-based Truncation and Selection](https://mlsanity.com/paper/2602.01518/) (2026)
- [Deep Kernel Fusion for Transformers](https://mlsanity.com/paper/2602.11808/) (2026)

## Built on by (most cited first, 3 of 3)

- [Towards a Data-Parameter Correspondence for LLMs: A Preliminary Discussion](https://mlsanity.com/paper/2604.17384/) (2026)
- [SonicSampler: Unified Tile-Aware Kernels for LLM Sampling and Speculative Verification](https://mlsanity.com/paper/2607.20475/) (2026)
- [The Cost of Conservation: Coordination-Memory Laws for Exact-Support Generation](https://mlsanity.com/paper/2609.26126/) (2026)

## Abstract (arXiv)

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. In tensor-parallel decoding, FlashSampling replaces the all-gather of logits with streaming peer-to-peer writes: This overlaps GPU-to-GPU communication with computation and HBM loads across up to 8 GPUs, with near-ideal scaling at large batch sizes. Our kernel is exact because argmax decomposes over partitions; grouped variants for online and tensor-parallel settings are exact by hierarchical factorization of the categorical distribution. FlashSampling demonstrates kernel-level speedups on decode workloads across 4 different datacenter GPUs (H100, H200, B200, B300), and in end-to-end vLLM experiments, it reduces time per output token by up to 10% on the models we test. These results show that exact sampling, with no approximation, can be integrated into the matmul itself, consolidating the bandwidth-bound sampling step in an efficient epilogue.