用分桶近似法提升top-k计算并行度，加速大模型推理
Approximate Top-$k$ for Increased Parallelism
- 分桶处理多个小规模top-k，突破传统算法串行瓶颈
- 在语言模型稀疏性任务中实现显著加速，保持精度损失可控
- 适合追求高并行的ML加速器场景，尤其大模型部署
我们评估了分桶近似top-k算法。精确计算top-k受限于向量上需聚合前k大值，难以在高度并行的机器学习加速器上高效运行。通过放宽对top-k精确性的要求，分桶算法可独立执行多个小型top-k操作，显著提升并行性。本文结合理论分析与下游任务实测，探索该类算法的设计选择。核心应用是语言模型中的稀疏性算法，常使用top-k筛选关键参数或激活值。我们还发布了适用于PyTorch的快速分桶top-k实现。
原文摘要 · Abstract (English)
We present an evaluation of bucketed approximate top-$k$ algorithms. Computing top-$k$ exactly suffers from limited parallelism, because the $k$ largest values must be aggregated along the vector, thus is not well suited to computation on highly-parallel machine learning accelerators. By relaxing the requirement that the top-$k$ is exact, bucketed algorithms can dramatically increase the parallelism available by independently computing many smaller top-$k$ operations. We explore the design choices of this class of algorithms using both theoretical analysis and empirical evaluation on downstream tasks. Our motivating examples are sparsity algorithms for language models, which often use top-$k$ to select the most important parameters or activations. We also release a fast bucketed top-$k$ implementation for PyTorch.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。