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) ఏర్పడుతుంది.

• క్వెరీలు, కీలు, విలువల సృష్టి. ప్రతి చిత్ర ప్యాచ్‌కు మూడు వేర్వేరు ప్రాతినిధ్యాలను సృష్టించడానికి మోడల్, ప్యాచ్ మ్యాట్రిక్స్ 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): సరిపోలిక దొరికితే ముందుకు పంపే అసలు దృశ్య ఫీచర్ డేటాను ఉంచుతుంది.

• అటెన్షన్ స్కోర్ల గణన. ప్రతి ప్యాచ్‌కు ప్రతి ఇతర ప్యాచ్‌తో సంబంధం ఎంత బలంగా ఉందో కొలవడానికి, మోడల్ అన్ని క్వెరీలు, కీల మధ్య డాట్ ప్రొడక్ట్‌ను లెక్కిస్తుంది:

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

Multi-Head 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 (అధిక బ్యాండ్‌విడ్త్ మెమరీ / 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లో అప్పటికప్పుడే అటెన్షన్ టైల్స్‌ను తిరిగి లెక్కిస్తుంది. HBM నుంచి పెద్ద మ్యాట్రిక్స్‌లను చదవడం కంటే SRAMలో అదనపు గణన చేయడం చాలా వేగవంతమని తేలింది. దీనివల్ల వాస్తవ గడిచిన సమయ పరంగా 2× నుంచి 4× వేగవృద్ధి లభిస్తూ, మెమరీ వినియోగం వర్గానుపాత O(N²) నుంచి రేఖీయ O(N)కు తగ్గుతుంది.