用于加速大语言模型推理的自选择注意力跨度
Self-Selected Attention Span for Accelerating Large Language Model Inference
AI总结：
针对大语言模型推理时注意力token增多导致效率低的问题，提出通过微调让模型自主选择每步最小注意力跨度并生成稀疏掩码，结合自定义CUDA内核实现推理吞吐量提升28%，为加速实际任务的自回归推理提供了端到端方案。
摘要：
大语言模型（LLM）能够解决具有挑战性的任务。然而，由于在生成新token时需要关注的token数量不断增加，它们在现代GPU上的推理计算效率极低。为解决这一效率问题，我们利用LLM的问题解决能力来优化其自身的推理阶段效率。我们通过两项具体任务进行验证：（a）复杂算术表达式求值，（b）新闻文章摘要。针对这两项任务，我们构建了自定义数据集来微调LLM。微调的目标有两方面：一是让LLM学会完成求值或摘要任务，二是训练它识别任务每一步所需的最小注意力跨度。最终，微调后的模型能够在推理过程中实时将这些自主识别的最小注意力跨度转换为稀疏注意力掩码。我们开发了自定义CUDA内核以利用缩减后的待关注上下文。实验表明，使用该自定义CUDA内核可将LLM推理吞吐量提升28%。我们的工作提供了端到端的验证，证明训练LLM自主选择注意力跨度能够加快其在解决实际任务时的自回归推理速度。
英文摘要：
Large language models (LLMs) can solve challenging tasks. However, their inference computation on modern GPUs is highly inefficient due to the increasing number of tokens they must attend to as they generate new ones. To address this inefficiency, we capitalize on LLMs' problem-solving capabilities to optimize their own inference-time efficiency. We demonstrate with two specific tasks: (a) evaluating complex arithmetic expressions and (b) summarizing news articles. For both tasks, we create custom datasets to fine-tune an LLM. The goal of fine-tuning is twofold: first, to make the LLM learn to solve the evaluation or summarization task, and second, to train it to identify the minimal attention spans required for each step of the task. As a result, the fine-tuned model is able to convert these self-identified minimal attention spans into sparse attention masks on-the-fly during inference. We develop a custom CUDA kernel to take advantage of the reduced context to attend to. We demonstrate that using this custom CUDA kernel improves the throughput of LLM inference by 28%. Our work presents an end-to-end demonstration showing that training LLMs to self-select their attention spans speeds up autoregressive inference in solving real-world tasks.
机构、总结、中文摘要均由AI分析而来，可能存在错误，请以论文正文为准。
