SonicSampler: Unified Tile-Aware Kernels for LLM Sampling and Speculative Verification

# SonicSampler: Unified Tile-Aware Kernels for LLM Sampling and Speculative Verification
Source: [https://arxiv.org/html/2607.20475](https://arxiv.org/html/2607.20475)
\\contribution

\[\*\]Equal Contribution\\correspondence\\codeTBA

\(May 2026\)

###### Abstract

Sampling in LLM inference comprises a combinatorial set of logit processing, token selection, and verification operations for speculative decoding\. However, existing implementations either accelerate only subsets of this pipeline, rely on multiple kernel launches, or assume homogeneous sampling behavior across a batch, limiting support for dynamic serving workloads and preventing efficient CUDA Graph execution\.

We presentSonicSampler, a unified suite of tile\-aware Triton kernels that vertically fuses the complete sampling pipeline into a fixed, workload\-aware execution model\. Our kernels support dynamic per\-request sampling behaviors, including grammar\-constrained decoding, repetition, frequency and presence penalties, logit bias, temperature scaling, top\-kk/ top\-pp/ min\-ppfiltering, and speculative verification \- within a single batched kernel while remaining fully CUDA Graph\-compatible\. Central to our approach is a novel hierarchical two\-stage top\-kkalgorithm that achieves up to10x speedupover competitive baselines and exploits the low\-entropy structure of LLM outputs to enable efficient selection over large vocabularies\.

Across heterogeneous speculative decoding workloads, SonicSampler achieves up to16x speedupover state\-of\-the\-art baselines while preserving flexible batched execution\.

## 1Introduction

The performance of large language model \(LLM\) inference is increasingly critical across latency\-sensitive applications, from real\-time speech interfaces to agentic workflows where smaller models must execute within tight time budgets\(defossez2024moshi;belcak2025small\)\. While much of the systems effort has focused on accelerating the model forward pass, the effectiveness of LLM generation also depends on theefficiencyandflexibilityof sampling–the process of converting model logits into discrete token decisions\. Sampling is not merely a post\-processing detail; rather, it governs output diversity, controllability, and structural validity\(holtzman2020curious;zhu2023penalty;willard2023outlines;dong2025xgrammarflexibleefficientstructured\)\. It also enables quality improvement strategies such as best\-of\-NNdecoding, where multiple diverse candidates are drawn and selected to enhance the final quality of the output\(wang2022self;chen2023universal;li2025selfmoa;wang2025effect\)\. Improving LLM serving therefore requires not only faster model execution but also a sampling pipeline that is both efficient and expressive under realistic deployment constraints\.

This sampling bottleneck is particularly acute in speculative decoding\(leviathan2023fast;chen2023accelerating\), where lightweight draft models propose candidate continuations that must be sampled and verified against the target distribution\. As drafters grow increasingly lightweight, e\.g\. Medusa heads\(cai2024medusa\), EAGLE\(li2024eagle\), and multi\-token prediction\(gloeckle2024mtp\), the relative cost of the downstream sampling pipeline grows correspondingly\. In practice, this pipeline is not a single primitive but a heterogeneous composition of operations, including grammar\-constrained masking, repetition, frequency and presence penalties, logit bias, temperature scaling, top\-kk/ top\-pp/ min\-ppfiltering, stochastic or greedy token selection, and speculative verification\(fan2018hierarchical;holtzman2020curious;nguyen2024minp;hewitt2022truncation\)\.

Existing kernel implementations fail to address this complexity adequately\. Some accelerate isolated components such as grammar masking, penalty application, or probability filtering, but leave the overall pipeline fragmented across multiple kernel launches, incurring additional overhead and intermediate memory traffic\. Others partially fuse sampling, yet assume homogeneous behavior across the batch and, therefore, cannot support the mixed greedy and stochastic configurations that arise under continuous batching\(yu2022orca\)\. Critically, many top\-kkimplementations are also incompatible with CUDA Graph, forfeiting the launch\-overhead amortization that is increasingly essential in production serving engines\(yu2022orca;kwon2023vllm;nvidia2023tensorrtllm;zheng2024sglang\)\.

In this work, we presentSonicSampler, a unified suite of tile\-aware Triton\(tillet2019triton\)kernels that vertically fuse the complete sampling pipeline, from logit processing through speculative verification, into a CUDA Graph\-compatible execution model\. Our kernels support bothsingularmode \(standard and draft\-model inference\) andmultistepmode \(verification\), with fully dynamic per\-request configurations within batched dispatch via compact bit\-level indicators\. Central to our approach is a reformulation of the top\-kkbottleneck\. While filtering nominally requires global reductions across vocabularies of2172^\{17\}–2182^\{18\}tokens, practical next\-token distributions are highly concentrated; thus, the effective support relevant for sampling is far smaller than the full vocabulary\. This motivates a bounded candidate set withk=128k\{=\}128, which is sufficient whenever the probability mass outside the retained set is negligible for the chosen truncation rule\. We verify this empirically in[Section˜4\.5](https://arxiv.org/html/2607.20475#S4.SS5), showing that top\-128128retains nearly all relevant mass and preserves downstream accuracy\.

Leveraging this, we introduce atwo\-stage hierarchical top\-kkalgorithm\. Stage 1 tiles across vocabulary blocks, applying the full logit\-processing prologue before reducing each tile tokkcandidates via an adaptive radix or bitonic selection, encoded under a monotonicity\- and stability\-preserving lexicographic scheme\. Stage 2 gathers and merges candidates across blocks into the final top\-kk\. This map\-reduce formulation maximizes arithmetic intensity by fusing compute\-dense prologues and epilogues with the reduction stages\.

As such, we summarize our contributions as follows:

- •We present a unified CUDA Graph\-compatible kernel suite that fuses logit processing, sampling, and speculative verification while supporting mixed greedy and stochastic decoding within a single batched workload\.
- •We introduce a two\-stage tile\-aware hierarchical top\-kkalgorithm that achieves up to10xspeedup over existing radix\- and bitonic\-based alternatives\.
- •We demonstrate sampling speedups of up to𝟏𝟔​𝐱\\mathbf\{16x\}on speculative decoding workloads over FlashInfer and other state\-of\-the\-art baselines\.

## 2Related Work

Existing sampling systems can be categorized along three axes: kernel fusion, support for heterogeneous sampling, and compatibility with CUDA Graph\. While prior work explores different points in this design space, none jointly satisfy all three requirements\. We provide a focused discussion of representative approaches here, and defer a more comprehensive review to[Appendix˜A](https://arxiv.org/html/2607.20475#A1)and present a feature comparison in[Appendix˜B](https://arxiv.org/html/2607.20475#A2)\.

#### Standalone sampling kernels\.

Kernel libraries such as FlashInfer\(flashinfer2024\)and XGrammar\(dong2025xgrammarflexibleefficientstructured\)provide efficient implementations of key components in the sampling pipeline, including top\-kk/top\-ppselection and grammar\-constrained decoding\. These works establish strong building blocks for production systems, though they operate as separate kernels when composed together\.

#### Fused sampling kernels\.

Some recent approaches\(modular2025max\)introduce partial fusion of sampling within a single kernel\. However, these designs assume homogeneous sampling configurations and cannot natively support mixed greedy and stochastic decoding within a batched workload, limiting their applicability in production serving\.

Orthogonal efforts such as MegaKernel\(cheng2025miragepersistentkernelcompiler\)fuse the model forward pass into a single execution unit, reducing kernel launch overhead, and are complementary as they target model execution rather than sampling\.

FlashSampling\(ruiz2026flashsampling\)reduces memory traffic by fusing the LM head with downstream sampling, targeting an inference stage complementary to post\-logit processing\. It also briefly notes a two\-stage decomposition of top\-kkfor future work\. Our work instead focuses on the system design required to realize such a decomposition efficiently in practice\.

#### Efficient top\-kkselection\.

A separate line of work focuses on accelerating top\-kkselection itself, spanning multi\-pass radix selection\(radik2024;wang2025tilelangcomposabletiledprogramming\), bitonic\-based streaming kernels\(tillet2019triton\), and hybrid or distribution\-aware schemes\(flashinfer2024;park2026qritahighperformancetopktopp\)\. These methods treat selection as an isolated primitive and do not account for its tight coupling with upstream logit transformations and downstream sampling or verification steps, missing opportunities for vertical fusion\.

Taken together, these works provide efficient primitives and partial fusion strategies for LLM inference and sampling\. SonicSampler builds on these insights by unifying logit processing, sampling, and verification within a CUDA Graph\-compatible execution model that supports heterogeneous batched workloads\.

## 3SonicSampler

This section presents SonicSampler in detail\. Our design is guided by three principles: \(1\) exploit the inherent structure of the sampling workload to maximize fusion opportunities, \(2\) provide theoretical grounding for algorithmic simplifications, and \(3\) support the full heterogeneity of production serving within a single batched dispatch\.

### 3\.1Background

To ground the design of SonicSampler, we first introduce the LLM sampling pipeline and its constituent operations\. These operations define the computational structure that our execution model seeks to optimize\.

LLM sampling transforms model logits into discrete tokens through a sequence of logit processing, truncation, sampling, and \(optionally\) speculative verification steps\. Given logits𝐱∈ℝV\\mathbf\{x\}\\in\\mathbb\{R\}^\{V\}, these operations produce a probability distribution and select the next token\.

#### Logit Processing\.

We view logit processing as a sequence of transformations applied to the raw logits, including grammar masking, repetition and frequency penalties, logit bias, and temperature scaling\(keskar2019ctrl;dong2025xgrammarflexibleefficientstructured;hinton2015distilling\)\. These operations act independently on each token and, therefore, form amap\-stylecomputation that is fully parallelizable across the vocabulary\.

#### Probability Truncation\.

Truncation methods such as top\-kk, top\-pp, and min\-ppselect a subset of candidate tokens\(fan2018hierarchical;holtzman2020curious;nguyen2024minp\)\. In contrast to logit processing, these operations depend on comparisons across tokens \(e\.g\., ranking or cumulative probability\) and thus requireglobal reductionover the vocabulary\. Notably, we show in[equation˜12](https://arxiv.org/html/2607.20475#A3.E12)\(Appendix\) that min\-ppcan be reformulated to avoid explicit softmax\.

#### Top\-kkSelection Primitives\.

Ranking\-based truncation on GPUs is typically realized through general\-purpose top\-kkselection primitives\.*Radix\-based selection*partitions keys by successive digit groups from most to least significant, recursing into the bucket containing thekk\-th element at each pass\(alabi2012fast\)\.*Bitonic\-based selection*builds on bitonic sorting networks\(batcher1968sorting\), whose regular, branch\-free structure maps naturally onto SIMT execution and especially effective for tile\-local, fixed\-size reductions\(satish2009designing;shanbhag2018efficient\)\.

#### Sampling\.

Token selection can be expressed as anarg⁡max\\arg\\maxover perturbed logits via Gumbel\-Max reparameterization\(jang2016categorical\), unifying stochastic and greedy sampling under a single formulation:x∗=arg⁡maxi⁡\(xi\+ϵi\)x^\{\*\}=\\arg\\max\_\{i\}\(x\_\{i\}\+\\epsilon\_\{i\}\), whereϵi\\epsilon\_\{i\}introduces stochasticity\.

#### Speculative Verification\.

Speculative decoding accelerates inference by having a draft model propose candidate tokens, which are accepted whenu⋅q​\(d\)≤p​\(d\)u\\cdot q\(d\)\\leq p\(d\)foru∼𝒰​\(0,1\)u\\sim\\mathcal\{U\}\(0,1\), or, in the case of a rejection, sampled from the residual distributionp~​\(x\)∝max⁡\{0,p​\(x\)−q​\(x\)\}\\tilde\{p\}\(x\)\\propto\\max\\\{0,\\,p\(x\)\-q\(x\)\\\}\(leviathan2023fast\)\.

These components exhibit a characteristic structure, i\.e\. element\-wise transformations followed by global reductions, which motivates the workload\-aware execution model introduced next\. Detailed formulations are deferred to[Appendix˜C](https://arxiv.org/html/2607.20475#A3)\.

![Refer to caption](https://arxiv.org/html/2607.20475v1/figures/sonic-sampler-central.png)Figure 1:SonicSampler’s two\-stage tile\-aware fused sampling pipeline for the Singular \(non\-speculative\) workload\.Stage 1\(top, left\-to\-right\) vertically fuses grammar masking, repetition penalties, logit bias, and temperature scaling with Lexicographic Encoding and per\-tile Local Top\-kkvia Optimized Bitonic or Adaptive Radix\-Bitonic selection \(right panel\), writing only the packedkkresults per tile to the HBM scratchpad as the sole cross\-kernel memory boundary\.Stage 2\(bottom, right\-to\-left\) re\-indexes across tiles, performs global Top\-kkreduction through bitonic merge networks with ascending / descending alternating tile layout, applies probability truncation, and emits the selected token via Gumbel\-Noise and ArgMax, completing a clockwise data loop through the shared HBM band\. The Adaptive Radix\-Bitonic strategy \(right\) screens on the 6 most\-significant bits via a single\-pass optimized reduction tree, branching to a 2\-Pass Radix Select on the lower 10 bits or falling back to an Optimized Bitonic selection\.

### 3\.2Workload\-Aware Execution Model

Existing implementations dispatch each sampling stage as a separate kernel, incurring launch overhead and intermediate materialization, while the top\-kkreduction itself executes as a single\-program streaming pass over the full vocabulary, thereby limiting occupancy toBBconcurrent programs and precluding tiled execution\. We decompose the global reduction into a*two\-stage hierarchical*operation over vocabulary tiles of sizeBNB\_\{N\}:

1. 1\.Stage 1launchesB⋅ZvB\\cdot Z\_\{v\}programs \(whereZv=⌈VBN⌉Z\_\{v\}=\\left\\lceil\\frac\{V\}\{B\_\{N\}\}\\right\\rceil\), each fusing the logit\-processing prologue with a tile\-local top\-kkreduction that emitskkpacked candidates to a scratchpad in global memory\.
2. 2\.Stage 2launchesBBprograms, each gathering theZv⋅kZ\_\{v\}\\cdot kpacked candidates and performing a cross\-tile merge followed by a vertically fused epilogue of probability truncation, Gumbel perturbation, andarg⁡max\\arg\\maxselection\.

This decomposition increases Stage 1 parallelism by a factor ofZvZ\_\{v\}relative to single\-stage formulations while confining all inter\-tile communication to a compactB⋅Zv⋅kB\\cdot Z\_\{v\}\\cdot kscratchpad of packeduint32entries\. The fusion boundaries are chosen so that no intermediate logit or probability vector is ever materialized at vocabulary scale between the two stages\.

### 3\.3Entropy Sufficiency for Bounded Top\-kk

The two\-stage decomposition requires bounding the number of candidates retained from each tile\. We establish that for well\-trained LLMs, a modest bound ofk=128k=128suffices to capture the effective probability mass under the practical top\-ppthreshold\.

#### Theoretical Analysis\.

Consider a probability distribution overVVtokens where massPPis concentrated inKKcandidates \(the top\-kkset\) and mass\(1−P\)\(1\-P\)is distributed over the remainingV−KV\-Ktokens\. The Shannon entropy of such a distribution is bounded by:

H≤H\(2\)​\(P\)\+P​log2⁡K\+\(1−P\)​log2⁡\(V−K\)H\\leq H^\{\(2\)\}\(P\)\+P\\log\_\{2\}K\+\(1\-P\)\\log\_\{2\}\(V\-K\)\(1\)whereH\(2\)​\(P\)=−P​log2⁡P−\(1−P\)​log2⁡\(1−P\)H^\{\(2\)\}\(P\)=\-P\\log\_\{2\}P\-\(1\-P\)\\log\_\{2\}\(1\-P\)is the binary entropy\. This follows from the Schur\-concavity of entropy, where the uniform distribution of mass within each partition maximizes entropy\. Rearranging forKKwhenV≫KV\\gg K, we’d have:

K≥PV\(1/P\)−1⋅2H/P=P​V⋅PerplexityVPK\\geq\\frac\{P\}\{V^\{\(1/P\)\-1\}\}\\cdot 2^\{H/P\}=PV\\cdot\\sqrt\[P\]\{\\frac\{\\text\{Perplexity\}\}\{V\}\}\(2\)ForK=128K=128andV=217V=2^\{17\}, the maximum satisfiable entropy is therefore given by:

H≤17−P​\(10\+log2⁡P\)H\\leq 17\-P\(10\+\\log\_\{2\}P\)\(3\)Now, withP=0\.95P=0\.95\(a typical top\-ppthreshold\), this yieldsH​\\lesssim​7\.6H\\lesssim 7\.6bits\. In the limiting case ofτ→∞\\tau\\to\\inftywhereH→log2⁡VH\\to\\log\_\{2\}V, we would requireK≥P​VK\\geq PV, but this regime is practically irrelevant\.

#### Empirical Grounding\.

du2025temperatureshow that optimal temperatures for multi\-sample inference typically satisfyτ<1\.0\\tau<1\.0, where higher temperatures increase improper token selection\. Consequently, well\-trained LLMs produce peaked distributions with probability mass concentrated in a small candidate set\. This places practical sampling in the regimeH≤7H\\leq 7bits, makingk=128k=128sufficient\. We validate this in[Section˜4\.5](https://arxiv.org/html/2607.20475#S4.SS5), showing negligible residual mass beyond top\-kkand no loss in downstream accuracy on GPQA\-Diamond\.

### 3\.4Two\-Stage Hierarchical Top\-kk

We decompose vocabulary\-scale top\-kkselection into two stages: tile\-local reduction in the first stage, followed by cross\-tile merging in the second\. This decomposition maximizes parallelism by launching independent programs across vocabulary tiles, each fusing the logit processing prologue with local candidate selection\. The tile\-local reduction admits two strategies, namely an optimized bitonic network and an adaptive radix\-bitonic filter \(see[Figure˜1](https://arxiv.org/html/2607.20475#S3.F1)\), which we detail after establishing the encoding that underlies both\.

#### Numerical Lexicographic Encoding\.

To enable correct ordering across the sign boundary while preserving index stability, we encodebfloat16values into a lexicographically ordered unsigned integer representation\. Givenx∈𝔽16x\\in\\mathbb\{F\}\_\{16\}, letβ​\(x\)\\beta\(x\)denote its bit\-reinterpretation asuint16andσ​\(x\)=𝟏​\[β​\(x\)≥215\]\\sigma\(x\)=\\mathbf\{1\}\[\\beta\(x\)\\geq 2^\{15\}\]be the sign indicator\. The encoding is thus given by:

ϕ​\(x\)=β​\(x\)⊕μ​\(σ​\(x\)\)s\.t\.μ​\(s\)=\{215,if​s=0216−1,otherwise\\phi\(x\)=\\beta\(x\)\\oplus\\mu\(\\sigma\(x\)\)\\hskip 19\.91684pt\\text\{s\.t\.\}\\hskip 19\.91684pt\\mu\(s\)=\\begin\{cases\}2^\{15\},&\\text\{if \}s=0\\\\ 2^\{16\}\-1,&\\text\{otherwise\}\\end\{cases\}This maps non\-negative values to\[215,216−1\]\[2^\{15\},2^\{16\}\-1\]via MSB toggle \(preserving order\) and negative values to\[0,215−1\]\[0,2^\{15\}\-1\]via bitwise complement \(reversing order to maintain monotonicity with increasing magnitude\)\. The encoding is self\-dual:

ϕ−1​\(y\)=y⊕μ​\(1−σ′​\(y\)\)s\.t\.σ′​\(y\)=𝟏​\[y≥215\]\\phi^\{\-1\}\(y\)=y\\oplus\\mu\(1\-\\sigma^\{\\prime\}\(y\)\)\\hskip 19\.91684pt\\text\{s\.t\.\}\\hskip 19\.91684pt\\sigma^\{\\prime\}\(y\)=\\mathbf\{1\}\[y\\geq 2^\{15\}\]For stable sorting, we pack the encoded values in the upper\-16 bits with the inverted indices in the lower half, i\.e\. the finalized transform being given by:

ψ​\(x,i,n\)=ϕ​\(x\)⋅216\+\(n−1\)⊕i\\psi\(x,i,n\)=\\phi\(x\)\\cdot 2^\{16\}\+\(n\-1\)\\oplus iwherennis the tile size andi∈\{0,…,n−1\}i\\in\\\{0,\\ldots,n\-1\\\}is the position index\. The index complement ensures descending index priority among tied values, guaranteeing stability\. We also generalize this transform to enable alternating relative indices, detailed in[Appendix˜D](https://arxiv.org/html/2607.20475#A4)\.

#### Optimized Bitonic Top\-kkSelection\.

Our primary reduction employs bitonic networks for their regular, data\-independent, and highly parallelizable comparison patterns\.[Algorithms˜1](https://arxiv.org/html/2607.20475#alg1)and[3](https://arxiv.org/html/2607.20475#alg3)present the critical pieces of our innnovation while supplementary and encapsulating portions of the overall algorithm \(HypercubeMerge,BitonicSelection\) are detailed in[Appendix˜E](https://arxiv.org/html/2607.20475#A5)\. TheCASprimitive \([Algorithm˜1](https://arxiv.org/html/2607.20475#alg1)\) expresses the conditional swap as explicitmin\\min/max\\maxoperations followed by a conditional select, rather than the conventional bitwise\-XOR formulation\. This enables the compiler to emit fusedIMNMX\.U32\.U32instructions that perform each compare\-and\-swap in a single cycle on SM90\+ hardware\.

BitonicTopK\([Algorithm˜8](https://arxiv.org/html/2607.20475#alg8)\) operates on a hypercube view of the packed tile in two phases:κ=log2⁡k\\kappa=\\log\_\{2\}kmerge stages sort everykk\-width slice, then amax\\max\-reduction along the fold axis reduces the tile fromBNB\_\{N\}toR⋅kR\\cdot kcandidates within the register file\. The remainingFFfolds are performed byBitonicFold\([Algorithm˜3](https://arxiv.org/html/2607.20475#alg3)\), which, at each iteration, reshapes into\[R/2i,2,k\]\[R/2^\{i\},\\,2,\\,k\]and transposes to\[R/2i,k,2\]\[R/2^\{i\},\\,k,\\,2\]before splitting\. This reshape\-transpose\-split sequence forces the compiler to emit aconvert\_layoutthat places a complete pair ofkk\-element sub\-tiles within each warp, compiling the subsequentmax\\maxto a warp\-localIMNMXreduction tree with no inter\-warp barriers\. While this sequence inevitably triggers an exchange through shared memory, i\.e\., aSTS\-BAR\-LDSchain, this re\-layout is effectively only triggered once, while subsequent sequences share the same layout and are compiled away\. The order parameterδ=block\_idmod2\\delta=\\texttt\{block\\\_id\}\\bmod 2alternates sort direction across tiles, enabling Stage 2 to merge pre\-sorted alternating sub\-tiles viaBitonicReduction\([Algorithm˜6](https://arxiv.org/html/2607.20475#alg6)\) without re\-sorting\.

Algorithm 1Compare\-and\-Swap1:Packed

𝐯∈ℕ2n\\mathbf\{v\}\\in\\mathbb\{N\}^\{2^\{n\}\}\(hypercube\), pair

ρ∈\{0,1\}2\\rho\\in\\\{0,1\\\}^\{2\}, flip

ff, dimension

dd, total

nn
2:Swapped

𝐯′∈ℕ2n\\mathbf\{v\}^\{\\prime\}\\in\\mathbb\{N\}^\{2^\{n\}\}
3:

o←n−d−1o\\leftarrow n\-d\-1⊳\\trianglerightOuter axis

4:

𝐫←𝐯⊕XorReduce​\(𝐯,axis=o\)\\mathbf\{r\}\\leftarrow\\mathbf\{v\}\\oplus\\textsc\{XorReduce\}\(\\mathbf\{v\},\\text\{axis\}=o\)
5:

σ←f⊕Expand​\(ρ,o,d\)\\sigma\\leftarrow f\\oplus\\textsc\{Expand\}\(\\rho,o,d\)
6:return

Select​\(σ≠0,max⁡\(𝐯,𝐫\),min⁡\(𝐯,𝐫\)\)\\textsc\{Select\}\(\\sigma\\neq 0,\\ \\max\(\\mathbf\{v\},\\mathbf\{r\}\),\\ \\min\(\\mathbf\{v\},\\mathbf\{r\}\)\)

Algorithm 2Screen\-Combine1:

\(ma,ca\),\(mb,cb\)∈ℕ×ℕ\(m\_\{a\},c\_\{a\}\),\(m\_\{b\},c\_\{b\}\)\\in\\mathbb\{N\}\\times\\mathbb\{N\}
2:

\(m′,c′\)∈ℕ×ℕ\(m^\{\\prime\},c^\{\\prime\}\)\\in\\mathbb\{N\}\\times\\mathbb\{N\}
3:

m′←max⁡\(ma,mb\)m^\{\\prime\}\\leftarrow\\max\(m\_\{a\},\\,m\_\{b\}\)⊳\\trianglerightIMNMX

4:

c^←Select​\(ma=mb,ca\+cb,cb\)\\hat\{c\}\\leftarrow\\textsc\{Select\}\(m\_\{a\}\{=\}m\_\{b\},\\,c\_\{a\}\+c\_\{b\},\\,c\_\{b\}\)⊳\\trianglerightIADD

5:

c′←Select​\(ma\>mb,ca,c^\)c^\{\\prime\}\\leftarrow\\textsc\{Select\}\(m\_\{a\}\{\>\}m\_\{b\},\\,c\_\{a\},\\,\\hat\{c\}\)
6:return

\(m′,c′\)\(m^\{\\prime\},c^\{\\prime\}\)

#### Adaptive Radix\-Bitonic Selection\.

As an alternative, we exploit a distributional property of LLM logits:∼99\.82%\{\\sim\}99\.82\\%of vocabulary tiles share a single mode in the upper 6 bits of theirϕ\\phi\-encoded values at unit scale, rising to99\.99%99\.99\\%with a power\-of\-two scaling factor \(e\.g\.,s=4s=4\)\. This permits radix\-based threshold computation over only the lower 10 bits\.

[Algorithm˜4](https://arxiv.org/html/2607.20475#alg4)describes the strategy\. The screening phase \(lines 1–4\) scales, encodes, and performs a single fused reduction viaScreenCombine\([Algorithm˜2](https://arxiv.org/html/2607.20475#alg2)\) that simultaneously computes the modal 6\-bit valueθ\\thetaand its frequencycc\. The combine operator tracks both quantities in one tree reduction, compiling to an interleavedIMNMX–IADDsequence that avoids a separate counting pass and halves inter\-warp synchronization relative to a two\-pass approach\.

Whenc≥kc\\geq k\(sparse path\),SparseThreshold\([Algorithm˜10](https://arxiv.org/html/2607.20475#alg10)\) performs two 5\-bit radix passes over 32\-bin warp\-local histograms\. Each pass invokesRadixPartition\([Algorithm˜12](https://arxiv.org/html/2607.20475#alg12)\), which identifies the pivot via a sum\-based count \(IADDtree\) and computes the margin via a masked\-max \(intra\-warpIMNMXtree\)\. Input values are XOR\-inverted before histogramming, transforming the suffix\-sum needed for descending top\-kkinto an ascending prefix\-sum compatible with the hardwareCumSumprimitive\. Whenc<kc<k\(dense path\), the algorithm falls back to the optimized bitonic top\-kk\([Algorithm˜7](https://arxiv.org/html/2607.20475#alg7)\)\.

Algorithm 3Bitonic Fold1:Block

𝐆∈ℕR×k\\mathbf\{G\}\\in\\mathbb\{N\}^\{R\\times k\}, pair

ρ\\rho, rows

RR, target

kk, folds

FF, order

δ\\delta
2:Reduced candidates

𝐲∈ℕk\\mathbf\{y\}\\in\\mathbb\{N\}^\{k\}
3:

n←log2⁡\(R⋅k\)\+1n\\leftarrow\\log\_\{2\}\(R\\cdot k\)\+1,

κ←log2⁡k\\kappa\\leftarrow\\log\_\{2\}k
4:

𝐁←𝐆\\mathbf\{B\}\\leftarrow\\mathbf\{G\}
5:for

i←1i\\leftarrow 1to

FFdo

6:

𝐁←Reshape​\(𝐁,\[R/2i,2,k\]\)\\mathbf\{B\}\\leftarrow\\textsc\{Reshape\}\(\\mathbf\{B\},\\ \[R/2^\{i\},\\ 2,\\ k\]\)⊳\\triangleright⊳\\trianglerightPair adjacent sub\-tiles along dim 1

7:

𝐋,𝐑←Split​\(Transpose​\(𝐁,\[0,2,1\]\)\)\\mathbf\{L\},\\mathbf\{R\}\\leftarrow\\textsc\{Split\}\(\\textsc\{Transpose\}\(\\mathbf\{B\},\\ \[0,2,1\]\)\)⊳\\triangleright⊳\\trianglerightTriggersconvert\_layout: each warp holds a full\(𝐋,𝐑\)\(\\mathbf\{L\},\\mathbf\{R\}\)pair

8:

𝐁←max⁡\(𝐋,𝐑\)\\mathbf\{B\}\\leftarrow\\max\(\\mathbf\{L\},\\ \\mathbf\{R\}\)⊳\\triangleright⊳\\trianglerightWarp\-localIMNMXtree, zero inter\-warpbar\.sync

9:

𝐇^←Reshape​\(𝐁,\[2\]n−i−1\)\\hat\{\\mathbf\{H\}\}\\leftarrow\\textsc\{Reshape\}\(\\mathbf\{B\},\\ \[2\]^\{n\-i\-1\}\)⊳\\triangleright⊳\\trianglerightHypercube view for bitonic merge

10:if

i<Fi<Fthen

11:

f←Expand​\(ρ,F−i−1,κ\)f\\leftarrow\\textsc\{Expand\}\(\\rho,\\ F\{\-\}i\{\-\}1,\\ \\kappa\)⊳\\triangleright⊳\\trianglerightIntermediate fold: parity\-dependent ordering

12:

𝐇^←HypercubeMerge​\(𝐇^,ρ,f,κ,n−i−1\)\\hat\{\\mathbf\{H\}\}\\leftarrow\\textsc\{HypercubeMerge\}\(\\hat\{\\mathbf\{H\}\},\\rho,f,\\kappa,n\{\-\}i\{\-\}1\)
13:else

14:

𝐇^←HypercubeMerge​\(𝐇^,ρ,δ,κ,n−i−1\)\\hat\{\\mathbf\{H\}\}\\leftarrow\\textsc\{HypercubeMerge\}\(\\hat\{\\mathbf\{H\}\},\\rho,\\delta,\\kappa,n\{\-\}i\{\-\}1\)⊳\\triangleright⊳\\trianglerightFinal fold: caller\-specified orderδ\\delta

15:endif

16:

𝐁←𝐇^\\mathbf\{B\}\\leftarrow\\hat\{\\mathbf\{H\}\}
17:endfor

18:return

Reshape​\(𝐁,\[k\]\)\\textsc\{Reshape\}\(\\mathbf\{B\},\\ \[k\]\)⊳\\triangleright⊳\\trianglerightFlatten collapsed hypercube tokkwinners

Algorithm 4Adaptive Radix\-Bitonic Selection1:Logits

𝐱∈ℝBN\\mathbf\{x\}\\in\\mathbb\{R\}^\{B\_\{N\}\}, indices

𝐈\\mathbf\{I\}, target

kk, scale

ss,block\_id

2:Top\-

kkcandidates written to scratchpad

YY
3:

𝐞←ϕ​\(𝐱⋅s\)\\mathbf\{e\}\\leftarrow\\phi\(\\mathbf\{x\}\\cdot s\)⊳\\trianglerightEncode scaled logits \(uint16\)

4:

𝐮←𝐞≫10\\mathbf\{u\}\\leftarrow\\mathbf\{e\}\\gg 10⊳\\trianglerightUpper 6 bits

5:

\(θ,c\)←Reduce​\(\(𝐮,1\),ScreenCombine\)\(\\theta,c\)\\leftarrow\\textsc\{Reduce\}\\bigl\(\(\\mathbf\{u\},\\,\\mathbf\{1\}\),\\ \\textsc\{ScreenCombine\}\\bigr\)⊳\\trianglerightFused max\-count scan

6:

𝐌←\(𝐮=θ\)\\mathbf\{M\}\\leftarrow\(\\mathbf\{u\}=\\theta\)
7:if

c≥kc\\geq kthen⊳\\trianglerightSparse path: radix over lower 10 bits

8:

θ,m←SparseThreshold​\(𝐞,𝐌,θ,k\)\\theta,m\\leftarrow\\textsc\{SparseThreshold\}\(\\mathbf\{e\},\\,\\mathbf\{M\},\\,\\theta,\\,k\)⊳\\trianglerightTwo\-pass 5\-bit radix \([Algorithm˜10](https://arxiv.org/html/2607.20475#alg10)\)

9:

RadixSelection​\(Y,𝐞,θ,m,𝐈,k,block\_id\)\\textsc\{RadixSelection\}\(Y,\\,\\mathbf\{e\},\\,\\theta,\\,m,\\,\\mathbf\{I\},\\,k,\\,\\textsc\{block\\\_id\}\)
10:else⊳\\trianglerightDense path: bitonic fallback

11:

𝐩←ψ​\(𝐞,𝐈,BN\)\\mathbf\{p\}\\leftarrow\\psi\(\\mathbf\{e\},\\,\\mathbf\{I\},\\,B\_\{N\}\)⊳\\trianglerightPack with inverted indices

12:

BitonicSelection​\(Y,𝐩,k,block\_id\)\\textsc\{BitonicSelection\}\(Y,\\,\\mathbf\{p\},\\,k,\\,\\textsc\{block\\\_id\}\)
13:endif

#### Stage 1: Tile\-Local Reduction\.

The first\-stage kernel tiles across batch rows and vocabulary blocks\. Each program instance loads aBNB\_\{N\}\-element logit tile, applies the logit\-processing prologue, i\.e\., grammar masking, repetition, frequency, and presence penalties, logit bias, and temperature scaling, conditioned on per\-request bit\-level indicators \([Section˜3\.5](https://arxiv.org/html/2607.20475#S3.SS5)\), and invokes either the bitonic or adaptive radix\-bitonic reduction\. The bitonic path writeskkpacked candidates in alternating ascending\-descending order keyed onblock\_idparity, while the adaptive path writeskkcandidates in coalesced order via acumsum\-basedscatter\. Full pseudocode is provided in[Algorithms˜11](https://arxiv.org/html/2607.20475#alg11)and[13](https://arxiv.org/html/2607.20475#alg13)within[Appendix˜E](https://arxiv.org/html/2607.20475#A5)\.

#### Stage 2: Cross\-Tile Merge\.

The second\-stage kernel gathersZv×kZ\_\{v\}\\times kpacked entries from the scratchpad, remaps tile\-local indices to absolute vocabulary positions, and performs a final cross\-tile top\-kkreduction\. When Stage 1 uses the bitonic path, the alternating sort order permits a directBitonicReduction\([Algorithm˜6](https://arxiv.org/html/2607.20475#alg6)\) that merges pre\-sorted sub\-tiles without re\-sorting; otherwise, a fullBitonicTopKis applied\. Following the merge, the kernel decodes the packed entries, applies the probability\-truncation epilogue \(top\-kk, top\-pp, min\-ppmasking via cumulative softmax\), adds pre\-drawn Gumbel noise, and selects the final token viaarg⁡max\\arg\\max\. For speculative verification, it additionally materializes the draft probabilities over the selected indices via fast, stable softmax computation\.

#### Complexity Analysis\.

LetVVdenote vocabulary size,Zv=⌈VBN⌉Z\_\{v\}=\\left\\lceil\\frac\{V\}\{B\_\{N\}\}\\right\\rceilthe number of tiles,kkthe candidate bound, andκ=log2⁡k\\kappa=\\log\_\{2\}k\. Stage 1 launchesB⋅ZvB\\cdot Z\_\{v\}programs\.

On the*bitonic path*, each program performsO​\(BN​log2⁡k\)O\(B\_\{N\}\\log^\{2\}k\)work: the initialκ\\kappamerge stages contribute∑s=1κs\\sum\_\{s=1\}^\{\\kappa\}sCAS passes overBNB\_\{N\}elements, i\.e\.,O​\(BN⋅κ2\)O\(B\_\{N\}\\cdot\\kappa^\{2\}\), dominating the subsequent fold phase with complexityO​\(BN⋅κ\)O\(B\_\{N\}\\cdot\\kappa\)\.

On the*radix path*\(sparse case\), each program performsO​\(BN\)O\(B\_\{N\}\)work via a single screening reduction and two histogram passes\. Stage 2 launchesBBprograms:O​\(Zv⋅k⋅log⁡k\)O\(Z\_\{v\}\\cdot k\\cdot\\log k\)per program when merging alternating sub\-tiles \(fold\-only reduction\), orO​\(Zv⋅k⋅log2⁡k\)O\(Z\_\{v\}\\cdot k\\cdot\\log^\{2\}k\)when applying a full bitonic top\-kkto unordered sub\-tiles\.

SinceZv⋅k≪VZ\_\{v\}\\cdot k\\ll V, Stage 1 dominates in both cases, yieldingO​\(B⋅V⋅log2⁡k\)O\(B\\cdot V\\cdot\\log^\{2\}k\)total work for the bitonic path andO​\(B⋅V\)O\(B\\cdot V\)for the radix path\. A single\-program streaming top\-kkperformsO​\(V⋅log⁡k\)O\(V\\cdot\\log k\)work per row but admits onlyBBconcurrent programs; our formulation trades a modest work factor oflog⁡k\\log k\(bitonic\) or none \(radix\) for aZv×Z\_\{v\}\{\\times\}increase in Stage 1 parallelism\.

### 3\.5Batch\-Level Dynamism via Bit\-Level Indicators

Figure 2:Bit\-level encoding of sampling operations in SonicSampler\.Production serving requires heterogeneous mixtures of greedy and stochastic sampling within a single batch\. Instead of host\-side conditioning and multiple kernel dispatches, which break CUDA Graph compatibility, SonicSampler encodes each request as a compact bit\-level indicator \(see[Figure˜2](https://arxiv.org/html/2607.20475#S3.F2)\) loaded once per program instance\. Consequently, the subsequent execution paths and parameter loads are conditioned on this indicator, allowing greedy and stochastic requests to coexist within a single batched launch without host\-side branching\.

Notably, greedy requests follow a streamlined series of localargmaxpaths, while stochastic requests execute the full pipeline\. On NVIDIA Hopper and newer GPUs, we further leverage Programmatic Dependent Launch \(PDL\) to overlap the two stages by triggering Stage 2 preparation as Stage 1 completes, effectively hiding Stage 2 launch latency at larger batch sizes\.

## 4Experiments

We evaluate SonicSampler across three key dimensions:

1. 1\.End\-to\-end decoding performanceunder realistic, heterogeneous workloads,
2. 2\.Kernel\-level efficiencyof the sampling pipeline, with a focus on our two\-stage top\-kkalgorithm, and
3. 3\.Correctness, ensuring that bounded top\-kksampling preserves distributional fidelity and downstream task accuracy\.

### 4\.1Experimental Setup

![[Uncaptioned image]](https://arxiv.org/html/2607.20475v1/figures/tps-compact.png)

Figure 3:End\-to\-end decode throughput with Eagle3 on Qwen3\-8B TP=4 with T=0\.6, top\-pp=0\.9, top\-kk=128\.
Experiments are run on NVIDIA B200 GPUs \(driver 575\.57\.08\) with Triton v3\.5\.1\. All floating\-point inputs, e\.g\., logits, usebfloat16, while integer\-based inputs useint32\. Latency is measured withdo\_bench\_cudagraphfrom thetriton\.testingsuite, with a 500 ms repetition time per configuration before taking the mean latency, synchronizing the GPU, flushing the CUDA cache, and invoking garbage collection between iterations\.

### 4\.2End\-to\-End Performance

![Refer to caption](https://arxiv.org/html/2607.20475v1/x1.png)Figure 4:Sampling latency breakdown by implementation across workloads with each panel corresponding to a decoding regime\. Stacked segments show the incremental cost of each sampling modifier\. The dashed lines for FlashInfer indicate that distribution truncation via top\-p and min\-p reduces pivot overhead, yielding non\-monotonic scaling\.We evaluate SonicSampler in end\-to\-end decoding to measure its impact on practical inference performance\. Experiments are conducted on Qwen3\-8B with Eagle3 using stochastic sampling \(τ=0\.6\\tau=0\.6,top\-​p=0\.9\\text\{top\-\}p=0\.9,top\-​k=128\\text\{top\-\}k=128\)\. Since TRT\-LLM does not natively sample during drafting, we extend it with its own target\-model sampling kernels for a fair comparison\.

As shown in[Figure˜3](https://arxiv.org/html/2607.20475#S4.F3), SonicSampler improves decoding throughput across all lookahead values, achieving 15\-17% speedups \(\+80\-120 TPS\) over the TRT\-LLM baseline\. The advantage grows withγ\\gamma, reflecting the increasing share of sampling in per\-token latency under speculative decoding\. By unifying the sampling pipeline into a single CUDA Graph\-compatible execution, SonicSampler reduces kernel launch overhead and intermediate memory traffic, translating directly into system\-level speedups\.

### 4\.3Sampling Latency Breakdown

[Figure˜4](https://arxiv.org/html/2607.20475#S4.F4)decomposes sampling latency across five decoding regimes: non\-speculative decoding \(±\\pmtop\-kklogprob\), singular draft step, and multi\-step verification \(±\\pmtop\-kklogprob\)\. Within each regime, we progressively enable sampling operations, from base sampling to temperature scaling, top\-kk, top\-pp, and min\-ppfiltering\. We compare against FlashInfer and two additional baselines:Naive, which mirrors SonicSampler’s pipeline using native PyTorch operations without dynamic workload support, andIndicator, itstorch\.compiled variant\.

SonicSampler consistently achieves the lowest latency across all regimes, yielding approximately 2\.5–4×\\timesover Indicator, 5–6×\\timesover Naive, and 10–16×\\timesover FlashInfer\. Baseline implementations incur progressively higher costs due to repeated kernel launches and intermediate memory traffic, leading to steadily increasing latency\. In particular, FlashInfer exhibits a pronounced increase in latency when top\-kkis applied\. In contrast, SonicSampler’s fused design absorbs these operations with minimal marginal overhead\.

![Refer to caption](https://arxiv.org/html/2607.20475v1/x2.png)Figure 5:Top\-kkkernel latency \(μ\\mus, log scale\) vs\. vocabulary sizeVVfork=128k=128across batch sizesB∈\{1,2,4,8,16,32\}B\\in\\\{1,2,4,8,16,32\\\}\. Triton\-TopK is limited toV≤215V\\leq 2^\{15\}; TileLang\-TopK toV≤217V\\leq 2^\{17\}\.
### 4\.4Top\-kkKernel Performance

We evaluate Top\-kkselection in isolation to quantify the efficiency of our two\-stage algorithm for large\-vocabulary sampling\. We use 2D logit tensors of shape\(B,V\)\(B,V\)with fixedK=128K=128, sweepingVVfrom2132^\{13\}to2182^\{18\}andB∈\{1,2,4,8,16,32\}B\\in\\\{1,2,4,8,16,32\\\}\. Latency is measured on inputs drawn from Zipf\(α=1\.1\\alpha=1\.1\) following prior work\(radik2024;park2026qritahighperformancetopktopp\); correctness is additionally verified on𝒰​\(128\.6,128\.7\)\\mathcal\{U\}\(128\.6,128\.7\)and all\-zero inputs\. All baselines are run under matched configurations \(precision, input shapes, and hardware\), e\.g\., TileLang uses its native compiler, while RadiK is evaluated as a precompiled CUDA C\+\+ binary\. For a subset of methods, we apply minimal adaptations to ensure correctness and compatibility with large vocabulary sizes andbfloat16inputs\. Notably, these modifications are conservative and, where measurable, do not degrade and, in some cases, slightly improve the baseline performance\.

#### FlashInfer\.

FlashInfer’s sampling kernel generates Gumbel noise internally from a random seed, which breaks CUDA graph compatibility\. We modify the kernel to accept precomputed Gumbel noise as an external input, enabling CUDA graph capture and ensuring a fair comparison with our approach\. We also adapt the kernel to operate onbfloat16inputs to match our evaluation precision\.

#### TileLang\-TopK\.

The original implementation supports onlyfloat32inputs\. We extend it to supportbfloat16by introducing an explicit cast path during input loading and promoting values tofloat32for intermediate computation\. We additionally increase the shared\-memory buffer size from 4,096 to 16,384 entries to avoid silent truncation of candidates at large vocabulary sizes\. Even with this adjustment, TileLang\-TopK remains limited toV≤217V\\leq 2^\{17\}due to shared\-memory capacity constraints\.

![[Uncaptioned image]](https://arxiv.org/html/2607.20475v1/x3.png)

Figure 6:Top\-kk=128 speedup over the best baseline\.

#### Triton\-TopK\.

The original implementation uses a fused value\-index packing scheme that restricts the maximum supported vocabulary size toV≤215V\\leq 2^\{15\}\. We retain this implementation as\-is and report configurations beyond this limit as unsupported\.

As shown in[Figure˜5](https://arxiv.org/html/2607.20475#S4.F5), SonicSampler achieves the lowest latency across all evaluated configurations and scales robustly with both vocabulary size and batch size\. In contrast, existing approaches either incur higher constant overhead or degrade more sharply as vocabulary size increases\. The advantage is most pronounced at moderate batch sizes and larger vocabularies, where global reduction increasingly dominates runtime\. At higher batch sizes, where the regimes become increasingly compute\-bound, the relative speedup narrows but remains substantial\.

#### Speedup\.

[Figure˜6](https://arxiv.org/html/2607.20475#S4.F6)summarizes the relative gains, showing that SonicSampler consistently outperforms the best across baselines by22–5×5\\timesin most configurations\. Detailed comparisons against FlashInfer and TileLang are deferred to[Appendix˜G](https://arxiv.org/html/2607.20475#A7)\. These results demonstrate that our two\-stage Top\-kkalgorithm provides both lower latency and more robust scaling, forming the foundation of SonicSampler’s end\-to\-end performance gains\.

Table 1:Normalized latency ofbitonic\(bit\) vs\.adaptive\(adp\) sorting strategies across vocab sizes and batch sizes on B200 \(k=128k=128\)\. All values are divided bymin⁡\(bitonic1,adaptive1\)\\min\(\\text\{bitonic\}\_\{1\},\\,\\text\{adaptive\}\_\{1\}\)for the corresponding architecture and vocab size, so 1\.00 anchors the fastestbatch\_size = 1configuration\.Boldmarks the lower \(better\) value per cell; ties bold both\.
#### Strategy Selection\.

Our top\-kkkernel supports two strategies, namely an optimized bitonic variant and an adaptive variant\. As shown in[Table˜1](https://arxiv.org/html/2607.20475#S4.T1), neither strategy dominates across all batch sizes\. Their relative performance varies with different batch sizes\. SonicSampler benchmarks both strategies offline and selects the faster variant for each batch size at runtime\. We also show a similar trend holding for H100 in[Table˜3](https://arxiv.org/html/2607.20475#A7.T3)\.

### 4\.5Effective Losslessness under Bounded Top\-kk

SonicSampler is exact for the full logit processing pipeline, including temperature scaling, top\-kkselection, frequency/presence penalties, etc\. The only potential approximation arises in nucleus sampling when the post\-top\-ppsupport exceeds 128 tokens, as our kernel assumes at most 128 nonzero entries\.

![Refer to caption](https://arxiv.org/html/2607.20475v1/figures/topk_mass_deepseek-v3.1-fp4.png)
![Refer to caption](https://arxiv.org/html/2607.20475v1/figures/topk_accuracy_deepseek-v3.1.png)

Figure 7:Effective losslessness and end\-to\-end validation\. \(a\) Residual probability mass outside top\-kkon DeepSeek v3\.1 logits\. \(b\) Fraction of steps where top\-kkmass exceeds top\-ppthresholds; \(c\) End\-to\-end accuracy of DeepSeek v3\.1 on GPQA\-Diamond\.#### Support truncation analysis\.

We quantify this effect on DeepSeek v3\.1\(deepseekai2025deepseekv3technicalreport\)logits\. As shown in[Figure˜7](https://arxiv.org/html/2607.20475#S4.F7)\(a\), the residual probability mass outside top\-kkdecays rapidly, reaching∼10−8\\sim 10^\{\-8\}atk=128k=128\.[Figure˜7](https://arxiv.org/html/2607.20475#S4.F7)\(b\) shows that atk=128k=128, over 98\.7% of decoding steps already exceedp≥0\.999p\\geq 0\.999, and 96\.1% exceedp≥0\.9999p\\geq 0\.9999, while standard thresholds \(e\.g\.,p=0\.9,0\.95p=0\.9,0\.95\) are effectively always satisfied\. Thus, although top\-ppunder bounded top\-kkis theoretically lossy, the excluded mass is negligible in practice, making our kernelseffectively lossless\. Results on other models can be found in[Appendix˜I](https://arxiv.org/html/2607.20475#A9)\.

#### End\-to\-end accuracy\.

To further validate that this has no practical impact, we evaluate end\-to\-end task performance\. As shown in[Figure˜7](https://arxiv.org/html/2607.20475#S4.F7)\(c\), on GPQA\-Diamond with DeepSeek v3\.1 \(temperature 0\.6, top\-pp0\.9\), SonicSampler matches the Torch\-based baseline\. The negligible difference, well within statistical variance, confirms that the bounded\-support assumption does not affect downstream accuracy, reinforcing that SonicSampler is effectively lossless in real workloads\.[Appendix˜H](https://arxiv.org/html/2607.20475#A8)presents results on other datasets\.

## 5Conclusion and Future Directions

We presented SonicSampler, a CUDA Graph\-compatible framework that fuses the full sampling pipeline into a single workload\-aware kernel, achieving the lowest kernel latency and consistent end\-to\-end speedups across diverse settings\. These results show that optimizing sampling directly yields system\-level gains and requires co\-design across operator efficiency, vertical fusion, and support for heterogeneous workloads\. Future work includes extending to beam search and deeper integration with upstream and downstream components\.

## References

## Appendix AExtended Related Work

We provide a more detailed survey of prior work on LLM sampling, top\-kkselection, and kernel fusion for completeness\.

#### FlashInfer\.

FlashInfer\(flashinfer2024\)provides GPU kernels for top\-kkand top\-ppsampling using a sorting\-free, pivot\-based formulation\. Kernels are organized around specific operations \(e\.g\.,min\_p\_sampling\_from\_probs\), with separate code paths for greedy and stochastic sampling\. The pivot selection procedure is data\-dependent, with performance characteristics that vary across temperature settings, and each iteration performs block\-level reduction over the full vocabulary\.

#### XGrammar\.

XGrammar\(dong2025xgrammarflexibleefficientstructured\)enables efficient grammar\-constrained decoding via byte\-level pushdown automaton masking, with CPU\-side mask generation overlapped with GPU inference to hide latency\. The grammar mask is computed on the CPU and transferred to the GPU, where it is applied to logits as a standalone kernel prior to sampling\.

#### Mojo / MAX\.

modular2025maxprovides a fused GPU kernel \(topk\_fused\_sampling\) that jointly handles top\-kkfiltering, nucleus sampling, temperature scaling, and stochastic selection in a single launch\. Mixed greedy/stochastic batches are routed entirely to the stochastic path when any non\-greedy request is present\. The fused kernel does not return intermediate probability distributions, which makes integration with speculative verification tricky\. MAX’s graph compiler supports operator fusion via MLIR\-based IR, with the sampling stage executing as a separate compiled graph from the model forward pass\.

#### MegaKernel\.

The MegaKernel approach\(cheng2025miragepersistentkernelcompiler\)pursues end\-to\-end kernel fusion by compiling multi\-layer LLM inference into a single megakernel, thereby reducing launch overhead and enabling cross\-operator pipelining\. While this is aligned with SonicSampler’s broader fusion perspective, MegaKernel focuses on the model forward pass rather than the sampling stage, and thus does not directly target challenges such as large\-vocabulary top\-kkreduction or the composition of heterogeneous sampling strategies \(e\.g\., greedy, nucleus, and grammar\-constrained sampling\) in dynamically batched settings\.

#### Top\-kkSelection Algorithms\.

Efficient top\-kkselection over modern LLM vocabularies \(V≥217V\\geq 2^\{17\}\) remains an open challenge, with existing approaches trading off between CUDA Graph compatibility, data type flexibility, and scalability:

\(1\) RadiK\.RadiK\(radik2024\)implements multi\-pass radix selection, where each radix pass is delimited by CPU\-GPU synchronization, rendering it incompatible with CUDA Graph capture\. Its public release additionally supports onlyfp32, limiting its applicability in mixed\-precision inference pipelines prevalent in modern deployments\.

\(2\) Triton Streaming Bitonic\.Triton’s streaming bitonic top\-kk, while being the most proximal to our implementation in that we share the underlying bitonic network implementation, is currently constrained toV<215V<2^\{15\}due to its fused value\-index packing scheme, which reserves 16 bits for indices within a wider sorting key\. Additionally, the block\-wise streaming design forces it to loop over vocabulary blocks within a single program, which limits opportunities for cross\-block parallelism\.

\(3\) FlashInfer Top\-kkVariants\.FlashInfer provides two top\-kkpaths selected by a heuristic wrapper that considerskk, vocabulary size, and data type\.FilteredTopKuses a single\-CTA, shared\-memory\-intensive 8\-bit coarse histogram with iterative radix refinement, but is constrained tok≤2048k\\leq 2048with 128KB shared memory, and is further limited to vocabulary sizes below dtype\-dependent thresholds \(16K forfp16/bf16, 32K forfp32\)\. Beyond these thresholds, the wrapper falls back toMulti\-CTA RadixTopK, which chunks rows across thread blocks with global histogram coordination, with associated atomic and synchronization costs\.

\(4\) TileLang\.TileLang\(wang2025tilelangcomposabletiledprogramming\)provides a tile\-aware multi\-pass radix selection implementation that proves competitive with our approach \([Section˜4](https://arxiv.org/html/2607.20475#S4)\)\. Its top\-kkkernel assigns a single thread block per batch element, iterating over vocabulary tiles sequentially within each block\. In contrast, our two\-stage decomposition distributes vocabulary blocks across independent thread blocks in the first phase, increasing SM occupancy, followed by a lightweight merge over partial results\.

\(5\) Qrita\.Qrita\(park2026qritahighperformancetopktopp\)combines Gaussianσ\\sigma\-truncation, which prunes the search space by thresholding logits to statistical outliers, with a quaternary pivot search that halves iteration count relative to binary search\. However, the single\-CTA Triton implementation limits SM utilization for large vocabularies, as each program must sequentially iterate over vocabulary blocks exceeding the thread block size\. Performance is also data\-dependent, withσ\\sigma\-truncation effectiveness conditioned on the logit distribution conforming to its Gaussian assumption\. As a standalone top\-kkprimitive, Qrita does not address fusion with the surrounding logit processing pipeline\.

#### Large\-Scale Selection\.

The need for fast top\-kkover largeNNextends beyond vocabulary selection\. DeepSeek\-V3\.2’s Lightning Indexer\(deepseekai2025deepseekv32pushingfrontieropen\)performs sparse attention token selection over context lengths approaching10510^\{5\}, driving demand for efficient top\-kkimplementations in frameworks such as TileLang and FlashInfer\. This underscores the broader applicability of scalable top\-kkalgorithms\. Yet within the sampling context, top\-kkis uniquely positioned: it is preceded by a dense chain of element\-wise logit transformations and followed by probability normalization and stochastic selection\. This structure presents an opportunity for vertical fusion that standalone top\-kkimplementations, designed as isolated primitives, cannot exploit\.

#### FlashSampling\.

FlashSampling\(ruiz2026flashsampling\)accelerates sampling by fusing the LM head with downstream sampling operations, reducing memory traffic between logit generation and token selection\. FlashSampling operates at the boundary between the model forward pass and token selection, whereas the post\-logit pipeline—logit processing, top\-kk/top\-ppfiltering, and speculative verification—constitutes a separate stage of the inference stack\.

Table 2:Feature comparison of sampling kernels\.\\checkmark\\checkmarkindicates full support,×\\timesindicates no support\. ‘Penalties’ denotes repetition, frequency, and presence penalties; ‘Truncation’ denotes top\-kk, top\-pp, and min\-ppprobability truncation; ‘Verify’ denotes speculative verification; ‘CG’ denotes CUDA Graph compatibility; and ‘Mixed G/S’ denotes mixed greedy/stochastic batch support\.

## Appendix BFeature Comparison

[Table˜2](https://arxiv.org/html/2607.20475#A1.T2)summarizes the feature coverage across existing implementations\. SonicSampler is the only implementation providing unified coverage of the complete sampling pipeline while maintaining CUDA Graph compatibility and supporting mixed greedy/stochastic batches\.

## Appendix CBackground: Detailed Formulation

We formalize the logit processing, sampling, and verification operations that constitute the LLM sampling pipeline\. Let𝐱∈ℝV\\mathbf\{x\}\\in\\mathbb\{R\}^\{V\}denote raw logits over a vocabulary of sizeVV, and𝐩=Softmax​\(𝐱\)\\mathbf\{p\}=\\text\{Softmax\}\(\\mathbf\{x\}\)the corresponding probability distribution\.

### C\.1Logit Processing Operations

#### Grammar Bit\-Masking\.

Constrained decoding enforces structural validity by masking logits corresponding to tokens that would violate grammatical constraints\. Given a binary mask𝐦∈\{0,1\}V\\mathbf\{m\}\\in\\\{0,1\\\}^\{V\}derived from a finite\-state machine state, we apply:

xi←\{xiif​mi=1−∞otherwisex\_\{i\}\\leftarrow\\begin\{cases\}x\_\{i\}&\\text\{if \}m\_\{i\}=1\\\\ \-\\infty&\\text\{otherwise\}\\end\{cases\}\(4\)

#### Repetition Penalties\.

To discourage degenerate repetition in open\-ended generation, three complementary penalties are applied to previously generated tokens\. Letci∈ℕ0c\_\{i\}\\in\\mathbb\{N\}\_\{0\}denote the number of times tokeniihas appeared in the decoded sequence so far, and letαrep≥1\\alpha\_\{\\mathrm\{rep\}\}\\geq 1,βfreq≥0\\beta\_\{\\mathrm\{freq\}\}\\geq 0, andγpres≥0\\gamma\_\{\\mathrm\{pres\}\}\\geq 0be the repetition, frequency, and presence penalty coefficients, respectively\.

The*multiplicative repetition penalty*contracts positive logits and amplifies negative logits for any previously generated token, preserving the sign\-dependent asymmetry needed to maintain the correct probability ordering:

x~i=\{xi/αrepif​ci\>0​and​xi\>0,xi⋅αrepif​ci\>0​and​xi≤0,xiotherwise\.\\tilde\{x\}\_\{i\}=\\begin\{cases\}x\_\{i\}\\,/\\,\\alpha\_\{\\mathrm\{rep\}\}&\\text\{if \}c\_\{i\}\>0\\;\\text\{and\}\\;x\_\{i\}\>0,\\\\ x\_\{i\}\\cdot\\alpha\_\{\\mathrm\{rep\}\}&\\text\{if \}c\_\{i\}\>0\\;\\text\{and\}\\;x\_\{i\}\\leq 0,\\\\ x\_\{i\}&\\text\{otherwise\}\.\\end\{cases\}\(5\)The*frequency penalty*then applies an additive reduction proportional to each token’s occurrence count, and the*presence penalty*applies a flat additive reduction to every token that has appeared at least once, yielding the final penalized logit:

xi←x~i−βfreq⋅ci−γpres⋅𝕀​\[ci\>0\]\.x\_\{i\}\\;\\leftarrow\\;\\tilde\{x\}\_\{i\}\\;\-\\;\\beta\_\{\\mathrm\{freq\}\}\\cdot\\,c\_\{i\}\\;\-\\;\\gamma\_\{\\mathrm\{pres\}\}\\cdot\\,\\mathbb\{I\}\[\\,c\_\{i\}\>0\\,\]\.\(6\)Settingαrep=1\\alpha\_\{\\mathrm\{rep\}\}=1,βfreq=0\\beta\_\{\\mathrm\{freq\}\}=0,γpres=0\\gamma\_\{\\mathrm\{pres\}\}=0recovers the unpenalized logits\.

#### Logit Bias\.

Request\-specific additive biases𝐛∈ℝV\\mathbf\{b\}\\in\\mathbb\{R\}^\{V\}:

xi←xi\+bix\_\{i\}\\leftarrow x\_\{i\}\+b\_\{i\}\(7\)

#### Temperature\.

Temperatureτ\>0\\tau\>0modulates distribution entropy:

xi←xi/τx\_\{i\}\\leftarrow x\_\{i\}/\\tau\(8\)Lower temperatures sharpen the distribution toward the mode; higher temperatures increase diversity\.

### C\.2Probability Truncation

#### Top\-kk\.

Restricts the candidate set to thekkhighest\-probability tokens, which is analogous to selecting thekkhighest valued logits, as Softmax is inherently a monotonic and order\-preserving transformation:

𝒯k=\{i:rank​\(xi\)≤k\}\\mathcal\{T\}\_\{k\}=\\\{i:\\text\{rank\}\(x\_\{i\}\)\\leq k\\\}\(9\)

#### Top\-pp\(Nucleus\)\.

Select the minimal set whose cumulative probability exceeds the thresholdpp\(holtzman2020curious\):

𝒯p=arg⁡minS⁡\|S\|s\.t\.∑i∈Spi≥p,pi≥pj​∀i∈S,j∉S\\mathcal\{T\}\_\{p\}=\\arg\\min\_\{S\}\|S\|\\quad\\text\{s\.t\.\}\\quad\\sum\_\{i\\in S\}p\_\{i\}\\geq p,\\;p\_\{i\}\\geq p\_\{j\}\\;\\forall i\\in S,j\\notin S\(10\)

#### Min\-pp\.

Establishes a probability floor relative to the maximum\. The naive formulation requires computing the full softmax:

𝒯θ=\{i:pi≥θ⋅maxj⁡pj\}\\mathcal\{T\}\_\{\\theta\}=\\\{i:p\_\{i\}\\geq\\theta\\cdot\\max\_\{j\}p\_\{j\}\\\}\(11\)However, we observe that this criterion admits a simplification in logit space\. Sinceln⁡pi=xi−ln​∑jexj\\ln p\_\{i\}=x\_\{i\}\-\\ln\\sum\_\{j\}e^\{x\_\{j\}\}, we have:

pi≥θ⋅maxj⁡pj\\displaystyle p\_\{i\}\\geq\\theta\\cdot\\max\_\{j\}p\_\{j\}\\quad⇔ln⁡pi≥ln⁡θ\+maxj⁡ln⁡pj\\displaystyle\\Leftrightarrow\\quad\\ln p\_\{i\}\\geq\\ln\\theta\+\\max\_\{j\}\\ln p\_\{j\}⇔xi≥M\+ln⁡θ\\displaystyle\\Leftrightarrow\\quad x\_\{i\}\\geq M\+\\ln\\theta\(12\)whereM=maxj⁡xjM=\\max\_\{j\}x\_\{j\}\. This reduces min\-ppfiltering to a simple threshold comparison against the maximum logit, avoiding explicit softmax computation until after truncation\. In fact, we leverage this critical simplification towards optimizing the min\-ppprocessing within SonicSampler\.

### C\.3Sampling

#### Gumbel\-Max Reparameterization\.

Sampling from a categorical distribution can be reformulated as a deterministicarg⁡max\\arg\\maxover perturbed logits\. If𝐠=\(g1,…,gV\)\\mathbf\{g\}=\(g\_\{1\},\\ldots,g\_\{V\}\)are i\.i\.d\. samples fromGumbel​\(0,1\)\\text\{Gumbel\}\(0,1\), then:

x∗=arg⁡maxi⁡\(gi\+ln⁡pi\)∼Categorical​\(𝐩\)x^\{\*\}=\\arg\\max\_\{i\}\(g\_\{i\}\+\\ln p\_\{i\}\)\\sim\\text\{Categorical\}\(\\mathbf\{p\}\)\(13\)Sincegi=−ln⁡zig\_\{i\}=\-\\ln z\_\{i\}wherezi∼Exp​\(1\)z\_\{i\}\\sim\\text\{Exp\}\(1\), and noting that the log\-partition function is constant across indices, we obtain:

x∗=arg⁡maxi⁡\(xi−ln⁡zi\)x^\{\*\}=\\arg\\max\_\{i\}\(x\_\{i\}\-\\ln z\_\{i\}\)\(14\)This formulation unifies stochastic sampling with thearg⁡max\\arg\\maxoperation used in greedy decoding, enabling a single code path where greedy selection corresponds tozi=1z\_\{i\}=1for allii\.

### C\.4Speculative Verification

In speculative decoding, a draft model proposes a sequence ofLLcandidate tokens\(d1,…,dL\)\(d\_\{1\},\\ldots,d\_\{L\}\)with probabilitiesq​\(dℓ\)q\(d\_\{\\ell\}\)\. The target model computes probabilitiesp​\(dℓ\)p\(d\_\{\\ell\}\)for each candidate\. Verification proceeds sequentially: tokendℓd\_\{\\ell\}is accepted if\(leviathan2023fast\):

uℓ≤p​\(dℓ\)q​\(dℓ\),uℓ∼Uniform​\(0,1\)u\_\{\\ell\}\\leq\\frac\{p\(d\_\{\\ell\}\)\}\{q\(d\_\{\\ell\}\)\},\\quad u\_\{\\ell\}\\sim\\text\{Uniform\}\(0,1\)\(15\)Upon rejection at positionℓ\\ell, a correction token is sampled from the residual distribution:

p~​\(x\)=normalize​\(max⁡\{0,p​\(x\)−q​\(x\)\}\)\\tilde\{p\}\(x\)=\\text\{normalize\}\\left\(\\max\\\{0,p\(x\)\-q\(x\)\\\}\\right\)\(16\)The first rejected position determines the acceptance length, and decoding continues from the corrected token\.

#### Draft\-to\-Target Mapping\.

Tree\-structured speculation methods such as EAGLE\-3\(li2025eagle\)require mapping between draft and target token indices during verification\. Letπ:\{1,…,L\}→\{1,…,L′\}\\pi:\\\{1,\\ldots,L\\\}\\to\\\{1,\\ldots,L^\{\\prime\}\\\}denote the mapping from draft positions to target positions\. This indirection is applied during probability lookup and token writeback in the verification kernel\.

## Appendix DGeneralized Packing Transform for Alternating Sub\-Tiles

The basic packing transformψ​\(x,i,n\)=ϕ​\(x\)⋅216\+\(n−1\)⊕i\\psi\(x,i,n\)=\\phi\(x\)\\cdot 2^\{16\}\+\(n\-1\)\\oplus iencodes a value–index pair into a singleuint32for tile\-local sorting\. In Stage 2, the scratchpad containsZv′⋅kZ\_\{v\}^\{\\prime\}\\cdot kpacked entries acrossZv′Z\_\{v\}^\{\\prime\}vocabulary blocks \(whereZv′=2⌈log2⁡Zv⌉Z\_\{v\}^\{\\prime\}=2^\{\\lceil\\log\_\{2\}Z\_\{v\}\\rceil\}\), and the bitonic fold expects alternating ascending–descending sub\-tiles\. We therefore assign*relative indices*across the full scratchpad such that the packed representation preserves each block’s sort order from Stage 1\.

#### Forward TransformΨ\\Psi\.

Letj∈\[0,Zv′\)j\\in\[0,Z\_\{v\}^\{\\prime\}\)be the block index,r∈\[0,k\)r\\in\[0,k\)the position within blockjj,πj=jmod2\\pi\_\{j\}=j\\bmod 2the block parity, andN=Zv′⋅kN=Z\_\{v\}^\{\\prime\}\\cdot kthe total scratchpad width\. Recall that Stage 1 writes blockjj’s candidates in ascending order whenπj=0\\pi\_\{j\}=0and descending order whenπj=1\\pi\_\{j\}=1\. Define the alternating mask

μ​\(j\)=\(N−1\)⊕\(\(1−πj\)​\(k−1\)\),\\mu\(j\)\\;=\\;\(N\-1\)\\oplus\\bigl\(\(1\-\\pi\_\{j\}\)\(k\-1\)\\bigr\),\(17\)and the generalized relative index

Ψ​\(j,r\)=\(j​k\+r\)⊕μ​\(j\)\.\\Psi\(j,\\,r\)\\;=\\;\(jk\+r\)\\oplus\\mu\(j\)\.\(18\)Expanding by parity and writingj¯=\(Zv′−1\)−j\\bar\{j\}=\(Z\_\{v\}^\{\\prime\}\-1\)\-jfor the bit\-complement ofjjwithin⌈log2⁡Zv′⌉\\lceil\\log\_\{2\}Z\_\{v\}^\{\\prime\}\\rceilbits:

Ψ​\(j,r\)=\{j¯​k\+rif​πj=0​\(ascending\),j¯​k\+\(k−1−r\)if​πj=1​\(descending\)\.\\Psi\(j,\\,r\)\\;=\\;\\begin\{cases\}\\bar\{j\}\\,k\+r&\\text\{if \}\\pi\_\{j\}=0\\ \\text\{\(ascending\)\},\\\\\[4\.0pt\] \\bar\{j\}\\,k\+\(k\-1\-r\)&\\text\{if \}\\pi\_\{j\}=1\\ \\text\{\(descending\)\}\.\\end\{cases\}\(19\)For ascending blocks,Ψ\\Psiincreases withrr, matching the ascending value order; for descending blocks,Ψ\\Psidecreases withrr, matching the descending value order\. The block complementj¯\\bar\{j\}arranges the sub\-tiles into the interleaved pattern expected by the bitonic fold network\.

Stage 2 replaces each entry’s tile\-local packed index withΨ​\(j,r\)\\Psi\(j,r\)while separately pre\-computing the absolute vocabulary indices

abs​\[j,r\]=j⋅BN\+\(\(pj,r∧0xFFFF\)⊕\(BN−1\)\),\\text\{abs\}\[j,\\,r\]\\;=\\;j\\cdot B\_\{N\}\+\\bigl\(\(p\_\{j,r\}\\ \\land\\texttt\{0xFFFF\}\)\\oplus\(B\_\{N\}\-1\)\\bigr\),\(20\)wherepj,rp\_\{j,r\}is the original packeduint32from Stage 1 and the XOR with\(BN−1\)\(B\_\{N\}\-1\)inverts the tile\-local index complement from the basic packing transformψ\\psi\.

#### Inverse TransformΨ−1\\Psi^\{\-1\}\.

After the bitonic reduction selectskkwinners with relative indices\{ℓ0′,…,ℓk−1′\}\\\{\\ell^\{\\prime\}\_\{0\},\\ldots,\\ell^\{\\prime\}\_\{k\-1\}\\\}, the original block and position are recovered as follows\. First, apply the total complement:

ℓ~=ℓ′⊕\(N−1\)\.\\tilde\{\\ell\}\\;=\\;\\ell^\{\\prime\}\\oplus\(N\-1\)\.\(21\)Then extract the block index and correct the within\-block position:

j^\\displaystyle\\hat\{j\}=ℓ~≫κ,\\displaystyle\\;=\\;\\tilde\{\\ell\}\\gg\\kappa,\(22\)r^\\displaystyle\\hat\{r\}=\(ℓ~∧\(k−1\)\)⊕\(\(1−j^mod2\)⋅\(k−1\)\),\\displaystyle\\;=\\;\\bigl\(\\tilde\{\\ell\}\\ \\land\(k\-1\)\\bigr\)\\oplus\\bigl\(\(1\-\\hat\{j\}\\bmod 2\)\\cdot\(k\-1\)\\bigr\),\(23\)whereκ=log2⁡k\\kappa=\\log\_\{2\}k\. The position correction in[equation˜23](https://arxiv.org/html/2607.20475#A4.E23)re\-applies the parity\-dependent index flip fromμ​\(j\)\\mu\(j\), yielding the original within\-block positionrrfor both parities\. The absolute vocabulary index is then recovered via a gather:

a=abs​\[j^,r^\]=abs​\[j^​k\+r^\]\.a\\;=\\;\\text\{abs\}\[\\hat\{j\},\\,\\hat\{r\}\]\\;=\\;\\text\{abs\}\\bigl\[\\hat\{j\}\\,k\+\\hat\{r\}\\bigr\]\.\(24\)

#### Correctness\.

We verify the round\-tripΨ−1∘Ψ=id\\Psi^\{\-1\}\\circ\\Psi=\\text\{id\}for both parities\. For an ascending block \(πj=0\\pi\_\{j\}=0\):

ℓ′\\displaystyle\\ell^\{\\prime\}=j¯​k\+r\\displaystyle=\\bar\{j\}k\+r\(25\)ℓ~\\displaystyle\\tilde\{\\ell\}=j​k\+\(k−1−r\)\\displaystyle=jk\+\(k\{\-\}1\{\-\}r\)\(26\)which yieldsj^=j\\hat\{j\}=jand:

r^=\(k−1−r\)⊕\(k−1\)=r\\hat\{r\}=\(k\{\-\}1\{\-\}r\)\\oplus\(k\{\-\}1\)=r\(27\)
For a descending block \(πj=1\\pi\_\{j\}=1\):

ℓ′\\displaystyle\\ell^\{\\prime\}=j¯​k\+\(k−1−r\)\\displaystyle=\\bar\{j\}k\+\(k\{\-\}1\{\-\}r\)\(28\)ℓ~\\displaystyle\\tilde\{\\ell\}=j​k\+r\\displaystyle=jk\+r\(29\)
which givesj^=j\\hat\{j\}=jandr^=r⊕0=r\\hat\{r\}=r\\oplus 0=r\. In both cases, the original\(j,r\)\(j,r\)pair is recovered exactly\.

## Appendix ESupplementary Algorithms

This appendix provides the complete pseudocode for the algorithms referenced but not fully detailed in the main text\.

Algorithm 5Hypercube Merge1:Hypercube

𝐯∈ℕ2n\\mathbf\{v\}\\in\\mathbb\{N\}^\{2^\{n\}\}, pair

ρ\\rho, flip

ff, stage

ss, total

nn
2:Merged

𝐯′∈ℕ2n\\mathbf\{v\}^\{\\prime\}\\in\\mathbb\{N\}^\{2^\{n\}\}
3:for

i←0i\\leftarrow 0to

s−1s\-1do

4:

𝐯←CAS​\(𝐯,ρ,f,s−i−1,n\)\\mathbf\{v\}\\leftarrow\\textsc\{CAS\}\(\\mathbf\{v\},\\rho,f,\\,s\-i\-1,\\,n\)
5:endfor

6:return

𝐯\\mathbf\{v\}

Algorithm 6Bitonic Reduction \(Cross\-Tile Merge\)1:Block

𝐁∈ℕR×k\\mathbf\{B\}\\in\\mathbb\{N\}^\{R\\times k\}, rows

RR, target

kk, total columns

NN
2:Reduced candidates

𝐲∈ℕk\\mathbf\{y\}\\in\\mathbb\{N\}^\{k\}
3:

n←log2⁡N\+1n\\leftarrow\\log\_\{2\}N\+1,

κ←log2⁡k\\kappa\\leftarrow\\log\_\{2\}k,

F←log2⁡RF\\leftarrow\\log\_\{2\}R
4:

ρ←\{0,1\}\\rho\\leftarrow\\\{0,1\\\}
5:

𝐲←BitonicFold\(𝐁,ρ,R,k,F,δ=1,n,κ\)\\mathbf\{y\}\\leftarrow\\textsc\{BitonicFold\}\(\\mathbf\{B\},\\,\\rho,\\,R,\\,k,\\,F,\\,\\delta\{=\}1,\\,n,\\,\\kappa\)
6:return

𝐲\\mathbf\{y\}

Algorithm 7Bitonic Selection \(Stage 1 Write\-Out\)1:Scratchpad

YY, batch

bb, block

jj, packed

𝐩∈ℕBN\\mathbf\{p\}\\in\\mathbb\{N\}^\{B\_\{N\}\}, target

kk
2:

kkcandidates written to

Y\[b,j⋅k:\(j\+1\)k\]Y\[b,j\\cdot k:\(j\{\+\}1\)k\]
3:

𝐲←BitonicTopK​\(𝐩,jmod2,k,BN\)\\mathbf\{y\}\\leftarrow\\textsc\{BitonicTopK\}\(\\mathbf\{p\},\\,j\\bmod 2,\\,k,\\,B\_\{N\}\)⊳\\trianglerightAlternating order

4:

Store​\(Y\+b⋅SY\+j⋅k,𝐲\)\\textsc\{Store\}\(Y\+b\\cdot S\_\{Y\}\+j\\cdot k,\\,\\mathbf\{y\}\)

Algorithm 8Bitonic Top\-kkSelection1:Packed values

𝐩∈ℕBN\\mathbf\{p\}\\in\\mathbb\{N\}^\{B\_\{N\}\}, order

δ∈\{0,1\}\\delta\\in\\\{0,1\\\}, target

kk
2:Top\-

kkcandidates

𝐲∈ℕk\\mathbf\{y\}\\in\\mathbb\{N\}^\{k\}
3:

n←log2⁡BNn\\leftarrow\\log\_\{2\}B\_\{N\},

κ←log2⁡k\\kappa\\leftarrow\\log\_\{2\}k
4:

F←n−κ−1F\\leftarrow n\-\\kappa\-1,

R←2FR\\leftarrow 2^\{F\}
5:

𝐇←Reshape​\(𝐩,\[2\]n\)\\mathbf\{H\}\\leftarrow\\textsc\{Reshape\}\(\\mathbf\{p\},\\ \[2\]^\{n\}\)⊳\\trianglerightHypercube view

6:

ρ←\{0,1\}\\rho\\leftarrow\\\{0,1\\\}
7:for

s←1s\\leftarrow 1to

κ\\kappado⊳\\trianglerightSortkk\-width slices

8:

f←Expand​\(ρ,n−s−1,s\)f\\leftarrow\\textsc\{Expand\}\(\\rho,\\ n\{\-\}s\{\-\}1,\\ s\)
9:

𝐇←HypercubeMerge​\(𝐇,ρ,f,s,n\)\\mathbf\{H\}\\leftarrow\\textsc\{HypercubeMerge\}\(\\mathbf\{H\},\\rho,f,s,n\)
10:endfor

11:

𝐇←max⁡\(𝐇,axis=F\)\\mathbf\{H\}\\leftarrow\\max\(\\mathbf\{H\},\\,\\text\{axis\}=F\)⊳\\trianglerightReduce folds

12:if

F\>0F\>0then

13:

f←Expand​\(ρ,F−1,κ\)f\\leftarrow\\textsc\{Expand\}\(\\rho,\\ F\{\-\}1,\\ \\kappa\)
14:

𝐇←HypercubeMerge​\(𝐇,ρ,f,κ,n−1\)\\mathbf\{H\}\\leftarrow\\textsc\{HypercubeMerge\}\(\\mathbf\{H\},\\rho,f,\\kappa,n\{\-\}1\)
15:

𝐆←Reshape​\(𝐇,\[R,k\]\)\\mathbf\{G\}\\leftarrow\\textsc\{Reshape\}\(\\mathbf\{H\},\\ \[R,\\,k\]\)
16:

𝐲←BitonicFold​\(𝐆,ρ,R,k,F,δ\)\\mathbf\{y\}\\leftarrow\\textsc\{BitonicFold\}\(\\mathbf\{G\},\\rho,R,k,F,\\delta\)
17:else

18:

𝐇←HypercubeMerge​\(𝐇,ρ,δ,κ,n−1\)\\mathbf\{H\}\\leftarrow\\textsc\{HypercubeMerge\}\(\\mathbf\{H\},\\rho,\\delta,\\kappa,n\{\-\}1\)
19:

𝐲←Reshape​\(𝐇,\[k\]\)\\mathbf\{y\}\\leftarrow\\textsc\{Reshape\}\(\\mathbf\{H\},\\ \[k\]\)
20:endif

21:return

𝐲\\mathbf\{y\}

Algorithm 9Radix Selection \(Stage 1 Write\-Out\)1:Scratchpad

YY, batch

bb, block

jj, encoded

𝐞\\mathbf\{e\}, threshold

θ\\theta, margin

mm, indices

𝐈\\mathbf\{I\}, target

kk, tile size

BNB\_\{N\}
2:

kkcandidates written to

Y\[b,j⋅k:\(j\+1\)k\]Y\[b,j\\cdot k:\(j\{\+\}1\)k\]
3:

𝐩←ψ​\(𝐞,𝐈,BN\)\\mathbf\{p\}\\leftarrow\\psi\(\\mathbf\{e\},\\,\\mathbf\{I\},\\,B\_\{N\}\)⊳\\trianglerightPack with inverted indices

4:

𝐪←\(𝐞=θ\)&\(CumSum​\(𝐞=θ\)≤m\)\\mathbf\{q\}\\leftarrow\(\\mathbf\{e\}=\\theta\)\\mathbin\{\\&\}\(\\textsc\{CumSum\}\(\\mathbf\{e\}=\\theta\)\\leq m\)⊳\\trianglerightMargin tie\-breaking

5:

𝐌←\(𝐞\>θ\)\|𝐪\\mathbf\{M\}\\leftarrow\(\\mathbf\{e\}\>\\theta\)\\mathbin\{\|\}\\mathbf\{q\}⊳\\trianglerightSelection mask

6:

pos←CumSum​\(𝐌\)−𝐌\\text\{pos\}\\leftarrow\\textsc\{CumSum\}\(\\mathbf\{M\}\)\-\\mathbf\{M\}⊳\\trianglerightCoalesced positions

7:

MaskedStore​\(Y\+b⋅SY\+j⋅k\+pos,𝐩,mask=𝐌\)\\textsc\{MaskedStore\}\(Y\+b\\cdot S\_\{Y\}\+j\\cdot k\+\\text\{pos\},\\,\\mathbf\{p\},\\,\\text\{mask\}\{=\}\\mathbf\{M\}\)

Algorithm 10Sparse Threshold \(Two\-Pass Radix\)1:Encoded

𝐞∈ℕBN\\mathbf\{e\}\\in\\mathbb\{N\}^\{B\_\{N\}\}, mask

𝐌\\mathbf\{M\}, initial threshold

θ0\\theta\_\{0\}, target

kk
2:Threshold

θ\\theta, margin

mm
3:

β←5\\beta\\leftarrow 5,

ℬ←32\\mathcal\{B\}\\leftarrow 32,

μ←ℬ−1\\mu\\leftarrow\\mathcal\{B\}\-1⊳\\triangleright5\-bit radix with 32 bins

4:⊳\\trianglerightPass 1: Mid\-span bits \[5:10\)

5:

𝐯←\(\(𝐞≫β\)&μ\)⊕μ\\mathbf\{v\}\\leftarrow\(\(\\mathbf\{e\}\\gg\\beta\)\\mathbin\{\\&\}\\mu\)\\oplus\\mu⊳\\trianglerightInvert for descending order

6:

h,m←RadixPartition​\(𝐯,𝐌,k,ℬ\)h,m\\leftarrow\\textsc\{RadixPartition\}\(\\mathbf\{v\},\\,\\mathbf\{M\},\\,k,\\,\\mathcal\{B\}\)
7:

θ←\(\(θ0≪β\)\|\(h⊕μ\)\)≪β\\theta\\leftarrow\(\(\\theta\_\{0\}\\ll\\beta\)\\mathbin\{\|\}\(h\\oplus\\mu\)\)\\ll\\beta
8:if

m\>0m\>0then⊳\\trianglerightPass 2: Low\-span bits \[0:5\)

9:

𝐌←𝐌&\(𝐯=h\)\\mathbf\{M\}\\leftarrow\\mathbf\{M\}\\mathbin\{\\&\}\(\\mathbf\{v\}=h\)
10:

𝐯←\(𝐞&μ\)⊕μ\\mathbf\{v\}\\leftarrow\(\\mathbf\{e\}\\mathbin\{\\&\}\\mu\)\\oplus\\mu
11:

t,m←RadixPartition​\(𝐯,𝐌,m,ℬ\)t,m\\leftarrow\\textsc\{RadixPartition\}\(\\mathbf\{v\},\\,\\mathbf\{M\},\\,m,\\,\\mathcal\{B\}\)
12:if

h=0h=0and

t=0t=0then⊳\\trianglerightDegenerate boundary case

13:

θ←θ\|\(μ−1\)\\theta\\leftarrow\\theta\\mathbin\{\|\}\(\\mu\-1\);

m←0m\\leftarrow 0
14:else

15:

θ←θ\|\(t⊕μ\)\\theta\\leftarrow\\theta\\mathbin\{\|\}\(t\\oplus\\mu\)
16:endif

17:else

18:

θ←θ\|μ\\theta\\leftarrow\\theta\\mathbin\{\|\}\\mu
19:endif

20:return

θ,m\\theta,\\,m

Algorithm 11Stage 1: Tile\-Local Reduction Kernel1:Logits

X∈ℝL×VX\\in\\mathbb\{R\}^\{L\\times V\}, indicators

𝐳∈ℕB\\mathbf\{z\}\\in\\mathbb\{N\}^\{B\}, scratchpad

YY, tile size

BNB\_\{N\}, target

kk, mode

∈\{bitonic,adaptive\}\\in\\\{\\texttt\{bitonic\},\\texttt\{adaptive\}\\\}
2:Packed top\-

kkcandidates per tile in

YY
3:

b←ProgramId​\(0\)b\\leftarrow\\textsc\{ProgramId\}\(0\),

j←ProgramId​\(1\)j\\leftarrow\\textsc\{ProgramId\}\(1\)
4:

𝐱←X\[b,j⋅BN:\(j\+1\)BN\]\\mathbf\{x\}\\leftarrow X\[b,\\,j\\cdot B\_\{N\}:\(j\{\+\}1\)B\_\{N\}\]⊳\\trianglerightLoad logit tile \(masked at boundary\)

5:⊳\\triangleright— Logit Processing Prologue —

6:

𝐱←GrammarMask​\(𝐱,zb\)\\mathbf\{x\}\\leftarrow\\textsc\{GrammarMask\}\(\\mathbf\{x\},z\_\{b\}\)
7:

𝐱←RepetitionPenalty​\(𝐱,zb\)\\mathbf\{x\}\\leftarrow\\textsc\{RepetitionPenalty\}\(\\mathbf\{x\},z\_\{b\}\)
8:

𝐱←LogitBias​\(𝐱,zb\)\\mathbf\{x\}\\leftarrow\\textsc\{LogitBias\}\(\\mathbf\{x\},z\_\{b\}\)
9:

𝐱←Temperature​\(𝐱,zb\)\\mathbf\{x\}\\leftarrow\\textsc\{Temperature\}\(\\mathbf\{x\},z\_\{b\}\)
10:⊳\\triangleright— Tile\-Local Top\-kkReduction —

11:if

zbz\_\{b\}indicates greedythen

12:

GreedyReduction​\(Y,b,j,𝐱\)\\textsc\{GreedyReduction\}\(Y,b,j,\\mathbf\{x\}\)⊳\\trianglerightScalararg⁡max\\arg\\max

13:elseifmode

==adaptivethen

14:

RadixBitonic​\(Y,b,j,𝐱,𝐈,k,s\)\\textsc\{RadixBitonic\}\(Y,b,j,\\mathbf\{x\},\\mathbf\{I\},k,s\)⊳\\triangleright[Algorithm˜4](https://arxiv.org/html/2607.20475#alg4)

15:else

16:

𝐩←ψ​\(𝐱,𝐈,BN\)\\mathbf\{p\}\\leftarrow\\psi\(\\mathbf\{x\},\\mathbf\{I\},B\_\{N\}\)
17:

BitonicSelection​\(Y,b,j,𝐩,k\)\\textsc\{BitonicSelection\}\(Y,b,j,\\mathbf\{p\},k\)⊳\\triangleright[Algorithm˜7](https://arxiv.org/html/2607.20475#alg7)

18:endif

Algorithm 12Radix Partition1:Values

𝐯\\mathbf\{v\}, mask

𝐌\\mathbf\{M\}, upper bound

uu, bins

ℬ\\mathcal\{B\}
2:Pivot

pp, margin

mm
3:

𝐡←CumSum​\(Histogram​\(𝐯,ℬ,mask=𝐌\)\)\\mathbf\{h\}\\leftarrow\\textsc\{CumSum\}\(\\textsc\{Histogram\}\(\\mathbf\{v\},\\,\\mathcal\{B\},\\,\\text\{mask\}\{=\}\\mathbf\{M\}\)\)
4:

𝐛←\(𝐡≤u\)\\mathbf\{b\}\\leftarrow\(\\mathbf\{h\}\\leq u\)
5:

p←Sum​\(𝐛\)p\\leftarrow\\textsc\{Sum\}\(\\mathbf\{b\}\)⊳\\triangleright⊳\\trianglerightIADD reduction tree

6:

t←max⁡\(Select​\(𝐛,𝐡,0\)\)t\\leftarrow\\max\(\\textsc\{Select\}\(\\mathbf\{b\},\\,\\mathbf\{h\},\\,\\mathbf\{0\}\)\)⊳\\triangleright⊳\\trianglerightMasked IMNMX tree

7:

m←Select​\(p\>0,u−t,u\)m\\leftarrow\\textsc\{Select\}\(p\>0,\\,u\-t,\\,u\)
8:return

\(p,m\)\(p,\\,m\)

Algorithm 13Stage 2: Cross\-Tile Merge and Selection Kernel1:Scratchpad

Y∈ℕL×Zv​kY\\in\\mathbb\{N\}^\{L\\times Z\_\{v\}k\}, indicators

𝐳\\mathbf\{z\}, Gumbel noise

𝐠\\mathbf\{g\}, parameters

\(kb,pb,nb\)\(k\_\{b\},p\_\{b\},n\_\{b\}\), target

kk, alternating flag

2:Selected token

t∈ℕt\\in\\mathbb\{N\}, log\-probability

w∈ℝw\\in\\mathbb\{R\}
3:

b←ProgramId​\(0\)b\\leftarrow\\textsc\{ProgramId\}\(0\)
4:if

zbz\_\{b\}indicates greedythen

5:

t←GreedyMaximum​\(Y,b\)t\\leftarrow\\textsc\{GreedyMaximum\}\(Y,b\)⊳\\trianglerightarg⁡max\\arg\\maxoverZvZ\_\{v\}scalar entries

6:else

7:⊳\\triangleright— Cross\-Tile Merge —

8:Load

\[Zv,k\]\[Z\_\{v\},k\]packed entries; remap tile\-local to absolute indices

9:ifalternatingthen

10:

𝐯,𝐣←BitonicReduction​\(Yb,Zv,k\)\\mathbf\{v\},\\mathbf\{j\}\\leftarrow\\textsc\{BitonicReduction\}\(Y\_\{b\},\\,Z\_\{v\},\\,k\)⊳\\triangleright[Algorithm˜6](https://arxiv.org/html/2607.20475#alg6)

11:else

12:

𝐯,𝐣←BitonicTopK​\(Yb,1,k,Zv⋅k\)\\mathbf\{v\},\\mathbf\{j\}\\leftarrow\\textsc\{BitonicTopK\}\(Y\_\{b\},\\,1,\\,k,\\,Z\_\{v\}\\cdot k\)
13:endif

14:Decode:

𝐯←ϕ−1​\(𝐯upper\)\\mathbf\{v\}\\leftarrow\\phi^\{\-1\}\(\\mathbf\{v\}\_\{\\text\{upper\}\}\);

𝐣←Gather​\(abs\. indices,𝐣\)\\mathbf\{j\}\\leftarrow\\textsc\{Gather\}\(\\text\{abs\.\\penalty 10000\\ indices\},\\,\\mathbf\{j\}\)
15:⊳\\triangleright— Probability Truncation Epilogue —

16:

𝐯←TopKMask​\(𝐯,kb\)\\mathbf\{v\}\\leftarrow\\textsc\{TopKMask\}\(\\mathbf\{v\},k\_\{b\}\)
17:

𝐯←TopPMask​\(𝐯,pb\)\\mathbf\{v\}\\leftarrow\\textsc\{TopPMask\}\(\\mathbf\{v\},p\_\{b\}\)⊳\\trianglerightCumulative softmax

18:

𝐯←MinPMask​\(𝐯,nb\)\\mathbf\{v\}\\leftarrow\\textsc\{MinPMask\}\(\\mathbf\{v\},n\_\{b\}\)⊳\\trianglerightlog\\log\-domain threshold \([equation˜12](https://arxiv.org/html/2607.20475#A3.E12)\)

19:⊳\\triangleright— Gumbel\-Max Token Selection —

20:

t←𝐣​\[arg⁡max⁡\(𝐯−𝐠​\[𝐣\]\)\]t\\leftarrow\\mathbf\{j\}\[\\arg\\max\(\\mathbf\{v\}\-\\mathbf\{g\}\[\\mathbf\{j\}\]\)\]
21:endif

22:return

tt

## Appendix FBaseline Implementation Details

We evaluate all baselines under matched configurations \(precision, input shapes, and hardware\)\. TileLang uses its native compiler, and RadiK is evaluated as a precompiled CUDA C\+\+ binary\. For a subset of methods, we apply minimal adaptations to ensure correctness and compatibility with large vocabulary sizes andbfloat16inputs; these modifications are conservative and, where measurable, do not degrade and in some cases slightly improve the baseline performance\.

#### FlashInfer\.

FlashInfer’s sampling kernel generates Gumbel noise internally from a random seed, which breaks CUDA graph compatibility\. We modify the kernel to accept precomputed Gumbel noise as an external input, enabling CUDA graph capture and ensuring a fair comparison with our approach\. We also adapt the kernel to operate onbfloat16inputs to match our evaluation precision\.

#### TileLang\-TopK\.

The original implementation supports onlyfloat32inputs\. We extend it to supportbfloat16by introducing an explicit cast path during input loading and promoting values tofloat32for intermediate computation\. We additionally increase the shared\-memory buffer size from 4,096 to 16,384 entries to avoid silent truncation of candidates at large vocabulary sizes\. Even with this adjustment, TileLang\-TopK remains limited toV≤217V\\leq 2^\{17\}due to shared\-memory capacity constraints\.

#### Triton\-TopK\.

The original implementation uses a fused value\-index packing scheme that restricts the maximum supported vocabulary size toV≤215V\\leq 2^\{15\}\. We retain this implementation as\-is and report configurations beyond this limit as unsupported\.

## Appendix GAdditional Top\-kkResults

[Figure˜8](https://arxiv.org/html/2607.20475#A7.F8)reports additional Top\-kkresults comparing SonicSampler against FlashInfer and TileLang\. This figure complements the main\-text evaluation by providing a more focused view of the strongest baseline implementations\. Across the tested settings, SonicSampler consistently outperforms both methods, with the relative gains generally increasing as the vocabulary size and workload scale grow\.

![Refer to caption](https://arxiv.org/html/2607.20475v1/x4.png)Figure 8:Top\-kkspeedup against FlashInfer and TileLang fork=128k=128on B200\.Table 3:Normalized latency ofbitonic\(bit\) vs\.adaptive\(adp\) sorting strategies across vocab sizes and batch sizes on H100 \(k=128k=128\)\. All values are divided bymin⁡\(bitonic1,adaptive1\)\\min\(\\text\{bitonic\}\_\{1\},\\,\\text\{adaptive\}\_\{1\}\)for the corresponding architecture and vocab size, so 1\.00 anchors the fastestbatch\_size = 1configuration\.Boldmarks the lower \(better\) value per cell; ties bold both\.
## Appendix HAdditional Accuracy Results

[Figure˜9](https://arxiv.org/html/2607.20475#A9.F9)shows that the introduction of the constrained Top\-kksetting \(k≤128k\\leq 128\) in SonicSampler does not degrade the model quality across benchmarks\. In GPQA, both configurations achieve nearly identical accuracy \(78\.0 vs\. 78\.2\), with overlapping error bars that indicate no statistically significant differences\.

On AIME 2024, we observe the same trend where average accuracy remains unchanged\. Importantly, these differences fall within the observed variance across runs, suggesting that the constrained sampling space does not negatively impact solution diversity or correctness\.

## Appendix IAdditional Probability Mass Validation under Bounded Top\-kk

This section extends the bounded Top\-kkanalysis from[Section˜4\.5](https://arxiv.org/html/2607.20475#S4.SS5)to additional models\. In the main text, we show on DeepSeek v3\.1 that the residual probability mass outside the retained Top\-kkset decays rapidly, making bounded Top\-kkeffectively lossless in practice\.[Figure˜10](https://arxiv.org/html/2607.20475#A9.F10)shows that the same qualitative behavior holds across GLM\-4\.7\(5team2025glm45agenticreasoningcoding\), GPT\-OSS\-120B\(openai2025gptoss120bgptoss20bmodel\), and Qwen3\-8B\(yang2025qwen3technicalreport\)\. Across these models, probability mass concentrates quickly in a small candidate set, providing further empirical support for usingk=128k=128in SonicSampler\.

![Refer to caption](https://arxiv.org/html/2607.20475v1/figures/topk_accuracy_deepseek-v3.1-gpqa.png)\(a\)DeepSeek\-V3\.1 \- GPQA
![Refer to caption](https://arxiv.org/html/2607.20475v1/figures/grouped_accuracy_deepseek-v3.1-aime24.png)\(b\)DeepSeek\-V3\.1 \- AIME 2024

Figure 9:End\-to\-end accuracy on DeepSeek\-V3\.1\. Error bars show standard deviation over 10 runs\.![Refer to caption](https://arxiv.org/html/2607.20475v1/figures/topk_mass_glm-4.7-fp4.png)\(a\)GLM\-4\.7
![Refer to caption](https://arxiv.org/html/2607.20475v1/figures/topk_mass_gpt-oss-120b.png)\(b\)GPT\-OSS\-120B
![Refer to caption](https://arxiv.org/html/2607.20475v1/figures/topk_mass_qwen3-8b.png)\(c\)Qwen3\-8B

Figure 10:Residual probability mass outside the Top\-kkset on additional models\. Across all three model families, probability mass decays rapidly withkk, indicating that bounded Top\-kkretains the effective support relevant for practical sampling thresholds\.