Self-Attention, CUDA C++ எதிர் Triton, FlashAttention: Transformer-களின் வேக ரகசியம்

நவீன Transformer-கள் Self-Attention matrix கணிதம், Triton போன்ற உயர்மட்ட GPU compiler-கள், FlashAttention போன்ற நினைவகம்-அறிந்த CUDA kernel-கள் ஆகியவற்றைச் சார்ந்து, பெரிய token தொடர்களை அதிகபட்ச hardware வேகத்தில் செயலாக்குகின்றன.

1. Vision Transformer-இல் Self-Attention கணிதமாக எவ்வாறு வேலை செய்கிறது

ஒரு Vision Transformer-இல் (ViT), ஒரு படம் (எ.கா., 224×224 pixel-கள்) 16×16 patch-களின் கட்டமாக வெட்டப்பட்டு, N = 196 சமப்படுத்தப்பட்ட patch-களை உருவாக்குகிறது. ஒவ்வொரு patch-ம் D பரிமாண (எ.கா., 768) vector-ஆக நேரியலாக (linearly) மாற்றப்பட்டு, positional embedding-உடன் இணைக்கப்பட்டு, X ∈ ℝ^(N × D) என்ற input matrix-ஐ உருவாக்குகிறது.

• Query, Key, Value-களை உருவாக்குதல். மாதிரி, patch matrix X-ஐ மூன்று கற்றல் weight matrix-களால் (W_Q, W_K, W_V ∈ ℝ^(D × d_k)) பெருக்கி, ஒவ்வொரு பட patch-க்கும் மூன்று வெவ்வேறு பிரதிநிதித்துவங்களை (representation) உருவாக்குகிறது:

Q=XW_(Q), K=XW_(K), V=XW_(V)

• Query (Q): ஒரு குறிப்பிட்ட பட patch எதைத் தேடுகிறது என்பதைக் குறிக்கிறது (எ.கா., 'நான் ஒரு ரோம அமைப்பு (fur texture), அருகில் உள்ள கண்கள் அல்லது காதுகளைத் தேடுகிறேன்').
• Key (K): ஒரு patch என்ன தகவலைக் கொண்டுள்ளது என்பதை அறிவிக்கிறது (எ.கா., 'நான் ஒரு கூர்மையான வளைந்த விளிம்பைக் கொண்டுள்ளேன்').
• Value (V): ஒரு பொருத்தம் ஏற்பட்டால் முன்னோக்கி அனுப்பப்படும் உண்மையான பார்வை அம்சப் பொருளை (feature payload) வைத்திருக்கிறது.

• Attention மதிப்பெண்களைக் கணக்கிடுதல். ஒவ்வொரு patch-ம் மற்ற ஒவ்வொரு patch-உடனும் எவ்வளவு வலுவாகத் தொடர்புடையது என்பதை அளக்க, மாதிரி எல்லா Query-களுக்கும் Key-களுக்கும் இடையிலான dot product-ஐக் கணக்கிடுகிறது:

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

sqrt(d_(k))-ஆல் வகுப்பது, பெரிய dot product-கள் வரவிருக்கும் softmax சார்பை மிகச் சிறிய gradient-கள் உள்ள பகுதிகளுக்குத் தள்ளுவதைத் தடுக்கிறது. இது ஒரு N × N attention matrix-ஐ உருவாக்குகிறது; அதில் (i, j) உறுப்பு, patch i என்பது patch j-ஐ எவ்வளவு கவனிக்க வேண்டும் என்பதை மதிப்பிடுகிறது.

• Softmax எடையிடலும் ஒன்றிணைத்தலும். Row-வாரியாக Softmax பயன்படுத்துவது மூல மதிப்பெண்களை மொத்தம் 1 ஆகும் நிகழ்தகவுகளாக மாற்றுகிறது. இந்த நிகழ்தகவுகளை Value matrix V-ஆல் பெருக்குவது ஒவ்வொரு patch-க்கும் கலப்புப் பிரதிநிதித்துவத்தை உருவாக்குகிறது:

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

Multi-Head Self-Attention-இல், இதே செயல்முறை பல head-களில் (எ.கா., 12 head-கள்) இணையாக இயங்குகிறது; வெவ்வேறு head-கள் வெவ்வேறு இடவியல் தூரங்கள், நிறங்கள் அல்லது பொருட்களின் பகுதிகளில் ஒரே நேரத்தில் கவனம் செலுத்த அனுமதிக்கிறது.

2. நடைமுறை Code ஒப்பீடு: மூல CUDA C++ எதிர் OpenAI Triton

GPU programming எவ்வாறு பரிணமித்தது என்பதைக் காண, N உறுப்புகளுக்கான எளிய element-wise vector கூட்டலை (C = A + B) கருதுங்கள்.

மூல CUDA C++ (Thread-மட்ட நிர்வாகம்)

CUDA C++-இல், ஒரு தனிப்பட்ட scalar thread-இன் கண்ணோட்டத்திலிருந்து நிரலாக்குகிறீர்கள். Thread ID-களைக் கைமுறையாகக் கணக்கிட்டு, array எல்லைகளைச் சரிபார்த்து, 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 (Block-மட்ட Python Programming)

Triton-இல், நேரடியாக Python-இல் ஒரு block (tile) தரவின் கண்ணோட்டத்திலிருந்து நிரலாக்குகிறீர்கள். Triton-இன் compiler thread synchronization, memory coalescing, vectorization ஆகியவற்றை உட்புறமாகத் தானாகவே கையாளுகிறது.

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 நினைவக நெரிசலைத் தீர்த்தல்

உயர்-தெளிவுத்திறன் படங்கள் அல்லது நீளமான உரைத் தொடர்களைச் செயலாக்கும்போது, attention matrix S = QK^T தொடர் நீளம் N-உடன் இருபடியாக (N × N) வளர்வதால், சாதாரண Self-Attention கடுமையான சுவரை எதிர்நோக்குகிறது.

• நினைவக Bandwidth நெரிசல் (HBM எதிர் SRAM). நவீன NVIDIA GPU-களில் இரு பிரதான நினைவக மட்டங்கள் உள்ளன:

• HBM (High Bandwidth Memory / VRAM): பெரியது (எ.கா., 80 GB), ஆனால் வாசிக்கவும் எழுதவும் ஒப்பீட்டளவில் மெதுவானது.
• SRAM (On-Chip Shared Memory): மிகச் சிறியது (எல்லா Streaming Multiprocessor-களிலும் மொத்தம் ~20 MB மட்டுமே), ஆனால் HBM-ஐ விட சுமார் 10× முதல் 20× மடங்கு வேகமானது.

சாதாரண attention-இல், GPU S = QK^T-ஐக் கணக்கிட்டு, பெரிய N × N matrix-ஐ மெதுவான HBM-இல் எழுதி, Softmax பயன்படுத்த மீண்டும் வாசித்து, நிகழ்தகவுகளை மீண்டும் HBM-இல் எழுதி, V-ஆல் பெருக்க மீண்டும் வாசிக்கிறது. Tensor Cores தமது பெரும்பாலான நேரத்தை நினைவகப் பரிமாற்றங்களுக்காக வெற்றே காத்திருக்கிறது.

• Tiling (N × N Matrix-ஐ ஒருபோதும் உருவாக்காமை). Tri Dao மற்றும் குழுவினர் (2022) அறிமுகப்படுத்திய FlashAttention, முழு attention கணக்கீட்டையும் வேகமான on-chip SRAM-க்குள் வைக்க kernel fusion-ஐப் பயன்படுத்துகிறது. அது Q, K, V-ஐ SRAM-க்குள் பொருந்தும் சிறிய block-களாக வெட்டி, tile-by-tile attention-ஐக் கணக்கிட்டு, இறுதி output matrix-ஐ மட்டுமே HBM-இல் எழுதுகிறது. பெரிய N × N இடைநிலை matrix ஒருபோதும் VRAM-இல் சேமிக்கப்படுவதில்லை.

• Online Softmax தந்திரம். வழக்கமாக, Softmax எந்த ஒரு உறுப்பையும் சீராக்குவதற்கு (normalize) முன் முழு row-இன் அதிகபட்ச மதிப்பையும் exponential-களின் கூட்டுத்தொகையையும் அறிந்திருக்க வேண்டும். FlashAttention, block-by-block ஓடும் அதிகபட்சங்களையும் normalization காரணிகளையும் கண்காணித்து, ஒவ்வொரு புதிய tile செயலாக்கப்படும்போதும் பகுதி output-களைக் கணிதமாக மீளளவிடும் (rescale) ஒரு online Softmax algorithm-ஐப் பயன்படுத்துகிறது.

• Backpropagation-இன் போது மறுகணக்கீடு. FlashAttention forward pass-இன் போது N × N attention matrix-ஐச் சேமிக்காததால், training-இன் போது ஒரு புத்திசாலித்தனமான சமரசத்தைப் பயன்படுத்துகிறது: HBM-இல் சிறிய normalization புள்ளியியல்களை மட்டுமே சேமித்து, backpropagation-இன் போது வேகமான SRAM-க்குள் attention tile-களை உடனுக்குடன் மறுகணக்கிடுகிறது. SRAM-இல் கூடுதல் கணிதம் செய்வது HBM-இலிருந்து பெரிய matrix-களை வாசிப்பதை விட மிக வேகமானது என்பது நிரூபணமாகியது; இது நினைவகப் பயன்பாட்டை இருபடி O(N²)-இலிருந்து நேரியல் O(N)-ஆகக் குறைத்து, 2× முதல் 4× wall-clock வேகத்தைத் தருகிறது.