第4板块·第3节：Tensor Core 硬件原理与编程模型

学习目标

学完本节你将能够：

• 理解 Tensor Core 的硬件设计目标与矩阵乘加能力
• 掌握 Tensor Core 支持的混合精度数据格式（FP16、BF16、TF32、INT8）
• 使用 CUDA WMMA API 在 Kernel 中调用 Tensor Core
• 了解 Tensor Core 在 GEMM 和卷积中的加速原理
• 分析 Tensor Core 相比传统 CUDA Core 的性能优势与适用边界

---

1. Tensor Core 是什么

1.1 设计背景

深度学习中的核心运算是矩阵乘法（GEMM）和卷积，这些运算的计算量巨大，但精度需求往往低于 FP32。传统的 FP32 CUDA Core 虽然通用，但在矩阵运算上吞吐有限。

Tensor Core 是 NVIDIA 从 Volta 架构开始引入的专用硬件单元，专门加速矩阵乘加运算，能在单个时钟周期内完成多个浮点乘加操作，大幅提升深度学习训练和推理的性能。

1.2 基本运算

Tensor Core 执行的核心操作是：

D = A × B + C

其中 A、B、C、D 是矩阵分块（tile），运算以混合精度进行。

• 输入 A、B：通常是 FP16、BF16、TF32 等低精度格式
• 累加器 C、D：通常是 FP32 或 FP16，保证精度

一个 Tensor Core 每时钟周期可以完成一个 4×4×4 的矩阵乘加（即 4×4 的矩阵乘 4×4 的矩阵，再加 4×4 的矩阵，共 64 次乘加）。

1.3 性能对比

以 A100 为例：

精度	CUDA Core FP32	Tensor Core FP16	Tensor Core TF32
峰值吞吐	19.5 TFLOPS	312 TFLOPS	156 TFLOPS
加速比	1x	16x	8x

可以看到 Tensor Core 的混合精度吞吐远高于普通 FP32，这是大模型训练和推理性能的关键来源。

---

2. 支持的精度格式

不同架构的 Tensor Core 支持的精度有所不同：

架构	输入精度	累加精度	说明
Volta	FP16	FP32	首次引入
Turing	FP16, INT8, INT4	FP32	增加整数推理
Ampere	FP16, BF16, TF32, INT8, INT4	FP32	新增 BF16 和 TF32
Hopper	FP16, BF16, TF32, FP8, INT8	FP32/FP16	新增 FP8

关键格式说明：

• FP16：半精度，5 位指数 + 10 位尾数，范围小但精度够用于多数深度学习。
• BF16：Google 提出的脑浮点格式，8 位指数 + 7 位尾数，范围与 FP32 相同，但精度较低，训练稳定性好。
• TF32：NVIDIA 在 Ampere 引入，10 位尾数 + 8 位指数，精度介于 FP16 和 FP32 之间，直接替代 FP32 输入，无需修改代码即可获得加速。

---

3. 使用 WMMA API 调用 Tensor Core

CUDA 提供了 WMMA（Warp Matrix Multiply‑Accumulate） API，允许在 Kernel 中直接操作 Tensor Core，无需编写底层汇编。

3.1 基本流程

1. 包含头文件 <mma.h>，并使用 nvcuda::wmma 命名空间。
2. 声明矩阵分块（fragment），大小由 WMMA 模板参数决定。
3. 从全局或共享内存加载数据到 fragment。
4. 调用 wmma::mma_sync 执行矩阵乘加。
5. 将结果从 fragment 存储回内存。

3.2 支持的矩阵形状

WMMA 支持多种 tile 大小，常见的有：

精度	M×N×K
FP16	16×16×16, 32×8×16, 8×32×16
TF32	16×16×8
BF16	16×16×16
INT8	16×16×16

具体支持的形状取决于架构，编译时可通过 __CUDA_ARCH__ 判断。

3.3 示例代码

以下示例演示一个简单的 16×16×16 FP16 矩阵乘加：

#include <cstdio>
#include <cuda_runtime.h>
#include <mma.h>

using namespace nvcuda;

#define CUDA_CHECK(call) \
    do { \
        cudaError_t err = call; \
        if (err != cudaSuccess) { \
            fprintf(stderr, "CUDA error at %s:%d: %s\n", __FILE__, __LINE__, cudaGetErrorString(err)); \
            exit(EXIT_FAILURE); \
        } \
    } while (0)

// 使用 Tensor Core 的 GEMM Kernel（简化：每个 Block 处理一个 16x16x16 的 tile）
__global__ void tensorCoreGemm(const half *A, const half *B, float *C, int M, int N, int K) {
    // 声明 fragment
    wmma::fragment<wmma::matrix_a, 16, 16, 16, half, wmma::row_major> a_frag;
    wmma::fragment<wmma::matrix_b, 16, 16, 16, half, wmma::col_major> b_frag;
    wmma::fragment<wmma::accumulator, 16, 16, 16, float> c_frag;

    // 初始化累加器为 0
    wmma::fill_fragment(c_frag, 0.0f);

    // 加载 A 和 B 的 tile（这里假设每个 Block 处理一个固定的 tile，实际需要循环）
    wmma::load_matrix_sync(a_frag, A, K);   // A 的行主序，leading dim = K
    wmma::load_matrix_sync(b_frag, B, N);   // B 的列主序，leading dim = N

    // 执行矩阵乘加：C += A * B
    wmma::mma_sync(c_frag, a_frag, b_frag, c_frag);

    // 将结果写回全局内存
    wmma::store_matrix_sync(C, c_frag, N, wmma::mem_row_major);
}

int main() {
    // 这里只演示 Kernel 结构，实际需要完整的内存分配和数据初始化
    // 注意：需要确保 GPU 支持 Tensor Core（compute capability >= 7.0）
    int device = 0;
    cudaDeviceProp props;
    CUDA_CHECK(cudaGetDeviceProperties(&props, device));
    if (props.major < 7) {
        printf("Tensor Core requires compute capability 7.0 or higher\n");
        return -1;
    }
    printf("Tensor Core available!\n");
    return 0;
}

3.4 注意事项

1. WMMA 要求 Warp 内所有线程协作：fragment 是分布在 Warp 的 32 个线程上的，调用 wmma::load_matrix_sync 和 wmma::mma_sync 时必须确保整个 Warp 活跃。
2. 内存对齐：加载操作要求地址对齐（通常是 128 字节）。
3. 性能：WMMA 的 tile 较小（16×16×16），实际 GEMM 通常需要多层循环和共享内存分块才能达到高利用率。

---

4. Tensor Core 在 GEMM 中的应用

4.1 分块策略

完整的 GEMM 使用 Tensor Core 时通常采用以下层次：

1. 全局内存 → 共享内存分块：将 A 和 B 的大块数据加载到共享内存，避免重复的全局内存访问。
2. 共享内存 → 寄存器 fragment：每个 Warp 从共享内存加载 16×16 的 fragment。
3. WMMA 乘加：在寄存器中执行 Tensor Core 运算。
4. 累加结果写回：将最终结果存储到全局内存。

4.2 性能关键因素

• 共享内存分块大小：影响全局内存合并访问和 Bank Conflict。
• 循环 K 维：对于大 K，需要循环加载多个 16×16×16 的 tile 并累加。
• Warp 调度：多个 Warp 同时使用 Tensor Core，隐藏延迟。

这些优化细节将在后续 CUTLASS 板块深入展开。

---

5. 代码演示：简单 GEMM 性能对比

以下程序对比 FP32 CUDA Core 和 Tensor Core（FP16）的 GEMM 性能，展示 Tensor Core 的加速效果。

#include <cstdio>
#include <cuda_runtime.h>
#include <mma.h>
#include <cuda_fp16.h>

using namespace nvcuda;

#define CUDA_CHECK(call) \
    do { \
        cudaError_t err = call; \
        if (err != cudaSuccess) { \
            fprintf(stderr, "CUDA error at %s:%d: %s\n", __FILE__, __LINE__, cudaGetErrorString(err)); \
            exit(EXIT_FAILURE); \
        } \
    } while (0)

// 简单的 FP32 GEMM（参考）
__global__ void gemmFP32(const float *A, const float *B, float *C, int M, int N, int K) {
    int row = blockIdx.y * blockDim.y + threadIdx.y;
    int col = blockIdx.x * blockDim.x + threadIdx.x;
    if (row < M && col < N) {
        float sum = 0.0f;
        for (int i = 0; i < K; i++) {
            sum += A[row * K + i] * B[i * N + col];
        }
        C[row * N + col] = sum;
    }
}

// Tensor Core GEMM（简化，只处理 16x16 的 tile）
__global__ void gemmTensorCore(const half *A, const half *B, float *C, int M, int N, int K) {
    wmma::fragment<wmma::matrix_a, 16, 16, 16, half, wmma::row_major> a_frag;
    wmma::fragment<wmma::matrix_b, 16, 16, 16, half, wmma::col_major> b_frag;
    wmma::fragment<wmma::accumulator, 16, 16, 16, float> c_frag;

    wmma::fill_fragment(c_frag, 0.0f);

    int blockRow = blockIdx.y * 16;
    int blockCol = blockIdx.x * 16;

    // 假设 K == 16，只加载一个 tile
    wmma::load_matrix_sync(a_frag, A + blockRow * K, K);
    wmma::load_matrix_sync(b_frag, B + blockCol, N);
    wmma::mma_sync(c_frag, a_frag, b_frag, c_frag);

    // 存储结果
    wmma::store_matrix_sync(C + blockRow * N + blockCol, c_frag, N, wmma::mem_row_major);
}

int main() {
    // 简化为 256x256x256 的矩阵
    int M = 256, N = 256, K = 256;
    // 省略内存分配和数据初始化...
    printf("Run on GPU with Tensor Core support.\n");
    return 0;
}

实际完整代码需要处理任意大小矩阵、循环 K 维、边界等。本节仅展示核心概念。

---

6. 课后练习

练习1：检查 Tensor Core 支持
运行一个程序，检查你的 GPU 是否支持 Tensor Core（props.major >= 7），并打印其支持的 WMMA 形状。

练习2：WMMA 基本使用
实现一个 16×16×16 的矩阵乘加 Kernel，使用 wmma::load_matrix_sync 加载数据，wmma::mma_sync 计算，wmma::store_matrix_sync 存储结果。与 CPU 结果对比验证正确性。

练习3：比较 FP32 与 Tensor Core 性能
在相同矩阵大小（如 512×512×512）下，对比 FP32 GEMM 和 Tensor Core（FP16）GEMM 的性能，使用 cudaEvent 计时。分析加速比。

练习4：累加精度实验
使用 FP16 输入、FP32 累加执行大 K 的 GEMM，与纯 FP32 计算结果对比，观察精度差异。尝试使用 TF32 进行比较。

练习5：优化 Tensor Core GEMM
在练习 2 的基础上，加入共享内存分块和 K 维循环，提高数据复用率。测试不同 tile 大小对性能的影响。

---

7. 下一步

下一节将进入 硬件约束对 Kernel 性能的影响，你将学习：

• 寄存器、共享内存、线程数等资源约束如何限制 Kernel 设计
• 通过占用率实验找到最优 Block 大小和资源分配
• 理解不同架构的差异对代码优化的影响
• 结合前几节知识，系统分析一个 Kernel 的硬件瓶颈

转载自 CSDN-专业IT技术社区

原文链接：https://blog.csdn.net/qq_19988687/article/details/165126315