Triton 从 0 到 1 完整学习教程

默认主题：OpenAI Triton ，即用于编写高性能 GPU kernel 的 Python DSL / 编译器。

如果你说的是 NVIDIA Triton Inference Server ，请看最后的附录 B。本文主线是 AI/HPC 中更常见的 Triton kernel 编程。

学前准备：你需要具备什么
Triton 是什么：它解决什么问题
环境安装与验证
GPU 编程基础：理解 Triton 的前提
Triton 编程模型：Program、Tile、Mask、Pointer
第一个 Triton Kernel：Vector Add
核心 API 与常用模式
实例一：Softmax
实例二：LayerNorm
实例三：Fused Bias + SiLU
实例四：矩阵乘法 Matmul，从能跑到高性能
实例五：Attention / FlashAttention 风格 kernel
Autotune：自动调优
性能优化方法论
调试、测试与排错
与 PyTorch 集成
学习路线与练习项目
常见坑 FAQ
附录 A：CUDA 程序员转 Triton 速查表
附录 B：如果你指的是 NVIDIA Triton Inference Server

1. 学前准备：你需要具备什么

学习 Triton 前，建议你至少具备以下基础：

必备基础

Python 基础
- 函数、装饰器、lambda
- 类型注解不是必须，但有助于理解代码
- 会用 pip / venv / conda
PyTorch 基础
- Tensor 是什么
- shape、stride、device、dtype
- torch.empty_like
- tensor.contiguous()
- CUDA tensor 的基本使用
基本线性代数
- 向量加法
- 矩阵乘法
- softmax
- normalization 的基本概念

建议但非必须

CUDA / GPU 编程基础
对计算机体系结构的基本理解：
- 内存层次
- cache
- 带宽
- 并行
对深度学习模型常见算子的理解：
- Linear / Matmul
- LayerNorm
- Softmax
- Attention

如果你完全没写过 GPU 程序，也可以学。Triton 相比 CUDA 更适合入门，因为它隐藏了大量线程级细节。

2. Triton 是什么：它解决什么问题

2.1 一句话定义

Triton 是一个用 Python 写高性能 GPU kernel 的编程框架。

你写的是一段 Python 风格的代码，Triton 会把它编译成可以在 GPU 上运行的高性能 kernel，例如 PTX / AMD GPU 代码。

2.2 为什么需要 Triton？

在深度学习系统中，性能瓶颈经常来自：

大量小算子频繁读写显存；
PyTorch 默认算子无法针对特定 shape 充分优化；
想融合多个算子，减少中间 tensor；
手写 CUDA 成本太高；
需要快速实验不同 tile size、pipeline、num_warps 等配置。

Triton 的目标是：

用比 CUDA 更高的抽象，写出接近 CUDA / cuBLAS / cuDNN 性能的 kernel。

2.3 Triton 与 CUDA 的核心区别

维度	CUDA	Triton
编程抽象	thread / warp / block	program / tile / block-level tensor
语言	C/C++	Python DSL
内存管理	手动较多	编译器辅助较多
shared memory	手动管理	编译器大量自动管理
software pipelining	手动或 CUTLASS	`num_stages` 等方式辅助
开发效率	较低	较高
极限性能	可非常高	很多场景接近手写 CUDA
适合人群	GPU 专业优化工程师	AI 工程师、编译器爱好者、高性能计算开发者

2.4 Triton 的核心思想：Tile 级编程

CUDA 通常让你思考：

每个 thread 处理哪个元素？

Triton 通常让你思考：

每个 program 处理哪一块 tile？

例如矩阵乘法：

text 复制代码

C = A @ B

你不会写每个线程怎么算一个元素，而是写：

text 复制代码

一个 Triton program 负责 C 的一个 BLOCK_M x BLOCK_N 子块。
它循环读取 A 的 BLOCK_M x BLOCK_K 子块，
以及 B 的 BLOCK_K x BLOCK_N 子块，
然后做 tile 级矩阵乘法并累加。

这就是 Triton 的核心编程模式。

2.5 Triton 的优势

Python 语法，学习成本低
tile 级抽象，更容易写 matmul / attention / norm / softmax
JIT 编译
支持 autotune
适合和 PyTorch 结合
对 LLM 算子优化非常有用
- fused attention
- flash attention
- layernorm / rmsnorm
- fused linear 后处理
- quantization/dequantization

2.6 Triton 的局限

不是所有 CUDA 特性都直接暴露；
某些极端优化场景仍不如手写 CUDA；
API 仍在演进；
对复杂控制流、动态 shape 的支持有限；
调试体验不如普通 Python；
某些高级硬件特性需要版本支持。

3. 环境安装与验证

本教程建议使用：

Linux / WSL2
NVIDIA GPU
Python 3.10+
PyTorch 2.x
Triton 3.x 或较新版本

Windows 原生支持有限，建议使用 WSL2 或 Linux。

AMD GPU 支持在持续发展中，本教程主要以 NVIDIA GPU 为例。

3.1 创建虚拟环境

bash 复制代码

python -m venv triton-env
source triton-env/bin/activate

如果你用 conda：

bash 复制代码

conda create -n triton python=3.11 -y
conda activate triton

3.2 安装 PyTorch

以 CUDA 12.1 为例：

bash 复制代码

pip install torch --index-url https://download.pytorch.org/whl/cu121

如果你的 CUDA 版本不同，请选择对应的 PyTorch wheel。

3.3 安装 Triton

bash 复制代码

pip install triton

某些 PyTorch 版本会自动安装 Triton 依赖，但显式安装通常更清楚。

3.4 验证环境

创建文件：

bash 复制代码

vim check_triton.py

内容：

python 复制代码

import torch
import triton

print("PyTorch version:", torch.__version__)
print("CUDA available:", torch.cuda.is_available())

if torch.cuda.is_available():
    print("GPU:", torch.cuda.get_device_name(0))

print("Triton version:", triton.__version__)

运行：

bash 复制代码

python check_triton.py

如果输出类似：

text 复制代码

PyTorch version: 2.x.x
CUDA available: True
GPU: NVIDIA A100-SXM4-80GB
Triton version: 3.x.x

说明环境基本可用。

3.5 可选环境变量

bash 复制代码

# Triton 编译缓存目录
export TRITON_CACHE_DIR=$HOME/.triton/cache

# 某些场景下可用 CPU 解释模式调试
export TRITON_INTERPRET=1

注意：TRITON_INTERPRET 不一定支持所有算子，主要用于调试和学习。

4. GPU 编程基础：理解 Triton 的前提

即使 Triton 隐藏了很多细节，你仍需要理解 GPU 的基本运行方式。

4.1 CPU 和 GPU 的区别

CPU：

text 复制代码

少量强核心
擅长复杂控制流、低延迟任务

GPU：

text 复制代码

大量弱核心
擅长高吞吐、规则并行的任务

例如：

python 复制代码

y = x + 1

如果有 1 亿个元素，每个元素做一次加法，那么 GPU 可以并行处理。

4.2 GPU 内存层次

简化版：

text 复制代码

寄存器 Register
  ↓
共享内存 Shared Memory
  ↓
L2 Cache
  ↓
全局显存 Global Memory / HBM

性能优化经常围绕：

尽量减少全局显存访问；
尽量让数据复用发生在更快的内存层级；
尽量让访存连续、合并；
尽量提高计算访存比。

4.3 Kernel 是什么？

Kernel 是一段在 GPU 上并行执行的函数。

例如：

text 复制代码

对数组每个元素加 1

CPU 写法：

python 复制代码

for i in range(N):
    y[i] = x[i] + 1

GPU 思想：

text 复制代码

启动 N 个并行执行单元，每个处理一个 i。

Triton 里则通常变成：

text 复制代码

启动多个 program，每个 program 处理一个 block/tile。

4.4 Grid、Program、Block

在 Triton 中，你可以粗略理解：

text 复制代码

Grid：一组 Triton program
Program：一个并行执行实例
Tile：一个 program 处理的数据块

例如：

text 复制代码

向量长度 N = 1,000,000
BLOCK_SIZE = 1024

需要大约 1000 / 1 = 977 个 program？
不对，1,000,000 / 1024 ≈ 977 个 program。

每个 program 处理 1024 个元素。

4.5 Warp 和 num_warps

在 NVIDIA GPU 中，warp 通常是 32 个线程一起执行。

Triton 不要求你手写线程，但你可以指定：

python 复制代码

num_warps=4

表示这个 kernel 大致使用多少个 warp 来执行一个 program。

经验：

elementwise / reduce：num_warps=4/8
matmul：num_warps=4/8/16
attention：常见 num_warps=4/8

4.6 num_stages 是什么？

num_stages 通常用于软件流水线，尤其是循环读取 K 维的 matmul / attention。

简单理解：

text 复制代码

当当前 tile 在计算时，提前加载后续 tile。

经验值：

text 复制代码

num_stages = 2 ~ 5

太大可能导致 shared memory 不够或寄存器压力大。

5. Triton 编程模型：Program、Tile、Mask、Pointer

这一章是 Triton 的核心心智模型。

5.1 一个 Triton kernel 的基本结构

python 复制代码

import triton
import triton.language as tl

@triton.jit
def my_kernel(...):
    pid = tl.program_id(axis=0)
    ...

解释：

@triton.jit：表示这是一个 Triton kernel；
tl.program_id(axis=0)：获取当前 program 在 grid 第 0 维的 id；
kernel 内部操作通常是 tile 级 tensor。

5.2 Grid 是什么？

Grid 指定启动多少个 program。

例如一维 grid：

python 复制代码

grid = (1024,)

二维 grid：

python 复制代码

grid = (128, 128)

也可以用 lambda：

python 复制代码

grid = lambda meta: (triton.cdiv(n, meta["BLOCK_SIZE"]),)

其中 meta 是编译期配置字典。

5.3 tl.constexpr：编译期常量

例如：

python 复制代码

BLOCK_SIZE: tl.constexpr

表示 BLOCK_SIZE 在编译时确定。

Triton 中 block shape 通常需要是编译期常量，而且通常建议是 2 的幂：

text 复制代码

16, 32, 64, 128, 256, 512, 1024

5.4 Pointer Arithmetic：指针运算

Triton 中你可以对指针做加法。

假设：

python 复制代码

x_ptr

是一个指向 GPU 内存的指针。

如果：

python 复制代码

offsets = tl.arange(0, BLOCK_SIZE)

那么：

python 复制代码

x_ptr + offsets

表示一组指针，分别指向：

text 复制代码

x_ptr + 0
x_ptr + 1
x_ptr + 2
...

5.5 tl.load 和 tl.store

读取数据：

python 复制代码

x = tl.load(x_ptr + offsets, mask=mask, other=0.0)

写入数据：

python 复制代码

tl.store(out_ptr + offsets, y, mask=mask)

重点：

mask 决定哪些地址合法；
other 是 mask 为 false 时的默认值；
不写 mask 可能导致非法内存访问。

5.6 Mask 是 Triton 的生命线

假设：

text 复制代码

N = 1000
BLOCK_SIZE = 1024

最后一个 program 会覆盖：

text 复制代码

offsets: 0 ~ 1023

但实际只有：

text 复制代码

0 ~ 999

合法。

所以必须：

python 复制代码

mask = offsets < n_elements

否则可能越界。

5.7 tile 级 shape 和 broadcast

一维：

python 复制代码

tl.arange(0, BLOCK_SIZE)

二维：

python 复制代码

offs_m = tl.arange(0, BLOCK_M)
offs_n = tl.arange(0, BLOCK_N)

构造二维索引：

python 复制代码

offs_m[:, None]
offs_n[None, :]

例如：

python 复制代码

ptrs = base + offs_m[:, None] * stride_m + offs_n[None, :] * stride_n

这是 Triton 里最常见的 2D tile 指针构造方式。

5.8 tl.sum / tl.max / tl.min

reduce 操作：

python 复制代码

total = tl.sum(x, axis=0)
max_val = tl.max(x, axis=0)

对于二维 tensor：

python 复制代码

row_sum = tl.sum(x, axis=1)

表示沿列方向 reduce，得到每行结果。

5.9 tl.dot：矩阵乘法核心

python 复制代码

acc += tl.dot(a, b)

其中：

text 复制代码

a: [BLOCK_M, BLOCK_K]
b: [BLOCK_K, BLOCK_N]
acc: [BLOCK_M, BLOCK_N]

tl.dot 是写 matmul / attention 的核心。

通常：

输入 a、b 可以是 fp16 / bf16；
累加器 acc 使用 fp32；
最后输出时再 cast 回 fp16 / bf16。

6. 第一个 Triton Kernel：Vector Add

我们从最简单的算子开始：向量加法。

目标：

python 复制代码

out = x + y

6.1 完整代码

python 复制代码

import torch
import triton
import triton.language as tl


@triton.jit
def add_kernel(
    x_ptr,
    y_ptr,
    out_ptr,
    n_elements,
    BLOCK_SIZE: tl.constexpr,
):
    # 当前 program 的 id
    pid = tl.program_id(axis=0)

    # 当前 program 负责的元素偏移
    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)

    # 防止越界
    mask = offsets < n_elements

    # 读取输入
    x = tl.load(x_ptr + offsets, mask=mask, other=0.0)
    y = tl.load(y_ptr + offsets, mask=mask, other=0.0)

    # 计算并写回
    tl.store(out_ptr + offsets, x + y, mask=mask)


def add(x: torch.Tensor, y: torch.Tensor):
    assert x.is_cuda and y.is_cuda
    assert x.shape == y.shape

    out = torch.empty_like(x)
    n = x.numel()

    # grid 是一个 lambda，meta 中可以拿到 BLOCK_SIZE
    grid = lambda meta: (triton.cdiv(n, meta["BLOCK_SIZE"]),)

    add_kernel[grid](
        x,
        y,
        out,
        n,
        BLOCK_SIZE=1024,
    )
    return out


if __name__ == "__main__":
    torch.manual_seed(0)

    x = torch.randn(1024 * 1024, device="cuda")
    y = torch.randn(1024 * 1024, device="cuda")

    out = add(x, y)
    ref = x + y

    torch.testing.assert_close(out, ref)
    print("Vector Add OK")

6.2 逐段解释

kernel 签名

python 复制代码

def add_kernel(
    x_ptr,
    y_ptr,
    out_ptr,
    n_elements,
    BLOCK_SIZE: tl.constexpr,
):

x_ptr、y_ptr、out_ptr 是 GPU 指针。

n_elements 是总元素数。

BLOCK_SIZE 是编译期常量。

获取 program id

python 复制代码

pid = tl.program_id(axis=0)

一维 grid 中，每个 program 拿到不同 pid。

例如：

text 复制代码

pid = 0 -> offsets 0 ~ 1023
pid = 1 -> offsets 1024 ~ 2047
pid = 2 -> offsets 2048 ~ 3071

计算 offsets

python 复制代码

offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)

tl.arange(0, BLOCK_SIZE) 生成：

text 复制代码

[0, 1, 2, ..., BLOCK_SIZE - 1]

加上 pid * BLOCK_SIZE 后就是当前 program 的全局偏移。

mask

python 复制代码

mask = offsets < n_elements

防止最后一个 block 越界。

load/store

python 复制代码

x = tl.load(x_ptr + offsets, mask=mask, other=0.0)
y = tl.load(y_ptr + offsets, mask=mask, other=0.0)
tl.store(out_ptr + offsets, x + y, mask=mask)

每个 program 处理 BLOCK_SIZE 个元素。

grid

python 复制代码

grid = lambda meta: (triton.cdiv(n, meta["BLOCK_SIZE"]),)

如果：

text 复制代码

n = 10000
BLOCK_SIZE = 1024

那么：

text 复制代码

ceil(10000 / 1024) = 10

启动 10 个 program。

6.3 benchmark

python 复制代码

import triton.testing as tt

x = torch.randn(1 << 26, device="cuda")
y = torch.randn_like(x)

ms = tt.do_bench(lambda: add(x, y))
print(f"triton add: {ms:.3f} ms")

ms = tt.do_bench(lambda: x + y)
print(f"torch add: {ms:.3f} ms")

do_bench 会自动处理 warmup 和 CUDA 同步。

6.4 这个 kernel 的性能关键点

Vector Add 是典型 memory-bound kernel。

对于 fp32：

text 复制代码

读取 x: 4N bytes
读取 y: 4N bytes
写入 out: 4N bytes
总计: 12N bytes

因此性能主要取决于显存带宽，而不是计算能力。

优化方向：

保证连续访存；
block size 不能太小；
避免不必要的分支；
避免额外中间 tensor。

7. 核心 API 与常用模式

这一章整理 Triton 里最常用的 API。

7.1 `tl.program_id`

python 复制代码

pid = tl.program_id(axis=0)

二维：

python 复制代码

pid_m = tl.program_id(axis=0)
pid_n = tl.program_id(axis=1)

7.2 `tl.arange`

python 复制代码

x = tl.arange(0, 64)

生成：

text 复制代码

[0, 1, ..., 63]

7.3 `tl.zeros`

python 复制代码

acc = tl.zeros((64, 64), dtype=tl.float32)

7.4 `tl.full`

python 复制代码

m = tl.full((64,), -1e30, dtype=tl.float32)

7.5 `tl.load`

python 复制代码

x = tl.load(ptr, mask=mask, other=0.0)

7.6 `tl.store`

python 复制代码

tl.store(ptr, value, mask=mask)

7.7 `tl.sum`

python 复制代码

s = tl.sum(x, axis=0)

二维：

python 复制代码

row_sum = tl.sum(x, axis=1)

7.8 `tl.max`

python 复制代码

m = tl.max(x, axis=1)

7.9 `tl.exp`

python 复制代码

y = tl.exp(x)

7.10 `tl.where`

python 复制代码

y = tl.where(mask, a, b)

等价：

text 复制代码

if mask:
    y = a
else:
    y = b

但作用于 tensor。

7.11 `tl.dot`

python 复制代码

acc += tl.dot(a, b)

常见 shape：

text 复制代码

a: [M, K]
b: [K, N]
acc: [M, N]

7.12 `tl.trans`

python 复制代码

b_t = tl.trans(b)

用于转置。

例如 attention：

python 复制代码

qk = tl.dot(q, tl.trans(k))

7.13 `tl.cdiv`

python 复制代码

n_blocks = tl.cdiv(N, BLOCK_N)

等价：

text 复制代码

ceil(N / BLOCK_N)

7.14 `tl.atomic_add`

原子加：

python 复制代码

tl.atomic_add(ptr, value, mask=mask)

适合 histogram、reduce 等，但通常性能不如 tree reduce / block reduce。

8. 实例一：Softmax

目标：

python 复制代码

y = softmax(x, dim=-1)

我们先写一个简单版本：每一行一个 program。

适合：

text 复制代码

n_cols 不是特别大

8.1 数学公式

对于一行向量 x：

text 复制代码

softmax(x_i) = exp(x_i - max(x)) / sum(exp(x_j - max(x)))

减最大值是为了数值稳定。

8.2 完整代码

python 复制代码

import torch
import triton
import triton.language as tl


@triton.jit
def softmax_kernel(
    out_ptr,
    in_ptr,
    in_row_stride,
    out_row_stride,
    n_cols,
    BLOCK_SIZE: tl.constexpr,
):
    row = tl.program_id(axis=0)

    cols = tl.arange(0, BLOCK_SIZE)
    mask = cols < n_cols

    in_ptrs = in_ptr + row * in_row_stride + cols

    # 读入一行
    x = tl.load(in_ptrs, mask=mask, other=-float("inf")).to(tl.float32)

    # 数值稳定 softmax
    x_max = tl.max(x, axis=0)
    y = tl.exp(x - x_max)
    denom = tl.sum(y, axis=0)
    y = y / denom

    out_ptrs = out_ptr + row * out_row_stride + cols
    tl.store(out_ptrs, y.to(out_ptr.dtype.element_ty), mask=mask)


def softmax(x: torch.Tensor):
    assert x.is_cuda
    assert x.ndim == 2

    n_rows, n_cols = x.shape
    y = torch.empty_like(x)

    # 找到 >= n_cols 的 2 的幂
    BLOCK_SIZE = 1 << (n_cols - 1).bit_length()

    num_warps = 4
    if BLOCK_SIZE >= 2048:
        num_warps = 8
    if BLOCK_SIZE >= 4096:
        num_warps = 16

    softmax_kernel[(n_rows,)](
        y,
        x,
        x.stride(0),
        y.stride(0),
        n_cols,
        num_warps=num_warps,
        BLOCK_SIZE=BLOCK_SIZE,
    )
    return y


if __name__ == "__main__":
    x = torch.randn(512, 1024, device="cuda", dtype=torch.float16)

    y = softmax(x)
    ref = torch.softmax(x.float(), dim=-1).to(x.dtype)

    torch.testing.assert_close(y, ref, rtol=1e-3, atol=1e-3)
    print("Softmax OK")

8.3 关键点解释

一行一个 program

python 复制代码

row = tl.program_id(axis=0)

grid 是：

python 复制代码

(n_rows,)

每个 program 处理一行。

BLOCK_SIZE 必须覆盖整行

这里简单版本要求：

text 复制代码

BLOCK_SIZE >= n_cols

否则一行装不下。

如果 n_cols 很大，需要使用分段 softmax / online softmax。

数值稳定

python 复制代码

x_max = tl.max(x, axis=0)
y = tl.exp(x - x_max)

防止 exp overflow。

dtype

python 复制代码

x = tl.load(...).to(tl.float32)

softmax 建议至少用 fp32 计算，即使输入是 fp16。

8.4 这个版本的限制

如果：

text 复制代码

n_cols = 131072

那么：

text 复制代码

BLOCK_SIZE = 131072

可能导致寄存器压力极大，甚至无法编译。

此时需要：

多 block 协作；
online softmax；
两次 pass；
更复杂的 reduce 策略。

9. 实例二：LayerNorm

目标：

python 复制代码

y = (x - mean) / sqrt(var + eps) * weight + bias

同样先写一个简单版本：每一行一个 program。

9.1 完整代码

python 复制代码

import torch
import triton
import triton.language as tl


@triton.jit
def layer_norm_kernel(
    X,
    Y,
    W,
    B,
    Mean,
    Rstd,
    stride,
    N,
    eps,
    BLOCK_SIZE: tl.constexpr,
):
    row = tl.program_id(axis=0)

    cols = tl.arange(0, BLOCK_SIZE)
    mask = cols < N

    # 读取一行
    x = tl.load(X + row * stride + cols, mask=mask, other=0.0).to(tl.float32)

    # mean
    mean = tl.sum(x, axis=0) / N

    # var
    x_hat = tl.where(mask, x - mean, 0.0)
    var = tl.sum(x_hat * x_hat, axis=0) / N

    rstd = 1.0 / tl.sqrt(var + eps)

    # weight / bias
    w = tl.load(W + cols, mask=mask, other=1.0).to(tl.float32)
    b = tl.load(B + cols, mask=mask, other=0.0).to(tl.float32)

    y = (x - mean) * rstd * w + b

    tl.store(Y + row * stride + cols, y.to(Y.dtype.element_ty), mask=mask)

    # 可选：保存统计量，backward 会用到
    tl.store(Mean + row, mean)
    tl.store(Rstd + row, rstd)


def layer_norm(
    x: torch.Tensor,
    weight: torch.Tensor,
    bias: torch.Tensor,
    eps: float = 1e-5,
):
    assert x.is_cuda
    assert x.ndim == 2

    M, N = x.shape
    y = torch.empty_like(x)

    mean = torch.empty(M, device=x.device, dtype=torch.float32)
    rstd = torch.empty(M, device=x.device, dtype=torch.float32)

    BLOCK_SIZE = 1 << (N - 1).bit_length()

    num_warps = 4
    if BLOCK_SIZE >= 2048:
        num_warps = 8
    if BLOCK_SIZE >= 4096:
        num_warps = 16

    layer_norm_kernel[(M,)](
        x,
        y,
        weight,
        bias,
        mean,
        rstd,
        x.stride(0),
        N,
        eps,
        BLOCK_SIZE=BLOCK_SIZE,
        num_warps=num_warps,
    )

    return y


if __name__ == "__main__":
    M, N = 1024, 4096
    x = torch.randn(M, N, device="cuda", dtype=torch.float16)
    weight = torch.randn(N, device="cuda", dtype=torch.float16)
    bias = torch.randn(N, device="cuda", dtype=torch.float16)

    y = layer_norm(x, weight, bias)

    ref = torch.nn.functional.layer_norm(
        x.float(),
        normalized_shape=(N,),
        weight=weight.float(),
        bias=bias.float(),
        eps=1e-5,
    ).to(x.dtype)

    torch.testing.assert_close(y, ref, rtol=1e-3, atol=1e-3)
    print("LayerNorm OK")

9.2 关键点

一行一个 program

python 复制代码

row = tl.program_id(axis=0)

LayerNorm 通常对最后一维做归一化。

fp32 计算

python 复制代码

x = ... .to(tl.float32)

Norm 类算子建议用 fp32 统计 mean / var。

保存 mean / rstd

python 复制代码

tl.store(Mean + row, mean)
tl.store(Rstd + row, rstd)

训练反向传播时通常需要它们。

限制

和 softmax 类似，简单版本要求：

text 复制代码

BLOCK_SIZE >= N

如果 hidden size 非常大，需要分段实现或更复杂策略。

10. 实例三：Fused Bias + SiLU

目标：

text 复制代码

y = x + bias
y = y * sigmoid(y)

即：

text 复制代码

SiLU(x + bias)

这是一个典型 elementwise fusion。

10.1 为什么要融合？

PyTorch 中如果分开写：

python 复制代码

y = x + bias
y = torch.nn.functional.silu(y)

会产生多次显存读写：

text 复制代码

读 x
读 bias
写 x + bias
读 x + bias
写 silu 结果

融合后：

text 复制代码

读 x
读 bias
计算
写 y

减少显存流量。

10.2 完整代码

python 复制代码

import torch
import triton
import triton.language as tl


@triton.jit
def bias_silu_kernel(
    Out,
    X,
    Bias,
    M,
    N,
    stride_xm,
    stride_xn,
    stride_om,
    stride_on,
    BLOCK_M: tl.constexpr,
    BLOCK_N: tl.constexpr,
):
    pid_m = tl.program_id(axis=0)
    pid_n = tl.program_id(axis=1)

    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)

    mask = (offs_m[:, None] < M) & (offs_n[None, :] < N)

    x_ptrs = X + offs_m[:, None] * stride_xm + offs_n[None, :] * stride_xn
    b_ptrs = Bias + offs_n

    x = tl.load(x_ptrs, mask=mask, other=0.0).to(tl.float32)
    b = tl.load(b_ptrs, mask=offs_n < N, other=0.0).to(tl.float32)

    y = x + b[None, :]

    # SiLU: y * sigmoid(y)
    y = y * (1.0 / (1.0 + tl.exp(-y)))

    out_ptrs = Out + offs_m[:, None] * stride_om + offs_n[None, :] * stride_on
    tl.store(out_ptrs, y.to(Out.dtype.element_ty), mask=mask)


def bias_silu(x: torch.Tensor, bias: torch.Tensor):
    assert x.is_cuda
    assert bias.is_cuda
    assert x.ndim == 2
    assert bias.ndim == 1
    assert x.shape[1] == bias.shape[0]

    M, N = x.shape
    out = torch.empty_like(x)

    BLOCK_M = 32
    BLOCK_N = 128

    grid = (
        triton.cdiv(M, BLOCK_M),
        triton.cdiv(N, BLOCK_N),
    )

    bias_silu_kernel[grid](
        out,
        x,
        bias,
        M,
        N,
        x.stride(0),
        x.stride(1),
        out.stride(0),
        out.stride(1),
        BLOCK_M=BLOCK_M,
        BLOCK_N=BLOCK_N,
        num_warps=4,
    )

    return out


if __name__ == "__main__":
    M, N = 4096, 4096

    x = torch.randn(M, N, device="cuda", dtype=torch.float16)
    bias = torch.randn(N, device="cuda", dtype=torch.float16)

    y = bias_silu(x, bias)

    ref = torch.nn.functional.silu(x.float() + bias.float()).to(x.dtype)

    torch.testing.assert_close(y, ref, rtol=1e-3, atol=1e-3)
    print("Bias + SiLU OK")

10.3 关键点

二维 grid

python 复制代码

grid = (
    triton.cdiv(M, BLOCK_M),
    triton.cdiv(N, BLOCK_N),
)

每个 program 处理一个 BLOCK_M x BLOCK_N tile。

bias broadcast

python 复制代码

b = tl.load(Bias + offs_n, mask=offs_n < N, other=0.0)
y = x + b[None, :]

bias 是 [N]，通过 b[None, :] broadcast 到 [BLOCK_M, BLOCK_N]。

SiLU 手工实现

python 复制代码

y = y * (1.0 / (1.0 + tl.exp(-y)))

避免依赖某些版本中不稳定的 API。

11. 实例四：矩阵乘法 Matmul，从能跑到高性能

矩阵乘法是 Triton 最重要的案例之一。

目标：

python 复制代码

C = A @ B

其中：

text 复制代码

A: [M, K]
B: [K, N]
C: [M, N]

11.1 Matmul 的 tile 思想

每个 program 负责 C 的一个子块：

text 复制代码

C_tile: [BLOCK_M, BLOCK_N]

计算方式：

text 复制代码

for k in range(0, K, BLOCK_K):
    A_tile = A[m:m+BLOCK_M, k:k+BLOCK_K]
    B_tile = B[k:k+BLOCK_K, n:n+BLOCK_N]
    C_tile += A_tile @ B_tile

这就是 tiled matmul。

11.2 简单版本 Matmul

先写一个教学版本，不追求极限性能。

python 复制代码

import torch
import triton
import triton.language as tl


@triton.jit
def matmul_kernel_simple(
    A,
    B,
    C,
    M,
    N,
    K,
    stride_am,
    stride_ak,
    stride_bk,
    stride_bn,
    stride_cm,
    stride_cn,
    BLOCK_M: tl.constexpr,
    BLOCK_N: tl.constexpr,
    BLOCK_K: tl.constexpr,
):
    pid_m = tl.program_id(axis=0)
    pid_n = tl.program_id(axis=1)

    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
    offs_k = tl.arange(0, BLOCK_K)

    a_ptrs = A + offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak
    b_ptrs = B + offs_k[:, None] * stride_bk + offs_n[None, :] * stride_bn

    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)

    for k in range(0, tl.cdiv(K, BLOCK_K)):
        a_mask = (offs_m[:, None] < M) & (offs_k[None, :] < K)
        b_mask = (offs_k[:, None] < K) & (offs_n[None, :] < N)

        a = tl.load(a_ptrs, mask=a_mask, other=0.0)
        b = tl.load(b_ptrs, mask=b_mask, other=0.0)

        acc += tl.dot(a, b)

        a_ptrs += BLOCK_K * stride_ak
        b_ptrs += BLOCK_K * stride_bk
        offs_k += BLOCK_K

    c_ptrs = C + offs_m[:, None] * stride_cm + offs_n[None, :] * stride_cn
    c_mask = (offs_m[:, None] < M) & (offs_n[None, :] < N)

    tl.store(c_ptrs, acc.to(C.dtype.element_ty), mask=c_mask)


def matmul_simple(a: torch.Tensor, b: torch.Tensor):
    assert a.is_cuda and b.is_cuda
    assert a.ndim == 2 and b.ndim == 2
    assert a.shape[1] == b.shape[0]

    a = a.contiguous()
    b = b.contiguous()

    M, K = a.shape
    K, N = b.shape

    c = torch.empty((M, N), device=a.device, dtype=a.dtype)

    BLOCK_M = 64
    BLOCK_N = 64
    BLOCK_K = 32

    grid = (
        triton.cdiv(M, BLOCK_M),
        triton.cdiv(N, BLOCK_N),
    )

    matmul_kernel_simple[grid](
        a,
        b,
        c,
        M,
        N,
        K,
        a.stride(0),
        a.stride(1),
        b.stride(0),
        b.stride(1),
        c.stride(0),
        c.stride(1),
        BLOCK_M=BLOCK_M,
        BLOCK_N=BLOCK_N,
        BLOCK_K=BLOCK_K,
        num_warps=4,
        num_stages=3,
    )

    return c


if __name__ == "__main__":
    M = N = K = 1024

    a = torch.randn(M, K, device="cuda", dtype=torch.float16)
    b = torch.randn(K, N, device="cuda", dtype=torch.float16)

    c = matmul_simple(a, b)
    ref = torch.matmul(a.float(), b.float()).to(torch.float16)

    torch.testing.assert_close(c, ref, rtol=1e-2, atol=1e-2)
    print("Simple Matmul OK")

11.3 简单版本解释

program 负责 C 的一个 tile

python 复制代码

pid_m = tl.program_id(axis=0)
pid_n = tl.program_id(axis=1)

A tile 指针

python 复制代码

a_ptrs = A + offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak

shape：

text 复制代码

[BLOCK_M, BLOCK_K]

B tile 指针

python 复制代码

b_ptrs = B + offs_k[:, None] * stride_bk + offs_n[None, :] * stride_bn

shape：

text 复制代码

[BLOCK_K, BLOCK_N]

K 维循环

python 复制代码

for k in range(0, tl.cdiv(K, BLOCK_K)):
    ...
    acc += tl.dot(a, b)

每次沿 K 维处理 BLOCK_K。

fp32 accumulator

python 复制代码

acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)

即使输入是 fp16，也常用 fp32 累加。

11.4 性能版 Matmul：加入 Autotune 和 L2 swizzle

真正高性能 matmul 通常还要做：

block size autotune；
num_warps autotune；
num_stages autotune；
program id swizzle 提升 L2 cache 命中；
更合理的 K 维 pipeline。

下面给一个较完整的 autotune 版本。

11.5 Autotuned Matmul 完整代码

python 复制代码

import torch
import triton
import triton.language as tl


def matmul_configs():
    return [
        triton.Config(
            {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 32, "GROUP_M": 8},
            num_warps=4,
            num_stages=3,
        ),
        triton.Config(
            {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 64, "GROUP_M": 8},
            num_warps=4,
            num_stages=3,
        ),
        triton.Config(
            {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 32, "GROUP_M": 8},
            num_warps=8,
            num_stages=3,
        ),
        triton.Config(
            {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 32, "GROUP_M": 8},
            num_warps=8,
            num_stages=3,
        ),
        triton.Config(
            {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 32, "GROUP_M": 8},
            num_warps=4,
            num_stages=4,
        ),
        triton.Config(
            {"BLOCK_M": 64, "BLOCK_N": 64, "BLOCK_K": 32, "GROUP_M": 8},
            num_warps=4,
            num_stages=5,
        ),
    ]


@triton.autotune(
    configs=matmul_configs(),
    key=["M", "N", "K"],
)
@triton.jit
def matmul_kernel(
    A,
    B,
    C,
    M,
    N,
    K,
    stride_am,
    stride_ak,
    stride_bk,
    stride_bn,
    stride_cm,
    stride_cn,
    BLOCK_M: tl.constexpr,
    BLOCK_N: tl.constexpr,
    BLOCK_K: tl.constexpr,
    GROUP_M: tl.constexpr,
):
    # -----------------------------
    # program swizzle，提高 L2 命中
    # -----------------------------
    pid = tl.program_id(axis=0)

    num_pid_m = tl.cdiv(M, BLOCK_M)
    num_pid_n = tl.cdiv(N, BLOCK_N)

    num_pid_in_group = GROUP_M * num_pid_n
    group_id = pid // num_pid_in_group

    first_pid_m = group_id * GROUP_M
    group_size_m = min(num_pid_m - first_pid_m, GROUP_M)

    pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)
    pid_n = (pid % num_pid_in_group) // group_size_m

    # -----------------------------
    # 计算当前 program 负责的 tile
    # -----------------------------
    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
    offs_k = tl.arange(0, BLOCK_K)

    a_ptrs = A + offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak
    b_ptrs = B + offs_k[:, None] * stride_bk + offs_n[None, :] * stride_bn

    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)

    for k in range(0, tl.cdiv(K, BLOCK_K)):
        a_mask = (offs_m[:, None] < M) & (offs_k[None, :] < K)
        b_mask = (offs_k[:, None] < K) & (offs_n[None, :] < N)

        a = tl.load(a_ptrs, mask=a_mask, other=0.0)
        b = tl.load(b_ptrs, mask=b_mask, other=0.0)

        acc += tl.dot(a, b)

        a_ptrs += BLOCK_K * stride_ak
        b_ptrs += BLOCK_K * stride_bk
        offs_k += BLOCK_K

    c_ptrs = C + offs_m[:, None] * stride_cm + offs_n[None, :] * stride_cn
    c_mask = (offs_m[:, None] < M) & (offs_n[None, :] < N)

    tl.store(c_ptrs, acc.to(C.dtype.element_ty), mask=c_mask)


def matmul(a: torch.Tensor, b: torch.Tensor):
    assert a.is_cuda and b.is_cuda
    assert a.ndim == 2 and b.ndim == 2
    assert a.shape[1] == b.shape[0]

    a = a.contiguous()
    b = b.contiguous()

    M, K = a.shape
    K, N = b.shape

    c = torch.empty((M, N), device=a.device, dtype=a.dtype)

    grid = lambda META: (
        triton.cdiv(M, META["BLOCK_M"]) * triton.cdiv(N, META["BLOCK_N"]),
    )

    matmul_kernel[grid](
        a,
        b,
        c,
        M,
        N,
        K,
        a.stride(0),
        a.stride(1),
        b.stride(0),
        b.stride(1),
        c.stride(0),
        c.stride(1),
    )

    return c


if __name__ == "__main__":
    M = N = K = 2048

    a = torch.randn(M, K, device="cuda", dtype=torch.float16)
    b = torch.randn(K, N, device="cuda", dtype=torch.float16)

    c = matmul(a, b)
    ref = torch.matmul(a.float(), b.float()).to(torch.float16)

    torch.testing.assert_close(c, ref, rtol=1e-2, atol=1e-2)
    print("Autotuned Matmul OK")

11.6 为什么使用一维 grid？

高性能 matmul 常把二维 program 映射到一维 grid：

python 复制代码

grid = num_pid_m * num_pid_n

这样更容易做 swizzle：

text 复制代码

让相邻 program 尽量复用相同的 A tile 或 B tile，
提高 L2 cache 命中率。

11.7 GROUP_M 是什么？

GROUP_M 控制 program 分组方式。

简化理解：

text 复制代码

把若干个 M 方向的 pid 分组，
让一个组内同时覆盖一段 N 方向，
从而增加 B/A tile 的 cache 复用。

常见值：

text 复制代码

4, 8, 16

11.8 benchmark Matmul

python 复制代码

import triton.testing as tt

M = N = K = 4096

a = torch.randn(M, K, device="cuda", dtype=torch.float16)
b = torch.randn(K, N, device="cuda", dtype=torch.float16)


def tflops(ms, M, N, K):
    flops = 2 * M * N * K
    return flops / (ms * 1e-3) / 1e12


ms_triton = tt.do_bench(lambda: matmul(a, b))
ms_torch = tt.do_bench(lambda: torch.matmul(a, b))

print(f"triton: {ms_triton:.3f} ms, {tflops(ms_triton, M, N, K):.2f} TFLOPS")
print(f"torch:  {ms_torch:.3f} ms, {tflops(ms_torch, M, N, K):.2f} TFLOPS")

注意：

不同 GPU 差异很大；
小 shape 时 autotune 开销明显；
torch.matmul 底层是 cuBLAS，很强，不必期望简单 Triton matmul 总是赢；
Triton 的价值常在于融合、自定义、特殊 layout、量化等场景。

12. 实例五：Attention / FlashAttention 风格 kernel

Attention 是 LLM 的核心算子。

标准 attention：

text 复制代码

Attention(Q, K, V) = softmax(Q @ K^T / sqrt(d)) @ V

直接实现的问题：

text 复制代码

Q @ K^T 会生成 [M, N] 大矩阵，
显存占用大，访存多。

FlashAttention 的核心思想：

text 复制代码

按 tile 计算 Q、K、V，
不 materialize 完整 attention matrix，
使用 online softmax 维护 max 和 sum。

下面写一个教学版 fused attention。

注意：这是教学版本，不是生产级 FlashAttention。

生产环境建议优先使用成熟实现，例如 PyTorch SDPA、FlashAttention、vLLM 等。

12.1 Online Softmax 思想

普通 softmax：

text 复制代码

max = max(score)
exp_scores = exp(score - max)
sum_exp = sum(exp_scores)
out = exp_scores / sum_exp

FlashAttention 风格中，我们按 K/V block 循环：

维护：

text 复制代码

m_i：当前最大 score
l_i：当前 exp sum
acc：当前输出累加器

每来一个新的 KV block：

text 复制代码

m_new = max(m_i, max(new_scores))
p = exp(new_scores - m_new)
alpha = exp(m_i - m_new)

l_i = l_i * alpha + sum(p)
acc = acc * alpha + p @ V_block
m_i = m_new

最后：

text 复制代码

out = acc / l_i

12.2 简化版 Attention Kernel

假设：

text 复制代码

Q: [M, D]
K: [N, D]
V: [N, D]
Out: [M, D]

单 head、无 batch。

D 建议是 16/32/64/128 等。

12.3 完整代码

python 复制代码

import torch
import triton
import triton.language as tl


@triton.jit
def attention_fwd(
    Q,
    K,
    V,
    Out,
    M,
    N,
    D: tl.constexpr,
    stride_qm,
    stride_qd,
    stride_kn,
    stride_kd,
    stride_vn,
    stride_vd,
    stride_om,
    stride_od,
    sm_scale,
    BLOCK_M: tl.constexpr,
    BLOCK_N: tl.constexpr,
    IS_CAUSAL: tl.constexpr,
):
    pid_m = tl.program_id(axis=0)

    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    offs_d = tl.arange(0, D)

    q_mask = offs_m < M

    q_ptrs = Q + offs_m[:, None] * stride_qm + offs_d[None, :] * stride_qd
    q = tl.load(q_ptrs, mask=q_mask[:, None], other=0.0)

    # online softmax 状态
    m_i = tl.full((BLOCK_M,), -1e30, dtype=tl.float32)
    l_i = tl.zeros((BLOCK_M,), dtype=tl.float32)
    acc = tl.zeros((BLOCK_M, D), dtype=tl.float32)

    # 遍历所有 KV block
    for start_n in range(0, N, BLOCK_N):
        offs_n = start_n + tl.arange(0, BLOCK_N)
        kv_mask = offs_n < N

        # Load K block: [BLOCK_N, D]
        k_ptrs = K + offs_n[:, None] * stride_kn + offs_d[None, :] * stride_kd
        k = tl.load(k_ptrs, mask=kv_mask[:, None], other=0.0)

        # QK^T: [BLOCK_M, BLOCK_N]
        qk = tl.dot(q, tl.trans(k)) * sm_scale

        valid = q_mask[:, None] & kv_mask[None, :]

        if IS_CAUSAL:
            causal = offs_m[:, None] >= offs_n[None, :]
            valid = valid & causal

        qk = tl.where(valid, qk, -float("inf"))

        # online softmax update
        m_new = tl.maximum(m_i, tl.max(qk, axis=1))

        p = tl.exp(qk - m_new[:, None])
        alpha = tl.exp(m_i - m_new)

        l_i = l_i * alpha + tl.sum(p, axis=1)
        acc = acc * alpha[:, None]

        # Load V block: [BLOCK_N, D]
        v_ptrs = V + offs_n[:, None] * stride_vn + offs_d[None, :] * stride_vd
        v = tl.load(v_ptrs, mask=kv_mask[:, None], other=0.0)

        # p: [BLOCK_M, BLOCK_N]
        # v: [BLOCK_N, D]
        acc += tl.dot(p.to(v.dtype), v)

        m_i = m_new

    # 防止无效行除 0
    l_i = tl.where(l_i == 0.0, 1.0, l_i)
    acc = acc / l_i[:, None]

    out_ptrs = Out + offs_m[:, None] * stride_om + offs_d[None, :] * stride_od
    tl.store(out_ptrs, acc.to(Out.dtype.element_ty), mask=q_mask[:, None])


def attention(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    causal: bool = False,
):
    assert q.is_cuda and k.is_cuda and v.is_cuda
    assert q.ndim == 2 and k.ndim == 2 and v.ndim == 2
    assert q.shape[1] == k.shape[1] == v.shape[1]

    q = q.contiguous()
    k = k.contiguous()
    v = v.contiguous()

    M, D = q.shape
    N = k.shape[0]

    # D 需要适合 tl.dot / tensor core
    assert D in {16, 32, 64, 128}, "D should be 16/32/64/128 in this tutorial kernel"

    out = torch.empty_like(q)

    BLOCK_M = 64
    BLOCK_N = 64

    grid = (triton.cdiv(M, BLOCK_M),)

    attention_fwd[grid](
        q,
        k,
        v,
        out,
        M,
        N,
        D,
        q.stride(0),
        q.stride(1),
        k.stride(0),
        k.stride(1),
        v.stride(0),
        v.stride(1),
        out.stride(0),
        out.stride(1),
        float(D ** -0.5),
        BLOCK_M=BLOCK_M,
        BLOCK_N=BLOCK_N,
        IS_CAUSAL=causal,
        num_warps=4,
        num_stages=2,
    )

    return out


def attention_ref(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    causal: bool = False,
):
    q = q.float()
    k = k.float()
    v = v.float()

    scale = q.shape[-1] ** -0.5
    scores = q @ k.T * scale

    if causal:
        M, N = scores.shape
        mask = torch.triu(
            torch.ones(M, N, dtype=torch.bool, device=scores.device),
            diagonal=1,
        )
        scores = scores.masked_fill(mask, float("-inf"))

    p = torch.softmax(scores, dim=-1)
    return (p @ v).to(q.dtype)


if __name__ == "__main__":
    M = N = 1024
    D = 64

    q = torch.randn(M, D, device="cuda", dtype=torch.float16)
    k = torch.randn(N, D, device="cuda", dtype=torch.float16)
    v = torch.randn(N, D, device="cuda", dtype=torch.float16)

    out = attention(q, k, v, causal=True)
    ref = attention_ref(q, k, v, causal=True)

    torch.testing.assert_close(out, ref, rtol=1e-2, atol=1e-2)
    print("Attention OK")

12.4 这个 Attention Kernel 的核心结构

每个 program 处理一个 Q block

python 复制代码

pid_m = tl.program_id(axis=0)

一个 program 负责：

text 复制代码

Q: [BLOCK_M, D]

遍历所有 KV block

python 复制代码

for start_n in range(0, N, BLOCK_N):

每次加载：

text 复制代码

K: [BLOCK_N, D]
V: [BLOCK_N, D]

QK^T

python 复制代码

qk = tl.dot(q, tl.trans(k)) * sm_scale

得到：

text 复制代码

[BLOCK_M, BLOCK_N]

causal mask

python 复制代码

causal = offs_m[:, None] >= offs_n[None, :]

表示：

text 复制代码

query position >= key position

未来 token 被 mask 掉。

online softmax

python 复制代码

m_new = tl.maximum(m_i, tl.max(qk, axis=1))
p = tl.exp(qk - m_new[:, None])
alpha = tl.exp(m_i - m_new)

这是 FlashAttention 的关键。

12.5 这个版本的局限

没有处理 multi-head / batch；
没有 skip 完全被 causal mask 掉的 KV block；
没有针对共享内存极限调优；
没有复杂 backward；
不支持任意 head_dim；
不适合直接用于生产训练。

但作为从 0 到 1 理解 FlashAttention，已经足够。

12.6 扩展到 multi-head / batch

常见做法：

把 batch/head 编入 grid：

python 复制代码

grid = (num_q_blocks, batch * heads)

给 Q/K/V 增加 batch/head stride；
用 tl.program_id(axis=1) 取 batch/head id；
根据 id 计算 base pointer。

伪代码：

python 复制代码

bh = tl.program_id(axis=1)
Q += bh * stride_qbh
K += bh * stride_kbh
V += bh * stride_vbh
Out += bh * stride_obh

13. Autotune：自动调优

Triton 的一个核心优势是 autotune。

13.1 为什么需要 autotune？

同一个 kernel，不同配置性能差异可能非常大：

text 复制代码

BLOCK_M = 64 vs 128
BLOCK_N = 64 vs 256
BLOCK_K = 32 vs 64
num_warps = 4 vs 8
num_stages = 2 vs 3 vs 4

在不同 GPU、不同 shape 下最优配置往往不同。

13.2 `@triton.autotune`

基本形式：

python 复制代码

@triton.autotune(
    configs=[
        triton.Config({"BLOCK_M": 64, "BLOCK_N": 64}, num_warps=4),
        triton.Config({"BLOCK_M": 128, "BLOCK_N": 128}, num_warps=8),
    ],
    key=["M", "N", "K"],
)
@triton.jit
def kernel(...):
    ...

含义：

枚举多个配置；
每个配置实际 benchmark；
选择最快配置；
key 变化时重新 autotune。

13.3 `key` 的作用

例如：

python 复制代码

key=["M", "N", "K"]

表示当 M/N/K 变化时重新选择配置。

如果 shape 动态变化很多，autotune 开销会变大。

生产环境中可以：

限制配置数量；
对 shape 分桶；
手动指定配置；
预热；
缓存 autotune 结果。

13.4 Config 里可以有什么？

python 复制代码

triton.Config(
    {
        "BLOCK_M": 128,
        "BLOCK_N": 128,
        "BLOCK_K": 32,
        "GROUP_M": 8,
    },
    num_warps=4,
    num_stages=3,
)

字典中的是 kernel 的 tl.constexpr 参数。

13.5 autotune 配置建议

elementwise / reduce

text 复制代码

BLOCK_SIZE: 512 / 1024 / 2048
num_warps: 4 / 8

matmul

text 复制代码

BLOCK_M: 64 / 128 / 256
BLOCK_N: 64 / 128 / 256
BLOCK_K: 16 / 32 / 64
num_warps: 4 / 8
num_stages: 2 / 3 / 4 / 5
GROUP_M: 4 / 8 / 16

attention

text 复制代码

BLOCK_M: 32 / 64 / 128
BLOCK_N: 32 / 64 / 128
num_warps: 4 / 8
num_stages: 1 / 2 / 3

head_dim 越大，block 通常要越小。

14. 性能优化方法论

Triton 性能优化不是靠猜，而是靠流程。

14.1 第一步：判断瓶颈类型

算子一般分两类：

Memory-bound

瓶颈在显存带宽。

常见算子：

text 复制代码

elementwise
reduce
softmax
layernorm
embedding
gather/scatter

优化方向：

减少读写次数；
融合算子；
连续访存；
合理 block size；
避免非对齐访问；
避免 atomics。

Compute-bound

瓶颈在计算单元。

常见算子：

text 复制代码

large matmul
conv
attention 的某些阶段

优化方向：

用 tensor core；
合理 tile size；
提高数据复用；
使用 fp16/bf16/fp8；
fp32 accumulator；
software pipelining；
L2 cache swizzle。

14.2 计算访存比

例如 Vector Add：

text 复制代码

FLOPs = N
Bytes = 12N for fp32
Arithmetic intensity ≈ 1/12

非常低，明显 memory-bound。

Matmul：

text 复制代码

FLOPs = 2MNK
Bytes 大约和 M,N,K 的 tile 复用有关

shape 越大越可能 compute-bound。

14.3 使用 do_bench

不要这样简单计时：

python 复制代码

import time
start = time.time()
fn()
print(time.time() - start)

因为 CUDA 是异步的。

应该：

python 复制代码

import triton.testing as tt

ms = tt.do_bench(lambda: fn())

或者：

python 复制代码

torch.cuda.synchronize()

14.4 用 Nsight Compute 分析

NVIDIA Nsight Compute 可以分析：

text 复制代码

SM throughput
DRAM throughput
L1/L2 hit rate
occupancy
register usage
shared memory usage

常见判断：

DRAM 带宽接近峰值：memory-bound；
SM 计算利用率高：compute-bound；
occupancy 很低：资源使用不当；
shared memory 用满：减小 block/stages；
register spill：减小 tile 或简化逻辑。

14.5 Triton 常见性能技巧

1. 尽量融合

坏：

python 复制代码

y = x + bias
y = silu(y)
y = dropout(y)

好：

text 复制代码

一个 kernel 内完成 add + silu + dropout

2. 保持连续访存

尽量：

python 复制代码

x = x.contiguous()

避免奇怪 stride。

3. 用 fp16/bf16 输入，fp32 累加

python 复制代码

acc = tl.zeros(..., dtype=tl.float32)
acc += tl.dot(a, b)

4. block size 不要太小

太小：

text 复制代码

program 太多
每个 program 工作量太少
启动和调度开销占比高

太大：

text 复制代码

寄存器/shared memory 压力大
occupancy 下降

5. num_stages 不是越大越好

num_stages 增大可以帮助 pipeline，但会消耗更多 shared memory。

如果报错 shared memory 不足，优先：

text 复制代码

减小 BLOCK_M/BLOCK_N/BLOCK_K
减小 num_stages

6. matmul 注意 L2 cache

使用 program swizzle / GROUP_M。

7. attention 中 causal 可以跳过 block

教学版循环所有 KV block，但优化版会根据 causal 限制 KV 范围：

text 复制代码

对于前面的 query，不需要遍历未来 KV。

15. 调试、测试与排错

Triton 的调试体验不如普通 Python，但可以系统化解决。

15.1 正确性测试原则

永远和参考实现对比。

例如：

python 复制代码

def ref(x):
    return torch.softmax(x.float(), dim=-1).to(x.dtype)

y = my_softmax(x)
torch.testing.assert_close(y, ref(x), rtol=1e-3, atol=1e-3)

建议测试多种 shape：

python 复制代码

shapes = [
    (1, 1),
    (17, 128),
    (1024, 1023),
    (4096, 4096),
]

尤其要测非整除 shape：

text 复制代码

M/N/K 不能被 BLOCK_SIZE 整除

因为这类 shape 最容易暴露 mask 问题。

15.2 常见错误一：忘记 mask

错误示例：

python 复制代码

x = tl.load(x_ptr + offsets)

如果 offsets 越界，可能直接 illegal memory access。

正确：

python 复制代码

mask = offsets < n_elements
x = tl.load(x_ptr + offsets, mask=mask, other=0.0)

15.3 常见错误二：stride 传错

例如 matmul：

python 复制代码

a.stride(0), a.stride(1)

必须和 layout 匹配。

如果 tensor 非 contiguous：

python 复制代码

a = a.contiguous()

否则可能性能差或逻辑复杂。

15.4 常见错误三：dtype 不匹配

例如：

text 复制代码

a: fp16
b: fp32

tl.dot 可能不支持。

建议：

python 复制代码

a = a.to(torch.float16)
b = b.to(torch.float16)

计算中：

python 复制代码

acc = tl.zeros(..., dtype=tl.float32)

15.5 常见错误四：block shape 不是 2 的幂

很多 Triton 操作要求 block shape 是 2 的幂。

15.6 常见错误五：tl.dot shape 不满足要求

tl.dot 通常适合：

text 复制代码

BLOCK_M / BLOCK_N / BLOCK_K 是 16 的倍数

例如：

text 复制代码

64, 128, 256

太小或太奇怪可能性能差或报错。

15.7 CPU 解释模式

可以试试：

bash 复制代码

TRITON_INTERPRET=1 python my_script.py

适合：

打印中间值；
检查逻辑；
排查 mask 问题。

限制：

很慢；
不一定支持所有算子；
与真实 GPU 行为可能有差异。

15.8 查看编译产物

Triton 编译链路大致是：

text 复制代码

Python AST
  ↓
Triton IR
  ↓
Triton GPU IR
  ↓
LLVM IR
  ↓
PTX / GPU binary

不同版本中导出方式略有差异。一般可以通过编译后的 kernel 对象查看类似：

text 复制代码

ttir
ttgir
llir
ptx

用于高级调试。

15.9 推荐调试流程

用小 shape；
用 fp32；
和 PyTorch CPU/GPU 参考实现对比；
打印中间 tile；
检查 mask；
检查 stride；
检查 dtype；
逐步放大 shape；
再开启 fp16 / autotune / 性能优化。

16. 与 PyTorch 集成

Triton kernel 最终通常要给 PyTorch 用。

16.1 最简单包装

python 复制代码

def my_op(x):
    y = torch.empty_like(x)
    grid = lambda meta: (...)
    my_kernel[grid](x, y, ..., BLOCK_SIZE=1024)
    return y

然后：

python 复制代码

y = my_op(x)

16.2 训练场景：需要 backward

如果你的算子要参与训练，需要写 backward。

常见方式：

python 复制代码

class MyOp(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x):
        ...
        ctx.save_for_backward(...)
        return y

    @staticmethod
    def backward(ctx, grad_y):
        ...
        return grad_x

然后 forward / backward 各自调用 Triton kernel。

16.3 保存中间量

例如 LayerNorm forward 保存：

text 复制代码

mean
rstd

backward 用：

text 复制代码

mean, rstd, grad_y, weight, x

来计算：

text 复制代码

grad_x, grad_weight, grad_bias

16.4 torch.compile 与 Triton

PyTorch 的 torch.compile 在某些后端中会生成 Triton kernel。

你可以用：

python 复制代码

model = torch.compile(model)

并观察 Inductor/Triton 生成的代码。

调试时可用：

bash 复制代码

TORCH_COMPILE_DEBUG=1 python train.py

具体行为取决于 PyTorch 版本。

16.5 自定义 op

如果你希望算子能进入 torch.export / torch.compile 图，可以考虑 PyTorch custom op 机制。

但对于学习 Triton，先从普通函数包装开始即可。

17. 学习路线与练习项目

这里给一个从 0 到 1 的实战路线。

17.1 第一阶段：入门

目标：能写、能跑、能验证。

Day 1

完成：

环境安装；
Vector Add；
理解 program_id、grid、mask、tl.load、tl.store。

练习：

text 复制代码

写一个 vector mul：
out = x * y

Day 2

完成：

elementwise fused add + relu；
一维 reduce sum；
row reduce。

练习：

text 复制代码

写 row_sum kernel：
输入 [M, N]
输出 [M]

Day 3

完成：

Softmax；
LayerNorm；
RMSNorm。

练习：

text 复制代码

写 RMSNorm：
y = x / sqrt(mean(x^2) + eps) * weight

17.2 第二阶段：核心算子

目标：掌握 tile 编程和 tl.dot。

Day 4

完成 simple matmul。

练习：

text 复制代码

支持 M/N/K 非整除。

Day 5

完成 autotuned matmul。

练习：

text 复制代码

对比不同 BLOCK_M/BLOCK_N/BLOCK_K 的性能。

Day 6

完成 fused bias + activation。

练习：

text 复制代码

写 fused linear bias relu：
不直接调用 torch.matmul，而是把 matmul + bias + relu 融合。

这会比较难，但非常有价值。

17.3 第三阶段：LLM 算子

目标：能写 attention 相关 kernel。

Day 7

完成教学版 attention。

练习：

text 复制代码

支持 causal=False 和 causal=True。

Day 8

扩展 attention：

text 复制代码

加入 batch/head 维度。

Day 9

写 RMSNorm + residual fusion：

text 复制代码

y = RMSNorm(x + residual)

Day 10

写 rotary embedding。

17.4 进阶项目

项目一：Fused Dropout + Residual + LayerNorm

目标：

text 复制代码

y = LayerNorm(x + dropout(x) + residual)

收益：

text 复制代码

减少多次显存读写。

项目二：Fused Cross Entropy

目标：

text 复制代码

log_softmax + nll_loss

重点：

text 复制代码

online log-sum-exp
避免 materialize 大 logits softmax

项目三：Dequantize Matmul

目标：

text 复制代码

A_fp16 @ dequant(W_int8)

适合学习：

text 复制代码

量化
访存优化
算子融合

项目四：Top-k / Sampling Kernel

目标：

text 复制代码

LLM sampling 中的 top-k / top-p

难点：

text 复制代码

reduce
sort/select
warp/block 协作

项目五：Persistent Kernel

目标：

text 复制代码

让少量 program 循环处理多个 tile，
减少 launch 开销，提高 tail 利用率。

适合 matmul / elementwise 小 shape。

18. 常见坑 FAQ

Q1：为什么我的 kernel 结果不对？

优先检查：

mask 是否覆盖所有边界；
stride 是否正确；
grid 是否足够大；
dtype 是否一致；
reduce axis 是否正确；
是否用了未初始化输出 tensor；
是否和参考实现使用相同 dtype。

Q2：为什么报 illegal memory access？

常见原因：

没有 mask；
mask 写错；
offsets 计算错误；
stride 错误；
pointer 偏移超出范围；
grid 过大或过小。

Q3：为什么 autotune 很慢？

因为每个配置都要编译并 benchmark。

解决：

减少 config 数量；
固定 shape 分桶；
对常见 shape 手动选择最佳 config；
提前 warmup；
缓存结果。

Q4：为什么我的 matmul 不如 torch.matmul？

很正常。

原因：

cuBLAS 极其成熟；
你的 config 未充分调优；
shape 不适合当前 block；
layout / stride 不佳；
GPU 架构差异；
没有针对硬件特性优化。

Triton 的优势常在：

text 复制代码

融合
自定义 epilogue
量化
特殊 attention
快速实验

Q5：BLOCK_SIZE 应该设多大？

没有固定答案。

一般：

text 复制代码

elementwise: 512 / 1024 / 2048
matmul M/N: 64 / 128 / 256
matmul K: 16 / 32 / 64
attention M/N: 32 / 64 / 128

需要 autotune。

Q6：num_warps 怎么选？

经验：

text 复制代码

小 tile: 4
中 tile: 4/8
大 tile: 8/16

但最好 autotune。

Q7：num_stages 怎么选？

经验：

text 复制代码

2 ~ 5

如果 shared memory 不足：

text 复制代码

降低 num_stages
降低 block size

Q8：Triton 能完全替代 CUDA 吗？

不能完全替代。

Triton 适合大多数 AI kernel 开发，但极端优化、特殊硬件指令、复杂调度等场景仍可能需要 CUDA / CUTLASS / PTX。

Q9：Triton 适合训练 kernel 吗？

适合，但你需要写 backward。

很多 forward 算子容易写，backward 更考验数学和 layout 理解。

Q10：学习 Triton 最重要的是什么？

不是背 API，而是建立这个思维：

text 复制代码

每个 program 处理一个 tile；
tile 内部做向量化计算；
边界用 mask；
循环加载依赖维度；
尽量复用数据；
尽量融合访存。

19. 附录 A：CUDA 程序员转 Triton 速查表

CUDA 概念	Triton 对应
`__global__` kernel	`@triton.jit`
`threadIdx` / `blockIdx`	`tl.program_id`
一个 thread 处理一个元素	一个 program 处理一个 tile
`__shared__`	编译器自动管理较多
`__syncthreads()`	通常不需要手写
manual vectorization	编译器可能自动处理
software pipeline	`num_stages`
tensor core MMA	`tl.dot`
launch bounds	`num_warps`
gridDim	grid
blockDim	不直接手写，受 `num_warps` 影响

20. 附录 B：如果你指的是 NVIDIA Triton Inference Server

如果你说的不是 OpenAI Triton，而是 NVIDIA Triton Inference Server，那它是另一个东西。

20.1 NVIDIA Triton Inference Server 是什么？

它是一个模型推理服务框架，用于在生产环境部署模型。

支持：

TensorRT
PyTorch
ONNX Runtime
TensorFlow
Python backend
自定义 backend

核心能力：

HTTP / gRPC 推理服务；
dynamic batching；
concurrent model execution；
model repository；
model versioning；
metrics；
ensemble pipeline。

20.2 快速启动

使用 Docker：

bash 复制代码

mkdir -p models

docker run --gpus all \
  -p 8000:8000 \
  -p 8001:8001 \
  -p 8002:8002 \
  -v $(pwd)/models:/models \
  nvcr.io/nvidia/tritonserver:latest \
  tritonserver --model-repository=/models

端口含义：

text 复制代码

8000: HTTP
8001: gRPC
8002: metrics

20.3 Model Repository 结构

示例：

text 复制代码

models/
└── my_model/
    ├── config.pbtxt
    └── 1/
        └── model.onnx

1 表示版本号。

20.4 简单 config.pbtxt 示例

protobuf 复制代码

name: "my_model"
platform: "onnxruntime_onnx"
max_batch_size: 8

input [
  {
    name: "input0"
    data_type: TYPE_FP32
    dims: [ 4 ]
  }
]

output [
  {
    name: "output0"
    data_type: TYPE_FP32
    dims: [ 4 ]
  }
]

20.5 Python client 示例

安装：

bash 复制代码

pip install tritonclient[all]

简单 HTTP client：

python 复制代码

import numpy as np
import tritonclient.http as httpclient

client = httpclient.InferenceServerClient(url="localhost:8000")

inputs = [
    httpclient.InferInput("input0", [1, 4], "FP32"),
]

data = np.random.randn(1, 4).astype(np.float32)
inputs[0].set_data_from_numpy(data)

outputs = [
    httpclient.InferRequestedOutput("output0"),
]

response = client.infer("my_model", inputs=inputs, outputs=outputs)

print(response.as_numpy("output0"))

20.6 什么时候用 Triton Inference Server？

适合：

模型上线服务；
多框架模型统一管理；
需要 dynamic batching；
需要 gRPC/HTTP 接口；
需要多模型并发；
需要 metrics / health / versioning。

不适合：

你只是想写 GPU kernel；
你想优化 PyTorch 算子；
你想写 FlashAttention；
你想做编译器研究。

这些场景应该用 OpenAI Triton。

总结：Triton 学习的关键路径

如果你想真正掌握 OpenAI Triton，建议按这个路径反复练习：

text 复制代码

1. 理解 GPU 并行和内存层次
2. 掌握 Triton 的 program/tile/mask 模型
3. 写 vector add / reduce / softmax / layernorm
4. 写 fused elementwise 算子
5. 写 tiled matmul
6. 使用 autotune
7. 写 attention
8. 做真实项目：fused norm、quant matmul、cross entropy、sampling
9. 用 profiler 分析性能
10. 不断对比 PyTorch / cuBLAS / FlashAttention

Triton 不是靠看会的，而是靠写会的。

建议你从下面的最小目标开始：

text 复制代码

今天：跑通 Vector Add。
明天：写 softmax。
第三天：写 layernorm。
一周内：写一个能跑的 tiled matmul。
两周内：写一个教学版 attention。

如果你能独立完成这些，并理解每个 kernel 的 mask、stride、tile、dtype、性能瓶颈，你就已经真正完成 Triton 从 0 到 1 的入门了。

Triton 从 0 到 1 完整学习教程

目录

1. 学前准备：你需要具备什么

必备基础

建议但非必须

2. Triton 是什么：它解决什么问题

2.1 一句话定义

2.2 为什么需要 Triton？

2.3 Triton 与 CUDA 的核心区别

2.4 Triton 的核心思想：Tile 级编程

2.5 Triton 的优势

2.6 Triton 的局限

3. 环境安装与验证

3.1 创建虚拟环境

3.2 安装 PyTorch

3.3 安装 Triton

3.4 验证环境

3.5 可选环境变量

4. GPU 编程基础：理解 Triton 的前提

4.1 CPU 和 GPU 的区别

4.2 GPU 内存层次

4.3 Kernel 是什么？

4.4 Grid、Program、Block

4.5 Warp 和 num_warps

4.6 num_stages 是什么？

5. Triton 编程模型：Program、Tile、Mask、Pointer

5.1 一个 Triton kernel 的基本结构

5.2 Grid 是什么？

5.3 tl.constexpr：编译期常量

5.4 Pointer Arithmetic：指针运算

5.5 tl.load 和 tl.store

5.6 Mask 是 Triton 的生命线

5.7 tile 级 shape 和 broadcast

5.8 tl.sum / tl.max / tl.min

5.9 tl.dot：矩阵乘法核心

6. 第一个 Triton Kernel：Vector Add

6.1 完整代码

6.2 逐段解释

kernel 签名

获取 program id

计算 offsets

mask

load/store

grid

6.3 benchmark

6.4 这个 kernel 的性能关键点

7. 核心 API 与常用模式

7.1 tl.program_id

7.2 tl.arange

7.3 tl.zeros

7.4 tl.full

7.5 tl.load

7.6 tl.store

7.7 tl.sum

7.8 tl.max

7.9 tl.exp

7.10 tl.where

7.11 tl.dot

7.12 tl.trans

7.13 tl.cdiv

7.14 tl.atomic_add

8. 实例一：Softmax

8.1 数学公式

8.2 完整代码

8.3 关键点解释

一行一个 program

BLOCK_SIZE 必须覆盖整行

数值稳定

dtype

8.4 这个版本的限制

9. 实例二：LayerNorm

9.1 完整代码

9.2 关键点

一行一个 program

fp32 计算

保存 mean / rstd

限制

10. 实例三：Fused Bias + SiLU

10.1 为什么要融合？

10.2 完整代码

7.1 `tl.program_id`

7.2 `tl.arange`

7.3 `tl.zeros`

7.4 `tl.full`

7.5 `tl.load`

7.6 `tl.store`

7.7 `tl.sum`

7.8 `tl.max`

7.9 `tl.exp`

7.10 `tl.where`

7.11 `tl.dot`

7.12 `tl.trans`

7.13 `tl.cdiv`

7.14 `tl.atomic_add`

13.2 `@triton.autotune`

13.3 `key` 的作用