SonicSampler: Unified Tile-Aware Kernels for LLM Sampling and Speculative Verification
Pragaash Ponnusamy · Shivam Sahni · Jue WANG · Tri Dao
SonicSampler fuses the entire LLM sampling and speculative verification pipeline into unified tile-aware kernels partitioned by an optimized 2-stage hierarchical top-k, achieving up to 16× speedup over competitive baselines.
摘要
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.