有这样一家主打“AI菜系”的知名餐厅——它的招牌菜是各种“模型大餐”(神经网络)。后厨里有两位不可或缺的大师傅:一位叫 cuBLAS,专精各种食材(矩阵)的切配与焯水——凡是矩阵乘法、向量运算这类基础活儿,找他一准没错;另一位叫 cuDNN,擅长按经典菜谱(卷积、池化、归一化这些深度学习算子)精准复刻每一道招牌菜。餐厅老板(PyTorch、TensorFlow 这些深度学习框架)从不亲自下厨,只需喊一嗓子“今天训练 100 份图像分类模型”,后厨的两位大师傅便会自动协作,把切配、配菜、装盘(张量计算)安排得明明白白。没有他们,后厨就只能靠纯手工(Python 循环)慢火细熬,一盘菜要等到猴年马月。这篇文章,就带你走进这两位大师傅的厨房,看看他们的刀工秘籍和配菜心法。

岗位关联

了解 cuBLAS 与 cuDNN,是通往 GPU 高性能计算和深度学习系统优化之路的必修课。几乎所有主流 AI 框架的训练与推理都踩在这两块基石之上,掌握它们,你才能真正看懂性能瓶颈在哪、优化空间在哪。

目录
  1. 可视化演示
  2. cuBLAS等级例程
  3. GEMM核心用法
  4. cuDNN算子全家桶
  5. Tensor Core加速
  6. 框架集成落地
  7. 性能调优实战
  8. AI 应用场景
  9. Python 演示
  10. 常见误区
  11. 面试题

可视化演示

动态展示二分查找在有序数组中通过不断折半缩小搜索范围定位目标值的过程,直观呈现O(log n)的高效搜索策略。

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

二分查找 - 初始状态 检查mid位置元素与目标值比较 命中目标值则搜索完成 mid=(lo+hi)/2 每轮折半缩小范围

cuBLAS等级例程

cuBLAS 这位“刀工大师傅”的刀法并非一种,而是分了三层功力。你可以想象一家餐厅的备菜间:Level 1 是“单根葱切段”——只处理一个向量,做缩放、点积、找最大值这类简单的食材预处理;Level 2 是“一把葱配一块肉”——让矩阵和向量配合,比如解线性方程组时反复用到的矩阵向量乘;Level 3 是真正的“全席大宴”——矩阵和矩阵直接对撞,这就是深度学习里最常点的硬菜 GEMM。

在 NVIDIA 的官方定义中,cuBLAS 的 Level 1 例程负责向量与向量之间的运算(如 cublasSdot 计算点积、cublasSaxpy 做标量乘加);Level 2 例程负责矩阵与向量运算(如 cublasSgemv 做矩阵向量乘);Level 3 例程负责矩阵与矩阵运算(如 cublasSgemm,即通用矩阵乘法)。等级越高,单次调用能“翻动的食材”就越多,计算密度也越高。这也是为什么在实际工程中,大家几乎只关注 Level 3 的 GEMM——哪怕一个卷积层,最终也会被底层改造成若干次 GEMM 来执行。

Level 1 例程看起来简单,但它们的价值在于“快”。例如求一个长度为 100 万的向量的 L2 范数,Python 里写 for 循环需要几千万次操作,而 cublasSnrm2 在 GPU 上用并行归约(parallel reduction)一次启动数千个线程,把求和拆成树状结构层层合并,几毫秒内就能给出结果。这就像后厨备菜时,几十个帮厨同时开工,每人只切一小把葱,最后统一收拢——效率远高于一个厨师从头切到尾。

回到正题:Level 3 的 GEMM 之所以成为 cuBLAS 的“招牌菜”,是因为它把计算密度推到了极致。矩阵 A 的每一行与矩阵 B 的每一列都可以独立求点积,这天然就是并行计算的完美形态。cuBLAS 内部还会把大矩阵切成小块(tile),让数据在共享内存里反复利用,避免频繁访问显存——这就是后厨里“一次性切好一桌菜的量,放在手边随取随用”的智慧。

y=αx+y(Level 1 的 axpy 操作:标量乘向量再加向量)\displaystyle y = \alpha x + y \quad \text{(Level 1 的 axpy 操作:标量乘向量再加向量)}
这个公式描述的是“把向量 x 放大 α 倍,再加到向量 y 上”,对应后厨里“把一筐土豆按每斤 2 元的价格算总价,再把账记到总账单上”。
  1. 标量 α这是一个普通的数,比如 2.5。它表示“放大倍数”,先单独算好,等待与 x 的每个分量相乘。
    α=2.5\displaystyle \alpha = 2.5
  2. 向量 x 与向量 yx 是“单价清单”,y 是“累计账单”,两者各有 N 个分量,一一对应。
    x=[x1,x2,…,xn]T,y=[y1,y2,…,yn]T\displaystyle x = [x_1, x_2, \dots, x_n]^T, \quad y = [y_1, y_2, \dots, y_n]^T
  3. 逐元素乘累加每个分量独立执行 y[i] = α * x[i] + y[i],GPU 上每个线程负责一个 i,N 个线程并行完成。
    yi←αxi+yi,i=1,2,…,n\displaystyle y_i \leftarrow \alpha x_i + y_i, \quad i = 1, 2, \dots, n

GEMM核心用法

当老板喊出“来一份全连接层”时,后厨的刀工大师傅实际上执行的是 GEMM——通用矩阵乘法。它的公式非常朴素:结果矩阵 C 的第 i 行第 j 列元素,等于 A 的第 i 行与 B 的第 j 列逐元素相乘再求和。但如果只是这么理解,你会发现和小学学的矩阵乘法没区别。GEMM 真正的“通用”之处在于它多了 α 和 β 两个缩放参数,能把“乘法”和“加法”合并成一次调用。

这里的 α 和 β 是标量。当 α = 1、β = 0 时,就是标准的 C = A × B;当 β = 1 时,变成了“累加模式”,适合把多个矩阵乘法的结果叠加到一起。在残差网络(ResNet)里,C = A × B + C 这种形式被大量使用——它让卷积输出的残差可以直接加到原特征图上,不需要额外的显存拷贝。cuBLAS 把这件事做进了硬件流水线里,一次内核启动就完成乘加两步,省掉了中间环节。

调用 cuBLAS 的 GEMM 并不是简单写一行 cublasSgemm(handle, ...) 就完事,关键坑位在内存布局。C/C++ 默认的行主序(row-major)与 cuBLAS 期望的列主序(column-major)需要小心处理——cuBLAS 为兼容 Fortran 传统,默认按列主序读取矩阵。这就像后厨的调料架:你习惯按“酱油→醋→盐”横着摆,但大师傅找调料是按“竖列”找的,如果直接递给他,他会全找错。解决方法是先把矩阵转置,或者用 CUBLAS_OP_T 这个转置标志告诉 cuBLAS“我给你的矩阵是行主序的,你按转置的方式读”。

完整的 API 调用流程可以分为四步:创建 cuBLAS 句柄(cuBLAS handle)、分配显存并拷贝输入矩阵、调用 cublasSgemm、拷贝结果回主机并销毁句柄。句柄就像后厨的“工作证”,所有 cuBLAS 函数都需要它来识别上下文。而 cublasSgemm 的完整签名包含 15 个参数:3 个转置标志、3 个维度(M、N、K)、2 个缩放因子、6 个指针与步长信息。初学者常在这里犯错——维度 M、N、K 的顺序和矩阵形状的对应关系必须精确匹配,否则就会出现“菜炒好了但端错桌”的尴尬。

C=α(A×B)+βC\displaystyle C = \alpha (A \times B) + \beta C
C[M][N]=α∑k=1KA[M][k]⋅B[k][N]+βC[M][N]\displaystyle C[M][N] = \alpha \sum_{k=1}^{K} A[M][k] \cdot B[k][N] + \beta C[M][N]
这个公式展示了 GEMM 的三个维度:M 是 A 的行数、K 是 A 的列数(也是 B 的行数)、N 是 B 的列数。整个矩阵乘法由 M × N 个独立的点积组成。
  1. 遍历 kK 是“内维”,点积的累加次数。A 的第 i 行有 K 个元素,B 的第 j 列也有 K 个元素,两者一一相乘再累加。
    acc=∑k=1KA[i][k]⋅B[k][j]\displaystyle \text{acc} = \sum_{k=1}^{K} A[i][k] \cdot B[k][j]
  2. 写入 C将累加结果乘以 α,再加上 β 倍的旧 C 值,写入输出位置。
    C[i][j]=α⋅acc+β⋅C[i][j]\displaystyle C[i][j] = \alpha \cdot \text{acc} + \beta \cdot C[i][j]
  3. 三个转置标志opA、opB 分别告诉 cuBLAS 是否把 A、B 转置后再参与运算,opC 通常为普通模式。用转置标志可以避免实际搬运数据。
    opA∈{N,T,C},opB∈{N,T,C}\displaystyle opA \in \{N, T, C\}, \quad opB \in \{N, T, C\}

cuDNN算子全家桶

如果说 cuBLAS 是刀工大师傅,那 cuDNN 就是“配菜总管”——他的工作不是处理基础食材,而是把深度学习里程碑式的经典菜谱(卷积、池化、归一化)做成标准化的流水线,让后厨(GPU)能批量生产。老板不用关心每道菜的细节,只要报菜名(调用 cuDNN API),配菜总管就能按最地道的方子做出来。

cuDNN 最核心的算子当属卷积(convolution)。卷积的数学定义看起来吓人,但它本质上是一个“滑动窗口的点积”:一个比输入小的卷积核(比如 3×3)在输入特征图上从左到右、从上到下滑动,每到一个位置,就把窗口内的像素与卷积核逐元素相乘再求和,产生输出特征图的一个像素。这个操作翻译成矩阵语言,恰好就是 cuBLAS 最擅长的 GEMM——也因此 cuDNN 内部常常把卷积拆成 GEMM 来执行:先把输入图像按窗口展开成一个大矩阵(im2col),再做一次大矩阵乘法。这就像配菜总管把切好的食材按固定顺序摆成“标准配菜矩阵”,刀工大师傅直接下锅翻炒。

除了卷积,cuDNN 还提供池化(pooling)算子——它像个“信息压缩器”。最大池化选窗口里的最大值,平均池化取窗口平均值;它们让特征图的尺寸缩小,同时保留最重要的信息。归一化(normalization)算子则负责把数据拉回到合适的数值范围,比如批归一化(Batch Normalization)会对一个 batch 内每个通道的数据做标准化,让均值归零、方差归一的分布,使深层网络的训练稳定下来。

O[c][h][w]=∑ch=1Cin∑kh=1Kh∑kw=1KwI[ch][h+kh−1][w+kw−1]×K[c][ch][kh][kw]+b[c]\displaystyle O[c][h][w] = \sum_{ch=1}^{C_{in}} \sum_{kh=1}^{K_h} \sum_{kw=1}^{K_w} I[ch][h + kh - 1][w + kw - 1] \times K[c][ch][kh][kw] + b[c]
这是 2D 卷积的完整公式:输出特征图的每个像素,等于输入通道上各窗口内像素与卷积核的逐元素乘积之和,再加偏置。
  1. 输入特征图 I 与卷积核 KI 的形状是 (C_in, H, W),K 的形状是 (C_out, C_in, K_h, K_w)。每个输出通道 c 对应一组独立的卷积核。
    I∈RCin×H×W,K∈RCout×Cin×Kh×Kw\displaystyle I \in \mathbb{R}^{C_{in} \times H \times W}, \quad K \in \mathbb{R}^{C_{out} \times C_{in} \times K_h \times K_w}
  2. 三重求和先遍历输入通道 ch,再遍历卷积核的高 kh 和宽 kw,把窗口内的每个像素与卷积核对应位置相乘并累加。
    ∑ch=1Cin∑kh=1Kh∑kw=1KwI[ch][h+kh−1][w+kw−1]⋅K[c][ch][kh][kw]\displaystyle \sum_{ch=1}^{C_{in}} \sum_{kh=1}^{K_h} \sum_{kw=1}^{K_w} I[ch][h+kh-1][w+kw-1] \cdot K[c][ch][kh][kw]
  3. 加偏置每个输出通道有一个偏置值 b[c],累加结束后加上它,得到最终输出像素。
    O[c][h][w]=acc+b[c]\displaystyle O[c][h][w] = \text{acc} + b[c]

Tensor Core加速

后厨的高效率离不开一套“秘密武器”——Tensor Core。它是一块专门为矩阵乘法设计的硬件加速器,就像餐厅引进了一台“智能切菜机器人”:传统 CUDA 核心是让一千个帮厨每人负责切一片菜叶,而 Tensor Core 则是让一个机器人一次同时处理 4×4 的菜块。官方名字叫 MMA(Matrix Multiply-Accumulate)指令,一条指令就能完成 4×4×4 的矩阵乘累加。

Tensor Core 的另一个关键设计是混合精度(mixed precision)。它用 FP16(半精度浮点数)做乘法和读取数据,用 FP32(单精度浮点数)做累加。这背后的逻辑是:AI 模型对数据的敏感度比传统的科学计算低——后厨做菜讲究“咸淡适中”,精确到小数点后 6 位反而没必要,FP16 的小数点后 3 位精度已经足够烹调出好味道。用 FP16 存储能省一半显存、带宽,用 Tensor Core 算 FP16 乘法又比 FP32 快一倍甚至更多,所以整体能实现“翻倍的出菜速度”。

Cfp32  +=  Afp16⋅Bfp16\displaystyle C_{fp32} \; += \; A_{fp16} \cdot B_{fp16}
混合精度的核心思想:用低精度做乘法以提速,用高精度累加以保精度。这就是“用便宜的碗装菜,用结实的锅炖汤”。
  1. FP16 输入A 和 B 以 FP16 格式参与乘法。每个数只占 2 字节,带宽压力减半,一次能搬运更多数据。
    Afp16,Bfp16∈R4×4\displaystyle A_{fp16}, B_{fp16} \in \mathbb{R}^{4 \times 4}
  2. FP32 累加器累加器 C 用 FP32 存储。因为加法累加会有误差累积,高精度的累加器能保证几百次累加后误差仍在可接受范围。
    Cfp32∈R4×4\displaystyle C_{fp32} \in \mathbb{R}^{4 \times 4}
  3. 一次 MMA 指令HMMA(Half-precision Matrix Multiply-Accumulate)指令在一个时钟周期内完成整个 4×4×4 运算,吞吐量是普通 CUDA 核心的 8 倍。
    HMMA.16816.F32.F16.F16.F32  Dfp32+=Afp16×Bfp16\displaystyle \text{HMMA.16816.F32.F16.F16.F32} \; D_{fp32} += A_{fp16} \times B_{fp16}

框架集成落地

老板喊一嗓子,后厨大师傅就动起来——这个“喊话”是怎么传到后厨的?答案是通过深度学习框架的底层引擎。PyTorch 和 TensorFlow 是老板(用户接口),但真正的大师傅是 cuDNN 和 cuBLAS。当你在 PyTorch 里写一行 torch.matmul(a, b) 或 self.conv1(x) 时,PyTorch 的底层调度器会先判断:这是一个 GEMM 吗?如果是,直接调用 cuBLAS 的 cublasSgemm;这是一个卷积吗?如果是,调用 cuDNN 的 cudnnConvolutionForward。你甚至可以在 PyTorch 里打印每个算子的底层调用名,会看到形如 cudnn_convolution、addmm(就是 GEMM 的 PyTorch 封装)这样的名字。

TensorFlow 的路径略有不同。除了直接调用 cuDNN 之外,TensorFlow 还有一个 XLA(Accelerated Linear Algebra)编译器,它会把多个算子融合成一个大算子,减少内核启动次数。这就像后厨的“合并订单”策略:把“切土豆+切胡萝卜+切青椒”三张小订单合并成一单“切配菜拼盘”,减少大师傅来回奔波的次数。在 TensorFlow 的 profiler 里,你会看到 Conv 和 MatMul 的执行时间占比往往超过 90%——它们就是 cuDNN 和 cuBLAS 在真正干活。

端到端的训练推理引擎,本质上就是一系列 cuBLAS/cuDNN 调用的排列组合。以训练一个 ResNet-50 为例:前向传播中的每个卷积层调用 cudnnConvolutionForward,全连接层调用 cublasSgemm;反向传播时再调用 cudnnConvolutionBackwardData 和 cudnnConvolutionBackwardFilter 来计算梯度和更新权重。整个过程就像是后厨在做流水线大餐——刀工大师傅切菜(GEMM),配菜总管按菜谱配菜(卷积),一系列操作环环相扣,最终端出训练好的模型。

Loss→forwardConv, ReLU, Pool, FC⏟cuDNN 算子 + cuBLAS GEMM→backward∇W=Convbwd+GEMMbwd\displaystyle \text{Loss} \xrightarrow{\text{forward}} \underbrace{\text{Conv, ReLU, Pool, FC}}_{\text{cuDNN 算子 + cuBLAS GEMM}} \xrightarrow{\text{backward}} \nabla W = \text{Conv}_{bwd} + \text{GEMM}_{bwd}
深度学习训练的一条完整流水线:每个算子都有前向与反向两个版本,框架自动调度 cuDNN/cuBLAS 完成全部计算。
  1. 前向传播每一层计算出输出特征图。卷积层用 cuDNN,全连接层用 cuBLAS GEMM。
    Xl+1=f(Wl∗Xl+bl)\displaystyle X_{l+1} = f(W_l * X_l + b_l)
  2. 损失函数计算模型输出与真实标签的误差。
    L=cross-entropy(y^,y)\displaystyle \mathcal{L} = \text{cross-entropy}(\hat{y}, y)
  3. 反向传播计算损失对每个权重的梯度,更新权重。cuDNN 提供 backward filter 和 backward data 两个原语分别对应权重梯度和输入梯度。
    W←W−η∇W\displaystyle W \leftarrow W - \eta \nabla W

性能调优实战

后厨的效率不是天生的,而是要调教的。cuDNN 的神奇之处在于:同一个卷积算子,它内部有多个不同的算法实现(如 implicit GEMM、FFT、Winograd),每种算法在不同输入尺寸、不同 GPU 架构上的性能差异可达数倍。cuDNN 内置了启发式(heuristic)算法选择器,根据输入形状和硬件型号“猜”一个较优的算法;但“猜”毕竟不如“跑一跑”来得准。为此 cuDNN 提供了 cudnnGetConvolutionForwardAlgorithm_v7 接口,可以遍历所有算法并对每个算法做基准测试(benchmark),找到最快的那个。

PyTorch 里对应地提供了 torch.backends.cudnn.benchmark = True 这个设置。开启后,PyTorch 会在第一次运行某个尺寸的卷积时,用一小段时间跑完所有候选算法,挑出最快的并缓存下来,之后的调用直接使用最优算法。这就像后厨新买了一口锅,配菜总管会先试一遍“大火爆炒、中火慢炖、小火煨汤”三种方式,选出做这道菜最快的那一种,然后固定下来。但要注意:如果输入尺寸频繁变化,benchmark 会反复重新寻找最优算法,反而浪费时间,此时应关闭它。

除了算法选择,内存带宽的优化同样关键。GPU 做矩阵运算时最怕的不是算不过来,而是“数据喂不上来”——后厨切菜再快,如果食材搬运工(显存总线)跟不上,厨师也只能空等。cuBLAS 的 GEMM 内部使用分块(tiling)和平铺(blocking)技术:先把大矩阵切成小块加载到共享内存,每个小块被反复使用几次后才丢弃,把对显存的访问次数降到最低。同样的道理,cuDNN 也常把多个小算子融合成一个内核(kernel fusion),避免中间结果写回显存再读出来,从而大幅减少访存开销。

Speedup≈TnaiveTtiled=DRAM visitsnaiveDRAM visitstiled≈K×M×NM×N/Ts=K×Ts\displaystyle \text{Speedup} \approx \frac{T_{\text{naive}}}{T_{\text{tiled}}} = \frac{\text{DRAM visits}_{\text{naive}}}{\text{DRAM visits}_{\text{tiled}}} \approx \frac{K \times M \times N}{M \times N / T_s} = K \times T_s
分块后的访存次数大约能减少 K 倍,T_s 是共享内存块的大小——这就是为什么同样的 GEMM 在优化前后性能可以相差一个数量级。
  1. 朴素访存每个输出元素需要读 A 一整行和 B 一整列,即 K 次访存,总共 K × M × N 次。每个 A 元素被读 N 次,每个 B 元素被读 M 次。
    visitsnaive=K×M×N\displaystyle \text{visits}_{\text{naive}} = K \times M \times N
  2. 分块访存将矩阵切成大小为 T_s 的块加载到共享内存,每个块只读一次,块内复用。访存次数降为 M × N / T_s 次全局内存读取。
    visitstiled=M×NTs\displaystyle \text{visits}_{\text{tiled}} = \frac{M \times N}{T_s}
  3. 理论加速比两者相除得到加速比,等于 K × T_s。实际还受共享内存容量和带宽限制,但方向明确——让数据复用次数最大化。
    Speedup∝K×Ts\displaystyle \text{Speedup} \propto K \times T_s

AI 应用场景

大规模图像分类训练加速

阶段:模型训练(训练引擎底层算子调度)

解决的问题:训练 ResNet-50 时,如果每个卷积层都用朴素 Python 循环实现,训练一个 epoch 需要数天;cuDNN 的卷积算子让每个卷积前向+反向在几十毫秒内完成,配合 cuBLAS 的全连接层 GEMM,训练时间从数天压缩到几小时。

具体案例:在一台 8×A100 的服务器上训练 ImageNet 分类模型。PyTorch 的 DistributedDataParallel 会把每个 batch 分到 8 张卡上,每张卡的卷积操作由 cuDNN 自动调度最优算法,全连接层由 cuBLAS GEMM 完成。开启 torch.backends.cudnn.benchmark = True 后,首次迭代会自动搜索最优卷积算法,后续迭代吞吐量提升约 30%。

推理引擎的算子融合优化

阶段:模型部署与推理优化

解决的问题:推理时延迟要求苛刻(如实时视频分析每帧 < 30ms)。TensorRT 这类推理引擎在 cuDNN 基础上做算子融合:把“卷积 + BatchNorm + ReLU”三个算子融合成一个内核,中间结果不写回显存,减少访存延迟。

具体案例:格物钛智能科技的安防摄像头端侧推理项目,将 YOLOv5s 从 PyTorch 导出为 ONNX 后,用 TensorRT 在 NVIDIA Jetson Orin 上构建引擎。TensorRT 底层调用 cuDNN 的卷积算子完成张量计算,把单帧推理延迟从 45ms 压到 18ms,同时让 GPU 利用率稳定在 85% 以上。

高性能计算中的自定义 GEMM 内核

阶段:底层算子开发与性能优化

解决的问题:cuBLAS 提供的标准 GEMM 是通用实现,但某些场景(如超大稀疏矩阵、特定分块的张量并行)需要自定义内核。理解 cuBLAS 的分块策略和 Tensor Core 指令,能指导开发者写出高效的定制内核。

具体案例:某大模型训练框架在实现张量并行时,需要在每个 Transformer 层的前向中执行两个大矩阵乘法(QKV 投影和输出投影)。开发者先用 cuBLAS cublasGemmEx 验证 FP16 混合精度的性能基线,再针对特定矩阵形状编写了融合偏置与残差的定制内核,最终将单层前向延迟降低了 22%。

Python 演示

python51 行
import torchimport torch.nn as nnimport time # 模拟一个包含卷积+全连接的小型网络# 底层实际调用 cuDNN(卷积)和 cuBLAS(GEMM/全连接)class SimpleNet(nn.Module):    def __init__(self):        super().__init__()        # cuDNN 卷积算子        self.conv1 = nn.Conv2d(3, 16, kernel_size=3, padding=1)        # cuDNN 池化算子        self.pool = nn.MaxPool2d(2)        # 底层是 cuBLAS GEMM        self.fc = nn.Linear(16 * 16 * 16, 10)     def forward(self, x):        x = torch.relu(self.conv1(x))        x = self.pool(x)        x = x.view(x.size(0), -1)        x = self.fc(x)        return x # 开启 cuDNN benchmark,让框架自动选择最快卷积算法torch.backends.cudnn.benchmark = True model = SimpleNet().cuda()optimizer = torch.optim.Adam(model.parameters(), lr=0.001)loss_fn = nn.CrossEntropyLoss() # 模拟训练迭代inputs = torch.randn(32, 3, 64, 64).cuda()labels = torch.randint(0, 10, (32,)).cuda() for epoch in range(3):    start = time.time()    for step in range(20):        optimizer.zero_grad()        outputs = model(inputs)        loss = loss_fn(outputs, labels)        loss.backward()        optimizer.step()    end = time.time()    print(f"Epoch {epoch}: loss={loss.item():.4f}, time={end - start:.3f}s") # 查看底层算子名:PyTorch 的 addmm(GEMM)与 cudnn_convolutionwith torch.autograd.profiler.profile(use_cuda=True) as prof:    outputs = model(inputs)    loss = loss_fn(outputs, labels)    loss.backward()print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=6))
目的:这段代码演示了 PyTorch 在训练小网络时如何自动调度 cuDNN(卷积、池化)与 cuBLAS(GEMM/全连接)。运行后你会看到 profiling 表中的 cudnn_convolution 和 addmm 条目,它们就是两大底层 GPU 库在幕后劳作的身影。
  1. torch.backends.cudnn.benchmark = True开启 cuDNN 的自动算法搜索,首轮迭代会做 benchmark,之后每轮都选用最快卷积实现。
  2. nn.Conv2d(3, 16, kernel_size=3, padding=1)创建一个卷积层,PyTorch 底层会调用 cudnnConvolutionForward 来执行 3×3 卷积。
  3. self.fc = nn.Linear(16 * 16 * 16, 10)全连接层,本质是 y = xW^T + b 的 GEMM 操作,PyTorch 调度到 cuBLAS 的 cublasSgemm。
  4. loss.backward()反向传播同样调用 cuDNN 的 cudnnConvolutionBackwardData/Filter 和 cuBLAS 的 cublasSgemm 来计算梯度。
  5. with torch.autograd.profiler.profile(use_cuda=True)PyTorch 内置的分析器,能显示每个底层算子的 CUDA 耗时,方便验证 cuDNN/cuBLAS 的调用情况。

常见误区

误区我以为 cuBLAS 就是 cuDNN,cuDNN 就是 cuBLAS,两个库解决的是同一个问题,只是名字不一样。

正解这两个库虽然有交集(cuDNN 内部会调用 cuBLAS 做 GEMM),但定位完全不同。cuBLAS 是 BLAS(基础线性代数库)的 GPU 版本,擅长向量、矩阵运算,是“通用数学工具”;cuDNN 是深度学习专用算子库,提供卷积、池化、归一化等高层神经网络原语,是“AI 菜系专属菜谱”。就像同一家后厨里,一位是全能刀工,一位是 AI 菜品配菜总管,两者配合但职责不同。

为何常见:因为两者的名字都以 cu 开头、都由 NVIDIA 发布、文档风格相似,且深度学习实践中两者往往同时出现,容易让人误以为是同一个东西的两个名字。

误区我听说 cuDNN 的 benchmark 模式能自动选最快算法,那把 `torch.backends.cudnn.benchmark = True` 一直开着肯定没错。

正解benchmark 模式会先以少量数据试跑所有候选算法,找到最快的再缓存。如果输入张量的形状固定,那确实值得;但如果图像尺寸频繁变化(比如处理不同分辨率的视频帧),每一轮新形状都需要重新 benchmark,额外的搜索时间可能超过节省的计算时间,反而更慢。这就像后厨每次来了不同大小的土豆都重新试一遍所有切法,浪费的时间比省下的还多。动态输入场景下建议关闭 benchmark,或用 torch.backends.cudnn.deterministic = True 保证确定性。

为何常见:官方文档和博客都在强调 benchmark 的性能收益,容易让人忽略其适应场景,形成“万能钥匙”的错觉。

误区FP16 只有大约 3 位有效十进制数字,用它训练模型,精度损失肯定很大,不能用在正式项目里。

正解混合精度的核心是“低精度计算 + 高精度累加”,乘法用 FP16 提速,累加用 FP32 保精度。更关键的是,PyTorch 的自动混合精度(AMP)机制还会对梯度做动态缩放(loss scaling),避免小梯度在 FP16 表示下被截断为 0。在绝大多数视觉和 NLP 任务上,混合精度能以几乎无损的精度换来 1.5~3 倍的速度提升。这就像后厨做菜:切菜用粗刀(FP16)快速完成,但调味(累加)时用精细天平(FP32)保证咸淡准确,最终成品质量不打折。

为何常见:人们容易孤立地看“FP16 精度低”这个事实,忽略了“混合”两者的智慧,以及 loss scaling 等辅助机制的兜底作用。

误区我直接把 C++ 里的二维数组指针传给 `cublasSgemm`,按行主序写好的数据应该没问题吧?矩阵怎么存不都是连续的吗?

正解内存布局(row-major vs column-major)搞错,结果矩阵会是正确答案的转置,甚至维度完全不匹配直接报错。cuBLAS 遵循 Fortran 时代的列主序约定:矩阵的“列”在内存中是连续的。C/C++ 常用的行主序数组直接传过去,会让 cuBLAS 把 A 看成 A^T。正确做法是设置 CUBLAS_OP_T 转置标志,或用 cublasLtMatmul 的 ld(leading dimension)参数显式描述布局。就像调料架是按“竖列”组织的,你按“横行”递给他,他只能拿到一堆错位的调料。

为何常见:BLAS 的 Fortran 历史背景对多数现代 C++ 开发者是知识盲区,且肉眼看起来“反正都是连续内存”,没有意识到“连续”的顺序方向才是关键。

误区显卡上有很多“核心”,Tensor Core 应该就是新一代的 CUDA 核心,数量越多越强。

正解两者是职责不同的硬件单元。CUDA 核心是通用的冯·诺依曼式计算单元,可以执行任意算术指令;Tensor Core 是专用的矩阵乘累加单元,只做矩阵乘法(MMA 指令),4×4 块一次完成。可以这样理解:CUDA 核心是普通厨师,能炒任何菜;Tensor Core 是专门的“宫保鸡丁机器人”,只做这一道菜但做得极快。程序中只有符合条件的 FP16/INT8 混合精度矩阵运算才能触发 Tensor Core,普通的标量操作不会用到它。

为何常见:市场上“核心数”的营销说法混在一起,Volta 之后的架构广告里 Tensor Core 和 CUDA 核心数量同时出现,容易让人误以为前者是后者的升级版。

面试题

Q1. cuBLAS 的 Level 3 例程主要处理哪种类型的运算?

单选题 · easy

  • A.向量与向量之间的运算
  • B.矩阵与向量之间的运算
  • C.矩阵与矩阵之间的乘法运算← 正确答案
  • D.标量与向量的乘法运算

答案:C。Level 1 是向量与向量(如点积、axpy),Level 2 是矩阵与向量(如 gemv),Level 3 才是矩阵与矩阵(如 GEMM)。相当于刀工大师傅的三个段位:切单根葱(L1)、配一盘菜(L2)、摆整桌宴(L3)。深度学习中最常用的 GEMM 属于 Level 3。

Q2. GEMM 公式 C = α(A × B) + βC 中,当 α=1、β=0 时表示什么含义?

单选题 · easy

  • A.只做累加,不做乘法
  • B.标准矩阵乘法 C = A × B,结果覆盖原 C← 正确答案
  • C.对 C 做缩放而不乘 A、B
  • D.转置矩阵 A 后再乘 B

答案:B。α=1、β=0 意味着输出 = 1×(A×B) + 0×C = A×B,即最朴素的标准矩阵乘法。β=0 时旧 C 被完全覆盖,β=1 时是累加模式。这就像后厨里“把炒好的菜装进新盘子”(覆盖)还是“倒在原有菜品上叠加”(累加)的区别。

Q3. cuBLAS 默认采用哪种内存布局方式读取矩阵?

单选题 · mid

  • A.行主序(row-major)
  • B.列主序(column-major)← 正确答案
  • C.对角线优先存储
  • D.哈希表随机存储

答案:B。cuBLAS 继承自 Fortran 时代的 BLAS,默认列主序。C/C++ 数组通常是行主序,直接传给 cuBLAS 会被错读成转置矩阵。解决方式是使用 CUBLAS_OP_T 转置标志或正确设置 leading dimension。这就像后厨调料架按竖列摆放,你按横排递过去就会拿错配料。

Q4. cuDNN 提供了以下哪些深度学习算子?

多选题 · mid

  • A.卷积算子(Convolution)← 正确答案
  • B.池化算子(Pooling)← 正确答案
  • C.归一化算子(Normalization)← 正确答案
  • D.内存分配算子(cudaMalloc)

答案:ABC。cuDNN 的核心全家桶包括卷积(前向+反向)、池化(最大/平均)、归一化(批归一化/层归一化)等高层算子。cudaMalloc 是 CUDA 运行时 API 提供的显存管理函数,不属于 cuDNN 的算子库。就像配菜总管负责配菜,但不负责买餐具。

Q5. Tensor Core 一次 HMMA 指令能完成多大规模的矩阵乘累加?

单选题 · mid

  • A.1×1
  • B.2×2
  • C.4×4×4← 正确答案
  • D.16×16×16

答案:C。第一代 Tensor Core 的 HMMA 指令在一个周期内完成 4×4×4 的矩阵乘累加:A(4×4) × B(4×4) 累加到 C(4×4)。它像“一块菜四刀切完”的机器人,而普通 CUDA 核心是“一刀一片”的帮厨。不同代际的 Tensor Core 形状可能不同,但经典入门级就是 4×4×4。

Q6. 以下哪个策略是 cuBLAS GEMM 分块优化(tiling)的主要目标?

单选题 · hard

  • A.减少共享内存的容量
  • B.减少对全局显存的访问次数← 正确答案
  • C.增加线程块的数量
  • D.把矩阵转置成列主序

答案:B。分块(tiling)将大矩阵切成小块加载到共享内存,块内数据被反复复用后才替换,大幅降低全局内存(DRAM)的访问次数。全局内存访存是 GPU 计算的主要瓶颈之一,减少访问次数直接提升吞吐。这就像后厨一次性把一桌菜需要的食材都备在手边,避免每次做菜都跑去仓库取。

Q7. 关于混合精度训练,以下说法正确的是?

多选题 · hard

  • A.乘法使用 FP16,累加使用 FP32← 正确答案
  • B.需要配合 loss scaling 防止梯度下溢← 正确答案
  • C.一定导致模型精度显著下降
  • D.能减少显存占用和带宽压力← 正确答案

答案:ABD。混合精度用 FP16 做乘法提速、FP32 累加保精度;loss scaling 在反向传播前放大梯度,防止小梯度在 FP16 下变为 0;FP16 占 2 字节,比 FP32 省一半显存带宽。C 选项错误——在主流视觉/NLP 任务中混合精度几乎不损失精度。这就像用粗刀快速切配、用精细天平精准调味,最终菜品质量不受影响。

Q8. 在 PyTorch 中开启 cuDNN 自动算法搜索的正确方式是?

单选题 · mid

  • A.torch.backends.cudnn.benchmark = True← 正确答案
  • B.torch.cuda.set_benchmark(True)
  • C.cudnn.enable_auto_tune()
  • D.torch.set_algorithm_search(True)

答案:A。PyTorch 通过 torch.backends.cudnn.benchmark = True 开启 cuDNN 的启发式算法搜索与缓存。首次遇到某尺寸输入时自动试跑所有卷积算法,之后复用最快方案。其他选项是编造的 API。它相当于后厨里试出最佳火候后固定下来,之后每道同尺寸菜都按这个火候做。

Q9. 以下哪个操作属于 cuDNN 池化算子的功能?

单选题 · easy

  • A.计算两个向量的点积
  • B.取滑窗内最大值或平均值← 正确答案
  • C.完成矩阵转置
  • D.生成随机数种子

答案:B。池化算子在滑窗内取最大值(最大池化)或平均值(平均池化),用于缩小特征图尺寸、保留主要信息。点积是 cuBLAS Level 1 的活儿,矩阵转置是 cuBLAS 的辅助功能,随机数生成属于 cuRAND 库。

Q10. 在动态输入尺寸的推理服务中,开启 cuDNN benchmark 可能带来什么影响?

多选题 · extreme

  • A.首次处理新形状时会额外消耗 benchmark 时间← 正确答案
  • B.每次推理都会自动选择最快的算法
  • C.可能导致 GPU 显存占用短暂增加← 正确答案
  • D.完全不影响推理延迟,多多益善

答案:AC。动态输入下,每遇到新形状都要重新试跑候选算法,产生额外延迟(A 对);benchmark 试跑需要为多个算法分配工作空间,短期显存占用上升(C 对)。它只在形状固定时选择最优算法并缓存(B 的“每次”不准确),因此 D 显然错误。这就像客人点的菜尺寸总变,厨师每次都重新试火候,反而耽误出菜。