CUDA → Triton → Liger-Kernel:GPU 编程的三层抽象

第 1 课 · 从手工雕琢到一行加速——每一层用多少"控制权"换多少"生产力"

这一课的唯一目标(tangible win)
学完后,你能判断一个给定的 GPU 优化需求该选哪层:手写 CUDA kernel、写 Triton 程序、还是直接 pip install liger-kernel 一行搞定。 也能对"为什么 torch.compile 生成的 kernel 是 Triton 而不是 CUDA"有清晰的解释。 如果你还没接触过 CUDA 的线程层级(grid/block/thread/warp),建议浏览 第 2 课:线程层级 建立心智模型——本课把 CUDA 作为底层起点来展开。

1. 三层抽象栈:从下往上,每层让渡控制权换生产力

GPU 上跑的任何计算最终都是CUDA kernel。但并非每个人都直接写 CUDA。 就像没人用汇编写 CRUD 应用——上层抽象让程序员专注于"做什么"而非"怎么做"。 下面是把三层叠在一起的全景图。从上往下是调用链;从下往上,每层拿控制权换生产力。

第 4 层 · 你主要接触的
HuggingFace / PyTorch 模型代码
model = AutoLigerKernelForCausalLM.from_pretrained(...) —— 一行 patch 掉底层,用优化 kernel 替换默认实现。完全不需要懂 GPU。
▲ 调用 / monkey-patch
第 3 层 · liger-kernel
Triton kernel 库(预写好的、精心调优的)
RMSNorm、RoPE、SwiGLU、CrossEntropy、FusedLinearCrossEntropy……每个都是优化过的 Triton 程序,带着 fusion/chunking 技巧。LinkedIn 开源,一行 apply_liger_kernel_to_llama() 生效。6
▲ 用 Triton 语言写的
第 2 层 · Triton
Python DSL + 编译器 → 生成 CUDA 代码
@triton.jit 写 kernel,block 级编程。编译器自动生成多线程调度、shared memory 分配、autotuning。torch.compile(Inductor) 的输出默认就是 Triton kernel。5
▲ 编译生成
第 1 层 · CUDA
C++ + NVIDIA 扩展 → GPU 硬件指令
__global__ + <<<>>> + thread/block/grid + shared/global memory。你控制一切——线程怎么排布、数据怎么搬、warp 怎么同步。这是第 2 课(线程层级)的核心内容。

关键洞察:你在这条栈里怎么定位,取决于你的目标。 如果目标是优化一个已知的 LLM 算子(RoPE、RMSNorm),liger-kernel 已经帮你写好了——直接用。 如果目标是想理解 torch.compile 是怎么工作的,你需要知道它背后的 Inductor 默认输出 Triton 而非 CUDA C++。 如果目标是写出比库里带的内核更快的新算法——那就自己写 CUDA 或 Triton。

2. CUDA:最高控制 = 最高成本

在第 2 课你会写一个让每个线程报"我是谁"的 kernel——代码很简单,但已经暴露了 CUDA 编程的所有核心复杂度

你控制什么如果搞错了
grid/block 大小(<<<B, T>>>occupancy 低下,SM 里一堆空闲 slot
shared memory 分配与同步(__syncthreads()数据竞争、死锁、bank conflict
全局内存访问模式非合并访问 → 带宽利用率从 ~90% 跌到 ~10%4
线程索引计算公式越界读写、算错负责的数据
warp 内分支路径warp divergence → 一半线程闲着

在 CUDA 的世界里,写一个"正确"的 kernel 只是及格线——写出"快"的 kernel 需要额外的 coalescing、shared tiling、register blocking、occupancy tuning……每一项都是独立的知识点。 Simon Boehm 的 matmul 优化日志一路从「一个线程算一个元素(naive)」优化到了 「1.3 TFLOPs 的缓存块 + 向量化」4——那是对 同一个数学运算十几轮迭代

量化感受
一个"能用"的 CUDA matmul kernel 大约 30 行;一个"接近 cuBLAS 速度"的 tuned matmul kernel 大约 300-500 行——且每一行背后都是对硬件行为的精细理解。4

3. Triton:把线程管理交给编译器

Triton(OpenAI 的开源项目,Philippe Tillet 等人 2019 年提出5)的核心创新: 你不需要为每一个线程写代码。

先看同一个 vector add 操作的 CUDA vs Triton 的代码量对比:

CUDA 版本(手动管理线程)
__global__ void vec_add(
  float *x, float *y, float *o, int n
) {
  int i = blockIdx.x * blockDim.x
        + threadIdx.x;
  if (i < n) o[i] = x[i] + y[i];
}

// 你还要手动: cudaMalloc
// cudaMemcpy, launch config,
// block size 选多少……
Triton 版本(描述 block 级逻辑)
@triton.jit
def vec_add(x_ptr, y_ptr, o_ptr,
            n_elements,
            BLOCK_SIZE: tl.constexpr):
  pid = tl.program_id(0)
  offs = pid*BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
  mask = offs < n_elements
  x = tl.load(x_ptr + offs, mask=mask)
  y = tl.load(y_ptr + offs, mask=mask)
  tl.store(o_ptr + offs, x + y, mask=mask)

3.1 Triton 替你干了什么

CUDA 里你手写的Triton 编译器自动搞定
block/thread 维度选择(128? 256? 320?)autotuning 自动搜索最优配置5
shared memory 分配 + __syncthreads()编译器自动插入——你只用标记哪些数据该放 shared
全局内存 coalescing编译器分析访存模式自动对齐
边界条件 if (i < n)mask=mask 自动处理越界
寄存器分配、指令调度LLVM 后端(Triton 编译链 → LLVM IR → PTX)5

核心哲学:你用较粗粒度的 program(概念上类似 CUDA 的 block)而非 thread 视角来写 kernel。 tl.program_id(0) 是"第几个程序"(类似 blockIdx),tl.arange(0, BLOCK_SIZE) 是"这个程序里的向量化工作单元"。 每个 program 内发生了什么——线程怎么切、数据怎么搬——编译器接管。7

torch.compile 的默认选择
PyTorch 的 Inductor backend 默认输出 Triton kernel,不是 CUDA C++。因为 Triton 在"开发效率"和"性能接近手写 CUDA"之间打了最甜的点—— 写 Python、自动 tuning、输出 PTX。这正是为什么 AI 编译器系列(ai-compiler)的第 2-3 课反复提到 Triton。

3.2 Triton 不完美——你得知道它的边界

一句话总结:手写 CUDA = 你当"硬件管理员";写 Triton = 你当"技术经理"——告诉编译器"这块数据要在 shared memory 上 tile",编译器去调度工人(thread)。

4. Liger-Kernel:连 kernel 都不用写

如果你已经理解了 CUDA(手写一切)和 Triton(编译器管线程),那 liger-kernel 就是一个 "请人帮你写了 Triton kernel" 的库。

LinkedIn 团队发现:LLM 训练中有几个高频算子——RMSNorm、RoPE、SwiGLU、CrossEntropy——占用了 大量显存和时间,但大家都在重复实现。于是他们把最优 Triton 实现开源了:

+20%
多 GPU 训练吞吐提升6
-60%
GPU 显存用量降低6
1 行
接入 HF 模型的代码量

4.1 它做了什么 CUDA/Triton 做不到的事

注意:liger-kernel 的 kernel 仍然是 Triton 写的。它的价值不在于"写了一种新语言",而在于 把 LLM 训练中最常见的优化模式封装好了。具体技法:

技法做什么为什么能省显存
Kernel Fusion 把 线性层 + CrossEntropy 合并成一个 kernel 中间结果不写回 HBM,在寄存器 / shared memory 里直接消费6
In-place 操作 RMSNorm 直接覆盖输入 tensor 省掉一份中间 buffer 的显存分配
Chunking 大 tensor 切成小块逐块算(像 FlashAttention 的 tiling) 峰值显存只等于一个 chunk 大小,不是整个 tensor6
Recomputation 前向时不存中间激活,反向时重算 用计算换显存(和 FlashAttention 同款思路)

和 FlashAttention 的相似点:FA 的核心优化——tiling + online softmax + recomputation——和 liger-kernel 的 chunking + in-place + fusion 是同一种哲学:在 kernel 里多做几步,少来回读 HBM。 FA 是 attention 专用优化;liger-kernel 把同样的思路搬到了 RMSNorm、RoPE、CrossEntropy 这些非 attention 算子上。6

4.2 怎么用:极限简单的接口

# === 方式 1:一行 monkey-patch(最常用)===
from liger_kernel.transformers import apply_liger_kernel_to_llama
apply_liger_kernel_to_llama()  # 替换 loss/RMSNorm/RoPE/SwiGLU
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3-8B")

# === 方式 2:AutoModel 自动 patch ===
from liger_kernel.transformers import AutoLigerKernelForCausalLM
model = AutoLigerKernelForCausalLM.from_pretrained("meta-llama/Llama-3-8B")

# === 方式 3:手拼 kernel(需要定制时)===
from liger_kernel.transformers import LigerFusedLinearCrossEntropyLoss
loss_fn = LigerFusedLinearCrossEntropyLoss()
loss = loss_fn(lm_head.weight, input, target)
loss.backward()

依赖极简:只需要 torch + triton(后面就是 CUDA),没有第三方库。8 这意味着你的训练环境只要装了 PyTorch + Triton,装 liger-kernel 就能用。

5. 三选一决策树:你的场景该用哪一层

  1. 这个算子的优化 kernel 已经有人在 liger-kernel(或类似库)里写好了?
    → 直接用 liger-kernel。不用自己写任何 kernel。
  2. 不在库里,但这是我独有的算子(比如 Open-Qwen2VL 的自定义 attention 变体)?
    → 用 Triton 写。block 级编程 + autotuning,开发速度远快于 CUDA。性能通常能达到手写 CUDA 的 80-95%。5
  3. Triton 的性能还不够(指标里卡着,差了 5-10%),或者 Triton 不支持某种硬件特性?
    → 手写 CUDA。这时候你对 warp-level 调度、register 分配、async copy(如 cp.async)有精确需求。

现实是:绝大多数人的优化需求落在第 1 或第 2 条。第 3 条(必须手写 CUDA)通常是 NVIDIA 自己、FlashAttention 核心团队、或者硬件厂商在维持。 这也是为什么"学 CUDA 是为了理解,但用的主要是 Triton / liger-kernel"是一个务实的学习策略。

6. 常见误区

误区 ① liger-kernel 是 Triton 的替代品
不是。liger-kernel 用 Triton 写的,不替代 Triton——它是"预写好的 Triton kernel 集合"。当你的算子在 liger 里没有时,你还是要自己写 Triton。 这三层是相互依赖的栈,不是三个"选一个"的平行选项。
误区 ② Triton 生成的代码一定比手写 CUDA 慢很多
对大多数 element-wise + tiled reduction 类算子,Triton 能达到手写 CUDA 80-95% 的性能。5 差距主要出现在需要精细控制 warp 内寄存器分配或 async copy 的场景——这些场景 Triton 的抽象层还不能完全覆盖。
误区 ③ torch.compile 出来的一定是 CUDA kernel
torch.compile 的默认后端 Inductor 生成 Triton kernel,不是 CUDA C++。 Triton 再把 kernel 编译成 CUDA(LLVM IR → PTX → SASS),但你实际"看到"和调试的是 Triton 层的输出(TORCH_LOGS 可以看到生成的 Triton 代码)。 详见 AI 编译器 · 第 2 课
误区 ④ liger-kernel 的每次调用都启动一个新 kernel
恰好相反。它的核心价值就是 fusion——把原来几个 kernel 合并成一个(如 FusedLinearCrossEntropy), 减少 kernel launch 和 HBM 往返。这也是它省显存+提速的根源。

📖 首选精读(primary source)

Liger Kernel 技术报告(arXiv:2410.10989)

这篇 17 页的技术报告精确解释了 kernel fusion、chunking、in-place 三大武器的工作机制,以及它们在不同模型(LLaMA-8B)上带来 +20%/−60% 的 benchmark 数据。读它,把本课的表格对照着看。

→ arxiv.org/abs/2410.10989

Triton 官方文档 · 第 1 章:Introduction

Triton 文档的开篇介绍,用 tiled matmul 为例展示了"用 CUDA vs 用 Triton"的代码差异。读完你会直观感受到为什么 block 级编程是一种思维升级。

→ triton-lang.org/programming-guide

引用与延伸

  1. NVIDIA, An Even Easier Introduction to CUDA — CUDA 基础(与本课对照看 CUDA 层)。
  2. NVIDIA, CUDA Refresher: The CUDA Programming Model — grid/block/thread 官方定义。
  3. Simon Boehm, How to Optimize a CUDA Matmul Kernel — 展示手写 CUDA 的真实复杂度:从 naive 到 tuned 的 12 轮迭代。
  4. Tillet et al., Triton: An Intermediate Language and Compiler for Tiled Neural Network Computations (MAPS 2019) — Triton 原始论文。块级抽象 vs CUDA 的线程级。
  5. Hsu et al., Liger Kernel: Efficient Triton Kernels for LLM Training (arXiv:2410.10989) — Liger 技术报告。fusion/chunking/in-place 的精确 benchmark。
  6. Triton 官方, Triton 文档 — 编程模型、installation、debugging。
  7. LinkedIn, Liger-Kernel GitHub — 源码、支持的模型列表、API。
💬 我是你的 CUDA 老师——随时追问
哪些概念还不清晰?想把课件里的 CUDA vs Triton 代码跑起来对比?想知道你实际业务里某个算子该用哪层? 直接问我。卡住时问,比自己硬猜更高效。