triton-cuda-reduce

Reduce 算子优化

适用于需要聚合多个值的归约操作

适用算子

基础归约: sum, mean, max, min, prod 归一化: softmax, logsoftmax, layernorm, batchnorm 统计: variance, std

通用归约策略

1. 块内归约 + 原子操作

@triton.jit
def reduction_kernel(input_ptr, output_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
    
    # 加载数据
    data = tl.load(input_ptr + offsets, mask=mask, other=0.0)
    
    # 块内归约
    block_sum = tl.sum(data, axis=0)
    
    # 原子操作写回全局内存
    tl.atomic_add(output_ptr, block_sum)

2. 数值稳定性处理

关键: 对于涉及 exp 的操作（softmax、logsoftmax），必须减去最大值防止溢出。

# 错误：直接 exp 可能溢出
exp_val = tl.exp(x)

# 正确：减去最大值
max_val = tl.max(x, axis=0)
exp_val = tl.exp(x - max_val)

特定算子优化

Softmax

标准 Softmax: output = exp(x - max(x)) / sum(exp(x - max(x)))

@triton.jit
def softmax_kernel(input_ptr, output_ptr, input_row_stride, output_row_stride,
                   n_rows, n_cols, BLOCK_SIZE: tl.constexpr):
    # 获取当前程序处理的行
    row_start = tl.program_id(0)
    row_step = tl.num_programs(0)

    for row_idx in tl.range(row_start, n_rows, row_step):
        # 计算当前行的起始指针
        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
        row = tl.load(input_ptrs, mask=mask, other=-float('inf'))

        # 数值稳定性：减去最大值
        row_minus_max = row - tl.max(row, axis=0)

        # 计算指数（CUDA 后端直接使用 tl.exp）
        numerator = tl.exp(row_minus_max)

        # 计算分母（归一化因子）
        denominator = tl.sum(numerator, axis=0)

        # 计算 softmax
        softmax_output = numerator / denominator

        # 存储结果
        output_row_start_ptr = output_ptr + row_idx * output_row_stride
        output_ptrs = output_row_start_ptr + col_offsets
        tl.store(output_ptrs, softmax_output, mask=mask)