★ FEATURED ARTICLE 1. 从DeepGEMM看矩阵乘法优化的核心思路1.1 这个项目到底在解决什么问题矩阵乘法说白了就是两个矩阵之间做乘加运算是深度学习和科学计算里最基础也最耗算力的操作之一。你训练一个模型前向传播、反向传播里面大量的计算都是矩阵乘法。谁把矩阵乘法做得更快谁就能在同样的硬件上跑更大的模型、更短的训练时间。DeepGEMM这个项目从名字就能看出来它聚焦的是深度学习场景下的通用矩阵乘法GEMM优化。GEMM是BLAS库里的一个经典操作全称General Matrix Multiply就是C alpha * A * B beta * C这种形式。深度学习里的全连接层、卷积层通过im2col转换、注意力机制里的QKV计算底层全都是GEMM。那为什么还要专门做一个DeepGEMM因为通用的BLAS库比如cuBLAS虽然优化得很好但它是面向所有场景的不会针对深度学习的特定矩阵形状、特定数据类型做极致优化。深度学习里的矩阵乘法有很多特殊之处比如矩阵形状往往比较规整batch size、hidden dim这些通常是2的幂次或者比较整的数比如经常用低精度FP16、BF16、FP8比如经常需要和bias加法、激活函数融合在一起做。这些特殊性就意味着有优化空间。DeepGEMM要做的就是针对深度学习的典型场景把GEMM的性能压榨到极致。它不是一个通用库而是一个专用库专门为深度学习里的矩阵乘法服务。1.2 为什么矩阵乘法优化这么重要我给大家算一笔账。假设你有一个BERT-base模型12层hidden size 768sequence length 512batch size 32。这里面有多少矩阵乘法每一层的自注意力有Q、K、V、O四个投影矩阵每个都是768x768输入是32x512x768。前馈网络有两个线性层768x3072和3072x768。光是一层矩阵乘法的浮点运算量就是QKV投影3 * 2 * 32 * 512 * 768 * 768 ≈ 58 GFLOPs注意力得分2 * 32 * 12 * 512 * 512 * 64 ≈ 12.9 GFLOPs输出投影2 * 32 * 512 * 768 * 768 ≈ 19.3 GFLOPsFFN第一层2 * 32 * 512 * 768 * 3072 ≈ 77.3 GFLOPsFFN第二层2 * 32 * 512 * 3072 * 768 ≈ 77.3 GFLOPs一层加起来大概245 GFLOPs12层就是接近3 TFLOPs。这还只是一个batch的前向传播。如果算上反向传播计算量还要翻倍。训练的时候要跑几十万甚至几百万个batch你算算总共需要多少算力。所以矩阵乘法的性能直接决定了模型训练和推理的速度。你把GEMM优化10%整个训练时间就能缩短接近10%。这就是为什么大厂都在拼命优化GEMM也是为什么DeepGEMM这种项目有价值。1.3 这个项目的技术路线选择DeepGEMM的技术路线我理解下来核心是几个点第一用CUDA C直接写kernel不用cuBLAS。这听起来很疯狂因为cuBLAS是NVIDIA官方优化了多少年的库你凭什么比它快但实际情况是cuBLAS要兼顾所有场景它的启发式算法选择不一定对你的特定形状最优。如果你知道你的矩阵形状就是固定的那么几种你完全可以针对这些形状手写kernel把性能做到比cuBLAS更好。第二用CUTLASS作为底层模板库。CUTLASS是NVIDIA开源的CUDA模板库专门用来写高性能GEMM kernel。它提供了很多现成的组件比如tile loading、warp-level MMA、epilogue fusion等等。DeepGEMM大概率是基于CUTLASS做的二次开发针对深度学习的场景做了定制。第三针对低精度做优化。现在深度学习训练和推理大量使用FP16、BF16甚至FP8。低精度矩阵乘法有特殊的优化技巧比如用Tensor Core做混合精度计算比如用更小的数据类型减少内存带宽压力。DeepGEMM应该在这些方面做了不少工作。第四kernel fusion。深度学习里的GEMM往往不是孤立的后面经常跟着bias加法、激活函数、dropout等等。如果每个操作都单独写一个kernelkernel launch的开销和中间结果的读写会浪费很多时间。DeepGEMM应该支持把GEMM和后续操作融合到一个kernel里减少内存访问。提示如果你之前没接触过GEMM优化建议先了解一下CUDA编程基础和Tensor Core的工作原理不然看DeepGEMM的代码会比较吃力。2. 核心细节解析与实操要点2.1 GEMM kernel的基本结构一个高性能的GEMM kernel结构上大概分成几个部分首先是tile划分。把输出矩阵C划分成若干个tile每个tile由一个thread block负责计算。比如输出是MxN的矩阵你可以划分成128x128的tile每个block算一个tile。tile的大小选择很关键太大了寄存器不够用太小了并行度不够。然后是数据加载。每个block需要加载A矩阵的一行tile和B矩阵的一列tile。这些数据从global memory加载到shared memory然后再从shared memory加载到寄存器。为什么要经过shared memory因为global memory的延迟很高直接访问效率低。shared memory的延迟低很多而且可以被同一个block里的所有线程共享。接着是计算。每个线程负责计算输出tile里的一个或几个元素。计算的时候需要把A和B的对应元素乘积累加起来。如果用Tensor Core就是调用warp-level的MMA指令一次算一个16x16x16或者更大的矩阵乘法。最后是写回。把计算结果从寄存器写回global memory。如果支持epilogue fusion还会在这里做bias加法、激活函数等操作。2.2 关键参数的选择与计算GEMM kernel的性能对参数非常敏感。我列几个最关键的参数以及怎么选参数含义典型值选择依据BLOCK_M输出tile的M维度大小128根据矩阵M大小和SM数量调整BLOCK_N输出tile的N维度大小128根据矩阵N大小和SM数量调整BLOCK_KK维度的分块大小32根据shared memory大小和寄存器压力调整THREADS每个block的线程数256通常是128或256太少并行度不够太多调度开销大STAGES流水线级数3-4根据shared memory大小和延迟隐藏需求调整这些参数怎么算我举个例子。假设你用A100 GPU有108个SM每个SM有192KB shared memory。你要算一个4096x4096x4096的GEMM。如果BLOCK_M128BLOCK_N128那么输出矩阵被划分成32x321024个tile。108个SM每个SM要算大概9.5个tile。这个并行度是够的。每个block需要加载A的128xBLOCK_K和B的BLOCK_Kx128的数据。如果BLOCK_K32那么A tile是128x32B tile是32x128。用FP16存储每个元素2字节那么A tile是8KBB tile是8KB。如果STAGES3那么需要(88)*348KB的shared memory。A100有192KB够用。寄存器方面每个线程要负责计算128x128/25664个输出元素。如果用FP32累加需要64个寄存器。加上其他临时变量大概需要100-150个寄存器。A100每个线程最多255个寄存器够用。注意这些参数不是固定的需要根据具体的矩阵形状和GPU型号做调优。我见过有人为了一个特定的矩阵形状把参数调了几十遍才找到最优组合。2.3 低精度矩阵乘法的特殊处理低精度矩阵乘法有几个坑要注意第一个是累加精度。用FP16做输入但累加的时候要用FP32。为什么因为FP16的表示范围有限如果累加也用FP16很容易溢出或者精度损失太大。Tensor Core的MMA指令默认就是FP16输入、FP32累加这个设计是有道理的。第二个是数据布局。Tensor Core对数据布局有要求。比如warp-level的MMA指令要求A矩阵是row-majorB矩阵是col-major或者有特定的swizzle模式。如果你从global memory加载的数据布局不对需要先做转换这会增加开销。第三个是量化误差。如果用FP8量化误差会更大。需要在训练的时候做量化感知训练QAT或者用一些补偿技术。推理的时候如果模型本身是用FP16训练的直接转FP8可能会有明显的精度下降。第四个是混合精度。有时候A用FP16B用FP8或者反过来。这种混合精度的矩阵乘法Tensor Core不一定直接支持需要做转换或者用多个指令组合。2.4 Kernel fusion的实现方式Kernel fusion是DeepGEMM的一个重要优化点。我详细说一下怎么实现假设你要算Y activation(A * B bias)。如果不做fusion你需要调用GEMM kernel算A*B结果写到临时buffer调用add kernel算A*B bias结果写到另一个buffer调用activation kernel算activation(...)结果写到Y这样做的问题是什么每次都要读写global memory。A*B的结果可能很大比如4096x4096的FP32矩阵就是64MB。读写三次就是384MB的内存流量。而A100的显存带宽大概是1.5TB/s384MB需要0.25ms。如果GEMM本身只需要1ms那这0.25ms就是25%的额外开销。如果做fusion在GEMM kernel的epilogue阶段直接做bias加法和activation就只需要写一次Y内存流量减少到64MB开销降到0.04ms。这就是fusion的价值。实现fusion的关键是在GEMM kernel计算完一个tile的结果后不要直接写回global memory而是先在寄存器里做bias加法和activation然后再写回。bias可以从global memory加载或者作为kernel参数传进来。activation函数可以直接在kernel里实现比如ReLU就是max(0, x)GELU稍微复杂一点但也可以实现。提示不是所有操作都适合fusion。如果activation函数特别复杂或者需要访问额外的数据fusion可能会增加寄存器压力反而降低性能。需要根据具体情况权衡。3. 实操过程与核心环节实现3.1 环境准备与依赖安装要复现DeepGEMM或者自己写一个类似的GEMM优化项目你需要准备以下环境硬件要求NVIDIA GPU建议算力7.0以上Volta架构及以上因为需要Tensor Core显存至少8GB建议16GB以上方便测试大矩阵如果要做多GPU测试需要NVLink或者PCIe 4.0以上软件要求CUDA Toolkit 11.0以上建议11.8或12.xCUTLASS 2.x或3.x建议用最新版本CMake 3.18以上C编译器建议GCC 9以上或Clang 10以上Python 3.8以上用于跑测试脚本安装步骤大概是这样# 安装CUDA Toolkit如果还没装 # 从NVIDIA官网下载对应版本的runfile然后执行 sudo sh cuda_12.1.0_530.30.02_linux.run # 克隆CUTLASS git clone https://github.com/NVIDIA/cutlass.git cd cutlass git checkout v3.0.0 # 或者最新版本 # 编译CUTLASS的示例确认环境没问题 mkdir build cd build cmake .. -DCUTLASS_NVCC_ARCHS80 # 根据你的GPU算力调整 make -j$(nproc)验证环境# 检查CUDA是否安装成功 nvcc --version # 检查GPU信息 nvidia-smi # 跑一个CUTLASS的示例 ./examples/00_basic_gemm/00_basic_gemm如果这些都没问题环境就准备好了。3.2 一个最简单的GEMM kernel实现我先带大家写一个最基础的GEMM kernel不用Tensor Core就用普通的CUDA核心。这样你能理解GEMM的基本原理然后再看DeepGEMM的优化就有感觉了。// 最简单的GEMM kernelC A * B // A: M x K, B: K x N, C: M x N __global__ void simple_gemm(const float* A, const float* B, float* C, int M, int N, int K) { // 每个线程计算C中的一个元素 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 k 0; k K; k) { sum A[row * K k] * B[k * N col]; } C[row * N col] sum; } }这个kernel的问题很明显每个线程都要从global memory读取A的一行和B的一列内存访问量是O(MNK)而计算量也是O(MNK)算术强度是1 FLOP/byte远低于A100的峰值大概15 FLOP/byte。所以这个kernel的性能很差大概只有峰值性能的1%不到。优化第一步用shared memory。#define BLOCK_SIZE 16 __global__ void shared_gemm(const float* A, const float* B, float* C, int M, int N, int K) { __shared__ float As[BLOCK_SIZE][BLOCK_SIZE]; __shared__ float Bs[BLOCK_SIZE][BLOCK_SIZE]; int row blockIdx.y * BLOCK_SIZE threadIdx.y; int col blockIdx.x * BLOCK_SIZE threadIdx.x; float sum 0.0f; // 分块加载 for (int k 0; k K; k BLOCK_SIZE) { // 加载A的tile if (row M k threadIdx.x K) As[threadIdx.y][threadIdx.x] A[row * K k threadIdx.x]; else As[threadIdx.y][threadIdx.x] 0.0f; // 加载B的tile if (k threadIdx.y K col N) Bs[threadIdx.y][threadIdx.x] B[(k threadIdx.y) * N col]; else Bs[threadIdx.y][threadIdx.x] 0.0f; __syncthreads(); // 计算 for (int i 0; i BLOCK_SIZE; i) { sum As[threadIdx.y][i] * Bs[i][threadIdx.x]; } __syncthreads(); } if (row M col N) C[row * N col] sum; }这个版本用了shared memory做tiling每个block加载一次A和B的tile然后复用BLOCK_SIZE次。算术强度提升到BLOCK_SIZE FLOP/byte如果BLOCK_SIZE16就是16 FLOP/byte比之前好很多。但16还是太小性能大概能到峰值的10-20%。优化第二步用Tensor Core。// 用warp-level MMA指令 // 需要包含mma.h #include mma.h using namespace nvcuda; __global__ void tensor_core_gemm(const half* A, const half* B, float* C, int M, int N, int K) { // 每个warp负责一个16x16的输出tile wmma::fragmentwmma::matrix_a, 16, 16, 16, half, wmma::row_major a_frag; wmma::fragmentwmma::matrix_b, 16, 16, 16, half, wmma::col_major b_frag; wmma::fragmentwmma::accumulator, 16, 16, 16, float c_frag; wmma::fill_fragment(c_frag, 0.0f); int warp_row (blockIdx.y * blockDim.y threadIdx.y) / 32; int warp_col (blockIdx.x * blockDim.x threadIdx.x) / 32; for (int k 0; k K; k 16) { wmma::load_matrix_sync(a_frag, A warp_row * 16 * K k, K); wmma::load_matrix_sync(b_frag, B k * N warp_col * 16, N); wmma::mma_sync(c_frag, a_frag, b_frag, c_frag); } wmma::store_matrix_sync(C warp_row * 16 * N warp_col * 16, c_frag, N, wmma::mem_row_major); }这个版本用了Tensor Core的WMMA API性能能到峰值的50-70%。但WMMA API比较高层不够灵活而且有额外的开销。真正的高性能GEMM会用更底层的MMA指令比如PTX里的mma.sync.aligned.m16n8k16。3.3 用CUTLASS写一个高性能GEMMCUTLASS把很多底层细节封装好了用它可以比较方便地写出高性能GEMM。我给大家看一个例子#include cutlass/gemm/device/gemm.h // 定义GEMM的类型 using Gemm cutlass::gemm::device::Gemm cutlass::half_t, // A的数据类型 cutlass::layout::RowMajor, // A的布局 cutlass::half_t, // B的数据类型 cutlass::layout::ColumnMajor, // B的布局 float, // C的数据类型 cutlass::layout::RowMajor, // C的布局 float, // 累加类型 cutlass::arch::OpClassTensorOp, // 用Tensor Core cutlass::arch::Sm80, // 目标架构 cutlass::gemm::GemmShape128, 128, 32, // Threadblock tile cutlass::gemm::GemmShape64, 64, 32, // Warp tile cutlass::gemm::GemmShape16, 8, 16, // MMA instruction shape cutlass::epilogue::thread::LinearCombination float, 1, float, float, // Epilogue cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle, // Swizzle 3 // Stages ; // 使用 Gemm gemm_op; cutlass::gemm::GemmCoord problem_size(M, N, K); Gemm::Arguments args(problem_size, {A, K}, // A的指针和leading dimension {B, N}, // B的指针和leading dimension {C, N}, // C的指针和leading dimension {C, N}, // D的指针和leading dimension {1.0f, 0.0f}); // alpha, beta gemm_op(args);这个代码看起来简单但CUTLASS在背后做了很多优化自动选择最优的tile大小、自动做shared memory swizzle、自动用cp.async做异步拷贝、自动做register blocking等等。性能通常能到峰值的80-90%。3.4 性能测试与调优写完kernel之后怎么知道性能好不好你需要做benchmark。测试方法// 用CUDA Event计时 cudaEvent_t start, stop; cudaEventCreate(start); cudaEventCreate(stop); // 预热 for (int i 0; i 10; i) { gemm_op(args); } cudaDeviceSynchronize(); // 计时 cudaEventRecord(start); for (int i 0; i 100; i) { gemm_op(args); } cudaEventRecord(stop); cudaEventSynchronize(stop); float milliseconds 0; cudaEventElapsedTime(milliseconds, start, stop); float avg_time milliseconds / 100; // 计算TFLOPS double flops 2.0 * M * N * K; double tflops flops / (avg_time * 1e-3) / 1e12; printf(Performance: %.2f TFLOPS\n, tflops);调优策略问题可能原因解决方法性能远低于峰值tile太小增大BLOCK_M/BLOCK_N性能远低于峰值寄存器溢出减小BLOCK_M/BLOCK_N或BLOCK_K性能远低于峰值shared memory不够减小STAGES或BLOCK_K性能波动大没有预热增加预热次数性能波动大GPU被其他进程占用检查nvidia-smi大矩阵性能好小矩阵差并行度不够用split-K或stream-K小矩阵性能好大矩阵差内存带宽瓶颈用更大的tile或更好的swizzle提示调优是一个迭代的过程。我建议你先用CUTLASS的profiler工具跑一遍看看它推荐什么参数然后在这个基础上微调。不要从零开始瞎试效率太低。4. 常见问题与排查技巧实录4.1 编译和链接问题问题一找不到CUTLASS头文件。这个通常是因为CMake没有正确配置CUTLASS的路径。解决方法是在CMakeLists.txt里加上set(CUTLASS_DIR /path/to/cutlass) include_directories(${CUTLASS_DIR}/include)或者用find_packagefind_package(CUTLASS REQUIRED) target_link_libraries(your_target PRIVATE CUTLASS::CUTLASS)问题二nvcc编译报错说某个模板参数不对。CUTLASS的模板参数很多而且不同版本的API可能有变化。我建议你直接看CUTLASS自带的example找到和你需求最接近的那个然后改。不要自己从头写容易踩坑。问题三链接的时候报undefined reference。这个通常是因为你用了CUTLASS的某些功能但没有链接对应的库。CUTLASS大部分是header-only的但有些功能比如profiler需要链接库。检查一下你的target_link_libraries有没有漏掉东西。4.2 运行时错误问题一CUDA error: an illegal memory access was encountered。这是最常见的错误通常是数组越界。检查你的指针计算特别是当M、N、K不是tile大小的整数倍时边界处理有没有问题。我踩过的一个坑在加载A的tile时如果row M我写了As[threadIdx.y][threadIdx.x] 0.0f但忘了检查k threadIdx.x K导致读取了A矩阵之外的内存。这种错误在矩阵大小是tile整数倍的时候不会出现但一旦不是整数倍就会崩。问题二结果不对但也不报错。这个通常是同步问题。GEMM kernel里有很多__syncthreads()如果漏了一个或者放错了位置就会导致数据竞争结果随机出错。排查方法先用小矩阵比如16x16x16测试如果小矩阵对大矩阵不对那大概率是同步问题。如果小矩阵也不对那就是计算逻辑有问题。问题三性能突然下降。如果你改了某个参数性能突然掉了一半那可能是寄存器溢出导致occupancy下降shared memory不够导致只能用一个stagetile大小和矩阵形状不匹配导致大量线程空转用nvcc --ptxas-options-v编译看看寄存器使用情况和shared memory使用情况。如果寄存器超过255或者shared memory超过硬件限制就会有问题。4.3 性能调优的常见误区误区一tile越大越好。不是的。tile越大寄存器压力越大occupancy越低。如果occupancy太低延迟隐藏不了性能反而下降。我见过有人把BLOCK_M和BLOCK_N都设成256结果寄存器溢出性能只有峰值的一半。误区二stages越多越好。stages越多流水线越深延迟隐藏越好。但stages多了shared memory占用也大。如果shared memory不够就只能减小tile反而得不偿失。一般来说3-4个stages是比较平衡的选择。误区三所有矩阵形状都用同一套参数。这是最大的误区。不同的矩阵形状最优参数可能完全不同。比如M1024, N1024, K1024和M4096, N4096, K4096最优的tile大小可能就不一样。DeepGEMM这种项目通常会针对几种典型的形状做专门的调优而不是用一个通用配置。误区四只看TFLOPS不看实际端到端性能。TFLOPS是kernel级别的指标但实际应用中GEMM只是整个模型的一部分。如果GEMM快了但数据搬运慢了端到端性能可能没提升。所以要看端到端的指标比如训练一个epoch需要多长时间推理一个batch需要多长时间。4.4 常见问题速查表现象可能原因排查方法解决方法编译报错找不到头文件路径配置错误检查CMakeLists加上include_directories运行时illegal memory access数组越界用compute-sanitizer检查边界条件结果不对但不报错同步问题用小矩阵测试检查__syncthreads性能远低于预期参数不合适用profiler分析调整tile/stages性能波动大没有预热增加预热次数预热10-100次大矩阵性能差并行度不够检查SM利用率用split-K小矩阵性能差启动开销大检查kernel launch时间用CUDA Graph精度不够累加精度低检查累加类型用FP32累加注意compute-sanitizer是NVIDIA提供的工具可以检测内存越界、数据竞争等问题。用法是compute-sanitizer ./your_program。虽然会慢很多但排查问题很有效。5. 从DeepGEMM延伸矩阵乘法优化的未来方向5.1 稀疏矩阵乘法深度学习模型里有很多稀疏性。比如ReLU激活之后大概有一半的神经元输出是0。如果能把这种稀疏性利用起来理论上可以把计算量减少一半。但稀疏矩阵乘法很难做因为稀疏的模式是不规则的GPU的SIMT架构对不规则计算不友好。现在有一些工作在做结构化稀疏比如2:4稀疏每4个元素里有2个是0这种稀疏模式比较规整Tensor Core可以直接支持。DeepGEMM未来可能会加入对稀疏矩阵乘法的支持。5.2 动态形状现在的GEMM kernel通常要求矩阵形状是固定的或者至少是tile大小的整数倍。但实际应用中矩阵形状往往是动态的比如NLP任务里sequence length是可变的。动态形状会导致kernel需要处理边界增加开销。有一些工作在做动态形状的GEMM比如用padding把形状补齐到tile大小的整数倍或者用更灵活的tile划分。DeepGEMM如果要做通用这个问题必须解决。5.3 多GPU和分布式单GPU的算力总是有限的。当模型大到单个GPU放不下的时候就需要多GPU甚至多节点。多GPU的GEMM涉及到通信比如all-reduce、all-gather。通信和计算的overlap是一个很大的优化空间。DeepGEMM目前看起来是单GPU的优化但未来可能会扩展到多GPU场景。这需要和通信库比如NCCL配合把通信和计算重叠起来。5.4 新硬件架构NVIDIA的GPU架构一直在演进。从Volta到Turing到Ampere到Hopper每一代都有新的特性。比如Hopper引入了TMATensor Memory Accelerator和Thread Block Cluster这些新特性可以进一步提升GEMM性能。DeepGEMM需要跟上硬件的发展针对新架构做优化。比如用TMA做异步数据加载用Cluster做跨block的数据共享。这些优化可以带来显著的性能提升。5.5 自动调优手动调优GEMM参数很费时间。自动调优Auto-tuning是一个趋势。比如用机器学习的方法预测最优参数或者用搜索算法自动找最优配置。CUTLASS的profiler就是一个自动调优工具但它需要手动跑。未来可能会有更智能的自动调优系统能够根据矩阵形状和硬件配置自动生成最优的kernel。这会大大降低GEMM优化的门槛。我个人在实际操作中的体会是GEMM优化是一个深不见底的领域。你每深入一层都会发现还有更多的优化空间。但这也是它的魅力所在。如果你对高性能计算感兴趣GEMM优化是一个很好的切入点。它涉及到计算机体系结构、并行计算、编译原理等多个领域的知识做透了会对你的技术能力有很大的提升。最后再分享一个小技巧如果你要优化一个特定的GEMM形状先用CUTLASS的profiler跑一遍看看它推荐什么配置。然后在这个配置的基础上手动调整几个关键参数比如BLOCK_K和STAGES。通常只需要调几次就能找到接近最优的配置。不要从零开始瞎试那样效率太低。
阅读完成 · 觉得有帮助？
