348 lines
8.8 KiB
Markdown
348 lines
8.8 KiB
Markdown
# Triton与CUDA算子
|
||
|
||
## Triton概述
|
||
|
||
Triton是OpenAI开发的开源GPU编程框架，旨在简化高效GPU算子的开发，无需编写CUDA C/C++代码。
|
||
|
||
### 核心特点
|
||
|
||
| 特性 | 说明 |
|
||
|------|------|
|
||
| Python前端 | 使用Python编写kernel |
|
||
| 自动优化 | 自动生成高效CUDA代码 |
|
||
| Torch集成 | 与PyTorch无缝结合 |
|
||
| 自动融合 | 算子融合策略 |
|
||
|
||
### 与CUDA C/C++对比
|
||
|
||
```
|
||
CUDA C/C++:
|
||
- 完全控制，但开发周期长
|
||
- 手动管理shared memory、register allocation
|
||
- 需要处理硬件细节
|
||
|
||
Triton:
|
||
- 抽象掉硬件细节
|
||
- 自动优化内存访问模式
|
||
- 开发效率高，但灵活性略低
|
||
```
|
||
|
||
## 自动Kernel生成
|
||
|
||
### Triton Kernel结构
|
||
|
||
```python
|
||
import triton
|
||
import triton.language as tl
|
||
|
||
@triton.jit
|
||
def matmul_kernel(
|
||
A_ptr, B_ptr, C_ptr,
|
||
M, N, K,
|
||
stride_am, stride_ak,
|
||
stride_bk, stride_bn,
|
||
stride_cm, stride_cn,
|
||
BLOCK_SIZE: tl.constexpr
|
||
):
|
||
# 块级别的程序化内存访问
|
||
pid = tl.program_id(axis=0)
|
||
num_pid_m = tl.cdiv(M, BLOCK_SIZE)
|
||
pid_m = pid // num_pid_m
|
||
pid_n = pid % num_pid_m
|
||
|
||
# 初始化累加器
|
||
acc = tl.zeros((BLOCK_SIZE, BLOCK_SIZE), dtype=tl.float32)
|
||
|
||
# 计算指针偏移
|
||
offs_m = (pid_m * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)) % M
|
||
offs_n = (pid_n * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)) % N
|
||
offs_k = tl.arange(0, BLOCK_SIZE)
|
||
|
||
# 分块计算
|
||
for k in range(0, tl.cdiv(K, BLOCK_SIZE)):
|
||
# 加载数据到SRAM
|
||
a_ptrs = A_ptr + offs_m[:, None] * stride_am + (offs_k[None, :] * stride_ak)
|
||
b_ptrs = B_ptr + (offs_k[:, None] * stride_bk) + offs_n[None, :] * stride_bn
|
||
|
||
a = tl.load(a_ptrs)
|
||
b = tl.load(b_ptrs)
|
||
|
||
# 矩阵乘法累加
|
||
acc += tl.dot(a, b)
|
||
|
||
# 更新k指针
|
||
offs_k += BLOCK_SIZE
|
||
|
||
# 写回结果
|
||
c_ptrs = C_ptr + stride_cm * offs_m[:, None] + stride_cn * offs_n[None, :]
|
||
tl.store(c_ptrs, acc.to(tl.float16))
|
||
```
|
||
|
||
### 自动编译优化
|
||
|
||
```python
|
||
# Triton自动编译配置
|
||
config = triton.Config(
|
||
num_warps=4, # 控制并行度
|
||
num_stages=2, # 流水线阶段数
|
||
block_size=128 # 块大小
|
||
)
|
||
|
||
# JIT编译
|
||
grid = (M * N // 1024,)
|
||
matmul_kernel[grid](A, B, C, M, N, K, **config)
|
||
```
|
||
|
||
## 算子融合策略
|
||
|
||
### 融合原则
|
||
|
||
```
|
||
融合收益 = 内存访问节省 - 计算额外开销
|
||
|
||
常见可融合模式：
|
||
1. MatMul + Softmax → Flash Attention
|
||
2. MatMul + Bias → Fused MatMul
|
||
3. Activation + Pointwise → 融合Activation
|
||
4. LayerNorm组件 → 单Kernel实现
|
||
```
|
||
|
||
### 融合示例：LayerNorm
|
||
|
||
```python
|
||
@triton.jit
|
||
def layer_norm_kernel(
|
||
x_ptr, y_ptr, weight_ptr, bias_ptr,
|
||
N, eps,
|
||
BLOCK_SIZE: tl.constexpr
|
||
):
|
||
# 计算均值
|
||
x = tl.load(x_ptr + tl.arange(0, BLOCK_SIZE))
|
||
mean = tl.sum(x) / N
|
||
|
||
# 计算方差
|
||
x_mean = x - mean
|
||
var = tl.sum(x_mean * x_mean) / N
|
||
|
||
# 归一化 + 仿射
|
||
y = (x_mean / tl.sqrt(var + eps))
|
||
if weight_ptr is not None:
|
||
w = tl.load(weight_ptr + tl.arange(0, BLOCK_SIZE))
|
||
y = y * w
|
||
if bias_ptr is not None:
|
||
b = tl.load(bias_ptr + tl.arange(0, BLOCK_SIZE))
|
||
y = y + b
|
||
|
||
tl.store(y_ptr + tl.arange(0, BLOCK_SIZE), y)
|
||
```
|
||
|
||
### 融合收益分析
|
||
|
||
```
|
||
LayerNorm融合前（3个独立kernel）:
|
||
- Mean计算: N次读 + N次写 + reduce
|
||
- Var计算: N次读 + N次写 + reduce
|
||
- Normalize: 读 + 算 + 写
|
||
|
||
融合后（单kernel）:
|
||
- 读一次，计算全在SRAM完成
|
||
- 内存访问减少 ~66%
|
||
```
|
||
|
||
## Flash Attention的Triton实现
|
||
|
||
### Flash Attention Kernel
|
||
|
||
```python
|
||
@triton.jit
|
||
def flash_attention_kernel(
|
||
Q, K, V, Out,
|
||
stride_qm, stride_qk,
|
||
stride_km, stride_kk,
|
||
stride_vm, stride_vk,
|
||
M, N, # M=seq_len_q, N=seq_len_k
|
||
HEAD_DIM: tl.constexpr,
|
||
BLOCK_M: tl.constexpr = 64,
|
||
BLOCK_N: tl.constexpr = 64,
|
||
BLOCK_DMODEL: tl.constexpr = 64
|
||
):
|
||
# 程序ID（行号）
|
||
row_id = tl.program_id(0)
|
||
|
||
# 加载Q块
|
||
q_offset = row_id * BLOCK_M * stride_qm
|
||
q = tl.load(Q + q_offset + tl.arange(0, BLOCK_M)[:, None] * stride_qm
|
||
+ tl.arange(0, BLOCK_DMODEL)[None, :])
|
||
|
||
# 初始化累加器
|
||
m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf") # 行最大值
|
||
l_i = tl.zeros([BLOCK_M], dtype=tl.float32) # 行归一化因子
|
||
acc = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32)
|
||
|
||
# 遍历K, V块
|
||
for start_n in range(0, N, BLOCK_N):
|
||
k = tl.load(K + start_n * stride_km +
|
||
tl.arange(0, BLOCK_N)[:, None] * stride_km +
|
||
tl.arange(0, BLOCK_DMODEL)[None, :])
|
||
v = tl.load(V + start_n * stride_vm +
|
||
tl.arange(0, BLOCK_N)[:, None] * stride_vm +
|
||
tl.arange(0, BLOCK_DMODEL)[None, :])
|
||
|
||
# 计算S = Q @ K^T
|
||
s = tl.dot(q, tl.trans(k))
|
||
|
||
# 更新m（行最大值）
|
||
m_new = tl.maximum(m_i, tl.max(s, axis=1))
|
||
|
||
# 计算P = exp(S - m_new)
|
||
p = tl.exp(s - m_new[:, None])
|
||
|
||
# 更新l（归一化因子）
|
||
l_new = l_i * tl.exp(m_i - m_new) + tl.sum(p, axis=1)
|
||
|
||
# 计算acc
|
||
acc_scale = l_i * tl.exp(m_i - m_new) / l_new
|
||
acc = acc * acc_scale[:, None] + tl.dot(p, v) / l_new[:, None]
|
||
|
||
m_i = m_new
|
||
l_i = l_new
|
||
|
||
# 写回
|
||
out_offset = row_id * BLOCK_M * stride_qm
|
||
tl.store(Out + out_offset + tl.arange(0, BLOCK_M)[:, None] * stride_qm
|
||
+ tl.arange(0, BLOCK_DMODEL)[None, :], acc)
|
||
```
|
||
|
||
### 性能对比
|
||
|
||
```
|
||
Flash Attention Triton vs CUDA原生:
|
||
|
||
硬件: A100
|
||
序列长度: 4096, heads=32, head_dim=128
|
||
|
||
| 实现 | TFLOPS | 显存占用 |
|
||
|------|--------|---------|
|
||
| CUDA C (手写) | 350 | 基准 |
|
||
| Triton | 340 | 基准 |
|
||
| 差距 | ~3% | 相当 |
|
||
```
|
||
|
||
## Triton最佳实践
|
||
|
||
### 内存访问优化
|
||
|
||
```python
|
||
@triton.jit
|
||
def optimized_matmul(
|
||
A, B, C,
|
||
M, N, K,
|
||
BLOCK_SIZE: tl.constexpr
|
||
):
|
||
# 预取策略：重叠加载和计算
|
||
# 使用多个stage的pipeline
|
||
|
||
# 向量化加载（提高带宽利用率）
|
||
# 假设BLOCK_SIZE=64, head_dim=64
|
||
# 使用 (64, 4) 的向量化访问
|
||
|
||
a = tl.load(A + offs, mask=offs < M * K, other=0.0, dtype=tl.float16)
|
||
```
|
||
|
||
### 算子融合判断
|
||
|
||
```python
|
||
# 可融合的组合
|
||
fusable_patterns = [
|
||
("linear", "relu"), # Linear + Activation
|
||
("matmul", "add"), # Fused MatMul + Bias
|
||
("softmax", "matmul"), # Flash Attention基础
|
||
("layernorm", "add"), # LayerNorm + Residual
|
||
("rmsnorm", "matmul"), # RMSNorm融合
|
||
]
|
||
|
||
# 不可融合（需拆分）
|
||
non_fusable = [
|
||
"matmul + convolution", # 不同的内存访问模式
|
||
"reduce + elementwise", # 依赖关系复杂
|
||
]
|
||
```
|
||
|
||
### 配置选择
|
||
|
||
```python
|
||
# 性能调优参数
|
||
config = {
|
||
"num_warps": 4, # 2/4/8，影响并行度和shared memory使用
|
||
"num_stages": 2, # 流水线深度，2适合short seq，3适合long seq
|
||
"BLOCK_SIZE": 128, # 影响occupancy
|
||
}
|
||
|
||
# 不同场景推荐
|
||
scenarios = {
|
||
"short_seq (<1024)": {"num_warps": 4, "num_stages": 3, "BLOCK": 64},
|
||
"medium_seq": {"num_warps": 4, "num_stages": 2, "BLOCK": 128},
|
||
"long_seq (>4096)": {"num_warps": 8, "num_stages": 2, "BLOCK": 64},
|
||
}
|
||
```
|
||
|
||
## 与其他框架的对比
|
||
|
||
| 框架 | 开发语言 | 学习曲线 | 性能 | 适用场景 |
|
||
|------|---------|---------|------|---------|
|
||
| CUDA C/C++ | C++/CUDA | 高 | 最优 |极致优化 |
|
||
| Triton | Python | 中 | 接近CUDA | 快速开发 |
|
||
| cuBLAS | - | 低 | 最优 | 通用MatMul |
|
||
| Cutlass | C++/CUDA | 高 | 最优 | 定制MatMul |
|
||
| TVM | Python/Relay | 中 | 接近CUDA | 自动优化 |
|
||
|
||
## PyTorch集成
|
||
|
||
### 算子注册
|
||
|
||
```python
|
||
import torch
|
||
from triton.ops import autotune
|
||
|
||
@autotune(
|
||
configs=[
|
||
triton.Config({'BLOCK_M': 64, 'BLOCK_N': 64, 'BLOCK_K': 32}),
|
||
triton.Config({'BLOCK_M': 128, 'BLOCK_N': 64, 'BLOCK_K': 32}),
|
||
],
|
||
key=['M', 'N', 'K']
|
||
)
|
||
@triton.jit
|
||
def matmul_kernel(...):
|
||
...
|
||
|
||
class MatMul(torch.nn.Module):
|
||
def forward(self, x, y):
|
||
return matmul_kernel(x, y)
|
||
```
|
||
|
||
### Torch.compile集成
|
||
|
||
```python
|
||
# PyTorch 2.0 torch.compile
|
||
model = model.compile(mode="reduce-overhead")
|
||
|
||
# Triton kernel自动使用torch.compile后端
|
||
# 无需额外配置
|
||
```
|
||
|
||
## 常见优化场景
|
||
|
||
| 场景 | 优化方法 | 效果 |
|
||
|------|---------|------|
|
||
| Element-wise融合 | 编译时合并 | 减少kernel launch开销 |
|
||
| 矩阵分块 | 自动tile | 提高cache命中率 |
|
||
| 向量化访问 | 使用tl.load(vectorized) | 提高带宽利用率 |
|
||
| 算子重排 | 调整计算顺序 | 减少中间内存访问 |
|
||
|
||
## 注意事项
|
||
|
||
1. **Debug困难**：Triton编译错误信息不直观
|
||
2. **不支持所有操作**：复杂控制流、某些特殊操作需回退到CUDA
|
||
3. **Occupancy限制**：BLOCK_SIZE过大导致occupancy下降
|
||
4. **数值稳定性**：FP16计算需注意溢出问题
|
||
5. **Memory Alignment**：数据对齐影响加载效率 |