triton-lang

사용 방법

설치 명령을 복사하고 자신의 에이전트 환경에 설치하기 전에 소스 패키지를 확인하세요.
필요한 도구, 인증 정보, 지원 환경은 원본 문서에서 확인하세요.
입력과 원하는 산출물을 설명하고 출처와 결과를 대조하세요.

npx skills add https://github.com/mohitmishra786/low-level-dev-skills --skill triton-lang

시작 전 확인

설치에는 모델 구독, API 한도, 파일 접근 권한이 포함되지 않습니다.
별점 수와 설치 수는 활동 지표이며 독립적인 품질·보안 검증이 아닙니다.

원본 Skill 문서

원본 저장소의 제목, 예시, 코드, 표, 링크, 이미지를 유지해 표시합니다.

Triton

Purpose

Guide agents through writing GPU kernels in OpenAI Triton: the @triton.jit decorator, block-oriented tl.load/tl.store with masking, atomic operations, shared memory via tl.constexpr, benchmarking with triton.testing.Benchmark, PyTorch integration, and debugging with barriers.

When to Use

Writing custom PyTorch ops faster than pure PyTorch but without raw CUDA
Prototyping fused kernels (e.g., softmax + scale + bias)
Comparing block sizes and warp counts with Triton's autotuner
Porting NumPy-style elementwise ops to GPU
Learning GPU programming with higher-level Python syntax
Benchmarking kernel variants systematically

Workflow

1. Minimal Triton kernel

python

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

Key concepts:

tl.program_id(0) — block index (like blockIdx.x)
tl.arange(0, BLOCK) — vector of thread indices within block
mask — predication for tail elements (no separate bounds kernel)
BLOCK: tl.constexpr — compile-time constant, enables unrolling

2. Load/store and masking

python

@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

Block pointers (Triton 2.x+) for structured 2D access:

python

@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. Atomic operations

python

@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)

Use atomics sparingly — they serialize memory updates. Prefer block-level reduction then single atomic per block.

4. Shared memory via constexpr

python

@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 as tl.constexpr lets the compiler allocate shared memory and unroll loops at compile time.

5. Autotuning

python

@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

Autotuner benchmarks configs on first run and caches the best for each key shape.

6. Benchmarking

python

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])

bash

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

7. PyTorch integration

python

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)

For torch.compile compatibility, register fake/meta kernels and avoid Python side effects in the JIT function.

8. Debugging

python

@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)

bash

# 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

Compare against PyTorch reference on small inputs before scaling up.

Common Problems

Symptom	Cause	Fix
`OutOfResources`	BLOCK too large	Reduce BLOCK or num_warps
Wrong results on tail	Missing mask	Add `mask=offs < n` to load/store
Slower than PyTorch	Suboptimal BLOCK	Use `@triton.autotune`
`CompilationError`	Type mismatch	Ensure consistent dtypes; use `.to(tl.float32)`
NaN in output	Uninitialized masked lanes	Pass `other=0.0` to masked loads
torch.compile fails	No fake kernel	Register `register_fake` meta function

Related Skills

skills/gpu/cuda — underlying CUDA concepts when Triton limits are hit
skills/gpu/cuda-profiling — Nsight profiling of Triton-compiled kernels
skills/gpu/gpu-memory-model — coalescing and occupancy theory
skills/gpu/hip-rocm — AMD GPU path (Triton supports ROCm)
skills/hpc/openmp — CPU parallelism alongside GPU kernels

사용 방법

시작 전 확인

Triton

Purpose

When to Use

Workflow

1. Minimal Triton kernel

2. Load/store and masking

3. Atomic operations

4. Shared memory via constexpr

5. Autotuning

6. Benchmarking

7. PyTorch integration

8. Debugging

Common Problems

Related Skills

관련 작업

더 많은 Skills

core-dumps

custom-allocators

ebpf

flamegraphs

triton-lang

사용 방법

시작 전 확인

Triton

Purpose

When to Use

Workflow

1. Minimal Triton kernel

2. Load/store and masking

3. Atomic operations

4. Shared memory via constexpr

5. Autotuning

6. Benchmarking

7. PyTorch integration

8. Debugging

Common Problems

Related Skills

관련 작업

python-patterns

docker-patterns

rust-patterns

rea-analysis-optimization

더 많은 Skills

core-dumps

custom-allocators

ebpf

flamegraphs