2. 向量加法性能调优

任务说明

本案例实现一个Triton向量加法内核，通过调整BLOCK_SIZE（每个线程块处理的元素数）与grid（线程块数量）来优化GPU上的内存吞吐量，并与cuBLAS的cublasSaxpy进行性能对比。最终目标是找到当前GPU上最优的块大小配置，并验证Triton内核在该配置下达到或接近cuBLAS的性能水平。

实现代码

import torch
import triton
import triton.language as tl
import triton.testing
import time

@triton.jit
def vector_add_kernel(X, Y, Z, N, BLOCK_SIZE: tl.constexpr):
    pid = tl.program_id(0)
    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = offsets < N
    x = tl.load(X + offsets, mask=mask)
    y = tl.load(Y + offsets, mask=mask)
    tl.store(Z + offsets, x + y, mask=mask)

def run_triton_vector_add(N, BLOCK_SIZE, grid):
    X = torch.randn(N, device='cuda', dtype=torch.float32)
    Y = torch.randn(N, device='cuda', dtype=torch.float32)
    Z = torch.empty_like(X)
    X_ref = X.clone()
    Y_ref = Y.clone()
    # 预热
    vector_add_kernel[(grid,)](X, Y, Z, N, BLOCK_SIZE)
    torch.cuda.synchronize()
    # 计时
    start = time.time()
    vector_add_kernel[(grid,)](X, Y, Z, N, BLOCK_SIZE)
    torch.cuda.synchronize()
    elapsed = time.time() - start
    # 验证正确性
    ref = X_ref + Y_ref
    assert torch.allclose(Z, ref, atol=1e-5), "Triton结果错误"
    return elapsed

def run_cublas_saxpy(N):
    X = torch.randn(N, device='cuda', dtype=torch.float32)
    Y = torch.randn(N, device='cuda', dtype=torch.float32)
    Z = torch.empty_like(X)
    alpha = 1.0
    # 预热
    torch.cuda.synchronize()
    start = time.time()
    torch.add(X, Y, out=Z)  # PyTorch的add在CUDA后端调用cuBLAS saxpy
    torch.cuda.synchronize()
    elapsed = time.time() - start
    return elapsed

def benchmark(N=1_000_000):
    print(f"向量长度: {N:,}")
    # cuBLAS基线
    cublas_time = run_cublas_saxpy(N)
    cublas_gbps = 3 * N * 4 / (cublas_time * 1e9)  # 读X,Y写Z共3N个float
    print(f"cuBLAS: {cublas_time*1e6:.1f} us, {cublas_gbps:.2f} GB/s")

    # 扫描不同BLOCK_SIZE
    results = []
    for BLOCK_SIZE in [128, 256, 512, 1024, 2048]:
        grid = (N + BLOCK_SIZE - 1) // BLOCK_SIZE
        t = run_triton_vector_add(N, BLOCK_SIZE, grid)
        gbps = 3 * N * 4 / (t * 1e9)
        results.append((BLOCK_SIZE, grid, t, gbps))
        print(f"Triton BLOCK={BLOCK_SIZE:5d} grid={grid:5d}: {t*1e6:.1f} us, {gbps:.2f} GB/s")

    # 找出最优
    best = min(results, key=lambda x: x[2])
    print(f"\n最优配置: BLOCK_SIZE={best[0]}, grid={best[1]}, "
          f"耗时={best[2]*1e6:.1f} us, 吞吐={best[3]:.2f} GB/s")
    print(f"相对cuBLAS性能比: {cublas_time/best[2]:.2f}x")

if __name__ == "__main__":
    benchmark()

关键实现说明

内核使用tl.arange生成连续偏移量，通过mask处理N不能被BLOCK_SIZE整除的边界情况。每个线程块处理BLOCK_SIZE个元素，grid由ceil(N/BLOCK_SIZE)决定。基准测试函数分别测量cuBLAS（通过PyTorch的torch.add调用）与不同块大小下Triton内核的执行时间，并计算有效内存带宽（读X、Y各N个float，写Z共N个float，总数据量3×N×4字节）。

运行与验证

# 安装依赖
pip install triton torch

# 运行基准测试
python vector_add_tune.py

判定标准：

1. 程序输出中Triton各配置的Z与X+Y的误差均在atol=1e-5内，无断言错误；
2. 对比不同BLOCK_SIZE的耗时，能观察到吞吐量差异，并找到最优配置；
3. 最优Triton配置的吞吐量应达到cuBLAS的80%以上（受GPU型号与驱动影响，通常为0.8x～1.2x）。

性能观测点

图中展示了数据流与性能观测点：Triton内核与cuBLAS读取相同的输入X、Y，输出Z需通过正确性校验；性能对比环节分别记录两者耗时，并换算为有效内存带宽。通过扫描不同BLOCK_SIZE，可观察块大小对内存事务效率的影响——过小的块导致调度开销占比高，过大的块可能造成寄存器溢出或占用率下降，最优值通常落在256～1024区间（取决于具体GPU架构）。