斯坦福大学 CS336 Lecture 05 GPU Principles and Distributed Training Fundamentals
1. Goals
①. Having a thorough understanding of the working principle and operation mechanism of GPU
②. Knowing how to make fast algorithms like FlashAttention
2. GPU Architecture
2.1 Why Parallelism
Dennard Scaling:随着晶体管变得越来越小,它们的功率密度保持不变,因此功率的使用与面积成比例
Moore's Law:集成电路芯片上所集成的电路的数目,每隔18个月就翻一番
以蓝线所示,尽管芯片的晶体管密度确实仍在持续增长,单线程性能的增长曲线已明显趋于平缓。这意味着我们无法依赖于计算速度的提升,为了获得更高的效率,必须引入并行化。
2.2 Difference between GPU and CPU
CPU:由于线程数量有限,单线程运行速度极高。针对低延迟进行优化
GPU:由数量极大的计算单元构成,注重并行运算。针对高数据吞吐进行优化。

例如 T1 到 T4 任务序列,即使 GPU 单个任务延迟可能更高,但仍先完成序列。
2.3 Layout of GPU

基本架构就是 GPU - SM - SP/CUDA Core 以及现代 GPU 配置的 Tensor Core 以及 RT Core (处于 SM 中)
SM(Streaming Multiprocessor):基本运算单元。(Triton 等工具的作用层级)
CUDA Core / SP (Streaming Processor):接受相同指令, 将其应用于不同的数据
Physical Layout MATTERS:当数据吞吐量如此之大时,内存与计算单元之间的物理距离开始产生重大影响。
内存单元距离 SM 越近,其访问速度就越快。Register、L1 Cache 以及 shared memory 就位于 SM 内部。而 L2 Cache 在 SM 外部,仍在 GPU 芯片上。 VRAM 在芯片旁,故数据读取速度逐渐降低。

2.4 Execution Model of GPU
为了实际编写针对 GPU 的高性能代码,需要深入理解 GPU 的实际执行模型工作机制,从线程块(blocks)、线程束(warps)以及线程(threads)三个层级进行深入考虑,其粒度依次细化

Thread:CUDA 中最小的执行单元,每个 Thread 独享一组 Register。SM 里有一大块“寄存器文件(Register File)”,里面包含成千上万个物理寄存器。
Warp:特定个 Thread 组成一个 Wrap(NVIDIA 为32,AMD 为64)。一个 Warp 里的所有 Thread 同时执行同一条指令,但处理不同的数据(SIMT,Single Instruction Multiple Threads)。GPU 硬件层面的调度单元。SM 里的调度器不是一个一个 Thread 去调度的,而是以 Warp 为单位去调度。
Block:由程序员定义的一组 Thread,数量可以自定。同一个 Block 里的所有 Thread,被分配到同一个 SM上执行。一个 Block 上的所有 Thread 共享 shared memory。L1 Cache 是 SM 级别的资源,它被该 SM 上同时执行的所有 Block(可能不止一个)共享。
一个 Block 里的 Thread 数量,会被硬件自动计算成 Warp 的整数倍。比如 33 个 Thread,硬件会分配 2 个 Warp(64 个 Thread),第二个 Warp 里只有 1 个 Thread 在干活,其余 31 个闲置。所以,Block 里的 Thread 数量最好是 32 的倍数。
2.5 About TPUs
based on Tensor Core
2.6 Strengths of the GPU Model
①. GPU 具有极好的延展性:为了获得更高的速率,不需要通过提高时钟效率,只需增加更多 SM 单元即可。
②. SIMT 架构适合处理矩阵运算等基础操作场景
③. 线程都极其轻量化,可以随时被暂停或重启

蓝线:GPU 与主机之间的数据传输带宽
绿线:全局内存的访问速度
灰线:计算速度
当前性能瓶颈已转向内存带宽,因为内存性能的提升速度远跟不上计算单元的演进,这种发展趋势在未来仍将持续
3. How to Make GPUs Go Faster

横坐标是参与乘法运算的矩阵的维度大小,纵轴则代表每秒执行的计算操作次数,可以理解为硬件利用率。
可以观察到,随着矩阵维度变大,硬件利用率不断增大。(类似于渲染 Draw Call,抵销任务调度等操作带来的额外开销)
但横纵坐标变化曲线表现出奇怪的波形,需要探究其原因。

上图的波形类似于屋顶线模型(Roofline Model):横轴为运算强度,衡量的是“每读入一个字节的数据,能进行多少次计算”。运算强度越高,说明计算越密集。纵轴为计算性能。
当我们考察性能时,会发现存在两种典型状态,第一种是内存瓶颈状态,对应上图中左侧曲线部分,大部分时间在等待数据从显存搬运到SM;第二种是计算瓶颈状态,对应图中右侧曲线部分,计算非常密集,已经把GPU的计算单元跑满了,接近硬件的理论峰值算力。
优化的最终目标:处于右侧区域,充分释放计算单元的效率
3.1 Control Divergence (not a memory issue)

由于 GPU 采用 SIMT 架构,所以一个 warp 中的所有 thread 会执行同样的命令,只是处理不同的数据。一个 warp 内部的条件语句会具有很强的破坏性,会强制暂停所有未执行主控制流的线程,破坏了 thread 的并行价值。所以,在大规模并行计算单元中应该尽量避免使用条件判断语句
3.2 Low Precision Computation
对一个 n 维向量做一次逐元素的 ReLU : x = max(0, x) 操作。若数据采用 fp32,则每进行一次 flop 运算(0 与 x 的比较)需要进行读写操作各一次(读入 x ,输出结果 ReLU(x) ),共两次 fp32,为 8 bytes。计算强度为 8 bytes / FLOP
如果采用 fp16,则计算强度为 4 bytes / FLOP。在保持算法和硬件不变的情况下,相当于免费获得了双倍的内存带宽。

上图是在混合精度训练过程中,哪些操作适合用什么精度数据。
3.3 Operation Fusion

在Lecture 04 4.3.2 DeepSeek v3 引入的 MLA 中,优化 kv-cache 的过程中就提到了 Operation Fusion。
3.4 Recomputation


(我去,简单的思路,极致的提升)
用富裕的计算资源换取紧缺的内存带宽其中
3.5 Memory Coalescing and DRAM

当读取某个内存地址时,系统会一次性返回整块连续的内存数据,而不只是单一数据,这种机制称为突发传输模式(burst mode),连续的整块内存称为突发传输块(burst section)。
其中原因是因为 DRAM 在读取数据时,要经历一个昂贵的阶段称为行激活:将电荷传递到放大器(Sense Amplifier)上,将电信号转化成可读的 0/1 信号。之后的列读取和 Burst Transfer 阶段耗时都很短。(类似于 Draw Call 和本章最开始的图所示,矩阵越大,硬件利用率越高是一样的。都是耗时很短的任务需要一个前置的代价很大的任务,为了节省时间,把多个小任务打包)
所以为了充分利用这种特性,要精心设计内存访问模式。于是引出了内存合并(Memory Coalescing):当同一个 Warp 里的 32 个线程同时访问全局内存(显存)时,如果它们访问的地址是连续的、对齐的,硬件就会把这些零散的访问请求合并成一次或少数几次大的内存事务,从而极大地提高内存带宽利用率。
以上图所示,若采用内存合并模式进行数据读取,比对数据进行随机读取,数据吞吐量为后者的四倍。

以上图所示,若要一系列 Threads 读取这个矩阵,有两种方式:一种是横向逐列读取,一种是纵向逐行读取。永远记得 Threads 之间执行同样的命令!横向读取时,四个 Thread 在 T0 时刻会分别读取到M_0_0,M_1_0,M_2_0,M_3_0,无法实现 Memory Coalescing。而纵向可以实现,极大提升读取效率。
3.6 Tiling
(其实原理就是矩阵的分块运算)

一个最简单的 W 维方阵相乘,一个 Thread 计算一个结果。如果不采用 Tiling,计算整个方阵需要的总读取次数为 2 * W * W * W 次(计算一个结果需要 2W,共有 W^2 个数据;换一个角度解释,方阵 M 中的一个元素要与 方阵N 中的 W 个元素相乘,方阵M、N共有 2 W^2 个元素)。
但如果采用 Tiling,例如将一个 N 维方阵切成若干个 T 维方阵:

而划分后,每次移动的最小单位从一个数字变成了一个 Tile。观察 M(0,0),它在 Tile_M(0,0) 中,它需要与 W/T 个 Tile_N 进行计算,故需要从 Global_Memory 中移到 shared_memory 中 W/T 次,故两个方阵中所有元素一共需要 2 * W^2 * W / T 次耗时大的读取。而 Tile 内部的计算,数据从 shared_memory 到计算单元,读取效率高。
所以,只要 shared_memory 足够大,能够容纳 tiling 后的矩阵块,就可以将全局内存的数据传输总量降低至原来的 1/T。
Complexity:
1.如何合理安排 Tile 大小

a. 实现内存合并(3.5)
b. 不能超过 shared_memory 能承载范围
c. 尽量分配平均。一个 Tile 分配给一个 SM,让 SM 间负载近似
2. Memory Alignment : Interaction between tiling and burst section

若 Tile 大小不是 burst section 大小的整数倍时,就容易出现这种错位的情况。需要通过 padding 来调整 Tile大小。
3.7 Summarize
回到本章最开始的图,重新分析一下。

①. 总趋势:类似于 Roofline,前期未达到硬件的计算瓶颈,后期已接近硬件的理论峰值算力;
②. 矩阵能否被 K 整除:主要涉及 3.6 中所提到的 Tiling 和 Burst Section 是否对齐的问题。
③. 断崖式下降:以横坐标为 1792 到 1793 的骤降为例,假设 Tile Size 为256 * 128,则 Tile 数为 (1792 / 256) * (1792 / 128) = 7 * 14 = 98。而当仅增长 1 到 1793 时,Tile 数为 8 * 15 = 120。以 A100 为例,它拥有 108 个 SM,这会带来新的计算消耗。
4. Flash Attention
传统 Attention 的问题:

1. 显存爆炸:对于 S 和 P 中间结果需要存储在显存中,当词序列较长时会出现 OOM 问题;
2.带宽瓶颈:读取 QKV (3 * N * d) -> 计算并写入 S (N * N) -> 读取 S 并计算 Softmax (N * N) -> 写入 P (N * N) -> 读取 P 和 V 并计算 O ( V: N * d) -> 写入 O (N * d)
总访存量:O(N² + N×d),计算量: O(N²×d) 。当d较小时,两者相近,GPU算力被严重浪费
4.1 Tiling for QKV Matrix

4.2 Incremental Computation of Softmax
稳定 Softmax 计算方法:

( 保证 ,数据稳定不会溢出)
传统 Softmax 无法在 Tile 中计算,需要全部元素信息。
数学原理很简单,就是分块处理,然后块与块之间比较做一次缩放就行。

两者结合起来就是完整的 FlashAttention 前向传播的过程。

更多推荐

所有评论(0)