Mode d’emploi
- Copiez la commande et examinez le paquet avant de l’installer dans votre environnement agent.
- Consultez le document original pour les outils, identifiants et environnements requis.
- Décrivez les entrées et le livrable souhaité ; comparez la sortie aux sources.
npx skills add https://github.com/mohitmishra786/low-level-dev-skills --skill triton-langAvant de commencer
- L’installation n’inclut ni abonnement au modèle, ni quota API, ni accès aux fichiers.
- Étoiles et installations reflètent l’activité, pas un test indépendant de qualité ou de sécurité.
Document Skill original
Rendu depuis le dépôt source en conservant titres, exemples, code, tableaux, liens et images.
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 outKey concepts:
tl.program_id(0)— block index (likeblockIdx.x)tl.arange(0, BLOCK)— vector of thread indices within blockmask— 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 valsBlock 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 ...
passAutotuner 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_cacheCompare 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 hitskills/gpu/cuda-profiling— Nsight profiling of Triton-compiled kernelsskills/gpu/gpu-memory-model— coalescing and occupancy theoryskills/gpu/hip-rocm— AMD GPU path (Triton supports ROCm)skills/hpc/openmp— CPU parallelism alongside GPU kernels
