Ascend上的PageAttention

简介: PageAttention旨在解决大型语言模型(LLM)服务中的内存管理低效问题,如内存碎片化、利用率低及缺乏灵活的内存共享机制。通过借鉴操作系统中的虚拟内存和分页技术,PageAttention实现了块级别的内存管理和灵活的KV cache共享机制,显著提高内存利用率,降低延迟,提升模型处理速度和性能。相比传统注意力机制,PageAttention通过分段处理序列,有效解决了长序列处理时的计算效率低下和内存过度使用问题。

1 PageAttention引入的原因

PageAttention的引入主要是为了解决大型语言模型(LLM)在服务过程中遇到的内存管理低效问题,具体原因如下:

  • 内存碎片化:传统的KV cache(Key-Value cache)管理方式在处理不同长度的请求时,预分配连续的块内存给每个请求,导致内存碎片化和冗余复制,造成内存浪费。
  • 内存利用率低:在实际使用中,KV cache的利用率只有20.4%到38.2%,说明当前的内存管理机制没有充分利用可用的内存空间。
  • 缺乏灵活的内存共享机制:现有的系统没有实现有效的KV cache共享,而这是许多大模型解码算法中所需要的。

2 PageAttention的方案

为了解决这些问题,PageAttention借鉴了操作系统中的虚拟内存和分页技术,实现了以下架构和优化:

  • 块级别的内存管理:PageAttention将请求的KV cache放入块中,每个块可以存储固定数量tokens的attention keys和values,不需要连续的空间。这减少了内部的碎片化,同时也不需要为所有的请求都预留固定数目的内存大小。
  • 灵活的KV cache共享机制:通过虚拟内存的技术,PageAttention使得KV cache的存储更加灵活,提高了利用率,允许跨请求共享,进一步减少内存使用。
  • 高效内存管理:vLLM大模型服务系统实现了近乎零浪费的KV cache内存机制,通过PageAttention,内存管理变得更加高效,减少了内存浪费。
  • 高吞吐量和低延迟:实验结果显示,vLLM比流行的大模型有更低的延时,大概是2~4倍。此外,越长的句子,越大的模型,越复杂的decoding算法,相对提升更明显。

PageAttention的实现架构包括:

  • Paged KV Cache:将KV cache分割成多个页面,每个页面存储一定数量的tokens,允许非连续存储,减少内存碎片。
  • CUDA并行计算:在GPU上利用CUDA的并行计算能力,PageAttention可以在GPU上实现高效的矩阵运算和注意力计算,进一步提升模型的处理速度和性能。
  • V1和V2版本:vLLM中有两个版本的PageAttention,使用一个简单的启发式方法来决定是使用V1还是V2版本。V1适合长度小于8192或者num_seqs * num_heads>512的情况。

通过这些架构和优化,PageAttention显著提高了大型语言模型的服务效率,降低了延迟,提升了整体性能。

3 与标准注意力机制的区别

  • 标准注意力机制通常需要一次性处理整个序列,这在序列很长时可能导致计算效率低下和内存使用过多。
  • Paged Attention 通过将序列分段处理,减轻了这些问题,使得模型能够在更长的序列上工作,同时保持较低的计算成本。

Paged Attention 是一种较为高级的技术,它的实现通常需要在具体的应用框架中找到相关的API支持,或者自己编写代码来实现这种机制。在华为昇腾AI处理器的API文档中提到的 PagedAttentionOperation 就是这样一种操作,它提供了分页注意力机制的具体实现方式。

4 Ascend上的实现

ascend上的PagedAttention实现为,它是可通过atb算子实现。参考如下:
ATB中的PagedAttention算子

相关文章
|
14天前
|
存储 人工智能
AscendC编程中的double buffer是什么?
在AI Core上,指令队列分为Vector(V)、Matrix(M)和存储移动指令队列(MTE2、MTE3)。这些队列的独立性和可并行性是double buffer优化的基础。通过将数据搬运与Vector计算并行执行,double buffer机制有效减少了Vector单元的等待时间,提高了其利用率。例如,Tensor1进行Compute时,Tensor2可同时执行CopyIn;当切换到Tensor2计算时,Tensor1执行CopyOut。此机制适用于大多数场景,但在数据搬运时间短或数据量小的情况下,性能提升可能有限。
|
1月前
|
人工智能 PyTorch 算法框架/工具
Ascend Extension for PyTorch是个what?
Ascend Extension for PyTorch 是针对华为昇腾处理器的PyTorch框架适配插件,旨在让PyTorch开发者能充分利用昇腾AI处理器的强大计算能力。此扩展通过最小化对原生PyTorch的改动,实现了对昇腾NPU的支持,包括动态图特性、自动微分等功能的完整继承,并提供了与原生PyTorch一致的使用体验。项目详情及源码可在昇腾社区获取。
Ascend Extension for PyTorch是个what?
|
1月前
|
机器学习/深度学习 自然语言处理 并行计算
SelfAttention在Ascend上的实现
Self-Attention(自注意力)机制是深度学习中用于自然语言处理的关键技术,由 Vaswani 等人在 2017 年提出。它通过让模型关注输入序列的不同部分,增强了对上下文的理解,是 Transformer 架构的核心。Self-Attention 包括 Query、Key 和 Value 三个组件,通过计算这些组件间的相似度,模型可以有效地捕捉长距离依赖关系,提高处理效率和模型表现。此外,多头注意力机制进一步提升了模型的表达能力。
|
1月前
|
存储 缓存 人工智能
Ascend上的FlashAttention实现
FlashAttention是优化Transformer模型计算效率和内存使用的技术,通过减少存储访问开销提升性能。它采用Tiling、Recomputation、分块SoftMax等策略,减少HBM访问,加速计算,并在昇腾AI处理器上实现了显著的性能提升。
|
1月前
|
机器学习/深度学习 人工智能 算法框架/工具
什么是CANN和Ascend C
CANN(Compute Architecture for Neural Networks)是华为推出的AI异构计算架构,支持多种AI框架如MindSpore、PyTorch等,适用于AI处理器与编程,旨在提升昇腾AI处理器的计算效率。CANN提供强大的图引擎、算子开发语言Ascend C、算子加速库AOL、集合通信库HCCL、毕昇编译器及Runtime运行时,支持快速构建AI应用,涵盖推理应用开发、模型训练和算子开发等关键功能。
|
1月前
|
测试技术 开发者 异构计算
AscendC从入门到精通系列(二)基于Kernel直调开发AscendC算子
本文介绍了AscendC算子的开发流程,包括核函数开发、算子类定义及其实现、核函数的CPU和NPU侧运行验证。通过具体示例`add_custom.cpp`,详细展示了如何使用Ascend C完成算子核函数的定义、初始化、数据搬运和计算过程,并提供了完整的CPU和NPU侧调用程序代码,帮助开发者理解和实践AscendC算子的开发。
|
1月前
|
分布式计算 并行计算 编译器
Ascend C的编程模型
Ascend C采用SPMD编程模型,实现多核并行计算。通过将数据切分为多个分片,分配给不同的计算核心(或称为block)并行处理,每个核心根据自身的block_idx执行相同代码但处理不同的数据分片,从而提高计算效率。此模型支持高效的数据并行处理,适用于大规模数据计算任务。
|
1月前
|
PyTorch API 算法框架/工具
AscendC从入门到精通系列(四)使用Pybind调用AscendC算子
本文介绍了如何通过Pybind11在PyTorch框架中调用自定义的Ascend C算子。首先,通过编写算子的C++实现和pybind11封装,将算子功能暴露给Python。接着,构建Python调用脚本,利用torch接口生成数据并调用封装好的算子模块。最后,通过CMake配置文件编译整个项目,实现从算子开发到测试的完整流程。
|
1月前
|
API C语言 开发者
AscendC从入门到精通系列(五)调用基于工程开发AscendC算子
单算子API调用方式是通过C语言API直接调用已编译的自定义算子。首先,需基于AscendC算子工程完成算子的定义与实现,并通过编译脚本部署。编译后,生成的头文件和动态库支持在应用程序中直接调用算子,包括初始化AscendCL、申请资源、数据传输、计算workspace、执行算子、同步等待及资源释放等步骤。编译算子调用程序时,需正确配置CMakeLists.txt,确保头文件和动态库的路径正确。
|
1月前
|
存储 人工智能 JSON
AscendC从入门到精通系列(三)基于自定义算子工程开发AscendC算子
本文介绍了基于Ascend C的自定义算子开发流程,涵盖从工程创建、代码编写、编译部署到运行验证的全过程。以动态shape的AddCustom算子为例,详细描述了如何利用CANN提供的工具msOpGen生成开发工程,实现算子核函数与host侧代码,以及如何编译、部署和测试自定义算子。