做算子优化这些年我一直有个体会真正卡住团队进度的往往不是算法本身而是底层那套矩阵计算的基础设施不够趁手。上个月调一个融合了残差和 LayerNorm 的线性层性能怎么都追不平 cuBLAS 的高峰最后翻出团队早期用 CATLASS 这套矩阵计算模板库做的老底子重新梳理才意识到问题根本不是算子写得不对而是数据流从设计之初就没有按模板提供的分块思路来。这篇文章就借这次复盘把 CATLASS 模板库的核心设计、GEMM 数据流拆解以及混合精度优化里那些文档不会明说的细节一次性讲透。CATLASS 是一套面向 CUDA 平台的高性能矩阵计算模板库核心思路是把 GEMM通用矩阵乘法这类算子彻底模板化——从分块大小、数据布局、指令派发到流水线策略全部变成编译期参数。它解决的痛点是直接写 CUDA kernel 虽然灵活但每次新业务场景都要从头优化一遍重复造轮子成本极高而直接用 cuBLAS 又很难做算子融合和定制。CATLASS 正好卡在中间既能通过拆解好的 Tile 迭代器、Warp 级指令封装拿到接近手工调优的性能又能让上层业务以组合模板的方式快速搭建新算子。适合谁看正在做深度学习推理引擎、算子库研发或者想搞懂高性能 GEMM 内部到底发生了什么的人。1. CATLASS 模板库定位与整体设计思路1.1 模板库到底解决的是什么问题做高性能计算的人都会遇到一个尴尬写一个能跑的 GEMM kernel 并不难难的是让它跑得跟 cuBLAS 一样快。cuBLAS 之所以快是因为它对每一种 shape、每一种数据类型、每一种架构都做了针对性调优而这份调优经验被固化成了庞大的代码库。CATLASS 这类模板库的思路就是把这套经验中与具体硬件相关的部分抽象成模板参数让使用者通过组合模板来生成高性能 kernel。模板参数化的核心收益是把计算逻辑和数据布局解耦。比如在 CATLASS 里一个矩阵乘法的模板参数会包含数据类型、矩阵布局行主序/列主序、Tile 大小、每个线程负责的元素数、是否启用双缓冲、是否使用 WMMA 张量核心指令等。这些参数全部在编译期确定编译器可以做完整的常量展开和循环展开最终生成的代码接近手工极致调优的版本同时又不失灵活性。举个例子要给一份 FP16 的矩阵乘法适配 Ampere 架构张量核心你只需要改一行模板参数把指令策略从 SIMT 换成 WMMA要给一个 128x128 的 Tile 改成 256x128也只是一个常量。这种改参数而不是改代码的体验对算子库团队来说极其重要——你从零写一个高性能 kernel 可能要两周调模板一天就能搞定。1.2 模板粒度粗了不够灵活细了编译爆炸CATLASS 设计上最值得琢磨的一点是模板粒度的取舍。模板参数如果过细每一个循环层次都暴露给用户确实灵活但代价是模板瞬间变得极其冗长编译时间暴涨报错信息人类几乎无法阅读如果过粗性能和定制能力又受限。CATLASS 的折中方案是采用三层映射结构线程块级负责 Tile 调度和 Shared Memory 切分Warp 级负责子 Tile 的排布线程级负责逐元素乘加和向量化加载。这种粒度划分背后是有道理的。以 CUDA 编程模型来看对性能影响最大的三个决策分别是数据如何从全局内存搬入共享内存涉及访问模式、向量化、数据如何从共享内存搬到寄存器涉及 bank conflict、寄存器重用、计算用哪种指令SIMT 指令、FMA 融合乘加还是 Tensor Core 的矩阵指令。CATLASS 把这三点拆成了不同层级的模板模块互相之间可以独立替换这是我用过之后觉得它比早期很多一个大类搞定一切的实现要高明的地方。粒度选择还有一个实践层面的考量——可读性。团队做模板库不是写完一锤子买卖后续要长期演进和维护。CATLASS 的模板分层让新成员能按线程块级→Warp 级→线程级的顺序逐步深入而不是一上来就要理解所有宏展开。这个设计对整个开源社区的高性能计算项目都有参考价值。1.3 核心抽象Tile、Iterator 与布局CATLASS 里我最常用到的是三个核心抽象Tile、Tile Iterator 和 Layout。Tile 是一块放进共享内存或寄存器中的子矩阵Tile Iterator 负责按预定模式从全局内存/共享内存搬运数据Layout 则描述矩阵元素在内存中的排列方式比如行主序下第 m 行第 n 列的元素在 offset 为 m * ldm n 的位置。这三者的关系可以这么理解Tile 是容器Iterator 是搬运机器人Layout 是仓库货架分布图。GEMM 的执行流程就是一个一个 Tile 被 Iterator 按 Layout 指引搬到共享内存和寄存器完成乘累加再把结果 Tile 写回全局内存。由于这三者都是模板参数你可以给同一个计算核心配上完全不同的 Layout 和 Iterator。比如处理非连续内存的稀疏场景你可以写一个自定义 Iterator而主计算循环一行都不用动。这种可替换性一度帮我省过很多事当时要给一个按列分块的业务矩阵写 GEMMnaive 做法是把列分块重排成连续内存额外一次数据拷贝但我直接写了个列访问 Iterator计算循环原封不动性能提升明显且省了一次全局内存读写。2. GEMM 数据流拆解从全局内存到寄存器的三层搬移2.1 分块计算Tiling是 GEMM 性能的基石GEMM 的计算公式是 C A * B C核心操作是乘累加。一个 1024x1024 的矩阵乘如果不做任何分块A 的每一行要跟 B 的每一列做内积意味着 A 的元素会被反复从全局内存读取 1024 次对每个输出列都要取一次。全局内存带宽是稀缺资源这种重复读取必然把性能压到内存带宽的天花板之下。分块计算的思路是把输出矩阵 C 切成若干个小的 Tile比如 128x128每个线程块只负责这一个 Tile 的计算。计算这个 Tile 只需要 A 的一个水平条带128 行和 B 的一个垂直条带128 列这两块数据在计算过程中会反复被汇编代码使用因此可以提前搬入共享内存或寄存器中做重复利用。这样一来全局内存的访问次数从每个输出元素读一次源数据降为每个 Tile 读一次源数据降幅大约为 Tile 尺寸的量级。CATLASS 在 Tiling 上的一个重要设计是 Tile 尺寸的选择。128x128 是一个经验上很均衡的尺寸共享内存占用合理FP16 下两个矩阵各 32KB加上主矩阵和额外开销正好压住 48KB~100KB 的范围线程块内并行度高可容纳 256 个线程寄存器利用率高。我用过超大 Tile256x256性能并不一定更好因为单个线程块要算的数据多了双缓冲流水线深度会被共享内存容量压住反而导致流水线气泡。2.2 三层搬移全局内存 → 共享内存 → 寄存器 → 计算CATLASS 的 GEMM 内层循环本质上是一条三级流水线。第一级把 A 和 B 的 Tile 从全局内存用向量化加载如 float4即连续读 16 字节搬入共享内存第二级把共享内存中的数据搬到线程各自的寄存器中并按乘累加的因子分配好每个线程需要计算的子区域第三级就是寄存器里的 FMAs融合乘加。这里面每一级的代价差距非常大。全局内存访问延迟一般几百个周期共享内存延迟约为几十个周期寄存器则零延迟。所以优化的核心思想是尽量把同一块数据在寄存器中反复使用避免频繁访问共享内存更不要频繁访问全局内存。CATLASS 把这一层逻辑做得非常细。在寄存器层级每个线程通常负责一个 8x8 或者 16x8 的微 Tile。以 FP16 利用 WMMA 张量核心为例每个线程通过一条mma.sync指令可以完成 16x16x16 的矩阵乘累加即一次指令做 4096 次乘累加。如果走普通 SIMT 的 FMA一次指令只能做一次乘累加。这也是混合精度性能和 FP32 拉开差距的本质原因之一——除了数值格式省了带宽指令本身的吞吐也完全不同。2.3 共享内存 Bank Conflict 规避与数据布局共享内存位于芯片内部吞吐远高于全局内存但它有 32 个 bank每个 bank 在单周期内只能返回一个 4 字节数据。如果同一周期内多个线程访问的地址映射到了同一个 bank称为 bank conflict这些访问会被串行化吞吐直接除以冲突次数。一个常见的 2-way bank conflict 就会让共享内存读取性能掉一半这个损耗在 GEMM 这种高频循环里会被放大成肉眼可见的性能滑坡。CATLASS 处理 bank conflict 的经典手法是 padding也就是在共享内存的 Tile 每一行末尾多塞几个元素把每行的实际 stride 从 64 字节调成 68 字节或类似值。这样原本按对角线访问模式的线程地址被错开避免多个线程同时命中同一 bank。我当初第一次跑自己的 GEMM kernel 时性能比 cuBLAS 低约三成用 Nsight Compute 看到 LDS加载共享内存指令的 bank conflict 统计高达 12%就是因为没有做 padding。换一个角度说bank conflict 的排查往往也是高性能算子优化中收益最直接的几个手段之一。你们如果发现自己写的 GEMM 或者 Implicit GEMM 卷积算子性能怎么调都上不去第一步不是怀疑指令选型而是打开 profiler 看内存访问那一档是不是有红黄指标。2.4 双缓冲与流水线用异步把延迟藏起来数据搬移和计算天然存在依赖要等数据到齐才能算算完才能搬下一批。如果串行执行等待数据的时间就得靠发虚。CATLASS 的解法是流水线化典型实现是双缓冲共享内存里开两份缓冲区一份用于当前计算一份用于预取下一批数据。计算当前 Tile 的同时通过cp.async指令异步把下一个 Tile 的数据从全局内存搬入另一份缓冲区两件事并行发生等待延迟被完全隐藏。Ampere 架构之后cp.async指令可以直接让数据绕开寄存器从全局内存异步拷贝到共享内存拷完自动触发完成机制不需要显式再走一次 LDS。用这套机制时流水线 stage 数的选择非常关键。stage2 是双缓冲stage3 是三级流水。stage 越多共享内存占用越大但抵抗延迟波动的能力越强。CATLASS 在部分 kernel 中默认 stage4因为它按最坏情况延迟来设计而对很多业务场景来说 stage2 就够。实际使用中我建议从 stage2 开始往上加观察实际吞吐变化。有的场景共享内存容量已经被 Tile 占得差不多硬塞 stage4 会导致 occupancy占用率下降反而不如 stage2 来的稳。混合精度下感受会更明显——低精度 Tile 本身占的内存少流水线深度提高带来的收益比 FP32 大得多。3. 混合精度优化的核心实践3.1 混合精度为什么能快算力与带宽双线并进混合精度Mixed Precision在训练和推理中的收益来自两个层面。第一是算力层面现代 GPU 的 Tensor Core 对 FP16/BF16 的矩阵乘累加吞吐通常可以达到 FP32 的数倍以上。以 A100 为例FP32 FMA 的标称算力约 19.5 TFLOPS而 Tensor Core FP16 可以到 312 TFLOPS差距接近 16 倍。第二是带宽层面FP16 每个元素只占 2 字节是 FP32 的一半全局内存和共享内存的单位时间搬运元素数量在同样带宽下直接翻倍。对于矩阵乘这类访存和计算并重的算子两个层面的收益可叠加。值得注意的是算力差 16 倍并不等于端到端快 16 倍。GEMM 如果是访存受限型矩阵比较小、无法充分复用数据主要的瓶颈在搬运Tensor Core 高算力完全使不上劲。比如一批 64x64 的小 GEMM性能差异可能只有 1.5 倍。混合精度优化的正确思路是先判断算子是 compute-bound 还是 memory-bound再决定是否值得上 Tensor Core 路径。我在实践中遇到不少团队直接把所有算子改成 FP16结果收益很有限反而还引入了精度风险就是忽略了这一步判据。3.2 FP16 / BF16 / TF32 三种格式的选型对比混合精度里最常碰到的三种低精度格式各有各的使用场景。FP16半精度1 位符号、5 位指数、10 位尾数。它能表示的最大值约 65504超过就上溢到无穷最小正常值约 6e-5更小的数会逐步损失精度一直到约 6e-8 的次正规数边界。FP16 的优势是 Tensor Core 支持最成熟、前后端软件栈适配最好劣势是指数范围窄训练时如果梯度或中间激活值超过 65504 或者跌到 1e-4 以下都会出问题。BF16Brain Floating Point1 位符号、8 位指数、7 位尾数。指数范围和 FP32 完全一致最大可表示约 3.4e38最小正常值约 1e-38基本不会出现上溢或下溢的问题。代价是尾数只有 7 位相对精度低。BF16 在训练场景下因为不容易炸又配合损失缩放已经成为主流训练格式但在推理场景如果模型权重数值本身很小比如某些蒸馏后的模型BF16 的尾数不足容易导致精度明显下降。TF32本质是在 Tensor Core 计算中把 FP32 数据截断为 10 位尾数参与矩阵乘。它不用改内存里的数据格式输入输出仍按 FP32 存储因此兼容性最好适合那些想换取 Tensor Core 加速又不愿改动数据通路的 FP32 模型。TF32 的精度约等于劣化 FP32只保留大约 10 位十进制有效数字对大多数深度学习模型在可接受范围内。我的选型经验训练首选 BF16 动态损失缩放推理首选 FP16如果模型精度敏感再做逐层精度检测来决定哪些层退回 FP32。TF32 适合基础设施层不想大动干戈、只是想白拿一部分算力的场景但因为内存占用没有减半带宽收益全无。格式指数位尾数位最大值最小正常值典型场景FP32823~3.4e38~1.2e-38默认基线、精度兜底FP1651065504~6.1e-5推理、混合精度训练BF1687~3.4e38~1.2e-38大规模训练TF32810~3.4e38~1.2e-38输入输出仍为 FP32 的加速场景3.3 损失缩放Loss Scaling混合精度不能绕过的一环FP16 训练中一个经典问题是梯度下溢。深度学习训练时梯度的数值往往远小于权重早期层梯度降到 1e-6 甚至更小很常见而 FP16 的最小正常值只有 6e-5更小的梯度直接变成 0参数就再也不更新了。解决办法是损失缩放在反向传播前把损失值乘以一个较大的系数通常 1024 或 2048梯度也会同比例放大FP16 能容纳的数值范围自动覆盖原本小到会消失的梯度反向传播完成后更新参数前再把梯度除以同样的系数。混合精度库一般在内部自动处理了损失缩放但用 CATLASS 这类底层模板库做自定义算子时这个环节要自己注意。我踩过一次坑用 FP16 写一个自定义 GEMM 核反向传播算出来的梯度在小数值区域明显偏小最后定位到是没有对中间梯度做任何 scale梯度在矩阵乘内部就截断成了 0。这里建议梯度的缩放因子遵循先放大、算完、再缩放的顺序不要在乘累加内部反复乘除因为每次缩放都会引入一次舍入误差。损失缩放还有一个细节是动态缩放。固定缩放因子 1024 在部分模型里不够用太大会让权重上溢太小挡不住梯度下溢更好的方案是动态检测连续一段时间没有出现 inf/NaN就适当调大缩放因子一旦检测到 inf/NaN立即调小并跳过本轮更新。PyTorch 的 GradScaler 就是这么做的底层模板库使用者可以仿照这个逻辑自己实现成本不高但效果明显。3.4 混合精度在矩阵乘法中的落点哪些算子能改不是所有算子都适合切混合精度。性能层面切低精度的收益来自算力翻倍和带宽减半所以 compute-bound 的大矩阵乘收益最大memory-bound 的 activation 类算子收益主要靠带宽。精度层面几何型算子做大量连续乘法和指数运算的对精度变化敏感规约型算子求和、求均值相对不敏感。以 Transformer 为例Attention 里的 QK 点积矩阵乘、以及 MLP 块里的两个大线性层是混合精度的最大受益者。它们都是大矩阵乘内存访问次数相对计算量占比低Tensor Core 高吞吐正好发挥。而 Softmax、LayerNorm 这类逐元素算子数值范围对精度影响大且本身不是计算密集型的通常保持 FP32或只在数据通路用 FP16 而将内部规约动态提升为 FP32。CATLASS 模板库因为把计算核心和规约算子分开设计做这种矩阵乘低精度、规约高精度的组合非常方便。实践上还有个经验混合精度代码里尽量保持计算时低精度、累加时高精度的模式。加法树用 FP32 累加器乘法和乘累加用 FP16这样既吃到了 Tensor Core 算力又把累加误差控制在 FP32 水平。这也是 CATLASS 混精度 kernel 里一个被写死的模板参数——如果你想改成全 FP16 累加不是不能改但绝大多数情况都不会更好。4. 混合精度踩坑实录与排查方法4.1 精度问题不一定是舍入先查指数范围溢出很多团队在混合精度上遇到的第一个问题不是性能不达标而是模型直接产出 NaN。常规思路会怀疑舍入误差、累加顺序、FMA 融合但排查了一圈下来FP16 的上溢才是最大嫌疑。FP16 最大只到 65504如果权重矩阵的初始值带几个偏差大的异常点或者激活值过了一个带大 scale 的层几乎必炸。排查方法是逐层打印中间激活值和权重的绝对最大值、绝对值最小值、绝对值均值对比 FP32 基线的同一层数值。如果 FP16 下某一层的 max 超过 65504或者 min 掉到 1e-4 以下基本可以锁定是范围问题而不是精度问题。解决手段有三条一是把该层数据先做 affine 归一化再运算二是在矩阵乘之前手动对输入做缩放三是干脆将该层退回 FP32。这三条路我都在实际项目里用过第一条路的性能损失最小因为归一化通常可以融合进前面的算子。4.2 性能问题bank conflict、occupancy 和指令混合混合精度 kernel 性能不如预期通常不是 Tensor Core 没发力而是旁边的数据通路拖了后腿。我排查得最多的三个方向是共享内存 bank conflict、寄存器溢出spill和线程块占用率occupancy过低。这三个问题有一个共同特点——不会导致结果错误所以更容易被忽略。排查思路是打开 Nsight Compute先看Memory Workload Analysis里的 LDS Bank Conflict 指标超过 5% 就值得优化。再看 Occupancy 那一栏如果实际值低于理论值大概率是共享内存或寄存器资源超了。最后看一眼 warp state 里的待命周期Stall如果大量周期是等待共享内存或全局内存数据回来说明流水线深度不足需要增加 stage。4.3 一个典型案例错误配置导致性能反而下降之前帮一个团队看一个 Transformer 推理引擎他们把 QKV 的 GEMM 全部切成 FP16FP32 基线大约 220 TFLOPS 的利用率切完反而只有 120 TFLOPS甚至不如基线。逐个参数排查后发现了两个问题叠加第一Tile 大小从 128x128 改成了 64x64因为共享内存里给双缓冲留的缓冲区被 FP16 数据塞得更大了看起来应该刚好能塞下但忘记了 CATLASS 还需要给共享内存预留一行 padding 和 stage 缓冲区结果 occupancy 掉了一半。第二warp 数量从 8 个减到 4 个后寄存器循环展开量不够全局内存加载频率变高流水线没法藏住延迟。修正方案很简单恢复 128x128 Tile把双缓冲 stage 从 4 降到 2给 padding 留出空间。改完性能直接回到接近 FP32 基线的 1.8 倍。这给团队的教训是混合精度不只是换数据类型它改变了一整条数据通路的资源占用面貌所有依赖资源的设计Tile、stage、寄存器数都要重新平衡。4.4 性能排查工具箱关键工具与基本使用流程我调 CATLASS kernel 时一般会按下面这个流程走用ncu --set full跑一次 Nsight Compute重点看 SOLSpeed of Light里的 Compute 和 Memory Throughput。如果两者都没超过 60%先假设 kernel 被延迟限制了看 warp stall 原因再往前推。如果 Compute 接近 90%、Memory 接近 90%那就是优化得很好的状态。接下来看 instructions 里的 LDS/UDS 数量、bank conflict 计数、global load/store 效率。最后用ncu --metrics launch__occupancy_limit_shared_mem,launch__occupancy_limit_registers确认资源瓶颈。数值精度排查方面写一个小工具脚本FP32 基线算一遍各层输出混合精度算一遍统计逐元素的绝对误差、相对误差、最大误差位置。通常相对误差在 1e-2 以内就可以接受如果超过 0.1需要定位是哪一层开始放大的。这个方法虽然土但在追踪精度在哪个算子被毁掉时非常有效。5. 从 GEMM 到更多算子模板库能力的延伸应用5.1 卷积的隐式 GEMMImplicit GEMM化CATLASS 的 GEMM 核心不止能算常规矩阵乘也是卷积算子的基础。卷积本质上是一个四维循环可以转换成矩阵乘形式把输入特征图按卷积窗口展开成矩阵im2col把卷积核展开成另一个矩阵两者相乘再重组。naive 的 im2col 需要额外拷贝数据内存开销大得吓人展开后数据量通常膨胀几十倍而隐式 GEMM 的思路是不真正展开数据而是在 GEMM 的 Tile Iterator 里直接根据卷积映射关系去取数据。这个思路在 CATLASS 里落地得很自然。Tile Iterator 本来就是按布局搬运数据的抽象你给它定义好第 m 行对应输入特征图的哪个位置、第 n 列对应卷积核的哪个位置它会自动在 GEMM 主循环里完成取数。这样做的好处是主循环计算逻辑完全复用 GEMM 的优化积累——Tiling、双缓冲、Tensor Core 指令全都是现成的。5.2 注意力机制算子的融合实践注意力机制里的 QK^T 和 Score·V 本质上也是矩阵乘定制的空间在于中间要穿插 Softmax并且要避免把整个 Score 矩阵写回全局内存。若完全按通用 GEMM 做QK^T 算完把 NxN 的 Score 矩阵写回全局内存再从全局内存读回来做 Softmax、再参与 PV 矩阵乘访存开销非常大。CATLASS 的做法是把两段矩阵乘和中间的 Softmax 融合成一个大的 kernel中间的 Score 留在寄存器或共享内存只借道 CPU 或小块内存做规约。这个方向最有名的实现是 FlashAttention它的核心贡献之一就是用分块思想避免大 Score 矩阵落回全局内存。CATLASS 用户完全可以照着这个思路实现一个定制化的 Attention kernelTile 迭代器照用 GEMM 的中间替换成融合了 Softmax 的 epilogue后处理。这种保留 GEMM 主循环、替换 epilogue的自定义模式是模板库最有价值的部分。5.3 自定义算子的复用模式与边界CATLASS 这类模板库给了使用者很大的自由但边界也要心里有数。适合复用的场景是算子主体是矩阵乘或点积类计算差异点在于取数方式和后处理逻辑。不适合复用的场景是算子数据依赖关系非常非线性比如复杂的稀疏路径、动态控制流极多的场景。这类算子用模板库强行套模板最后只能越套越复杂不如直接写 NVCC kernel。我给团队定的经验准则是如果算子中有超过 80% 的工作本质是计算一个输出 Tile 若干输入 Tile 的乘累加就用 CATLASS 改否则直接手写。手写时可以借用 CATLASS 里的分块、双缓冲思路但不需要强行上模板。很多团队项目失败是因为过早地想把所有算子都用模板库统一了结果收益没看到只看到编译时间和排查成本的上升。说到调试CATLASS 模式下的模板代码报错信息向来不友好。一个通用技巧是按C 模板类实例化展开的方式去读报错先把缺失的类型、常量推断出来再反查是哪个模板参数类型不匹配。并且确保编译时打开-G调试模式和-lineinfo选项这样反汇编和性能分析的定位会准确得多。关于这套模板库的实际使用我个人的体会是它确实不是最快让人出活的那条路初期学习成本比直接调 cuBLAS 高不少但一旦把 Tile 迭代器、流水线、双缓冲、混合精度这些概念吃透后续写任何高性能算子都会上一个台阶。最后再分享一个调试小技巧当 kernel 性能不符合预期时先别急着搜代码把 Nsight Compute 的 SOL 页面截图对照 cuBLAS 同规格算子的 SOL 图找差异区间通常一眼就能看出是计算瓶颈还是访存瓶颈省掉大量盲目尝试的时间。这套方法论比背任何库的 API 都值钱。 深入解析CATLASS模板库：GEMM数据流与混合精度优化实战
做算子优化这些年我一直有个体会真正卡住团队进度的往往不是算法本身而是底层那套矩阵计算的基础设施不够趁手。上个月调一个融合了残差和 LayerNorm 的线性层性能怎么都追不平 cuBLAS 的高峰最后翻出团队早期用 CATLASS 这套矩阵计算模板库做的老底子重新梳理才意识到问题根本不是算子写得不对而是数据流从设计之初就没有按模板提供的分块思路来。这篇文章就借这次复盘把 CATLASS 模板库的核心设计、GEMM 数据流拆解以及混合精度优化里那些文档不会明说的细节一次性讲透。CATLASS 是一套面向 CUDA 平台的高性能矩阵计算模板库核心思路是把 GEMM通用矩阵乘法这类算子彻底模板化——从分块大小、数据布局、指令派发到流水线策略全部变成编译期参数。它解决的痛点是直接写 CUDA kernel 虽然灵活但每次新业务场景都要从头优化一遍重复造轮子成本极高而直接用 cuBLAS 又很难做算子融合和定制。CATLASS 正好卡在中间既能通过拆解好的 Tile 迭代器、Warp 级指令封装拿到接近手工调优的性能又能让上层业务以组合模板的方式快速搭建新算子。适合谁看正在做深度学习推理引擎、算子库研发或者想搞懂高性能 GEMM 内部到底发生了什么的人。1. CATLASS 模板库定位与整体设计思路1.1 模板库到底解决的是什么问题做高性能计算的人都会遇到一个尴尬写一个能跑的 GEMM kernel 并不难难的是让它跑得跟 cuBLAS 一样快。cuBLAS 之所以快是因为它对每一种 shape、每一种数据类型、每一种架构都做了针对性调优而这份调优经验被固化成了庞大的代码库。CATLASS 这类模板库的思路就是把这套经验中与具体硬件相关的部分抽象成模板参数让使用者通过组合模板来生成高性能 kernel。模板参数化的核心收益是把计算逻辑和数据布局解耦。比如在 CATLASS 里一个矩阵乘法的模板参数会包含数据类型、矩阵布局行主序/列主序、Tile 大小、每个线程负责的元素数、是否启用双缓冲、是否使用 WMMA 张量核心指令等。这些参数全部在编译期确定编译器可以做完整的常量展开和循环展开最终生成的代码接近手工极致调优的版本同时又不失灵活性。举个例子要给一份 FP16 的矩阵乘法适配 Ampere 架构张量核心你只需要改一行模板参数把指令策略从 SIMT 换成 WMMA要给一个 128x128 的 Tile 改成 256x128也只是一个常量。这种改参数而不是改代码的体验对算子库团队来说极其重要——你从零写一个高性能 kernel 可能要两周调模板一天就能搞定。1.2 模板粒度粗了不够灵活细了编译爆炸CATLASS 设计上最值得琢磨的一点是模板粒度的取舍。模板参数如果过细每一个循环层次都暴露给用户确实灵活但代价是模板瞬间变得极其冗长编译时间暴涨报错信息人类几乎无法阅读如果过粗性能和定制能力又受限。CATLASS 的折中方案是采用三层映射结构线程块级负责 Tile 调度和 Shared Memory 切分Warp 级负责子 Tile 的排布线程级负责逐元素乘加和向量化加载。这种粒度划分背后是有道理的。以 CUDA 编程模型来看对性能影响最大的三个决策分别是数据如何从全局内存搬入共享内存涉及访问模式、向量化、数据如何从共享内存搬到寄存器涉及 bank conflict、寄存器重用、计算用哪种指令SIMT 指令、FMA 融合乘加还是 Tensor Core 的矩阵指令。CATLASS 把这三点拆成了不同层级的模板模块互相之间可以独立替换这是我用过之后觉得它比早期很多一个大类搞定一切的实现要高明的地方。粒度选择还有一个实践层面的考量——可读性。团队做模板库不是写完一锤子买卖后续要长期演进和维护。CATLASS 的模板分层让新成员能按线程块级→Warp 级→线程级的顺序逐步深入而不是一上来就要理解所有宏展开。这个设计对整个开源社区的高性能计算项目都有参考价值。1.3 核心抽象Tile、Iterator 与布局CATLASS 里我最常用到的是三个核心抽象Tile、Tile Iterator 和 Layout。Tile 是一块放进共享内存或寄存器中的子矩阵Tile Iterator 负责按预定模式从全局内存/共享内存搬运数据Layout 则描述矩阵元素在内存中的排列方式比如行主序下第 m 行第 n 列的元素在 offset 为 m * ldm n 的位置。这三者的关系可以这么理解Tile 是容器Iterator 是搬运机器人Layout 是仓库货架分布图。GEMM 的执行流程就是一个一个 Tile 被 Iterator 按 Layout 指引搬到共享内存和寄存器完成乘累加再把结果 Tile 写回全局内存。由于这三者都是模板参数你可以给同一个计算核心配上完全不同的 Layout 和 Iterator。比如处理非连续内存的稀疏场景你可以写一个自定义 Iterator而主计算循环一行都不用动。这种可替换性一度帮我省过很多事当时要给一个按列分块的业务矩阵写 GEMMnaive 做法是把列分块重排成连续内存额外一次数据拷贝但我直接写了个列访问 Iterator计算循环原封不动性能提升明显且省了一次全局内存读写。2. GEMM 数据流拆解从全局内存到寄存器的三层搬移2.1 分块计算Tiling是 GEMM 性能的基石GEMM 的计算公式是 C A * B C核心操作是乘累加。一个 1024x1024 的矩阵乘如果不做任何分块A 的每一行要跟 B 的每一列做内积意味着 A 的元素会被反复从全局内存读取 1024 次对每个输出列都要取一次。全局内存带宽是稀缺资源这种重复读取必然把性能压到内存带宽的天花板之下。分块计算的思路是把输出矩阵 C 切成若干个小的 Tile比如 128x128每个线程块只负责这一个 Tile 的计算。计算这个 Tile 只需要 A 的一个水平条带128 行和 B 的一个垂直条带128 列这两块数据在计算过程中会反复被汇编代码使用因此可以提前搬入共享内存或寄存器中做重复利用。这样一来全局内存的访问次数从每个输出元素读一次源数据降为每个 Tile 读一次源数据降幅大约为 Tile 尺寸的量级。CATLASS 在 Tiling 上的一个重要设计是 Tile 尺寸的选择。128x128 是一个经验上很均衡的尺寸共享内存占用合理FP16 下两个矩阵各 32KB加上主矩阵和额外开销正好压住 48KB~100KB 的范围线程块内并行度高可容纳 256 个线程寄存器利用率高。我用过超大 Tile256x256性能并不一定更好因为单个线程块要算的数据多了双缓冲流水线深度会被共享内存容量压住反而导致流水线气泡。2.2 三层搬移全局内存 → 共享内存 → 寄存器 → 计算CATLASS 的 GEMM 内层循环本质上是一条三级流水线。第一级把 A 和 B 的 Tile 从全局内存用向量化加载如 float4即连续读 16 字节搬入共享内存第二级把共享内存中的数据搬到线程各自的寄存器中并按乘累加的因子分配好每个线程需要计算的子区域第三级就是寄存器里的 FMAs融合乘加。这里面每一级的代价差距非常大。全局内存访问延迟一般几百个周期共享内存延迟约为几十个周期寄存器则零延迟。所以优化的核心思想是尽量把同一块数据在寄存器中反复使用避免频繁访问共享内存更不要频繁访问全局内存。CATLASS 把这一层逻辑做得非常细。在寄存器层级每个线程通常负责一个 8x8 或者 16x8 的微 Tile。以 FP16 利用 WMMA 张量核心为例每个线程通过一条mma.sync指令可以完成 16x16x16 的矩阵乘累加即一次指令做 4096 次乘累加。如果走普通 SIMT 的 FMA一次指令只能做一次乘累加。这也是混合精度性能和 FP32 拉开差距的本质原因之一——除了数值格式省了带宽指令本身的吞吐也完全不同。2.3 共享内存 Bank Conflict 规避与数据布局共享内存位于芯片内部吞吐远高于全局内存但它有 32 个 bank每个 bank 在单周期内只能返回一个 4 字节数据。如果同一周期内多个线程访问的地址映射到了同一个 bank称为 bank conflict这些访问会被串行化吞吐直接除以冲突次数。一个常见的 2-way bank conflict 就会让共享内存读取性能掉一半这个损耗在 GEMM 这种高频循环里会被放大成肉眼可见的性能滑坡。CATLASS 处理 bank conflict 的经典手法是 padding也就是在共享内存的 Tile 每一行末尾多塞几个元素把每行的实际 stride 从 64 字节调成 68 字节或类似值。这样原本按对角线访问模式的线程地址被错开避免多个线程同时命中同一 bank。我当初第一次跑自己的 GEMM kernel 时性能比 cuBLAS 低约三成用 Nsight Compute 看到 LDS加载共享内存指令的 bank conflict 统计高达 12%就是因为没有做 padding。换一个角度说bank conflict 的排查往往也是高性能算子优化中收益最直接的几个手段之一。你们如果发现自己写的 GEMM 或者 Implicit GEMM 卷积算子性能怎么调都上不去第一步不是怀疑指令选型而是打开 profiler 看内存访问那一档是不是有红黄指标。2.4 双缓冲与流水线用异步把延迟藏起来数据搬移和计算天然存在依赖要等数据到齐才能算算完才能搬下一批。如果串行执行等待数据的时间就得靠发虚。CATLASS 的解法是流水线化典型实现是双缓冲共享内存里开两份缓冲区一份用于当前计算一份用于预取下一批数据。计算当前 Tile 的同时通过cp.async指令异步把下一个 Tile 的数据从全局内存搬入另一份缓冲区两件事并行发生等待延迟被完全隐藏。Ampere 架构之后cp.async指令可以直接让数据绕开寄存器从全局内存异步拷贝到共享内存拷完自动触发完成机制不需要显式再走一次 LDS。用这套机制时流水线 stage 数的选择非常关键。stage2 是双缓冲stage3 是三级流水。stage 越多共享内存占用越大但抵抗延迟波动的能力越强。CATLASS 在部分 kernel 中默认 stage4因为它按最坏情况延迟来设计而对很多业务场景来说 stage2 就够。实际使用中我建议从 stage2 开始往上加观察实际吞吐变化。有的场景共享内存容量已经被 Tile 占得差不多硬塞 stage4 会导致 occupancy占用率下降反而不如 stage2 来的稳。混合精度下感受会更明显——低精度 Tile 本身占的内存少流水线深度提高带来的收益比 FP32 大得多。3. 混合精度优化的核心实践3.1 混合精度为什么能快算力与带宽双线并进混合精度Mixed Precision在训练和推理中的收益来自两个层面。第一是算力层面现代 GPU 的 Tensor Core 对 FP16/BF16 的矩阵乘累加吞吐通常可以达到 FP32 的数倍以上。以 A100 为例FP32 FMA 的标称算力约 19.5 TFLOPS而 Tensor Core FP16 可以到 312 TFLOPS差距接近 16 倍。第二是带宽层面FP16 每个元素只占 2 字节是 FP32 的一半全局内存和共享内存的单位时间搬运元素数量在同样带宽下直接翻倍。对于矩阵乘这类访存和计算并重的算子两个层面的收益可叠加。值得注意的是算力差 16 倍并不等于端到端快 16 倍。GEMM 如果是访存受限型矩阵比较小、无法充分复用数据主要的瓶颈在搬运Tensor Core 高算力完全使不上劲。比如一批 64x64 的小 GEMM性能差异可能只有 1.5 倍。混合精度优化的正确思路是先判断算子是 compute-bound 还是 memory-bound再决定是否值得上 Tensor Core 路径。我在实践中遇到不少团队直接把所有算子改成 FP16结果收益很有限反而还引入了精度风险就是忽略了这一步判据。3.2 FP16 / BF16 / TF32 三种格式的选型对比混合精度里最常碰到的三种低精度格式各有各的使用场景。FP16半精度1 位符号、5 位指数、10 位尾数。它能表示的最大值约 65504超过就上溢到无穷最小正常值约 6e-5更小的数会逐步损失精度一直到约 6e-8 的次正规数边界。FP16 的优势是 Tensor Core 支持最成熟、前后端软件栈适配最好劣势是指数范围窄训练时如果梯度或中间激活值超过 65504 或者跌到 1e-4 以下都会出问题。BF16Brain Floating Point1 位符号、8 位指数、7 位尾数。指数范围和 FP32 完全一致最大可表示约 3.4e38最小正常值约 1e-38基本不会出现上溢或下溢的问题。代价是尾数只有 7 位相对精度低。BF16 在训练场景下因为不容易炸又配合损失缩放已经成为主流训练格式但在推理场景如果模型权重数值本身很小比如某些蒸馏后的模型BF16 的尾数不足容易导致精度明显下降。TF32本质是在 Tensor Core 计算中把 FP32 数据截断为 10 位尾数参与矩阵乘。它不用改内存里的数据格式输入输出仍按 FP32 存储因此兼容性最好适合那些想换取 Tensor Core 加速又不愿改动数据通路的 FP32 模型。TF32 的精度约等于劣化 FP32只保留大约 10 位十进制有效数字对大多数深度学习模型在可接受范围内。我的选型经验训练首选 BF16 动态损失缩放推理首选 FP16如果模型精度敏感再做逐层精度检测来决定哪些层退回 FP32。TF32 适合基础设施层不想大动干戈、只是想白拿一部分算力的场景但因为内存占用没有减半带宽收益全无。格式指数位尾数位最大值最小正常值典型场景FP32823~3.4e38~1.2e-38默认基线、精度兜底FP1651065504~6.1e-5推理、混合精度训练BF1687~3.4e38~1.2e-38大规模训练TF32810~3.4e38~1.2e-38输入输出仍为 FP32 的加速场景3.3 损失缩放Loss Scaling混合精度不能绕过的一环FP16 训练中一个经典问题是梯度下溢。深度学习训练时梯度的数值往往远小于权重早期层梯度降到 1e-6 甚至更小很常见而 FP16 的最小正常值只有 6e-5更小的梯度直接变成 0参数就再也不更新了。解决办法是损失缩放在反向传播前把损失值乘以一个较大的系数通常 1024 或 2048梯度也会同比例放大FP16 能容纳的数值范围自动覆盖原本小到会消失的梯度反向传播完成后更新参数前再把梯度除以同样的系数。混合精度库一般在内部自动处理了损失缩放但用 CATLASS 这类底层模板库做自定义算子时这个环节要自己注意。我踩过一次坑用 FP16 写一个自定义 GEMM 核反向传播算出来的梯度在小数值区域明显偏小最后定位到是没有对中间梯度做任何 scale梯度在矩阵乘内部就截断成了 0。这里建议梯度的缩放因子遵循先放大、算完、再缩放的顺序不要在乘累加内部反复乘除因为每次缩放都会引入一次舍入误差。损失缩放还有一个细节是动态缩放。固定缩放因子 1024 在部分模型里不够用太大会让权重上溢太小挡不住梯度下溢更好的方案是动态检测连续一段时间没有出现 inf/NaN就适当调大缩放因子一旦检测到 inf/NaN立即调小并跳过本轮更新。PyTorch 的 GradScaler 就是这么做的底层模板库使用者可以仿照这个逻辑自己实现成本不高但效果明显。3.4 混合精度在矩阵乘法中的落点哪些算子能改不是所有算子都适合切混合精度。性能层面切低精度的收益来自算力翻倍和带宽减半所以 compute-bound 的大矩阵乘收益最大memory-bound 的 activation 类算子收益主要靠带宽。精度层面几何型算子做大量连续乘法和指数运算的对精度变化敏感规约型算子求和、求均值相对不敏感。以 Transformer 为例Attention 里的 QK 点积矩阵乘、以及 MLP 块里的两个大线性层是混合精度的最大受益者。它们都是大矩阵乘内存访问次数相对计算量占比低Tensor Core 高吞吐正好发挥。而 Softmax、LayerNorm 这类逐元素算子数值范围对精度影响大且本身不是计算密集型的通常保持 FP32或只在数据通路用 FP16 而将内部规约动态提升为 FP32。CATLASS 模板库因为把计算核心和规约算子分开设计做这种矩阵乘低精度、规约高精度的组合非常方便。实践上还有个经验混合精度代码里尽量保持计算时低精度、累加时高精度的模式。加法树用 FP32 累加器乘法和乘累加用 FP16这样既吃到了 Tensor Core 算力又把累加误差控制在 FP32 水平。这也是 CATLASS 混精度 kernel 里一个被写死的模板参数——如果你想改成全 FP16 累加不是不能改但绝大多数情况都不会更好。4. 混合精度踩坑实录与排查方法4.1 精度问题不一定是舍入先查指数范围溢出很多团队在混合精度上遇到的第一个问题不是性能不达标而是模型直接产出 NaN。常规思路会怀疑舍入误差、累加顺序、FMA 融合但排查了一圈下来FP16 的上溢才是最大嫌疑。FP16 最大只到 65504如果权重矩阵的初始值带几个偏差大的异常点或者激活值过了一个带大 scale 的层几乎必炸。排查方法是逐层打印中间激活值和权重的绝对最大值、绝对值最小值、绝对值均值对比 FP32 基线的同一层数值。如果 FP16 下某一层的 max 超过 65504或者 min 掉到 1e-4 以下基本可以锁定是范围问题而不是精度问题。解决手段有三条一是把该层数据先做 affine 归一化再运算二是在矩阵乘之前手动对输入做缩放三是干脆将该层退回 FP32。这三条路我都在实际项目里用过第一条路的性能损失最小因为归一化通常可以融合进前面的算子。4.2 性能问题bank conflict、occupancy 和指令混合混合精度 kernel 性能不如预期通常不是 Tensor Core 没发力而是旁边的数据通路拖了后腿。我排查得最多的三个方向是共享内存 bank conflict、寄存器溢出spill和线程块占用率occupancy过低。这三个问题有一个共同特点——不会导致结果错误所以更容易被忽略。排查思路是打开 Nsight Compute先看Memory Workload Analysis里的 LDS Bank Conflict 指标超过 5% 就值得优化。再看 Occupancy 那一栏如果实际值低于理论值大概率是共享内存或寄存器资源超了。最后看一眼 warp state 里的待命周期Stall如果大量周期是等待共享内存或全局内存数据回来说明流水线深度不足需要增加 stage。4.3 一个典型案例错误配置导致性能反而下降之前帮一个团队看一个 Transformer 推理引擎他们把 QKV 的 GEMM 全部切成 FP16FP32 基线大约 220 TFLOPS 的利用率切完反而只有 120 TFLOPS甚至不如基线。逐个参数排查后发现了两个问题叠加第一Tile 大小从 128x128 改成了 64x64因为共享内存里给双缓冲留的缓冲区被 FP16 数据塞得更大了看起来应该刚好能塞下但忘记了 CATLASS 还需要给共享内存预留一行 padding 和 stage 缓冲区结果 occupancy 掉了一半。第二warp 数量从 8 个减到 4 个后寄存器循环展开量不够全局内存加载频率变高流水线没法藏住延迟。修正方案很简单恢复 128x128 Tile把双缓冲 stage 从 4 降到 2给 padding 留出空间。改完性能直接回到接近 FP32 基线的 1.8 倍。这给团队的教训是混合精度不只是换数据类型它改变了一整条数据通路的资源占用面貌所有依赖资源的设计Tile、stage、寄存器数都要重新平衡。4.4 性能排查工具箱关键工具与基本使用流程我调 CATLASS kernel 时一般会按下面这个流程走用ncu --set full跑一次 Nsight Compute重点看 SOLSpeed of Light里的 Compute 和 Memory Throughput。如果两者都没超过 60%先假设 kernel 被延迟限制了看 warp stall 原因再往前推。如果 Compute 接近 90%、Memory 接近 90%那就是优化得很好的状态。接下来看 instructions 里的 LDS/UDS 数量、bank conflict 计数、global load/store 效率。最后用ncu --metrics launch__occupancy_limit_shared_mem,launch__occupancy_limit_registers确认资源瓶颈。数值精度排查方面写一个小工具脚本FP32 基线算一遍各层输出混合精度算一遍统计逐元素的绝对误差、相对误差、最大误差位置。通常相对误差在 1e-2 以内就可以接受如果超过 0.1需要定位是哪一层开始放大的。这个方法虽然土但在追踪精度在哪个算子被毁掉时非常有效。5. 从 GEMM 到更多算子模板库能力的延伸应用5.1 卷积的隐式 GEMMImplicit GEMM化CATLASS 的 GEMM 核心不止能算常规矩阵乘也是卷积算子的基础。卷积本质上是一个四维循环可以转换成矩阵乘形式把输入特征图按卷积窗口展开成矩阵im2col把卷积核展开成另一个矩阵两者相乘再重组。naive 的 im2col 需要额外拷贝数据内存开销大得吓人展开后数据量通常膨胀几十倍而隐式 GEMM 的思路是不真正展开数据而是在 GEMM 的 Tile Iterator 里直接根据卷积映射关系去取数据。这个思路在 CATLASS 里落地得很自然。Tile Iterator 本来就是按布局搬运数据的抽象你给它定义好第 m 行对应输入特征图的哪个位置、第 n 列对应卷积核的哪个位置它会自动在 GEMM 主循环里完成取数。这样做的好处是主循环计算逻辑完全复用 GEMM 的优化积累——Tiling、双缓冲、Tensor Core 指令全都是现成的。5.2 注意力机制算子的融合实践注意力机制里的 QK^T 和 Score·V 本质上也是矩阵乘定制的空间在于中间要穿插 Softmax并且要避免把整个 Score 矩阵写回全局内存。若完全按通用 GEMM 做QK^T 算完把 NxN 的 Score 矩阵写回全局内存再从全局内存读回来做 Softmax、再参与 PV 矩阵乘访存开销非常大。CATLASS 的做法是把两段矩阵乘和中间的 Softmax 融合成一个大的 kernel中间的 Score 留在寄存器或共享内存只借道 CPU 或小块内存做规约。这个方向最有名的实现是 FlashAttention它的核心贡献之一就是用分块思想避免大 Score 矩阵落回全局内存。CATLASS 用户完全可以照着这个思路实现一个定制化的 Attention kernelTile 迭代器照用 GEMM 的中间替换成融合了 Softmax 的 epilogue后处理。这种保留 GEMM 主循环、替换 epilogue的自定义模式是模板库最有价值的部分。5.3 自定义算子的复用模式与边界CATLASS 这类模板库给了使用者很大的自由但边界也要心里有数。适合复用的场景是算子主体是矩阵乘或点积类计算差异点在于取数方式和后处理逻辑。不适合复用的场景是算子数据依赖关系非常非线性比如复杂的稀疏路径、动态控制流极多的场景。这类算子用模板库强行套模板最后只能越套越复杂不如直接写 NVCC kernel。我给团队定的经验准则是如果算子中有超过 80% 的工作本质是计算一个输出 Tile 若干输入 Tile 的乘累加就用 CATLASS 改否则直接手写。手写时可以借用 CATLASS 里的分块、双缓冲思路但不需要强行上模板。很多团队项目失败是因为过早地想把所有算子都用模板库统一了结果收益没看到只看到编译时间和排查成本的上升。说到调试CATLASS 模式下的模板代码报错信息向来不友好。一个通用技巧是按C 模板类实例化展开的方式去读报错先把缺失的类型、常量推断出来再反查是哪个模板参数类型不匹配。并且确保编译时打开-G调试模式和-lineinfo选项这样反汇编和性能分析的定位会准确得多。关于这套模板库的实际使用我个人的体会是它确实不是最快让人出活的那条路初期学习成本比直接调 cuBLAS 高不少但一旦把 Tile 迭代器、流水线、双缓冲、混合精度这些概念吃透后续写任何高性能算子都会上一个台阶。最后再分享一个调试小技巧当 kernel 性能不符合预期时先别急着搜代码把 Nsight Compute 的 SOL 页面截图对照 cuBLAS 同规格算子的 SOL 图找差异区间通常一眼就能看出是计算瓶颈还是访存瓶颈省掉大量盲目尝试的时间。这套方法论比背任何库的 API 都值钱。