一文掌握大模型推理加速：FlashAttention与PagedAttention技术详解，零基础小白收藏这一篇就够了！！

文章详细解析了FlashAttention和PagedAttention两种大模型推理加速技术。FlashAttention通过分块计算和重计算机制将内存访问次数从O(n²)降至O(n)，显著提升计算效率；PagedAttention借鉴操作系统虚拟内存思想，将KV Cache分割为固定大小页面，通过逻辑页表映射管理非连续显存，解决碎片化问题。两者分别从计算效率和内存管理两个维度优化大模型推理性能

健忘的派大星

604人浏览 · 2025-09-24 15:29:18

健忘的派大星 · 2025-09-24 15:29:18 发布

前言

本文旨在深度研究两种业界领先的注意力加速技术：FlashAttention 和 PagedAttention。 将详细阐述 FlashAttention 的 I/O 感知设计、在线 Softmax 的数学原理与推导、以及其在硬件层面的优化；同时，将深入解析 PagedAttention 借鉴操作系统虚拟内存思想所实现的 KV Cache 管理机制。

在大语言模型（LLM）的推理过程中，Attention 是计算和内存消耗的主要瓶颈。FlashAttention 和 PagedAttention 作为两项革命性的优化技术，分别从计算效率和内存管理两个维度显著提升了LLM的推理性能。

FlashAttention通过分块计算和重计算机制，将注意力计算的内存访问次数从O(n²)降至O(n)，在A100 GPU上可实现高达7.6倍的训练速度提升；

PagedAttention 借鉴操作系统的内存分页思想，将KV缓存分割为固定大小的页面，通过逻辑页表映射管理非连续物理显存，解决了传统KV缓存的碎片化问题，在高并发场景下可使vLLM的吞吐量比原生HuggingFace Transformers高出24倍。

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

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

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

Transformer模型架构

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

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

带宽瓶颈：现代GPU的内存层次结构中，高带宽内存（HBM）提供较大的存储容量（如A100的40-80GB）但速度较慢（约1.5-2.0TB/s），而片上静态随机访问存储器（SRAM）速度更快（约19TB/s）但容量极小（每个流式多处理器仅192KB）。注意力机制的计算速度远超内存速度，当考虑Transformer时，大多数操作受限于内存访问效率，成为性能提升的关键限制因素。

并行化限制：传统注意力机制难以高效利用 GPU 的并行计算能力，尤其是在处理不同长度的序列时。同时，KV Cache的动态增长特性使得批量处理变得复杂，进一步降低了GPU利用率。

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

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

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

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

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

核心原理

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

分块策略： FlashAttention将Q、K、V矩阵分割为多个小块，块大小B_q和B_k的选择需平衡 SRAM 容量与计算效率。具体来说，将Q分成Tr块，每块大小为Br×d；将K和V分成Tc块，每块大小为Bc×d。在计算时，逐块加载Q、K、V到SRAM，计算局部注意力得分，再通过累积统计量（最大值和指数和）得到全局结果。
重计算机制：在反向传播中，FlashAttention不需要存储O(n²)的中间注意力矩阵，而是通过存储输出O、softmax归一化统计量（最大值m和logsumexp L）来重新计算梯度。

FlashAttention-2 改进

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

具体改进包括：

消除频繁的系数更新，减少对非矩阵乘法运算的需求
在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[4] 针对Hopper架构（如H100 GPU）的硬件特性进行了深度优化，通过异步执行和低精度计算实现了突破性性能提升。

其核心改进包括：

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

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

硬件加速的低精度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倍数值误差

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

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

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

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

内存浪费与碎片化：大部分请求的实际序列长度远小于最大长度，导致大量预分配的内存被浪费（内部碎片）。同时，不同请求序列长度各异，使得内存管理变得困难，容易产生无法被新请求利用的小块内存（外部碎片）。
并发吞吐量低：由于内存浪费严重，一块 GPU 能容纳的并发请求数受限，导致吞吐量低下。
复制开销大：在某些复杂的采样策略（如 Beam Search）中，需要复制和管理多个序列的 KV Cache，这在连续内存布局下意味着昂贵的 memcpy 操作。

现有系统中的KV缓存管理。存在三种类型的内存浪费——保留、内部碎片和外部碎片——它们会阻止其他请求进入内存。每个内存插槽中的token代表其KV缓存。请注意，相同的token在不同位置时可以具有不同的KV缓存。

核心原理

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

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

KV缓存分块 (Paging) ：PagedAttention 将每个序列的 KV Cache 划分为固定大小的块（Block 或 Page）。每个块可以存储固定数量（例如 16 或 32 个）token 的 Key 和 Value 。
块表 (Block Table) ：系统为每个请求维护一个“块表”（类似于操作系统的页表），该表存储了逻辑块到物理块的映射关系。逻辑块是连续的，代表了 token 在序列中的位置；而物理块则是 GPU 显存中实际存储 KV 数据的、非连续的内存块。
动态按需分配：当模型生成新 token 时，调度器只需为这个新 token 分配一个新的物理块，并更新该序列的块表即可。这种按需分配的方式避免了预先分配大块连续内存所造成的浪费。当一个请求结束时，其占用的所有物理块都可以被回收并用于其他请求。

具有 PagedAttention 的请求的示例生成过程

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

对多个输出进行采样的请求示例生成过程

性能优势

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

四、FlashAttention与PagedAttention的区别与联系

FlashAttention和PagedAttention虽然都针对Transformer模型的性能瓶颈，但它们的优化方向、技术手段和适用场景有明显差异，同时也存在互补关系。

特性	FlashAttention	PagedAttention
主要优化方向	计算效率（减少HBM访问次数）	显存管理（消除碎片化，提高利用率）
适用阶段	训练/推理中的注意力计算	推理阶段的KV Cache管理
核心技术	分块计算 + 核函数融合 + 安全softmax	显存分页 + 逻辑页表映射 + 写时复制
内存占用	减少中间矩阵（QK^T）的显存占用	减少KV Cache的碎片化和浪费
计算延迟	通过分块和核函数融合降低延迟	通过内存复用减少分配开销

技术联系：

虽然两者优化方向不同，但它们在实际应用中形成了互补关系。FlashAttention优化了单次注意力计算的效率，而PagedAttention解决了多请求场景下的显存管理问题。 在长序列推理中，两者结合使用可以同时解决计算和内存两个维度的瓶颈。

实现差异：

FlashAttention 通过CUDA内核融合（kernel fusion）将多个操作（矩阵乘法、softmax、dropout等）集成到一个内核中执行，减少内存访问次数。而PagedAttention 则通过块表（block table）管理逻辑块与物理块的映射关系，实现动态内存分配和共享。

适用场景：

FlashAttention 更适合处理长序列的训练和推理场景，因为它显著降低了注意力计算的内存访问次数。而PagedAttention在高并发推理服务中表现更佳，因为它有效解决了显存碎片化和内存共享问题。

写在最后

FlashAttention 和 PagedAttention 是大模型时代背景下，针对注意力机制不同性能瓶颈的两个里程碑式创新。

FlashAttention 是一种计算优化技术。它通过 I/O 感知的算法设计，将计算瓶颈从内存带宽转移回计算本身，是底层算法与硬件结构协同优化的典范。它主要在训练和单次长序列推理中发挥巨大作用。

PagedAttention 是一种内存管理技术。它通过精细化的内存分页和调度，解决了推理服务中 KV Cache 的管理难题，是系统工程层面的重大突破。它主要在高并发推理服务中提升系统吞吐量和资源利用率。

它们代表了大模型推理优化的两个重要方向：计算效率和内存管理。 它们的结合使用为当前LLM推理性能的大幅提升做出了关键贡献，也将继续引领未来大模型推理技术的发展。

最后

为什么要学AI大模型

当下，⼈⼯智能市场迎来了爆发期，并逐渐进⼊以⼈⼯通⽤智能（AGI）为主导的新时代。企业纷纷官宣“ AI+ ”战略，为新兴技术⼈才创造丰富的就业机会，⼈才缺⼝将达 400 万！

DeepSeek问世以来，生成式AI和大模型技术爆发式增长，让很多岗位重新成了炙手可热的新星，岗位薪资远超很多后端岗位，在程序员中稳居前列。

在这里插入图片描述

与此同时AI与各行各业深度融合，飞速发展，成为炙手可热的新风口，企业非常需要了解AI、懂AI、会用AI的员工，纷纷开出高薪招聘AI大模型相关岗位。
在这里插入图片描述
最近很多程序员朋友都已经学习或者准备学习 AI 大模型，后台也经常会有小伙伴咨询学习路线和学习资料，我特别拜托北京清华大学学士和美国加州理工学院博士学位的鲁为民老师给大家这里给大家准备了一份涵盖了AI大模型入门学习思维导图、精品AI大模型学习书籍手册、视频教程、实战学习等录播视频 全系列的学习资料，这些学习资料不仅深入浅出，而且非常实用，让大家系统而高效地掌握AI大模型的各个知识点。