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结构

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))

自动编译优化

# 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

@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

@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最佳实践

内存访问优化

@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)

算子融合判断

# 可融合的组合
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",   # 依赖关系复杂
]

配置选择

# 性能调优参数
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集成

算子注册

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集成

# PyTorch 2.0 torch.compile
model = model.compile(mode="reduce-overhead")
 
# Triton kernel自动使用torch.compile后端
# 无需额外配置

常见优化场景

场景	优化方法	效果
Element-wise融合	编译时合并	减少kernel launch开销
矩阵分块	自动tile	提高cache命中率
向量化访问	使用tl.load(vectorized)	提高带宽利用率
算子重排	调整计算顺序	减少中间内存访问

注意事项

Debug困难：Triton编译错误信息不直观
不支持所有操作：复杂控制流、某些特殊操作需回退到CUDA
Occupancy限制：BLOCK_SIZE过大导致occupancy下降
数值稳定性：FP16计算需注意溢出问题
Memory Alignment：数据对齐影响加载效率

Huajishe Notes

Explorer

2-Triton与CUDA算子

Triton与CUDA算子

Triton概述

核心特点

与CUDA C/C++对比

自动Kernel生成

Triton Kernel结构

自动编译优化

算子融合策略

融合原则

融合示例：LayerNorm

融合收益分析

Flash Attention的Triton实现

Flash Attention Kernel

性能对比

Triton最佳实践

内存访问优化

算子融合判断

配置选择

与其他框架的对比

PyTorch集成

算子注册

Torch.compile集成

常见优化场景

注意事项

Graph View

Table of Contents