英伟达的GPU(3)

关注

原创最新推荐文章于 2026-10-15 19:11:04 发布 1.6k 阅读 AI 写同款· GEO 优化 ›

CUDA——wmma Tensor Core编程 share memory空间的申请 __shared__ int smem[1024] load_matrix_sync 从share memory load一个matrix到fragment中 wmma::load_matrix_sync(fragment<matrix_a/matrix_b, M, N, K, DType, row_major/col_major>& frag, const DType* p, unsigned ldm) 该接口是一个warp级别的操作， . 阅读详情

上节内容：英伟达的GPU(2) (qq.com)

书接上文，上文我们讲到CUDA编程体系和硬件的关系，也留了一个小问题CUDA core以外的矩阵计算能力是咋提供的

本节介绍一下Tensor Core

上节我们介绍了CUDA core，或者一般NPU，CPU执行矩阵运算的逻辑，基本就是矩阵的一条横向量*另一个矩阵的列向量（逻辑上可以这么认为）

如上面的图所示，左边代表了Pascal架构就是P架构的时候 CUDA core 来处理矩阵运算的逻辑，蓝色的矩阵和紫色的矩阵分别代表两个矩阵，然后他俩做点积的时候，基本就是一个横向量*一个列向量。

我们把这个过程细化一下：

两个矩阵A和B，他俩点乘等于一个C，写出来其实就是这样的

I和J代表行和列的角标，k就是算到第几轮计算了。

这个好理解吧。

虽然宏观上我们说是向量点乘向量，但是微观上，其实还是一个格子对一个格子的算。

CUDA Core 实现矩阵乘法

矩阵分块：将大矩阵划分成适合 CUDA 核心处理的小块（block）。通常每个 block 是一个二维块，其中包含多个线程（thread）。例如，16x16 或 32x32 的 block 大小是常见的选择。
线程分配：每个线程块中的线程负责计算结果矩阵 C 中一个小块的元素。例如，一个 16x16 的 block 会有 256 个线程，每个线程计算 C 中一个 16x16 小块中的一个元素。
并行计算：每个线程独立执行矩阵乘法的部分计算。具体来说，每个线程计算一个元素Cij，它需要遍历矩阵 A 的第 i 行和矩阵 B 的第 j 列，进行乘法和累加操作。
共享内存：为了提高性能，CUDA 核心利用共享内存。共享内存是一种高速缓存，允许同一个 block 内的线程共享数据。（这我后面讲Cache和显存那块会细讲）矩阵的分块计算过程中，子矩阵会被加载到共享内存中，减少全局内存访问次数，提高计算效率。

具体计算步骤

用代码表示：

分配线程和块：
- 定义网格（grid）和块（block）的尺寸。（这块看不懂的，去看我上一节讲的CUDA编程线程分级体系）
- 将计算任务分配给每个块和线程。
加载数据到共享内存：
- 每个线程块加载一小块矩阵 A 和 B 到共享内存中。
- 这些小块矩阵被多次重复使用，减少对全局内存的访问。
计算并累加结果：
- 每个线程计算其负责的结果矩阵 C 中一个元素。
- 进行多次小块矩阵乘法的累加，直到完成整个矩阵乘法运算。
写回结果：
- 计算完成后，将结果写回全局内存中的结果矩阵 C。

上面对比图右边那个操作就是Tensor Core的矩阵计算操作，先不解释，就光看就比左边猛很多，对吧？它就不是行列级别了，就是直接矩阵和矩阵运算了，其实当时V系列第一次出Tensor core的时候是很让人惊艳的，但是到了现在，大多数NPU都支持MXM（matrix 乘模块），但是当年V系列推出的时候还是很惊艳的，现在其实也很猛，但是主要是连年的性能提高。

Tensor Core除了对比图中展示的，直接矩阵*矩阵，在一个单位的时钟里面能提供尽可能多的计算能力以外，还有就是可以支持16和32的混合精度能力

如上图所示，在V100刚出的时候就推出了这个功能。

每个 Tensor Core 一个计算周期能执行 4x4x4 GEMM，就相当于64 个 FMA。

比如对于运算D=A*B+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 将计算结果从累加器片段存储回全局内存。
- 结果矩阵的子块被写回到指定的内存位置。