Triton
Write fast GPU kernels in Python.
Triton is a language and compiler from OpenAI for writing GPU kernels in Python. You program blocks of data instead of individual threads, and the compiler handles memory coalescing, shared memory and tensor core use. PyTorch's torch.compile generates Triton kernels, and FlashAttention-style kernels are often written in it.
Why learn Triton?
When PyTorch is too slow for one operation and CUDA C++ feels like too much, Triton is the middle path most AI teams now use.
- Python syntax: kernels are Python functions with a decorator.
- Fast: often matches hand-written CUDA for matmul, softmax and attention.
- Portable: the same kernel runs on NVIDIA and AMD GPUs.
Set up and run your first program
Triton ships with PyTorch on Linux. You need an NVIDIA (compute capability 8.0+ recommended) or supported AMD GPU.
pip install torch triton
python -c "import triton; print(triton.__version__)"
python hello.pyimport torch, triton, triton.language as tl
@triton.jit
def hello(out_ptr):
tl.store(out_ptr, 42)
out = torch.empty(1, device="cuda", dtype=torch.int32)
hello[(1,)](out)
print("Hello, AI! Triton kernel wrote", out.item())The block programming model
A Triton kernel runs as a grid of programs. Each program handles one block of data. tl.program_id(0) tells a program which block it owns, and tl.arange builds the offsets inside that block.
import torch, triton, triton.language as tl
@triton.jit
def add_kernel(x_ptr, y_ptr, out_ptr, n, BLOCK: tl.constexpr):
pid = tl.program_id(0)
offs = pid * BLOCK + tl.arange(0, BLOCK)
mask = offs < n # guard the last block
x = tl.load(x_ptr + offs, mask=mask)
y = tl.load(y_ptr + offs, mask=mask)
tl.store(out_ptr + offs, x + y, mask=mask)
n = 98_432
x, y = torch.rand(n, device="cuda"), torch.rand(n, device="cuda")
out = torch.empty_like(x)
grid = (triton.cdiv(n, 1024),)
add_kernel[grid](x, y, out, n, BLOCK=1024)
print(torch.allclose(out, x + y))Loads, stores and masks
tl.load and tl.store move whole blocks between GPU memory and registers. Masks stop out-of-bounds reads; other= sets the value for masked lanes.
@triton.jit
def relu_kernel(x_ptr, out_ptr, n, BLOCK: tl.constexpr):
offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
mask = offs < n
x = tl.load(x_ptr + offs, mask=mask, other=0.0)
tl.store(out_ptr + offs, tl.maximum(x, 0.0), mask=mask)Fusing a softmax
The big win in Triton is fusion: read data once, do several steps in registers, write once. This row-wise softmax does max, exp, sum and divide in a single kernel.
@triton.jit
def softmax_kernel(out_ptr, in_ptr, stride, n_cols, BLOCK: tl.constexpr):
row = tl.program_id(0)
cols = tl.arange(0, BLOCK)
mask = cols < n_cols
x = tl.load(in_ptr + row * stride + cols, mask=mask, other=-float("inf"))
x = x - tl.max(x, axis=0)
num = tl.exp(x)
y = num / tl.sum(num, axis=0)
tl.store(out_ptr + row * stride + cols, y, mask=mask)
def softmax(x):
rows, cols = x.shape
out = torch.empty_like(x)
softmax_kernel[(rows,)](out, x, x.stride(0), cols, BLOCK=triton.next_power_of_2(cols))
return outMatrix multiply with tl.dot
tl.dot multiplies two blocks and uses tensor cores automatically. Each program computes one tile of the output by looping over K.
@triton.jit
def matmul_kernel(a, b, c, M, N, K, sam, sak, sbk, sbn, scm, scn,
BM: tl.constexpr, BN: tl.constexpr, BK: tl.constexpr):
pm, pn = tl.program_id(0), tl.program_id(1)
rm = pm * BM + tl.arange(0, BM)
rn = pn * BN + tl.arange(0, BN)
rk = tl.arange(0, BK)
acc = tl.zeros((BM, BN), dtype=tl.float32)
for k in range(0, K, BK):
A = tl.load(a + rm[:, None] * sam + (k + rk)[None, :] * sak,
mask=(rm[:, None] < M) & ((k + rk)[None, :] < K), other=0.0)
B = tl.load(b + (k + rk)[:, None] * sbk + rn[None, :] * sbn,
mask=((k + rk)[:, None] < K) & (rn[None, :] < N), other=0.0)
acc += tl.dot(A, B)
tl.store(c + rm[:, None] * scm + rn[None, :] * scn, acc.to(tl.float16),
mask=(rm[:, None] < M) & (rn[None, :] < N))Autotuning and benchmarking
Block sizes change speed a lot. @triton.autotune tries several configs and keeps the fastest; triton.testing.do_bench times a kernel.
@triton.autotune(
configs=[triton.Config({"BLOCK": b}, num_warps=w) for b in (256, 1024, 4096) for w in (4, 8)],
key=["n"],
)
@triton.jit
def scale_kernel(x_ptr, out_ptr, n, s, BLOCK: tl.constexpr):
offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
m = offs < n
tl.store(out_ptr + offs, tl.load(x_ptr + offs, mask=m) * s, mask=m)
x = torch.rand(1 << 24, device="cuda"); out = torch.empty_like(x)
grid = lambda meta: (triton.cdiv(x.numel(), meta["BLOCK"]),)
ms = triton.testing.do_bench(lambda: scale_kernel[grid](x, out, x.numel(), 2.0))
print(f"{ms:.3f} ms, {3 * x.numel() * 4 / ms / 1e6:.0f} GB/s")Triton in the AI stack
Triton sits between PyTorch and the GPU. torch.compile turns PyTorch code into Triton kernels automatically, libraries such as Liger Kernel and Unsloth use hand-written Triton to speed up LLM training, and attention kernels like FlashAttention have Triton versions. Learning it shows you exactly why memory bandwidth and tensor cores matter.
The essential libraries
See the Triton code torch.compile writes
Set an environment variable and PyTorch prints the generated kernels, a great way to learn.
# TORCH_LOGS="output_code" python fused.py
import torch
def gelu_bias(x, b):
return torch.nn.functional.gelu(x + b)
f = torch.compile(gelu_bias)
x = torch.randn(4096, 4096, device="cuda")
b = torch.randn(4096, device="cuda")
print(torch.allclose(f(x, b), gelu_bias(x, b), atol=1e-5))An agent tool backed by a Triton kernel
Triton kernels are not agents, but an agent can use them as tools. Here a local model decides when to call a GPU softmax kernel (from the fusion chapter) through the OpenAI-compatible API.
import json, torch
from openai import OpenAI
from softmax import softmax # the Triton kernel above
client = OpenAI(base_url="http://localhost:11434/v1", api_key="ollama")
tools = [{"type": "function", "function": {"name": "gpu_softmax",
"description": "Softmax of a list of numbers, computed on the GPU",
"parameters": {"type": "object", "properties": {"values": {"type": "array", "items": {"type": "number"}}},
"required": ["values"]}}}]
def gpu_softmax(values):
x = torch.tensor([values], device="cuda", dtype=torch.float32)
return [round(v, 4) for v in softmax(x)[0].tolist()]
messages = [{"role": "user", "content": "Turn the scores 2.0, 1.0 and 0.1 into probabilities."}]
for _ in range(5):
msg = client.chat.completions.create(model="llama3.1", messages=messages, tools=tools).choices[0].message
messages.append(msg)
if not msg.tool_calls:
print(msg.content); break
for call in msg.tool_calls:
result = gpu_softmax(**json.loads(call.function.arguments))
messages.append({"role": "tool", "tool_call_id": call.id, "content": json.dumps(result)})In real systems, Triton speeds up the model itself; the agent logic stays in plain Python.
ollama pull llama3.1, and keep it running. The code talks to its OpenAI-compatible API at http://localhost:11434/v1, so the same program also works with vLLM, llama.cpp server or a cloud API by changing the URL, model name and adding an API key.Triton learning roadmap
Tick items off as you go. Your progress is saved in this browser.
Beginner
Intermediate
Advanced
Projects to build
Triton cheat sheet
| Task | Syntax |
|---|---|
| Decorator | @triton.jit |
| Block id | pid = tl.program_id(0) |
| Offsets | pid * B + tl.arange(0, B) |
| Load | tl.load(p + offs, mask=m) |
| Launch | kern[grid](args, B=1024) |
| Matmul | acc += tl.dot(a, b) |
Common mistakes
- Block sizes passed to
tl.arangemust be powers of two andtl.constexpr. - Forgetting the mask reads or writes past the end of the tensor.
- Python features like lists and dicts are not allowed inside a kernel.
Where to go next
- Triton tutorials — official
- Triton GitHub — source
- GPU MODE — kernel lectures
On TensorCore.si: CUDA volume, What is a tensor core?.
Keyboard: [ previous language · ] next language · / search the library