FlashSampling

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 $\arg\max$ 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.

Standard categorical sampling computes the full $[B,V]$ logits tensor with a matmul, writes it to HBM, and then launches separate kernels that read it back to normalize and sample. During decoding, these extra kernels and the logits round-trip add HBM traffic but no useful computation. FlashSampling instead samples inside the matmul epilogue using the Gumbel-Max trick: exact sampling needs only the index of the largest perturbed logit, so no softmax, prefix sum, or normalized probabilities are required.

The fused primitive

1. Tile-by-tile logits on chip. Compute logits over the vocabulary one tile at a time in on-chip memory, never writing the full tensor to HBM.
2. Add Gumbel noise. Perturb each logit with i.i.d. Gumbel noise, so that the per-row argmax is an exact draw from the softmax distribution (the Gumbel-Max trick).
3. Keep one maximizer per tile. For each row and each vocabulary tile, retain only the maximizing index and its value.
4. Reduce over tiles. A small final reduction across tiles selects the global maximizer per row, which is the sampled token.

The result is exact, not approximate: the maximum over the whole vocabulary equals the maximum over the tile-local maxima.

Tensor-parallel decoding

When the LM head is sharded across GPUs along the vocabulary, a conventional sampler must all-gather the logits after the matmul. FlashSampling writes each tile-local candidate to the other GPUs with peer-to-peer writes from inside the matmul epilogue, so communication overlaps with the matmul's compute and HBM loads. Grouped variants that pass a local sample and a log-mass per shard are exact by hierarchical factorization of the categorical distribution.

• Kernel speedups. For batch sizes $B \le 64$, FlashSampling is faster than all three baselines (compiled PyTorch multinomial sampling, FlashInfer top-k/top-p sampling (FI1), and FlashInfer Gumbel-Max sampling (FI2)) on H100, H200, B200, and B300. Its peak speedups are 2.23× over multinomial sampling (B300) and 1.74× over FI1 (B200).
• Tensor parallelism. FlashSampling is the fastest method at every tensor-parallel size for batch sizes 16 and 64, and it follows the ideal $1/\mathrm{TP}$ scaling closely at batch size 256.
• End-to-end vLLM. On B200, time per output token drops by up to 10.2% (Qwen3-1.7B), 8.7% (Qwen3-8B), 2.9% (Qwen3-32B, TP2), and 2.7% (Llama-3.3-70B, TP2).
• Exactness. A chi-squared goodness-of-fit test finds no significant difference from the PyTorch reference sampler, and Qwen3-1.7B reaches 89.4% on GSM8K with FlashSampling versus 89.6% with the baseline ($p = 0.776$).

If you find this work useful, please cite:

@article{ruiz2026flashsampling,
  title   = {FlashSampling: Fast and Memory-Efficient Exact Sampling},
  author  = {Ruiz, Tomas and Qin, Zhen and Zhang, Yifan and Shen, Xuyang and Zhong, Yiran and Wang, Mengdi},
  journal = {arXiv preprint arXiv:2603.15854},
  year    = {2026}
}

Copy BibTeX

(* denotes equal contribution, † denotes corresponding authors)