Self-attention, CUDA kernels සහ FlashAttention

නූතන Transformers විශාල ටෝකන අනුක්‍රම උපරිම දෘඪාංග වේගයෙන් සැකසීමට Self-Attention න්‍යාස ගණිතය, Triton වැනි ඉහළ මට්ටමේ GPU සම්පාදක, සහ FlashAttention වැනි මතකය පිළිබඳ සවිඥානික CUDA kernels මත රඳා පවතී.

1. Vision Transformer එකක Self-Attention ගණිතමය වශයෙන් ක්‍රියා කරන ආකාරය

Vision Transformer (ViT) එකක, රූපයක් (උදා: පික්සල 224×224) 16×16 පැච් වල ජාලකයකට කපා, පැතලි කළ පැච් N = 196 ක් නිපදවයි. සෑම පැච් එකක්ම රේඛීයව D (උදා: 768) මානයේ දෛශිකයකට ප්‍රක්ෂේපණය කර positional embedding එකක් සමඟ ඒකාබද්ධ කර, ආදාන න්‍යාසය 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 ශ්‍රිතය කුඩා gradients ඇති ප්‍රදේශවලට තල්ලු කිරීම වැළකේ. මෙය 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 අතින් ගණනය කළ යුතු අතර, අරා සීමා පරීක්ෂා කළ යුතු අතර, host පැත්තේ launch වින්‍යාස සඳහා boilerplate කේතය ලිවිය යුතුය.

#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 හි සම්පාදකය පසුබිමේදී නූල් සමමුහුර්ත කිරීම, memory coalescing සහ දෛශිකකරණය ස්වයංක්‍රීයව හසුරුවයි.

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 (චිපය මත Shared Memory): කුඩා (සියලු Streaming Multiprocessors හරහා මුළු ~20 MB) නමුත් HBM ට වඩා ආසන්න වශයෙන් 10× සිට 20× දක්වා වේගවත්.

සම්මත අවධානයේදී, GPU S = QK^T ගණනය කර, විශාල N × N න්‍යාසය මන්දගාමී HBM වෙත ලියා, Softmax යෙදීමට නැවත කියවා, සම්භාවිතා නැවත HBM වෙත ලියා, V ගෙන් ගුණ කිරීමට නැවත කියවයි. Tensor Cores ඔවුන්ගේ බොහෝ කාලය මතක මාරුවීම් සඳහා නිෂ්ක්‍රීයව බලා සිටීමෙන් ගත කරයි.

• Tiling (N × N න්‍යාසය කිසිවිටෙක ද්‍රව්‍යමය නොකිරීම). Tri Dao සහ අනෙකුත් අය (2022) හඳුන්වා දුන් FlashAttention, සම්පූර්ණ අවධාන ගණනය වේගවත් චිපය මත SRAM තුළ තබාගැනීමට kernel ඒකාබද්ධකරණය (kernel fusion) භාවිත කරයි. එය Q, K සහ V SRAM තුළ ගැලපෙන කුඩා කුට්ටිවලට පෙත්තක් කර, අවධානය ටයිල්-ටයිල් ගණනය කර, අවසන් ප්‍රතිදාන න්‍යාසය පමණක් HBM වෙත ලියයි. විශාල N × N අතරමැදි න්‍යාසය කිසිවිටෙක VRAM වෙත සුරකිනු නොලැබේ.

• Online Softmax උපක්‍රමය. සාමාන්‍යයෙන්, Softmax හට තනි මූලද්‍රව්‍යයක් සාමාන්‍යකරණය කිරීමට පෙර මුළු පේළිය පුරා උපරිම අගය සහ ඝාතකවල එකතුව දැනගැනීම අවශ්‍ය වේ. FlashAttention online Softmax ඇල්ගොරිතමයක් භාවිත කරන අතර, එය කුට්ටියෙන්-කුට්ටියට ධාවන උපරිම අගයන් සහ සාමාන්‍යකරණ සාධක ලුහුබැඳ, නව ටයිල් එකක් සැකසෙන විට අර්ධ ප්‍රතිදාන ගණිතමය ලෙස නැවත පරිමාණය කරයි.

• Backpropagation අතරතුර නැවත ගණනය කිරීම. FlashAttention ඉදිරි පාස් එකේදී N × N අවධාන න්‍යාසය සුරකින්නේ නැති නිසා, පුහුණුවේදී එය දක්ෂ සම්මුතියක් භාවිත කරයි: HBM හි සුරකින්නේ කුඩා සාමාන්‍යකරණ සංඛ්‍යාලේඛන පමණක් වන අතර, backpropagation අතරතුර වේගවත් SRAM තුළ අවධාන ටයිල් එකවර නැවත ගණනය කරයි. SRAM තුළ අමතර ගණිතය කිරීම HBM වෙතින් විශාල න්‍යාස කියවීමට වඩා බොහෝ වේගවත් බව ඔප්පු වන අතර, මතක භාවිතය වර්ගාත්මක O(N²) සිට රේඛීය O(N) දක්වා හැකිලෙමින් 2× සිට 4× දක්වා wall-clock වේග වැඩිවීම් ලබාදේ.