PaperlayerView exact revision on arXiv

Source-linked AI summary

FlashSampling: Fast and Memory-Efficient Exact Sampling

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

arXiv:2603.15854v2cs.LGcs.AIcs.CL

TL;DR

Large-vocabulary decoding makes exact categorical sampling costly because conventional pipelines materialize logits and perform separate memory- and communication-heavy kernels. FlashSampling fuses Gumbel-Max sampling into the LM-head matmul, and the paper reports consistent decode speedups, including up to 2.23× over Multinomial Sampling and near-ideal scaling to 8 GPUs. Its main scope limitation is reduced efficiency versus cuBLAS at large batch sizes.

  • Problem

    Large-vocabulary decoding incurs extra memory traffic and kernels from materializing logits, while tensor parallelism adds logits all-gather communication.

  • Method

    FlashSampling computes logits tile-by-tile on chip, adds Gumbel noise, retains tile-local maxima, and reduces them without materializing logits in HBM.

  • Results

    FlashSampling is faster than all baselines for B ≤64, reaching 2.23× versus Multinomial Sampling on B300 and 1.74× versus FI1 on B200.

  • Takeaways & Limitations

    Exact categorical sampling can be integrated into the matmul epilogue, eliminating logits materialization and overlapping distributed communication with computation.

  • Takeaways & Limitations

    The Triton matmul implementation becomes less efficient than cuBLAS at large batch sizes.

Abstract

from arXiv · show

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.

1 Introduction

Large-vocabulary sampling becomes a memory and communication bottleneck because conventional pipelines materialize logits and launch separate sampling work. FlashSampling fuses exact sampling into the LM-head epilogue, reducing data movement while supporting distributed decoding.

  • Motivation: Sampling can consume over 10% of single-GPU token-generation time and 20–38% in tensor-parallel settings.The bottleneck is attributed primarily to separate kernels that materialize, normalize, and scan logits.
  • Motivation: Conventional sampling writes the full [B, V] logits tensor to HBM, then rereads it for normalization and sampling.These operations add memory traffic without useful model computation because logits are discarded after one sample.
  • FlashSampling: FlashSampling computes logits tile-by-tile on chip, adds Gumbel noise, retains one candidate per tile, and performs a lightweight reduction.The full logits tensor is never materialized in HBM.
  • FlashSampling: In tensor-parallel decoding, FlashSampling overlaps cross-GPU communication with matmul computation and HBM loads, scaling near-ideally to 8 GPUs.It uses streaming peer-to-peer communication rather than gathering the full logits tensor.
  • Exactness: The fused tiled kernel is exact pathwise, while grouped, online, and distributed variants are exact in distribution through hierarchical factorization.The distinction is between argmax decomposition over vocabulary tiles and factorization through group log-masses.

2 Background

The background frames sampling as a Gumbel-Max argmax problem rather than a required softmax computation, while distributed logits introduce all-gather communication overhead.

  • Notation: Transformed logits are denoted e_l, and each row is assumed to contain at least one finite entry.Without a finite entry, the target categorical distribution is undefined.
  • Sampling pipelines: Conventional pipelines compute, transform, normalize, and sample logits, requiring at least one logits write and reread when logits reach HBM.Softmax followed by inverse-CDF sampling is given as an example.
  • Distributed logits: Vocabulary-sharded tensor-parallel decoding requires an all-gather of logit shards before sampling.The per-GPU communication is of order B · V, with aggregate costs proportional to the number of GPUs.
  • Gumbel-Max: The Gumbel-Max trick samples exactly by adding i.i.d. Gumbel noise to logits and taking the argmax.This avoids explicitly forming probabilities or a softmax.

3 FlashSampling

FlashSampling performs exact Gumbel-Max sampling inside the LM-head matmul by retaining tile-local maxima and reducing them, thereby eliminating the logits round-trip to HBM.

  • Core algorithm: FlashSampling maintains the largest perturbed score and its index in a single online pass, requiring no softmax, normalization constant, or prefix sum.The algorithm computes s_i = e_l_i + g_i and returns the argmax.
  • Core algorithm: Each vocabulary tile keeps a local maximizer, and a second reduction selects the global maximizer across tiles.Only the current best score and index need to be retained per row.
  • Two-stage design: The fused kernel computes batch and vocabulary tiles on chip, applies transforms and Gumbel noise, and writes only one candidate per row and tile.Stage 2 reduces the candidate buffer to one sample per row.
  • Exactness: Exactness follows because tiled processing finds the same perturbed-logit maximizer as a full Gumbel-Max pass without forming probabilities.The method therefore introduces no approximation.
  • Multi-GPU communication: In multi-GPU decoding, per-tile peer-to-peer writes replace a post-GEMM all-gather and overlap communication with computation and HBM loads.A cross-rank barrier precedes the reduction stage because the writes are not collective operations.
  • IO cost model: The IO-model speedup is the ratio of baseline and fused data-movement costs, with the simplified expression approaching 1 + 2B for current LLMs.The model predicts larger speedups at larger batch sizes and for smaller hidden dimensions, while measured speedups exceed its predictions.

4 Experiments

Experiments evaluate FlashSampling through kernel microbenchmarks and end-to-end vLLM decoding across multiple GPUs, batch sizes, tensor-parallel configurations, and model sizes. FlashSampling consistently improves performance, especially in memory-bound regimes, while its advantage narrows when computation dominates.

  • Setup: Kernel microbenchmarks span four NVIDIA datacenter GPUs, decode-centric configurations, and batch sizes from 1 to 256.The evaluation uses BF16 inputs and weights, with additional larger-dimension results showing the same qualitative trends.
  • Single-GPU Results: FlashSampling is faster than all baselines across all batch sizes on B200, with peak speedups of 2.23× versus Multinomial Sampling and 1.74× versus FI1.FI2 gains are smaller because it already uses Gumbel-Max sampling.
  • Batch-Size Trend: At B = 256, FlashSampling’s advantage narrows because GEMM efficiency matters more and the workload becomes less dominated by memory-bound sampling.The same qualitative trend appears for D = 8192, with the crossover occurring earlier.
  • Multi-GPU Results: Across TP ∈{1, 2, 4, 8}, FlashSampling is faster at memory-bound batch sizes and closely follows ideal speedup at batch size 256.Per-tile P2P writes overlap matmul computation and effectively hide GPU-to-GPU communication, unlike baselines that perform an all-gather afterward.
  • Performance Analysis: Separate sampling kernels are the primary speedup opportunity: FlashSampling absorbs sampling into the matmul at 2–6% of kernel time, while baseline sampling runtime grows with batch size.Avoiding the logits write and reread alone saves at most 6% of traffic; the advantage narrows partly because Triton matmul becomes less efficient than cuBLAS.
  • End-to-End Evaluation: End-to-end vLLM TPOT reductions reach 10.2% for Qwen3-1.7B and 8.7% for Qwen3-8B, while Qwen3-32B and Llama-3.3-70B achieve peaks of 2.9% and 2.7%.The smaller gains on larger models occur because attention and FFN layers dominate decode time.

5 Related Work

FlashSampling extends fused matmul-epilogue techniques to exact inference-time categorical sampling, exploiting Gumbel-Max decomposability rather than approximation. It differs from related sampling and fusion methods by avoiding pre-materialized logits.

  • Prior fusion work primarily targets attention, training-time cross-entropy, MLPs, RNNs, and whole-model inference rather than exact decoding sampling.
  • EVT auto-generates fused GEMM epilogues, while related TPU work fuses approximate top-k selection; FlashSampling instead achieves exact sampling without approximation.
  • FlashSampling applies IO-aware matmul-plus-epilogue fusion to inference-time sampling, where exactness follows from Gumbel-Max decomposability.
  • Existing efficient sampling methods such as FlashInfer, Qrita, Min-p, SIMPLE, and sampled softmax operate on pre-materialized logits or trade exactness for speed.

6 Conclusion

FlashSampling performs exact categorical sampling without materializing the full logits tensor in HBM, using tiled argmax decomposition and grouped log-masses for online and distributed variants. It is most effective in memory-bound decoding, while its Triton matmul is less efficient than cuBLAS at large batch sizes and lower-precision support remains unfinished.

  • FlashSampling keeps logits on chip and avoids writing the full [B, V] logits tensor to HBM during exact categorical sampling.
  • Exactness follows from argmax decomposition over vocabulary tiles, while grouped log-masses provide exact online and distributed variants.
  • In multi-GPU decoding, FlashSampling overlaps cross-GPU communication with logit computation and HBM loads, scaling near-ideally to 8 GPUs.
  • At large batch sizes, the Triton matmul becomes less efficient than cuBLAS, narrowing the advantage despite portability to platforms such as AMD GPUs.
  • The top-k extension is proven correct but not implemented, and lower-precision inputs including FP8 and MXFP4 are not yet supported.

D Theoretical Analysis of FlashSampling

FlashSampling establishes exactness for fused, grouped, online, and distributed sampling by combining partition-wise maximization with hierarchical categorical factorization. Its extensions cover streaming groups, tensor-parallel shards, truncated supports, and masking, while top-p remains non-decomposable except after top-k reduction.

  • Exactness of grouped sampling: Grouped variants are exact by sampling a group from its log-masses and then sampling within the selected group.This follows from categorical factorization over partitioned groups.
  • Online and distributed variants: Online FlashSampling preserves exactness by merging each streamed group with the running sample using a binary mass-weighted choice.Zero-mass groups can be skipped, and induction over groups yields an exact full-categorical sample.
  • Online and distributed variants: Tensor-parallel FlashSampling replaces vocabulary-sized logit all-gather with local samples and shard log-masses followed by a final exact shard selection.Each rank treats its vocabulary shard as a group and returns only a local sample plus its shard log-mass.
  • Exactness arguments: Grouped Gumbel-Max relies on partition-wise maxima and Gumbel max-stability, whereas the fused two-stage kernel needs only the deterministic maximum decomposition.The grouped outer sample can use fresh Gumbels or explicitly computed group maxima.
  • Exactness of fused sampling: The fused tiled kernel is exact because the global maximizer equals the maximum among tile-local maximizers.With continuous Gumbel noise, the maximizer is unique almost surely, so tile-wise reduction returns the same index as full-row argmax.
  • Extensions and scope: Top-k extends through local candidate reduction, masking preserves exactness over restricted supports, and top-p requires global softmax, sorting, and cumulative summation.The implementation of these integrated strategies is left for future work.

E Kernel Microbenchmark Data

Kernel microbenchmarks show that FlashSampling is strongest in the small-batch decode regime, with its advantage narrowing as GEMM efficiency becomes more important.

  • Kernel microbenchmark data: FlashSampling is strongest in small-batch decode workloads, while its advantage narrows when workloads become more GEMM-efficiency dominated.The smaller-configuration results report the same qualitative pattern across configurations.

F Multi-GPU Runtime Values

Across multi-GPU kernel measurements, FlashSampling generally delivers the lowest runtime, although its advantage narrows at larger batch sizes and one configuration favors FI2.

  • Multi-GPU runtime values: FlashSampling attains the lowest runtime in every measured cell except B=256, TP=1, where FI2 is marginally faster.Table 6 uses minimum kernel runtime in microseconds across TP ∈ {1, 2, 4, 8} and B ∈ {16, 64, 256}.
  • Multi-GPU runtime values: At B≥128, FlashSampling’s advantage narrows as cuBLAS GEMM efficiency becomes increasingly important.The larger-configuration benchmark attributes the narrowing advantage to workload and GEMM-efficiency effects.

G Absolute/Relative TPOT Results for vLLM Evaluation

The vLLM evaluation reports median time-per-output-token measurements and relative speedups across four Qwen3 and Llama models, using single- or two-GPU tensor parallelism according to model size.

  • vLLM evaluation: The evaluation measures median TPOT in milliseconds across Qwen3-1.7B, Qwen3-8B, Qwen3-32B, and Llama-3.3-70B.The smaller models use TP1, while the larger models use TP2.
  • vLLM evaluation: TPOT speedup is reported as (1 − FlashSampling/baseline), with standard deviation across five runs and peak values marked per model.Absolute TPOT values are provided separately from the relative speedups.

H Roofline Analysis and Bandwidth Utilization

FlashSampling targets the memory-bandwidth-bound LM-head regime by fusing exact Gumbel-Max sampling into tile-wise computation. Its grouped and distributed variants preserve exactness while reducing materialization and communication overhead.

  • The LM-head projection is memory-bandwidth-bound at small batch sizes because arithmetic intensity equals B, with the weight matrix dominating traffic.
  • FlashSampling achieves higher bandwidth utilization than all baselines in the B200 decode regime and avoids the logits round-trip.Near the ridge point, performance flattens below the compute ceiling, where cuBLAS outperforms Triton.
  • FlashSampling generates Gumbel noise and performs streaming argmax over perturbed logits, retaining only tile-level candidates before final reduction.Gumbel noise is generated as g = −log(−log u) with u ∼ Unif(0, 1).
  • Grouped variants produce exact samples from the target categorical distribution without approximation.Parallel and sequential formulations process vocabulary groups independently or one at a time while preserving exact sampling.
  • Distributed tensor-parallel variants communicate local summaries rather than all logits, using O(1) scalars per rank after local shard computation.The distributed formulation computes local log-masses and samples, then combines these summaries across ranks.

K Logits-Store Ablation

The logits-store ablation isolates the cost of writing computed logits back to HBM and compares measured overhead with the cost-model prediction. The measured overhead is slightly larger but follows the predicted trend.

  • The ablation changes only a single kernel flag, isolating logits data-movement overhead without other kernel modifications.The measured overhead tracked the cost-model trend closely despite being slightly larger than predicted.
  • Table 9 measures the relative slowdown from enabling logits storage in the fused kernel and compares it with the predicted 2B/D overhead.The measurement averages five runs on a B200 GPU.
  • The core sampler does not require log Z, while returning it is optional and adds numerically stable log-sum-exp work to the fused epilogue.This extra epilogue work is why log Z is treated as an optional feature rather than part of the core design.
Loading 2603.15854v2…