arXiv:2501.02573cs.LGcs.CL2025-01

首个支持指数衰减因果线性注意力的PyTorch库,加速大模型推理。

LeetDecoding: A PyTorch Library for Exponentially Decaying Causal Linear Attention with CUDA Implementations

  • 提供指数衰减因果线性注意力的高效计算工具
  • 集成CUDA实现,显著提升GPU推理速度
  • 无需懂底层编程,适合大模型开发者快速上手

机器学习社区在加速基于Transformer的大语言模型(LLM)方面取得进展,其中一种有前景的方法是用指数衰减因果线性注意力替代生成式预训练变压器(GPT)中的原始因果注意力。本文提出LeetDecoding,这是首个为该基础算子提供大量计算例程的Python包。其发布源于当前缺乏对该算子复杂度的清晰理解、现有计算方法的全面整合(常分散于不同领域),以及用于GPU快速推理的CUDA实现。LeetDecoding设计易于与现有线性注意力LLM集成,支持研究人员对新计算方法进行基准测试和评估。使用LeetDecoding无需掌握GPU编程或底层复杂度分析,旨在让大模型从业者轻松使用。源代码可在GitHub仓库获取,用户可通过命令'pip install leet-decoding'安装。

原文摘要 · Abstract (English)

The machine learning and data science community has made significant while dispersive progress in accelerating transformer-based large language models (LLMs), and one promising approach is to replace the original causal attention in a generative pre-trained transformer (GPT) with \emph{exponentially decaying causal linear attention}. In this paper, we present LeetDecoding, which is the first Python package that provides a large set of computation routines for this fundamental operator. The launch of LeetDecoding was motivated by the current lack of (1) clear understanding of the complexity regarding this operator, (2) a comprehensive collection of existing computation methods (usually spread in seemingly unrelated fields), and (3) CUDA implementations for fast inference on GPU. LeetDecoding's design is easy to integrate with existing linear-attention LLMs, and allows for researchers to benchmark and evaluate new computation methods for exponentially decaying causal linear attention. The usage of LeetDecoding does not require any knowledge of GPU programming and the underlying complexity analysis, intentionally making LeetDecoding accessible to LLM practitioners. The source code of LeetDecoding is provided at \href{https://github.com/Computational-Machine-Intelligence/LeetDecoding}{this GitHub repository}, and users can simply install LeetDecoding by the command \texttt{pip install leet-decoding}.

大模型加速线性注意力CUDA优化PyTorch库

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。