triton-lang

Triton

用途

指导 Agent 使用 OpenAI Triton 编写 GPU 内核：@triton.jit 装饰器、带掩码的面向块的 tl.load/tl.store、原子操作、通过 tl.constexpr 使用共享内存、使用 triton.testing.Benchmark 进行基准测试、与 PyTorch 集成，以及通过屏障进行调试。

适用场景

• 编写比纯 PyTorch 更快、但无需直接使用 CUDA 的自定义 PyTorch 算子
• 为融合内核制作原型（例如 softmax + scale + bias）
• 使用 Triton 自动调优器比较块大小和 warp 数量
• 将 NumPy 风格的逐元素操作移植到 GPU
• 使用更高级的 Python 语法学习 GPU 编程
• 系统地对不同内核变体进行基准测试

工作流

1. 最小 Triton 内核

import torch
import triton
import triton.language as tl

@triton.jit
def add_kernel(x_ptr, y_ptr, out_ptr, n, BLOCK: tl.constexpr):
    pid = tl.program_id(0)
    offsets = pid * BLOCK + tl.arange(0, BLOCK)
    mask = offsets < n
    x = tl.load(x_ptr + offsets, mask=mask)
    y = tl.load(y_ptr + offsets, mask=mask)
    tl.store(out_ptr + offsets, x + y, mask=mask)

def add(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    n = x.numel()
    out = torch.empty_like(x)
    grid = lambda meta: (triton.cdiv(n, meta["BLOCK"]),)
    add_kernel[grid](x, y, out, n, BLOCK=1024)
    return out

关键概念：

• tl.program_id(0) — 块索引（类似 blockIdx.x）
• tl.arange(0, BLOCK) — 块内线程索引向量
• mask — 对尾部元素进行谓词处理（无需单独的边界内核）
• BLOCK: tl.constexpr — 编译期常量，可启用循环展开

2. 加载/存储与掩码

@triton.jit
def masked_load_example(ptr, n, BLOCK: tl.constexpr):
    pid = tl.program_id(0)
    offs = pid * BLOCK + tl.arange(0, BLOCK)
    mask = offs < n
    # masked load returns 0 for masked-off lanes
    vals = tl.load(ptr + offs, mask=mask, other=0.0)
    return vals

用于结构化 2D 访问的块指针（Triton 2.x+）：

@triton.jit
def matvec_kernel(a_ptr, x_ptr, y_ptr, M, N, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr):
    pid_m = tl.program_id(0)
    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    acc = tl.zeros((BLOCK_M,), dtype=tl.float32)
    for start_n in range(0, N, BLOCK_N):
        offs_n = start_n + tl.arange(0, BLOCK_N)
        a = tl.load(a_ptr + offs_m[:, None] * N + offs_n[None, :])
        x = tl.load(x_ptr + offs_n)
        acc += tl.sum(a * x[None, :], axis=1)
    tl.store(y_ptr + offs_m, acc, mask=offs_m < M)

3. 原子操作

@triton.jit
def atomic_histogram(data_ptr, hist_ptr, n, BLOCK: tl.constexpr):
    pid = tl.program_id(0)
    offs = pid * BLOCK + tl.arange(0, BLOCK)
    mask = offs < n
    data = tl.load(data_ptr + offs, mask=mask)
    bucket = (data % 256).to(tl.int32)
    tl.atomic_add(hist_ptr + bucket, 1, mask=mask)

请谨慎使用原子操作，因为它们会串行化内存更新。优先采用块级归约，然后每个块仅执行一次原子操作。

4. 通过 constexpr 使用共享内存

@triton.jit
def reduce_kernel(x_ptr, out_ptr, n, BLOCK: tl.constexpr):
    pid = tl.program_id(0)
    offs = pid * BLOCK + tl.arange(0, BLOCK)
    mask = offs < n
    x = tl.load(x_ptr + offs, mask=mask, other=0.0)
    # Block reduction
    x = tl.sum(x, axis=0)
    tl.atomic_add(out_ptr, x)

将 BLOCK 设为 tl.constexpr，可让编译器分配共享内存，并在编译期展开循环。

5. 自动调优

@triton.autotune(
    configs=[
        triton.Config({"BLOCK": 128}, num_warps=4),
        triton.Config({"BLOCK": 256}, num_warps=4),
        triton.Config({"BLOCK": 512}, num_warps=8),
    ],
    key=["n"],
)
@triton.jit
def tuned_kernel(x_ptr, y_ptr, out_ptr, n, BLOCK: tl.constexpr):
    # ... kernel body ...
    pass

自动调优器会在首次运行时对各配置进行基准测试，并为每种 key 形状缓存最佳配置。

6. 基准测试

from triton.testing import Benchmark

def benchmark_add():
    n = 1024 * 1024
    x = torch.randn(n, device="cuda")
    y = torch.randn(n, device="cuda")

    def triton_add():
        return add(x, y)

    def torch_add():
        return x + y

    bench = Benchmark(
        x_names=["n"],
        x_vals=[2**i for i in range(10, 24)],
        line_arg="provider",
        line_vals=["triton", "torch"],
        line_names=["Triton", "PyTorch"],
        plot_name="add-bench",
        args={},
    )
    bench.run(lambda n, provider: {
        "triton": lambda: add(x[:n], y[:n]),
        "torch": lambda: x[:n] + y[:n],
    }[provider](), quantiles=[0.5, 0.9])

# Quick timing in REPL
import triton.testing as tt
ms = tt.do_bench(lambda: add(x, y))
print(f"{ms:.3f} ms")

7. PyTorch 集成

import torch
from torch.library import custom_op

@custom_op("mylib::triton_add", mutates_args=())
def triton_add_op(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    return add(x, y)

@triton_add_op.register_fake
def _(x, y):
    return torch.empty_like(x)

# Use in model
class MyModule(torch.nn.Module):
    def forward(self, x, y):
        return triton_add_op(x, y)

为兼容 torch.compile，请注册 fake/meta 内核，并避免在 JIT 函数中产生 Python 副作用。

8. 调试

@triton.jit
def debug_kernel(x_ptr, n, BLOCK: tl.constexpr):
    pid = tl.program_id(0)
    offs = pid * BLOCK + tl.arange(0, BLOCK)
    mask = offs < n
    x = tl.load(x_ptr + offs, mask=mask)
    # Synchronize threads within block for inspection
    tl.debug_barrier()
    tl.store(x_ptr + offs, x * 2.0, mask=mask)

# Dump generated PTX/LLVM IR
TRITON_PRINT_AUTOTUNING=1 python script.py
# Set cache dir to inspect compiled kernels
export TRITON_CACHE_DIR=/tmp/triton_cache

在扩大规模前，先使用小型输入与 PyTorch 参考实现进行比较。

常见问题

症状	原因	解决方法
OutOfResources	BLOCK 过大	减小 BLOCK 或 num_warps
尾部结果错误	缺少掩码	在加载/存储中添加 mask=offs < n
比 PyTorch 慢	BLOCK 配置欠佳	使用 @triton.autotune
CompilationError	类型不匹配	确保 dtype 一致；使用 .to(tl.float32)
输出中出现 NaN	掩码通道未初始化	为掩码加载传入 other=0.0
torch.compile 失败	缺少 fake 内核	注册 register_fake 元函数

相关技能

• skills/gpu/cuda — Triton 遇到限制时参考底层 CUDA 概念
• skills/gpu/cuda-profiling — 对 Triton 编译的内核进行 Nsight 性能分析
• skills/gpu/gpu-memory-model — 合并访问与占用率理论
• skills/gpu/hip-rocm — AMD GPU 路径（Triton 支持 ROCm）
• skills/hpc/openmp — 与 GPU 内核配合使用的 CPU 并行计算