CUDA Kernel 调优实战：从 Naive 矩阵乘法到 Tensor Core 加速的 18 倍提升

一、GPU 算力未被充分利用的"暗黑物质"

写一个 GPU 上的矩阵乘法只需要 20 行 CUDA 代码，但要让这段代码跑出接近硬件理论峰值的性能，则需要深入理解 SM（Streaming Multiprocessor）的调度机制、Shared Memory 的 Bank Conflict、以及 Tensor Core 的矩阵分块策略。实测中，一个 Naive 实现的 INT8 矩阵乘法（M=N=K=4096）在 A100 上仅能达到 8.2 TFLOPS，而经过 Tensor Core 优化后的实现可以达到 148 TFLOPS——18 倍的性能差距。

这 18 倍的差距就是被浪费的"暗黑物质"——GPU 的算力确实在那里，但大部分 CUDA Kernel 因为访存模式和线程调度效率低，无法真正触达硬件上限。就像羽毛球比赛中，有力气不等于能把力量有效地传递到球上——握拍角度、挥拍轨迹、击球时机缺一不可。

二、从 SM 到 Tensor Core 的计算层次

flowchart TD subgraph GPU[GPU 架构层次] A[Grid<br/>整体计算任务] --> B1[Block 0<br/>映射到 SM 0] A --> B2[Block 1<br/>映射到 SM 1] A --> B3[Block N<br/>映射到 SM N] B1 --> C1[Warp 0..31<br/>32 线程一组] C1 --> D1[Tensor Core<br/>单周期 4×4×4 MMA] B1 --> E1[Shared Memory<br/>每个 SM 最大 164KB] B1 --> F1[Register File<br/>每个 SM 65536 × 32bit] end style D1 fill:#e74c3c,color:#fff style E1 fill:#e67e22,color:#fff

CUDA 的执行模型从粗到细分为 Grid → Block → Warp → Thread。性能优化的核心在于最大化 SM 的占用率（Occupancy）和最小化全局内存（Global Memory）访问。Tensor Core 是 Volta 架构引入的专用矩阵乘加单元，通过mma.sync指令可在单个时钟周期内完成 4×4×4 的矩阵乘累加。

三、逐步优化：从 Naive 到 Tensor Core

// matmul_tensorcore.cu —— 逐步优化的矩阵乘法 CUDA Kernel // ===== Level 0: Naive 实现（仅返回正确结果） ===== // 问题：每个线程从 Global Memory 读取 2K 个元素 // Global Memory 带宽 (~1.5TB/s) 是主要瓶颈 __global__ void matmul_naive( 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; // 每个线程迭代 K 次，每次 2 次 Global Memory 读取 for (int k = 0; k < K; k++) { sum += A[row * K + k] * B[k * N + col]; } C[row * N + col] = sum; } } // 性能：M=N=K=4096 → ~8.2 TFLOPS (A100 FP32) // ===== Level 1: Shared Memory Tiling ===== // 将 A 和 B 的分块缓存到 Shared Memory 中 // 每个线程块处理一个 TILE_SIZE × TILE_SIZE 的输出子矩阵 #define TILE_SIZE 32 __global__ void matmul_tiled( const float* A, const float* B, float* C, int M, int N, int K ) { // Shared Memory 分配：每个 Block 的两块输入 Tile __shared__ float As[TILE_SIZE][TILE_SIZE]; __shared__ float Bs[TILE_SIZE][TILE_SIZE]; int bx = blockIdx.x, by = blockIdx.y; int tx = threadIdx.x, ty = threadIdx.y; int row = by * TILE_SIZE + ty; int col = bx * TILE_SIZE + tx; float sum = 0.0f; // 循环遍历 K 维度，每次处理一个 TILE_SIZE 的切片 for (int tile = 0; tile < (K + TILE_SIZE - 1) / TILE_SIZE; tile++) { // 协作加载 A 和 B 的分块到 Shared Memory // 每个线程加载一个元素，128 个线程并发加载 int a_col = tile * TILE_SIZE + tx; if (row < M && a_col < K) As[ty][tx] = A[row * K + a_col]; else As[ty][tx] = 0.0f; int b_row = tile * TILE_SIZE + ty; if (b_row < K && col < N) Bs[ty][tx] = B[b_row * N + col]; else Bs[ty][tx] = 0.0f; __syncthreads(); // 确保所有线程加载完成 // 在 Shared Memory 上做内积计算 // Shared Memory 带宽 ~12TB/s（vs Global ~1.5TB/s） for (int k = 0; k < TILE_SIZE; k++) { sum += As[ty][k] * Bs[k][tx]; } __syncthreads(); // 确保计算完成后再加载下一批 } if (row < M && col < N) { C[row * N + col] = sum; } } // 性能：~45 TFLOPS (5.5× 提升) // ===== Level 2: Tensor Core MMA 指令 ===== // 使用 wmma (Warp Matrix Multiply-Accumulate) API // Tensor Core 以 Warp 为单位执行 16×16×16 的矩阵乘加 #include <cuda_fp16.h> #include <mma.h> using namespace nvcuda; #define WMMA_M 16 #define WMMA_N 16 #define WMMA_K 16 __global__ void matmul_tensorcore( const half* A, const half* B, float* C, int M, int N, int K ) { // 声明 Tensor Core 的 Fragment（存储在寄存器中） wmma::fragment<wmma::matrix_a, WMMA_M, WMMA_N, WMMA_K, half, wmma::row_major> a_frag; wmma::fragment<wmma::matrix_b, WMMA_M, WMMA_N, WMMA_K, half, wmma::col_major> b_frag; wmma::fragment<wmma::accumulator, WMMA_M, WMMA_N, WMMA_K, float> c_frag; // 初始化累加器为零 wmma::fill_fragment(c_frag, 0.0f); int warpM = (blockIdx.x * blockDim.x + threadIdx.x) / 32; int warp_row = warpM * WMMA_M; int warp_col = blockIdx.y * WMMA_N; // K 维度循环，每次处理 WMMA_K=16 的切片 for (int k = 0; k < K; k += WMMA_K) { // 从 Global Memory 加载矩阵分块到 Fragment // 利用 Tensor Core 的硬件矩阵加载指令 wmma::load_matrix_sync(a_frag, A + warp_row * K + k, K); wmma::load_matrix_sync(b_frag, B + k * N + warp_col, N); // 单条指令执行 16×16×16 的矩阵乘累加 // 硬件在每个时钟周期处理 4×4×4 的子块 wmma::mma_sync(c_frag, a_frag, b_frag, c_frag); } // 将结果写回 Global Memory wmma::store_matrix_sync(C + warp_row * N + warp_col, c_frag, N, wmma::mem_row_major); } // 性能：~148 TFLOPS (18× vs Naive, 3.3× vs Tiled) // ===== 辅助函数：FP32 → FP16 转换 ===== // Tensor Core 原生支持 FP16 输入、FP32 累加 // 精度损失 < 0.1%，对于推理场景完全可接受

优化效果汇总

四、何时不需要 Tensor Core 优化

推理 vs 训练：Tensor Core 优化在推理场景中的收益取决于矩阵规模。对于 Batch Size=1 的推理，矩阵乘法可能不是主要瓶颈（瓶颈在内存加载模型权重），Tensor Core 优化的收益有限。

FP32 精度要求：Tensor Core 的 FP16 计算在累加时使用 FP32，但输入精度受限于 FP16。对于对数值精度极度敏感的场景（如科学计算），需评估 FP16 的精度损失。

开发成本：从 Naive 到 Tensor Core 优化需要数天甚至数周的调优时间。如果只是将 30 秒的推理优化到 10 秒，可以考虑降低优先度。

五、总结

CUDA Kernel 优化的核心路径是：Shared Memory Tiling（减少 Global Memory 访问）→ Bank Conflict 消除（提升 SM 内部带宽）→ Tensor Core（利用专用硬件）。18 倍的性能差距说明：GPU 的算力上限很高，但触达上限需要精心的 Kernel 设计和调优。

优化路线建议：

1. 先用ncu（Nsight Compute）分析 Kernel 的瓶颈（Memory Bound vs Compute Bound）；
2. Memory Bound → 引入 Shared Memory Tiling；
3. Compute Bound → 考虑 Tensor Core / Warp-level 优化；
4. 使用nvprof/nsys验证每一步优化的实际收益。