Self-Attention, CUDA कर्नल आणि FlashAttention

आधुनिक Transformers Self-Attention मॅट्रिक्स गणित, Triton सारखे उच्च-स्तरीय GPU कंपायलर आणि FlashAttention सारखे मेमरी-जागरूक CUDA कर्नल यांच्यावर अवलंबून असतात जेणेकरून प्रचंड टोकन सिक्वन्सवर हार्डवेअरच्या कमाल वेगाने काम होऊ शकते.

1. Vision Transformer मध्ये Self-Attention गणितीयदृष्ट्या कसे काम करते

Vision Transformer (ViT) मध्ये, प्रतिमा (उदा. 224×224 पिक्सेल) 16×16 पॅचच्या ग्रिडमध्ये कापली जाते, त्यामुळे N = 196 सपाट केलेले पॅच मिळतात. प्रत्येक पॅच D परिमाणाच्या (उदा. 768) व्हेक्टरमध्ये रेषीयरीत्या प्रोजेक्ट केला जातो आणि पोझिशनल एम्बेडिंगसह जोडला जातो, जेणेकरून इनपुट मॅट्रिक्स X ∈ R^(N × D) तयार होते.

• Query, Key आणि Value तयार करणे. मॉडेल पॅच मॅट्रिक्स X ला तीन शिकल्या गेलेल्या वजन मॅट्रिक्स (W_Q, W_K, W_V ∈ R^(D × d_k)) ने गुणते जेणेकरून प्रत्येक प्रतिमा पॅचसाठी तीन वेगळी प्रतिनिधित्वे तयार होतात:

Q=XW_(Q), K=XW_(K), V=XW_(V)

• Query (Q): विशिष्ट प्रतिमा पॅच काय शोधत आहे हे दर्शवते (उदा. "मी केसांसारखा टेक्स्चर आहे, जवळ डोळे किंवा कान शोधत आहे").
• Key (K): पॅचमध्ये कोणती माहिती आहे हे जाहीर करते (उदा. "माझ्याकडे तीक्ष्ण वक्र कडा आहे").
• Value (V): जुळणी झाल्यास पुढे पाठवले जाणारे प्रत्यक्ष दृश्य फीचर पेलोड धरते.

• अटेन्शन स्कोअर मोजणे. प्रत्येक पॅच इतर प्रत्येक पॅचशी किती संबंधित आहे हे मोजण्यासाठी, मॉडेल सर्व Query आणि Key मधील डॉट प्रोडक्ट मोजते:

Scores=(QK^(T))/(sqrt(d_(k)))

sqrt(d_(k)) ने भाग दिल्यामुळे मोठे डॉट प्रोडक्ट पुढे येणाऱ्या softmax फंक्शनला अतिशय लहान ग्रेडियंट असलेल्या प्रदेशात ढकलत नाहीत. यामुळे N × N अटेन्शन मॅट्रिक्स मिळते जिथे एन्ट्री (i, j) हे स्कोअर करते की पॅच i ने पॅच j कडे किती लक्ष द्यायला हवे.

• Softmax वेटिंग आणि एकत्रीकरण. पंक्तीनुसार Softmax लागू केल्यामुळे कच्चे स्कोअर अशा संभाव्यतांत बदलतात ज्यांची बेरीज 1 होते. या संभाव्यतांना Value मॅट्रिक्स V ने गुणल्यामुळे प्रत्येक पॅचसाठी मिश्रित प्रतिनिधित्व मिळते:

Attention(Q,K,V)=softmax((QK^(T))/(sqrt(d_(k))))V

मल्टि-हेड Self-Attention मध्ये, हीच प्रक्रिया अनेक हेडवर समांतर चालते (उदा. 12 हेड), ज्यामुळे वेगवेगळे हेड एकाच वेळी वेगवेगळ्या अवकाशी अंतर, रंग किंवा वस्तूंच्या भागांवर लक्ष केंद्रित करू शकतात.

2. प्रत्यक्ष कोड तुलना: कच्चे CUDA C++ विरुद्ध OpenAI Triton

GPU प्रोग्रामिंग कशी विकसित झाली हे पाहण्यासाठी, N घटकांसाठी साधी एलिमेंट-वाइझ व्हेक्टर बेरीज (C = A + B) विचारात घ्या.

कच्चे CUDA C++ (थ्रेड-स्तरीय व्यवस्थापन)

CUDA C++ मध्ये, तुम्ही एका स्केलर थ्रेडच्या दृष्टिकोनातून प्रोग्राम करता. तुम्हाला थ्रेड ID स्वतः मोजावे लागतात, अॅरेच्या मर्यादा तपासाव्या लागतात आणि होस्ट-साइड लॉन्च कॉन्फिगरेशनचा बॉयलरप्लेट कोड लिहावा लागतो.

#include <cuda_runtime.h>

// Device kernel: executed by thousands of individual threads
__global__ void vector_add_cuda(const float* A, const float* B, float* C, int N) {
    // Manually compute global thread index from block and thread IDs
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    
    // Guard against out-of-bounds memory access
    if (idx < N) {
        C[idx] = A[idx] + B[idx];
    }
}

// Host launch setup
void launch_cuda(const float* d_A, const float* d_B, float* d_C, int N) {
    int threads_per_block = 256;
    int blocks_per_grid = (N + threads_per_block - 1) / threads_per_block;
    vector_add_cuda<<<blocks_per_grid, threads_per_block>>>(d_A, d_B, d_C, N);
}

OpenAI Triton (ब्लॉक-स्तरीय Python प्रोग्रामिंग)

Triton मध्ये, तुम्ही Python मध्येच डेटाच्या ब्लॉकच्या (टाइल) दृष्टिकोनातून प्रोग्राम करता. Triton चा कंपायलर थ्रेड सिंक्रोनायझेशन, मेमरी कोअलेसिंग आणि व्हेक्टरायझेशन आपोआप आतून सांभाळतो.

import triton
import triton.language as tl

@triton.jit
def vector_add_triton(A_ptr, B_ptr, C_ptr, N, BLOCK_SIZE: tl.constexpr):
    # Identify which block of data this program instance handles
    pid = tl.program_id(axis=0)
    
    # Create a vector of offsets for the entire block at once
    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = offsets < N
    
    # Load chunks of memory into fast SRAM, add them, and store back
    a = tl.load(A_ptr + offsets, mask=mask)
    b = tl.load(B_ptr + offsets, mask=mask)
    tl.store(C_ptr + offsets, a + b, mask=mask)

# Launch directly from PyTorch tensors
def launch_triton(A, B, C):
    N = A.numel()
    grid = lambda meta: (triton.cdiv(N, meta['BLOCK_SIZE']),)
    vector_add_triton[grid](A, B, C, N, BLOCK_SIZE=1024)

3. FlashAttention: Transformer चा मेमरी अडथळा दूर करणे

प्रमाण Self-Attention उच्च-रिझोल्यूशन प्रतिमा किंवा लांब मजकूर सिक्वन्सवर काम करताना गंभीर भिंतीला आपटते कारण अटेन्शन मॅट्रिक्स S = QK^T सिक्वन्स लांबी N सह द्विघातीय रीतेने (N × N) वाढते.

• मेमरी बँडविड्थ अडथळा (HBM विरुद्ध SRAM). आधुनिक NVIDIA GPU मध्ये दोन मुख्य मेमरी स्तर असतात:

• HBM (High Bandwidth Memory / VRAM): मोठी (उदा. 80 GB) पण वाचण्यासाठी-लिहिण्यासाठी तुलनेने संथ.
• SRAM (ऑन-चिप शेअर्ड मेमरी): लहान (सर्व Streaming Multiprocessors मिळून एकूण ~20 MB) पण HBM पेक्षा साधारण 10× ते 20× जलद.

प्रमाण अटेन्शनमध्ये, GPU S = QK^T मोजते, प्रचंड N × N मॅट्रिक्स संथ HBM मध्ये लिहिते, Softmax लागू करण्यासाठी ती परत वाचते, संभाव्यता HBM मध्ये परत लिहिते आणि V ने गुणण्यासाठी त्या पुन्हा वाचते. Tensor Cores त्यांचा बहुतेक वेळ मेमरी ट्रान्सफरची वाट पाहत निष्क्रिय घालवतात.

• टायलिंग (N × N मॅट्रिक्स कधीही प्रत्यक्षात न तयार करणे). Tri Dao व सहकारी (2022) यांनी सादर केलेले FlashAttention कर्नल फ्यूजन वापरून संपूर्ण अटेन्शन गणना जलद ऑन-चिप SRAM मध्येच ठेवते. ते Q, K आणि V चे SRAM मध्ये बसणारे छोटे ब्लॉक करते, टाइल-बाय-टाइल अटेन्शन मोजते आणि फक्त अंतिम आउटपुट मॅट्रिक्स HBM मध्ये परत लिहिते. प्रचंड N × N मधली मॅट्रिक्स कधीही VRAM मध्ये साठवली जात नाही.

• ऑनलाइन Softmax युक्ती. सामान्यतः, Softmax साठी कोणताही एक घटक प्रमाणित करण्यापूर्वी संपूर्ण पंक्तीतील कमाल मूल्य आणि घातांकांची बेरीज माहित असणे आवश्यक असते. FlashAttention ऑनलाइन Softmax अल्गोरिदम वापरते जी ब्लॉक-बाय-ब्लॉक चालू कमाल मूल्ये आणि प्रमाणीकरण घटकांचा मागोवा ठेवते, आणि प्रत्येक नवीन टाइलवर प्रक्रिया होताच आंशिक आउटपुटला गणितीयरीत्या पुन्हा स्केल करते.

• बॅकप्रॉपगेशनदरम्यान पुनर्गणना. FlashAttention फॉरवर्ड पासमध्ये N × N अटेन्शन मॅट्रिक्स साठवत नसल्यामुळे, ते प्रशिक्षणादरम्यान हुशार तडजोड करते: ते HBM मध्ये फक्त छोटे प्रमाणीकरण आकडे साठवते आणि बॅकप्रॉपगेशनदरम्यान जलद SRAM मध्ये अटेन्शन टाइल्स तत्काळ पुन्हा मोजते. SRAM मध्ये अतिरिक्त गणना करणे हे HBM मधून मोठ्या मॅट्रिक्स वाचण्यापेक्षा खूप जलद ठरते, ज्यामुळे 2× ते 4× वॉल-क्लॉक वेगवाढ मिळते आणि मेमरी वापर द्विघातीय O(N²) वरून रेषीय O(N) पर्यंत कमी होतो.