SonicSampler: Unified Tile-Aware Kernels for LLM Sampling and Speculative Verification
06:00 · July 24, 2026 · arXiv cs.AI RSS
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 present $\textbf{SonicSampler}$, 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-$k$ / top-$p$ / min-$p$ filtering, 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-$k$ algorithm that achieves up to $\textbf{10x speedup}$ over 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 to $\textbf{16x speedup}$ over state-of-the-art baselines while preserving flexible batched execution.
Summary
SonicSampler addresses a persistent bottleneck in LLM inference by fusing the full sampling pipeline—logit processing, token selection, and speculative verification—into a single set of tile-aware Triton kernels. Existing approaches typically handle only isolated steps such as top-k filtering or grammar masking, issue multiple kernel launches, or enforce uniform sampling parameters across a batch. These limitations increase memory traffic, prevent CUDA Graph capture, and restrict use in continuous-batching environments where requests may require different penalties, temperature settings, or constrained decoding rules.
The kernels support fully dynamic per-request configurations within one batched dispatch. Grammar-constrained decoding, repetition and frequency penalties, logit bias, temperature scaling, and top-k/top-p/min-p filtering are all applied on the fly, alongside verification steps required by speculative decoding. A compact bit-level indicator encodes each request’s sampling mode, allowing mixed greedy and stochastic behavior without separate kernel launches or host-side branching.
At the core of the design is a two-stage hierarchical top-k algorithm that exploits the low-entropy structure of LLM output distributions. The first stage processes vocabulary tiles independently, applying the full logit-processing prologue before reducing each tile to a bounded set of candidates. The second stage merges these candidates across tiles. By keeping the candidate set small (k = 128) and fusing compute-intensive operations with the reduction, the method achieves up to 10× speedup over competitive radix- and bitonic-based baselines while remaining compatible with large vocabularies. In heterogeneous speculative-decoding workloads the same kernels deliver up to 16× end-to-end improvement over state-of-the-art libraries such as FlashInfer.
Why it matters
This research is highly relevant for AI infrastructure researchers and engineers in the Netherlands looking to optimize LLM inference and reduce compute costs. The significant speedups in sampling and speculative decoding align with European goals for sustainable and efficient AI deployment.
