سیلف اٹینشن، CUDA کرنلز اور FlashAttention

جدید Transformers بڑی ٹوکن ترتیبوں کو ہارڈویئر کی چوٹی کی رفتار پر پروسیس کرنے کے لیے Self-Attention کی میٹرکس ریاضی، Triton جیسے اعلیٰ سطح کے GPU کمپائلرز، اور FlashAttention جیسے میموری سے باخبر CUDA کرنلز پر انحصار کرتے ہیں۔

1. وژن ٹرانسفارمر میں Self-Attention ریاضیاتی طور پر کیسے کام کرتی ہے

Vision Transformer (ViT) میں، ایک تصویر (مثلاً 224×224 پکسلز) کو 16×16 پیچز کی گرڈ میں کاٹا جاتا ہے، جس سے N = 196 پھیلائے گئے پیچز بنتے ہیں۔ ہر پیچ کو لکیری طور پر ڈائمینشن D (مثلاً 768) کے ویکٹر میں پروجیکٹ کیا جاتا ہے اور پوزیشنل ایمبیڈنگ کے ساتھ ملایا جاتا ہے، جس سے ان پٹ میٹرکس X ∈ R^(N × D) بنتا ہے۔

• Queries، Keys اور Values بنانا۔ ماڈل پیچ میٹرکس 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): اصل بصری فیچر کا مواد رکھتی ہے جو میچ ہونے پر آگے بھیجا جائے گا۔

• توجہ کے اسکور کا حساب۔ یہ ناپنے کے لیے کہ ہر پیچ دوسرے ہر پیچ سے کتنا گہرا تعلق رکھتا ہے، ماڈل تمام Queries اور Keys کے درمیان ڈاٹ پروڈکٹ نکالتا ہے:

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++ میں، آپ ایک اسکیلر تھریڈ کے نقطۂ نظر سے پروگرام کرتے ہیں۔ آپ کو تھریڈ IDs خود حساب کرنے، ارے کی حدیں جانچنے، اور ہوسٹ کی طرف سے لانچ کنفیگریشن کا بوائلر پلیٹ کوڈ لکھنا پڑتا ہے۔

#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 GPUs میں میموری کے دو اہم درجے ہوتے ہیں:

• 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 کے اندر توجہ کی ٹائلز فوراً دوبارہ حساب کرتی ہے۔ SRAM میں اضافی ریاضی کرنا HBM سے بڑے میٹرکسز پڑھنے سے کہیں تیز ثابت ہوتا ہے، جس سے 2× سے 4× تک وال کلاک رفتار ملتی ہے اور میموری کا استعمال دوگنی طاقت O(N²) سے گھٹ کر خطی O(N) رہ جاتا ہے۔