Optimizarea unui model AI nu începe cu promisiunea că „GPU-ul va fi mai rapid”, ci cu o întrebare măsurabilă: ce operații repetate mută prea multe date între memoria globală și unitățile de calcul? Triton, limbajul și compilatorul folosit pentru a scrie kerneluri GPU în stil Python, este util exact aici: când o secvență mică de operații PyTorch este corectă, dar plătește costuri intermediare prea mari. Scopul nu este să înlocuim PyTorch peste tot, ci să testăm disciplinat dacă o fuziune locală merită.
Exemplul clasic din documentația Triton este softmax-ul fuzionat. Varianta naivă calculează maximul, scăderea pentru stabilitate numerică, exponentul, suma și împărțirea în pași separați. Fiecare pas poate citi și scrie tensori intermediari. Un kernel fuzionat poate încărca un bloc de date, poate face reducerea și normalizarea în aceeași execuție și poate scrie rezultatul o singură dată. Ideea generală se aplică și în modele reale: normalizări, activări, bias plus transformări element-wise sau secvențe scurte unde memoria, nu calculul brut, domină timpul.
Primul pas este să alegi o țintă îngustă. Nu optimiza întregul model. Rulează profilarea în PyTorch și caută o zonă stabilă, apelată des, cu forme de tensor previzibile. O operație rară sau una care primește dimensiuni foarte variabile poate să nu justifice costul unui kernel personalizat. Notează exact hardware-ul, versiunea PyTorch, versiunea Triton, tipul de precizie și dimensiunile tensorilor. Fără aceste detalii, benchmark-ul nu poate fi reprodus.
Al doilea pas este să definești comportamentul de referință în PyTorch. Implementarea PyTorch trebuie să fie clară, stabilă numeric și testată. Pentru softmax, asta înseamnă scăderea maximului înainte de exponentiere, nu o formulă fragilă. Pentru normalizări, verifică axele, broadcasting-ul și tipul de acumulare. Referința este adevărul funcțional: kernelul Triton trebuie să producă rezultate apropiate, nu doar să pară mai rapid.
Abia apoi proiectezi kernelul. În Triton, gândești în blocuri: ce porțiune din tensor încarcă fiecare program, ce mască protejează accesul la margini, ce dimensiune de bloc se potrivește în memoria rapidă și câte warp-uri folosești. Documentația Triton arată că anumite blocuri trebuie rotunjite la puteri ale lui doi și protejate prin măști. Acesta este un detaliu important pentru corectitudine: optimizarea nu trebuie să citească în afara datelor reale.
Fuziunea trebuie să fie explicită. În loc să materializezi tensorul intermediar după fiecare pas, păstrezi valorile în registru sau în memoria on-chip cât timp calculezi rezultatul final. Pentru o secvență de tip „bias + activare + scalare”, kernelul citește intrarea și parametrii, aplică operațiile în ordine și scrie ieșirea. Pentru o reducere, cum este softmax sau o normalizare, kernelul trebuie să calculeze întâi statistica necesară, apoi rezultatul final. Aici apar limitele: dacă rândul nu încape eficient în bloc, dacă accesul la memorie este prost aliniat sau dacă ocuparea GPU scade, Triton poate să nu bată PyTorch.
Verificarea corectitudinii vine înaintea vitezei. Rulează teste pe dimensiuni mici, dimensiuni nealiniate, valori extreme, tensori cu precizii diferite și cazuri de margine. Compară cu torch.allclose sau o metodă echivalentă, cu toleranțe adaptate preciziei. Pentru float32 poți cere toleranțe mai stricte decât pentru float16 sau bfloat16. Dacă rezultatul diferă, nu ajusta benchmark-ul ca să arate bine; repară kernelul.
Benchmark-ul corect trebuie să evite trei capcane: compilarea inclusă în timp, lipsa sincronizării GPU și un singur rezultat convenabil. Triton oferă triton.testing.do_bench, cu încălzire și repetări, iar PyTorch recomandă torch.utils.benchmark.Timer pentru comparații mai robuste decât timeit simplu. Pe GPU, operațiile sunt asincrone, deci sincronizarea contează. Rulează o încălzire, măsoară de mai multe ori, raportează mediană sau cuantile, nu doar „cel mai bun timp”.
Compară mere cu mere. Aceleași dimensiuni, același dtype, același device, aceleași date de intrare și aceeași ieșire verificată. Include și implementarea PyTorch idiomatică, nu doar o variantă intenționat naivă. Dacă PyTorch 2.x cu torch.compile sau operatorii nativi deja fuzionează bine cazul tău, kernelul personalizat poate aduce puțin sau chiar poate pierde. Rezultatele depind de GPU, lățimea de bandă a memoriei, dimensiuni, precizie, layout, driver și implementare.
Concluzia practică este simplă: un kernel Triton personalizat merită doar când ai o problemă măsurată, o referință corectă și un benchmark reproductibil. Pentru echipe mici, valoarea nu este în a scrie kerneluri peste tot, ci în a optimiza punctual blocajele reale. Dacă poți arăta că fuziunea reduce mișcările inutile de date și păstrează aceeași ieșire numerică, ai o optimizare legitimă. Dacă nu, PyTorch rămâne alegerea mai sigură, mai ușor de întreținut și mai portabilă.
