FlashSampling: Fast and Memory-Efficient Exact Sampling

• URL: http://arxiv.org/abs/2603.15854v1
• Date: Mon, 16 Mar 2026 19:37:08 GMT
• Title: FlashSampling: Fast and Memory-Efficient Exact Sampling
• Abstract summary: FlashSampling is an exact sampling primitive that fuses sampling into the LM-head matmul and never materializes the logits tensor in.<n>Across H100, H200, B200, and B300 GPUs, FlashSampling speeds up kernel-level decode workloads.<n>In end-to-end vLLM experiments, it reduces time per output token by up to $19%$ on the models we test.
• Score: 62.5203057469482
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• 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. The fused tiled kernel is exact because $\argmax$ decomposes over a partition; grouped variants for online and tensor-parallel settings are exact by hierarchical factorization of the categorical distribution. Across H100, H200, B200, and B300 GPUs, FlashSampling speeds up kernel-level decode workloads, and in end-to-end vLLM experiments, it reduces time per output token by up to $19%$ on the models we test. These results show that exact sampling, with no approximation, can be integrated into the matmul itself, turning a bandwidth-bound postprocessing step into a lightweight epilogue. Project Page: https://github.com/FlashSampling/FlashSampling.

Related papers

• Conditioning Tree-Based Diffusions and Flows for Probabilistic Tabular Regression [0.0]
  Tree-based diffusion models fit flexible conditional predictive distributions for regression without a neural density estimator.<n>We show these defaults are the binding constraint: what a gradient-boosted ensemble actually solves is a supervised regression problem.<n>We present DiffGBM, which makes them explicit along two axes.
  arXiv  Detail & Related papers  (2026-07-30T22:08:02Z)
• Speculative Correction: Draft-then-Refine Decoding for Diffusion Language Models [50.758327861505784]
  We study a simple plug-and-play inference pattern: first generate a complete draft, then refine the full response using bidirectional diffusion.<n>We evaluate two configurations using LLaDA2.1-Flash and LLaDA2.1-Mini.<n>Flash-Flash improves GSM8K-384 accuracy from 0.848 to 0.899 while running 1.20 times faster than the selected Flash block-autoregressive baseline.<n>Mini-Flash provides useful quality-latency trade-offs, including MATH-384 performance of 0.294 versus 0.300 for Flash while running 2.17 times faster.
  arXiv  Detail & Related papers  (2026-07-21T17:58:10Z)
• Accelerating Discrete Diffusion Models with Parallel-In-Time Sampling [55.388363730120325]
  We parallelize the mainstream $$-leaping algorithm for absorbing discrete diffusion in a Continuous-Time Markov Chain (CTMC) framework.<n>By leveraging the continuous-time integral form of the $$-leaping algorithm and the Picard method, we achieve parallel-in-time sampling acceleration.<n>Our research expands the potential of discrete diffusion models for efficient parallel inference, with broader implications for applications such as molecular structure and language generation.
  arXiv  Detail & Related papers  (2026-07-01T10:59:33Z)
• SonicSampler: Unified Tile-Aware Kernels for LLM Sampling and Speculative Verification [27.358689380362964]
  We present $textbfSonicSampler$, a unified suite of tile-aware Triton kernels.<n>Our kernels support dynamic per-request sampling behaviors, including grammar-constrained decoding.<n>Across speculative decoding, SonicSampler achieves up to $textbf16x speedup over state-of-the-art baselines.
  arXiv  Detail & Related papers  (2026-05-24T23:01:50Z)
• Stochastic Sparse Attention for Memory-Bound Inference [19.301894658575502]
  We present Additive No-mulT Attention (SANTA), a method that sparsifies value-cache access by sampling $S ll n_k$ indices from the post-softmax distribution.<n>We also propose Bernoulli $qKmathsfT$ sampling as a complementary technique to sparsify the score stage.
  arXiv  Detail & Related papers  (2026-05-03T14:44:14Z)
• Is Your Diffusion Sampler Actually Correct? A Sampler-Centric Evaluation of Discrete Diffusion Language Models [14.764619905977739]
  We introduce a sampler-centric oracle framework that replaces learned denoisers with an exact Hidden Markov Model posterior derived from a ground-truth Markov chain.<n>We show that few-step discrete diffusion samplers are not distributionally correct even under an oracle denoiser, with transition-level mismatch that vanishes only as the number of steps approaches the sequence length.
  arXiv  Detail & Related papers  (2026-02-23T09:06:13Z)
• Constrained and Composite Sampling via Proximal Sampler [2.087898608419977]
  We study two log-concave sampling problems: constrained sampling and composite sampling.<n>The main challenge is enforcing feasibility without degrading mixing.<n>In composite sampling, the target is proportional to $exp(-f(x)-h(x))$ with closed and convex $f$ and $h$.
  arXiv  Detail & Related papers  (2026-02-16T05:36:36Z)
• Beyond GEMM-Centric NPUs: Enabling Efficient Diffusion LLM Sampling [14.471123653746275]
  Diffusion Large Language Models (dLLMs) introduce iterative denoising to enable parallel token generation.<n>Our design employs lightweight non-GEMM vector primitives, in-place memory reuse strategies, and a decoupled mixed-precision memory hierarchy.
  arXiv  Detail & Related papers  (2026-01-28T15:37:50Z)
• Upsample Anything: A Simple and Hard to Beat Baseline for Feature Upsampling [38.24831571443335]
  Upsample Anything restores low-resolution features to high-resolution, pixel-wise outputs without any training.<n>It runs in only $approx0.419 texts$ per 224x224 image and achieves state-of-the-art performance on semantic segmentation, depth estimation, and both depth and probability map upsampling.
  arXiv  Detail & Related papers  (2025-11-20T12:27:53Z)
• dParallel: Learnable Parallel Decoding for dLLMs [77.24184219948337]
  Diffusion large language models (dLLMs) offer parallel token prediction and lower inference latency.<n>Existing open-source models still require nearly token-length decoding steps to ensure performance.<n>We introduce dParallel, a simple and effective method that unlocks the inherent parallelism of dLLMs for fast sampling.
  arXiv  Detail & Related papers  (2025-09-30T16:32:52Z)
• Sampling-Efficient Test-Time Scaling: Self-Estimating the Best-of-N Sampling in Early Decoding [48.55176091771745]
  Test-time scaling enhances large language model performance by allocating additional compute resources during inference.<n>Best-of-N (BoN) sampling serves as a common sampling-based scaling technique.<n>We propose Self-Truncation Best-of-N (ST-BoN), a decoding method that avoids fully generating all N samples.<n>In terms of cost, ST-BoN achieves the same performance as Full-BoN while saving computational cost by 70%-80%, and under the same cost, it can improve accuracy by 3-4 points.
  arXiv  Detail & Related papers  (2025-03-03T11:21:01Z)
• Boundary-aware Decoupled Flow Networks for Realistic Extreme Rescaling [49.215957313126324]
  Recently developed generative methods, including invertible rescaling network (IRN) based and generative adversarial network (GAN) based methods, have demonstrated exceptional performance in image rescaling. However, IRN-based methods tend to produce over-smoothed results, while GAN-based methods easily generate fake details. We propose Boundary-aware Decoupled Flow Networks (BDFlow) to generate realistic and visually pleasing results.
  arXiv  Detail & Related papers  (2024-05-05T14:05:33Z)
• Distributed Matrix-Based Sampling for Graph Neural Network Training [0.0]
  We propose a matrix-based bulk sampling approach that expresses sampling as a sparse matrix multiplication (SpGEMM) and samples multiple minibatches at once. When the input graph topology does not fit on a single device, our method distributes the graph and use communication-avoiding SpGEMM algorithms to scale GNN minibatch sampling. In addition to new methods for sampling, we introduce a pipeline that uses our matrix-based bulk sampling approach to provide end-to-end training results.
  arXiv  Detail & Related papers  (2023-11-06T06:40:43Z)
• Accelerating Training and Inference of Graph Neural Networks with Fast Sampling and Pipelining [58.10436813430554]
  Mini-batch training of graph neural networks (GNNs) requires a lot of computation and data movement. We argue in favor of performing mini-batch training with neighborhood sampling in a distributed multi-GPU environment. We present a sequence of improvements to mitigate these bottlenecks, including a performance-engineered neighborhood sampler. We also conduct an empirical analysis that supports the use of sampling for inference, showing that test accuracies are not materially compromised.
  arXiv  Detail & Related papers  (2021-10-16T02:41:35Z)
• Cherry-Picking Gradients: Learning Low-Rank Embeddings of Visual Data via Differentiable Cross-Approximation [53.95297550117153]
  We propose an end-to-end trainable framework that processes large-scale visual data tensors by looking emphat a fraction of their entries only. The proposed approach is particularly useful for large-scale multidimensional grid data, and for tasks that require context over a large receptive field.
  arXiv  Detail & Related papers  (2021-05-29T08:39:57Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.