Kerneluri Triton pentru GPU, scrise în Python în PyTorch & Deep Learning Practic

import triton
import triton.language as tl

@triton.jit
def kernel(x_ptr, out_ptr, n, BLOCK_SIZE: tl.constexpr):
    pid = tl.program_id(axis=0)
    offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    masca = offs < n
    x = tl.load(x_ptr + offs, mask=masca)
    tl.store(out_ptr + offs, x * 2, mask=masca)

grid = lambda meta: (triton.cdiv(n, meta["BLOCK_SIZE"]),)
kernel[grid](x, out, n, BLOCK_SIZE=1024)

Ghidul de programare Triton descrie modelul: în CUDA scrii un program scalar rulat de blocuri de fire, în Triton scrii un program care operează pe blocuri de date (blocked program, scalar threads). Compilatorul face automat optimizări care în CUDA se scriu de mână, între care coalescing-ul accesului la memorie, alocarea și sincronizarea memoriei shared și vectorizarea. Versiunea curentă pe PyPI este 3.8.0 (pip install triton, roți pentru CPython 3.10–3.14).

Documentația torch.compiler spune că TorchInductor, backend-ul implicit al lui torch.compile, folosește Triton ca element de bază pentru GPU-urile NVIDIA, AMD și Intel. Deci pentru fuziunea operațiilor obișnuite nu scrii Triton de mână; kerneluri proprii se scriu pentru operații pe care compilatorul nu le generează bine.

Un kernel este o funcție decorată cu @triton.jit. Fiecare instanță își află poziția cu tl.program_id, calculează deplasările cu tl.arange și citește / scrie cu tl.load / tl.store, cu o mască pentru ultimul bloc incomplet. Kernelul se lansează cu kernel[grid](...); tensorii PyTorch se transmit ca pointeri către primul element, iar parametrii tl.constexpr se dau cu nume. Lansarea este asincronă. Din PyTorch 2.3, kernelurile Triton proprii pot fi apelate din funcții compilate cu torch.compile; pentru integrare completă (AOTInductor, antrenare cu register_autograd) tutorialul recomandă torch.library.triton_op cu wrap_triton.

Concepte cheie și caracteristici

Concept / Parametru	Descriere practică
@triton.jit	Marchează o funcție ca kernel Triton.
tl.program_id(axis)	Indicele instanței de program pe axa dată a grilei.
tl.arange(start, end)	Vector de indici; mărimea blocului trebuie cunoscută la compilare (tl.constexpr).
tl.load(ptr, mask=None, other=None) / tl.store(ptr, val, mask=None)	Citire / scriere din memoria GPU, cu mască pentru elementele în afara limitelor.
triton.cdiv(a, b)	Împărțire rotunjită în sus; numărul de blocuri din grilă.
kernel[grid](...)	Lansare; grid este un tuplu sau o funcție de meta-parametri.
triton.autotune(configs, key)	Alege automat configurația; cu torch.compile sunt acceptate doar configs, key, restore_value, reset_to_zero.
torch.library.triton_op / wrap_triton	Operator PyTorch bazat pe kerneluri Triton, compatibil cu torch.compile și AOTInductor.
triton.testing.do_bench(fn)	Măsoară timpul unei funcții pe GPU.

Exemple practice și scenarii reale

Exemplul 1: Adunarea a doi vectori (tutorialul oficial Triton)

import torch
import triton
import triton.language as tl

DEVICE = triton.runtime.driver.active.get_active_torch_device()

@triton.jit
def add_kernel(x_ptr, y_ptr, output_ptr, n_elements, BLOCK_SIZE: tl.constexpr):
    pid = tl.program_id(axis=0)
    block_start = pid * BLOCK_SIZE
    offsets = block_start + tl.arange(0, BLOCK_SIZE)
    mask = offsets < n_elements
    x = tl.load(x_ptr + offsets, mask=mask)
    y = tl.load(y_ptr + offsets, mask=mask)
    tl.store(output_ptr + offsets, x + y, mask=mask)

def add(x, y):
    output = torch.empty_like(x)
    n_elements = output.numel()
    grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),)
    add_kernel[grid](x, y, output, n_elements, BLOCK_SIZE=1024)
    return output

x = torch.rand(98432, device=DEVICE)
y = torch.rand(98432, device=DEVICE)
print(torch.max(torch.abs(add(x, y) - (x + y))))

Exemplul 2: Kernel propriu apelat din torch.compile

Din rețeta PyTorch pentru kerneluri Triton definite de utilizator (PyTorch 2.3+).

@torch.compile(fullgraph=True)
def add_fn(x, y):
    output = torch.zeros_like(x)
    n_elements = output.numel()
    grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),)
    add_kernel[grid](x, y, output, n_elements, BLOCK_SIZE=4)
    return output

Exemplul 3: Operator PyTorch cu triton_op

from torch.library import triton_op, wrap_triton

@triton_op("mylib::aduna", mutates_args={})
def aduna(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    out = torch.empty_like(x)
    n = x.numel()
    wrap_triton(add_kernel)[(triton.cdiv(n, 1024),)](x, y, out, n, BLOCK_SIZE=1024)
    return out

z = torch.compile(aduna)(x, y)

Greșeli frecvente și capcane

• Să omiți masca la tl.load / tl.store. Ultimul bloc citește și scrie dincolo de capătul tensorului când numărul de elemente nu e multiplu de BLOCK_SIZE.
• Să compari timpul unui kernel fără sincronizare. Lansarea este asincronă; folosește triton.testing.do_bench sau torch.cuda.synchronize().
• Să scrii de mână în Triton operații pe care torch.compile le fuzionează deja. Profilează întâi modelul compilat.

Noțiuni și module înrudite

• torch-compile
• torch-scaled-dot-product-attention
• torch-profiler-cuda-memory