FlashSampling: Fast and Memory-Efficient Exact Sampling

TL;DR

FlashSampling is a fused, IO-aware kernel that does exact categorical sampling directly inside the LM-head matmul — logits are computed tile-by-tile on chip, Gumbel noise is added in-tile, and only the per-tile argmax survives to a small final reduction. The logits tensor is never materialized in HBM. Exactness comes from two algebraic facts: argmax decomposes over a partition (so per-tile argmax of logits + Gumbel is the global argmax of the same), and the categorical distribution factorizes hierarchically for grouped online / tensor-parallel variants. Across H100, H200, B200, and B300, the kernel speeds up decode microbenchmarks; in end-to-end vLLM integration it cuts time per output token by up to 19% on the tested models.

Key claims

Large-vocabulary decoding’s softmax-then-sample postprocessing is bandwidth-bound: it materializes a B × V logits tensor in HBM after the LM-head matmul, then runs separate softmax + sample kernels that re-read it [Abstract, §1].
The Gumbel-Max trick lets sampling be re-cast as argmax over logits + Gumbel, which decomposes over a partition of the vocabulary — so an argmax over a vocabulary tile commutes with the global argmax, making a fused tile-streaming sampler mathematically exact, not approximate [§2.3, §3].
Grouped variants for online (chunked) and tensor-parallel settings are exact by hierarchical factorization of the categorical distribution — the partition-argmax property extends across both the vocabulary tile axis and the TP-rank axis [Abstract, §3].
In tensor-parallel decoding, FlashSampling replaces the all-gather of the full logits tensor with streaming peer-to-peer writes of only the per-shard argmax; GPU-to-GPU communication overlaps with computation and HBM loads, with near-ideal scaling reported up to 8 GPUs at large batch [§3, §4].
Implementation is in Triton; experiments run on Modal cloud across four NVIDIA datacenter GPUs (H100, H200, B200, B300) spanning Hopper and Blackwell, with PyTorch 2.11 / CUDA 13.0 / Triton 3.6 [§5].
End-to-end vLLM integration reduces time per output token by up to 19% on the tested models — i.e. the savings are not just a kernel microbenchmark artefact but propagate through the full inference stack [Abstract, §5].
The contribution is dataflow, not algorithm: output samples are exactly distributed as the original categorical, with no approximation [Abstract, §1].

Method

A standard decode step runs four ops back-to-back: (1) LM-head matmul logits = h @ W^T writing a B × V tensor to HBM; (2) softmax over V; (3) Gumbel/temperature scaling; (4) sample / argmax. Each of (2)–(4) is bandwidth-bound — it touches B × V once for nothing more than a row-wise reduction.

FlashSampling collapses all four into one tile-streamed epilogue on top of the matmul. For each output row, the kernel iterates over vocabulary tiles: compute the logit tile on-chip, add per-element Gumbel noise sampled from the tile’s RNG state, and keep only the per-row (max_value, argmax_index) over the tile. After all tiles are processed the kernel emits a small reduction (the partition argmax) to produce the sampled token id. Because argmax decomposes over partitions, the result is bit-identical (up to RNG seeding) to the unfused exact sampler.

For tensor-parallel decoding, each rank does the same fused tile-stream over its vocabulary shard, producing a per-rank (local_max, local_argmax). Ranks exchange only these two scalars per row (streaming peer-to-peer writes), and a final cross-rank argmax yields the sampled token — replacing the standard all_gather(logits) of an B × V tensor with 2 × B scalars and overlapping the exchange with the per-shard compute.

For online (chunked) settings the same hierarchical factorization gives an exact streaming sampler that can finalize the sample as soon as later tiles can no longer overtake the running maximum.

Results

Reported across H100, H200, B200, and B300 (Hopper + Blackwell):

Kernel-level decode microbenchmarks: FlashSampling speeds up per-token decode workloads across all four architectures [Abstract, §5].
End-to-end vLLM: time per output token reduced by up to 19% on the tested models [Abstract, §5].
Tensor parallelism: near-ideal scaling up to 8 GPUs at large batch sizes; the all-gather of logits is replaced with streaming P2P writes of per-rank max/argmax pairs [§3, §4].
An OpenReview-stage earlier draft (a precursor “Group-Gumbel-Max” variant) cited up to 3.84× faster sampling and ~18× reduced memory on its own benchmarks; the arxiv version’s headline is the vLLM 19% E2E number rather than the standalone-kernel speedup.

Why it’s interesting

FlashSampling is the third filed instance of the same IO-aware kernel-design playbook on the wiki, applied to a new bottleneck — enough that the wiki now has a IO-Aware Kernel Design concept page anchoring the recipe. The template — identify HBM materialization of an intermediate tensor as the bottleneck, fuse the produce/consume pair into a streaming kernel with on-chip state, exploit an algebraic decomposition that lets the reduction commute with tiling — was set by FlashAttention; the wiki has it applied to the dense softmax(QK^T)V in FlashAttention-4: Algorithm and Kernel Pipelining Co-Design for Asymmetric Hardware Scaling (which also targets H100/B200) and to Lloyd’s k-means in Flash-KMeans: Fast and Memory-Efficient Exact K-Means (which collapses the N×K distance matrix via an online argmin + segmented-reduction update). FlashSampling does the equivalent move at the output of the LM head: argmax-over-partitions plays the same role for sampling that softmax-over-blocks plays for attention, and the categorical- distribution hierarchical factorization is what lets the tensor- parallel variant ship per-rank (max, argmax) instead of the full logits all-gather.

It also bookends Lost in Backpropagation: The LM Head is a Gradient Bottleneck from the opposite end of the LM head: Lost-in-Backprop shows the backward pass through the LM head destroys 95–99% of the gradient norm and bottlenecks pretraining; FlashSampling shows the forward sampling step through the same LM head is the inference-decode bottleneck. The LM-head layer keeps getting upgraded from “trivial output projection” to a load-bearing primitive that deserves its own training-time and inference-time engineering.