الانتباه الذاتي ونوى CUDA وFlashAttention

تعتمد محولات Transformer الحديثة على حسابات مصفوفات الانتباه الذاتي ومترجمات 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: حل اختناق ذاكرة Transformer

يصطدم الانتباه الذاتي القياسي بعائق شديد عند معالجة صور عالية الدقة أو تسلسلات نصية طويلة، لأن مصفوفة الانتباه S = QK^T تنمو تربيعيا (N × N) مع طول التسلسل N.

• اختناق عرض نطاق الذاكرة (HBM مقابل SRAM). تضم وحدات NVIDIA GPU الحديثة مستويين رئيسيين من الذاكرة:

• HBM (ذاكرة ذات عرض نطاق عال / VRAM): كبيرة (مثلا 80 GB)، لكن القراءة منها والكتابة إليها بطيئتان نسبيا.
• SRAM (الذاكرة المشتركة على الشريحة): صغيرة جدا (~20 MB إجمالا عبر جميع Streaming Multiprocessors)، لكنها أسرع من HBM بنحو 10× إلى 20×.

في الانتباه القياسي، تحسب 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، مما يحقق تسريعا زمنيا فعليا بمقدار 2× إلى 4× مع تقليص استخدام الذاكرة من التربيعي O(N²) إلى الخطي O(N).