triton-lang

Triton

目的

エージェントが OpenAI Triton で GPU カーネルを記述できるように、@triton.jit デコレーター、マスク処理を使ったブロック単位の tl.load／tl.store、アトミック操作、tl.constexpr による共有メモリ、triton.testing.Benchmark によるベンチマーク、PyTorch との統合、バリアを使ったデバッグについて案内します。

使用する場面

• 生の CUDA を使わずに、純粋な PyTorch より高速なカスタム 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)

アトミック操作はメモリ更新を直列化するため、必要最小限にとどめてください。まずブロック単位でリダクションし、各ブロックにつき 1 回だけアトミック操作を行う方法を優先してください。

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 の meta 関数を登録します

関連スキル

• 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 並列処理