​​LLM推理效率的范式转移:FlashAttention与PagedAttention正在重塑AI部署的未来​

本文涉及的产品
多模态交互后付费免费试用,全链路、全Agent
简介: 本文深度解析FlashAttention与PagedAttention两大LLM推理优化技术:前者通过分块计算提升注意力效率,后者借助分页管理降低KV Cache内存开销。二者分别从计算与内存维度突破性能瓶颈,显著提升大模型推理速度与吞吐量,是当前高效LLM系统的核心基石。建议收藏细读。

本文较长,建议点赞收藏,以免遗失。

在大语言模型(LLM)的推理过程中,Attention机制是计算和内存消耗的主要瓶颈。FlashAttention和PagedAttention作为两项革命性优化技术,分别从计算效率和内存管理两个维度显著提升了LLM的推理性能。今天我将深度解析这两种注意力加速技术及其优化策略,希望对你有所帮助,如有遗漏,欢迎交流。

一、Transformer中的Attention机制与性能瓶颈

Transformer的核心在于自注意力机制,它通过计算序列中各token之间的关系权重,实现全局交互。标准自注意力机制的计算公式为:

ef11b2bc1f21b785a8d59e20528d1ad5__fallback_source=1&height=1280&mount_node_token=NgvfdJId8oDqPlxPMAtcepYdnvh&mount_point=docx_image&policy=allow_down&width=1280.jpg

其中Q(查询)、K(键)、V(值)是输入序列经过线性变换后的三个矩阵,dₖ为特征维度。

f685b404375489e1167b798d4be0e6fc.png

自注意力机制的时间复杂度为O(n²d),空间复杂度为O(n²+Nd),其中n为序列长度,d为特征维度。这种二次复杂度导致长序列处理时计算量和内存占用急剧增加,成为模型扩展的主要障碍。

在推理过程中,KV Cache需要存储所有已生成token的键值对,以便后续计算。传统实现要求KV Cache使用连续的显存空间,但实际序列长度各不相同,导致大量碎片化内存无法被有效利用。研究表明,由于碎片化和过度预留,现有系统浪费了60%-80%的显存。

(ps:这里补充一个知识点,如果你对KV Cache的工作原理和机制不了解,我这里也整理了一份小白都能看懂的技术文档,粉丝朋友自行查阅:《小白也能看懂的LLMs中的KV Cache,视觉解析》

二、FlashAttention:I/O感知的精确注意力算法

FlashAttention是一种针对Transformer模型中注意力机制的计算优化技术,由斯坦福DAWN实验室于2022年提出。其并非注意力的近似算法,而是一种在数学上与标准注意力等价,但实现方式上完全不同的精确注意力算法。

6969.jpg

FlashAttention的核心目标是减少对HBM的访问次数,最大限度地利用GPU上速度更快的片上SRAM,特别适用于处理长序列任务。

标准注意力机制使用HBM来存储、读取和写入注意力分数矩阵。具体步骤为将这些从HBM加载到GPU的片上SRAM,然后执行注意力机制的单个步骤,然后写回HBM,并重复此过程。

而Flash Attention采用分块计算(Tiling)技术,将大型注意力矩阵划分为多个块(tile),在SRAM中逐块执行计算。

核心原理

FlashAttention的核心思想是分块计算(tiling)和重计算(recomputation)。将输入矩阵Q、K、V分割成小块,每次只处理其中一块数据,在GPU的片上SRAM中进行计算,避免频繁访问高带宽内存(HBM)。同时,通过安全softmax(safe softmax)技术处理数值稳定性问题,并在反向传播中利用存储的中间统计量重新计算梯度,而非存储整个注意力矩阵。

分块策略将Q、K、V矩阵分割为多个小块,块大小的选择需平衡SRAM容量与计算效率。具体来说,将Q分成Tr块,每块大小为Br×d;将K和V分成Tc块,每块大小为Bc×d。在计算时,逐块加载Q、K、V到SRAM,计算局部注意力得分,再通过累积统计量(最大值和指数和)得到全局结果。

重计算机制在反向传播中不需要存储O(n²)的中间注意力矩阵,而是通过存储输出O、softmax归一化统计量(最大值m和logsumexp L)来重新计算梯度。

FlashAttention-2改进

2023年提出的FlashAttention-2进一步优化了工作分区(work partitioning),减少了非矩阵乘法运算的需求,并实现了序列长度上的并行化。

image.png

具体改进包括:消除频繁的系数更新,减少对非矩阵乘法运算的需求;在CUDA线程块内分配工作到不同warp上,减少通信和共享内存读写;仅存储logsumexp统计量,而非同时存储max和sum。

FlashAttention在A100 GPU上实现了显著的性能提升:BERT模型训练速度比MLPerf 1.1记录高出15%;GPT-2训练速度比HuggingFace实现高3倍,比Megatron高1.7倍;FlashAttention-2在A100上达到理论FLOPs的50-73%,训练速度提升至225 TFLOPs/s。

FlashAttention-3改进

2024年提出的FlashAttention-3针对Hopper架构(如H100 GPU)的硬件特性进行了深度优化,通过异步执行和低精度计算实现了突破性性能提升。

其核心改进包括:生产者-消费者异步,通过定义一个warp-specialized软件流水线方案,利用数据移动和Tensor Cores的异步执行,将生产者和消费者分为不同的warps,从而延长算法隐藏内存和指令发出延迟的能力。

image.png

在异步块状GEMM下隐藏softmax,通过重叠低吞吐量的非GEMM操作(如浮点乘加和指数运算)与异步WGMMA指令进行块矩阵乘法,重写FlashAttention-2算法以规避softmax和GEMM之间的某些顺序依赖性。

image.png

硬件加速的低精度GEMM调整前向传播算法以针对FP8 Tensor Cores进行GEMM,几乎翻倍了测量的TFLOPs/s。这需要弥合WGMMA在块布局一致性方面的要求,使用块量化和非相干处理来减轻由于转换为FP8精度而导致的精度损失。

在H100 GPU上实现里程碑式突破:FP16性能达到740 TFLOPs/s(75%理论峰值),较FlashAttention-2提升1.5-2.0倍;FP8性能接近1.2 PFLOPs/s,首次突破PetaFLOP级注意力计算;数值精度方面,FP8版本比基线FP8注意力降低2.6倍数值误差。

image.png

三、PagedAttention:高效的KV Cache内存管理

PagedAttention是一种显存管理优化技术,由UC Berkeley团队提出并集成到vLLM框架中。它借鉴了操作系统中虚拟内存和分页的经典思想,将KV Cache分割成固定大小的页面(pages),允许在非连续内存空间中存储连续的KV张量,从而有效解决显存碎片化问题。

核心问题:KV Cache管理的挑战

传统的KV Cache管理方式通常为每个请求预分配一个连续的大块内存,其大小等于模型的最大序列长度。这种方式存在严重问题:

内存浪费与碎片化:大部分请求的实际序列长度远小于最大长度,导致大量预分配的内存被浪费(内部碎片)。同时,不同请求序列长度各异,使得内存管理变得困难,容易产生无法被新请求利用的小块内存(外部碎片)。

并发吞吐量低:由于内存浪费严重,一块GPU能容纳的并发请求数受限,导致吞吐量低下。

复制开销大:在某些复杂的采样策略(如Beam Search)中,需要复制和管理多个序列的KV Cache,这在连续内存布局下意味着昂贵的memcpy操作。

image.png

核心原理

PagedAttention将每个序列的KV缓存进行分块,每个块包含固定长度的token的KV对。通过页面表(page table)维护逻辑地址到物理地址的映射,实现动态内存分配和共享。

111.gif

▲图示:PagedAttention:KV Cache 被划分成块。块在内存空间上不需要连续

PagedAttention将KV Cache划分成块,块在内存空间上不需要连续。系统为每个请求维护一个"块表"(类似于操作系统的页表),该表存储了逻辑块到物理块的映射关系。逻辑块是连续的,代表了token在序列中的位置;而物理块则是GPU显存中实际存储KV数据的、非连续的内存块。

当模型生成新token时,调度器只需为这个新token分配一个新的物理块,并更新该序列的块表即可。这种按需分配的方式避免了预先分配大块连续内存所造成的浪费。当一个请求结束时,其占用的所有物理块都可以被回收并用于其他请求。

222.gif
▲图示:具有 PagedAttention 的请求的示例生成过程

PagedAttention通过写时复制(Copy-on-Write,CoW)实现高效的内存共享。当多个序列共享同一个提示(prompt)时,它们的KV缓存页面可以指向同一个物理块。系统跟踪物理块的引用计数,当某个序列尝试修改共享页面时,才触发复制到新页面并更新引用计数。这种机制在并行采样和集束搜索中特别有效,可将内存使用量降低55%,吞吐量提升2.2倍。

333.gif
▲图示:对多个输出进行采样的请求示例生成过程

性能优势

vLLM在相同延迟水平下,吞吐量较基线系统提升2-4倍,长序列、大模型场景优势更显著(如OPT-175B在Alpaca数据集上提升3.58倍)。内存浪费趋近于零,而基线系统浪费达61.8%-79.6%。在束搜索中内存共享节省37.6%-66.3%,共享前缀场景减少冗余计算16.2%-30.5%。块大小设为16时平衡并行性与碎片率,重计算较交换降低20%恢复开销。

写在最后

  • FlashAttention是一种计算优化技术。它通过I/O感知的算法设计,将计算瓶颈从内存带宽转移回计算本身,是底层算法与硬件结构协同优化的典范。它主要在训练和单次长序列推理中发挥巨大作用。
  • PagedAttention是一种内存管理技术。它通过精细化的内存分页和调度,解决了推理服务中KV Cache的管理难题,是系统工程层面的重大突破。它主要在高并发推理服务中提升系统吞吐量和资源利用率。

它们代表了大模型推理优化的两个重要方向:计算效率和内存管理,它们的结合使用使当前LLM推理性能的大幅提升。好了,今天的分享就到这里,如果对你有所帮助,记得点个小红心,我们下期见。

目录
相关文章
|
9天前
|
消息中间件 人工智能 运维
事件驱动重塑 AI 数据链路:阿里云 EventBridge 发布 AI ETL 新范式
“一个简单的数据集成任务,开始时总是轻松愉快的,但随着业务扩展,数据源越来越多,格式越来越乱,整个数据链路就会变得一团糟。”陈涛在演讲中指出了当前 AI 数据处理的普遍困境。扩展难、运维难、稳定性差,这三大挑战已成为制约 AI 应用创新和落地的关键瓶颈。针对这些痛点,在2025云栖大会期间,阿里云重磅发布了事件驱动 AI ETL 新范式,其核心产品 EventBridge 通过深度集成 AI 能力,为开发者提供了一套革命性的解决方案,旨在彻底改变 AI 时代的数据准备与处理方式。
|
21天前
|
人工智能 Ubuntu 前端开发
Dify部署全栈指南:AI从Ubuntu配置到HTTPS自动化的10倍秘籍
本文档介绍如何部署Dify后端服务及前端界面,涵盖系统环境要求、依赖安装、代码拉取、环境变量配置、服务启动、数据库管理及常见问题解决方案,适用于开发与生产环境部署。
282 1
|
15天前
|
人工智能 JSON 运维
告别假期规划的“人肉运维”!这个AI提示词,帮你“一键部署”十一完美旅行
十一黄金周将至,开发者们如何用“工程师思维”高效规划假期?本文推出一款AI旅行规划神器——结构化提示词,无需代码,只需填空,即可生成专业级定制攻略。从痛点分析到即用模板,涵盖行程、预算、避坑指南,助你一键生成完美假期方案,把时间留给诗和远方。
270 124
|
21天前
|
存储 人工智能 NoSQL
用Context Offloading解决AI Agent上下文污染,提升推理准确性
上下文工程是将AI所需信息(如指令、数据、工具等)动态整合到模型输入中,以提升其表现。本文探讨了“上下文污染”问题,并提出“上下文卸载”策略,通过LangGraph实现,有效缓解长文本处理中的信息干扰与模型幻觉,提升AI代理的决策准确性与稳定性。
95 2
用Context Offloading解决AI Agent上下文污染,提升推理准确性
|
27天前
|
机器学习/深度学习 人工智能 前端开发
解决推理能力瓶颈,用因果推理提升LLM智能决策
从ChatGPT到AI智能体,标志着AI从对话走向自主执行复杂任务的能力跃迁。AI智能体可完成销售、旅行规划、外卖点餐等多场景任务,但其发展受限于大语言模型(LLM)的推理能力。LLM依赖统计相关性,缺乏对因果关系的理解,导致在非确定性任务中表现不佳。结合因果推理与内省机制,有望突破当前AI智能体的推理瓶颈,提升其决策准确性与自主性。
125 6
解决推理能力瓶颈,用因果推理提升LLM智能决策
|
21天前
|
人工智能 监控 数据可视化
如何破解AI推理延迟难题:构建敏捷多云算力网络
本文探讨了AI企业在突破算力瓶颈后,如何构建高效、稳定的网络架构以支撑AI产品化落地。文章分析了典型AI IT架构的四个层次——流量接入层、调度决策层、推理服务层和训练算力层,并深入解析了AI架构对网络提出的三大核心挑战:跨云互联、逻辑隔离与业务识别、网络可视化与QoS控制。最终提出了一站式网络解决方案,助力AI企业实现多云调度、业务融合承载与精细化流量管理,推动AI服务高效、稳定交付。
|
21天前
|
存储 人工智能 关系型数据库
阿里云AnalyticDB for PostgreSQL 入选VLDB 2025:统一架构破局HTAP,Beam+Laser引擎赋能Data+AI融合新范式
在数据驱动与人工智能深度融合的时代,企业对数据仓库的需求早已超越“查得快”这一基础能力。面对传统数仓挑战,阿里云瑶池数据库AnalyticDB for PostgreSQL(简称ADB-PG)创新性地构建了统一架构下的Shared-Nothing与Shared-Storage双模融合体系,并自主研发Beam混合存储引擎与Laser向量化执行引擎,全面解决HTAP场景下性能、弹性、成本与实时性的矛盾。 近日,相关研究成果发表于在英国伦敦召开的数据库领域顶级会议 VLDB 2025,标志着中国自研云数仓技术再次登上国际舞台。
135 0
|
6月前
|
机器学习/深度学习 存储 缓存
加速LLM大模型推理,KV缓存技术详解与PyTorch实现
大型语言模型(LLM)的推理效率是AI领域的重要挑战。本文聚焦KV缓存技术,通过存储复用注意力机制中的Key和Value张量,减少冗余计算,显著提升推理效率。文章从理论到实践,详细解析KV缓存原理、实现与性能优势,并提供PyTorch代码示例。实验表明,该技术在长序列生成中可将推理时间降低近60%,为大模型优化提供了有效方案。
985 15
加速LLM大模型推理,KV缓存技术详解与PyTorch实现