用分块搜索提速文本生成,效率翻倍还省一半内存。
Qrita: High-performance Top-k and Top-p using Pivot-based Truncation and Selection
- 基于分片的搜索机制,先缩小候选范围再精准定位。
- 实测推理吞吐最高提升40%,内存占用减半且输出不变。
- 适合追求高效推理的开发者,已成vLLM默认采样器。
尽管在模型采样中至关重要,针对大词汇量的Top-k和Top-p算法高效实现仍是重大挑战。现有方法通常依赖排序,导致GPU计算与内存开销巨大,或采用随机方法改变输出结果。本文提出Qrita,一种基于分片截断与选择的高效Top-k和Top-p算法。Qrita通过两项关键技术实现优化:1. 基于高斯分布的sigma截断,大幅缩减词表搜索空间;2. 四元分片搜索结合重复处理机制,将搜索迭代次数减半并保证确定性输出。我们使用Triton实现Qrita,评估其在SGLang和FlashInfer等高性能LLM引擎的Top-k/Top-p内核上的表现,结果显示端到端服务吞吐最高提升1.4倍,内存使用减半,且输出与排序算法完全一致。Qrita现已作为vLLM GPU执行路径的默认采样器,其三元版本代码见https://github.com/vllm-project/vllm/blob/main/vllm/v1/sample/ops/topk_topp_triton.py。
原文摘要 · Abstract (English)
Despite their importance in model sampling, efficient implementation of Top-k and Top-p algorithms for large vocabularies remains a significant challenge. Existing approaches often rely on sorting, which incurs significant computation and memory overhead on GPUs, or on stochastic approaches that alter the algorithm's output. In this work, we propose Qrita, an efficient Top-k and Top-p algorithm based on a pivot-based truncation and selection. Qrita leverages pivot-based search for both Top-k and Top-p with two key techniques: 1. Gaussian-based sigma-truncation, which greatly reduces the search space of the vocabulary, and 2. Quaternary pivot search with duplication handling, which halves the number of pivot search iterations and guarantees deterministic output. We implement Qrita using Triton and evaluate its performance against the Top-k and Top-p kernels of high-performance LLM execution engines such as SGLang and FlashInfer, improving end-to-end serving throughput up to 1.4 times with half the memory usage, while providing the same output as the sorting-based algorithms. Qrita is now the default Top-k and Top-p sampler for the GPU execution path of vLLM, and a ternary implementation of Qrita is available at https://github.com/vllm-project/vllm/blob/main/vllm/v1/sample/ops/topk_topp_triton.py.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。