FlashSampling: Fast and Memory-Efficient Exact Sampling

Tomas Ruiz; Zhen Qin; Yifan Zhang; Xuyang Shen; Yiran Zhong; Mengdi Wang

cs.LGMar 16, 2026

FlashSampling: Fast and Memory-Efficient Exact Sampling

Authors: Tomas Ruiz, Zhen Qin, Yifan Zhang, Xuyang Shen, Yiran Zhong, Mengdi Wang

Organizations: LMU Munich · FlashSampling · Princeton University

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. 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 $\text{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.

Figures & tables

Figure 1: Conventional multinomial sampling (left) materializes the full

[B,V]

logits tensor in HBM between the matmul and the sampler. FlashSampling (right) fuses sampling into the matmul epilogue, followed by a lightweight reduction over vocabulary tiles. Logits are computed tile-by-tile in on-chip memory, perturbed with Gumbel noise, and reduced without ever writing the full logits tensor to HBM. Red arrows denote HBM traffic; green arrows denote on-chip data movement.

Figure 2: Relative performance on B200. Left: Versus Multinomial Sampling. Right: Versus FlashInfer FI1 and FI2. The error bars indicate min-max relative speedups across 10 runs of 100 iterations each.

Figure 3: Runtime (lower is better) of FlashSampling and the three baselines at TP=1, 2, 4, 8. The batch sizes (16, 64, 256) are selected to include compute- and memory-bound decode regimes. The multi-GPU microbenchmarks took under 5 hours total to run. The data is tabulated in Table 6 (Appendix F ).

\toprule	FlashSampling		Multinomial Sampling		FI2 (Gumbel-Max)
\cmidrule (lr)2-3 \cmidrule (lr)4-5 \cmidrule (lr)6-7 $B$	matmul (%)	sampl. (%)	matmul (%)	sampl. (%)	matmul (%)	sampl. (%)
\midrule 1	97.7	2.3	93.7	6.3	94.7	5.3
16	97.7	2.3	87.1	12.9	93.4	6.6
64	93.6	6.0	71.3	28.7	88.6	11.4
256	93.4	6.2	73.1	26.9	88.2	11.8
\bottomrule

Table 1: Sampling as a percentage of total kernel time. A high fraction spent on sampling rather than matmul is an indicator of inefficient sampling implementation. FlashSampling’s sampling fraction stays low because it is fused into the matmul epilogue; the baselines’ fraction grows with batch size

B

. Bold marks the highest sampling fraction for each method.

Figure 4: Sampling runtime (left) and matmul runtime (right) in

\mu

s vs. batch size. Lower is better.

Figure 5: TPOT vs. concurrency on B200 across four model sizes. Qwen3-1.7B (TP1) and Qwen3-8B (TP1) see the largest reductions (up to

10\%

and

7

–

9\%

, respectively). Qwen3-32B (TP2) and Llama-3.3-70B (TP2) achieve smaller but consistent gains (peaks of

2.9\%

and

2.7\%

), since attention and FFN dominate decode time at these scales.

Appendix figures & tables10 assets

Supplementary material from the paper’s appendix.

Appendix

\toprule Level	Capacity	Bandwidth
\midrule Registers/SRAM	256 KB / SM	$>$ 100 TB/s
L2 cache	50 MB	${\sim}$ 12 TB/s
HBM3	80 GB	3.35 TB/s
\bottomrule

Appendix

Table 2: GPU memory hierarchy (H100 SXM) [ NVIDIA, 2022 , NVIDIA, 2024 ] .

\toprule	H100	H200	B200	B300
\midrule Architecture	Hopper	Hopper	Blackwell	Blackwell
HBM capacity (GB)	80	141	192	288
HBM bandwidth (TB/s)	3.35	4.8	8.0	8.0
Peak BF16 dense (TFLOP/s)	989	989	2,250	2,250
Ops:byte ratio	295	206	281	281
\bottomrule

Appendix

Table 3: GPU specifications. Peak BF16 TFLOP/s are dense (without structured sparsity), since the LM-head matmul is a dense GEMM. The ops:byte ratio (peak compute / bandwidth) contextualizes the crossover between bandwidth- and compute-limited regimes, although the exact crossover is kernel-dependent.

\toprule	vs. Multinomial Sampling				vs. FI1 (top- $k$ /top- $p$ )				vs. FI2 (Gumbel-Max)
\cmidrule (lr)2-5 \cmidrule (lr)6-9 \cmidrule (lr)10-13 $B$	H100	H200	B200	B300	H100	H200	B200	B300	H100	H200	B200	B300
\midrule 1	1.29	1.35	1.58	1.55	1.34	1.39	1.65	1.57	1.20	1.24	1.35	1.33
2	1.28	1.34	1.52	1.50	1.32	1.44	1.58	1.56	1.19	1.22	1.32	1.29
4	1.28	1.35	1.58	1.54	1.32	1.42	1.62	1.61	1.19	1.23	1.37	1.32
8	1.31	1.38	1.58	1.55	1.35	1.48	1.63	1.60	1.19	1.24	1.37	1.33
16	1.35	1.44	1.66	1.63	1.37	1.49	1.70	1.65	1.20	1.25	1.39	1.35

Appendix

Table 4: FlashSampling relative speedup vs. three baselines on the smaller configuration (

D{=}4096

V{=}151\text{k}

). Values

>1

indicate FlashSampling is faster; bold marks the peak per GPU within each baseline. The numbers are medians over 100 iterations. Each column takes under 10 minutes to run.

\toprule	vs. Multinomial Sampling				vs. FI1 (top- $k$ /top- $p$ )				vs. FI2 (Gumbel-Max)
\cmidrule (lr)2-5 \cmidrule (lr)6-9 \cmidrule (lr)10-13 $B$	H100	H200	B200	B300	H100	H200	B200	B300	H100	H200	B200	B300
\midrule 1	1.23	1.26	1.45	1.41	1.20	1.25	1.36	1.30	1.14	1.13	1.21	1.19
2	1.22	1.23	1.41	1.39	1.21	1.21	1.33	1.29	1.13	1.11	1.21	1.18
4	1.22	1.24	1.35	1.34	1.20	1.21	1.31	1.28	1.13	1.12	1.15	1.13
8	1.23	1.25	1.37	1.36	1.21	1.23	1.30	1.29	1.13	1.12	1.14	1.13
16	1.24	1.27	1.40	1.39	1.21	1.25	1.32	1.30	1.14	1.13	1.15	1.15

Appendix

Table 5: FlashSampling speedup vs. three baselines on the larger configuration (

D{=}8192

V{=}128\text{k}

). Values

>1

indicate FlashSampling is faster; bold marks the peak per GPU within each baseline. At

B{\geq}128

the advantage narrows and cuBLAS GEMM efficiency becomes increasingly important.

\toprule $B$	Method	TP=1	TP=2	TP=4	TP=8
\midrule \multirow 4*16	FlashSampling	333.8	203.8	125.8	87.1
	FI1	438.3	321.5	254.1	215.1
	FI2	385.0	268.4	200.7	164.2
	Multinomial	461.8	344.1	274.5	241.7
\midrule \multirow 4*64	FlashSampling	344.1	215.9	136.2	96.4
	FI1	494.8	383.1	313.3	258.0

Appendix

Table 6: Minimum kernel runtime (

\mu

s, lower is better) across tensor-parallel sizes

\text{TP}\in\{1,2,4,8\}

and batch sizes

B\in\{16,64,256\}

for input shape (

D{=}8192

V{=}128\text{k}

). Bold marks the fastest method per (

B

, TP) cell.

\toprule	Qwen3-1.7B (TP1)		Qwen3-8B (TP1)		Qwen3-32B (TP2)		Llama-3.3-70B (TP2)
\cmidrule (lr)2-3 \cmidrule (lr)4-5 \cmidrule (lr)6-7 \cmidrule (lr)8-9 $B$	Base	FlashSampling	Base	FlashSampling	Base	FlashSampling	Base	FlashSampling
\midrule 1	1.82	1.68	3.91	3.62	7.73	7.64	14.21	14.10
2	1.81	1.67	3.93	3.61	7.85	7.66	14.38	14.24
4	1.81	1.66	4.01	3.68	7.91	7.79	14.60	14.42
8	1.83	1.67	4.06	3.74	8.00	7.82	14.77	14.62
16	1.84	1.68	4.24	3.87	8.30	8.08	15.15	14.76

Appendix

Table 7: Median TPOT (ms) over 5 runs on B200 for baseline and FlashSampling. Lower is better.

B

is the batch size.

\toprule $B$	Qwen3-1.7B (TP1)	Qwen3-8B (TP1)	Qwen3-32B (TP2)	Llama-3.3-70B (TP2)
\midrule 1	$7.5\pm 0.2$ %	$7.3\pm 0.2$ %	$1.2\pm 0.0$ %	$0.8\pm 0.0$ %
2	$7.7\pm 0.1$ %	$7.9\pm 0.2$ %	$2.3\pm 0.1$ %	$1.0\pm 0.0$ %
4	$8.3\pm 0.0$ %	$8.1\pm 0.1$ %	$1.5\pm 0.0$ %	$1.2\pm 0.0$ %
8	$8.5\pm 0.0$ %	$8.1\pm 0.1$ %	$2.3\pm 0.0$ %	$1.0\pm 0.0$ %
16	$8.8\pm 0.1$ %	$\mathbf{8.7}\pm 0.3$ %	$2.7\pm 0.0$ %	$2.5\pm 0.3$ %
32	$\mathbf{10.2}\pm 0.8$ %	$8.5\pm 0.3$ %	$\mathbf{2.9}\pm 0.2$ %	$\mathbf{2.7}\pm 0.1$ %

Appendix

Table 8: TPOT speedup in % (larger is better) computed as (

1-\text{FlashSampling}/\text{baseline}

), and standard deviation across 5 runs.

B

is the batch size. Bold marks the peak per model. Absolute TPOT values are in Appendix G . Each column takes under 2 hours to run.

Figure 6: Roofline (left) and HBM bandwidth utilization (right) on B200 (higher is better). Left: all methods track the memory-bound slope for

B\leq 64

; FlashSampling sits above baselines because it avoids the logits round-trip. Close to the ridge point (

\mathrm{AI}\approx 281

), performance flattens below the compute ceiling, where cuBLAS outperforms Triton. Right: FlashSampling achieves higher bandwidth utilization than all baselines in the decode regime, confirming that fusion removes overhead rather than shifting it.

\toprule	$D{=}8192$ , $V{=}128\text{k}$		$D{=}4096$ , $V{=}152\text{k}$
\cmidrule (lr)2-3 \cmidrule (lr)4-5 $B$	Predicted	Measured	Predicted	Measured
\midrule 1	0.02%	0.5%	0.05%	0.9%
4	0.10%	0.7%	0.20%	1.2%
16	0.39%	1.6%	0.78%	2.1%
64	1.56%	3.6%	3.13%	4.8%
128	3.13%	5.4%	6.25%	8.9%

Appendix

Table 9: Predicted vs. measured overhead of storing logits to HBM. Predicted:

2B/D

. Measured: relative slowdown of the fused kernel with the logits store enabled, averaged over 5 runs (B200 GPU).

\toprule Asset	Reference	License
\midrule PyTorch	Paszke et al. [2019]	\href https://github.com/pytorch/pytorch/blob/main/LICENSEBSD-3-Clause
Triton	Tillet et al. [2019]	\href https://github.com/triton-lang/triton/blob/main/LICENSEMIT
vLLM	Kwon et al. [2023]	\href https://github.com/vllm-project/vllm/blob/main/LICENSEApache-2.0
FlashInfer	Ye et al. [2025]	\href https://github.com/flashinfer-ai/flashinfer/blob/main/LICENSEApache-2.0
Qwen3 (1.7B, 8B, 32B)	Yang et al. [2025]	\href https://huggingface.co/Qwen/Qwen3-8B/blob/main/LICENSEApache-2.0
Llama-3.3-70B	Grattafiori et al. [2024]	\href https://www.llama.com/llama3_3/license/Llama 3.3 Community License

Appendix

Table 10: Licenses of existing assets used in this paper.

LLM Inference Acceleration

LLM inference acceleration develops algorithms and systems that reduce language-model latency, memory use, and serving cost while preserving output quality.

See all 1,084 papers

Stochastic Sparse Attention for Memory-Bound Inference
May 3, 2026Kyle Lee, Corentin Delacour, Kevin Callahan-Coray +5GPU Acceleration Self-Attention
Department of Electrical and Computer Engineering, University of California, Santa Barbara, CA, USA · Department of Electrical and Computer Engineering, University of Michigan, Ann Arbor, MI, USA · Department of Electrical and Computer Engineering, Carnegie Mellon University, Pittsburgh, PA, USA
Early-Bird Decoding: Accelerating Diffusion LLMs with Learnable Block Sizes and Parallel Sampling
Sep 15, 2026Lixuan Wei, Wei Zhou, Jianwen Wu +3Diffusion Language Model Inference Parallel Decoding
Harvard University · Georgia Institute of Technology · Purdue University
Multi-Head Low-Rank Attention
Mar 2, 2026Songtao Liu, Hongwu Peng, Zhiwei Zhang +2Long-Context Language Model Inference LLM Inference Acceleration
The Pennsylvania State University · University of Connecticut · Carnegie Mellon University +1

May 3, 2026J/K move · Enter open · S save

Stochastic Sparse Attention for Memory-Bound Inference

Kyle Lee, Corentin Delacour, Kevin Callahan-Coray, Kyle Jiang+4

Department of Electrical and Computer Engineering, University of California, Santa Barbara, CA, USA · Department of Electrical and Computer Engineering, University of Michigan, Ann Arbor, MI, USA · Department of Electrical and Computer Engineering, Carnegie Mellon University, Pittsburgh, PA, USA

GPU Acceleration Self-Attention KV Caching Long-Context Language Model Inference KV-Cache Management LLM Inference Acceleration+1

Open paper AskPDF arXiv

Autoregressive decoding becomes bandwidth-limited at long contexts, as generating each token requires reading all

n_k

key and value vectors from KV cache. We present Stochastic Additive No-mulT Attention (SANTA), a method that sparsifies value-cache access by sampling

S \ll n_k

indices from the post-softmax distribution and aggregates only those value rows. This yields an unbiased estimator of the post-softmax value aggregation while replacing value-stage multiply-accumulates with gather-and-add. We introduce stratified and systematic sampling to design variance-reduced, GPU-friendly variants. Evaluated on Llama-3.1-8B-Instruct at 32k-token contexts, S

^2

ANTA matches baseline accuracy while achieving up to

1.5\times

decode-step attention-kernel speedup over FlashInfer and FlashDecoding on an NVIDIA RTX 6000 Ada. In batched long-context generation, these kernel gains translate to up to

1.25\times

end-to-end decode-latency speedup. Finally, we propose Bernoulli

qK^\mathsf{T}

sampling as a complementary technique to sparsify the score stage, reducing key-feature access through stochastic ternary queries. Both methods are complementary to upstream quantization, low-rank projection, KV-cache compression, and KV-cache selection methods. Together, they point toward sparse, multiplier-free, and energy-efficient inference. We open-source our kernels at: https://github.com/OPUSLab/SANTA.git

Abstract

LLM Inference Acceleration

Explore similar work

Stochastic Sparse Attention for Memory-Bound Inference

Early-Bird Decoding: Accelerating Diffusion LLMs with Learnable Block Sizes and Parallel Sampling

Multi-Head Low-Rank Attention