خودتوجهی، کرنل‌های CUDA و FlashAttention

ترنسفورمرهای امروزی به محاسبات ماتریسی خودتوجهی، کامپایلرهای سطح بالای GPU مانند Triton و کرنل‌های CUDA آگاه از حافظه مانند FlashAttention متکی‌اند تا دنباله‌های عظیم توکن را با بیشترین سرعت سخت‌افزار پردازش کنند.

1. خودتوجهی از نظر ریاضی در یک ترنسفورمر بینایی چگونه کار می‌کند

در ترنسفورمر بینایی (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)

• پرس‌وجو (Q): نشان می‌دهد یک تکه خاص تصویر به دنبال چیست (مثلا «من بافت خز هستم و دنبال چشم یا گوش در نزدیکی می‌گردم»).
• کلید (K): اعلام می‌کند یک تکه چه اطلاعاتی دارد (مثلا «من حاوی یک لبه منحنی تیز هستم»).
• مقدار (V): داده واقعی ویژگی بصری را نگه می‌دارد که در صورت تطابق به جلو فرستاده می‌شود.

• محاسبه امتیازهای توجه. مدل برای سنجش شدت ارتباط هر تکه با تمام تکه‌های دیگر، ضرب داخلی میان همه پرس‌وجوها و کلیدها را محاسبه می‌کند:

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

تقسیم بر sqrt(d_(k)) مانع می‌شود ضرب‌های داخلی بزرگ، تابع softmax بعدی را به ناحیه‌هایی با گرادیان‌های بسیار کوچک ببرند. حاصل یک ماتریس توجه N × N است که درایه (i, j) آن مشخص می‌کند تکه i باید چقدر به تکه j توجه کند.

• وزن‌دهی و تجمیع با Softmax. اعمال Softmax به‌صورت سطری، امتیازهای خام را به احتمال‌هایی تبدیل می‌کند که مجموعشان 1 است. ضرب این احتمال‌ها در ماتریس مقدار V برای هر تکه یک بازنمایی ترکیبی تولید می‌کند:

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

در خودتوجهی چندسر، همین فرایند به‌صورت موازی در چندین سر (مثلا 12 سر) اجرا می‌شود و به سرهای مختلف اجازه می‌دهد هم‌زمان بر فاصله‌های فضایی، رنگ‌ها یا بخش‌های متفاوت اشیا تمرکز کنند.

2. مقایسه عملی کد: CUDA C++ خام در برابر OpenAI Triton

برای دیدن تکامل برنامه‌نویسی GPU، جمع ساده عنصر‌به‌عنصر دو بردار (C = A + B) با N عنصر را در نظر بگیرید.

CUDA C++ خام (مدیریت در سطح رشته)

در CUDA C++ از دید یک رشته اسکالر منفرد برنامه می‌نویسید. باید شناسه رشته‌ها را دستی محاسبه کنید، حدود آرایه را بررسی کنید و پیکربندی‌های تکراری راه‌اندازی در سمت میزبان را بنویسید.

#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: حل گلوگاه حافظه ترنسفورمر

خودتوجهی استاندارد هنگام پردازش تصاویر با وضوح بالا یا دنباله‌های طولانی متن به مانعی جدی می‌رسد، زیرا ماتریس توجه S = QK^T با طول دنباله N به‌صورت درجه دوم (N × N) رشد می‌کند.

• گلوگاه پهنای باند حافظه (HBM در برابر SRAM). ‏GPUهای امروزی NVIDIA دو رده اصلی حافظه دارند:

• HBM (حافظه با پهنای باند بالا / VRAM): بزرگ (مثلا 80 GB)، اما خواندن و نوشتن آن نسبتا کند است.
• SRAM (حافظه اشتراکی روی تراشه): بسیار کوچک (در مجموع ~20 MB در تمام Streaming Multiprocessorها)، اما تقریبا 10× تا 20× سریع‌تر از HBM است.

در توجه استاندارد، GPU ابتدا S = QK^T را محاسبه می‌کند، ماتریس عظیم N × N را در HBM کند می‌نویسد، آن را برای اعمال Softmax دوباره می‌خواند، احتمال‌ها را باز در HBM می‌نویسد و برای ضرب در V دوباره می‌خواند. Tensor Cores بیشتر زمانشان را بی‌کار و منتظر انتقال حافظه می‌گذرانند.

• کاشی‌بندی (بدون ساختن کامل ماتریس N × N). ‏FlashAttention که Tri Dao و همکاران (2022) معرفی کردند، از ادغام کرنل‌ها استفاده می‌کند تا کل محاسبه توجه در SRAM سریع روی تراشه بماند. Q و K و V را به بلوک‌های کوچکی که در SRAM جا می‌گیرند تقسیم می‌کند، توجه را کاشی‌به‌کاشی محاسبه می‌کند و فقط ماتریس خروجی نهایی را در HBM می‌نویسد. ماتریس میانی عظیم N × N هرگز در VRAM ذخیره نمی‌شود.

• ترفند Softmax برخط. در حالت عادی، Softmax پیش از نرمال‌سازی هر عنصر باید بیشینه مقدار و مجموع مقادیر نمایی کل یک سطر را بداند. FlashAttention از الگوریتم Softmax برخط استفاده می‌کند که بیشینه‌های جاری و ضریب‌های نرمال‌سازی را بلوک‌به‌بلوک دنبال می‌کند و با پردازش هر کاشی جدید، مقیاس خروجی‌های جزئی را به‌صورت ریاضی تنظیم می‌کند.

• محاسبه مجدد در پس‌انتشار. چون FlashAttention در گذر رو به جلو ماتریس توجه N × N را ذخیره نمی‌کند، هنگام آموزش از موازنه‌ای هوشمندانه استفاده می‌کند: فقط آمار کوچک نرمال‌سازی را در HBM نگه می‌دارد و هنگام پس‌انتشار، کاشی‌های توجه را همان موقع در SRAM سریع دوباره محاسبه می‌کند. انجام محاسبات بیشتر در SRAM بسیار سریع‌تر از خواندن ماتریس‌های بزرگ از HBM از آب درمی‌آید و در عین کاهش مصرف حافظه از درجه دوم O(N²) به خطی O(N)، زمان واقعی اجرا را 2× تا 4× سریع‌تر می‌کند.