层标准化 （Layer Normalization）
在本节中，我们将使用 Triton 编写一个比 PyTorch 实现运行更快的高性能层标准化 (layer normalization) 内核。
计算内核
import pytest
import torch
import triton
import triton.language as tl
import torch_npu
@triton.jit
def _layer_norm_fwd_fused(
X, # 输入指针
Y, # 输出指针
W, # 权重指针
B, # 偏差指针
Mean, # 均值指针
Rstd, # 1/std 指针
stride, # 指针移动一行应该增加多少元素
N, # X 的列数
eps, # 用于避免除以 0 的 epsilon
BLOCK_SIZE: tl.constexpr,
):
# 映射程序 id 到对应计算的 X 和 Y 的行
row = tl.program_id(0)
Y += row * stride
X += row * stride
# 计算均值
mean = 0
_mean = tl.zeros([BLOCK_SIZE], dtype=tl.float32)
for off in range(0, N, BLOCK_SIZE):
cols = off + tl.arange(0, BLOCK_SIZE)
a = tl.load(X + cols, mask=cols < N, other=0.).to(tl.float32)
_mean += a
mean = tl.sum(_mean, axis=0) / N
# 计算方差
_var = tl.zeros([BLOCK_SIZE], dtype=tl.float32)
for off in range(0, N, BLOCK_SIZE):
cols = off + tl.arange(0, BLOCK_SIZE)
x = tl.load(X + cols, mask=cols < N, other=0.).to(tl.float32)
x = tl.where(cols < N, x - mean, 0.)
_var += x * x
var = tl.sum(_var, axis=0) / N
rstd = 1 / tl.sqrt(var + eps)
# 写入 mean / rstd
tl.store(Mean + row, mean)
tl.store(Rstd + row, rstd)
# 归一化并应用线性变换
for off in range(0, N, BLOCK_SIZE):
cols = off + tl.arange(0, BLOCK_SIZE)
mask = cols < N
w = tl.load(W + cols, mask=mask)
b = tl.load(B + cols, mask=mask)
x = tl.load(X + cols, mask=mask, other=0.).to(tl.float32)
x_hat = (x - mean) * rstd
y = x_hat * w + b
# 写入输出
tl.store(Y + cols, y, mask=mask)
使用 Triton 自定义的 LayerNorm 实现方式
@torch.inference_mode()
def layer_norm(x, normalized_shape, weight, bias, eps=1e-5):
# 分配与输入相同形状和数据类型的输出张量
y = torch.empty_like(x)
# 将输入 x 展平成二维形状 [-1, feature_dim] 以便处理最后一个维度
x_arg = x.reshape(-1, x.shape[-1])
M, N = x_arg.shape
mean = torch.empty((M, ), dtype=torch.float32, device=x.device)
rstd = torch.empty((M, ), dtype=torch.float32, device=x.device)
BLOCK_SIZE = 1024
# enqueue kernel
kernel = _layer_norm_fwd_fused[(M, )]( # M 表示 block 数目，launch grid=(M,)
x_arg, y, weight, bias, mean, rstd, # 输入输出及中间量
x_arg.stride(0), N, eps,
BLOCK_SIZE=BLOCK_SIZE)
# 返回归一化后的输出结果
return y
# 前向传播时调用层归一化
def _layer_norm(M, N, dtype, eps=1e-5, device='npu'):
# 构造数据
x_shape = (M, N)
w_shape = (x_shape[-1], )
weight = torch.rand(w_shape, dtype=dtype, device=device, requires_grad=True)
bias = torch.rand(w_shape, dtype=dtype, device=device, requires_grad=True)
x = -2.3 + 0.5 * torch.randn(x_shape, dtype=dtype, device=device)
dy = .1 * torch.randn_like(x)
x.requires_grad_(True)
# 前向传播
y_tri = layer_norm(x, w_shape, weight, bias, eps)
y_ref = torch.nn.functional.layer_norm(x, w_shape, weight, bias, eps).to(dtype)
# 判断是否近似
assert torch.allclose(y_tri, y_ref, atol=1e-2, rtol=0)
print(f"y_tri: {y_tri}")
print(f"y_ref: {y_ref}")
print(f"Layer Normalization {M},{N} {dtype} PASSED!")
# 执行测试
if __name__ == '__main__':
_layer_norm(128, 128, torch.float16)
_layer_norm(128, 128, torch.bfloat16)
_layer_norm(128, 128, torch.float32)
结果
y_tri: tensor([[ 0.2512, 0.0647, 0.8389, ..., 2.3652, 1.5039, 1.1904],
[ 1.0908, 1.5391, 0.2269, ..., 1.6846, 1.0996, 0.9614],
[-0.2974, 0.5918, 0.3225, ..., 2.2891, -0.8418, 0.6885],
...,
[ 0.5225, -0.0068, 0.4968, ..., -1.1221, 1.7422, 0.6143],
[ 0.4463, 1.2441, 0.2224, ..., 2.2969, -0.3311, 0.6177],
[-0.0113, 0.8423, 0.3696, ..., 1.3838, 1.2471, 0.8750]],
device='npu:0', dtype=torch.float16)
y_ref: tensor([[ 0.2512, 0.0647, 0.8389, ..., 2.3652, 1.5039, 1.1904],
[ 1.0908, 1.5391, 0.2269, ..., 1.6846, 1.0996, 0.9614],
[-0.2974, 0.5918, 0.3225, ..., 2.2891, -0.8418, 0.6885],
...,
[ 0.5225, -0.0068, 0.4968, ..., -1.1221, 1.7422, 0.6143],
[ 0.4463, 1.2441, 0.2224, ..., 2.2969, -0.3311, 0.6177],
[-0.0113, 0.8423, 0.3696, ..., 1.3838, 1.2471, 0.8750]],
device='npu:0', dtype=torch.float16, grad_fn=<NativeLayerNormBackward0>)
Layer Normalization 128,128 torch.float16 PASSED!
y_tri: tensor([[-0.4180, 0.9648, 0.8633, ..., 0.7656, 0.8438, 0.3633],
[ 0.4453, 0.5352, 0.9102, ..., 1.1875, -0.0562, 0.5391],
[ 1.3125, 0.9961, 0.9219, ..., 0.9688, 0.0025, 0.5156],
...,
[-0.1426, 0.6289, 0.9609, ..., 0.9648, -0.1260, -0.1270],
[ 1.1641, 0.6680, 0.8281, ..., 0.9258, 0.9062, 0.1768],
[-0.2129, 0.7109, 0.9141, ..., 0.7891, -0.0767, 0.5156]],
device='npu:0', dtype=torch.bfloat16)
y_ref: tensor([[-0.4180, 0.9648, 0.8633, ..., 0.7656, 0.8438, 0.3633],
[ 0.4453, 0.5352, 0.9102, ..., 1.1875, -0.0562, 0.5391],
[ 1.3125, 0.9961, 0.9219, ..., 0.9688, 0.0025, 0.5156],
...,
[-0.1426, 0.6289, 0.9609, ..., 0.9648, -0.1260, -0.1270],
[ 1.1641, 0.6680, 0.8281, ..., 0.9258, 0.9062, 0.1768],
[-0.2129, 0.7109, 0.9141, ..., 0.7891, -0.0767, 0.5156]],
device='npu:0', dtype=torch.bfloat16, grad_fn=<NativeLayerNormBackward0>)
Layer Normalization 128,128 torch.bfloat16 PASSED!
y_tri: tensor([[-0.2980, 0.2922, 0.6481, ..., 0.9786, 0.7304, 0.8982],
[ 1.5911, 0.0474, 0.6518, ..., 0.8013, 0.2435, 1.3748],
[ 1.3024, 0.6265, 0.6473, ..., 0.8423, 0.0984, -1.1839],
...,
[-0.2195, 0.1359, 0.6461, ..., 0.8319, 1.0899, 1.5015],
[ 0.6371, 0.3687, 0.6530, ..., 0.9359, 0.0818, 0.6499],
[ 0.1178, 0.3639, 0.6475, ..., 0.7221, 0.4622, 1.4510]],
device='npu:0')
y_ref: tensor([[-0.2980, 0.2922, 0.6481, ..., 0.9786, 0.7304, 0.8982],
[ 1.5911, 0.0474, 0.6518, ..., 0.8013, 0.2435, 1.3748],
[ 1.3024, 0.6265, 0.6473, ..., 0.8423, 0.0984, -1.1839],
...,
[-0.2195, 0.1359, 0.6461, ..., 0.8319, 1.0899, 1.5015],
[ 0.6371, 0.3687, 0.6530, ..., 0.9359, 0.0818, 0.6499],
[ 0.1178, 0.3639, 0.6475, ..., 0.7221, 0.4622, 1.4510]],
device='npu:0', grad_fn=<NativeLayerNormBackward0>)
Layer Normalization 128,128 torch.float32 PASSED!
“Layer Normalization 128,128 torch.float16 PASSED!、
Layer Normalization 128,128 torch.bfloat16 PASSED!
Layer Normalization 128,128 torch.float32 PASSED!” 表明Triton和PyTorch上float16、bfloat16、float32数据类型的输出结果完全一致。