Triton Kernel 实现练习#

本页内容8 个部分

本文整理另一个实验工作区的一组 Triton kernel 练习，所列 .py 脚本未随本站发布；路径是当时环境示例。代码块包含局部片段与公式伪代码，不是可直接运行的完整实现。推荐学习顺序是：

vector_add
row_wise_sum / row_wise_max
softmax
online_softmax
matmul
advanced_matmul
fused_activation / approx_activation
layernorm
RMSNorm
fused_residual_rmsnorm
attention
flash_attention

运行脚本时，这个环境里建议用：

/workspace/isaaclab/isaaclab.sh -p <file.py>

例如：

/workspace/isaaclab/isaaclab.sh -p flash_attention.py

1. Triton 基础模式#

大多数文件都遵循同一个结构：

@triton.jit
def xxx_kernel(..., BLOCK_SIZE: tl.constexpr):
    pid = tl.program_id(axis=0)
    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = offsets < n_elements
    x = tl.load(x_ptr + offsets, mask=mask, other=0.0)
    ...
    tl.store(out_ptr + offsets, y, mask=mask)

几个核心概念：

tl.program_id(axis=...): 当前 program/block 的编号。
tl.arange(0, BLOCK_SIZE): 在一个 program 内生成向量 offset。
mask: 防止越界读写，也可以表达 causal mask 这类算法语义。
tl.load(..., mask=..., other=...): masked load，越界位置用 other 填。
tl.store(..., mask=...): masked store，只写合法位置。
tl.constexpr: 编译期常量，Triton 用它生成专门化代码。
triton.next_power_of_2: 常用于 reduction/softmax，让 block 大小适合 Triton 向量化。
triton.testing.do_bench: 计时 benchmark。

2. 逐元素运算与归约#

2.1 Vector Add#

文件：vector_add.py

算子：

out = x + y

学习重点：

最基本的一维 elementwise kernel。
一个 program 处理一段连续元素。
mask = offsets < n_elements 是为了处理尾部不足一个 block 的情况。

核心模式：

offsets = block_id * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
x = tl.load(x_ptr + offsets, mask=mask)
y = tl.load(y_ptr + offsets, mask=mask)
tl.store(out_ptr + offsets, x + y, mask=mask)

适合先理解 Triton 的指针、offset、mask、grid。

运行：

/workspace/isaaclab/isaaclab.sh -p vector_add.py

2.2 Row-wise Sum#

文件：row_wise_sum.py

算子：

out[i] = sum_j x[i, j]

学习重点：

一行一个 program。
每个 program 内做一次 vector reduction。
tl.sum(values, axis=0) 对当前 block 内的向量求和。

核心公式：

values = tl.load(x_ptr + row * n_cols + cols, mask=mask, other=0.0)
result = tl.sum(values, axis=0)

这个是 softmax、LayerNorm、RMSNorm 的基础。

运行：

/workspace/isaaclab/isaaclab.sh -p row_wise_sum.py

2.3 Row-wise Max#

文件：row_wise_max.py

算子：

out[i] = max_j x[i, j]

学习重点：

和 row-wise sum 同结构。
padding 位置要用 -inf，否则 max 会被错误影响。

核心模式：

values = tl.load(x_ptr + row * n_cols + cols, mask=mask, other=-float("inf"))
result = tl.max(values, axis=0)

这个是 stable softmax 的第一步。

运行：

/workspace/isaaclab/isaaclab.sh -p row_wise_max.py

2.4 Softmax#

文件：softmax.py

算子：

softmax(x)_j = exp(x_j) / sum_k exp(x_k)

文件里有两个版本：

softmax_kernel: 使用 tl.softmax(values)。
stable_softmax_kernel: 手写稳定 softmax。

稳定版本公式：

m = max(x)
numerator = exp(x - m)
denominator = sum(numerator)
out = numerator / denominator

学习重点：

softmax 必须处理数值稳定性。
masked load 的无效位置用 -inf。
如果 n_cols 不是 2 的幂，BLOCK_SIZE = next_power_of_2(n_cols) 会补齐，所以必须用 mask。

运行：

/workspace/isaaclab/isaaclab.sh -p softmax.py

2.5 Online Softmax#

文件：online_softmax.py

目的：

当一行非常长，不能或不想一次性把整行放进一个 block 时，用 online softmax 分块扫描。

online softmax 维护两个状态：

m: 当前已经扫描部分的最大值
d: 当前分母 sum(exp(x - m))

每来一个 block：

m_new = max(m_old, block_max)
d_new = d_old * exp(m_old - m_new) + sum(exp(block - m_new))

文件里有两个版本：

online_softmax: 单 kernel 内循环扫描整行。
split_online_softmax: 多 kernel 分块统计 partial，再 reduce，再写结果。

学习重点：

这是 FlashAttention 的前置知识。
FlashAttention 其实就是把 online softmax 和 P @ V 融在一起。

运行：

/workspace/isaaclab/isaaclab.sh -p online_softmax.py

3. 矩阵乘法与融合#

3.1 Matmul#

文件：matmul.py

算子：

C = A @ B

学习重点：

二维 tile。
tl.dot(a, b)。
grid 通常是 (num_m_tiles * num_n_tiles,) 或二维 grid。
K 维循环累加。

典型结构：

acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
for k in range(0, K, BLOCK_K):
    a = tl.load(...)
    b = tl.load(...)
    acc += tl.dot(a, b)
tl.store(c, acc)

运行：

/workspace/isaaclab/isaaclab.sh -p matmul.py

3.2 Advanced Matmul#

文件：advanced_matmul.py

在基础 matmul 上增加：

C = A @ B
C = A @ B + bias

学习重点：

bias fusion。
避免单独的 C + bias kernel。
体会算子融合的价值：少一次全局内存读写。

运行：

/workspace/isaaclab/isaaclab.sh -p advanced_matmul.py

3.3 Fused Activation#

文件：fused_activation.py

算子：

out = silu(x) * gate + residual

其中：

silu(x) = x * sigmoid(x)

学习重点：

elementwise fusion。
一个 kernel 完成多个 elementwise 操作。
对比 eager PyTorch 和 torch.compile。

为什么融合会快：

未融合: 读 x -> 写 silu -> 读 silu/gate -> 写 tmp -> 读 tmp/residual -> 写 out
融合:   读 x/gate/residual -> 写 out

运行：

/workspace/isaaclab/isaaclab.sh -p fused_activation.py

3.4 Approx Activation#

文件：approx_activation.py

包含：

exact GELU。
approximate GELU。
fused residual + bias + GELU。

GELU 常见公式：

gelu_exact(x) = 0.5 * x * (1 + erf(x / sqrt(2)))

近似 GELU：

gelu_approx(x) = 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3)))

学习重点：

数学近似和性能/误差权衡。
fused kernel 的显存收益。
benchmark 里可以观察 exact 和 approx 的速度差异。

运行：

/workspace/isaaclab/isaaclab.sh -p approx_activation.py

4. 归一化与残差#

4.1 LayerNorm#

文件：layernorm.py

算子：

mean = mean(x)
var = mean((x - mean) ** 2)
out = (x - mean) * rsqrt(var + eps) * weight + bias

学习重点：

一行一个 program。
每行做两个 reduction：sum(x) 和 sum((x - mean)^2)。
padding 位置不能污染 variance。

关键代码：

values = tl.load(..., mask=mask, other=0.0).to(tl.float32)
mean = tl.sum(values, axis=0) / n_cols
centered = tl.where(mask, values - mean, 0.0)
variance = tl.sum(centered * centered, axis=0) / n_cols

centered = tl.where(mask, values - mean, 0.0) 的作用：

如果 n_cols 不是 2 的幂，Triton 会补到 BLOCK_SIZE。无效位置不能参与 variance，否则会把 0 - mean 也平方进去。

运行：

/workspace/isaaclab/isaaclab.sh -p layernorm.py

4.2 RMSNorm#

文件：RMSNorm.py

算子：

inv_rms = rsqrt(mean(x ** 2) + eps)
out = x * inv_rms * weight

和 LayerNorm 的区别：

LayerNorm: 减 mean，再除 std
RMSNorm:   不减 mean，只除 rms

RMSNorm 少一个 mean/variance 中心化步骤，所以更简单，也通常更快。很多 LLM 使用 RMSNorm。

学习重点：

row-wise square sum reduction。
fp32 累加，然后输出 cast 回原 dtype。
对比 PyTorch eager 公式。

运行：

/workspace/isaaclab/isaaclab.sh -p RMSNorm.py

4.3 Fused Residual RMSNorm#

文件：fused_residual_rmsnorm.py

算子：

residual_out = x + residual
out = residual_out * rsqrt(mean(residual_out ** 2) + eps) * weight

接口：

output, residual_out = fused_residual_rmsnorm(x, residual, weight)

为什么这个例子重要：

Transformer block 里经常出现 residual add + norm：

x = x + residual
y = rmsnorm(x)

如果不融合，会先写一次 x + residual，再读回来做 RMSNorm。融合后一个 kernel 内完成：

读 x/residual/weight
算 x + residual
算 rms
写 output 和 residual_out

文件里还比较了：

手写 Triton fused kernel。
PyTorch eager 公式。
torch.compile(..., fullgraph=True, mode="max-autotune")。

学习重点：

算子融合不只是 elementwise，也可以融合 reduction 前后的数据流。
torch.compile 会加速 eager PyTorch，但不一定稳定超过手写 Triton。
返回 output 和 residual_out 两个 tensor 时，compiler 需要真实写出 residual。

运行：

/workspace/isaaclab/isaaclab.sh -p fused_residual_rmsnorm.py

5. Attention 与性能边界#

5.1 Naive Attention#

文件：attention.py

标准 attention：

scores = q @ k.transpose(-1, -2) * scale
probs = softmax(scores)
out = probs @ v

输入布局：

q, k, v: [batch, heads, seq_len, head_dim]

文件中分成三个 Triton kernel：

qk_matmul_kernel: 计算 QK^T。
softmax_kernel: 对 scores 的最后一维做 softmax。
pv_matmul_kernel: 计算 P @ V。

支持：

naive_attention(q, k, v, causal=False)
naive_attention(q, k, v, causal=True)

causal mask：

mask = n_offsets[None, :] <= m_offsets[:, None]

含义是第 i 个 query 只能看 j <= i 的 key，不能看未来 token。

注意这里有两类 mask：

bounds_mask: 防止越界读写
causal mask: 表达算法语义

naive attention 的缺点：

显式分配 scores: [batch * heads, seq_len, seq_len]
显式分配 probs:  [batch * heads, seq_len, seq_len]

当 seq_len 很大时，中间矩阵显存开销是 O(seq_len^2)。

运行：

/workspace/isaaclab/isaaclab.sh -p attention.py

5.2 FlashAttention-style Forward#

文件：flash_attention.py

目标：

不显式写出完整的 scores 和 probs，而是在一个 kernel 内按 K/V block 扫描。

核心状态：

m: 每个 query 当前见过的最大 score
l: 每个 query 当前 softmax 分母
acc: 每个 query 当前累积的 softmax(scores) @ V

每扫描一个 K/V block：

scores = Q_block @ K_block.T * scale
m_new = max(m_old, max(scores))
alpha = exp(m_old - m_new)
p = exp(scores - m_new)
l_new = l_old * alpha + sum(p)
acc_new = acc_old * alpha + p @ V_block

最后：

out = acc / l

这就是 online softmax 和 attention value accumulation 的融合。

支持：

flash_attention(q, k, v, causal=False)
flash_attention(q, k, v, causal=True)

和 naive attention 的关键区别：

naive: scores/probs 落地到 HBM
flash: scores/probs 只在 SRAM/register 中短暂存在

显存收益：

省掉 [batch * heads, seq_len, seq_len] 的 scores/probs 中间矩阵

当前版本定位：

forward only。
教学版。
支持 causal/non-causal。
未实现 dropout、backward、variable length、GQA/MQA。

运行：

/workspace/isaaclab/isaaclab.sh -p flash_attention.py

6. 常见优化方向#

6.1 调 block size#

常见参数：

BLOCK_M
BLOCK_N
BLOCK_K
BLOCK_D
BLOCK_SIZE
num_warps
num_stages

不同 shape 最优参数不同。比如：

小 hidden size: 更小 block，减少浪费
大 hidden size: 更大 block，提高算术强度，但寄存器压力会上升

6.2 减少 HBM 读写#

多数深度学习 kernel 都容易 memory-bound。融合常常比单独优化算术更有效。

例子：

residual + RMSNorm
bias + activation
matmul + bias
attention softmax + PV

6.3 使用 online reduction#

当 reduction 维度很长时，可以用 online 形式避免一次性加载完整行。

典型例子：

online_softmax
flash_attention

6.4 精度策略#

常见选择：

输入: fp16/bf16
累加: fp32
输出: fp16/bf16

这样通常在速度和数值稳定性之间比较稳。

6.5 对比 baseline 要分清楚#

PyTorch baseline 可能有几种：

eager PyTorch: 最容易写，但经常不是最优
torch.compile: 可能融合/优化，但不一定稳定超过手写 Triton
官方 fused kernel: 最接近生产级 baseline

所以 benchmark 里看到 Triton 赢 eager PyTorch，不代表已经赢了所有生产级实现。

7. 实现与验证边界#

尾部 masked sum 通常填 0，masked max / softmax 通常填 −∞；全 mask 行需显式约定行为。LayerNorm 的方差应按真实有效维度求均值，不能把 padding 当样本。GELU 公式中的幂是数学记号，Python 中应写 x ** 3，不是 x ^ 3。

参考：Triton 官方教程。本站本轮只检查片段语法与公式；GPU 编译、数值对照和性能结果仍需在相应环境运行。

Triton Kernel 实现练习#

1. Triton 基础模式#

2. 逐元素运算与归约#

2.1 Vector Add#

2.2 Row-wise Sum#

2.3 Row-wise Max#

2.4 Softmax#

2.5 Online Softmax#

3. 矩阵乘法与融合#

3.1 Matmul#

3.2 Advanced Matmul#

3.3 Fused Activation#

3.4 Approx Activation#

4. 归一化与残差#

4.1 LayerNorm#

4.2 RMSNorm#

4.3 Fused Residual RMSNorm#

5. Attention 与性能边界#

5.1 Naive Attention#

5.2 FlashAttention-style Forward#

6. 常见优化方向#

6.1 调 block size#

6.2 减少 HBM 读写#

6.3 使用 online reduction#

6.4 精度策略#

6.5 对比 baseline 要分清楚#

7. 实现与验证边界#

延伸阅读#