GPU架构与计算系统优化加速

1. GPU架构概览

1.1 GPU vs CPU

CPU（中央处理器）和GPU（图形处理器）在设计理念和擅长任务上存在显著差异：

• CPU (Central Processing Unit):
  • 拥有强大的控制单元和高效的上下文切换能力。
  • 设计用于通用计算，能够灵活、高效地处理各种复杂且多样化的任务序列。
  • 控制能力相对更重要，擅长处理串行任务和需要复杂逻辑判断的任务。
  • 特点是低延迟，单个核心性能强大。
• GPU (Graphics Processing Unit):
  • 最初是为了游戏领域的图形渲染而发明，主要处理大量的矩阵运算。
  • 是一种特化的计算处理器，设计目标是进行大规模并行计算。
  • 并行计算能力相对更重要，擅长处理可高度并行化的任务。
  • 特点是高吞吐量，拥有数千个小核心。
  • 这种卓越的并行计算能力使其成为大规模机器学习应用的基础。

1.2 GPU硬件架构

GPU的硬件架构围绕着高效的并行计算设计，其核心是流式多处理器（Streaming Multiprocessor, SM）。

1.2.1 SM (Streaming Multiprocessor)

• SM 是 GPU 的核心计算单元。一个GPU通常包含多个SM。
• 每个SM能够调度多个warp并同时执行。
• 编程模型中的每一个Block必然在一个SM上执行，但一个SM可以执行多个Block（以时间片轮转等方式）。
• A100 GPU拥有108个SM。
• One SM in A100 的结构组成：
  • 包含4个更小的计算单元。
  • 包含128个FP32 / INT32 CUDA核心。
  • 包含64个FP64 (double) CUDA核心。
  • 包含4个Tensor Core。
  • 拥有可配置的Shared Memory，通常为164 KB / SM。
  • 拥有寄存器（Registers），通常为64K 32-bit registers / SM。

1.2.2 CUDA Core

• CUDA Core 是GPU中执行基本算术运算的执行单元。
• 它们处理诸如整数加法、浮点运算等指令。
• 每个SM包含多个CUDA核心，这些核心协同工作来执行线程的指令。

1.2.3 Tensor Core

• Tensor Core 是专门用于加速矩阵乘法和其他张量操作的硬件单元。
• 它们在深度学习任务中表现出色，尤其在处理大型矩阵运算时能提供显著的性能提升。
• 每个A100的SM包含4个Tensor Core。

1.3 GPU计算能力

GPU的计算能力主要通过其执行融合乘加（FMA）操作的速度来衡量。

1.3.1 乘加操作 (FMA)

• 乘加操作（Fused Multiply-Add, FMA） 是一种在一个硬件指令中同时执行一次乘法和一次加法运算的操作。
• 它是现代神经网络中最频繁的操作，因为神经网络的核心计算大量涉及矩阵乘法（乘法）和累加（加法）。
• FMA操作相比分开执行乘法和加法，具有更高的效率和精度。

1.3.2 算力评估指标

• GPU的算力通常使用每秒FMA次数作为评估指标。
• 一次FMA操作被认为包含两个运算（一个乘法和一个加法）。
• A100 GPU 的时钟频率（clock rate）为 1.41 GHz。
• A100的理论峰值算力计算示例：
  • 假设所有SM（108个）满载运算。
  • 这个理论峰值在实际中无法达到，因为存在各种限制（如内存带宽、指令调度等）。
• GPU计算能力的来源:
  • Source: CUDA Programming Guide (说明这些概念和数据来源自NVIDIA官方文档)。

2. GPU编程模型与执行机制

2.1 CUDA编程模型核心概念

NVIDIA CUDA（Compute Unified Device Architecture）是一个并行计算平台和编程模型，它允许开发者使用C、C++、Fortran等语言编写程序，并在NVIDIA GPU上执行。

2.1.1 Kernel

• Kernel 是在GPU上执行的函数。
• 它类似于在CPU上运行的普通函数，但其执行是高度并行的，由数千个GPU线程同时执行。

2.1.2 Grid

• Grid (网格) 是一次Kernel运行的执行范围。
• 它由多个Block（线程块） 组成，用于划分整个计算任务的执行规模。
• Grid是最高层级的抽象，代表了所有参与计算的线程块的集合。

2.1.3 Block

• Block (线程块) 是线程的组织单位，也是GPU中并行计算的基本调度单元。
• 每个线程块包含一组线程（例如，最多1024个线程）。
• 一个Block内部的线程可以在Shared Memory中相互通信，并可以通过同步点（__syncthreads()）进行同步。
• 一个Block内部的线程保证在同一时间、同一SM上执行（但并非所有线程同时执行）。

2.1.4 Thread

• Thread (线程) 是GPU执行的最基本计算单元。
• 每个线程执行Kernel代码的一个独立副本，但通常操作不同的数据。

2.1.5 Shared Memory

• Shared Memory 是一种位于片上、Block内线程之间共享的内存。
• 它的访问速度非常快，类似于CPU的L1缓存，是实现Block内线程高效协作的关键。

2.1.6 Global Memory

• Global Memory 是所有线程（无论属于哪个Block）都可以访问的内存空间。
• 它位于GPU的主存储器（如HBM）上，容量最大，但访问延迟相对较高。

2.2 从编程模型到硬件执行：SIMT

GPU的执行模型被称为SIMT，它将编程模型中的抽象概念（Grid、Block、Thread）映射到底层的硬件单元（SM、Warp）。

2.2.1 Warp

• Warp 是GPU的基本执行单位，它由一组并行执行的线程组成。
• NVIDIA GPU中一个warp通常包含32个线程。
• 这些线程会在同一时间被调度，并执行相同的指令（SIMT模式），但操作不同的数据。

2.2.2 SIMT (Single Instruction, Multiple Threads)

• SIMT（Single Instruction, Multiple Threads）类似于SIMD（Single Instruction, Multiple Data）。
• SIMD 要求所有数据通道严格执行相同的指令。
• SIMT 则更加灵活，它允许不同线程根据条件分支执行不同的路径（而SIMD必须通过掩码实现）。这意味着如果warp内的线程遇到分支，它们会串行执行分支的两部分，然后重新汇合。
• SIMT是GPU实现大规模并行计算的核心机制。

2.2.3 线程块与SM的映射

• 编程模型中的线程块（Block） 会被映射到GPU的SM上。
• 每个SM可以同时处理多个线程块（通过时间片轮转或多路复用）。
• 这些线程块会共享SM内的共享内存、L1缓存等资源。

2.2.4 Warp调度

• SM内部有多个warp调度器，它们负责管理warp的执行。
• 调度器在多个warp之间进行切换，以隐藏内存访问延迟和其他计算瓶颈，从而提高SM的利用率。
• A100 GPU的每个SM支持4个并发warp，并可以有64个在运行的warp。
• 它可以在一个时钟周期内进行快速、轻量级的上下文切换。

2.2.5 启动示例：1000维向量相加

假设我们要实现一个1000维向量的相加操作：

int N = 1000;
int blockSize = 256; // 每个Block包含256个线程
int gridSize = (N + blockSize - 1) / blockSize; // 计算需要的Block数量
// Kernel启动语法：<<<gridDim, blockDim>>>(args...)
addKernel<<<gridSize, blockSize>>>(a, b, c, N);

在这个例子中，addKernel 将在GPU上执行，gridSize 决定了有多少个Block被创建，blockSize 决定了每个Block有多少个线程。

2.3 矩阵乘法Kernel示例

• 矩阵乘法是GPU上的一个典型并行计算任务。
• 它涉及大量的乘加操作，非常适合GPU的并行架构。
• 一个优化的矩阵乘法Kernel会充分利用GPU的内存层级结构（寄存器、共享内存、全局内存），并通过分块（tiling）等技术来提高数据复用率和带宽利用率。
• 示例链接: https://dlsyscourse.org/

3. GPU内存层级架构与访问模式

3.1 内存层级结构

GPU的内存系统是一个复杂的层级结构，旨在平衡容量、带宽和延迟。物理上越靠近GPU计算核心的内存，其空间通常越小，但带宽越大，访问延迟也越低。

• 示意图描述: 原始资料中提供了一张展示GPU内存层级结构的示意图，从最靠近计算核心（上部）到最远的（底部）排列，并标注了容量、带宽（相对值）和延迟（相对值）。

  该图从上到下展示了以下内存层级：

  • Register File: 位于最顶端，最接近计算核心。
  • L1 Cache & Shared Memory: 紧随其后。
  • L2 Cache: 位于L1/Shared Memory之下。
  • High Bandwidth Memory (HBM): GPU的主存储器，容量最大，但距离计算核心相对较远。
  • Host Memory (通过PCIE): 位于最底部，是CPU的主内存，通过PCIe总线与GPU连接，访问延迟最高，带宽最低。

  图中标注了各层级内存的相对带宽和容量：

  • 256kB Register File per SM (27MB total)：每个SM 256KB，总共27MB。
  • 192kB L1 Cache & Shared Memory per SM (20MB total)：每个SM 192KB，总共20MB。
  • 40MB L2 cache shared across all SMs：所有SM共享40MB。
  • 80GB High Bandwidth Memory：80GB，具有1,935GB/s的带宽（标注为1x）。
  • Host memory through PCIE：带宽标注为0.02x，显著低于HBM。

  带宽的相对值表示，从Host Memory到HBM有约50倍的提升 (1/0.02)。从HBM到L2 Cache，再到L1/Shared Memory，乃至Register File，带宽呈指数级增长，同时延迟急剧降低。【若需查看原始图片详情，请参考原文中的“GPU内存层级架构图”】

3.1.1 Register File

• 特性: 位于GPU核心内部，速度最快，容量最小。
• 作用: 用于存储线程的局部变量，每个线程拥有独立的寄存器。访问延迟极低，通常在一个时钟周期内完成。

3.1.2 L1 Cache & Shared Memory

• 特性: 位于每个SM内部，速度非常快，容量适中。
• 作用:
  • L1 Cache: 作为SM的局部缓存，用于加速对Global Memory的访问。
  • Shared Memory: 供同一个Block内的线程共享，实现快速的线程间通信和数据复用。其访问延迟远低于Global Memory。

3.1.3 L2 Cache

• 特性: 通常由所有SM共享，容量比L1/Shared Memory大，速度介于L1/Shared Memory和HBM之间。
• 作用: 作为GPU的全局缓存，用于缓存Global Memory的数据，进一步减少对HBM的访问。

3.1.4 HBM (High Bandwidth Memory)

• 特性: GPU的主存储器，容量最大，带宽极高，但访问延迟相对较高。
• 作用: 存储大量的计算数据和模型参数，是GPU与CPU之间数据交换的主要介质。

3.1.5 Host Memory (通过PCIE)

• 特性: CPU的主内存，通过PCIe总线与GPU连接。
• 作用: GPU与CPU进行数据传输的通道。其带宽最低，延迟最高，是GPU计算中常见的性能瓶颈。

3.2 内存访问特性

3.2.1 带宽与延迟关系

• 物理限制: 物理上，光速上限决定了信息传输的最小延迟。因此，内存距离计算核心越远，访问延迟自然越高。
• 设计权衡: 内存层级的设计是容量、带宽和延迟之间的权衡。大容量、低成本的内存（如HBM）通常具有较高延迟，而小容量、高成本的内存（如寄存器、L1/Shared Memory）则具有极低延迟和极高带宽。

3.2.2 HBM访问模式与连续性

高效利用HBM（和其他内存）的关键在于优化访问模式。

• HBM访问模式图示描述: 原始资料中提供了一张关于HBM访问模式的示意图，展示了从"burst"（突发）到"one read per page"（每页一次读取）的效率变化。

  该图可能演示了随着内存访问模式的变化，数据传输效率的差异：

  • Burst模式: 当访问是连续的、对齐的，内存控制器能够一次性读取一大块数据（burst），效率最高。图中"burst部分占大头"可能指的是这种理想情况。
  • One read per burst with negligible page reads: 访问模式仍然较好，但可能不是完全对齐，或者存在一些小的间隙，但页内读取效率依然很高。
  • One read per page: 当访问变得稀疏，每次读取都需要访问一个新的内存页时，开销会显著增加，因为每次页访问都会引入额外的延迟。图中"足够稀疏之后page read开销占大头，变成one read per page"描述了效率最低的情况。【若需查看原始图片详情，请参考原文中的“HBM访问模式图”】

• 连续访问的重要性: 为了最大化HBM的利用率和降低访问延迟，应尽可能地进行连续访问，从而避免频繁的页交换（换页）和不必要的开销。

• 大量并发的warp通过合并访存降低访问开销:

  • CUDA硬件具备一种称为合并访存（coalesced access） 的机制。当一个warp中的32个线程同时访问Global Memory时，如果它们的内存地址是连续的，硬件能够将这些独立的访问请求合并成一个或少数几个大的内存事务，从而显著提高访存效率。
  • Example:
    • 一个线程获取float2数据类型（包含两个浮点数）的向量元素，需要 $2 \times 4 = 8$ 字节。
    • 一个warp（32个线程）如果连续访问，理论上会读取 $8 \times 32 = 256$ 字节。
    • 一个SM中有4个同时运行的warp，如果它们的数据访问也是连续的，可能一次性读取 $256 \times 4 = 1024$ 字节。这个1024字节的量，通常与一个内存行（cache line或内存事务的基本单位）的字节数相匹配，从而实现了高效的内存合并访问。
  • 结论: 我们需要确保跨线程束/线程的内存访问是连续的，这样CUDA才能更好地调度线程束，将多个线程的访存请求合并成高效的事务。

• 线程独立性:

  • 为了让并行度带来的收益最大化，一个重要的设计原则是：线程之间尽可能少同步，尽可能独立。
  • 这意味着每个线程应尽可能地独立完成自己的计算任务，减少对其他线程的依赖和同步点。
  • 通过将调度工作交给GPU（warp调度器），可以更好地隐藏延迟，提升并发效率。即“你忙你的，我忙我的”，让GPU调度器在不同warp之间快速切换，充分利用计算资源。

4. GPU计算系统优化思路

4.1 核心优化目标

GPU计算系统的优化是一个多维度的挑战，主要围绕以下四个核心目标展开：

1. 延迟 (Latency): 指完成一个操作所需的时间。在GPU中，特别是内存访问延迟，是影响性能的关键因素。
2. 内存 (Memory): 指内存的容量、访问速度和访问模式。高效的内存使用能减少瓶颈。
3. 带宽 (Bandwidth): 指单位时间内数据传输的量。高带宽是GPU处理大规模数据的基础。
4. 算力 (Compute): 指GPU执行计算操作的原始能力（如FLOPS）。

• 重要提示: 在GPU优化中，一个普遍且至关重要的认知是：算力在上述问题中是最最不重要的部分！首先要充分利用算力——提高GPU利用率！

  这意味着，即使GPU拥有强大的原始计算能力，如果数据不能及时、有效地到达计算核心，或者计算任务不够并行化以充分利用所有核心，那么这些算力就无法被利用。因此，优化的重点往往是解决数据流问题（延迟、内存、带宽），以确保算力能够得到充分利用。

4.2 GPU利用率目标

• GPU利用率 是衡量GPU计算资源被实际利用程度的指标。它表示GPU的计算单元在给定时间内处于繁忙状态的比例。
• 对于复杂的、内存密集型的通用GPU计算任务来说，达到高利用率是极具挑战性的。
• 实际表现:
  • GPU利用率 ~50% = 表现不错！ 这表明即使是一半的利用率，对于许多工作负载来说也已经是相当好的结果。
  • GPU利用率 ~75% = 非常厉害！ 达到75%或更高的利用率通常意味着代码已经经过了高度优化，并且能够高效地利用GPU的并行处理能力和内存子系统。

5. 具体优化目标与策略

5.1 优化目标：隐藏延迟

延迟 (Latency) 是指完成一个操作所需的时间。在GPU计算中，数据从全局内存（HBM）传输到计算核心的延迟是主要的瓶颈之一。为了克服这一瓶颈，GPU的设计和编程模型都强调通过大量的并发任务来隐藏延迟。

5.1.1 大量并发隐藏延迟

• 基本原理: 当一个计算单元（如一个warp或SM）因为等待数据从慢速内存（如HBM）加载而停顿时，GPU调度器可以立即切换到另一个准备好执行的计算单元，从而避免处理器空闲。通过并行执行足够的任务，GPU可以确保总是有可执行的任务，从而“隐藏”掉数据加载的延迟。
• 例子: 文档中提到“这就是为什么执行4个却有64个wrap同时工作的原因”，这指的是在一个SM中，虽然可能只有少数（如4个）warp在同时进行计算，但调度器会管理更多（如64个）处于不同执行阶段的warp。当一个warp因内存访问而阻塞时，调度器可以快速切换到另一个就绪的warp，保持计算单元的忙碌。

5.1.2 线程块与SM的映射策略

• 充分利用SM: “使用单个线程块启动函数只能将工作分配给一个SM（流处理器组）；为了充分利用具有多个SM的GPU，需要启动多个线程块。” 这强调了为了最大化GPU的计算能力，必须将任务分解为足够多的线程块（blocks），以便这些线程块能够并行地调度到GPU的多个SM上执行。
• 隐藏延迟: 即使在单个SM内部，启动多个线程块也很有益。当一个线程块的warp因为内存访问而阻塞时，SM可以切换到处理另一个线程块的warp，从而隐藏延迟。

5.1.3 Warp调度机制

• 快速上下文切换: GPU的warp调度器能够在单个时钟周期内从一个线程束切换到下一个。这种超快速的上下文切换是隐藏内存访问延迟的关键。当一个warp发起内存请求并进入等待状态时，调度器会立即选择另一个已就绪的warp来执行，从而避免SM的空闲。
• 并发量: A100 GPU的每个SM支持4个并发warp（即可以在任意时刻实际执行指令的warp数量），但可以同时管理64个在运行中的warp（处于不同状态，包括等待内存、计算等）。这种高并发度确保了在任何给定时间都有足够的warp可以被调度执行。

5.1.4 并行访存与合并访问

• HBM访问模式: 文本指出HBM访问模式应“尽可能的连续访问，从而避免换页！” 这意味着当多个线程访问内存时，如果它们的访问地址是连续的，GPU的内存控制器可以将其合并为一个大的突发访问（burst access），从而更高效地利用内存带宽。
  • 图片描述 (HBM访问模式图示): 原始资料中HBM访问模式的图示可能展示了随着访问连续性（或稀疏性）的变化，内存访问开销的组成。在非常连续的情况下，突发传输（burst）占主导；在适度稀疏时，每次突发可能伴随少量页访问；而在高度稀疏时，页访问开销占大头，可能导致每次读取都需要一个独立的页访问。
• 合并访存: “大量并发的wrap通过合并访存降低访问开销！”
  • 示例: 一个线程获取 float2 数据类型（8 bytes）。一个warp（32个线程）如果访问连续的内存，可以一次性获取 8 bytes/thread * 32 threads = 256 bytes。一个SM中有4个运行的warp，如果它们访问连续的内存，可以一次性获取 1024 bytes。如果这正好是一行的字节数，那么这将是一个非常高效的访问模式。
  • 重要性: 确保跨线程束/线程的内存访问是连续的，这样CUDA才能更好地调度线程束，实现高效的内存合并（memory coalescing）。内存合并能够显著减少内存事务的数量，从而提高访存效率。

5.1.5 线程独立性

• 最大限度利用并行度: “让并行度带来的收益最大化的另一个重要设计：threads之间尽可能少同步，尽可能独立！”
• 原理: 线程之间频繁的同步（例如，__syncthreads()）会引入额外的开销，并可能导致一些线程等待其他线程完成，从而降低并行度。理想情况下，每个线程应独立完成其任务，只在必要时进行同步或通信。将调度交给GPU，让其自由选择就绪的线程执行，可以最大化吞吐量。
• 示例: 基于CNN的物体识别通常涉及大量独立的卷积操作，每个卷积可以在不同的GPU线程上并行执行，且彼此之间依赖性较低，是并行计算的良好应用场景。

5.2 优化目标：提升带宽利用率

即使算力强大，如果数据传输速度跟不上，GPU的计算单元也会因为等待数据而空闲。这就是带宽瓶颈。

5.2.1 带宽与计算的Trade-off

• 定义:
  • $T_{mem}$：用于访问内存的时间。
  • $T_{math}$：用于执行数学运算的时间。
• 重叠与总时间: 如果不同线程的内存访问和数学运算可以重叠（即通过流水线（Pipelining） 技术），那么函数的总执行时间是 $max(T_{mem}, T_{math})$。
• 性能限制因素:
  • 如果 $T_{math} > T_{mem}$，则该函数是数学（计算）受限的。这意味着计算时间是瓶颈，提高计算能力会有更大收益。
  • 如果 $T_{mem} > T_{math}$，则该函数是内存（带宽）受限的。这意味着数据传输时间是瓶颈，提高内存带宽会有更大收益。

5.2.2 算法算术强度与设备计算密集度

为了量化这种Trade-off，我们引入算术强度（Arithmetic Intensity） 的概念。

• 算术强度 (Arithmetic Intensity):

  • 算法相关的性质，定义为 $\#ops / \#bytes$。
  • 表示每传输一个字节的数据，执行了多少次浮点运算（operations）。

• 设备计算密集度 (Device Compute Intensity):

  • 设备相关的性质，定义为 $BW_{math} / BW_{mem}$。
  • $BW_{math}$ 是设备的峰值浮点运算能力（FLOP/s）。
  • $BW_{mem}$ 是设备的峰值内存带宽（Byte/s）。
  • 代表了充分利用设备算力所需的最低算术强度。例如，A100 GPU的设备计算密集度约为100，意味着每传输1字节数据，需要执行100次浮点运算才能完全利用其算力。

• 关系:

  • 当任务计算密集度（即算法的算术强度） > 设备计算密集度时：任务是计算受限的。充分利用了设备算力，瓶颈在于计算单元的速度。
  • 当任务计算密集度 < 设备计算密集度时：任务是带宽受限的。未充分利用设备算力，瓶颈在于内存传输速度。

• 推导: 原始不等式：$T_{math} > T_{mem}$ 等价于：$\frac{\#ops}{BW_{math}} > \frac{\#bytes}{BW_{mem}}$ 通过代数运算，可以重构为：$\frac{\#ops}{\#bytes} > \frac{BW_{math}}{BW_{mem}}$ 左侧是算法的算术强度，右侧是设备的计算密集度。

5.2.2.3 Roofline Plot模型

• 图片描述 (Roofline Plot): 原始资料中的Roofline Plot是一个二维图表，横轴通常是算术强度 (Arithmetic Intensity)，纵轴是性能 (Performance) (通常以FLOP/s为单位)。
  • 两条线:
    • 一条是水平线 (Peak Performance)，代表设备的峰值计算能力（$BW_{math}$），这是算力上限。
    • 一条是斜线 (Memory Bandwidth Limit)，代表由内存带宽决定的性能上限，其斜率由设备的内存带宽（$BW_{mem}$）决定。
  • 交点: 两条线的交点即为设备计算密集度。
  • 工作区域:
    • 当算术强度小于交点时，性能受限于内存带宽（处于斜线部分）。
    • 当算术强度大于交点时，性能受限于计算能力（处于水平线部分）。
  • 目标: 理想情况下，我们希望算法的性能位于Roofline的“屋顶”部分，即达到计算或内存的理论峰值。

5.2.3 带宽瓶颈的解决方案

问题: 如果你目前任务的处于带宽瓶颈，怎么办？

• A. 换用更大的显存：通常不能解决带宽瓶颈，因为更大的容量不一定意味着更快的传输速度。
• B. 换用更快的显存：正确。更快的显存（如升级到更高带宽的HBM）直接增加了 $BW_{mem}$，从而减少了设备计算密集度（$BW_{math} / BW_{mem}$），使得任务更容易达到计算受限区域。
• C. 换用更快的计算芯片：如果已经是带宽瓶颈，意味着计算单元已经在等待数据，此时再提高计算能力（增加 $BW_{math}$）只会让设备计算密集度更高，并不能解决瓶颈。
• D. 使用更低的精度进行计算：正确。例如，从FP32降到FP16或BF16。这会增加任务计算密集度。因为同样的数据量（例如，一个16位浮点数只占2字节），可以进行更多的运算（如16位乘加），从而提高了 #ops / #bytes 的比值，使得任务更容易达到计算受限区域。

总结: 解决带宽瓶颈的根本方法是提高算法的算术强度（例如，通过降低精度、重用数据）或提高内存带宽。

• 挑战: “对于FLOPs非常大的先进GPU，Arithmetic Intensity想达到其设备计算密集度（100 for A100）还是挺难的！”这意味着很多实际应用在A100上仍然是带宽受限的。
• 解决方案: “只有足够大的矩阵运算才能够满足”。大型矩阵运算通常具有很高的算术强度。
• 批处理 (Batch Processing): “因此批处理训练可以带来吞吐率提升。” 批处理通过一次处理更多的数据，可以有效地增加操作数与字节传输的比率，从而提高任务的算术强度，使其更接近计算受限。
  • 图片描述 (Batch size vs. Throughput): 原始资料中可能包含一张图，展示了随着 (batch_size * dim_size) * (dim_size * 1024) 在GeForce RTX 3090上执行时，吞吐量如何随批处理大小增加而提升的例子。这图会说明，在较小的批处理大小下，性能受限于内存带宽，吞吐量较低；而随着批处理大小增加，算术强度提高，吞吐量也随之提升，直到达到计算能力的瓶颈。

5.3 优化目标：高效内存使用

高效利用内存是优化GPU性能的关键。

• 与设备计算密集度关联: “数据传输越高效，任务越容易充分利用算力。” 通过减少实际传输的数据量，或者提高数据的复用率，可以有效地提高算法的算术强度，使其更接近或达到设备计算密集度，从而摆脱带宽瓶颈。

5.3.1 多级内存访问与复用

• 策略: “增加片上内存/缓存的复用度，从而减少了实际的设备计算密集度！”
  • GPU具有多级内存层级结构（寄存器、L1缓存、共享内存、L2缓存、HBM）。
  • 片上内存（如寄存器、共享内存、L1/L2缓存）速度远快于片外内存（HBM）。
  • 通过精心设计算法，使数据尽可能在片上内存中多次使用，可以大幅减少对慢速HBM的访问次数，从而降低平均内存访问延迟和带宽需求。

5.3.2 矩阵乘法的分块优化示例

矩阵乘法（C = A * B）是展示内存优化的经典案例。

5.3.2.1 寄存器级别优化

• 原始矩阵乘法问题: 在最简单的实现中，对于结果矩阵 $C$ 中的每一个元素 $C_{ij}$，都需要重新读取 $A$ 矩阵的第 $i$ 行和 $B$ 矩阵的第 $j$ 列。
• 开销: 如果每次都重新读取，总的内存访问开销非常高，达到 $2N^3$（对于 $N \times N$ 的矩阵）。
• 寄存器分块 (Tiling): “当V = 1时候，该运算相当于原始矩阵乘法，对每一个C中的单元重新读取对应的行和列，因此开销为2*N^3.” (此处的 V=1 可能指的是没有在寄存器中进行分块复用。) 寄存器级别的优化就是将一小块数据（如矩阵 $A$ 的一小行和矩阵 $B$ 的一小列）加载到寄存器中，然后在寄存器中完成多次乘加操作，而不是每次都从全局内存中读取。

5.3.2.2 寄存器-共享内存级别优化

• 层级分块: 这是一个更高级的优化策略，结合了共享内存和寄存器。
• 思想:
  1. 将大矩阵 $A$ 和 $B$ 分割成若干个子块。
  2. 每个线程块负责计算 $C$ 矩阵的一个子块。
  3. 为了计算 $C$ 的一个子块，线程块会分批次地将 $A$ 和 $B$ 对应的子块从全局内存加载到共享内存 (Shared Memory)。
  4. 一旦数据进入共享内存，该线程块内的所有线程都可以高速访问这些数据。
  5. 每个线程再从共享内存中加载更小的数据块到自己的寄存器 (Registers) 中进行计算。
• 优势: 共享内存的访问速度远快于全局内存，且可以在一个线程块内高效复用。通过这种分层的数据移动和复用，可以显著减少对全局内存的访问次数，极大地提高内存访问效率，从而提升整体性能。
• 示例: “向量相乘生成累加矩阵 L>V” (这可能是一个具体的矩阵乘法分块算法的命名或描述，例如 L 表示加载到共享内存的块大小，V 表示加载到寄存器的块大小，且 L 大于 V)。

6. FlashAttention案例分析

FlashAttention 是针对 Transformer 模型中的注意力（Attention）层在GPU上进行高效计算而提出的一系列优化方法，旨在解决标准Attention计算中由于内存访问模式不佳导致的性能瓶颈。

6.1 Transformer Attention层问题

在 Transformer 模型中，Attention 机制的核心计算涉及 Q（Query）、K（Key）、V（Value）矩阵。对于一个上下文长度（context length）为 $N$ 的序列，Attention 的计算公式通常包含 $N \times N$ 大小的中间矩阵 $S$ 和 $P$： $\text{Attention}(Q, K, V) = \text{softmax}(\frac{QK^T}{\sqrt{d_k}})V$

其中，$S = QK^T$ 是一个 $N \times N$ 的分数矩阵，$P = \text{softmax}(\frac{QK^T}{\sqrt{d_k}})$ 是注意力权重矩阵。

随着 $N$（上下文长度）的增大，主要问题如下：

• 内存开销巨大：中间结果 $S$ 和 $P$ 矩阵的大小是 $N \times N$，其存储和传输开销会随着 $N$ 的平方增长，迅速消耗 GPU 的高带宽内存（HBM）。
• HBM访存频率过高：标准Attention计算需要频繁地在HBM上读写 $N \times N$ 大小的中间矩阵，导致 HBM 成为性能瓶颈。
• 任务计算密集度低：由于大量的内存访问，导致实际的计算量相对于数据传输量较少，即任务计算密集度很低。
• GPU利用率不高：在内存带宽限制下，GPU的算力无法得到充分利用，导致GPU利用率低下。

6.2 FlashAttention核心思想

FlashAttention 的核心在于避免显式地构建和存储完整的 $N \times N$ 大小的注意力矩阵 $S$ 和 $P$，从而减少 HBM 的读写量。

6.2.1 分块思想 (Tiling)

• 核心策略：通过将注意力计算分解成更小的块（tile），并在 GPU 的片上内存（如 SRAM/Shared Memory）中进行计算。
• 跳过中间结果：在每个块的计算过程中，直接从输入的 Q、K、V 矩阵生成最终的输出 $O$ 矩阵，而不显式地将 $S$ 和 $P$ 写入 HBM。这避免了 $N \times N$ 矩阵的存储和传输开销。
• 目的：减少对 HBM 的访问，提高数据局部性，充分利用片上内存的高带宽和低延迟特性。

6.2.2 重计算 (Recomputation)

• 策略：FlashAttention 在前向传播时不保存中间结果 $S$ 和 $P$。
• 反向传播处理：在反向传播（BP）时，如果需要 $S$ 和 $P$ 来计算梯度，则选择重新计算这些中间结果，而不是从 HBM 中读取它们。
• 权衡：这种方法用额外的计算成本（重新计算 $S$ 和 $P$）来换取显著的内存带宽节省，因为重新计算通常比从 HBM 读写更划算，尤其是在带宽受限的情况下。

6.3 FlashAttention具体实现

6.3.1 Partial Updates

FlashAttention通过分块（Tiling）的方式，将 Q、K、V 沿 token 维度切分，并在一个 GPU 块（block）内完成局部的注意力计算。

• 图像描述（结合PPT内容）：
  • 想象一个大的 $Q$ 矩阵，它被水平切分成若干行，每一行代表一个 token 的查询向量。
  • 同样，$K$ 矩阵被水平切分，每个行代表一个 token 的键向量。
  • 一个 GPU block 负责处理 $Q$ 的一个“块”与 $K$ 的所有“块”之间的注意力计算。
  • 例如，一个 block 可能负责计算 $Q_i$ （Q矩阵的第i个分块）对所有 $K_j$ （K矩阵的所有分块）的注意力，并逐步累积输出 $O_i$。
  • 在计算过程中，每一个小分块的 $Q$ 和 $K$ 会被加载到片上内存（如 Shared Memory），计算出局部的注意力分数和输出，然后将这些局部结果合并。
  • token 维度切分：这意味着我们将序列长度 $N$ 沿着某个维度（通常是行维度）进行切分，分配给不同的 GPU block 来并行处理。
• 目标：通过在片上内存中完成大部分计算，减少对慢速 HBM 的依赖，从而实现更高的吞吐量。

6.3.2 FlashAttention on GPT2

FlashAttention在实际模型如GPT2上的应用表明，它可以显著提升训练速度和降低内存消耗，尤其是在处理长序列时。

6.4 FlashAttention的挑战

6.4.1 Block间同步问题

• 在FlashAttention v1的设计中，一个 GPU block 可能需要读取 Q 的一个部分，然后遍历所有的 K 和 V 分块，逐一计算并更新输出。
• 这导致了Load/write synchronization issue（加载/写入同步问题）：不同的 Block 可能尝试写入输出矩阵的同一部分或需要协调对共享数据的访问，这会引入复杂的同步机制和潜在的竞争条件，从而影响效率。
• 图像描述（结合PPT内容）：
  • ...... One block：PPT图示可能显示一个Block在处理 Q_i 时，需要与所有 K_j 交互。
  • 这意味着一个 Block 的输出并非完全独立，它可能需要等待其他 Block 完成其对应的 K/V 分块的处理，或者在写入输出时与其他 Block 协调。
  • 这种设计使得Block之间不能尽可能相互独立地写入，降低了并行度。

6.5 FlashAttention v2的改进

FlashAttention v2 主要解决了 FlashAttention v1 中的 Block 间同步和数据依赖问题，进一步提高了并行效率。

6.5.1 解决方案：循环结构调整

• 核心改进：FlashAttention v2 的主要优化在于交换了内外循环的顺序（“Just swap the outer/inner loop!"）。
• FlashAttention (v1)：每个 block 计算所有 tokens “询问部分上下文tokens” 的工作。这意味着所有 blocks 完成才能得到完整的 scores 以及输出 O 的任何子值。
• FlashAttention v2：每个 block 计算部分 tokens “询问所有上下文tokens” 的工作。
• 目的：通过这种调整，每个 GPU block 负责计算输出矩阵 $O$ 的不同部分（One block for one O_i），使得Block之间的数据写入尽可能地相互独立。
• 效果：减少了 Block 间的同步需求和潜在的写冲突，进一步提升了并行效率和 GPU 利用率。

6.5.2 FlashAttention v2与FlashAttention的对比

• FlashAttention v1：一个 block 可能负责处理 Q 的某一部分，然后它需要遍历 K 的所有分块。要得到最终的输出 $O$，需要将所有 block 的计算结果合并。这种模式下，所有 blocks 完成才能得到完整的 scores 以及输出 $O$ 的任何子值。
• FlashAttention v2：一个 block 负责计算输出 $O$ 的某一部分 O_i。为了得到 O_i，这个 block 会读取 Q 的对应部分 Q_i，并遍历 K 和 V 的所有分块。因此，每个 block 处理输出矩阵 $O$ 的不同部分。这种设计使各个 block 的输出计算更加独立。

总结：FlashAttention v2 的优化方向是让每个 block 独立地计算输出矩阵的特定部分，从而最小化 block 之间的依赖和同步开销。

• 参考原始论文获取更多细节：
  • FlashAttention: https://arxiv.org/abs/2205.14135
  • FlashAttention v2: https://arxiv.org/abs/2307.08691