Skip to content

Commit c25932f

Browse files
AndSonderAndSonder
committed
add triton 编程范式 doc
1 parent 582c219 commit c25932f

File tree

2 files changed

+782
-0
lines changed

2 files changed

+782
-0
lines changed

‎docs/18_triton/01_triton_programming_paradigms/README.md‎

Lines changed: 358 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,358 @@
1+
# Triton 编程范式入门
2+
3+
## 前言
4+
5+
此文档面向「有 CUDA 编程基础」的同学，通过 CUDA vs Triton 对比的方式，帮助你快速理解 Triton 的编程范式。本文将从最简单的向量加法入手，展示两种编程模型的核心差异，并通过实际代码示例帮助你建立从 CUDA 到 Triton 的思维转换。
6+
7+
期望在阅读完本文后，你能够理解 Triton 的核心抽象思想，掌握其基本语法，并能够将简单的 CUDA Kernel 改写为 Triton 版本。
8+
9+
## 一、从向量加法看两种编程范式
10+
11+
向量加法是最简单的并行任务，非常适合用来对比 CUDA 和 Triton 的编程思想差异。我们先从任务本身说起。
12+
13+
### 1. 任务描述
14+
15+
给定两个长度为 N 的向量 `x` 和 `y`，计算 `out[i] = x[i] + y[i]`，输出新的向量 `out`。这是一个典型的 element-wise 操作，每个输出元素的计算都是独立的，天然适合并行化。
16+
17+
### 2. CUDA 实现：线程级并行
18+
19+
在 CUDA 中，你需要思考：**每个线程处理哪个元素？**
20+
21+
```cuda
22+
__global__ void vector_add_cuda(float *x, float *y, float *out, int n) {
23+
// 每个线程计算自己的全局索引
24+
int idx = blockIdx.x * blockDim.x + threadIdx.x;
25+
26+
// 边界检查
27+
if (idx < n) {
28+
out[idx] = x[idx] + y[idx]; // 每个线程处理一个元素
29+
}
30+
}
31+
32+
// Host 端调用
33+
int main() {
34+
int N = 10000000;
35+
// ... 省略内存分配和数据初始化 ...
36+
37+
int blockSize = 256; // 每个 Block 有 256 个线程
38+
int numBlocks = (N + blockSize - 1) / blockSize;
39+
40+
vector_add_cuda<<<numBlocks, blockSize>>>(cuda_x, cuda_y, cuda_out, N);
41+
cudaDeviceSynchronize();
42+
43+
return 0;
44+
}
45+
```
46+
47+
在 CUDA 的编程模型中，你需要**将任务分解到每一个线程。每个线程就像一个独立的工人**，负责处理一个元素。线程通过 `blockIdx` 和 `threadIdx` 这两个内置变量计算出自己负责的全局索引 `idx`，然后访问对应位置的数据。由于数组长度 `n` 可能不是 Block 大小的整数倍，我们需要用 `if (idx < n)` 来做边界检查，但这可能导致 Warp Divergence，影响性能。
48+
49+
50+
### 3. Triton 实现：数据块级并行
51+
52+
在 Triton 中，你需要思考：**每个 Program Instance 处理哪批元素？**
53+
54+
```python
55+
import torch
56+
import triton
57+
import triton.language as tl
58+
59+
@triton.jit
60+
def vector_add_triton(
61+
x_ptr, # 输入向量 x 的指针
62+
y_ptr, # 输入向量 y 的指针
63+
out_ptr, # 输出向量 out 的指针
64+
n_elements, # 向量长度
65+
BLOCK_SIZE: tl.constexpr, # 每个 Program 处理的元素数（编译时常量）
66+
):
67+
# 1. 获取当前 Program 的 ID（类似 blockIdx.x）
68+
pid = tl.program_id(axis=0)
69+
70+
# 2. 计算当前 Program 负责的元素范围
71+
block_start = pid * BLOCK_SIZE
72+
offsets = block_start + tl.arange(0, BLOCK_SIZE)
73+
# offsets 是一个向量！例如：[0, 1, 2, ..., BLOCK_SIZE-1]
74+
75+
# 3. 创建边界检查的 mask（向量化的边界检查）
76+
mask = offsets < n_elements
77+
78+
# 4. 向量化加载数据（一次加载 BLOCK_SIZE 个元素）
79+
x = tl.load(x_ptr + offsets, mask=mask)
80+
y = tl.load(y_ptr + offsets, mask=mask)
81+
82+
# 5. 向量化计算
83+
out = x + y
84+
85+
# 6. 向量化存储
86+
tl.store(out_ptr + offsets, out, mask=mask)
87+
88+
89+
# Host 端调用
90+
def add(x: torch.Tensor, y: torch.Tensor):
91+
assert x.is_cuda and y.is_cuda, "输入必须在 GPU 上"
92+
assert x.shape == y.shape, "输入形状必须一致"
93+
94+
out = torch.empty_like(x)
95+
n_elements = out.numel()
96+
97+
# 计算需要多少个 Program Instance
98+
grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']),)
99+
100+
# 启动 Kernel
101+
vector_add_triton[grid](x, y, out, n_elements, BLOCK_SIZE=1024)
102+
103+
return out
104+
105+
106+
# 使用示例
107+
if __name__ == "__main__":
108+
N = 10_000_000
109+
x = torch.randn(N, device='cuda', dtype=torch.float32)
110+
y = torch.randn(N, device='cuda', dtype=torch.float32)
111+
112+
out = add(x, y)
113+
114+
# 验证正确性
115+
expected = x + y
116+
assert torch.allclose(out, expected), "结果不匹配！"
117+
```
118+
119+
相比之下，Triton 的编程模型更像是把线程们组织成团队。每个 Program Instance 处理一批元素（BLOCK_SIZE 个），而不是单个元素。你通过 `tl.arange` 生成一个向量化的偏移量数组，这个数组包含了当前 Program 要处理的所有元素的索引。边界检查也变成了向量化的操作：`mask = offsets < n_elements` 会生成一个布尔向量，标记哪些位置是有效的。这种向量化的 mask 机制可以有效避免 Warp Divergence。实际上，**Triton 中的所有操作都是向量化的（SIMD），这是它和 CUDA 最本质的区别**。
120+
121+
122+
## 二、核心概念对比
123+
124+
### 1. 概念映射表
125+
126+
| CUDA 概念 | Triton 概念 | 说明 |
127+
|-----------|-------------|------|
128+
| `__global__` | `@triton.jit` | Kernel 函数标记 |
129+
| Thread | **不存在** | Triton 不暴露线程概念 |
130+
| Block | **Program Instance** | 并行执行的基本单元 |
131+
| `blockIdx.x` | `tl.program_id(axis=0)` | 获取当前 Block/Program 的 ID |
132+
| `threadIdx.x` | **不存在** | Triton 自动向量化 |
133+
| `blockDim.x` | `BLOCK_SIZE` | 每个 Block 处理的元素数 |
134+
| `idx = blockIdx.x * blockDim.x + threadIdx.x` | `offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)` | 索引计算方式不同 |
135+
| `if (idx < n)` | `mask = offsets < n` | 边界检查方式不同 |
136+
| `x[idx]` | `tl.load(x_ptr + offsets, mask=mask)` | 内存访问方式不同 |
137+
| `out[idx] = value` | `tl.store(out_ptr + offsets, value, mask=mask)` | 内存写入方式不同 |
138+
| `__syncthreads()` | **自动处理** | Triton 不需要手动同步 |
139+
140+
141+
### 2. 线程索引 vs 向量化偏移
142+
143+
这是 CUDA 和 Triton 最核心的区别。
144+
145+
#### CUDA：标量索引
146+
147+
```cuda
148+
// 每个线程计算自己的索引（标量）
149+
int idx = blockIdx.x * blockDim.x + threadIdx.x;
150+
151+
// 访问一个元素
152+
float val = x[idx];
153+
```
154+
155+
在 CUDA 中，你的思考方式是："我是第 `idx` 号线程，我处理第 `idx` 个元素"。这是一种非常直观的一对一映射关系。
156+
157+
#### Triton：向量偏移
158+
159+
```python
160+
# 每个 Program 计算一批偏移量（向量）
161+
pid = tl.program_id(axis=0)
162+
offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
163+
164+
# 访问一批元素
165+
x = tl.load(x_ptr + offsets, mask=mask)
166+
```
167+
168+
而在 Triton 中，你需要转变思维："我是第 `pid` 号 Program，我处理第 `[start, start+1, ..., start+BLOCK_SIZE-1]` 批元素"。这里最关键的区别在于，CUDA 的 `idx` 是一个标量（单个整数），而 Triton 的 `offsets` 是一个向量（整数数组）。
169+
170+
171+
### 3. 边界检查：`if` vs `mask`
172+
173+
#### CUDA 的方式
174+
175+
```cuda
176+
if (idx < n) {
177+
out[idx] = x[idx] + y[idx]; // 可能导致 Warp Divergence
178+
}
179+
```
180+
181+
使用标量 `if` 的问题在于，同一个 Warp 内的 32 个线程可能会因为边界条件走不同的分支。比如在最后一个 Block 中，前面的线程满足 `idx < n` 继续执行，而后面的线程不满足条件被过滤掉。这种分支分歧（Divergence）会导致 Warp 内的线程无法同步执行，从而降低性能。
182+
183+
#### Triton 的方式
184+
185+
```python
186+
mask = offsets < n_elements # 向量化比较，生成布尔向量
187+
188+
x = tl.load(x_ptr + offsets, mask=mask) # 只加载有效位置
189+
out = x + y
190+
tl.store(out_ptr + offsets, out, mask=mask) # 只存储有效位置
191+
```
192+
193+
Triton 的 `mask` 机制则完全不同。`mask = offsets < n_elements` 是一个向量化的比较操作，会生成一个布尔向量。在 `tl.load` 和 `tl.store` 中使用 mask 时，编译器会生成 predicated instructions（带谓词的指令），这是 GPU 硬件原生支持的特性。每个线程都有独立的 predicate 寄存器，可以在不引起分支的情况下选择性地执行指令。这样既保证了正确性，又避免了 Warp Divergence，代码还更简洁。
194+
195+
### 4. 具体示例：处理 10000 个元素
196+
197+
假设我们要处理 `N = 10000` 个元素。
198+
199+
#### CUDA 的执行方式
200+
201+
```
202+
配置：blockSize = 256, numBlocks = 40
203+
204+
Grid
205+
├── Block[0] Block[1] ... Block[39]
206+
│ ├── Thread[0] ├── Thread[0] ├── Thread[0]
207+
│ │ 处理 idx=0 │ 处理 idx=256 │ 处理 idx=9984
208+
│ ├── Thread[1] ├── Thread[1] ├── Thread[1]
209+
│ │ 处理 idx=1 │ 处理 idx=257 │ 处理 idx=9985
210+
│ ├── ... ├── ... ├── ...
211+
│ └── Thread[255] └── Thread[255] └── Thread[255]
212+
│ 处理 idx=255 处理 idx=511 处理 idx=10239 (越界！)
213+
214+
总共启动: 40 × 256 = 10240 个线程
215+
其中 240 个线程因为 if (idx < 10000) 被过滤掉
216+
```
217+
218+
#### Triton 的执行方式
219+
220+
```
221+
配置：BLOCK_SIZE = 1024, numPrograms = 10
222+
223+
Grid
224+
├── Program[0] Program[1] ... Program[9]
225+
│ 处理元素 处理元素 处理元素
226+
│ [0~1023] [1024~2047] [9216~10239]
227+
│ (其中 10000~10239 被 mask 过滤)
228+
229+
总共启动: 10 个 Program Instance
230+
每个 Program 处理 1024 个元素（向量化）
231+
Triton 内部会自动映射到合适的线程配置
232+
```
233+
234+
从这个例子可以看出，CUDA 启动了 10240 个线程，你需要思考"我是第几号线程"。而 Triton 只启动了 10 个 Program Instance，你要思考的是"我处理哪批数据"。这种抽象层次的提升，让代码更简洁，也更容易理解。
235+
236+
## 三、Grid 配置对比
237+
238+
### CUDA 的 Grid 配置
239+
240+
```cuda
241+
int blockSize = 256; // 每个 Block 有 256 个线程
242+
int numBlocks = (n + blockSize - 1) / blockSize; // 向上取整
243+
244+
my_kernel<<<numBlocks, blockSize>>>(args);
245+
```
246+
247+
在 CUDA 中，你需要同时指定 `numBlocks`（有多少个 Block）和 `blockSize`（每个 Block 有多少个 Thread）。`blockSize` 的选择受到 SM 资源的限制，通常设置为 128/256/512。如果设置得太大，会因为寄存器和 Shared Memory 的限制导致 Occupancy 下降。
248+
249+
### Triton 的 Grid 配置
250+
251+
```python
252+
BLOCK_SIZE = 1024 # 每个 Program Instance 处理 1024 个元素
253+
254+
grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']),)
255+
256+
my_kernel[grid](args, BLOCK_SIZE=BLOCK_SIZE)
257+
```
258+
259+
Triton 的 Grid 配置则简单得多，你只需要指定 `BLOCK_SIZE`（每个 Program 处理多少元素），Grid 的大小会通过 `triton.cdiv(n, BLOCK_SIZE)` 自动计算。这里有一个重要的概念需要澄清：Triton 的 `BLOCK_SIZE` 并不等同于 CUDA 的 `blockDim.x`。CUDA 的 `blockDim.x` 表示线程数，而 Triton 的 `BLOCK_SIZE` 表示元素数。Triton 编译器会根据你指定的 `BLOCK_SIZE`，自动选择合适的线程配置来映射这些元素。
260+
261+
根据经验，Triton 的 `BLOCK_SIZE` 通常设置为 1024/2048/4096，比 CUDA 的 `blockDim` 要大得多。建议从 1024 开始尝试，然后根据实际性能进行调整。
262+
263+
264+
## 四、课后练习
265+
266+
请打开 `01_exercises.py` 完成以下三个练习：练习 1 实现 AXPY 操作（$Z = \alpha \cdot X + Y$），巩固基本的向量化加载和存储；练习 2 测试不同 `BLOCK_SIZE` 的性能影响，理解为什么 Triton 的最优 `BLOCK_SIZE` 比 CUDA 的 `blockDim` 要大；练习 3 实现 1D 卷积，体会如何用向量化方式处理滑动窗口操作。每个练习都包含了测试函数和思考题。
267+
268+
## 五、常见问题 FAQ
269+
270+
### Q1: Triton 的 BLOCK_SIZE 应该设置多大？
271+
272+
**A**: 如果你有 CUDA 经验，需要注意 CUDA 的经验值在这里不适用。CUDA 的 `blockDim.x` 通常设置为 128/256/512，而 Triton 的 `BLOCK_SIZE` 通常要大得多，一般是 1024/2048/4096。这是因为 Triton 的 `BLOCK_SIZE` 表示的是元素数，而不是线程数。建议从 1024 开始尝试，然后根据性能 profiling 的结果进行调整。影响最优 `BLOCK_SIZE` 的因素包括：寄存器使用量、Shared Memory 大小、以及数据复用程度。
273+
274+
275+
### Q2: Triton 内部到底有没有线程？性能会比 CUDA 差吗？
276+
277+
**A**: Triton 内部是有线程的，只是抽象层次更高，不暴露给程序员。Triton 编译器会将你写的向量化代码编译成高效的 PTX（GPU 汇编），最终还是在 GPU 的线程上执行。在性能方面，对于简单算子（如 element-wise 操作），Triton 的性能可以接近手写的优化 CUDA 代码；对于复杂算子（如 Flash Attention），Triton 可以达到优化后 CUDA 的 95% 以上的性能。但在开发效率方面，Triton 远远领先于 CUDA。
278+
279+
280+
### Q3: mask 操作会导致性能下降吗？（类似 Warp Divergence）
281+
282+
**A**: Triton 的 `mask` 是向量化的，编译器会生成 predicated instructions（带谓词的指令），不会像 CUDA 的标量 `if` 那样导致严重的 Warp Divergence。性能损失通常可以忽略。从技术细节来看，现代 GPU 支持 predicated execution，每个线程都有独立的 predicate 寄存器。Triton 编译器会自动将 `mask` 映射到这些硬件特性，因此可以在不引入分支的情况下实现条件执行。
283+
284+
285+
### Q4: 什么时候不能用 Triton？
286+
287+
**A**: 以下场景建议使用 CUDA：
288+
1. 需要显式管理 Shared Memory 布局（如手动消除 Bank Conflicts）
289+
2. 需要使用 Warp-level primitives（`__shfl_`, `__ballot_`, `__syncwarp`）
290+
3. 需要动态并行（Dynamic Parallelism）
291+
4. 算法严重依赖线程间细粒度通信
292+
5. 需要与现有 CUDA 代码库深度集成
293+
294+
### Q5: 如何从 CUDA 代码迁移到 Triton？
295+
296+
**A**: 五步迁移法：
297+
298+
1. **识别数据访问模式**：你的 Kernel 主要做什么？（element-wise、reduction、matmul？）
299+
300+
2. **改变思维方式**：
301+
- ❌ "每个线程处理第 `idx` 个元素"
302+
- ✅ "每个 Program 处理第 `[start...end]` 批元素"
303+
304+
3. **转换索引计算**：
305+
```cuda
306+
// CUDA
307+
int idx = blockIdx.x * blockDim.x + threadIdx.x;
308+
```
309+
```python
310+
# Triton
311+
pid = tl.program_id(0)
312+
offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
313+
```
314+
315+
4. **转换边界检查**：
316+
```cuda
317+
// CUDA
318+
if (idx < n) { ... }
319+
```
320+
```python
321+
# Triton
322+
mask = offsets < n
323+
```
324+
325+
5. **转换内存访问**：
326+
```cuda
327+
// CUDA
328+
float x = x_ptr[idx];
329+
out[idx] = result;
330+
```
331+
```python
332+
# Triton
333+
x = tl.load(x_ptr + offsets, mask=mask)
334+
tl.store(out_ptr + offsets, result, mask=mask)
335+
```
336+
337+
338+
## 六、学习检查清单
339+
340+
完成本小节后，你应该能够理解 Triton 的 Program Instance 概念以及它和 CUDA Block 的本质区别，掌握用 `tl.arange` 生成向量化偏移来代替 CUDA 的线程索引计算，理解 Triton 的 `mask` 如何优雅地避免 Warp Divergence，能够将简单的 CUDA element-wise Kernel 改写为 Triton 版本，并深刻理解 Triton 的抽象级别是 Block-level 而非 Thread-level。
341+
342+
下表展示了从 CUDA 思维到 Triton 思维的转变：
343+
344+
| ❌ CUDA 思维 | ✅ Triton 思维 |
345+
|-------------|----------------|
346+
| "这个线程处理第 `idx` 个元素" | "这个 Program 处理第 `[start, ..., end]` 批元素" |
347+
| "`idx` 是一个整数" | "`offsets` 是一个整数向量" |
348+
| "用 `if (idx < n)` 检查边界" | "用 `mask = offsets < n` 向量化检查" |
349+
| "直接访问 `x[idx]`" | "显式 `tl.load(x_ptr + offsets)`" |
350+
351+
352+
## 附：参考资料
353+
354+
- [Triton 官方教程](https://triton-lang.org/main/getting-started/tutorials/index.html)
355+
- [OpenAI Triton GitHub](https://github.com/openai/triton)
356+
- [CUDA Programming Guide](https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html)
357+
358+
**下一步**：完成所有练习后，进入 **Module 02: 内存与数据搬运**，学习更复杂的内存访问模式！

0 commit comments

Comments
 (0)