# 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**：数据对齐影响加载效率