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

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 present **SonicSampler**, a unified suite of tile-aware Triton kernels that vertically fuses the common production 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, entirely within a fixed set of batched kernels while remaining fully CUDA Graph-compatible. Central to our approach is a novel hierarchical two-stage top-$k$ algorithm that exploits the low-entropy structure of LLM outputs to enable efficient selection over large vocabularies and achieves up to **10x speedup** on 128K vocab sizes over competitive baselines. Across heterogeneous speculative decoding workloads, SonicSampler achieves up to **16x speedup** on draft sampling and up to **10x speedup** on target sampling and verification over state-of-the-art baselines while preserving flexible batched execution.