triton-cuda-basics

Triton CUDA 编程基础

1. 核心概念

内核 (Kernel)

• 定义: 使用 @triton.jit 装饰的 Python 函数，编译后在 GPU 上并行执行
• 特点: 每个内核实例处理数据的一个子集，通过程序 ID 区分

网格 (Grid) 与块 (Block)

• 网格: 内核启动时的并行维度配置，如 (num_blocks_x, num_blocks_y)
• 块: 每个程序实例处理的数据块大小，如 BLOCK_SIZE = 1024
• 关系: grid_size = ceil(total_elements / block_size)

内存层次

• 全局内存 (Global Memory): 主内存（HBM），所有程序可访问，延迟高，带宽大
• 共享内存 (Shared Memory): SM 内共享，延迟低，容量有限（通常 48-164 KB/SM）
• 寄存器 (Registers): 每个线程私有，最快访问

CUDA GPU 架构要点

• SM (Streaming Multiprocessor): GPU 基本计算单元
• Warp: 32 个线程为一组并行执行
• Tensor Core: 专用矩阵计算单元（Ampere/Hopper 架构）

2. 标准内核结构（五步模式）

所有 Triton 内核都遵循相同的五步结构模式：

@triton.jit
def standard_kernel(
    output_ptr, input_ptr, n_elements, 
    BLOCK_SIZE: tl.constexpr,
):
    # 1. 获取程序 ID 和计算偏移
    pid = tl.program_id(0)
    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    
    # 2. 创建边界掩码
    mask = offsets < n_elements
    
    # 3. 加载数据
    data = tl.load(input_ptr + offsets, mask=mask)
    
    # 4. 执行计算
    result = compute_function(data)
    
    # 5. 存储结果
    tl.store(output_ptr + offsets, result, mask=mask)

3. 内核启动方式

函数形式

def launch_kernel(input_tensor, output_tensor):
    BLOCK_SIZE = 1024  
    grid = (triton.cdiv(input_tensor.numel(), BLOCK_SIZE),)
    
    kernel[grid](
        output_tensor, input_tensor, input_tensor.numel(),
        BLOCK_SIZE=BLOCK_SIZE,
    )

ModelNew 类格式（推荐）

class ModelNew(torch.nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, input_tensor):
        output_tensor = torch.empty_like(input_tensor)
        BLOCK_SIZE = 1024  
        grid = (triton.cdiv(input_tensor.numel(), BLOCK_SIZE),)
        
        kernel[grid](
            output_tensor, input_tensor, input_tensor.numel(),
            BLOCK_SIZE=BLOCK_SIZE,
        )
        return output_tensor

4. 边界处理

使用 mask 处理边界

# 基本边界检查
offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
mask = offsets < n_elements
data = tl.load(ptr + offsets, mask=mask, other=0.0)