如果你已经在NVIDIA卡上做过几轮LLM推理性能优化大概率会有一种感觉瓶颈从来不在算力而在怎么把数据喂给Tensor Core。无论是MHA里的QKV投影还是FFN的两个大GEMM又或者是KV Cache相关的特殊算子底层都绕不开一个核心问题如何写出一个在特定GPU上跑满带宽和算力的矩阵乘法。CUTLASS这套模板库解决的就是这个问题。它不是一个pip install就完事的黑盒而是一套让GEMM、Conv、Attention算子作者能精确控制每一笔访存的C模板体系。这篇文章我会从源码结构、CuTe DSL、Collective编程模型、工程能力和AI推理落地几个维度做一次深度拆解。内容稍微长但看到最后你应该能自己判断什么时候该直接用cuBLAS什么时候该拿起CUTLASS手写。1. 源码地图CUTLASS 3.x的分层设计与目录导航1.1 为什么我不建议你从gemm.hpp开始读源码很多第一次接触CUTLASS的人会本能地打开include/cutlass/gemm/device/gemm.h然后在一千多行模板代码里迷失。这不是你的问题——CUTLASS 2.x的gemm.hpp本来就是面向模板特化设计的大量逻辑隐藏在platform/和thread/路径下读起来极其痛苦。CUTLASS 3.x转型之后源码结构已经清晰得多。顶层设计上整个库拆成三大块cute/CuTe DSL负责描述张量布局、线程映射和数据搬运是整个3.x的核心抽象cutlass/gemm/GEMM相关的高层组合包括kernel/、collective/、tile_scheduler.hpp等cutlass/epilogue/输出阶段的collective实现处理累加器的缩放、偏置、激活函数和写入。如果你现在拿到3.7的主线代码我建议的阅读顺序是先看cute/layout.hpp和cute/tensor.hpp把Layout和Tensor这两个概念吃透再去读cutlass/gemm/collective/下面针对sm90的几个Collective实现最后才回头看gemm_universal.hpp因为那个文件只是把不同Collective粘到一起的壳。1.2 三个核心层Kernel、Collective、CuTeCUTLASS 3.x的架构本质上是三层结构第一层是CuTe它定义了Lambda、Layout、Tensor、MMA_Atom、Copy_Atom这些基础类型。你可以把CuTe理解成一套数据排布的数学语言它不关心你是做GEMM还是做Attention只负责回答一个问题给定一个逻辑坐标系如何映射到全局内存、共享内存或寄存器文件的物理坐标。第二层是Collective它在CuTe的基础上封装了具体的计算和搬运策略。比如CollectiveMma处理主循环里的矩阵乘法CollectiveEpilogue处理累加器写回。这一层会用到sm90的TMA、wgmma也会用到sm80的cp.async和mma.sync。第三层是Kernel层。它只做两件事决定Tile的调度方式比如Simple、Persistent还是StreamK然后调用Collective完成实际计算。gemm_universal.hpp里的CUTLASS_DEVICE入口函数基本就是把这几个组件按顺序排起来。这种分层最大的好处是你不需要为了测试一个新的Mainloop去重写整个GEMM。只要CuTe里描述的布局对Collective能编译过Kernel层基本不需要动。1.3 读懂源码前先理解的三个硬件事实直接读CUTLASS源码之前必须先确认自己理解下面三件事不然代码看起来就是天书第一Tensor Core是有固定 shape 的。以sm80的mma.sync.aligned.m16n8k8为例一条指令一次算16行8列8深数据必须按特定布局放在线程的寄存器里。CUTLASS里的MMA_Atom就是把这些指令包成一个个原子CuTe的make_tiled_mma再把原子在warp内重复排布覆盖更大的块。第二sm90的TMATensor Memory Accelerator是一台独立于线程的异步拷贝引擎。它可以通过一个描述符把全局内存里的数据批量搬进共享内存完全不需要线程一条条访问。CUTLASS 3.x里大量实现在MainloopSm90TmaGmmaWarpSpecialized这样的名字中核心配置就是由warp specialized分工Producer warp负责TMA搬运Consumer warp负责wgmma计算。第三sm90引入了线程块簇Cluster多个CTA可以共享TMEM和DSMEM空间。CUTLASS的ClusterShape模板参数就是在描述这个协作粒度。看懂这三个硬件特性你再看CUTLASS源码时就不会总问为什么要写得这么复杂了。2. CuTe DSL原理Layout代数把数据排布变成了可计算对象2.1 Layout不是数组是函数Shape Stride的世界大多数从PyTorch或NumPy过来的开发者对张量的直觉是多维数组而CuTe的Layout完全换了一个心智模型Layout本质上是一个从逻辑坐标到内存偏移的纯函数。using namespace cute; // 一个128行64列的row-major矩阵 auto layout LayoutShape_128, _64, Stride_64, _1{}; // 把ptr包装成一个Tensor逻辑形状和物理布局绑定 float* ptr ...; auto tensor make_tensor(ptr, layout);Shape_128,_64表示逻辑维度是128和64Stride_64,_1表示第0维步长64、第1维步长1。坐标(i, j)对应的内存偏移就是i * 64 j * 1。这个模型看起来朴素但它的强大之处在于可以表达各种奇怪的排布交错数据、pad后的共享内存、按转置读取等。在CUTLASS里的一切计算本质上都是Layout的组合与变换。Shared Memory里的一块张量其实就是一个指针加上一个Layout寄存器里的一个Tensor也就是一个arrayfloat, N绑上一个Layout。把这个想法贯彻到底整个GEMM就变成了几个张量之间的搬运和计算而CuTe的算法会帮你自动推导每个线程在哪个寄存器上操作。2.2 布局操作与TiledMMA怎么把一个大GEMM切给warpCuTe最核心的价值是把MMA如何在warp内展开这个难题变成了几个Layout操作。先看一个典型例子。假设你的Tile是128x128每个warp要算其中一小块。你首先需要local_tile从大Tensor里切出一个子块然后用make_tiled_mma把一个MMA_Atom和一份warp映射组合起来得到一个TiledMMA。// sm80上的一条fp16 MMA原子 using MMA_Atom MMA_AtomSM80_16x8x8_F16F16F32F32_TN; // 把原子在warp维度重复8次覆盖更大的形状 auto tiled_mma make_tiled_mma( MMA_Atom{}, LayoutShape_8, _2, Stride_2, _1{} // 当前示例中的线程映射 );这个LayoutShape_8,_2, Stride_2,_1就是CuTe里最神奇的部分它决定了32个线程如何映射到MMA原子内部的Tile坐标。tiled_mma随后可以通过partition_A、partition_B和partition_C自动把共享内存里的A、B矩阵和累加器切分成每个线程负责的寄存器片段。你需要做的事只剩三步第一用local_tile切出共享内存中的Tile第二用tiled_mma.partition_*取出当前线程的寄存器张量第三调用mma.callback或gemm循环完成累加。背后的坐标推导和线程对应关系全部由CuTe的Layout代数完成。这也是为什么CUTLASS 3.x的代码看起来比2.x更数学它不在代码里到处写线程ID与行列坐标的if判断而是通过Layout的计算自动完成映射几乎找不到需要手工取模的代码。2.3 进阶组合、补全、互逆与coalesce这几个操作解决了什么问题CuTe里有一组非常值得深入研究的Layout操作读懂了它们你才算真正理解CuTe DSL。第一个是composition布局组合。它的作用是把两个布局串成一个先经过第一个布局做一个坐标变换再经过第二个布局做内存映射。在partition_A等操作里composition是底层的主力它能把从大Tile中取局部块和按线程摊位映射到寄存器这两个步骤无缝衔接。第二个是complement补布局。补布局是生成一个双射的关键手段。假设你已经定义了一个从逻辑坐标到寄存器坐标的映射你还需要一个补齐的布局来描述剩余那些线程或内存应该如何排列。很多看起来很绕的源码其实都是在算某个Layout的complement。第三个是coalesce合并维度。它把相邻的连续维度合并成更大的步长在分析内存访问是否连续时非常有用。做性能优化时经常需要检查一个全局内存Tensor经过多次partition之后在内存地址上到底有没有保持相邻线程访问相邻地址用coalesce一眼就能看出来。我之前在Jetson Orin上调一个自定义算子时就因为忽视coalesce导致全局内存访存路径呈跳变状态带宽利用率只剩三成。后来用CuTe把布局打出来一看Stride完全不连续改完Layout定义后性能直接翻倍。2.4 CuTe不止做GEMM它是张量计算的编译器基础设施很多读者容易把CuTe局限在GEMM的辅助工具这个印象里这是对CuTe DSL最大的误解。CuTe实际是一个通用的张量映射算法库。它并不在意你的计算到底是GEMM还是FlashAttention也不在乎你最后是调用wgmma还是普通的ldmatrix。只要能描述成从全局内存搬数据到共享内存再从共享内存搬到寄存器再算一次CuTe的原子和布局操作就都可以用。官方仓库里的cute/atom/mma_traits_sm90.hpp、cute/atom/copy_traits_sm90.hpp分别封装了sm90的MMA指令和TMA拷贝指令。你在自研Attention、MOE、甚至非矩阵类算子时依然可以复用CuTe来做数据排布只是把计算替换为你自己的CUDA C逻辑。这也是CUTLASS 3.x和2.x一个很大的区别3.x不再是一个只会做GEMM的库而更像一个面向GPU算子作者的领域特定语言和基础设施。理解了这一点你在做AI推理自定义内核时就能把CuTe用得更顺手。3. Collective编排层Mainloop、Epilogue与TileScheduler怎么配合3.1 从collective_mma.hpp看主循环的四个阶段在CUTLASS 3.x里一个GEMM主循环Mainloop的完整执行可以被拆成四个阶段。第一阶段是全局内存到共享内存的搬运。对sm90来说这一步通常由TMA发起只需要几行代码构造一个copy操作。TMA描述符在kernel开头创建一次后续每一轮迭代只需要更新偏移量数据就到了共享内存。第二阶段是共享内存到寄存器的加载。在非sm90路径或特殊布局下这一步由ldmatrix或普通ld.shared完成在sm90的wgmma路径下共享内存数据甚至可以直接被wgmma指令消费减少了中间寄存器拷贝。第三阶段是MMA计算。CUTLASS把连续的若干次MMA组合成一个GMMA操作在warpgroup级别执行。CollectiveMma会维护一个寄存器累加器数组每次迭代把共享内存中的A、B分块乘进去。第四阶段是累加器的更新与传递。累加器经过多轮K迭代后最终需要交给Epilogue。这里有一个非常容易被忽略的问题Mainloop和Epilogue之间累加器到底是留在寄存器里还是先写回共享内存还是通过TMEM中转——这个决策直接决定了性能上限也是MainloopSm90TmaGmmaWarpSpecialized这类实现中大量barrier和warp specialization代码存在的意义。理解这四个阶段后再读collective_mma.hpp就会轻松很多因为代码的注释和函数命名基本就是围绕这条流水线写的。3.2 Epilogue不是写回Global Memory这么简单很多人在研究CUTLASS时会主攻Mainloop却轻视Epilogue。但对于AI推理来说Epilogue往往才是最耗时、最需要灵活性的部分。现代推理算子几乎不会直接输出累加器的原始int32或fp32结果。要么需要乘以一个scale再转成fp16/bf16要么需要加上bias再做SiLU或ReLU要么需要做GELU、Softmax之外的更复杂融合。CUTLASS 3.x的CollectiveEpilogue提供了统一的Visitor机制允许你在输出前插入多个变换步骤。我见过最典型的一个优化案例是一个模型的前向里有GEMM加LayerNorm的操作。传统做法是GEMM算完把结果写回全局内存再读出来做LayerNorm。用CUTLASS的Epilogue则可以直接在累加器上完成部分归约和缩放然后把结果写回。表面上看只是省了一次全局读写实测在A100上端到端算子延迟能降20%以上。所以在做推理落地时不要只盯着GEMM本身认真想想哪些算子可以融进Epilogue。这是CUTLASS带来的额外优化空间。3.3 TileSchedulerStreamK和Persistent调度决定性能下限与上限GEMM性能不仅取决于单个Tile算得快不快还取决于整个GPU上的Tile调度得是否均匀。CUTLASS 2.x的经典实现是CTA按M、N格子切分每个CTA算一个Tile遇到M很小、K很大的形状时会有大量CTA结束后SM空转。CUTLASS 3.x引入了tile_scheduler.hpp里面至少有三类调度器值得关注普通Tile调度每个CTA只负责一个固定的输出Tile实现简单适合形状规整、M足够大的场景Persistent流水线调度CTA以Persistent方式驻留在SM上通过原子索引动态领取下一个Tile减少了CTA启动和回收的开销StreamK调度在K维度上把工作切碎多个CTA并行处理同一输出Tile的部分K通过原子加做归约。这种调度特别适合小M、大K的Decode场景。实际推理落地时我建议不要一开始就追求StreamK。它虽然能解决负载均衡问题但原子归约和最后阶段的累加会带来额外的精度和调度风险。先用普通调度把基线跑通再对照Profiling结果决定要不要上StreamK。选调度器时还有一个很容易忽略的点TileScheduler的模板参数会改变kernel的全局状态布局。同一个CollectiveMma换调度器后共享内存占用和寄存器用量都可能变。所以调度和Tile尺寸必须一起sweep不要单独调某一个。4. 工程能力评测CUTLASS真正拉开差距的地方4.1 构建系统、examples和测试框架的组织方式CUTLASS的工程能力在开源算子库里是第一梯队。这不光是代码能跑层面的工程能力而是大规模模板库可持续维护的工程能力。首先是构建系统。CUTLASS使用CMake但与普通项目不同它编译的不是一个可执行文件而是大量测试和example。官方提供了一套CMakeLists.txt和tools/library/scripts用来生成不同架构、不同数据类型的GEMM配置组合最后封装成libcutlass.so或libcutlass.a。你甚至可以只生成自己关心的几个kernel节省大量编译时间。其次是examples。examples/目录下每个子目录都是一个小而完整的独立示例从最简单的00_basic_gemm到sm90的50_hopper系列。我第一次接触CUTLASS 3.x时就是照着examples/47_hopper_gemm_with_collective_builder跑通的。这些example不是玩具它们会直接调用CUTLASS不同层次的API告诉你完整的调用链长什么样。测试框架方面CUTLASS维护了test/目录单元测试覆盖了CuTe的布局运算、各个架构下的MMA指令、Collective和Epilogue的组合。更关键的是CUTLASS的测试里有一个与cuBLAS对比的reference模块你在自研kernel时可以直接复用它的验证逻辑。我建议你把test/unit/cute和test/unit/gemm下的代码当作额外的文档来读很多API的边界条件写得很清楚。4.2 Profiling工具链sweep、JSON、与Nsight Compute的配合CUTLASS在tools/profiler下自带一个cutlass_profiler这是我在调GEMM时最常用的工具。它能自动遍历你指定的多个配置组合跑一轮并输出Kernel时间、TOPS、带宽利用率等指标。例如./cutlass_profiler --kernelssm90_xmma_gemm_f16f16_bf16_bf16_f32_tn_f16_tensor_op_f32 \ --m4096 --n4096 --k4096 \ --m-shape128,128,256 --n-shape128,256,512这个Profiler的价值在于它能快速帮你建立不同TileShape、不同ClusterShape下性能如何变化的全局视图而不是靠猜。拿到初筛结果后再用Nsight Computencu做单kernel深挖。重点看四类指标访存指标SMEM bank conflict、global load命中率、TMA带宽利用率计算指标Tensor Core pipe利用率、MMA每周期发射数调度指标warp stall原因、barrier等待时间、Occupancy是否被SMEM限制发射指标issue slot利用率、Control flow overhead。CUTLASS的kernel在ncu下有非常好的symbol名和源码行号映射基本能做到哪一行代码导致哪个stall的定位。这一点比很多闭源库或手写kernel强太多。4.3 源码质量与迭代速度什么样的社区生态在支撑它我长期跟踪CUTLASS的commit历史后一个直观感受是它其实是在为NVIDIA硬件做算子落地的工程验证。每次新卡发布前CUTLASS都会提前拿到新指令集的支持并在主线里更新。3.x主线的源码质量可以概括为注释密度高、模板边界清晰、但阅读曲线陡峭。很多核心模板的注释甚至比代码还长重点解释了这个模板参数组合对应的是哪个硬件功能。尤其要提的是CUTLASS对sm90和Blackwell的代码路径会明确标注这是为哪一代卡优化不会默认所有架构走同一份代码。不过也要说句公道话CUTLASS的代码风格并不是所有人都会喜欢。因为它大量使用CRTP、tag dispatch、ndexpr和折叠表达式新手阅读时会有很大的心理负担。但一旦你习惯了模板元编程的表达方式你会觉得这套体系远比密密麻麻的CUDA C代码更易于维护。对于AI推理团队来说与其把CUTLASS看成一个拿来即用的算子库不如把它看成一个拥有NVIDIA官方维护的、可持久跟随硬件更新的算子开发框架。这意味着你的自定义内核在下一代GPU出来时有更大概率通过升级CUTLASS版本而不是重写来获得支持。5. AI推理落地从LLM算子到生产后端的完整链路5.1 PreFill与Decode的两种GEMM特征差异LLM推理环境下的GEMM和你在Profile里看到的大矩阵乘法是两回事。如果你直接拿固定形状的大矩阵测CUTLASS性能实战时大概率会失望因为真实推理的GEMM形状分布极其不均匀。PreFill阶段的特点是M等于序列长度和batch的乘积可能从几百到几千K和N对应隐藏层宽度通常在4096或8192。这个阶段单个GEMM的形状仍然比较大Tensor Core利用率能撑起来。优化重点放在TMA流水线的stage数量、TileShape与M的匹配以及如何在多buffer之间隐藏TMA延迟。Decode阶段则是完全不同的世界。M通常只有1到几十实际计算退化成GEMV。这种情况下直接跑一个标准GEMM kernel是极度浪费的。一样的Tile大部分位置都在搬运空气。这时候你要么修改GEMM的调度让一个CTA处理多个输出行要么在K维度上做split和reduce要么把逻辑转成批量GEMV来减少并行浪费。我在实际项目里PreFill阶段能稳定拿到cuBLAS的85%到95%性能但Decode阶段如果不做定制往往只能拿到50%左右。所以如果你要自研推理算子第一个要问自己的问题就是我的场景是PreFill多还是Decode多两者对kernel的要求差得很远。5.2 量化、窄类型与现代TCINT8/FP8在CUTLASS里的用法AI推理落地绕不开量化。它在CUTLASS里的体现主要是窄类型GEMM 对应缩放因子的组合。CUTLASS 3.x对sm90的FP8支持很有代表性。它分别提供E4M3和E5M2两种FP8格式的MMA原子并在Epilogue里集成了从FP32累加器转成FP8输出的逻辑。实际使用中你需要额外处理scale对齐因为FP8的指数范围有限A和B的scale通常要单独计算再转化到全局的scale里。INT8方面CUTLASS同样保持了对sm80以来Per-Column和Per-Output量化的支持。很多推理框架里的Weight-Only量化比如AWQ和GPTQ后端都依赖这类能力权重被预先量化和重排激活在运行时使用用很小的额外开销完成反量化。我在整合INT8 GEMM到推理管线时一个重要的教训是不要只盯着kernel的TOPS还要看两头的转换开销。权重量化后如果不在kernel启动前做re-orderkernel内部会频繁命中bank conflict或需要额外的通道shuffle。CUTLASS的CutlassLayout提供了预置的卷积权重重排工具但很多人根本不知道要用它。量化模型的精度验证又是另一个坑。建议在集成前先用一组固定的输入跑一遍reference fp16结果再跑量化后的结果拉出每个输出通道的最大绝对误差。如果误差集中某几个通道大概率是scale选择有问题而不是GEMM kernel本身的问题。5.3 决策边界何时直接用cuBLAS/TensorRT何时必须自研很多新读者热情高涨觉得有了CUTLASS就能自己写GEMM超过cuBLAS。实际上大多数场景直接调cuBLAS或TensorRT是更理性的方案。我把自己的决策逻辑总结成四个问题你的算子形状是否足够非常规比如变长序列、稀疏模式、特殊融合逻辑cuBLAS没有一个直接对应的API。这时候值得自研。你是否需要对多核/多卡做算法级协同比如跨苏格块的Attention、多查询聚合的MOE路由这种算子之上的控制流用CUTLASS写更自然。你是否需要和已有的推理引擎深度绑定比如TensorRT-LLM的自定义plugin你需要在它的框架约束下写算子这时候自研内核往往比塞进cuBLAS更灵活。你的团队是否有持续的Tuning时间CUTLASS不是一次写完就完事的你需要针对不同GPU型号做TileShape和调度策略的适配。如果只做一次性优化用cuBLAS性价比反而更高。对我而言CUTLASS最适合的位置是为那些cuBLAS覆盖不到、或者融合需求极强的算子提供一个可控的基础设施而不是替代cuBLAS本身。同样如果你的后端框架已经用了TensorRT-LLM或NVIDIA NIM它底层可能已经调了CUTLASS或cuBLAS你再去重写反而会失去框架层面的优化联动。5.4 部署实战清单环境检查、精度验证、性能三步最后给你一份我在生产环境做CUTLASS算子落地时的检查清单每条都是踩过坑之后总结的。环境检查这一关很多人会忽略。编译器版本、CUDA版本和GPU驱动版本必须匹配。比如sm90的某些TMA指令需要CUDA 12.x以上和对应驱动版本如果你在Jetson Orin这类嵌入式设备上构建还要确认JetPack SDK自带哪一版CUDA是否支持你需要的架构。曾有人问为什么nvidia-smi都报错显示驱动通信失败这种环境问题解决之前任何CUTLASS内核都不用谈。然后是精度验证。不要只看平均误差小。在GEMM场景下最危险的错误是某些K维度分块的累加顺序改变导致个别输出通道出现无法接受的偏差。我建议用三组数据验证全零输入、全一输入、随机正态分布输入。全零和全一能快速暴露Index或Layout的错位随机输入则测试数值范围。性能验证方面做两次就够了但每次都要和正确的基线比。第一次和cuBLAS比看是否达到期望的量级第二次和相同CUTLASS内核但不同Tile配置比做一轮小规模sweep。注意sweep时要固定L2 cache影响多次跑同一kernel取中位数或最小值避免被随机噪声误导。最后是集成阶段。CUTLASS内核一旦被放进推理服务它就和普通CUDA kernel一样需要处理stream、多卡、内存池和跨kernel间的依赖。一个容易被忽略的点是如果多个推理请求并发kernel的SMEM占用会直接影响同一SM上能并发几个block这会反过来影响端到端吞吐。所以最终部署前一定要在满负载并发场景下重新做一次性能验证不要只信单kernel的Profile报告。我在实际项目中见过太多次这样的场景单测一切正常Profiling数据非常漂亮一上线并发负载就崩。原因就是TileSize和共享内存分配只考虑了单kernel工况没有考虑多个block在SM上的并发竞争。最后说点个人体会做算子优化的这些年CUTLASS给我的最大帮助不是直接拿来用的GEMM而是让我的脑子里有了一套如何把计算映射到硬件的思维框架。CuTe的Layout代数让我在写任何高性能kernel之前都会先思考数据的物理排布Collective分层让我把搬运和计算当成两个独立问题去设计StreamK和Persistent调度让我明白算法层面的负载均衡和硬件层面的指令流水一样重要。如果你刚开始接触CUTLASS不必急着追新鲜的内核实现。先把CuTe的Layout概念吃透再对照examples/00_basic_gemm和examples/47_hopper_gemm_with_collective_builder把流程跑通。最后再用cutlass_profiler做一轮自己的sweep对比一下官方默认配置和cuBLAS的差异。这不只是一次源码阅读更像是一次对GEMM到底是怎么在GPU上跑的完整祛魅。等你丢下gemm.hpp能用自己的话解释L2 Cache residency、TMA pipeline和warp specialization之间的关系恭喜你你已经真正走出调用框架的阶段进入了设计算子的阶段。