Ascend上的FlashAttention实现

简介: FlashAttention是优化Transformer模型计算效率和内存使用的技术,通过减少存储访问开销提升性能。它采用Tiling、Recomputation、分块SoftMax等策略,减少HBM访问,加速计算,并在昇腾AI处理器上实现了显著的性能提升。

1 FlashAttention

FlashAttention是一种优化Transformer模型计算效率和内存使用的技术。它通过减少存储访问开销(Memory Access Cost,MAC),而非降低FLOPS(浮点运算次数),来提升性能。

2 前述知识点

涉及到内存访问,肯定与计算的硬件架构有关系。

从GPU架构进行解析,参考如下博客:
大模型推理加速技术的学习路线是什么
首先,我们将探讨GPU架构,特别是其内存层次结构。我们确定了两个重要模式:计算限制(compute bound)和内存限制(memory bound),并讨论了大规模Transformer推理受内存限制的原因。大部分优化都基于Transformer推理受内存限制这一基本事实,例如只要我们提高FLOP利用率,就能提高效率。

2.1 GPU架构

GPU架构总体如下图所示:
image.png

基础部分:DRAM(动态随机存取存储器)、L2缓存和SM(流处理器单元)

  • 与CPU对比
    • SM类似于CPU核心,但具有更高级的并行性;
    • L2缓存和DRAM类似于CPU的L2缓存和DRAM
    • 在Flash Attention论文中,L2缓存被称为SRAM(静态随机存取存储器)
  • A100 80G SXM
    • 08个SM,DRAM容量为80GB,有40M L2缓存

SM内部包含什么?

  • L1缓存:指令和数据
  • 张量核心:进行矩阵乘法运算的地方。回想一下,神经网络计算基本上就是巨大批量的矩阵乘法。

GPU编程基础

在执行model.generate(prompt)时,我们进行以下操作:

  • 内存访问:
    • 从高带宽内存(HBM)加载模型权重 -> L2缓存 -> 传输到SM(流处理器单元)
  • 计算:
    • 在SM中执行矩阵乘法,SM请求张量核心执行计算
  • A100:
    • 108个SM,DRAM容量为80G,40M L2缓存
    • bf16张量核心:每秒312万亿浮点运算(TFLOPS)
    • DRAM内存带宽为2039GB/秒 = 2.039T/秒
  • 如果模型很大,我们将其分割到多个GPU上,比如两个由NVLink连接的GPU
    • NVLink 300GB/秒 = 0.3T/秒
    • 我们大致观察了速度层次结构。尽管不能直接比较,但它们的数量级差异是我们需要优化的主要方面:
    • 312T(SM计算) > 2.03T(DRAM内存访问) > 0.3T=300G(NVLink跨设备通信) > 60G(PCIe跨设备通信)
  • 这意味着,如果我们希望速度更快,我们应该尽力:
    • 充分利用SM
    • 减少单个GPU的内存访问(因为它比计算慢得多),减少GPU之间的通信(因为它甚至比内存访问还要慢)。

计算限制与内存限制

如何确定我们是否充分利用了SM呢?我们通过以下方式检查是否计算或内存限制:

定义每字节GPU操作 = flop / 内存带宽

  • A100 = 312 / 2.039
  • 定义计算强度 = 计算 / 内存访问
  • 如果计算强度大,说明程序更会受到计算限制;如果计算强度较小,则更受内存限制。
  • 增加批次大小会将行为从内存限制变为计算限制。
  • 内核融合:减少了内存访问操作,因为我们将多个操作合并为一个操作。

2.2 Transformer推理

内存布局
image.png

正如我们所看到的,为了在bf16格式下运行一个13B模型,我们大约只有10GB的内存来存储kv缓存。这意味着:

  • 不能使用太大型的批次(尽管我们希望使用更大的批次大小以提高效率)
  • 也不能处理太长的序列,尽管我们确实希望能够处理长度为100k的序列。

3 FlashAttention的策略

FlashAttention的核心策略包括:

  • Tiling(平铺/切分):将注意力矩阵分解成更小的子矩阵,分别计算,确保每个子矩阵的大小适合SRAM(静态随机存取存储器)的存储能力,从而减少对HBM(高带宽内存)的访问。
  • Recomputation(重算):在反向传播时,不存储所有中间状态,而是在需要时重新计算,节省内存。
  • 分块SoftMax:解决标准SoftMax在分块计算中的问题,确保整个Flash Attention的正确性。
  • 优化显存交换:减少SRAM与HBM之间的数据交换,加速计算。
    这些策略共同作用,使FlashAttention在保持计算精度的同时,显著提高计算速度和内存效率

4 Ascend 上的FlashAttention

昇腾异构计算架构CANN针对昇腾AI处理器的片上内存和缓存大小,以及数据搬运通路,基于Ascend C算子编程语言优化实现FlashAttention融合算子,充分利用片上缓存,提升Attention处理性能。根据实测,在一些典型场景中CANN的FlashAttention算子相比小算子取得了5倍以上的性能提升,开发者可直接调用相关算子API接口使能大模型极致性能优化。

可参考:
基于Ascend C的FlashAttention算子性能优化最佳实践-技术干货-昇腾社区

相关实践学习
部署Stable Diffusion玩转AI绘画(GPU云服务器)
本实验通过在ECS上从零开始部署Stable Diffusion来进行AI绘画创作,开启AIGC盲盒。
相关文章
|
2月前
|
机器学习/深度学习 存储 并行计算
Ascend上的PageAttention
PageAttention旨在解决大型语言模型(LLM)服务中的内存管理低效问题,如内存碎片化、利用率低及缺乏灵活的内存共享机制。通过借鉴操作系统中的虚拟内存和分页技术,PageAttention实现了块级别的内存管理和灵活的KV cache共享机制,显著提高内存利用率,降低延迟,提升模型处理速度和性能。相比传统注意力机制,PageAttention通过分段处理序列,有效解决了长序列处理时的计算效率低下和内存过度使用问题。
|
并行计算 C++ 异构计算
Nvidia 并行计算架构 CUDA 分析(一)——CUDA 简介
    CUDA(Compute Unified Device Architecture,统一计算设备架构)是由 NVIDIA 推出的通用并行计算架构,该架构使 GPU 能够解决复杂的计算问题。
5216 0
|
20天前
|
安全 应用服务中间件 API
Ascend+FastAPI+ Uvicorn 实现推理
FastAPI 是一个高性能的 Python Web 框架,专为构建 RESTful API 设计,支持异步编程和数据验证。Uvicorn 是基于 ASGI 的轻量级异步服务器,与 FastAPI 结合使用可提供高效、高并发的服务。在生产环境中,二者配合 Nginx 可实现稳定安全的部署。示例代码展示了如何通过 FastAPI 和 Uvicorn 启动服务,并在 Ascend 平台上运行推理模型,处理自然语言任务。 简介字数:239
|
2月前
|
机器学习/深度学习 自然语言处理 并行计算
SelfAttention在Ascend上的实现
Self-Attention(自注意力)机制是深度学习中用于自然语言处理的关键技术,由 Vaswani 等人在 2017 年提出。它通过让模型关注输入序列的不同部分,增强了对上下文的理解,是 Transformer 架构的核心。Self-Attention 包括 Query、Key 和 Value 三个组件,通过计算这些组件间的相似度,模型可以有效地捕捉长距离依赖关系,提高处理效率和模型表现。此外,多头注意力机制进一步提升了模型的表达能力。
|
2月前
|
人工智能 PyTorch 算法框架/工具
Ascend Extension for PyTorch是个what?
Ascend Extension for PyTorch 是针对华为昇腾处理器的PyTorch框架适配插件,旨在让PyTorch开发者能充分利用昇腾AI处理器的强大计算能力。此扩展通过最小化对原生PyTorch的改动,实现了对昇腾NPU的支持,包括动态图特性、自动微分等功能的完整继承,并提供了与原生PyTorch一致的使用体验。项目详情及源码可在昇腾社区获取。
Ascend Extension for PyTorch是个what?
|
2月前
|
机器学习/深度学习 人工智能 算法框架/工具
什么是CANN和Ascend C
CANN(Compute Architecture for Neural Networks)是华为推出的AI异构计算架构,支持多种AI框架如MindSpore、PyTorch等,适用于AI处理器与编程,旨在提升昇腾AI处理器的计算效率。CANN提供强大的图引擎、算子开发语言Ascend C、算子加速库AOL、集合通信库HCCL、毕昇编译器及Runtime运行时,支持快速构建AI应用,涵盖推理应用开发、模型训练和算子开发等关键功能。
|
2月前
|
分布式计算 并行计算 编译器
Ascend C的编程模型
Ascend C采用SPMD编程模型,实现多核并行计算。通过将数据切分为多个分片,分配给不同的计算核心(或称为block)并行处理,每个核心根据自身的block_idx执行相同代码但处理不同的数据分片,从而提高计算效率。此模型支持高效的数据并行处理,适用于大规模数据计算任务。
|
2月前
|
PyTorch Shell API
Ascend Extension for PyTorch的源码解析
本文介绍了Ascend对PyTorch代码的适配过程,包括源码下载、编译步骤及常见问题,详细解析了torch-npu编译后的文件结构和三种实现昇腾NPU算子调用的方式:通过torch的register方式、定义算子方式和API重定向映射方式。这对于开发者理解和使用Ascend平台上的PyTorch具有重要指导意义。
|
6月前
|
并行计算 Linux TensorFlow
【Deepin 20系统+Tensorflow 2】Linux系统解决Could not load dynamic library ‘libcudart.so.10.0‘
本文描述了在Deepin 20系统中使用TensorFlow 2时遇到GPU未被利用的问题,并给出了相关的调试日志信息。
63 2
|
机器学习/深度学习 人工智能 测试技术
利用 UMA 使硬件加速器可直接用于 TVM
利用 UMA 使硬件加速器可直接用于 TVM
176 0

热门文章

最新文章