Paper deep dive
FlashSampling: Fast and Memory-Efficient Exact Sampling
Tomas Ruiz, Zhen Qin, Yifan Zhang, Xuyang Shen, Yiran Zhong, Mengdi Wang
Intelligence
Status: succeeded | Model: google/gemini-3.1-flash-lite-preview | Prompt: intel-v1 | Confidence: 95%
Last extracted: 3/22/2026, 5:26:06 AM
Summary
FlashSampling is a memory-efficient, exact sampling primitive for large-vocabulary LLM decoding that fuses sampling into the LM-head matmul. By computing logits tile-by-tile on-chip and applying the Gumbel-Max trick, it avoids materializing the full logits tensor in HBM, reducing memory bandwidth bottlenecks and improving performance in decode workloads.
Entities (5)
Relation Signals (3)
FlashSampling → utilizes → Gumbel-Max trick
confidence 100% · The fused tiled kernel is exact because argmax decomposes over a partition
FlashSampling → optimizes → LM-head matmul
confidence 95% · fuses sampling into the LM-head matmul
FlashSampling → reduceslatencyin → vLLM
confidence 90% · in end-to-end vLLM experiments, it reduces time per output token by up to 19%
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
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: this https URL.
Tags
Links
- Source: https://arxiv.org/abs/2603.15854v1
- Canonical: https://arxiv.org/abs/2603.15854v1
Trouble viewing inline? Open PDF directly →
Full Text
54,421 characters extracted from source content.
Expand or collapse full text
FlashSampling: Fast and Memory-Efficient Exact Sampling Tomas Ruiz1 Zhen Qin311footnotemark: 1 Yifan Zhang222footnotemark: 2 Xuyang Shen3 Yiran Zhong3 Mengdi Wang2† 1LMU Munich 2Princeton University 3FlashSampling Equal contribution; †Corresponding authors. ‡Revised: March 16, 2026. (February 28, 202633footnotemark: 3) 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 *arg\,max 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%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 1 Introduction Sampling from a categorical distribution is a small mathematical operation, but in large-categorical systems, it can become an expensive inner-loop primitive. Modern LLM serving stacks invoke sampling repeatedly during autoregressive decoding, often on outputs with tens or hundreds of thousands of categories (kwon2023efficient; ye2025flashinfer; maddison2014astar; huijben2022review). Recent measurements confirm the cost: sampling can account for over 10% of token generation time even on a single GPU (key2024approximate), and 20–38% in tensor-parallel settings where logits must be gathered across ranks (zhao2025simpledisaggregatingsamplinggpu). The bottleneck is usually not arithmetic, but the chain of separate kernels that materialize, normalize, and scan the logits tensor. At decode time, the LM-head projection already streams a large [V,D][V,D] weight matrix from HBM. When the active batch is small, this projection is typically memory-bandwidth bound. Materializing the resulting [B,V][B,V] logits tensor, launching extra kernels to normalize and sample from it, and then discarding it adds extra memory traffic and synchronization but no useful model computation. In this regime, the separate sampler is pure overhead (dao2022flashattention; wijmans2025cutyourlosses). Throughout, B denotes batch size and V denotes the number of categories, such as vocabulary size. Figure 1: Conventional multinomial sampling (left) materializes the full [B,V][B,V] logits tensor in HBM between the matmul and the sampler. FlashSampling (right) fuses sampling into the matmul epilogue, followed by a lightweight reduction over vocabulary tiles. Logits are computed tile-by-tile in on-chip memory, perturbed with Gumbel noise, and reduced without ever writing the full logits tensor to HBM. Red arrows denote HBM traffic; green arrows denote on-chip data movement. Standard pipelines write logits to HBM and read them back for sampling, even though logits are immediately discarded after one sample is drawn. Exact sampling is often described as “compute softmax, then sample”, which obscures the fact that exact sampling does not require forming probabilities at all. For large vocabularies, streaming and tensor-parallel settings turn sampling into a memory and communication problem if full logits must be materialized or gathered. In this work, we introduce FlashSampling, which computes logits tile-by-tile on chip and writes only one candidate per row and per vocabulary tile, followed by a lightweight reduction. Exact sampling needs only the index of the largest perturbed logit, so there is no need to form a softmax, a prefix sum, or normalized probabilities; the method introduces no approximation. A simple hierarchical factorization yields exact online and distributed variants that keep only small summaries in flight and communicate only small summaries across ranks. Our contributions can be summarized as follows: 1. FlashSampling, a simple fused exact sampler. We introduce a two-stage design that computes logits tile-by-tile in the LM-head epilogue, adds Gumbel noise on chip, and stores only one candidate per row and per vocabulary tile instead of materializing the full [B,V][B,V] logits tensor. 2. A clean exactness argument. We separate the two ingredients used in the paper: the fused tiled kernel is exact pathwise by argmax *arg\,max decomposition over vocabulary tiles, while grouped, online, and distributed variants are exact in distribution by hierarchical factorization through group log-masses. 3. A systems analysis and evaluation. We show why raw logits-byte savings alone are too small to explain the measured speedups, and we demonstrate consistent gains in the memory-bandwidth-bound decode regime across four NVIDIA GPUs and in end-to-end vLLM evaluation. 2 Background Notation. Let [V]:=1,…,V[V]:=\1,…,V\. Let ℓ~∈(ℝ∪−∞)V ∈(R∪\-∞\)^V denote transformed logits after any deterministic operations such as additive bias, temperature scaling, or masking. We assume that each row has at least one finite entry; otherwise, the target categorical distribution is undefined. The target distribution is p(i)=exp(ℓ~i)∑j=1Vexp(ℓ~j).p(i)= ( _i) _j=1^V ( _j). Raw logits ℓ are the special case ℓ~=ℓ = . We denote i.i.d. standard Gumbel variables by gi∼Gumbel(0,1)g_i (0,1). Because the Gumbel law is continuous, ties occur with probability zero, so argmax *arg\,max is unique almost surely. 2.1 Why Sampling Is Expensive at Scale A common materialized-logits pipeline first computes transformed logits, then forms probabilities, and finally samples from those probabilities. One representative example is softmax followed by prefix-sum sampling: GEMM(produce logits)→write logits to HBM→read logits for sampling.GEMM(produce logits)\;→\;write logits to HBM\;→\;read logits for sampling. Algorithm 1 summarizes this pattern. Algorithm 1 One common materialized-logits sampling pipeline 1:Hidden state ∈ℝD h ^D, LM-head weights ∈ℝV×D W ^V× D, optional deterministic transforms 2:Sampled index i⋆∈1,…,Vi ∈\1,…,V\ 3:ℓ← ← W h ⊳ GEMM: compute logits and write to HBM 4:ℓ~←transform(ℓ) ( ) ⊳ temperature, bias, mask; read/write HBM 5:m←maxiℓ~im← _i _i ⊳ pass 1 over transformed logits 6:Z←∑i=1Vexp(ℓ~i−m)Z← _i=1^V ( _i-m) ⊳ pass 2 over transformed logits 7:pi←exp(ℓ~i−m)/Zp_i← ( _i-m)/Z for all i ⊳ write probabilities 8:ci←∑j=1ipjc_i← _j=1^ip_j for all i ⊳ prefix sum 9:Draw u∼Unif(0,1)u (0,1) 10:i⋆←mini:ci≥ui ← \i:c_i≥ u\ ⊳ search 11:return i⋆i Not every implementation uses exactly these kernels, but any materialized-logits baseline pays the same structural costs: at least one logits write, at least one logits reread, and extra sampling work after the GEMM. Decode regime. In autoregressive decoding, B is typically small. The LM-head projection is then often memory-bandwidth bound because it repeatedly streams the large [V,D][V,D] weight matrix from HBM. Materializing [B,V][B,V] logits and reading them back for sampling adds multiple avoidable HBM round-trips in the most latency-sensitive part of the decode loop (kwon2023efficient; ye2025flashinfer). 2.2 GPU Memory Hierarchy Table 1 summarizes the GPU memory hierarchy. On-chip memory (registers, SRAM) is orders of magnitude faster than HBM but far smaller. FlashSampling exploits this gap by keeping logits in registers/SRAM and never writing the full logits tensor to HBM. Table 1: GPU memory hierarchy (H100 SXM) (nvidia_h100_whitepaper; nvidia_h100_datasheet). Level Capacity Bandwidth Registers/SRAM 256 KB / SM >> 100 TB/s L2 cache 50 MB ∼ 12 TB/s HBM3 80 GB 3.35 TB/s 2.3 The Gumbel-Max Trick The classical Gumbel-Max trick states that exact categorical sampling can be performed by adding i.i.d. Gumbel noise and taking an argmax *arg\,max: Theorem 2.1 (Gumbel-Max). Let ℓ~∈(ℝ∪−∞)V ∈(R∪\-∞\)^V have at least one finite entry, and let gii=1V\g_i\_i=1^V be i.i.d. Gumbel(0,1)Gumbel(0,1). Then i⋆=argmaxi∈[V](ℓ~i+gi)⟹ℙ(i⋆=i)=eℓ~i∑j=1Veℓ~j.i = *arg\,max_i∈[V] ( _i+g_i ) (i =i)= e _i _j=1^Ve _j. This classical result goes back to gumbel1954statistical and is widely used in machine learning (maddison2014astar; huijben2022review). The trick extends to sampling without replacement via the Gumbel-Top-k method (pmlr-v97-kool19a). The key point for this paper is simple: exact sampling does not require an explicit softmax. It only requires the index of the largest perturbed logit. 3 FlashSampling We now describe FlashSampling from simplest to most practical form. The core algorithm is intentionally simple and introduces no approximation: maintain the largest perturbed score seen so far and its index. 3.1 Exact Sampling via Online Gumbel-Max Given transformed logits ℓ~∈(ℝ∪−∞)V ∈(R∪\-∞\)^V, exact sampling from Cat(softmax(ℓ~))Cat(softmax( )) is: i⋆=argmaxi∈[V](ℓ~i+gi),gi∼Gumbel(0,1) i.i.d.i = *arg\,max_i∈[V] ( _i+g_i ), g_i (0,1) i.i.d. Algorithm. Generate i.i.d. Gumbels, compute si=ℓ~i+gis_i= _i+g_i, and return i⋆=argmaxisii = *arg\,max_is_i. The computation can be performed online in a single pass that maintains only the current best score and its index, analogous to the online normalizer calculation for softmax (milakov2018online). No softmax, no normalization constant, and no prefix sum are required (see Algorithm B.1 in the Appendix). Systems implication. Sampling reduces to a single reduction over perturbed logits. This naturally fits GPU reductions and removes the extra normalization and prefix-sum work used by common softmax-based pipelines. Simplicity. The online algorithm keeps only two running state variables per row: the current best perturbed score and the corresponding index. This simplicity is what makes fusion with the LM-head epilogue practical. GPU parallelization. Each threadblock can process one contiguous vocabulary chunk, or vocabulary tile. The block computes perturbed scores for that chunk, keeps only the tile-local maximizer, and a small second-stage reduction selects the global maximizer across vocabulary tiles. 3.2 FlashSampling for LM-Head Sampling We now consider the common case where logits are produced by GEMM: =⊤∈ℝB×V, Y= H W ^B× V, where ∈ℝB×D H ^B× D are hidden states and ∈ℝV×D W ^V× D are LM-head weights. We wish to sample one index per row from Cat(softmax(b,:))Cat(softmax( Y_b,:)), possibly after deterministic transforms such as temperature scaling, additive bias, or masking. Goal: avoid materializing Y. FlashSampling performs sampling inside the matmul kernel and writes only one candidate per row and per vocabulary tile, never the full [B,V][B,V] logits tensor: • Stage 1 (fused kernel): compute one batch tile and one vocabulary tile on chip, apply deterministic transforms, add Gumbel noise, and keep the tile-local maximizer for each row. • Stage 2 (reduction): reduce over vocabulary-tile candidates to obtain one global sample per row. Algorithm 2 FlashSampling fused matmul-sample (two-stage): one candidate per row and per vocabulary tile, followed by reduction 1:Hidden states ∈ℝB×D H ^B× D, LM-head weights ∈ℝV×D W ^V× D, temperature τ>0τ>0, optional mask/bias, RNG key 2:Samples ⋆∈1,…,VB i ∈\1,…,V\^B 3:Stage 1 (fused kernel): for each batch tile ℬB and vocabulary tile tT_t in parallel 4:Initialize accumulator (t)∈ℝ|ℬ|×|t|←0 A^(t) ^|B|×|T_t|← 0 5:for d0=1,1+Ktile,…,Dd_0=1,1+K_tile,…,D do 6: Load ℬ,d0:d0+Ktile−1 H_B,\,d_0:d_0+K_tile-1 and t,d0:d0+Ktile−1 W_T_t,\,d_0:d_0+K_tile-1 into on-chip memory 7: (t)←(t)+ℬ,d0:d0+Ktile−1(t,d0:d0+Ktile−1)⊤ A^(t)← A^(t)+ H_B,\,d_0:d_0+K_tile-1 ( W_T_t,\,d_0:d_0+K_tile-1 ) 8:end for 9:for each output element (b,i)∈ℬ×t(b,i) ×T_t do 10: y~b,i←transform(Ab,i(t)) y_b,i \! (A^(t)_b,i ) ⊳ temperature, bias, mask 11: Draw ub,i∈(0,1)u_b,i∈(0,1) and set gb,i←−log(−logub,i)g_b,i←- \! (- u_b,i ) 12: sb,i←y~b,i+gb,is_b,i← y_b,i+g_b,i 13:end for 14:for each row b∈ℬb do 15: (mb(t),jb(t))←argmaxi∈tsb,i(m_b^(t),j_b^(t))← *arg\,max_i _ts_b,i 16: idxb(t)←idx_b^(t)← global vocabulary index corresponding to jb(t)j_b^(t) 17: Write (mb(t),idxb(t))(m_b^(t),idx_b^(t)) to HBM 18:end for 19: 20:Stage 2 (reduction): for each row b 21:t⋆←argmaxtmb(t)t ← *arg\,max_tm_b^(t) 22:ib⋆←idxb(t⋆)i_b _b^(t ) 23:return ⋆ i Why the two-stage design is simple. The fused stage does all expensive work in the matmul epilogue. The second stage is only an argmax *arg\,max over a small candidate buffer of shape roughly [B,#vocab tiles][B,\#vocab tiles]. This design is easy to implement and already captures most of the benefit in the decode regime. Why this avoids softmax. The algorithm never forms probabilities and never computes an explicit softmax. Exactness follows because it computes the same maximizer of the perturbed logits that a full Gumbel-Max pass would compute. Tensor-parallel fusion. When the vocabulary is sharded across ranks, each rank can run the fused kernel on its local shard and return only small summaries rather than all local logits. In the grouped formulation below, these summaries are a local sample and a local log-mass. No O(V)O(V) all-gather of logits is required. RNG determinism. For reproducibility, RNG streams are indexed by the logical output position (b,i)(b,i) using a counter-based RNG (e.g. Philox), so each random number is a deterministic function of a key and a counter. Uniform variates are mapped to the open interval (0,1)(0,1) to avoid infinities in the Gumbel transform g=−log(−logu)g=- (- u). Numerical precision. GEMM accumulation and perturbed scores are computed in FP32 for stability, even when inputs are FP16 or BF16. Gumbel noise is likewise generated in FP32 to avoid numerical error in the logarithms. The overhead is minor compared with the GEMM itself. 4 Theoretical Analysis of FlashSampling This section separates the two exactness arguments used in the paper. The fused tiled kernel is exact pathwise: once perturbed scores are formed, the global maximizer is exactly the maximizer of the tile-local maxima. Grouped, online, and distributed variants are exact in distribution: they rely on hierarchical factorization through group log-masses. 4.1 Group-Gumbel-Max: Hierarchical Exact Sampling Partition [V][V] into m disjoint groups kk=0m−1\G_k\_k=0^m-1; the groups need not have equal size. For any group with at least one finite transformed logit, define Lk=log∑i∈kexp(ℓ~i)=logsumexp(ℓ~k).L_k\;=\; _i _k ( _i)\;=\;logsumexp( _G_k). If a group contains no finite transformed logit, then Lk=−∞L_k=-∞, the group has zero probability mass, and it can be skipped. After discarding zero-mass groups, the categorical distribution factorizes as ℙ(K=k)⏟choose group∝exp(Lk),ℙ(I=i∣K=k)⏟choose within group∝exp(ℓ~i)for i∈k. P(K=k)_choose group (L_k), P(I=i K=k)_choose within group ( _i) i _k. Thus exact sampling from the full categorical can be implemented by first choosing a group using the logits Lk\L_k\ and then sampling within the chosen group. Parallel FlashSampling. Suppose logits arise from a linear projection = y= W x, where ∈ℝV×D W ^V× D and ∈ℝD x ^D. Let k∈ℝ|k|×D W_G_k ^|G_k|× D be the block of rows indexed by group kG_k, so k=k∈ℝ|k| y_k= W_G_k x ^|G_k| are the group logits. Parallel FlashSampling computes groups independently: each group with nonzero mass computes (i) an exact local sample zk∼Cat(softmax(k))z_k (softmax( y_k)) and (i) its group log-mass Lk=logsumexp(k)L_k=logsumexp( y_k). The algorithm then samples K∼Cat(softmax())K (softmax( L)) and returns zKz_K mapped to its global index. This is exact by direct factorization. Online FlashSampling. When memory is the primary constraint, FlashSampling can stream groups one at a time and maintain only a running log-mass and a running sample. Suppose the current running state is (Lrun,z)(L_run,z) and the next nonzero-mass group has log-mass LkL_k and exact local sample zkz_k. Define Lnew=log(eLrun+eLk).L_new= (e^L_run+e^L_k ). Then replace z by zkz_k with probability eLkeLrun+eLk=eLk−Lnew=11+eLrun−Lk, e^L_ke^L_run+e^L_k=e^L_k-L_new= 11+e^L_run-L_k, and otherwise keep z. Section 4.4 proves that this binary merge rule preserves exactness by induction. 4.2 Distributed FlashSampling for Tensor-Parallel Vocabularies In tensor-parallel LM heads, the vocabulary dimension is sharded across n GPUs. Naively, each GPU computes local logits and then an all-gather concatenates the full V logits before sampling, incurring communication proportional to the vocabulary size per row. FlashSampling treats shards as groups: each rank returns (i) a local exact sample from its shard, if its shard has nonzero mass for that row, and (i) the shard log-mass LkL_k. A final exact categorical sample over the shard log-masses chooses which rank provides the global sample. Communication therefore scales with the number of shards, not the number of vocabulary entries. 4.3 A Unifying View: Max-Stability of Grouped Gumbel Perturbations Group-Gumbel-Max and FlashSampling both rely on the same structural fact: max decomposes over partitions. For grouped variants we additionally use the max-stability of Gumbel perturbations. Lemma 4.1 (Gumbel max-stability under grouping). Let gii=1V\g_i\_i=1^V be i.i.d. Gumbel(0,1)Gumbel(0,1) and let kk=0m−1\G_k\_k=0^m-1 be a partition of [V][V]. Assume each group under discussion contains at least one finite transformed logit. Define Mk=maxi∈k(ℓ~i+gi),Ik=argmaxi∈k(ℓ~i+gi),Lk=log∑i∈keℓ~i.M_k\;=\; _i _k( _i+g_i), I_k\;=\; *arg\,max_i _k( _i+g_i), L_k\;=\; _i _ke _i. Then: 1. Mk∼Gumbel(Lk,1)M_k (L_k,1), 2. Mk\M_k\ are independent across disjoint groups, 3. ℙ(Ik=i)=eℓ~i/∑j∈keℓ~jP(I_k=i)=e _i/ _j _ke _j for i∈ki _k. Proof 4.2. For any real t, ℙ(Mk≤t)=∏i∈kℙ(gi≤t−ℓ~i)=∏i∈kexp(−e−(t−ℓ~i))=exp(−e−(t−Lk)),P(M_k≤ t)= _i _kP(g_i≤ t- _i)= _i _k (-e^-(t- _i) )= (-e^-(t-L_k) ), which is the CDF of Gumbel(Lk,1)Gumbel(L_k,1). Independence follows because the groups are disjoint and the underlying Gumbels are independent. The within-group argmax probabilities are exactly the Gumbel-Max trick applied to the restricted transformed logits. Consequence. For grouped variants, selecting a group by argmaxkMk *arg\,max_kM_k is equivalent in distribution to applying Gumbel-Max directly to the group logits Lk\L_k\. The outer group sample may therefore use fresh independent Gumbels, or it may reuse explicitly computed group maxima. For the fused two-stage kernel in Algorithm 2, exactness does not rely on max-stability: once the perturbed scores xi=ℓ~i+gix_i= _i+g_i have been formed, exactness is simply the deterministic identity maxixi=maxtmaxi∈txi. _ix_i= _t _i _tx_i. 4.4 Exactness of Group-Gumbel-Max The correctness of grouped FlashSampling rests on two facts: exact group factorization, and the binary merge rule used by the online variant. Lemma 4.3 (Exact group factorization). Let [V][V] be partitioned into groups kk=0m−1\G_k\_k=0^m-1, and discard any zero-mass groups. Define Lk=log∑i∈kexp(ℓ~i)L_k= _i _k ( _i). If we sample K∼Cat(softmax())K (softmax( L)) and then sample I∣(K=k)∼Cat(softmax(ℓ~k))I (K=k) (softmax( _G_k)), the marginal distribution of I equals Cat(softmax(ℓ~))Cat(softmax( )). Proof 4.4. For any i∈ki _k, ℙ(I=i)=ℙ(K=k)ℙ(I=i∣K=k)=eLk∑seLs⋅eℓ~i∑j∈keℓ~j=eℓ~i∑j=1Veℓ~j.P(I=i)=P(K=k)\,P(I=i K=k)= e^L_k _se^L_s· e _i _j _ke _j= e _i _j=1^Ve _j. Lemma 4.5 (Binary merge rule). Let A,B⊆[V]A,B [V] be disjoint and suppose both have nonzero mass. Define LA=log∑i∈Aeℓ~i,LB=log∑i∈Beℓ~i.L_A= _i∈ Ae _i, L_B= _i∈ Be _i. Suppose ZA∼Cat(softmax(ℓ~A))Z_A (softmax( _A)), ZB∼Cat(softmax(ℓ~B))Z_B (softmax( _B)), and an independent Bernoulli choice selects B with probability eLB/(eLA+eLB)e^L_B/(e^L_A+e^L_B). Returning ZBZ_B when B is selected and ZAZ_A otherwise yields an exact sample from Cat(softmax(ℓ~A∪B))Cat(softmax( _A∪ B)). Proof 4.6. For any i∈Ai∈ A, ℙ(Z=i)=ℙ(choose A)ℙ(ZA=i)=eLAeLA+eLB⋅eℓ~i∑j∈Aeℓ~j=eℓ~i∑j∈A∪Beℓ~j.P(Z=i)=P(choose A)\,P(Z_A=i)= e^L_Ae^L_A+e^L_B· e _i _j∈ Ae _j= e _i _j∈ A∪ Be _j. The same calculation for i∈Bi∈ B gives ℙ(Z=i)=eLBeLA+eLB⋅eℓ~i∑j∈Beℓ~j=eℓ~i∑j∈A∪Beℓ~j.P(Z=i)= e^L_Be^L_A+e^L_B· e _i _j∈ Be _j= e _i _j∈ A∪ Be _j. Hence Z∼Cat(softmax(ℓ~A∪B))Z (softmax( _A∪ B)). Theorem 4.7 (Exactness of hierarchical FlashSampling). Algorithms B.2, B.3, and B.4 return an exact sample from Cat(softmax(ℓ~))Cat(softmax( )). Proof 4.8. For the parallel and distributed variants, Lemma 4.3 shows that it suffices to sample the group or shard index from logits Lk\L_k\ and then sample within the chosen group; both steps are exact. For the online variant, initialize with an exact sample from the first nonzero-mass group. Each subsequent update merges the current union with the next nonzero-mass group using Lemma 4.5. An induction over the streamed groups therefore yields an exact sample from the full categorical distribution. 4.5 Exactness of Tile-Wise FlashSampling Reduction FlashSampling also relies on a simpler structural lemma: the global maximum equals the maximum of the tile-local maxima. Lemma 4.9 (Max over vocabulary tiles). Let xii=1V\x_i\_i=1^V be real numbers and let ss=0ntile−1\T_s\_s=0^n_tile-1 be a partition of [V][V] into vocabulary tiles. For each tile, define ms=maxi∈sxi,ı^s∈argmaxi∈sxi,m_s= _i _sx_i, _s∈ *arg\,max_i _sx_i, where ı^s _s is a global index in sT_s. Then maxi∈[V]xi=maxsms. _i∈[V]x_i= _sm_s. Moreover, for any s⋆∈argmaxsmss ∈ *arg\,max_sm_s, the chosen index ı^s⋆ _s is a global maximizer. Conversely, every global maximizer lies in some tile s⋆∈argmaxsmss ∈ *arg\,max_sm_s. Proof 4.10. The identity for the maximum value is immediate: maxi∈[V]xi=maxsmaxi∈sxi=maxsms. _i∈[V]x_i= _s _i _sx_i= _sm_s. If s⋆∈argmaxsmss ∈ *arg\,max_sm_s, then xı^s⋆=ms⋆=maxixix_ _s =m_s = _ix_i, so ı^s⋆ _s is a global maximizer. Conversely, if i⋆i is any global maximizer, then its tile s⋆s satisfies ms⋆=xi⋆=maxixim_s =x_i = _ix_i, hence s⋆∈argmaxsmss ∈ *arg\,max_sm_s. Applying Lemma 4.9 to xi=ℓ~i+gix_i= _i+g_i justifies the two-stage fused design in Algorithm 2. Because the Gumbel variables are continuous, the global maximizer is unique almost surely, so the tile-wise reduction returns exactly the same index as a full row-wise argmax *arg\,max with probability one. 4.6 Top-k, Nucleus Sampling, and Masking Practical decoding often uses truncated supports, and the tiled structure of FlashSampling naturally accommodates most of them. • Top-k: The Group-Gumbel-Max decomposition extends directly to top-k via the Gumbel-Top-k trick (pmlr-v97-kool19a). Each tile computes top-k candidates locally (logits and indices), and a second stage reduces all per-tile candidates into a global top-k. Sampling from the final k candidates can be done with multinomial or Gumbel-Max sampling. • Top-p (nucleus): Unlike top-k, nucleus sampling (Holtzman2020The) requires a global softmax followed by a sorted cumulative sum, neither of which decomposes into independent tile-local work. However, top-p can be applied after top-k on the reduced candidate set of only k elements, where softmax, sorting, and cumulative summation are negligible. This sequential top-k-then-top-p strategy is used in practice by vLLM***https://github.com/vllm-project/vllm/blob/v0.16.0/vllm/v1/sample/ops/topk_topp_sampler.py#L264-L279,†https://github.com/vllm-project/vllm/blob/v0.16.1rc0/vllm/v1/sample/ops/topk_topp_triton.py#L956, FlashInfer‡https://github.com/flashinfer-ai/flashinfer/blob/v0.6.3/flashinfer/sampling.py#L1069-L1072, and other SOTA top-k top-p algorithms (park2026qritahighperformancetopktopp). • Masking: Forbidden indices (e.g. banned tokens, grammar constraints) are supported by setting their logits to −∞-∞ before perturbation, which preserves exactness over the restricted support. While the FlashSampling theory allows integrating these sampling strategies, we leave the implementation to future work. 4.7 Cost Model: Bandwidth, Kernels, and Overhead We outline a simple model to reason about speedups. Materialized baseline (lower bound). For a BF16 baseline that materializes logits, the GEMM must at least read W and H and write Y once; sampling must then read Y at least once again. An optimistic lower bound on arithmetic intensity is therefore Imat(B)≈2BVD2(VD+BD+2BV)=BVDVD+BD+2BVFLOP/byte,I_mat(B)\;≈\; 2BVD2(VD+BD+2BV)= BVDVD+BD+2BV /byte, where the denominator counts mandatory BF16 traffic only. Real softmax-based samplers usually make more than one pass over the materialized logits, so the true baseline intensity is lower. Fused matmul + sampling. If sampling is fused into the GEMM epilogue so that the logits write and reread are removed, then, up to lower-order terms from the small candidate buffer, Ifused(B)≈2BVD2(VD+BD)=BV+BFLOP/byte.I_fused(B)\;≈\; 2BVD2(VD+BD)= BVV+B /byte. Thus fusion raises the effective arithmetic intensity. Incremental traffic saved by fusion. Relative to a fused kernel, any materialized baseline incurs at least one write and one reread of the [B,V][B,V] logits tensor. In BF16 this minimal extra traffic is 4BV4BV bytes. Compared with the mandatory LM-head weight read of 2VD2VD bytes, the extra fraction is 4BV2VD=2BD. 4BV2VD= 2BD. For the small configuration (D=4096D=4096), this ratio is 0.049%0.049\% at B=1B=1, 3.125%3.125\% at B=64B=64, and 6.25%6.25\% at B=128B=128. Thus raw logits-byte savings alone are too small to explain the largest measured speedups. The main gains come from eliminating extra sampling kernels, global-memory round-trips through those kernels, and their launch and synchronization overhead. In the memory-bandwidth-bound decode regime, these extra kernels are pure overhead. At B=1B=1 on the small configuration, the minimal avoided logits round-trip is 4BV=4⋅1⋅151,936=607,744 bytes≈0.608 MB.4BV=4· 1· 151,936=607,744 bytes≈ 0.608 MB. At 88 TB/s, this corresponds to only 7.6×10−57.6× 10^-5 ms. The observed latency gap therefore cannot be explained by raw HBM bandwidth alone. 5 Experiments We evaluate FlashSampling at two levels: kernel-level microbenchmarks that isolate fused matmul-plus-sample across four GPU architectures, and end-to-end vLLM integration that measures autoregressive decode latency. All benchmarks use the open-source FlashSampling Triton implementation (ruiz_fmms_repo). 5.1 Setup Hardware. Kernel microbenchmarks are run on four NVIDIA GPUs spanning two architecture generations. Table 2 summarizes their specifications. All GPUs are provisioned via Modal cloud. Table 2: GPU specifications. Peak BF16 TFLOP/s are dense (without structured sparsity), since the LM-head matmul is a dense GEMM. The ops:byte ratio (peak compute / bandwidth) contextualizes the crossover between bandwidth- and compute-limited regimes, although the exact crossover is kernel-dependent. H100 H200 B200 B300 Architecture Hopper Hopper Blackwell Blackwell HBM capacity (GB) 80 141 192 288 HBM bandwidth (TB/s) 3.35 4.8 8.0 8.0 Peak BF16 dense (TFLOP/s) 989 989 2,250 2,250 Ops:byte ratio 295 206 281 281 Software. PyTorch 2.10.0, CUDA 13.0, Triton 3.6, and FlashInfer 0.6.3. All kernels are warmed up for 25 iterations before timing. Workload configuration. The main text focuses on the decode-centric configuration D=4,096,V=151,936,D=4,096, V=151,936, which matches models such as Qwen3-8B and Qwen3-235B-A22B MoE. We sweep batch sizes B∈1,2,4,8,16,32,64,128,256B∈\1,2,4,8,16,32,64,128,256\. Additional results for a larger configuration show the same qualitative trends (Appendix A). Baselines. 1. Multinomial Sampling. This baseline materializes the logits using a matmul (cuBLAS), followed by sampling with softmax and multinomial. We apply torch.compile to it, which improves speed by 14% on average over PyTorch eager (range: 7–30% across GPUs and batch sizes). Unless explicitly stated, all references to Multinomial Sampling refer to the compiled version. 2. FI1 (FlashInfer top-k/top-p). top_k_top_p_sampling_from_logits§https://docs.flashinfer.ai/api/sampling.html, a sampling kernel used by vLLM for top-k/top-p decode. Logits are also materialized using cuBLAS. 3. FI2 (FlashInfer Gumbel-Max). sampling_from_logits§ ‣ 2, FlashInfer’s exact Gumbel-Max sampler on pre-materialized logits. Logits materialized using cuBLAS. 5.2 Standalone Logits Sampling Standalone FlashSampling applies Gumbel-Max to pre-materialized logits. This is algorithmically close to FI2, which also uses Gumbel-Max on materialized logits. We therefore focus on the fused setting, which is the primary systems contribution: FlashSampling’s advantage comes from eliminating the logits materialization and the sampling pass. 5.3 Fused Matmul and Sampling Table 3 reports FlashSampling speedups relative to the three baselines (D=4096D=4096, V=151kV=151k). All numbers are median latency over 100 timed iterations. Table 3: FlashSampling speedup vs. three baselines (D=4096D=4096, V=151kV=151k). Values >1>1 indicate FlashSampling is faster; bold marks the peak per GPU within each baseline. FI1: FlashInfer top-k/top-p kernel. FI2: FlashInfer Gumbel-Max kernel. vs. Multinomial Sampling vs. FI1 (top-k/top-p) vs. FI2 (Gumbel-Max) B H100 H200 B200 B300 H100 H200 B200 B300 H100 H200 B200 B300 1 1.25 1.28 1.46 1.46 1.30 1.35 1.51 2.45 1.18 1.20 1.32 1.31 2 1.24 1.26 1.46 1.43 1.30 1.38 1.56 2.18 1.16 1.17 1.30 1.28 4 1.24 1.27 1.47 1.44 1.32 1.41 1.61 2.37 1.16 1.18 1.32 1.30 8 1.27 1.31 1.53 1.50 1.33 1.44 1.61 2.46 1.17 1.19 1.33 1.30 16 1.30 1.37 1.57 1.54 1.35 1.45 1.65 2.47 1.18 1.21 1.36 1.33 32 1.38 1.47 1.68 1.64 1.37 1.47 1.68 2.47 1.20 1.22 1.38 1.34 64 1.60 1.60 1.84 1.81 1.46 1.47 1.67 2.52 1.26 1.21 1.39 1.37 128 1.64 1.37 1.89 1.98 1.35 1.12 1.55 2.28 1.12 0.91 1.27 1.33 256 1.29 1.17 1.58 1.65 1.04 0.98 1.32 2.04 0.81 0.78 1.07 1.09 Figure 2: Relative performance on B300. Left: FlashSampling vs. the Multinomial Sampling (baseline =1=1). Right: FlashSampling vs. FlashInfer FI1 and FI2 (baseline =1=1). FlashSampling is faster than the Multinomial Sampling across all shown batch sizes, faster than FI1 throughout, and faster than FI2 in the decode regime. Key observations. 1. FlashSampling is consistently faster in the decode regime. For B≤64B≤ 64, FlashSampling is faster than all three baselines on all four GPUs. In this regime, the peak speedup vs. Multinomial Sampling is 1.84×1.84× and the peak speedup vs. FI1 is 2.52×2.52×. 2. The gain is primarily from fusion. Speedups over FI2 are smaller than speedups over Multinomial Sampling or FI1 because FI2 already uses Gumbel-Max. The remaining gain therefore comes mainly from eliminating logits materialization and sampling overhead (Section 5.4). 3. The advantage narrows at larger batch sizes. As batch size grows, GEMM efficiency matters more and the workload becomes less dominated by memory-bandwidth-bound postprocessing. The larger-configuration appendix shows the same qualitative trend, with the crossover occurring earlier. 5.4 Interpreting the Batch-Size Trend The cost model in Section 4.7 showed that HBM savings from avoiding the logits write and reread alone are small (≤6%≤6\% of traffic). Figure 3 reveals a larger effect: the baselines’ separate sampling kernels are expensive, and their runtime grows steeply with batch size, while FlashSampling absorbs sampling into the matmul at negligible cost (Table 4: 2–6% of kernel time). Eliminating these separate kernels is the primary source of speedup. The advantage narrows at large batch sizes because FlashSampling’s Triton matmul becomes less efficient than cuBLAS (right panel), partially offsetting the sampling savings. Note that Triton is platform-agnostic (AMD, Intel GPUs, etc.), so the cuBLAS gap is a trade-off for portability. Profiling was performed on an RTX 3090 using Nsight Compute and Proton. Table 4: Sampling as a percentage of total kernel time. A high fraction spent on sampling rather than matmul is an indicator of inefficient sampling implementation. FlashSampling’s sampling fraction stays low because it is fused into the matmul epilogue; the baselines’ fraction grows with batch size B. Bold marks the highest sampling fraction for each method. FlashSampling Multinomial Sampling FI2 (Gumbel-Max) B matmul (%) sampl. (%) matmul (%) sampl. (%) matmul (%) sampl. (%) 1 97.7 2.3 93.7 6.3 94.7 5.3 16 97.7 2.3 87.1 12.9 93.4 6.6 64 93.6 6.0 71.3 28.7 88.6 11.4 256 93.4 6.2 73.1 26.9 88.2 11.8 Figure 3: Sampling runtime (left) and matmul runtime (right) in μ vs. batch size. Lower is better. 5.5 Roofline Analysis and Bandwidth Utilization The LM-head projection is memory-bandwidth-bound at small batch sizes because arithmetic intensity equals B (the weight matrix dominates traffic). Figure 4 confirms this on H100. Figure 4: Roofline (left) and HBM bandwidth utilization (right) on H100. Left: all methods track the memory-bound slope for B≤64B≤ 64; FlashSampling sits slightly above baselines because it avoids the logits round-trip. Close to the ridge point (AI≈295AI≈ 295), performance flattens below the compute ceiling, where cuBLAS outperforms Triton. Right: FlashSampling achieves higher bandwidth utilization than all baselines in the decode regime, confirming that fusion removes overhead rather than shifting it. Appendix D shows the same pattern on B200. 5.6 End-to-End vLLM Evaluation In this section, we demonstrate the end-to-end speedups achieved by FlashSampling on LLM inference. We integrate FlashSampling into vLLM (kwon2023efficient) by replacing the LM-head projection and the sampling step. We benchmark TPOT using problems from the AIME22-24 dataset¶https://huggingface.co/datasets/AI-MO/aimo-validation-aime. vLLM uses continuous batching, so the effective batch size varies dynamically during serving. We use vllm bench sweep serve with --max-concurrency=B to implement the batch size, and set --request-rate=B for requests to follow a Poisson process at B requests per second. We rerun the benchmark 5 times for each batch size, compare TPOT between baseline and FlashSampling, and report the median TPOT reduction across the 5 runs. Experiments run on a single B200 GPU with four models spanning a range of sizes and architectures. Key observation: The speedups are proportional to the decoding time spent on the LM head compared to attention and FFN. This explains the highest speedups on Qwen3-1.7B, which sees up to 19%19\% TPOT reduction. For Qwen3-32B and gpt-oss-120b, attention and FFN layers dominate decode time, so the speedups are smaller. Table 5: TPOT speedup (%) computed as (1−FlashSampling/baseline1-FlashSampling/baseline), and standard deviation across 5 runs. B is the maximum number of concurrent requests. Bold marks the peak per model. B Qwen3-1.7B Qwen3-8B Qwen3-32B gpt-oss-120b 1 10.8±0.310.8± 0.3 % 2.8±0.92.8± 0.9 % 1.9±1.41.9± 1.4 % 0.3±0.80.3± 0.8 % 2 12.8±0.612.8± 0.6 % 6.9±1.76.9± 1.7 % −1.8±0.5-1.8± 0.5 % 1.7±0.51.7± 0.5 % 4 14.5±0.414.5± 0.4 % 3.7±0.13.7± 0.1 % 1.6±0.01.6± 0.0 % 2.4±0.32.4± 0.3 % 8 15.2±0.415.2± 0.4 % 3.6±0.43.6± 0.4 % 1.1±0.11.1± 0.1 % 2.1±0.12.1± 0.1 % 16 17.7±9.017.7± 9.0 % 3.0±0.43.0± 0.4 % 1.1±0.11.1± 0.1 % 1.8±0.41.8± 0.4 % 32 9.7±5.19.7± 5.1 % 4.4±1.74.4± 1.7 % 1.3±0.31.3± 0.3 % 1.5±2.01.5± 2.0 % 64 18.7±6.818.7± 6.8 % 5.2±2.25.2± 2.2 % 1.2±0.31.2± 0.3 % 1.6±0.81.6± 0.8 % Figure 5: TPOT vs. concurrency on B200 for all four models. Top row: Qwen3-1.7B (up to 19%19\% reduction) and Qwen3-8B (roughly 33–7%7\%). Bottom row: Qwen3-32B and gpt-oss-120b, where gains are smaller because attention and FFN dominate decode time. 5.7 Empirical Correctness Verification Kernel Level: To verify sampling correctness, we compare samples from FlashSampling to the reference PyTorch implementation using a chi-squared goodness-of-fit test on 5,000 samples, and find no statistically significant difference. End-to-end Level: We run FlashSampling on 1,319 questions from the GSM8K dataset using Qwen3-1.7B and check the answers with a LLM judge. FlashSampling achieves 89.4%89.4\% accuracy versus 89.6%89.6\% for the baseline. This difference is not statistically significant (p=0.776), according to a paired bootstrap test. This is consistent with exact sampling. One cannot use greedy sampling here, since it would disable FlashSampling. 6 Related Work Gumbel-Max and Extensions. The Gumbel-Max trick for exact categorical sampling dates to gumbel1954statistical and was formalized by maddison2014astar. jang2017gumbelsoftmax introduced the Gumbel-Softmax relaxation for differentiable discrete sampling, which complements our focus on exact sampling. huijben2022review surveys the broader Gumbel-Max literature. pmlr-v97-kool19a extend the trick to top-k sampling without replacement, and qi2020fastgumbel study fast Gumbel variate generation. ahmed2026entropyaligneddecodinglmsbetter modify the sampling distribution via entropy-aware reweighting and use Gumbel-Max as a subroutine. FlashSampling contributes a systems-oriented hierarchical decomposition for exact online and distributed sampling in LLM inference, preserving the original distribution exactly. IO-Aware Kernel Fusion. FlashAttention (dao2022flashattention) showed that avoiding materialization of the attention matrix can substantially reduce HBM traffic, with subsequent work improving parallelism (dao2024flashattention) and exploiting hardware asynchrony (shah2024flashattention). Cut Your Losses (wijmans2025cutyourlosses), Liger Kernel (hsu2025ligerkernel), and dong2025projectionpredictionlogitsscalable apply the same idea to training-time cross-entropy by fusing the LM-head matmul with the loss computation. The same matmul-plus-epilogue fusion pattern appears in MLP layers (zhang2026deepkernelfusiontransformers), RNNs (poppel2025flashrnn), and whole-model inference (nrusimha2025flashformerwholemodelkernelsefficient). At the compiler level, EVT (chen2024evt) auto-generates fused GEMM epilogues via CUTLASS, and samaga2025fasterapproxtopkharnessing fuse approximate top-k selection into the matmul on TPUs. FlashSampling applies this methodology to a different domain: inference-time sampling, exploiting domain-specific structure (Gumbel-Max decomposability), and achieving exactness (no approximations). Efficient LLM Sampling FlashInfer (ye2025flashinfer) provides optimized GPU kernels for attention and sampling in LLM serving, including sorting-free rejection sampling for top-k/top-p. Qrita (park2026qritahighperformancetopktopp) achieves 2× throughput over prior sampling kernels via pivot-based selection. Min-p sampling (minh2025turning) proposes a dynamic truncation method that, like top-p, requires probability computation before truncation. SIMPLE (zhao2025simpledisaggregatingsamplinggpu) offloads sampling to the CPU, motivated by the same bottleneck FlashSampling addresses. Sampled softmax (rawat2019sampled) reduces large-vocabulary cost by computing the loss over a random subset, trading exactness for speed. All these methods operate on pre-materialized logits, while FlashSampling avoids materializing them entirely and introduces no approximation. 7 Conclusion We presented FlashSampling, a simple fused design for exact categorical sampling that avoids materializing the [B,V][B,V] logits tensor in HBM. The key ideas are straightforward: exact sampling does not require an explicit softmax, the fused tiled kernel is exact by argmax *arg\,max decomposition over vocabulary tiles, and grouped log-masses yield exact online and distributed variants. The method introduces no approximation: it produces exact samples from the target categorical distribution. Empirically, FlashSampling is most effective in the memory-bandwidth-bound decode regime, where it removes pure sampling overhead and turns sampling into a lightweight epilogue. Acknowledgement We sincerely thank Yongye Zhu, Zhuoqing Song, and Mayank Mishra for their helpful discussions and constructive feedback. We used large language models to assist in polishing the writing of this work. References [section] [section]l1 Appendix A Additional Kernel Results for the Large Configuration For completeness, Table 6 reports the larger-configuration kernel results deferred from the main text. The same qualitative pattern appears: FlashSampling is strongest in the small-batch decode regime, while the advantage narrows once the workload becomes more GEMM-efficiency dominated. Table 6: FlashSampling speedup vs. three baselines on the larger configuration (D=8192D=8192, V=128kV=128k). Values >1>1 indicate FlashSampling is faster; bold marks the peak per GPU within each baseline. At B≥128B≥128 the advantage narrows and cuBLAS GEMM efficiency becomes increasingly important. vs. Multinomial Sampling vs. FI1 (top-k/top-p) vs. FI2 (Gumbel-Max) B H100 H200 B200 B300 H100 H200 B200 B300 H100 H200 B200 B300 1 1.22 1.22 1.43 1.39 1.21 1.17 1.31 1.51 1.13 1.09 1.20 1.18 2 1.20 1.19 1.38 1.35 1.18 1.19 1.32 1.77 1.12 1.08 1.20 1.17 4 1.20 1.19 1.34 1.33 1.18 1.20 1.23 1.75 1.12 1.09 1.13 1.13 8 1.21 1.21 1.32 1.30 1.21 1.22 1.27 1.72 1.13 1.09 1.13 1.12 16 1.23 1.24 1.35 1.34 1.22 1.23 1.30 1.77 1.14 1.11 1.13 1.13 32 1.25 1.31 1.42 1.39 1.22 1.28 1.34 1.79 1.14 1.16 1.18 1.16 64 1.36 1.39 1.54 1.52 1.28 1.30 1.42 1.88 1.18 1.17 1.26 1.25 128 1.29 0.99 1.44 1.49 1.14 0.86 1.26 1.64 1.03 0.76 1.12 1.14 256 0.88 0.76 1.15 1.17 0.80 0.72 1.08 1.44 0.70 0.64 0.92 0.90 Appendix B FlashSampling Algorithm Pseudocode This appendix collects detailed pseudocode for the FlashSampling variants described in the main text. Streaming Gumbel-Max (standalone logits). Algorithm B.1 presents the basic one-pass streaming Gumbel-Max sampler over pre-materialized logits. Algorithm B.1 Gumbel-Max sampling (standalone logits): streaming argmax over perturbed logits 1:Logits ℓ∈ℝV ^V, RNG state 2:Sample index i⋆∈1,…,Vi ∈\1,…,V\ 3:m←−∞m←-∞, i⋆←1i ← 1 4:for i=1i=1 to V do 5: g←Gumbel(0,1)g← Gumbel(0,1) ⊳ via g=−log(−logu)g=- (- u), u∼Unif(0,1)u (0,1) 6: s←ℓi+gs← _i+g 7: if s>ms>m then 8: m←sm← s, i⋆←i ← i 9: end if 10:end for 11:return i⋆i Parallel Group-Gumbel-Max. Algorithm B.2 extends streaming Gumbel-Max to a group-parallel setting where each group is processed by an independent threadblock. Algorithm B.2 FlashSampling (parallel): Group-Gumbel-Max over groups 1:Input ∈ℝd x ^d, weight matrix ∈ℝd×V W ^d× V, group size g (so V=mgV=mg), RNG state 2:Sample index z∈1,…,Vz∈\1,…,V\ and optional log-normalizer ℓZ=logsumexp() _Z=logsumexp( y) 3:for k=0k=0 to m−1m-1 in parallel do 4: k←k⊤∈ℝg y_k← W_k x ^g 5: zk←argmaxj∈[g](yk,j−log(−loguk,j))z_k← *arg\,max_j∈[g] (y_k,j- (- u_k,j) ) ⊳ uk,j∼Unif(0,1)u_k,j\! \!Unif(0,1) 6: Lk←logsumexp(k)L_k ( y_k) 7:end for 8:k⋆←argmaxk∈[m](Lk−log(−logu¯k))k ← *arg\,max_k∈[m] (L_k- (- u_k) ) ⊳ u¯k∼Unif(0,1) u_k\! \!Unif(0,1) 9:z←k⋆g+zk⋆z← k g+z_k ⊳ map group-local index to global vocabulary index 10:ℓZ←logsumexp([L0,…,Lm−1]) _Z ([L_0,…,L_m-1]) ⊳ optional 11:return (z,ℓZ)(z, _Z) Sequential/online Group-Gumbel-Max. Algorithm B.3 provides a memory-efficient variant that streams groups one at a time. Algorithm B.3 FlashSampling (sequential/online): streaming Group-Gumbel-Max with O(g)O(g) working memory 1:Input ∈ℝd x ^d, weight matrix ∈ℝd×V W ^d× V, group size g (so V=mgV=mg), RNG state 2:Sample index z∈1,…,Vz∈\1,…,V\ and optional log-normalizer ℓZ _Z 3:Initialize with the first group. 4:0←0⊤∈ℝg y_0← W_0 x ^g 5:L0←logsumexp(0)L_0 ( y_0) 6:z0←argmaxj∈[g](y0,j−log(−logu0,j))z_0← *arg\,max_j∈[g] (y_0,j- (- u_0,j) ) ⊳ u0,j∼Unif(0,1)u_0,j\! \!Unif(0,1) 7:z←z0z← z_0, ℓ←L0 ← L_0 8:for k=1k=1 to m−1m-1 do 9: k←k⊤∈ℝg y_k← W_k x ^g 10: Lk←logsumexp(k)L_k ( y_k) 11: ℓnew←logsumexp([ℓ,Lk]) _new ([ ,\,L_k]) 12: preplace←exp(Lk−ℓnew)p_replace← (L_k- _new) ⊳ =eLkeℓ+eLk= e^L_ke +e^L_k 13: Draw u∼Unif(0,1)u (0,1) 14: if u<preplaceu<p_replace then 15: zk←argmaxj∈[g](yk,j−log(−loguk,j))z_k← *arg\,max_j∈[g] (y_k,j- (- u_k,j) ) ⊳ sample within selected group 16: z←kg+zkz← kg+z_k 17: end if 18: ℓ←ℓnew ← _new 19:end for 20:ℓZ←ℓ _Z← ⊳ optional 21:return (z,ℓZ)(z, _Z) Distributed Group-Gumbel-Max. Algorithm B.4 extends FlashSampling to tensor-parallel vocabularies sharded across multiple GPUs. Algorithm B.4 FlashSampling (distributed, tensor-parallel vocab): communicate O(1)O(1) scalars per rank 1:World size n. Rank k∈0,…,n−1k∈\0,…,n-1\ holds shard (k)∈ℝd×(V/n) W^(k) ^d×(V/n) covering vocab indices k⋅V/n+1,…,(k+1)⋅V/n\k· V/n+1,…,(k+1)· V/n\. Input ∈ℝd x ^d, RNG state. 2:Global sample index z∈1,…,Vz∈\1,…,V\ (and optional ℓZ _Z) 3:On each rank k: 4: compute local logits (k)←((k))⊤∈ℝV/n y^(k)←( W^(k)) x ^V/n 5: compute local log-mass Lk←logsumexp((k))L_k ( y^(k)) 6: sample local index z~k∼Cat(softmax((k))) z_k (softmax( y^(k))) ⊳ e.g., via Gumbel-Max / Group-Gumbel-Max / fused kernel 7:All-gather (Lk,z~k)k=0n−1\(L_k, z_k)\_k=0^n-1 to a coordinator (or perform an equivalent reduction) 8:Sample winning rank k⋆←argmaxk∈[n](Lk−log(−logu¯k))k ← *arg\,max_k∈[n] (L_k- (- u_k) ) ⊳ u¯k∼Unif(0,1) u_k\! \!Unif(0,1) 9:z←k⋆⋅(V/n)+z~k⋆z← k ·(V/n)+ z_k ⊳ convert rank-local index to global 10:Optionally ℓZ←logsumexp([L0,…,Ln−1]) _Z ([L_0,…,L_n-1]) 11:return z (and ℓZ _Z) Appendix C Numerically Stable and Fast Gumbel Generation Gumbel noise can be generated as g=−log(−logu)g=- (- u) with u∼Unif(0,1)u (0,1). In GPU kernels, two issues matter: • Numerical stability: avoid u=0u=0 or u=1u=1 which lead to infinities. • Throughput: the cost of generating random numbers and computing logs should not dominate. Practical recipe. Given a 32-bit RNG output r∈0,…,232−1r∈\0,…,2^32-1\, map to u=r+1232+1∈(0,1),u= r+12^32+1∈(0,1), then compute g=−log(−logu)g=- (- u). Many GPU RNG libraries (e.g. Philox, XORWOW) support generating floats in (0,1)(0,1) directly; the above mapping is a safe fallback. Approximate log options. If exactness in the distribution is required, the Gumbel generation must be statistically correct. However, using fast approximate log implementations can introduce small distortions. FlashSampling supports two modes: • Exact-math mode: use standard log for high fidelity. • Fast-math mode: use approximate logs for speed, with empirical validation that sampling bias remains negligible for target applications. The sampling remains algorithmically exact with respect to the generated Gumbels; any bias comes from numeric approximations. Appendix D Roofline and Bandwidth Utilization on B200 Figure 6 shows the roofline and bandwidth utilization on B200. The same pattern holds: FlashSampling tracks the memory-bound slope more closely and achieves higher bandwidth utilization than all baselines in the decode regime. Figure 6: Roofline (left) and HBM bandwidth utilization (right) on B200 (D=4096D=4096, V=151kV=151k). The pattern matches H100 (Figure 4): FlashSampling uses bandwidth more efficiently in the memory-bound decode regime and narrows at large batch sizes where cuBLAS GEMM efficiency dominates. Appendix E Returning Log-Normalizers or Max Values Some applications need logZ=log∑jeℓ~j Z= _je _j, for example to compute log-probabilities. The core FlashSampling sampler does not need logZ Z, but it can be added as an optional mode by accumulating a numerically stable log-sum-exp alongside sampling. In fused settings, this requires extra work in the epilogue, so we treat it as an optional feature rather than part of the core design.