FlashSampling

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.

FlashSampling enables efficient categorical sampling by fusing the operation into the language model head matmul, eliminating memory overhead and reducing decoding time by up to 19%.

Standard categorical sampling first computes the full logits tensor over the vocabulary, writes it to HBM, and then samples from it — an expensive memory round-trip that dominates decode latency at scale. FlashSampling collapses this into a single fused kernel built on the Gumbel-Max trick.

The fused primitive

1. Tile-by-tile logits on chip. Compute logits over the vocabulary one tile at a time, keeping them in fast on-chip memory — never writing the full tensor to HBM.
2. Add Gumbel noise. Perturb each on-chip 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 single maximizing index and its value.
4. Reduce over tiles. A small final reduction across tiles selects the global maximizer per row — the sampled token.

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)