triton-cuda-api

Triton CUDA API 参考手册

本文档提供 Triton 核心 API 的详细参考，包括函数签名、参数说明和使用示例。

1. 内核装饰器

@triton.jit

@triton.jit
def kernel_function(...):
    pass

• 作用: 将 Python 函数编译为 GPU 内核
• 约束: 函数内部不能使用 return、break、continue 语句

2. 程序 ID 与网格 API

tl.program_id(axis)

pid = tl.program_id(axis)  # axis: 0, 1, or 2

• 参数: axis - 维度轴 (0, 1, 2)
• 返回: 当前程序在该轴上的 ID
• 用途: 确定当前程序块处理的数据范围

tl.num_programs(axis)

num_pids = tl.num_programs(axis)  # axis: 0, 1, or 2

• 参数: axis - 维度轴 (0, 1, 2)
• 返回: 该轴上的总程序数
• 用途: 计算网格大小和边界条件

triton.cdiv(a, b)

grid_size = triton.cdiv(total_elements, block_size)

• 参数: a, b - 被除数和除数
• 返回: 向上取整的除法结果
• 用途: host 侧使用，计算启动网格大小

3. 内存操作 API

tl.load(pointer, mask=None, other=None, boundary_check=None)

data = tl.load(ptr + offsets, mask=mask, other=0.0)

• 参数:
  • pointer: 内存指针
  • mask: 布尔掩码，True 表示有效位置
  • other: 掩码为 False 时的默认值
  • boundary_check: 边界检查维度 (0, 1) 或 None
• 返回: 加载的张量数据
• 用途: 从全局内存加载数据

tl.store(pointer, value, mask=None, boundary_check=None)

tl.store(ptr + offsets, result, mask=mask)

• 参数:
  • pointer: 内存指针
  • value: 要存储的值
  • mask: 布尔掩码，True 表示有效位置
  • boundary_check: 边界检查维度 (0, 1) 或 None
• 用途: 将数据存储到全局内存

tl.make_block_ptr(base, shape, strides, offsets, block_shape, order)

block_ptr = tl.make_block_ptr(
    base=ptr,                    # 基础指针
    shape=(M, N),                # 完整矩阵形状
    strides=(stride_m, stride_n), # 步长
    offsets=(start_m, start_n),   # 当前块偏移
    block_shape=(BLOCK_M, BLOCK_N), # 块形状
    order=(1, 0)                 # 内存布局顺序
)

• 参数:
  • base: 基础内存指针
  • shape: 完整张量的形状
  • strides: 每个维度的步长
  • offsets: 当前块的起始偏移
  • block_shape: 当前块的大小
  • order: 内存布局顺序 (1, 0) 表示行主序
• 返回: 块指针对象
• 用途: 高效访问 2D 数据块

tl.advance(ptr, offsets)

block_ptr = tl.advance(block_ptr, (BLOCK_M, 0))

• 参数:
  • ptr: 块指针
  • offsets: 各维度的偏移量
• 返回: 移动后的块指针
• 用途: 移动块指针到下一个位置