CUDA系统拆解-16-CUTLASS、Triton、cuBLAS 与 FlashAttention：高性能实现都在做什么

本文是「CUDA系统拆解」系列第 16 篇。
系列导读：CUDA系统拆解-00-导读：从编程模型到 AI 推理系统的学习路线
上一篇：CUDA系统拆解-15-Tensor Core、WMMA 与 MMA：矩阵乘指令路径怎么打通
下一篇：CUDA系统拆解-17-vLLM、TensorRT-LLM 与 Continuous Batching：CUDA 为什么最终连到推理系统

1. 这篇解决什么问题

这一篇不是教你背 API，而是帮你建立一张高性能实现路线图。核心问题有 5 个：

  • cuBLAS / cuBLASLt、CUTLASS、Triton、FlashAttention 各自在解决什么问题。
  • 为什么 GPU 上很多高性能实现都喜欢从 GEMM 思维出发。
  • 为什么 FlashAttention 的核心是 IO-aware，而不只是“把几个算子 fusion 起来”。
  • 什么时候应该优先调库，什么时候才值得自己定制 kernel。
  • 这些路线在 AI 推理工程里分别处在什么位置，有什么边界和取舍。

如果这篇只记住一句话，那就是：

高性能实现不只是“算得快”，更关键的是“数据怎么分块、怎么复用、怎么少搬、怎么把硬件峰值真正喂满”。

2. 先记住的核心结论

  • cuBLAS / cuBLASLt 代表标准高性能库路线，适合标准 GEMM 和它的常见变体。
  • CUTLASS 代表结构化构建高性能 GEMM / Tensor Core kernel 的路线，重点在分层组织和可定制性。
  • Triton 代表更高生产力的 GPU kernel DSL 路线，特别适合 fusion、特殊 layout、快速迭代的定制算子。
  • FlashAttention 代表 IO-aware 的专用算法与 kernel 共设计路线，核心是减少 attention 中间结果的显存读写。
  • 很多高性能实现都从 GEMM 思维出发，因为它天然适合 tiling、数据复用、shared memory staging 和 Tensor Core。
  • 真实工程里这几条路线通常不是互相替代，而是分工协作。

3. 正文讲解

3.1 四条高性能路线的总图

可以先把这四条路线粗略分成四类：

  • cuBLAS / cuBLASLt：标准高性能库
  • CUTLASS：高性能 kernel 的结构化积木
  • Triton：高生产力的定制 kernel 语言和编译器
  • FlashAttention：面向 attention 的 IO-aware 专用算法与实现

它们不是同一层的东西，所以不要直接问“谁替代谁”。更合理的问题是：当前这个瓶颈属于哪一类问题，应该用哪条路线解决。

3.2 为什么高性能实现常从 GEMM 思维出发

在 GPU 上，GEMM 是最重要的高性能母问题之一。原因很简单：

  • 计算模式规则，容易分块
  • 数据复用高，适合做 tiling
  • 算术强度高，容易把 GPU 算力压出来
  • 很适合走 shared memory + registers + Tensor Core 这条路径

而 AI 推理里很多看起来不完全一样的计算，最后都能尽量改写成 GEMM 风格：

  • 线性层
  • MLP
  • attention 里的 QK^T
  • attention 里的 PV
  • 一些卷积实现

所以高性能实现经常围绕这些问题组织：

  • 一个 block 负责哪一块输出
  • 输入 tile 怎么搬运
  • shared memory 怎么做 staging
  • registers 怎么保存局部累加
  • warp / Tensor Core 怎么参与计算
  • epilogue 是否顺手融合 bias、activation、scale

一旦你接受“很多高性能 kernel 本质上是在做 GEMM 风格的数据流设计”，后面看 cuBLAS、CUTLASS、Triton 甚至 FlashAttention 都会顺很多。

3.3 cuBLAS / cuBLASLt：为什么标准问题优先调库

cuBLAS 的核心价值不是“能算矩阵乘法”，而是它把标准线代问题做成了成熟、稳定、跨架构持续优化的高性能方案。

标准 GEMM 很难随手超越，原因不在于公式复杂，而在于工程细节极多：

  • tile 形状怎么选
  • layout 和对齐怎么处理
  • Tensor Core 路径怎么映射
  • shared memory 和 registers 怎么配合
  • prefetch、pipeline、epilogue 怎么组织
  • 不同 shape、dtype、架构下怎么选更优 kernel

所以工程上一个很实际的原则是：

如果问题本质上还是标准 GEMM 或 batched GEMM，先看 cuBLAS / cuBLASLt 能不能解决。

这里 cuBLASLt 特别重要，因为现代 AI workload 往往不只是“裸 GEMM”，还会关心：

  • bias
  • activation
  • layout
  • quantized path
  • 更灵活的 epilogue
  • heuristic 和 autotuning

这也是为什么很多推理系统真正依赖的不是最老的 cuBLAS 接口，而是更现代的 cuBLASLt。

3.4 CUTLASS：高性能 kernel 不是一段代码，而是一套层次结构

CUTLASS 的价值不在于“又多了一个 GEMM 库”，而在于它把高性能 GEMM / Tensor Core kernel 的组织方式拆成了可复用的结构。

最重要的不是具体模板，而是这套思路：

  • threadblock-level tile
  • warp-level tile
  • instruction-level tile
  • global memory -> shared memory -> registers 的数据路径
  • mainloop 与 epilogue 的拆分

换句话说，CUTLASS 在教你一件事：

高性能 kernel 不是“把 for-loop 搬到 GPU 上”，而是把计算层次、存储层次和执行层次一起对齐。

所以 CUTLASS 很适合两类人：

  • 想看懂高性能 GEMM / Tensor Core kernel 是怎么搭出来的人
  • 需要在接近库级性能的前提下做深度定制的人

它在工程里的定位通常是：

  • 不想从零赤手写所有底层细节
  • 但标准库表达力不够
  • 需要更细地控制 tile、epilogue、pipeline、layout

3.5 Triton：为什么它特别适合定制和 fusion

Triton 的核心不是“比 CUDA 更底层”，而是“用更高效的方式表达 tile 级数据流和定制 kernel 结构”。

它特别受欢迎，是因为现实里有很多算子处在一个尴尬区间：

  • 直接调库不完全合适
  • 手写 CUDA 成本高、调试慢、维护重
  • 但性能又不能退化太多

这正是 Triton 的甜点区：

  • fused operator
  • pointwise + reduction 组合
  • 特殊 layout 下的自定义 kernel
  • 快速试验不同 tiling 和 autotune 配置

它的一个关键优势，是很适合把多步逻辑收进一个 kernel，减少不必要的 global memory 往返。例如：

  • fused softmax
  • fused layernorm
  • fused activation
  • 一些 attention 周边算子

也就是说，Triton 的价值常常不只是“写起来更短”，而是更容易把你真正想要的 tile 级数据流写对。

但不要神化它。Triton 仍然吃性能基本功：

  • memory access pattern
  • register pressure
  • tile 设计
  • Tensor Core 映射
  • fusion 收益是否真的大于复杂度

所以真正强的不是“会写 Triton”，而是既懂 CUDA 性能原理，又会用 Triton 更快把这些原理落成代码。

3.6 FlashAttention：核心不是 fusion，而是 IO-aware

FlashAttention 和前面三者不完全同类，因为它首先是一条专门针对 attention 的算法与实现路线。

标准 attention 的直觉流程是：

  1. 计算 S = QK^T
  2. 对 S 做缩放、mask、softmax
  3. 再算 O = softmax(S)V

问题在于，S 往往是一个很大的中间矩阵。序列一长，这个中间结果的显存读写就会变得非常重。于是瓶颈很容易不在 FLOPs，而在 IO。

FlashAttention 的核心思想是：

  • 按块处理 Q / K / V
  • 尽量把中间状态留在片上路径里
  • 不去物化完整的 attention score 矩阵
  • 用 online softmax 保证分块处理时仍能得到精确结果

所以它真正优化的是：

  • HBM / global memory traffic
  • attention 中间状态的物化成本
  • 长序列下的 IO 压力

这就是为什么说它是 IO-aware。
它当然也会用到 Tensor Core、shared memory、warp 协作，但这些只是手段。真正的一层思想是：attention 不能只从“算”看，还必须从“搬”看。

3.7 四条路线不是竞争关系，而是分工关系

更实用的理解方式是：

  • cuBLAS / cuBLASLt：标准问题的首选
  • CUTLASS：需要理解和构建高性能 GEMM/Tensor Core 路线时的结构化基础
  • Triton：需要高生产力地写定制 kernel、尤其是 fusion 时的强工具
  • FlashAttention：当 attention 的核心矛盾在 IO，而不是普通 GEMM 本身时的专用路线

所以真实推理系统通常是混用的。例如：

  • 线性层主干走 cuBLASLt
  • 一些 fused normalization / activation / special layout kernel 用 Triton
  • 某些深度定制矩阵核参考或基于 CUTLASS
  • attention 主路径走 FlashAttention 或类似的 paged attention 变体

4. 和 AI 推理的关系

站在推理工程角度，这四条路线的差别，本质上是四种 trade-off：

  • 标准化 vs 定制化
  • 生产力 vs 极限控制力
  • 通用性 vs 算子特化
  • 算力峰值 vs IO 路径优化

在推理系统里，真正重要的不是“我会不会某个工具”，而是你能不能判断：

  • 当前瓶颈更偏 compute-bound 还是 memory-bound
  • 是标准 GEMM 问题，还是 fusion 问题，还是 attention 的 IO 问题
  • 是继续调库更划算，还是值得投入定制 kernel

这也是 AI infra / 推理岗位真正看重的能力：不是工具崇拜，而是代价模型和工程判断。

5. 常见误区

  • Triton 不是“自动比 CUDA 更快”，它只是更容易把某些定制 kernel 写出来。
  • CUTLASS 不是简单的“另一个调用库”，它更像高性能 GEMM/Tensor Core kernel 的结构化框架。
  • FlashAttention 不只是把几个算子融合在一起，它的核心是重构 attention 的数据流，减少 IO。
  • 不是所有问题都该自己写 kernel；标准 GEMM 问题优先调 cuBLAS / cuBLASLt 往往才是成熟工程选择。
  • 不是所有 attention 优化都等价于 FlashAttention；只有当核心矛盾真在中间状态物化和 IO 时，这条路线才最有价值。

6. 复习自测

  • 为什么说很多 GPU 高性能实现都能从 GEMM 思维出发？
  • cuBLAS / cuBLASLt、CUTLASS、Triton、FlashAttention 分别适合解决哪类问题？
  • 为什么标准 GEMM 通常优先调库，而不是自己写 CUDA kernel？
  • CUTLASS 最值得学的到底是某个模板参数，还是那套分层结构？为什么？
  • Triton 为什么特别适合 fusion 和快速迭代？它的边界又在哪里？
  • 为什么 FlashAttention 的核心是 IO-aware，而不是简单 fusion？
  • 在一个真实推理系统里，什么时候该继续调库，什么时候才值得走 Triton 或更深的定制路线？

系列导航