想象一下,你是一家大型快递公司的分拣中心主管。每天,成千上万的包裹需要按照目的地分拣、打包、装车。传统的做法是:每个分拣员一次只能处理一个包裹,从传送带上取下、查看标签、放到对应的区域。这就是普通 CPU/GPU 核心的工作方式——一次处理一个数据。 现在,突然来了一个巨大的订单:上百万个包裹需要同时分拣。如果你还靠单个分拣员一个个处理,那得等到天荒地老。于是,你引入了一条“自动化分拣流水线”——它不是一个分拣员,而是一整条能够同时处理“一排排包裹”的智能传送带系统。这就是 Tensor Core 的思想:不再一次处理一个数据,而是一次处理整个“数据矩阵”。 在 AI 的世界里,深度学习模型的每一次推理和训练,本质上都在做一件核心的事情:矩阵乘法(GEMM,General Matrix Multiplication)。无论是图像识别、语音助手还是大语言模型,底层都在疯狂地进行矩阵运算。Tensor Core 就是英伟达 GPU 上专门为这种“批量包裹分拣”设计的超级流水线,它让矩阵乘法的速度比传统核心提升了数倍。 接下来,让我们走进这个“快递分拣中心”,看看 Tensor Core 是如何工作的,以及我们如何写出高效的 GEMM 核函数,让 AI 计算飞起来。

岗位关联

掌握 Tensor Core 与 GEMM 优化,是通往高性能 AI 计算世界的入场券。无论你是做模型训练、推理部署还是 AI 芯片设计,理解矩阵乘法的加速原理,都能让你在性能调优时直击要害,成为团队中能解决“算力瓶颈”的那个人。

目录
  1. 可视化演示
  2. GEMM计算原语
  3. 传统CUDA瓶颈
  4. Tensor Core硬件
  5. 混合精度计算
  6. 分块策略设计
  7. 共享内存优化
  8. 内存访问优化
  9. 双缓冲流水线
  10. 高效GEMM核函数
  11. AI 应用场景
  12. Python 演示
  13. 常见误区
  14. 面试题

可视化演示

展示GEMM分块计算流程(矩阵分块加载到共享内存→Tensor Core MMA→结果写回)、混合精度MMA(FP16输入FP32累加与五种精度类型算力对比)、共享内存双缓冲流水线(重叠访存与计算),三模式切换理解GEMM加速原理。

本节配有交互动效(进入视口后自动加载,约 3KB)。

'+C1+'

GEMM计算原语

在快递分拣中心的故事里,GEMM(General Matrix Multiplication,通用矩阵乘法)就像是“分拣任务书”——它规定了要从哪些货架(矩阵 A)取货、按照什么组合方式(矩阵 B)打包,最终送到哪些出口(矩阵 C)。所有深度学习模型,无论是卷积神经网络还是 Transformer,底层都在反复执行这张“任务书”。

GEMM 的数学形式非常简洁:给定一个 M×K 的矩阵 A 和一个 K×N 的矩阵 B,计算得到一个 M×N 的矩阵 C,其中 C 中的每个元素是 A 的某一行与 B 的某一列的对应元素乘积之和。看起来简单,但计算量却十分惊人——需要 M×N×K 次乘加运算。对于一个 1024×1024×1024 的矩阵乘法,就是超过 10 亿次运算!

GEMM 之所以被称为“计算原语”,是因为它是构建一切 AI 应用的基石。理解了如何加速 GEMM,就理解了如何加速整个 AI 计算。而 Tensor Core 正是为这张“任务书”量身定制的超级流水线。

Cmn=∑k=1KAmk×BknC_{mn} = \sum_{k=1}^{K} A_{mk} \times B_{kn}
这个公式表示:结果矩阵 C 中第 m 行第 n 列的元素,等于 A 矩阵第 m 行与 B 矩阵第 n 列逐元素相乘后求和。
  1. 结果元素 $C_{mn}$C 矩阵中第 m 行、第 n 列的单个元素,就像分拣中心某条出口线上最终打包好的一个包裹。
    CmnC_{mn}
  2. 求和符号 $\sum_{k=1}^{K}$表示从 k=1 到 k=K 的所有项加起来,就像把 K 个货架上的商品都取一遍。
    ∑k=1K\sum_{k=1}^{K}
  3. 乘积 $A_{mk} \times B_{kn}$A 矩阵第 m 行第 k 列的元素乘以 B 矩阵第 k 行第 n 列的元素,相当于从货架 A 取第 k 件、从货架 B 取第 k 件配对。
    Amk×BknA_{mk} \times B_{kn}

传统CUDA瓶颈

回到快递分拣中心,如果只有一群“单件分拣员”——每个分拣员一次只能从传送带上取一个包裹、查看标签、放到对应区域——那么即使有几千个分拣员同时工作,处理上百万个包裹仍然需要很长时间。这正是传统 CUDA 核心的困境:每个核心一次只能执行一条标量指令,处理一个数据元素。

假设我们要计算一个 4096×4096 的矩阵乘法,需要大约 680 亿次乘加运算。传统 CUDA 核心虽然数量众多(比如 A100 有 6912 个),但每个核心每周期只能执行一次乘加运算。即使以 1.5 GHz 的频率运行,也需要数毫秒才能完成——对于需要反复执行数万亿次矩阵乘法的深度学习训练来说,这远远不够。

另一方面,传统核心的访存开销巨大。每计算一个输出元素,需要从内存中读取 A 的一行和 B 的一列,共 2K 个数据。这些数据从全局内存(相当于巨大的中央仓库)传输到计算单元(相当于分拣台),路途遥远、延迟很高。计算一个元素要等很久的数据,导致计算单元经常处于“等待包裹”的空闲状态。

这就是传统 CUDA 核心在 GEMM 上的两大瓶颈:计算密度低(一次只算一个元素)和访存开销高(数据搬运时间远大于计算时间)。就像分拣中心里,分拣员大部分时间在等包裹,而不是在分拣。

Tensor Core硬件

为了解决单件分拣效率低下的问题,你引进了“自动化分拣线”——一整条流水线,一次能从传送带上同时抓取 16 件包裹,自动完成识别、分类、打包的全部流程。在 GPU 中,这条“分拣线”就是 Tensor Core。

Tensor Core 是英伟达从 Volta 架构开始引入的专用矩阵计算单元,到 Ampere 架构(A100)已经非常成熟。它不是取代传统的 CUDA 核心,而是作为“专业分拣线”与“普通分拣员”共存。每个 Tensor Core 本质上是一个 4×4×4 的矩阵乘加阵列,能够通过一条指令(HMMA——Half-precision Matrix Multiply-Accumulate)完成 4×4 矩阵与 4×4 矩阵的乘法,并将结果累加到 4×4 的累加矩阵上。

关键之处在于:Tensor Core 通过一条指令执行了 64 次乘法和 48 次加法,而传统核心需要 112 条指令才能完成同样的工作。这就像自动化分拣线一次处理 16 个包裹,而普通分拣员只能一次搬一个。A100 GPU 上每个流式多处理器(SM)包含 4 个 Tensor Core,整个 GPU 共有 432 个 Tensor Core,理论算力高达 312 TFLOPS(FP16 稀疏),而传统 FP32 算力只有 19.5 TFLOPS——相差 16 倍!

在实际编程中,我们不需要直接操作 Tensor Core 的底层电路,而是通过 CUDA 提供的 wgmma(Ampere)或 mma(Volta/Turing)指令来调用它。而 cuBLAS 等库已经帮我们封装好了这些指令,我们只需要调用库函数即可。但理解原理,才能写出更优的代码。

D4×4=A4×4×B4×4+C4×4D_{4\times4} = A_{4\times4} \times B_{4\times4} + C_{4\times4}
这条指令描述 Tensor Core 一次执行的操作:4×4 的 A 矩阵乘以 4×4 的 B 矩阵,再加上一个 4×4 的 C 矩阵,得到 4×4 的 D 矩阵。
  1. 矩阵乘法 $A_{4\times4} \times B_{4\times4}$两个 4×4 矩阵相乘,包含 64 次乘法和 48 次加法,但 Tensor Core 用硬件电路在一个周期内完成。
    A4×4×B4×4A_{4\times4} \times B_{4\times4}
  2. 累加操作 $+ C_{4\times4}$将乘法的结果加到已有的累加矩阵上,这对于需要多次累加的 GEMM 非常关键,因为不需要每次重新从头开始。
    +C4×4+ C_{4\times4}
  3. 结果矩阵 $D_{4\times4}$最终的输出,4×4 的矩阵,里面每个元素都是对应位置乘积累加的结果。
    D4×4D_{4\times4}

混合精度计算

分拣流水线有一个特殊的设计:它使用“快速标签”(FP16 半精度)来快速识别包裹,但最终记录包裹重量时使用“精确秤”(FP32 单精度)。这就是 Tensor Core 的混合精度计算:输入使用 FP16,累加使用 FP32。

为什么要这样做?因为 FP16 只占 16 位(2 字节),FP32 占 32 位(4 字节)。同样的显存带宽,FP16 可以搬运两倍的数据量。同时,FP16 的矩阵相乘在硬件上可以实现更高的吞吐量——数据“更小”意味着同样面积的芯片可以放更多的计算单元。

但是,直接用 FP16 累加会有精度问题。想象一下:你用精确到克的秤去秤每个包裹(FP16),但如果直接把多个包裹的重量相加,误差会累积。而且 FP16 能表示的数值范围有限(最大约 65504),如果累加结果太大就会溢出。所以 Tensor Core 采用“FP16 输入 + FP32 累加”的方式:输入数据用半精度(省空间、提速度),但累加过程用全精度(保精度、防溢出)。

混合精度带来的加速是惊人的:A100 的 FP16 Tensor Core 算力是 FP32 的 8 倍。当然,这种加速是有代价的——并非所有数据都适合用 FP16 表示。如果矩阵元素的值域过大,FP16 的有限精度可能不足以区分相近的值。因此,实际应用中常常配合“损失缩放”(loss scaling)等技术,先将梯度放大再输入 FP16,计算完再缩小回来。这些细节让混合精度从“能用”变成“好用”。

C~mn=∑k=1KFP32(FP16(Amk)×FP16(Bkn))\tilde{C}_{mn} = \sum_{k=1}^{K} \text{FP32}\left( \text{FP16}(A_{mk}) \times \text{FP16}(B_{kn}) \right)
公式说明:将 A 和 B 的元素转为 FP16 进行乘法,乘积转为 FP32 进行累加,最终结果以 FP32 精度存储。
  1. FP16 转换 $\text{FP16}(A_{mk})$将 A 矩阵的元素从原本的格式转为半精度(16 位浮点数),这就像用快速标签代替详细清单来加快分拣速度。
    FP16(Amk)\text{FP16}(A_{mk})
  2. 半精度乘法 $\times \text{FP16}(B_{kn})$两个半精度数相乘,结果是一个半精度或更高精度的乘积,Tensor Core 硬件内部会自动处理。
    ×FP16(Bkn)\times \text{FP16}(B_{kn})
  3. FP32 累加 $\text{FP32}(\cdot)$将乘积转为单精度浮点数再进行累加,避免误差累积,就像每个包裹最终都要用精确秤重新称重。
    FP32(⋅)\text{FP32}(\cdot)

分块策略设计

分拣中心要处理的是海量包裹,如果让流水线直接面对整个仓库,传送带会被堵死。聪明的做法是:把大订单拆成小批次,每批包裹先运到流水线旁边的一个“临时操作台”(共享内存),操作台里的包裹被流水线快速处理完后,再换下一批。这就是分块(Tiling)策略。

在 GEMM 中,我们把巨大的矩阵 A(M×K)和 B(K×N)切成小块:A 切成大小为 BM×BK 的块,B 切成 BK×BN 的块,这样每个线程块(CTA,Cooperative Thread Array)只需要负责计算一个大小为 BM×BN 的输出块。这个输出块的计算需要 A 的一个 BM×BK 的块和 B 的一个 BK×BN 的块。

分块的核心价值在于数据复用。如果直接计算,每个输出元素需要读取 A 的 K 个元素和 B 的 K 个元素;对于 M×N 的输出,总计需要读取 2×M×N×K 个数据。但如果分成大小为 BM×BN 的块,每个块只需要读取 (BM×BN×K) 个数据(从全局内存),而每个数据会被重复使用 BM×BN 次。比如,A 的一个元素会被同一行的 BN 个输出元素共用,B 的一个元素会被同一列的 BM 个输出元素共用。

用快递的比喻来说:如果没有分块,每次计算一个包裹都要回仓库取一次货;有了分块,一次性把一批相关的货全部搬到操作台,流水线连续处理很多个包裹才需要回一次仓库。分块大小(BM、BN、BK)的选择是 GEMM 调优的关键参数。常用的配置如 128×128×8 或 64×64×16,需要在“块太大导致占不满并行单元”和“块太小导致数据复用率低”之间找平衡。

共享内存优化

分拣中心的“临时操作台”——共享内存——是整个流程效率的关键。共享内存在 GPU 中是一块片上高速存储,距离计算单元极近,访问延迟比全局内存低一个数量级,带宽也高得多。

在每个线程块(CTA)中,我们先把 A 块和 B 块的数据从全局内存加载到共享内存中,然后所有线程反复从共享内存读取数据进行计算。这样,全局内存只需要被访问一次(加载到共享内存),之后的所有 K 次数据复用都从共享内存获取。这就像把包裹从中央仓库搬运到流水线旁的临时操作台,流水线操作员不需要每次跑去仓库。

共享内存的容量是有限的——A100 上每个 SM 有 163KB 的共享内存。因此,我们需要合理安排分块大小,确保 A 块、B 块和输出块都能放得下。同时,共享内存的使用还需要注意“同步”问题:当一个线程块内的线程正在使用共享内存中的数据时,不能有其他线程同时修改这些数据。CUDA 提供了 __syncthreads() 屏障来确保所有线程都到达同一执行点后才能继续。

一个经典的技术是使用共享内存上的矩阵转置。由于全局内存是按行存储的,访问 B 矩阵的一列会导致非合并的内存访问(每个线程读取不同行的同一列)。一个有效的优化是:将 B 块从全局内存加载到共享内存时同时进行转置,这样后续的计算可以在共享内存中按行访问 B 块的列,既保证了加载时的合并访问,也保证了计算时的高效访问。这就像在操作台上把包裹重新排列成流水线最容易抓取的方式。

共享内存的优化效果立竿见影:通过将数据加载到共享内存并复用,全局内存的访问量可以从 O(M×N×K) 降低到 O((M×N×K)/(BM×BN))。例如,对于 4096³ 的 GEMM 和 128×128 的块大小,全局内存访问量减少为原来的约 1/16384。

内存访问优化

在快递分拣中心,如果每个操作员都去仓库的不同角落取不同种类的包裹,仓库的出入口就会拥堵。但是如果操作员们按顺序、连续地取一排排连续的包裹,传送带就能高效运转。这就是“合并内存访问”(coalesced memory access)的思想。

在 GPU 中,全局内存以 32 字节(或更大)的粒度进行访问。当一个 warp(32 个线程)执行一条加载指令时,如果这 32 个线程访问的地址是连续的,那么 GPU 可以合并成一次或少数几次内存事务,高效地完成加载。反之,如果地址分散(比如访问一个矩阵的列),则需要多次内存事务,效率大幅下降。

另一个内存优化要点是避免存储体冲突(bank conflict)。共享内存在物理上被分成 32 个存储体(bank),每个存储体每周期只能响应一次访问。如果多个线程同时访问同一个存储体,就会出现冲突,访问被序列化。一个常见的解决技巧是“填充”(padding):给共享内存数组每行多加几个元素,使每行的起始地址错开,从而避免不同行相同位置的元素落在同一存储体。

还有数据对齐:矩阵的数据在内存中的起始地址应该对齐到 16 字节或更大的粒度,这样 GPU 可以用更高效的方式加载。CUDA 提供的 __align__ 或 cudaMallocPitch 等接口可以帮助实现对齐。在分块加载时,确保每个线程块加载的数据是连续对齐的,也能显著提升加载效率。

内存事务数=⌈覆盖的地址跨度事务粒度⌉\text{内存事务数} = \left\lceil \frac{\text{覆盖的地址跨度}}{\text{事务粒度}} \right\rceil
公式说明:当 32 个线程访问的内存地址跨度越大,需要的内存事务就越多,效率越低。合并访问让跨度最小化,所有线程的一次加载只需要一次事务。
  1. 地址跨度 $\text{覆盖的地址跨度}$32 个线程访问的内存地址的最小值到最大值的距离。如果连续访问,跨度就小。
    覆盖的地址跨度\text{覆盖的地址跨度}
  2. 向上取整 $\lceil \cdot \rceil$即使只有 33 字节的跨度,也需要 2 个 32 字节的事务。连续访问让跨度刚好等于数据大小。
    ⌈⋅⌉\lceil \cdot \rceil

双缓冲流水线

分拣中心的终极优化是让流水线“永不停歇”:当传送带把一批新包裹运到操作台时,流水线正在处理上一批包裹;当当前批处理完,下一批已经就位,无需等待。这就是双缓冲(double buffering)流水线技术。

在 GEMM 核函数中,每个 K 迭代需要完成三个步骤:从全局内存加载 A 块和 B 块到共享内存 → 同步等待所有线程完成加载 → 计算该块的部分乘积并累加。如果串行执行,加载和计算的时间是相加的——GPU 在加载时无事可做,在计算时内存总线空闲。

双缓冲的解决方案是使用两块共享内存缓冲区:当缓冲区 A 的数据正在被计算时,缓冲区 B 正在从全局内存加载下一批数据。计算完成后,交换角色:等待 B 加载完成,然后计算 B 的同时,预加载下一批数据到 A。这样,内存加载和矩阵计算在时间上重叠,GPU 的计算单元和内存带宽都被充分利用。

实现双缓冲的难点在于同步管理。我们需要确保:计算当前缓冲区时不会覆盖正在加载的缓冲区;加载新数据到空闲缓冲区时,该缓冲区中的所有数据都已经被消费完。CUDA 中通常使用 __syncthreads() 协同屏障来同步线程块内的所有线程,或者使用 cp.async 指令(Ampere 架构支持)在硬件层面实现异步拷贝,避免显式同步的负担。

在 A100 上,通过双缓冲和异步拷贝,GEMM 的 Tensor Core 利用率可以达到 85% 以上。这就像一条精心调校的自动化分拣线,传送带、操作台、流水线无缝衔接,每一秒都在高效产出。

T总=T加载+max⁡(T计算,T加载)×(N块−1)T_{\text{总}} = T_{\text{加载}} + \max(T_{\text{计算}}, T_{\text{加载}}) \times (N_{\text{块}} - 1)
公式说明:理想情况下,加载时间被计算时间“掩盖”(overlap),总时间从“加载+计算串行”变为“首块加载+后续计算主导”。
  1. 首次加载 $T_{\text{加载}}$第一个块的数据必须加载完成才能开始计算,这是不可避免的启动开销,就像流水线开机后第一批包裹需要时间走到操作台。
    T加载T_{\text{加载}}
  2. 重叠部分 $\max(T_{\text{计算}}, T_{\text{加载}})$后续每一块的时间是计算时间和加载时间的较大者,因为两者并行——计算这一块的同时加载下一块。
    max⁡(T计算,T加载)\max(T_{\text{计算}}, T_{\text{加载}})
  3. 块数 $N_{\text{块}} - 1$总共有 N_{\text{块}} 个数据块,除了第一个块之外,其余块都能通过双缓冲实现加载与计算重叠。
    N块−1N_{\text{块}} - 1

高效GEMM核函数

掌握了前七节的所有关键技术,我们可以开始组装一条完整的“超级分拣流水线”——一个高效的 GEMM 核函数。在英伟达最新的 Hopper 架构(H100)中,GEMM 的优化已经高度封装在 cuBLAS 等库中,但理解手写核函数的过程仍然是学习 GPU 优化的最佳途径。

一个高效的 Tensor Core GEMM 核函数通常包含以下层级:最外层是“线程块级”的分块调度(对应思维导图中的“块级复用”和“CTA协作”),中间是“warp 级”的 Tensor Core 调用(每个 warp 执行 wgmma 指令完成矩阵乘法),最内层是“线程级”的数据搬运和寄存器缓存。

以 A100 上的 HGEMM(Half-precision GEMM)为例,典型配置如下:线程块大小为 256(8 warps),每个线程块计算 128×128 的输出块,分块大小 BK=32(每步加载 32 层 K)。共享内存双缓冲区,每块 128×32×2 字节(FP16),两块共 32KB。每个 warp 负责计算 64×32 的输出子块,通过 4 条 wgmma 指令完成。通过 cp.async 指令异步加载数据,实现加载与计算的重叠。

核函数性能的衡量标准是“计算效率”——实际达到的 FLOPS 除以硬件峰值 FLOPS。一个优化良好的 GEMM 核函数应该达到 80% 以上的效率。NVIDIA 的 CUTLASS 库提供了一个模板化的 GEMM 实现框架,将分块、加载、计算等步骤抽象为可组合的组件,是目前手写核函数的最佳参考。

对于大多数 AI 从业者来说,不需要真的从零手写 GEMM 核函数——cuBLAS 已经非常高效。但理解这些优化技术,能帮助你在两种场景下受益:一是当你需要调用 cuBLAS 时,你能理解为什么某个参数(如 cublasGemmEx 的数学模式)会影响性能;二是当你遇到 cuBLAS 不支持的矩阵形状或特殊格式时,你能用 CUTLASS 或手写核函数做定制化优化。

AI 应用场景

深度学习训练加速

阶段:模型训练

解决的问题:在大规模深度学习训练中,每个批次的前向传播和反向传播都涉及大量矩阵乘法。传统 CUDA 核心计算速度慢,导致训练一个大型模型需要数周时间。Tensor Core 的混合精度计算和矩阵乘法加速能力,可以将训练时间缩短数倍,同时保持模型精度。

具体案例:在训练 GPT-3 这样的大语言模型时,需要处理 1750 亿参数和数万亿的 token。使用 A100 的 Tensor Core 和 FP16 混合精度训练,相比纯 FP32 训练可以将吞吐量提升约 4 倍,同时显存占用减半,使得更大模型的训练成为可能。

推理服务部署

阶段:模型推理

解决的问题:在线推理服务(如聊天机器人、图像识别 API)要求低延迟和高吞吐量。每次推理都需要执行矩阵乘法,Tensor Core 可以显著降低推理延迟,提升单卡并发服务能力,降低部署成本。

具体案例:使用 TensorRT 优化一个 BERT 模型在 T4 GPU 上的推理。相比未优化的实现,TensorRT 利用 Tensor Core 的 INT8 量化推理,将吞吐量提升约 6 倍,延迟降低约 5 倍,使得单卡可以支撑数千个并发请求。

科学计算加速

阶段:数值模拟

解决的问题:许多科学计算(分子动力学、流体力学、量子化学)的核心也是 GEMM。传统 CPU 和 GPU 标量核心的计算密度不足,Tensor Core 为科学计算提供了新的加速路径。

具体案例:在药物分子动力学模拟中,使用 GROMACS 软件配合 A100 的 Tensor Core 混合精度模式,可以将模拟速度提升 2-3 倍,使得在合理时间内完成微秒级的蛋白质折叠模拟成为可能,加速新药研发流程。

Python 演示

python72 行
import numpy as npimport time def gemm_naive(A, B, C, M, N, K):    """CPU 上的朴素矩阵乘法,作为性能基线"""    for i in range(M):        for j in range(N):            s = C[i, j]            for k in range(K):                s += A[i, k] * B[k, j]            C[i, j] = s    return C def gemm_tiled_numpy(A, B, C, M, N, K, block_size=64):    """用 numpy 分块模拟 GPU 分块策略,降低缓存 miss"""    # 将 C 清零    C[:] = 0    for i0 in range(0, M, block_size):        i1 = min(i0 + block_size, M)        for j0 in range(0, N, block_size):            j1 = min(j0 + block_size, N)            accum = np.zeros((i1 - i0, j1 - j0), dtype=np.float32)            for k0 in range(0, K, block_size):                k1 = min(k0 + block_size, K)                A_blk = A[i0:i1, k0:k1]                B_blk = B[k0:k1, j0:j1]                accum += A_blk @ B_blk            C[i0:i1, j0:j1] = accum    return C def mixed_precision_gemm(A_fp16, B_fp16, M, N, K, block=64):    """模拟 Tensor Core 的 FP16 输入 + FP32 累加"""    A = A_fp16.astype(np.float16)    B = B_fp16.astype(np.float16)    C = np.zeros((M, N), dtype=np.float32)    for i0 in range(0, M, block):        i1 = min(i0 + block, M)        for j0 in range(0, N, block):            j1 = min(j0 + block, N)            accum = np.zeros((i1 - i0, j1 - j0), dtype=np.float32)            for k0 in range(0, K, block):                k1 = min(k0 + block, K)                A_blk = A[i0:i1, k0:k1].astype(np.float16)                B_blk = B[k0:k1, j0:j1].astype(np.float16)                accum += A_blk.astype(np.float32) @ B_blk.astype(np.float32)            C[i0:i1, j0:j1] = accum    return C if __name__ == "__main__":    M, N, K = 256, 256, 256    np.random.seed(0)    A = np.random.randn(M, K).astype(np.float32) / np.sqrt(K)    B = np.random.randn(K, N).astype(np.float32) / np.sqrt(K)    C_naive = np.zeros((M, N), dtype=np.float32)    C_tiled = np.zeros((M, N), dtype=np.float32)    C_ref = A.astype(np.float32) @ B.astype(np.float32)     t0 = time.perf_counter()    gemm_naive(A, B, C_naive, M, N, K)    t1 = time.perf_counter()    gemm_tiled_numpy(A, B, C_tiled, M, N, K, block_size=64)    t2 = time.perf_counter()    C_mp = mixed_precision_gemm(A, B, M, N, K, block=64)    t3 = time.perf_counter()     err1 = np.abs(C_naive - C_ref).max()    err2 = np.abs(C_tiled - C_ref).max()    err3 = np.abs(C_mp - C_ref).max()    print(f"naive  : {t1-t0:.4f}s  max_err={err1:.6f}")    print(f"tiled  : {t2-t1:.4f}s  max_err={err2:.6f}")    print(f"mixed  : {t3-t2:.4f}s  max_err={err3:.6f}")    print("注: numpy 已高度优化, 分块优势在缓存受限时更明显")
目的:这段代码演示了从朴素矩阵乘法到分块优化、再到混合精度计算的三个层次,帮助理解 Tensor Core GEMM 的核心优化思想如何在 CPU 上进行模拟。
  1. for i in range(M)朴素实现的三个嵌套循环,对应最直接的三重循环遍历。
  2. s += A[i, k] * B[k, j]内层循环逐元素乘加,每次从头读取 A 的行和 B 的列。
  3. accum += A_blk @ B_blk分块优化的关键:利用 numpy 的高效矩阵乘法处理子块。
  4. A_blk.astype(np.float16)模拟 Tensor Core 的 FP16 输入转换。
  5. accum.astype(np.float32)计算时提升为 FP32,模拟混合精度累加。
  6. err1 = np.abs(C_naive - C_ref).max()验证朴素实现的正确性,最大误差应接近 0。

常见误区

误区很多人以为有了 Tensor Core 之后,传统的 CUDA 核心就没用了,所有计算都应该用 Tensor Core 完成。

正解Tensor Core 只擅长矩阵乘法和卷积等特定形态的运算,而 CUDA 核心是通用的标量计算单元。在 GEMM 核函数中,Tensor Core 负责矩阵乘法的主体计算,但数据搬运、地址计算、格式转换、激活函数等操作仍然需要 CUDA 核心参与。两者是协作关系:CUDA 核心负责“打杂”,Tensor Core 负责“重活”。

为何常见:看到 Tensor Core 的算力数字(TFLOPS)远高于 CUDA 核心,容易产生“碾压”的印象,忽视了 GPU 执行流程中各种辅助操作的必要性。

误区很多初学者认为使用 FP16 计算必然导致精度下降,会影响模型效果,因此不敢用混合精度训练。

正解Tensor Core 的混合精度方案采用 FP16 输入 + FP32 累加,累加过程保持全精度,所以精度损失主要来自输入的量化。对于大多数深度学习模型,配合损失缩放(loss scaling)技术,混合精度训练可以达到与全 FP32 训练几乎相同的精度。此外,FP16 还能减少显存占用、提升训练速度,是平衡速度与精度的工程艺术。

为何常见:直觉上“精度越低误差越大”,但忽略了深度学习中噪声的存在和损失缩放等补偿技术的作用。

误区不少开发者在 GPU 上调用 cuBLAS 库计算矩阵乘法后,就认为性能调优工作已结束,完全信任库的默认配置。

正解cuBLAS 虽然封装了高效实现,但默认配置并不总是最优。矩阵的形状(M、N、K 的比例)、数据布局(是否转置)、批处理大小、输入输出精度等因素都会影响最优算法选择。例如,对于小矩阵或非标准形状的矩阵,可能需要通过 cublasGemmEx 指定不同的算法(CUBLAS_GEMM_DEFAULT_TENSOR_OP 等)来获得最佳性能。理解 Tensor Core 和分块原理,才能做出正确的调优决策。

为何常见:库函数“拿来即用”的便利性,让开发者误以为库已经处理了所有情况,忽略了库需要配合具体场景做配置的本质。

误区在编写 GEMM 核函数时,初学者容易认为分块越大,数据复用率越高,性能就越好。

正解分块大小受限于硬件资源:共享内存容量是有限的(A100 为 163KB/SM),寄存器数量有限,线程块数量受 SM 资源限制。块太大导致单个 SM 能容纳的线程块减少,无法用足够的线程隐藏访存延迟;块太小又导致复用不足。实际设计中需要通过调优找到最优值,常见的做法是使用自动调优工具(如 CUTLASS 的 autotuner)搜索最优配置。

为何常见:只看数据复用率而忽略了硬件资源约束,是初学者做计算优化时常犯的典型错误。

面试题

Q1. GEMM 是深度学习中最重要的计算原语,它的字母展开是什么?

单选题 · easy

  • A.General Matrix Multiplication← 正确答案
  • B.Generic Memory Management
  • C.Gradient Model Merging
  • D.Global Memory Mapping

答案:A。GEMM 的全称是 General Matrix Multiplication(通用矩阵乘法)。它是矩阵乘法的通用形式,深度学习模型的卷积和全连接层最终都可以转化为 GEMM 运算。其他选项都是干扰项,与矩阵乘法的定义无关。

Q2. Tensor Core 与传统 CUDA 核心在计算方式上的根本区别是什么?

单选题 · mid

  • A.Tensor Core 频率更高
  • B.Tensor Core 一次执行矩阵乘加运算← 正确答案
  • C.Tensor Core 显存更大
  • D.Tensor Core 只支持整数运算

答案:B。Tensor Core 的核心理念是“单指令多数据”——通过一条指令(如 HMMA)完成 4×4×4 的矩阵乘加运算,而传统 CUDA 核心一次只能执行一个标量乘加。频率、显存不是区分 Tensor Core 和 CUDA 核心的关键特征,Tensor Core 也不只支持整数运算。

Q3. Tensor Core 采用的混合精度策略是什么?

单选题 · mid

  • A.FP32 输入,FP64 累加
  • B.FP16 输入,FP32 累加← 正确答案
  • C.INT8 输入,FP16 累加
  • D.BF16 输入,FP16 累加

答案:B。Tensor Core 的标准混合精度模式是 FP16(半精度)输入,FP32(单精度)累加。这种设计既利用了 FP16 的小体积带来更高的数据吞吐和计算密度,又通过 FP32 累加保证了数值稳定性,避免误差累积。其他选项中的精度组合不是 Tensor Core 的主流配置。

Q4. 在 GEMM 核函数中,分块(Tiling)策略带来哪些好处?

多选题 · hard

  • A.提高数据复用率← 正确答案
  • B.减少全局内存访问次数← 正确答案
  • C.增加 GPU 核心频率
  • D.使数据可以放入共享内存← 正确答案

答案:ABD。分块策略将大矩阵切成小块,使得块内数据可以加载到共享内存中复用,从而提高数据复用率(A)、减少全局内存访问次数(B)、使数据能够放入高速共享内存(D)。但分块不会改变 GPU 核心频率——那是硬件特性。

Q5. 共享内存在 GEMM 中的作用是什么?

单选题 · easy

  • A.存放最终计算结果
  • B.暂存分块数据供计算单元复用← 正确答案
  • C.替代全局内存存储所有数据
  • D.提高 GPU 核心频率

答案:B。共享内存是 GPU 上的一块片上高速存储,在 GEMM 中主要用来暂存从全局内存加载的分块数据,供线程块内的计算单元反复读取,从而减少昂贵的全局内存访问。它容量有限,不可能存储所有数据,也无法提高核心频率。

Q6. 以下哪些措施可以有效提升 GEMM 核函数的性能?

多选题 · mid

  • A.使用双缓冲让加载和计算重叠← 正确答案
  • B.保证线程访问连续内存实现合并访问← 正确答案
  • C.增加全局面内存访问次数
  • D.使用 __syncthreads() 保证数据同步← 正确答案

答案:ABD。双缓冲能隐藏内存延迟(A),合并访问能提高内存带宽利用率(B),正确的同步能确保数据一致性(D)。增加全局内存访问次数(C)显然会降低性能。这几项对应文章第 6-8 节的核心优化技术。

Q7. 在 A100 GPU 上,一个 Tensor Core 的一次 HMMA 指令可以完成多少次乘加运算?

单选题 · hard

  • A.16 次
  • B.32 次
  • C.64 次← 正确答案
  • D.112 次

答案:C。一个 HMMA 指令执行 4×4 矩阵乘以 4×4 矩阵,包含 64 次乘法和 48 次加法,共 112 次浮点运算。但题目问的是乘加运算(FMA)次数,即每个乘加算一次,共 64 次。这也是 Tensor Core 比标量核心快一个数量级的关键。

Q8. 为什么 Tensor Core 使用 FP32 累加而不是 FP16 累加?

单选题 · mid

  • A.FP32 累加更快
  • B.防止累加过程误差累积和溢出← 正确答案
  • C.FP32 数据占内存更小
  • D.因为 CUDA 核心不支持 FP16

答案:B。FP16 的指数位和尾数位都较短,表示范围和精度都有限。在 GEMM 中需要累加 K 项乘积,如果累加器也用 FP16,误差会不断累积,且可能超出最大表示范围(65504)。FP32 累加提供充足的位数来携带累积值,速度损失不大但精度显著提升。

Q9. 关于共享内存双缓冲技术,以下说法正确的有哪些?

多选题 · hard

  • A.使用两块缓冲区交替加载和计算← 正确答案
  • B.可以让内存加载和矩阵计算重叠进行← 正确答案
  • C.不需要考虑线程块内同步问题
  • D.使用 cp.async 可以实现异步数据拷贝← 正确答案

答案:ABD。双缓冲使用两块缓冲区(乒乓球式)角色互换,让加载和计算重叠,隐藏内存延迟(A、B 正确)。在 Ampere 架构上可以用 cp.async 指令实现异步拷贝,避免显式同步(D 正确)。但双缓冲并没有消除同步问题——计算完一个缓冲区后仍需要 __syncthreads() 确保所有线程都完成消费,之后才能覆盖该缓冲区(C 错误)。

Q10. 深度学习训练中,使用 Tensor Core 混合精度相比传统 FP32 的主要优势是什么?

单选题 · easy

  • A.模型精度更高
  • B.训练速度更快且显存占用更少← 正确答案
  • C.无需梯度计算
  • D.不需要损失函数

答案:B。Tensor Core 混合精度训练的核心优势是速度——输入数据从 FP32 换成 FP16 后,数据量减半,内存带宽压力降低,Tensor Core 的 FP16 计算吞吐量是 FP32 的数倍。同时显存占用也减少。但模型精度不会更高(通常持平),梯度计算和损失函数仍然必要。