中文

LeetDecoding:一个带有 CUDA 实现的指数衰减因果线性注意力 PyTorch 库

机器学习 2025-01-07 v1 计算与语言 数学软件

摘要

机器学习与数据科学社区在加速基于 Transformer 的大语言模型(LLM)方面取得了显著但较为分散的进展,一种极具前景的方法是用指数衰减因果线性注意力替换生成式预训练 Transformer(GPT)中原始的因果注意力。在本文中,我们提出了 LeetDecoding,这是第一个为这一基础算子提供大量计算例程的 Python 包。LeetDecoding 的推出是由于目前缺乏:(1) 对该算子复杂度的清晰理解;(2) 对现有计算方法的全面收集(通常散布于看似不相关的领域中);以及 (3) 用于 GPU 快速推理的 CUDA 实现。LeetDecoding 的设计易于与现有的线性注意力 LLM 集成,并允许研究人员对指数衰减因果线性注意力的新计算方法进行基准测试和评估。LeetDecoding 的使用不需要任何 GPU 编程和底层复杂度分析知识,旨在让 LLM 从业者能够便捷地使用。LeetDecoding 的源代码提供于该 GitHub 仓库,用户可通过命令 \texttt{pip install leet-decoding} 直接安装。

关键词

引用

@article{arxiv.2501.02573,
  title  = {LeetDecoding: A PyTorch Library for Exponentially Decaying Causal Linear Attention with CUDA Implementations},
  author = {Jiaping Wang and Simiao Zhang and Qiao-Chu He and Yifan Chen},
  journal= {arXiv preprint arXiv:2501.02573},
  year   = {2025}
}

备注

The source code of LeetDecoding is hosted at https://github.com/Computational-Machine-Intelligence/LeetDecoding