Nvidia 算子开发工具箱

最近在做深度学习算子开发相关的工作，绕不开要理解 Nvidia 这边的算子生态。cuDNN、CUTLASS、CuTe、CUB、CCCL……一堆名词，网上的资料要么太浅（就是个名词解释），要么太深（直接给你看源码）。所以我决定从头捋一遍，搞清楚每个库要解决的问题是什么、用什么方式解决的，知其然，也知其所以然。

一、先理解根本矛盾：算力 vs 带宽

理解 GPU 软件栈，必须先理解一个根本矛盾。

现代 GPU 的计算能力增长非常快。以 H100 SXM5（Hopper 架构）为例，Tensor Core 的 FP16 稠密算力是 989 TFLOPS，但 HBM3 显存带宽只有 3.35 TB/s。

两个数字一除：

989 TFLOPS / 3.35 TB/s ≈ 295

意思是：GPU 每从显存读 1 个字节，可以做 590 次浮点运算。如果你每算一次就去显存读一次数据，GPU 里 99% 的时间都在等数据，计算单元闲置。

这个比值叫 arithmetic intensity（计算强度），是理解所有 GPU 优化的核心概念。

优化的目标，不只是减少访存次数，而是让每次读进来的数据被尽量多地复用，让计算单元一直忙着。

二、GPU 的内存层次结构

GPU 内部有一套内存层次，从快到慢（以下数据以 H100 SXM5 为参考）：

寄存器（Register File）
  └── 每个线程私有，延迟约 1 cycle，每 SM 共 64K 个 32-bit 寄存器

共享内存（Shared Memory / L1）
  └── block 内所有线程共享，延迟约 30 cycle，最大 228 KB/SM

L2 Cache
  └── 全 GPU 共享，延迟约 200 cycle，容量 50 MB

全局内存（Global Memory / HBM）
  └── 显存，延迟约 500 cycle，容量 80 GB

离计算单元越近，越快，但越小。

关键直觉：Tensor Core 算完一次 16×16 矩阵乘只需要约 16 个 cycle，但去 global memory 取一次数据要等约 500 个 cycle。如果每算一次就取一次数据，Tensor Core 几乎一直在空转。

Shared memory 是解法——提前把数据搬进来，放在离计算单元只有 5 个 cycle 的地方，让它连续不断地计算。

三、tile 思想：数据复用的核心手段

先看不用 tile 会怎样

矩阵乘法 C = A × B，最朴素的写法是：每个线程负责计算 C 的一个元素，每次需要的数据都直接去 global memory 读。

__global__ void naive_matmul(float* A, 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;

    float sum = 0;
    for (int k = 0; k < K; k++) {
        sum += A[row * K + k]    // 每次去 global memory 读
             * B[k   * N + col]; // 每次去 global memory 读
    }
    C[row * N + col] = sum;
}

问题出在哪里？A 的同一行数据，会被 N 列的线程各读一遍。B 的同一列数据，会被 M 行的线程各读一遍。同一份数据被重复从 global memory 读了 N 次或 M 次，每次都要等 800 cycle，带宽全部浪费在重复搬运上。

tile 的解法

核心思想：一组线程协作，把一小块数据（tile）搬进 shared memory，然后大家反复读这块数据，让每次 global memory 读取被复用很多次。

#define TILE 32

__global__ void tiled_matmul(float* A, float* B, float* C, int N) {
    __shared__ float As[TILE][TILE]; // 从 A 搬进来的 tile
    __shared__ float Bs[TILE][TILE]; // 从 B 搬进来的 tile

    int row = blockIdx.y * TILE + threadIdx.y;
    int col = blockIdx.x * TILE + threadIdx.x;
    float sum = 0;

    for (int t = 0; t < N/TILE; t++) {
        // 第一步：每个线程搬一个元素，协作完成整块 tile 的搬运
        As[threadIdx.y][threadIdx.x] = A[row * N + (t*TILE + threadIdx.x)];
        Bs[threadIdx.y][threadIdx.x] = B[(t*TILE + threadIdx.y) * N + col];

        __syncthreads(); // 等所有线程搬完

        // 第二步：从 shared memory 计算，延迟只有 5 cycle
        for (int k = 0; k < TILE; k++)
            sum += As[threadIdx.y][k] * Bs[k][threadIdx.x];

        __syncthreads(); // 用完再同步，才能进下一轮
    }
    C[row * N + col] = sum;
}

效果：同样是 1024×1024 的矩阵乘，TILE=32 时，每个数据被复用 32 次，global memory 流量从 8GB 降到 256MB。Tensor Core 利用率从不到 1% 提升到接近 80%。

Grid / Block / Thread 三层结构

写 CUDA kernel 之前还需要理解 GPU 的线程组织方式，三层从大到小：

• Grid：启动一个 kernel 时，所有 block 的集合，代表整个任务
• Block：一组线程，共享同一块 shared memory
• Thread：最小单位，每个线程有自己的寄存器

dim3 gridDim(4, 4);   // Grid 里有 4×4 = 16 个 block
dim3 blockDim(16, 16); // 每个 block 里有 16×16 = 256 个线程
kernel<<<gridDim, blockDim>>>(args);

每个线程通过坐标确定自己负责哪个数据：

int row = blockIdx.y * blockDim.y + threadIdx.y;
int col = blockIdx.x * blockDim.x + threadIdx.x;

公式逻辑和”电影院座位编号”完全一样：排号 × 每排座位数 + 列号。

四、Tensor Core：专门做矩阵乘法的硬件单元

普通 CUDA Core 每个 cycle 做一次乘加（a × b + c）。Tensor Core 是一块专用硬件电路，一条指令让 32 个线程协作，一次完成 16×16×16 的矩阵乘法，等效于 4096 次乘加，速度大约是 CUDA Core 的 8~16 倍。

但 Tensor Core 有个硬性要求：数据必须按特定格式分布在参与计算的线程的寄存器里，才能一次性读入做计算。这个要求从 Volta 架构引入 Tensor Core 起就存在，但每代细节不同——参与协作的线程数从 Volta 的 8 个，到 Ampere 的 32 个（一个 warp），再到 Hopper 的 128 个（一个 warp group），指令名也从 mma.sync 变成了 wgmma。

问题在于：shared memory 里的数据是按行列顺序存的，不符合 Tensor Core 要求的格式。中间需要一个”重排”步骤，用 ldmatrix 指令把数据从 shared memory 搬到寄存器时，按 Tensor Core 要求的方式摆好。

这个”怎么摆”的规则，就是 CuTe 要解决的核心问题。

五、各个库存在的意义

有了上面的背景，现在来看每个库封装的是什么。

CUB（属于 CCCL）

CUB 解决的是 kernel 内部最常见的套路：block 内的规约（reduce）、扫描（scan）、排序。

比如 block 内 32 个线程各持有一个数，要算总和，手写需要 tree reduction，大约 10 行代码，还要仔细处理同步点。用 CUB：

using BlockReduce = cub::BlockReduce<float, 32>;
__shared__ typename BlockReduce::TempStorage temp;

float total = BlockReduce(temp).Sum(my_value); // 一行

CUB 属于 CCCL（CUDA Core Compute Libraries），CCCL 还包括 Thrust（GPU 版 STL）和 libcu++（GPU 上的 C++ 标准库），但对 kernel 开发来说 CUB 是最直接相关的部分。

CuTe

CuTe 解决的是描述”数据排布”的问题。

核心抽象是 Layout = (Shape, Stride)：

• Shape：矩阵的形状，比如 (2, 4) 表示 2 行 4 列
• Stride：步长，每个维度走一步，地址跳多少

Layout = (Shape(2,4), Stride(4,1))

取 A[row][col] 的地址 = row*4 + col*1

这和手写 A[row * 4 + col] 完全等价，只是现在变成了数据结构，可以被代码操作和组合。

好处在于：换一个 Stride，同一块内存就有了不同的”视角”——行优先、列优先、Tensor Core 要求的特殊格式，都只是不同的 Stride，不需要搬数据。换代 GPU 时，Tensor Core 要求的 Layout 变了，只改 Stride 的数值，代码逻辑不用动。

CUTLASS

CUTLASS 把整个 tiled matmul 的流程模板化：tile 循环、搬运、Tensor Core 计算、写回，全部封装。你只需要指定可变的部分：

using Gemm = cutlass::gemm::device::Gemm<
    float, ColumnMajor,    // A 的类型和排布
    float, RowMajor,       // B 的类型和排布
    float, RowMajor,       // C 的类型和排布
    Shape<128, 128, 32>,   // tile 大小
    Mma,                   // 计算指令（自动匹配当前 GPU）
    EpilogueReLU           // 算完之后做什么
>;

Gemm gemm;
gemm(args);

换 GPU 只改 Mma 那一行，换 tile size 只改 Shape，加自定义 epilogue 只改最后一行。手写需要的几百行地址计算和流水线编排，全部不用碰。

cuDNN

cuDNN 是最高层的封装，面向框架（PyTorch、TensorFlow）开发者。你只需要描述”我要做什么”，不需要描述”怎么做”：

cudnnConvolutionForward(handle, plan, x, w, y);

cuDNN 内部自己选算法（implicit GEMM？Winograd？FFT？）、选 tile size、选指令。代价是灵活性低——不支持自定义算子组合，Conv 之后想加一个特殊激活函数，就只能改用 CUTLASS。

六、总结

GPU 硬件     → 算力快，内存慢，需要 tile + shared memory 解决
裸 CUDA      → 你自己管一切，灵活但痛苦
CUB          → block 内规约套路，不用手写了
CuTe         → 用 (Shape, Stride) 描述数据排布，换 GPU 不改逻辑
CUTLASS      → 整个 tile 流水线模板化，只填可变参数
cuDNN        → 只说"我要 Conv"，其余全不用管

每一层库封装的，都是”如果自己写，会写几百行且换卡就报废”的部分。抽象层次越高，灵活度越低，但开发成本也越低。用哪层取决于你的需求：标准算子用 cuDNN，自定义融合算子用 CUTLASS，kernel 内部的套路用 CUB。

附：和华为昇腾软件栈的对比

理解了 Nvidia 这套分层之后，再看华为昇腾的软件栈会清晰很多。华为的整体架构叫 CANN（Compute Architecture for Neural Networks），是对标 CUDA + cuDNN 的全栈软件层。两边的对应关系大致如下：

几个值得注意的地方：

CATLASS 对应 CUTLASS，是华为开源的矩阵计算模板库，配合 Ascend C 使用，解决的是同一个问题：把 tile 循环、数据搬运、流水线这些固定套路封装掉，让开发者只关注算子逻辑本身。仓库已在 Gitee/GitCode 开源，可以直接看源码。

Ascend C 对应的层级比 CUTLASS 更低一些，更接近裸 CUDA kernel 的位置——直接操作片上内存（Local Memory）和计算单元（Cube/Vector Core）。CATLASS 是构建在 Ascend C 之上的模板层，两者的关系类似 CUTLASS 和裸 CUDA 的关系。

ACLNN 对应 cuDNN，是面向框架开发者的高层算子库。PyTorch 在昇腾上跑，底层算子基本都走 ACLNN。

早期的 TBE（Tensor Boost Engine） 是华为用 Python DSL 写自定义算子的方案，存在性能天花板和调试困难的问题。2023 年 CANN 7.0 之后华为推出 Ascend C，用 C++ 替代了 TBE，开发体验和性能都有明显提升。

总体来说，Nvidia 这套分层更细、更模块化，每一层的边界很清晰。华为这边的分层相对粗一些，但对于研究竞品来说，用 Nvidia 的框架去对应华为的每个组件，是一个很好的快速定位方式。