Cum 25 de linii de SQL devin un kernel FlashAttention

Cum funcționează optimizerul?

asta? Înțelegerea mecanismului de bază e esențială pentru a aprecia valoarea acestei inovații. Atenția, componenta centrală a oricărui model transformer, se bazează pe formula ``softmax (Q·Kᵀ/√d)·V``.

În această ecuație, ieșirea fiecărui token reprezintă o medie ponderată a vectorilor de valoare ``V``, unde ponderarea se bazează pe compatibilitatea dintre interogările ``Q`` ale tokenului curent și cheile ``K`` ale tokenilor anteriori.

În contextul acestui proiect, matricele ``Q``, ``K`` și ``V`` sunt tratate ca tabele SQL cu coloane ``(t, d, x)``, unde ``t`` indică poziția tokenului, ``d`` indică dimensiunea caracteristicilor, iar ``x`` conține valorile.

Cum se validează performanța?

Calculul atenției se efectuează prin două sume bazate pe operații de join: prima calculează scorurile prin înmulțire element cu element a ``Q`` și ``K`` pe axa caracteristicilor, iar a doua calculează media ponderată a ``V`` folosind greutăți obținute din exp (scor).

Când un motor SQL execută această interogare în forma ei naivă, rezultatul este un plan de execuție care stochează toate scorurile N×N în memorie. Acest comportament e ineficient pentru secvențe lungi.

Totuși, algoritmul FlashAttention rezolvă această problemă prin procesare în blocuri, evitând stocarea completă a matricei de scoruri. Einfold, ca optimizer de interogări care rescrie planurile SQL în algoritmi de matrice, identifică automat acest tip de plan. Mai mult, o singură regulă generală în einfold fuzionează orice medie ponderată cu o funcție softmax a cărei scoruri derivă dintr-o operație de join.

Această regulă se aplică nu doar la atenție, ci și la alte interogări cu structură similară, demonstrând puterea de generalizare a optimizerului.

Cum se testează gradientele?

Proiectul include un script de benchmarking, disponibil la calea ``benchmarks/ml/01_attention. py``, care validează corectitatea și performanța implementării.

Scriptul rulează mai multe teste pe mai multe motoare de baze de date, numite „hosts” în terminologia einfold. Printre aceștia se numără DataFusion, DataFusion cu extensia einfold activă, și DuckDB. Validarea se face prin compararea rezultatelor obținute din SQL cu cele generate de JAX, biblioteca Google pentru programe diferențiabile de matrice.

Primul tip de test verifică trecerea înainte. Fiecare variantă a scrierii modelului se execută pe fiecare motor și se compară cu rezultatul JAX.

Cum se tratează măștile?

Al doilea tip de test verifică calculul gradientelor.

Se definește o funcție de pierdere care citește ieșirea modelului, iar gradientul ei față de fiecare valoare din ``Q``, ``K`` și ``V`` se calculează printr-o singură interogare SQL.

Pentru asta se folosește ddx, alt proiect de la xqlsystems, care diferențiază interogări SQL. JAX se folosește ca referință pentru validarea gradientelor calculate.

Doar pentru atenție?

Un aspect critic al implementării este tratarea măștilor. În atenția cauzală, perechile de tokeni care nu ar trebui să fie vizibili unuia față de celălalt se tratează ca linii lipsă în rezultatul interogării, nu ca linii cu greutate zero.

Acest lucru e esențial pentru corectitatea matematică, deoarece o greutate zero înmulțită cu o valoare infinită ar produce un rezultat nan (Not a Number), care ar rupe calculul.

Prin tratarea măștilor ca linii lipsă, această problemă se evită, iar semantica SQL garantează că orice plan mai rapid trebuie să respecte această regulă.

De ce e esențial să înțelegi mecanismul de bază al atenției?

Ca control pentru a demonstra generalitatea optimizerului, se include și estimatorul Nadaraya–Watson, o tehnică de regridding care folosește o medie ponderată cu un nucleu Gaussian. Acesta nu e atenție, dar are o structură identică din punctul de vedere al optimizerului: o medie ponderată cu greutăți calculate ca exp (scor).

Orice regulă care accelerează atenția trebuie neapărat să accelereze și această interogare de control, demonstrând că optimizarea nu e specifică atenției, ci se aplică oricărei interogări cu structură similară.

Scriptul de testare include o verificare specifică a comportamentului măștilor. Se creează o matrice de testare cu 16 tokeni, iar la poziția 9 se setează o valoare infinită în matricea ``V``.

Dacă măștile sunt implementate corect ca linii lipsă, această valoare infinită ar trebui să afecteze doar interogările care le e permis să o vadă, adică cele cu index mare sau egal cu 9.

Cum se execută?

Testul verifică că niciun alt token nu e afectat, confirmând corectitatea implementării.

Un alt test important verifică că einfold identifică automat planurile fuzionate.

Fiecare variantă a scrierii modelului, precum și interogarea de control, trebuie să se executeze ca unul dintre pipeline-urile fuzionate de einfold. Acest lucru se verifică prin inspectarea planurilor de execuție generate de motor. Dacă optimizerul nu identifică automat fuziunea, înseamnă că nu e suficient de general.