Triton-Lang在Transformer优化加速中的实践 | 得物技术

size = 1024
x = torch.rand(size, device='cuda')
y = torch.rand(size, device='cuda')
output_torch = x + y

import triton.language as tl

@triton.jit
def add_kernel(x_ptr,  # 第一个输入向量的指针
               y_ptr,  # 第二个输入向量的指针
               output_ptr,  # 输出向量的指针
               n_elements,  # 向量长度
               BLOCK_SIZE: tl.constexpr,  # 每个线程块处理的元素数量
               ):
    # 有多个'程序'处理不同的数据, 用pid标识当前是哪个程序
    pid = tl.program_id(axis=0)  
    # 计算当前程序所需要的数据的偏置
    block_start = pid * BLOCK_SIZE
    offsets = block_start + tl.arange(0, BLOCK_SIZE)
    # 创建一个掩码以防止内存操作超出范围
    mask = offsets < n_elements
    # 从 DRAM 加载 x 和 y
    x = tl.load(x_ptr + offsets, mask=mask)
    y = tl.load(y_ptr + offsets, mask=mask)
    output = x + y
    # 将计算结果output写回 DRAM
    tl.store(output_ptr + offsets, output, mask=mask)

@triton.jit
def add_kernel(x_ptr, y_ptr, o_ptr, n_elements, num_blocks_per_CTA, BLOCK_SIZE: tl.constexpr,):
    pid = tl.program_id(axis=0)  
    program_offsets = pid * num_blocks_per_CTA * BLOCK_SIZE 
    offsets = program_offsets + tl.arange(0, BLOCK_SIZE)
    
    for i in range(num_blocks_per_CTA):
        mask = offsets < n_elements
        x = tl.load(x_ptr + offsets, mask=mask)
        y = tl.load(y_ptr + offsets, mask=mask)
        output = x + y
        tl.store(o_ptr + offsets, output, mask=mask)
        offsets += BLOCK_SIZE

def add(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    output = torch.empty_like(x)
    assert x.is_cuda and y.is_cuda and output.is_cuda
    n_elements = output.numel()
    # SPMD启动网格，表示并行运行的内核实例数。
    grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']), )
    # 调用内核函数
    add_kernel[grid](x, y, output, n_elements, BLOCK_SIZE=1024)
    # 我们返回一个指向z的句柄，但是，由于`torch.cuda.synchronize()`尚未被调用，内核此时仍在异步运行。
    return output

torch.manual_seed(0)
size = 98432
x = torch.rand(size, device='cuda')
y = torch.rand(size, device='cuda')
output_triton = add(x, y)

@triton.testing.perf_report(
    triton.testing.Benchmark(
        x_names=['size'],  # Argument names to use as an x-axis for the plot.
        x_vals=[2**i for i in range(12, 28, 1)],  # Different possible values for `x_name`.
        x_log=True,  # x axis is logarithmic.
        line_arg='provider',  # Argument name whose value corresponds to a different line in the plot.
        line_vals=['triton', 'torch'],  # Possible values for `line_arg`.
        line_names=['Triton', 'Torch'],  # Label name for the lines.
        styles=[('blue', '-'), ('green', '-')],  # Line styles.
        ylabel='GB/s',  # Label name for the y-axis.
        plot_name='vector-add-performance',  # Name for the plot. Used also as a file name for saving the plot.
        args={},  # Values for function arguments not in `x_names` and `y_name`.
    ))
def benchmark(size, provider):
    x = torch.rand(size, device='cuda', dtype=torch.float32)
    y = torch.rand(size, device='cuda', dtype=torch.float32)
    quantiles = [0.5, 0.2, 0.8]
    if provider == 'torch':
        ms, min_ms, max_ms = triton.testing.do_bench(lambda: x + y, quantiles=quantiles)
    if provider == 'triton':
        ms, min_ms, max_ms = triton.testing.do_bench(lambda: add(x, y), quantiles=quantiles)
    gbps = lambda ms: 3 * x.numel() * x.element_size() * 1e-9 / (ms * 1e-3)
    return gbps(ms), gbps(max_ms), gbps(min_ms)

benchmark.run(print_data=True, show_plots=True, save_path="./")

os.environ["TRITON_INTERPRET"]=1

# Pytorch
import torch
from typing import Tuple

@torch.jit.script
def block_matrix_multiplication(A: torch.Tensor, B:torch.Tensor, 
                                M: int, N: int, K: int, 
                                BLOCK_SIZE_M: int, BLOCK_SIZE_N: int, BLOCK_SIZE_K: int) -> torch.Tensor:
    C = torch.zeros((M, N), dtype=torch.float32)
    for m in range(0, M, BLOCK_SIZE_M):
        for n in range(0, N, BLOCK_SIZE_N):
            acc = torch.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=torch.float32)
            for k in range(0, K, BLOCK_SIZE_K):
                a = A[m : m+BLOCK_SIZE_M, k : k+BLOCK_SIZE_K]
                b = B[k : k+BLOCK_SIZE_K, n : n+BLOCK_SIZE_N]
                acc += torch.matmul(a, b)
            C[m : m+BLOCK_SIZE_M, n : n+BLOCK_SIZE_N] = acc
    return C

# 用法示例
A = torch.rand(100, 100)
B = torch.rand(100, 100)
result = block_matrix_multiplication(A, B, 100, 100, 100, 16, 16, 16)

# Triton kernel
@triton.jit
det matmul_kernel(A, B, C, M, N, K, 
    BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr):
    # 2d grid
    pid_m = tl.program_id(0)
    pid_n = tl.program_id(1)
    offsets_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    mask_m = offsets_m < M
    offsets_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
    mask_n = offsets_n < N
    
    # 2d tile
    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
    for start_k in range(0, K, BLOCK_K):
        offsets_k = start_k + tl.arange(0, BLOCK_K)
        mask_k = offsets_k < K
        
        a_ptrs = A + offsets_m[:, None]*K + offsets_k[None, :]
        mask_a = mask_m[:, None] & mask_k[None, :]
        b_ptrs = B + offsets_k[:, None]*N + offsets_n[None, :]
        mask_b = mask_k[:, None] & mask_n[None, :]
    
        a = tl.load(a_ptrs, mask=mask_a, other=0)
        b = tl.load(b_ptrs, mask=mask_b, other=0)
        acc += tl.dot(a, b)
    c_ptrs = C + offsets_m[:, None]*N + offsets_n
    mask_c = mask_m[:, None] & mask_n[None, :]
    tl.store(c_ptrs, acc, mask = mask_c) 
    
# grid = (tl.cdiv(M, BLOCK_M), tl.cdiv(N, BLOCK_N), 1)

for m in range(0, m, BLOCK_M):
    for n in range(0, n, BLOCK_N): 
        CTA(m, n) ....

for n in range(0, N, BLOCK_N):
    for m in range(0, M, BLOCK_M): 
        CTA(m, n) ....

pid = tl.program_id(axis=0)
num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
num_pid_in_group = GROUP_SIZE_M * num_pid_n
group_id = pid // num_pid_in_group
first_pid_m = group_id * GROUP_SIZE_M
group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
pid_m = first_pid_m + (pid % group_size_m)
pid_n = (pid % num_pid_in_group) // group_size_m

def rotate_half(x):
    """Rotates half the hidden dims of the input."""
    x1 = x[..., : x.shape[-1] // 2]
    x2 = x[..., x.shape[-1] // 2 :]
    return torch.cat((-x2, x1), dim=-1)

q_embed = (q*cos) + (rotate_half(q)*sin)
k_embed = (k*cos) + (rotate_half(k)*sin)

div, mod = divmod(n, GROUP_SIZE)
n_groups = div + (mod != 0)
grid = (n_rows, n_groups)

def _rope_embedding(
    Q,     Q_row_stride,
    cos, cos_row_stride,
    sin, sin_row_stride,
    seqlen,
    head_dim      : tl.constexpr,
    n_heads       : tl.constexpr,
    BLOCK_SIZE    : tl.constexpr,
):
    """
        Calculates the RoPE Embedding quickly
        RoPE is Q * cos + rotate_half(Q) * sin
        See our blog post for more info
    """
    GROUP_SIZE = 4
    row_position  = tl.program_id(0)
    group_head_position = tl.program_id(1)
    col_offsets  = tl.arange(0, BLOCK_SIZE)
    half_head_dim = head_dim // 2
    mask = col_offsets < half_head_dim

    sin1 = tl.load(sin + (row_position % seqlen)*sin_row_stride + \
                   half_head_dim*0 + col_offsets, mask = mask, other = 0)
    cos1 = tl.load(cos + (row_position % seqlen)*cos_row_stride + \
                   half_head_dim*0 + col_offsets, mask = mask, other = 0)

    head_start = group_head_position * GROUP_SIZE
    head_end = min((head_start + GROUP_SIZE), n_heads)

    for k in range(head_start, head_end):
        offs_q1 = row_position * Q_row_stride + k * head_dim + col_offsets
        offs_q2 = row_position * Q_row_stride + k * head_dim + col_offsets + half_head_dim

        Q1 = tl.load(Q + offs_q1, mask = mask, other = 0).to(sin1.dtype)
        Q2 = tl.load(Q + offs_q2, mask = mask, other = 0).to(sin1.dtype)

        tl.store(Q + offs_q1, Q1*cos1 - Q2*sin1, mask = mask)
        tl.store(Q + offs_q2, Q2*cos1 + Q1*sin1, mask = mask)