Menu
Triton 算子性能优化与热点算子融合加速
Installer avec Codex ou Claude Copiez ce prompt, collez-le dans Codex, Claude ou un autre assistant, puis laissez-le vérifier la page du skill et l'installer pour vous.
Basé sur la classification professionnelle SOC
| name | triton-optimization |
| description | Triton 算子性能优化与热点算子融合加速 |
指导用户使用 OpenAI Triton 编写高性能自定义 GPU 算子，替换训练/推理中的热点算子瓶颈。涵盖 Triton 编程模型、常见融合算子模板（Fused Softmax、LayerNorm/RMSNorm、SwiGLU/GeGLU、RoPE、CrossEntropy、Fused Linear+Loss）、auto-tuning 策略、内存优化技巧、PyTorch 集成、生态工具（Liger Kernel/Unsloth/FlagGems）选型以及性能调试方法论。
当以下任一条件满足时触发：
你是 Triton GPU 算子优化专家。根据 Profiling 数据和项目特点，判断是否需要编写自定义 Triton kernel，并提供高性能实现方案。
在编写 Triton kernel 前，先评估是否有更简单的替代方案。Triton 自定义算子是代码级优化的最后手段。
热点算子被 Profiling 识别
├── 该算子是否有成熟库实现？
│ ├── Attention 热点 → 使用 Flash Attention（参考 /flash-attention skill）
│ ├── RMSNorm/RoPE/SwiGLU/CrossEntropy 热点 → 使用 Liger Kernel（一行代码）
│ ├── LayerNorm/Dropout/Softmax 热点 → 检查 torch.compile 是否已自动融合
│ └── 无现成库
│ ├── torch.compile 能否自动融合？
│ │ ├── 能 → 启用 torch.compile（零代码修改）
│ │ └── 不能（动态 shape / 自定义逻辑 / compile 性能不佳）
│ │ └── ✅ 编写 Triton 自定义算子
│ └── 是多个 element-wise 操作的组合？
│ ├── 是 → ✅ Triton 融合 kernel 收益显著
│ └── 否（纯 GEMM 等）→ cuBLAS/CUTLASS 通常更优，Triton 不一定更快
└── 是否为多个相邻算子的组合？（如 Linear + Activation + Norm）
└── ✅ Triton 融合 kernel 可消除中间 tensor 显存分配
| 适用场景（Triton 优势大） | 不适用场景（其他方案更优） |
|---|---|
| Element-wise 融合（多个逐元素操作合并为一个 kernel） | 纯矩阵乘法（cuBLAS/CUTLASS 更优） |
| Reduction 融合（Softmax、LayerNorm、RMSNorm 等） | 已有高度优化的 CUDA 库（cuDNN、NCCL） |
| Fused Linear + Activation + Norm | 非计算密集型操作（如纯内存拷贝） |
| Fused Linear + Loss（消除大 logits tensor） | 简单操作 torch.compile 已能自动融合 |
| 自定义 Attention 变体（标准 FA 不覆盖的 pattern） | I/O 密集型瓶颈（应优化数据管道而非算子） |
| Quantized GEMM（INT8/FP8 矩阵乘） | 通信瓶颈（应优化 NCCL 而非计算） |
| 优化类型 | 典型加速比 | 典型显存节省 | 示例 |
|---|---|---|---|
| 多个 element-wise 融合 | 2-5x | 消除中间 tensor | Activation + Dropout 融合 |
| Reduction kernel 优化 | 1.5-3x | 中间结果不写回 HBM | Fused Softmax、Fused LayerNorm |
| Fused Linear + Loss | 2-4x | 50-80%（消除大 logits tensor） | Fused Linear CrossEntropy |
| Fused Norm + Activation | 2-3x | 减少 kernel launch | Fused RMSNorm + SwiGLU |
| 自定义量化 GEMM | 2-4x（vs FP16） | 模型大小减半+ | INT8/FP8 Matmul |
Triton 是一个基于 block（tile）的 GPU 编程模型，使用 Python 语法编写 GPU kernel，由编译器自动处理 shared memory、warp 调度等底层细节。
| 概念 | 说明 |
|---|---|
| Program | 一个 Triton kernel 实例，对应 CUDA 中的一个 block |
| Block（Tile） | kernel 一次处理的数据块大小（如 BLOCK_SIZE=1024 个元素） |
| program_id | 类似 CUDA blockIdx，标识当前 block 在 grid 中的位置 |
| tl.load / tl.store | 从 HBM 读写数据，支持 mask 处理边界 |
| tl.arange | 生成 block 内的偏移量序列 |
| @triton.jit | 标记函数为 Triton kernel |
| @triton.autotune | 自动搜索最优 block/warp/pipeline 配置 |
| tl.constexpr | 编译期常量，用于 block size 等参数 |
import torch
import triton
import triton.language as tl
@triton.jit
def add_kernel(
x_ptr, y_ptr, output_ptr,
n_elements,
BLOCK_SIZE: tl.constexpr,
):
# 当前 block 处理的数据范围
pid = tl.program_id(axis=0)
block_start = pid * BLOCK_SIZE
offsets = block_start + tl.arange(0, BLOCK_SIZE)
# mask 处理越界（最后一个 block 可能不满）
mask = offsets < n_elements
# 从 HBM 读取 → 计算 → 写回 HBM
x = tl.load(x_ptr + offsets, mask=mask)
y = tl.load(y_ptr + offsets, mask=mask)
output = x + y
tl.store(output_ptr + offsets, output, mask=mask)
def add(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
output = torch.empty_like(x)
n_elements = output.numel()
# grid lambda：根据 meta 参数（如 BLOCK_SIZE）动态计算 grid
grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']),)
add_kernel[grid](x, y, output, n_elements, BLOCK_SIZE=1024)
return output
关键模式：
tl.program_id(axis=0) 获取 block IDtl.arange(0, BLOCK_SIZE) 生成 block 内偏移mask = offsets < n_elements 处理边界tl.load(..., mask=mask) 带 mask 的安全加载kernel[grid](...) 启动 kernel，grid 可以是 lambda| 要点 | 说明 | 实践建议 |
|---|---|---|
| 合并访问 (Coalescing) | 同一 warp 内的线程访问连续内存地址 | 确保最内层维度是连续的，使用行主序 |
| BLOCK_SIZE 对齐 | 块大小应为 2 的幂且 >= 64 | 常用值：64, 128, 256, 512, 1024 |
| 减少 HBM 读写 | 算子融合的核心价值在于减少中间结果写回 HBM | 多个操作合并在同一个 kernel 内完成 |
| 向量化加载 | Triton 编译器自动向量化 tl.load | 保持数据对齐有助于生成更优的向量化代码 |
| Eviction Policy | 控制 L2 cache 驱逐策略 | 对不会复用的数据用 eviction_policy='evict_first' |
以下是训练场景中最常见的 Triton 融合算子模板，按优化收益排序。
适用场景：Attention 中 softmax 占比高，且无法使用 Flash Attention（如自定义 attention 变体）。
核心优化：将 max reduction、exp、sum reduction、normalize 四步融合为一个 kernel，中间结果仅保留在 SRAM（寄存器/shared memory）中，避免多次 HBM 读写。
@triton.jit
def fused_softmax_kernel(
output_ptr, input_ptr,
input_row_stride, output_row_stride,
n_cols,
BLOCK_SIZE: tl.constexpr,
):
row_idx = tl.program_id(0)
row_start_ptr = input_ptr + row_idx * input_row_stride
col_offsets = tl.arange(0, BLOCK_SIZE)
input_ptrs = row_start_ptr + col_offsets
mask = col_offsets < n_cols
# 1. Load entire row to SRAM
row = tl.load(input_ptrs, mask=mask, other=-float('inf'))
# 2. Fused: max → subtract → exp → sum → divide
row_max = tl.max(row, axis=0)
numerator = tl.exp(row - row_max) # 数值稳定
denominator = tl.sum(numerator, axis=0)
softmax_output = numerator / denominator
# 3. Write back
output_row_start_ptr = output_ptr + row_idx * output_row_stride
tl.store(output_row_start_ptr + col_offsets, softmax_output, mask=mask)
性能要点：
n_cols 能放入一个 block 的场景（如 head_dim <= 8192）n_cols 需要分块 + 在线 softmax 算法（参考 Flash Attention）适用场景：LLaMA、Qwen、Mistral 等主流 LLM 使用 RMSNorm 替代 LayerNorm。
@triton.jit
def fused_rms_norm_kernel(
output_ptr, input_ptr, weight_ptr,
stride, n_cols, eps,
BLOCK_SIZE: tl.constexpr,
):
row_idx = tl.program_id(0)
row_start = input_ptr + row_idx * stride
col_offsets = tl.arange(0, BLOCK_SIZE)
mask = col_offsets < n_cols
# Load row and weight
row = tl.load(row_start + col_offsets, mask=mask, other=0.0).to(tl.float32)
weight = tl.load(weight_ptr + col_offsets, mask=mask, other=0.0).to(tl.float32)
# RMSNorm: x * weight / sqrt(mean(x^2) + eps)
row_sq = row * row
mean_sq = tl.sum(row_sq, axis=0) / n_cols
rrms = 1.0 / tl.sqrt(mean_sq + eps)
normed = row * rrms * weight
# Write back (can cast to fp16/bf16 here)
output_row_start = output_ptr + row_idx * stride
tl.store(output_row_start + col_offsets, normed, mask=mask)
融合扩展：RMSNorm + Residual Add（前一层的 residual 与 norm 融合，减少一次 HBM 读写）
# Fused: output = RMSNorm(x + residual) * weight
# 同时写回 normed output 和 x + residual（给后续 residual 用）
row = tl.load(input_ptr + offsets, mask=mask).to(tl.float32)
residual = tl.load(residual_ptr + offsets, mask=mask).to(tl.float32)
x = row + residual
# ... RMSNorm 计算 ...
tl.store(output_ptr + offsets, normed, mask=mask)
tl.store(residual_ptr + offsets, x, mask=mask) # in-place 更新 residual
适用场景：LLaMA/Qwen 等使用 SwiGLU (SiLU Gate)、Gemma 使用 GeGLU (GELU Gate)。
@triton.jit
def fused_swiglu_kernel(
output_ptr, gate_ptr, up_ptr,
n_elements,
BLOCK_SIZE: tl.constexpr,
):
pid = tl.program_id(0)
offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
mask = offsets < n_elements
gate = tl.load(gate_ptr + offsets, mask=mask).to(tl.float32)
up = tl.load(up_ptr + offsets, mask=mask).to(tl.float32)
# SwiGLU: silu(gate) * up = gate * sigmoid(gate) * up
sigmoid_gate = tl.sigmoid(gate)
result = gate * sigmoid_gate * up
tl.store(output_ptr + offsets, result, mask=mask)
进阶融合：SwiGLU + RMSNorm（在 MLP block 内将 norm 和 activation 融合）
适用场景：几乎所有主流 LLM 都使用 RoPE 位置编码。
@triton.jit
def fused_rope_kernel(
q_ptr, k_ptr, # Q, K tensors
cos_ptr, sin_ptr, # 预计算的 cos/sin 表
output_q_ptr, output_k_ptr,
seq_len, head_dim,
stride_seq, stride_head,
BLOCK_SIZE: tl.constexpr,
):
pid_seq = tl.program_id(0) # 序列位置
pid_head = tl.program_id(1) # head 索引
half_dim = head_dim // 2
offsets = tl.arange(0, BLOCK_SIZE)
mask = offsets < half_dim
# 加载 Q 的前半和后半
base = pid_seq * stride_seq + pid_head * stride_head
q_first = tl.load(q_ptr + base + offsets, mask=mask).to(tl.float32)
q_second = tl.load(q_ptr + base + half_dim + offsets, mask=mask).to(tl.float32)
# 加载 cos/sin
cos = tl.load(cos_ptr + pid_seq * half_dim + offsets, mask=mask).to(tl.float32)
sin = tl.load(sin_ptr + pid_seq * half_dim + offsets, mask=mask).to(tl.float32)
# RoPE: [x1, x2] → [x1*cos - x2*sin, x1*sin + x2*cos]
out_first = q_first * cos - q_second * sin
out_second = q_first * sin + q_second * cos
tl.store(output_q_ptr + base + offsets, out_first, mask=mask)
tl.store(output_q_ptr + base + half_dim + offsets, out_second, mask=mask)
# 对 K 做同样的操作（可融合在同一 kernel 中）
# ...
优化要点：Q 和 K 的 RoPE 融合在同一 kernel 中，减少一半的 kernel launch 开销。
适用场景：LLM 训练中 CrossEntropy 需要处理 (batch*seq, vocab_size) 的大 logits tensor，vocab_size 通常 32K-128K+，显存占用巨大。
核心优化：Online Softmax + Log + NLL 融合，避免实例化完整的 softmax 中间结果。
@triton.jit
def fused_cross_entropy_kernel(
loss_ptr, logits_ptr, labels_ptr,
n_cols, # vocab_size
logits_row_stride,
ignore_index,
BLOCK_SIZE: tl.constexpr,
):
row_idx = tl.program_id(0)
label = tl.load(labels_ptr + row_idx)
# Skip ignored labels
if label == ignore_index:
tl.store(loss_ptr + row_idx, 0.0)
return
logits_start = logits_ptr + row_idx * logits_row_stride
# Online softmax: 分块计算 max 和 sum(exp)
m = -float('inf') # running max
d = 0.0 # running sum(exp)
target_logit = 0.0
for block_start in range(0, n_cols, BLOCK_SIZE):
col_offsets = block_start + tl.arange(0, BLOCK_SIZE)
mask = col_offsets < n_cols
logits = tl.load(logits_start + col_offsets, mask=mask, other=-float('inf')).to(tl.float32)
# 提取 target logit
is_target = col_offsets == label
target_logit += tl.sum(tl.where(is_target, logits, 0.0))
# Online max and sum update
block_max = tl.max(logits, axis=0)
new_m = tl.maximum(m, block_max)
d = d * tl.exp(m - new_m) + tl.sum(tl.exp(logits - new_m), axis=0)
m = new_m
# Loss = log(sum(exp)) + max - target_logit
loss = tl.log(d) + m - target_logit
tl.store(loss_ptr + row_idx, loss)
性能收益：
(batch*seq, vocab_size) 的 softmax 中间 tensor → 显存大幅降低LigerCrossEntropyLoss 就是基于此原理适用场景：LLM 的最后一层 lm_head (linear) + cross_entropy_loss。这是单一最高价值的 Triton 融合优化。
核心原理：将 logits = linear(hidden, weight) 和 loss = cross_entropy(logits, labels) 融合，按 vocab_size 分块计算，每个 block 计算一部分 logits 并立即用于 online softmax，永远不实例化完整的 logits tensor。
# 伪代码框架（完整实现参考 Liger Kernel 的 LigerFusedLinearCrossEntropyLoss）
@triton.jit
def fused_linear_cross_entropy_kernel(
loss_ptr, hidden_ptr, weight_ptr, labels_ptr,
hidden_dim, vocab_size,
BLOCK_V: tl.constexpr, # vocab 分块大小
BLOCK_H: tl.constexpr, # hidden 分块大小
):
row_idx = tl.program_id(0)
label = tl.load(labels_ptr + row_idx)
m = -float('inf')
d = 0.0
target_logit = 0.0
# 分块遍历 vocab 维度
for v_start in range(0, vocab_size, BLOCK_V):
# 计算 hidden[row] @ weight[v_start:v_start+BLOCK_V].T
# 得到部分 logits（仅 BLOCK_V 个值，不需要完整 vocab_size）
partial_logits = tl.zeros([BLOCK_V], dtype=tl.float32)
for h_start in range(0, hidden_dim, BLOCK_H):
# 分块矩阵乘（在 SRAM 中）
h = tl.load(hidden_ptr + row_idx * hidden_dim + h_start + tl.arange(0, BLOCK_H))
w = tl.load(weight_ptr + (v_start + tl.arange(0, BLOCK_V)[:, None]) * hidden_dim
+ h_start + tl.arange(0, BLOCK_H)[None, :])
partial_logits += tl.sum(w * h[None, :], axis=1)
# Online softmax update（与 3.5 相同）
# ...
loss = tl.log(d) + m - target_logit
tl.store(loss_ptr + row_idx, loss)
显存收益分析：
| 模型 | vocab_size | seq_len | batch | 标准方式 logits 显存 | Fused 方式 logits 显存 |
|---|---|---|---|---|---|
| LLaMA-3-8B | 128,256 | 2048 | 4 | ~2 GB (BF16) | 0 GB（永不实例化） |
| Qwen-2-7B | 151,936 | 4096 | 2 | ~2.3 GB | 0 GB |
| LLaMA-3-70B | 128,256 | 8192 | 1 | ~2 GB | 0 GB |
推荐：在大多数场景下，直接使用 Liger Kernel 的
LigerFusedLinearCrossEntropyLoss而非手写。仅当 Liger 不支持特定 loss 变体时才需自行实现。
Triton 的 @triton.autotune 装饰器可自动搜索最优的 kernel 配置参数。
@triton.autotune(
configs=[
triton.Config({'BLOCK_M': 128, 'BLOCK_N': 256, 'BLOCK_K': 64}, num_stages=3, num_warps=8),
triton.Config({'BLOCK_M': 64, 'BLOCK_N': 256, 'BLOCK_K': 32}, num_stages=4, num_warps=4),
triton.Config({'BLOCK_M': 128, 'BLOCK_N': 128, 'BLOCK_K': 32}, num_stages=4, num_warps=4),
triton.Config({'BLOCK_M': 128, 'BLOCK_N': 64, 'BLOCK_K': 32}, num_stages=4, num_warps=4),
triton.Config({'BLOCK_M': 64, 'BLOCK_N': 128, 'BLOCK_K': 32}, num_stages=4, num_warps=4),
triton.Config({'BLOCK_M': 64, 'BLOCK_N': 64, 'BLOCK_K': 32}, num_stages=5, num_warps=2),
],
key=['M', 'N', 'K'], # 触发重新搜索的参数（shape 变化时重新 tune）
)
@triton.jit
def matmul_kernel(
a_ptr, b_ptr, c_ptr,
M, N, K,
BLOCK_M: tl.constexpr,
BLOCK_N: tl.constexpr,
BLOCK_K: tl.constexpr,
):
# ... kernel body
pass
| 参数 | 说明 | 调优建议 |
|---|---|---|
BLOCK_* | Tile 大小 | 必须是 2 的幂；增大 → 更多数据复用但寄存器压力增大 |
num_warps | 每个 block 使用的 warp 数 | 2/4/8；计算密集型用 8，内存密集型用 4 |
num_stages | 软件流水线阶段数 | 2-5；增大可隐藏内存延迟但增加寄存器使用 |
key | 触发重新 tune 的参数 | 通常设为影响 kernel 行为的 shape 参数 |
Element-wise / Reduction kernel：
configs=[
triton.Config({'BLOCK_SIZE': 1024}, num_warps=4),
triton.Config({'BLOCK_SIZE': 2048}, num_warps=4),
triton.Config({'BLOCK_SIZE': 4096}, num_warps=8),
triton.Config({'BLOCK_SIZE': 8192}, num_warps=8),
]
矩阵乘（GEMM）kernel：
configs=[
triton.Config({'BLOCK_M': 128, 'BLOCK_N': 256, 'BLOCK_K': 64, 'GROUP_M': 8}, num_stages=3, num_warps=8),
triton.Config({'BLOCK_M': 64, 'BLOCK_N': 256, 'BLOCK_K': 32, 'GROUP_M': 8}, num_stages=4, num_warps=4),
triton.Config({'BLOCK_M': 128, 'BLOCK_N': 128, 'BLOCK_K': 32, 'GROUP_M': 8}, num_stages=4, num_warps=4),
triton.Config({'BLOCK_M': 256, 'BLOCK_N': 64, 'BLOCK_K': 32, 'GROUP_M': 8}, num_stages=4, num_warps=4),
]
矩阵乘中，Tile 的遍历顺序影响 L2 cache 命中率。Grouped ordering 比简单的行优先遍历显著提升缓存命中：
# 在 GEMM kernel 中，使用 grouped ordering 提升 L2 cache 利用率
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
原理：将输出矩阵 C 的 tile 按 GROUP_M 行为一组遍历，同组内的 tile 共享 A 矩阵的相同行，提高 L2 cache 中 A 数据块的复用率。
import torch
import triton
import triton.language as tl
class FusedRMSNorm(torch.autograd.Function):
@staticmethod
def forward(ctx, x, weight, eps=1e-6):
# 保存反向传播需要的数据
output = torch.empty_like(x)
n_cols = x.shape[-1]
n_rows = x.numel() // n_cols
# 选择 BLOCK_SIZE（需 >= n_cols 且为 2 的幂）
BLOCK_SIZE = triton.next_power_of_2(n_cols)
fused_rms_norm_kernel[(n_rows,)](
output, x, weight,
x.stride(-2), n_cols, eps,
BLOCK_SIZE=BLOCK_SIZE,
)
ctx.save_for_backward(x, weight)
ctx.eps = eps
ctx.BLOCK_SIZE = BLOCK_SIZE
return output
@staticmethod
def backward(ctx, grad_output):
x, weight = ctx.saved_tensors
# 反向传播 kernel（同样用 Triton 实现）
grad_x, grad_weight = fused_rms_norm_backward(
grad_output, x, weight, ctx.eps, ctx.BLOCK_SIZE
)
return grad_x, grad_weight, None
# 用户接口
def fused_rms_norm(x, weight, eps=1e-6):
return FusedRMSNorm.apply(x, weight, eps)
import torch.nn as nn
class TritonRMSNorm(nn.Module):
"""Drop-in 替换 PyTorch RMSNorm"""
def __init__(self, hidden_size, eps=1e-6):
super().__init__()
self.weight = nn.Parameter(torch.ones(hidden_size))
self.eps = eps
def forward(self, x):
return fused_rms_norm(x, self.weight, self.eps)
# 在模型中替换
def patch_model_with_triton_rmsnorm(model):
"""将模型中所有 RMSNorm 替换为 Triton 版本"""
for name, module in model.named_modules():
if isinstance(module, (nn.RMSNorm,)) or type(module).__name__ == 'RMSNorm':
parent_name = '.'.join(name.split('.')[:-1])
child_name = name.split('.')[-1]
parent = model.get_submodule(parent_name) if parent_name else model
triton_norm = TritonRMSNorm(module.weight.shape[0], getattr(module, 'eps', 1e-6))
triton_norm.weight = module.weight # 共享权重
setattr(parent, child_name, triton_norm)
return model
Triton kernel 天然与 torch.compile 兼容（Triton 本身就是 torch.compile/Inductor 的后端）：
model = MyModel()
model = patch_model_with_triton_rmsnorm(model)
model = torch.compile(model) # Triton 自定义 kernel 不会被 compile 重写
注意：
torch.compile不会重新编译已有的@triton.jitkernel，两者互补。compile 优化模型其他部分的图级融合，Triton 优化你手写的热点 kernel。
在手写 Triton kernel 前，优先考虑成熟的 Triton 生态工具：
| 工具 | 覆盖范围 | 接入成本 | 推荐场景 |
|---|---|---|---|
| Liger Kernel | RMSNorm, RoPE, SwiGLU, GeGLU, CrossEntropy, FusedLinearCrossEntropy, FusedLinear+DPO/ORPO/SimPO/CPO/KTO Loss, KLDiv, JSD | 一行代码 apply_liger_kernel_to_llama() | 首选。HuggingFace 模型训练，吞吐 +20%，显存 -60% |
| Unsloth | 全套 LLM 训练 Triton kernel（Attention, MLP, Norm, Embedding, Loss） | tuner_backend: unsloth（框架集成） | LoRA 微调场景，速度 +70%，显存 -50% |
| FlagGems | 300+ 通用 Triton 算子（替换 PyTorch ATen 算子） | import flag_gems; flag_gems.enable() | 通用算子加速，支持 AMD ROCm |
| Triton-Puzzles | 学习资源 | — | Triton 编程学习 |
# 方式 1：自动 Monkey-Patch（推荐）
from liger_kernel.transformers import AutoLigerKernelForCausalLM
model = AutoLigerKernelForCausalLM.from_pretrained("meta-llama/Llama-3-8B", ...)
# 方式 2：按模型 Patch
from liger_kernel.transformers import apply_liger_kernel_to_llama
apply_liger_kernel_to_llama() # 必须在 model = AutoModelForCausalLM.from_pretrained() 之前调用
# 方式 3：HuggingFace Trainer 一行开启
training_args = TrainingArguments(use_liger_kernel=True, ...)
# 方式 4：单独使用 Liger 的 Triton 算子
from liger_kernel.ops.rms_norm import LigerRMSNormFunction
from liger_kernel.ops.rope import LigerRopeFunction
from liger_kernel.ops.swiglu import LigerSiLUMulFunction
from liger_kernel.ops.cross_entropy import LigerCrossEntropyFunction
from liger_kernel.ops.fused_linear_cross_entropy import LigerFusedLinearCrossEntropyFunction
Liger Kernel 支持的模型：LLaMA, Qwen2, Mistral, Gemma, Phi3, GLM, InternVL, Mixtral, Mllama, DeepSeek 等 20+ 主流模型。
import flag_gems
flag_gems.enable() # 自动替换 PyTorch 算子为 Triton 实现
# 之后正常写 PyTorch 代码，底层自动使用 Triton kernel
output = torch.softmax(x, dim=-1) # 实际调用 FlagGems 的 Triton softmax
output = torch.layer_norm(x, normalized_shape) # Triton layer_norm
需要优化的算子是什么？
├── RMSNorm/RoPE/SwiGLU/CrossEntropy/FusedLinear+Loss
│ └── 使用 Liger Kernel（一行代码，生产验证）
├── LoRA 微调整体加速
│ └── 使用 Unsloth（LlamaFactory/ms-swift 集成）
├── 通用 PyTorch 算子加速（softmax/layernorm/gelu/dropout 等）
│ └── 使用 FlagGems
├── 以上工具不覆盖的自定义算子
│ └── 手写 Triton kernel（参考本 skill 第三步模板）
└── 以上工具不覆盖 + 性能要求极致 + 有 CUDA 经验
└── 考虑 CUDA C++ / CUTLASS（超出本 skill 范围）
import triton
@triton.testing.perf_report(
triton.testing.Benchmark(
x_names=['N'],
x_vals=[2**i for i in range(10, 18)], # 1K to 128K
line_arg='provider',
line_vals=['triton', 'torch'],
line_names=['Triton', 'PyTorch'],
styles=[('blue', '-'), ('green', '-')],
ylabel='GB/s',
plot_name='rmsnorm-performance',
args={'M': 4096},
)
)
def benchmark(M, N, provider):
x = torch.randn(M, N, device='cuda', dtype=torch.float16)
weight = torch.randn(N, device='cuda', dtype=torch.float16)
if provider == 'triton':
ms = triton.testing.do_bench(lambda: fused_rms_norm(x, weight))
else:
rms_norm = torch.nn.RMSNorm(N).cuda().half()
ms = triton.testing.do_bench(lambda: rms_norm(x))
gbps = lambda ms: 2 * x.numel() * x.element_size() * 1e-9 / (ms * 1e-3)
return gbps(ms)
benchmark.run(print_data=True, save_path='./bench/')
def verify_triton_kernel(triton_fn, torch_fn, *args, atol=1e-2, rtol=1e-2):
"""验证 Triton kernel 与 PyTorch 参考实现的数值一致性"""
# 使用相同输入
triton_out = triton_fn(*args)
torch_out = torch_fn(*args)
# 数值对比
if torch.allclose(triton_out, torch_out, atol=atol, rtol=rtol):
print(f"PASS: max diff = {(triton_out - torch_out).abs().max().item():.6f}")
else:
diff = (triton_out - torch_out).abs()
print(f"FAIL: max diff = {diff.max().item():.6f}, "
f"mean diff = {diff.mean().item():.6f}")
# 定位差异位置
idx = diff.argmax()
print(f" Worst element: triton={triton_out.flatten()[idx]:.6f}, "
f"torch={torch_out.flatten()[idx]:.6f}")
# 梯度对比（训练场景必须验证）
# ... 类似对比反向传播梯度
验证 Checklist：
| # | 检查项 | 方法 |
|---|---|---|
| 1 | 前向数值精度 | torch.allclose(triton_out, torch_out, atol=1e-2, rtol=1e-2) |
| 2 | 反向梯度精度 | 对比 grad_input, grad_weight |
| 3 | 边界条件 | 测试 n_elements 不被 BLOCK_SIZE 整除的情况 |
| 4 | 数据类型 | 分别测试 FP16、BF16、FP32 |
| 5 | 大尺寸输入 | 测试实际训练中的 shape（如 seq_len=8192, hidden=4096） |
| 6 | 训练稳定性 | 在完整训练中运行 100+ step 确认 loss 曲线正常 |
| 问题 | 可能原因 | 排查方法 | 解决方案 |
|---|---|---|---|
| Triton kernel 比 PyTorch 慢 | BLOCK_SIZE 不合理 | triton.autotune 搜索 | 调整 config 搜索空间 |
| 显存未减少 | 中间 tensor 仍然被分配 | torch.cuda.memory_snapshot() | 检查 wrapper 是否保存了不必要的 tensor |
| 编译时间过长 | 首次 JIT 编译 | 正常现象（后续有缓存） | TRITON_CACHE_DIR 持久化缓存 |
| 精度下降 | FP16 累加精度不足 | 对比 FP32 参考实现 | reduction 使用 .to(tl.float32) 累加 |
| 无法 autotune | config 导致 OOM | 减小 BLOCK_SIZE 搜索上限 | 移除过大的 config |
| kernel launch overhead | 输入太小不值得 kernel | 添加 fallback 逻辑 | 小输入走 PyTorch，大输入走 Triton |
当 /gpu-training-optimizer 在 Step 6（代码级优化）中决定使用 Triton 自定义算子时，按以下流程执行：
1. PyTorch Profiler 识别 Top-K 热点算子（如 aten::layer_norm 占 15%）
2. 检查是否有成熟替代（Liger Kernel / FlagGems / torch.compile）
3. 若无替代 → 确认算子类型：
├── Element-wise → 参考模板 3.3 (SwiGLU) 或 3.4 (RoPE)
├── Reduction → 参考模板 3.1 (Softmax) 或 3.2 (RMSNorm)
├── Linear + Something → 参考模板 3.6 (Fused Linear + Loss)
└── 其他 → 基于 2.2 基础模板开发
4. 编写 Triton kernel + torch.autograd.Function wrapper
5. 数值验证（第七步 7.2）
6. 性能 benchmark（第七步 7.1）
7. 在 fork 项目中集成（Monkey-Patch 或 Module 替换）
8. 完整训练验证（运行 100+ step，对比 loss 曲线）
# 文件: ${WORK_DIR}/project/triton_kernels/__init__.py
# 新建 triton_kernels 目录存放自定义 kernel
# 文件: ${WORK_DIR}/project/triton_kernels/fused_rmsnorm.py
# 存放 Triton kernel 实现
# 文件: ${WORK_DIR}/project/patch_model.py
# 集成脚本，在训练入口处调用
from triton_kernels.fused_rmsnorm import TritonRMSNorm, patch_model_with_triton_rmsnorm
# 在训练脚本的模型初始化之后、训练之前插入：
model = patch_model_with_triton_rmsnorm(model)
每个 Triton 优化都应在 v{N}_optimization-changes.md 中记录：
### X. 自定义 Triton Fused RMSNorm
**对应瓶颈**: PyTorch Profiler 显示 `aten::_native_rms_norm` 占总耗时 12%
**修改文件**:
- 新增: `triton_kernels/fused_rmsnorm.py` — Triton kernel 实现
- 修改: `train.py:L85` — 添加 `patch_model_with_triton_rmsnorm(model)` 调用
**验证结果**:
- 数值精度: max diff = 1.2e-4 (FP16), PASS
- 性能: RMSNorm 算子耗时 3.2ms → 1.1ms (-66%)
- 显存: 峰值显存 42.1GB → 41.3GB (-0.8GB)
**收益**: 该算子占总耗时从 12% 降至 4.5%
| API | 用途 | 常用场景 |
|---|---|---|
tl.load(ptr, mask, other) | 从 HBM 加载数据 | 所有 kernel |
tl.store(ptr, value, mask) | 写回 HBM | 所有 kernel |
tl.arange(start, end) | 生成偏移序列 | block 内索引 |
tl.program_id(axis) | 获取 block ID | grid 定位 |
tl.max(x, axis) | 规约求最大值 | Softmax、Norm |
tl.sum(x, axis) | 规约求和 | Norm、Loss |
tl.exp(x) | 指数函数 | Softmax |
tl.log(x) | 对数函数 | CrossEntropy |
tl.sigmoid(x) | Sigmoid | SiLU/SwiGLU |
tl.where(cond, x, y) | 条件选择 | Mask 操作 |
tl.dot(a, b) | Tile 矩阵乘 | GEMM kernel |
tl.atomic_add(ptr, val) | 原子加 | 梯度累加 |
tl.cdiv(a, b) | 向上取整除 | Grid 计算 |
tl.constexpr | 编译期常量 | Block size |
triton.next_power_of_2(n) | 下一个 2 的幂 | BLOCK_SIZE 选择 |
GPU 性能优化中，访存模式的优化常常比计算优化收益更大。许多 kernel 慢的根因不是"算不动"，而是"访存烂"。
GPU 中同一 warp 内的线程应该访问连续的内存地址，这样硬件可以把多次访问合并为一次事务。
@triton.jit
def good_access_kernel(x_ptr, out_ptr, N: tl.constexpr):
pid = tl.program_id(0)
# ✅ 好：连续线程访问连续地址
offsets = pid * N + tl.arange(0, N)
x = tl.load(x_ptr + offsets)
tl.store(out_ptr + offsets, x)
@triton.jit
def bad_access_kernel(x_ptr, out_ptr, N: tl.constexpr, stride: tl.constexpr):
pid = tl.program_id(0)
# ❌ 差：跨步访问，每个线程隔 stride 个元素，无法合并
offsets = pid * N + tl.arange(0, N) * stride
x = tl.load(x_ptr + offsets)
tl.store(out_ptr + offsets, x)
Triton 中的实践：
tl.arange 天然产生连续偏移，直接用即可保证 coalescedtl.arange 遍历GPU 硬件支持单条指令加载 128 bit 数据（如 float4 = 4 个 float32、half8 = 8 个 float16）。Triton 编译器在 BLOCK_SIZE 为 2 的幂且访存连续时通常能自动向量化，但需要确保：
Triton 对 shared memory 的使用是隐式的（编译器自动决定哪些数据放 shared memory），但理解 bank conflict 有助于选择更好的 BLOCK_SIZE：
num_stages 参数控制 Triton 编译器的软件流水线深度——在计算当前 tile 的同时，提前异步加载下一个 tile：
triton.Config({'BLOCK_M': 128, 'BLOCK_N': 128, 'BLOCK_K': 32}, num_stages=4, num_warps=4)
# ^^^^^^^^^^^
# num_stages=4 表示最多同时有 4 个阶段的数据在 pipeline 中
# = 当前计算 tile_k + 预取 tile_{k+1} + 预取 tile_{k+2} + 预取 tile_{k+3}
| num_stages | 效果 | 适用场景 |
|---|---|---|
| 1 | 无流水线（load 完再 compute） | 极简 kernel、调试 |
| 2 | 双缓冲（经典 double buffering） | 大多数 memory-bound kernel |
| 3-5 | 深度流水线（隐藏更多延迟） | A100+，register 充裕时 |
| > 5 | 寄存器压力过大，通常适得其反 | 极少使用 |
与 L2 Cache 的关系：num_stages 增大意味着同时在 shared memory 中缓存更多 tile，这会增加 register 和 shared memory 使用量。需要在延迟隐藏和occupancy之间找平衡。
Occupancy（占用率）指 SM 上实际驻留的 warp 数与理论最大 warp 数之比。它不是越高越好，但太低通常意味着性能问题。
| 资源 | 对 Occupancy 的影响 | Triton 中的控制方式 |
|---|---|---|
| Register 使用量 | 每线程 register 越多 → 每 SM 可驻留的 block 越少 → occupancy 下降 | 减小 BLOCK_SIZE、减少临时变量、降低 num_stages |
| Shared Memory 使用量 | 每 block shared memory 越大 → 同 SM 可驻留的 block 越少 | 减小 BLOCK 维度（尤其 GEMM 的 BLOCK_M/N/K） |
| Block Size (num_warps) | 每 block 线程越多 → 单 block 消耗更多资源 | 调整 num_warps（2/4/8） |
当编译器分配的寄存器超过硬件上限，超出部分会 spill 到 local memory（实际走 global memory + cache），性能大幅下降：
spill loads 和 spill stores 指标，或 Triton 编译时的 register 使用报告| kernel 类型 | 推荐 num_warps | 推荐 BLOCK_SIZE | occupancy 目标 |
|---|---|---|---|
| Element-wise | 4 | 1024-4096 | > 50% |
| Reduction (Norm/Softmax) | 4-8 | 与 hidden_dim 对齐 | > 30%（可接受低些，因为是计算密集） |
| GEMM | 4-8 | BLOCK_M/N: 64-256, BLOCK_K: 32-64 | > 40% |
| Fused Linear+Loss | 4-8 | BLOCK_M: 64-128, BLOCK_V: vocab 分块 | > 30% |
Memory-bound kernel 对 occupancy 更敏感：更多活跃 warp 可以在等待内存返回时切换到其他 warp 执行，从而隐藏内存延迟。Compute-bound kernel 对 occupancy 不太敏感，因为计算本身就能让 pipeline 保持忙碌。
GPU 的真实执行单位是 warp（32 个线程）。虽然 Triton 抽象了大部分底层细节，但理解 warp 行为有助于写出更高效的 kernel。
同一 warp 内的线程必须执行相同指令。如果分支导致 warp 内线程走不同路径，这些路径会被串行执行：
@triton.jit
def kernel(x_ptr, out_ptr, N: tl.constexpr):
pid = tl.program_id(0)
offsets = pid * N + tl.arange(0, N)
x = tl.load(x_ptr + offsets, mask=offsets < total_size, other=0.0)
# ⚠️ tl.where 比 Python if/else 更友好于 warp
# tl.where 本质上是 predication（两个分支都算，根据条件选结果），无 divergence
result = tl.where(x > 0, x, x * 0.01) # LeakyReLU
tl.store(out_ptr + offsets, result)
最佳实践：
tl.where 替代分支——predication 模式下无 divergencetl.load(..., mask=...) ）本身不产生 divergence（Triton 编译器做了特殊处理）Triton 通过 tl.reduce 和 tl.scan 暴露了底层的 warp shuffle 操作：
# tl.sum 内部使用 warp shuffle 实现高效的 block-level reduction
# 比显式的 shared memory reduction 更快（避免 shared memory 读写）
row_sum = tl.sum(x, axis=-1) # 编译为 warp shuffle + cross-warp reduce
tl.reduce（sum/max/min）→ 编译为 warp shuffle + cross-warp synctl.scan（cumsum/cumprod）→ 编译为 warp shuffle 的 prefix scanTriton 编译器会自动对 for 循环做展开优化以提高指令级并行（ILP）。但展开程度可以通过代码结构影响：
@triton.jit
def matmul_kernel(..., BLOCK_K: tl.constexpr):
acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
# 这个循环会被编译器展开——因为 BLOCK_K 是 constexpr
for k in range(0, K, BLOCK_K):
a = tl.load(a_ptr + ...)
b = tl.load(b_ptr + ...)
acc = tl.dot(a, b, acc)
展开的收益：
展开的风险：
实践建议：
tl.constexpr 声明的循环上界会触发完全展开——对短循环（< 32 次迭代）有效| 工具 | 配合方式 |
|---|---|
| Flash Attention | FA 处理 Attention，Triton 处理 Norm/Activation/Loss/RoPE，两者互补 |
| Liger Kernel | Liger 本身基于 Triton，用 Liger 覆盖常见算子，仅 Liger 不覆盖的才手写 |
| torch.compile | compile 不重写 @triton.jit kernel，两者可同时使用 |
| DeepSpeed / FSDP | 分布式框架与 Triton 算子完全兼容（算子在单卡上执行，分布式管理权重分片） |
| 混合精度 (AMP) | Triton kernel 内部精度自行控制，不受 AMP autocast 影响；建议 reduction 路径用 FP32 |
| CUDA Graph | Triton kernel 可被 CUDA Graph 捕获，但首次编译不能在 graph 内 |
| PyTorch Profiler | 自定义 Triton kernel 会以 triton_* 名称出现在 profiler 中 |
| Nsight Systems | Triton kernel 可在 timeline 中以 CUDA kernel 形式可见 |
# Triton 随 PyTorch 2.0+ 自动安装
pip install torch # Triton 已包含在内
# 单独安装（特定版本）
pip install triton
# 验证安装
python -c "import triton; print(triton.__version__)"
| 要求 | 最低版本 | 推荐版本 |
|---|---|---|
| PyTorch | 2.0 | 2.4+ |
| Triton | 2.0 | 3.0+（随 PyTorch 自动安装） |
| CUDA | 11.6 | 12.1+ |
| GPU | Volta (V100)+ | Ampere (A100)+ |
| Python | 3.8 | 3.10+ |
注意：Triton 的 CUDA backend 支持 V100 及以上 GPU。与 Flash Attention（仅 Ampere+）不同，Triton 可在 V100 上运行，但 Ampere+ 的 tensor core 指令可能带来更好的性能。Triton 也支持 AMD ROCm 后端。
DeepSpeed ZeRO 显存优化与训练加速
Flash Attention 版本选型与最佳配置
提交代码到远程仓库
NVIDIA Nsight Systems 系统级 GPU Profiling
Nsight Systems 自动检测、安装与升级
PyTorch Profiler 性能分析，提供算子级瓶颈定位