书接上文,上文我们讲到CUDA编程体系和硬件的关系,也留了一个小问题CUDA core以外的矩阵计算能力是咋提供的?
本节介绍一下Tensor Core 。上节我们介绍了CUDA core,或者一般NPU,CPU执行矩阵运算的逻辑,基本就是矩阵的一条横向量另一个矩阵的列向量(逻辑上可以这么认为)
如上面的视频所示,左边代表了Pascal架构就是P架构的时候 CUDA core 来处理矩阵运算的逻辑,蓝色的矩阵和紫色的矩阵分别代表两个矩阵,然后他俩做点积的时候,基本就是一个横向量一个列向量。
我们把这个过程细化一下:两个矩阵A和B,他俩点乘等于一个C,写出来其实就是这样的
I和J代表行和列的角标,k就是算到第几轮计算了。虽然宏观上我们说是向量点乘向量,但是微观上,其实还是一个格子对一个格子的算。
CUDA Core 实现矩阵乘法
(1) 矩阵分块:将大矩阵划分成适合 CUDA 核心处理的小块(block)。通常每个 block 是一个二维块,其中包含多个线程(thread)。例如,16x16 或 32x32 的 block 大小是常见的选择。
(2)线程分配:每个线程块中的线程负责计算结果矩阵 C 中一个小块的元素。例如,一个 16x16 的 block 会有 256 个线程,每个线程计算 C 中一个 16x16 小块中的一个元素。
(3)并行计算:每个线程独立执行矩阵乘法的部分计算。具体来说,每个线程计算一个元素Cij,它需要遍历矩阵 A 的第 i 行和矩阵 B 的第 j 列,进行乘法和累加操作。
(4)共享内存:为了提高性能,CUDA 核心利用共享内存。共享内存是一种高速缓存,允许同一个 block 内的线程共享数据。(这我后面讲Cache和显存那块会细讲)矩阵的分块计算过程中,子矩阵会被加载到共享内存中,减少全局内存访问次数,提高计算效率。
具体计算步骤
(1)分配线程和块:
a. 定义网格(grid)和块(block)的尺寸。(这块看不懂的,去看我上一节讲的CUDA编程线程分级体系)
b. 将计算任务分配给每个块和线程。
(2)加载数据到共享内存:
a. 每个线程块加载一小块矩阵 A 和 B 到共享内存中。
b. 这些小块矩阵被多次重复使用,减少对全局内存的访问。
(3)计算并累加结果:
a. 每个线程计算其负责的结果矩阵 C 中一个元素。
b. 进行多次小块矩阵乘法的累加,直到完成整个矩阵乘法运算。
(4)写回结果:
计算完成后,将结果写回全局内存中的结果矩阵 C。
上面视频右边那个操作就是Tensor Core的矩阵计算操作,先不解释,就光看就比左边猛很多,对吧?它就不是行列级别了,就是直接矩阵和矩阵运算了,其实当时V系列第一次出Tensor core的时候是很让人惊艳的,但是到了现在,大多数NPU都支持MXM(matrix 乘模块),但是当年V系列推出的时候还是很惊艳的,现在其实也很猛,但是主要是连年的性能提高。
Tensor Core除了视频中展示的,直接矩阵矩阵,在一个单位的时钟里面能提供尽可能多的计算能力以外,还有就是可以支持16和32的混合精度能力。
如上图所示,在V100刚出的时候就推出了这个功能。
每个 Tensor Core 一个计算周期能执行 4x4x4 GEMM,就相当于64 个 FMA。
比如对于运算D=AB+C,其中A、B、C 和 D 是 4×4矩阵。矩阵乘法输入 A 和 B 是 FP16矩阵,而累加矩阵 C 和 D我就不非得要求是 FP16,我是FP16还是FP32 矩阵都行。
这个对于CUDA Core来讲,也不是做不到的,你可以手动实现可以通过 CUDA 代码手动实现混合精度计算,例如使用 FP16 数据类型进行部分计算,然后转换为 FP32 进行累加等。但是这么做,第一是墨迹,多出一步增加复杂度和延迟,第二是没专门硬件给你优化啊,因为CUDA Core我们第一章讲过,固定的精度,多少就是多少。
所以对于混合精度,现在也是LLM训练必备的能力了,从某种意义上讲,在NV上想支持,Tensor Core就是必须的了。
再就是A100和后面的型号的sparse matrix的压缩
说白了就是 稀疏矩阵,NV的Tensor core给你做的话,能把0给压没了,你矩阵变小了,算的不就快了吗,这也是为什么大家看NV的datasheet总看不懂的原因
比如上图,看着这么猛,实际上都是按稀疏矩阵算的,所以我们正常算的时候都按1半算,这也就是大家一聊H100就说900多的原因。
Tensor Core是怎么处理矩阵计算的呢?
Tensor Core 矩阵乘法运算
还是假设有两个矩阵 A 和 B,它们的乘积是矩阵 C。Tensor Core 的主要特点是支持 WMMA(Warp Matrix Multiply-Accumulate)操作,这是一个特定的 CUDA 函数,用于执行矩阵乘法和累加。
Tensor Core 计算步骤
- 分配线程和块:
使用 Warp(通常是 32 个线程)来分配计算任务。
一个 Warp 负责计算结果矩阵 C 的一个 16x16 子矩阵。 - 加载数据到共享内存:
将矩阵 A 和 B 的子矩阵块加载到共享内存中。
这些子矩阵块在共享内存中进行矩阵乘法运算。 - 执行矩阵乘法和累加操作:
使用 WMMA API 来执行矩阵乘法和累加操作。
Tensor Core 在一个时钟周期内执行多个浮点运算。 - 写回结果:
计算完成后,将结果写回全局内存中的结果矩阵 C。
这里就得提一嘴WMMA了
WMMA(Warp Matrix Multiply-Accumulate)是 NVIDIA 为 Tensor Core 提供的专用 API,用于在 CUDA 中执行高效的矩阵乘法和累加操作。WMMA API 主要特点和工作原理如下:
WMMA API 的主要特点
-
高效的矩阵运算:
WMMA API 专门优化了矩阵乘法和累加操作,能够在一个时钟周期内执行多个浮点运算,从而显著提高性能。
利用 Tensor Core 的硬件支持,实现高吞吐量的计算。 -
支持混合精度计算:
WMMA API 支持混合精度计算,即输入矩阵可以使用半精度浮点数(FP16),而计算和输出可以使用单精度浮点数(FP32)。
这种方式不仅提高了计算速度,还在一定程度上保持了计算精度。 -
Warp级别的操作:
WMMA API 在 Warp 级别(通常是 32 个线程)进行操作。每个 Warp 负责计算结果矩阵中的一个 16x16 子矩阵。
通过并行执行多个 Warp,实现大规模并行计算。 -
片段操作:
WMMA API 引入了片段(fragment)的概念,用于存储子矩阵和累加器。
片段在寄存器中进行存储和操作,减少了对全局内存的访问,从而提高了性能。
WMMA API 的工作流程
-
声明和初始化片段:
使用 wmma::fragment 声明用于存储矩阵块和累加器的片段。
使用 fill_fragment 对累加器片段进行初始化。 -
加载矩阵数据到片段:
使用 load_matrix_sync 将全局内存中的矩阵数据加载到片段中。
这些数据将被加载到共享内存或寄存器中,以便快速访问和计算。 -
执行矩阵乘法和累加操作:
使用 mma_sync 执行矩阵乘法和累加操作。
该函数将两个输入矩阵片段相乘,并将结果累加到累加器片段中。 -
存储结果到全局内存:
使用 store_matrix_sync 将计算结果从累加器片段存储回全局内存。
结果矩阵的子块被写回到指定的内存位置。
最后值得一说的就是CUDA core 和Tensor Core支持的精度不一样,不是啥下游任务两个都可以做,还是得看具体支持。