AI & Teknologi

Self-attention, kernel CUDA, dan FlashAttention

2026-11-15

Blok memori mikrokontroler sebagai ilustrasi, bukan GPU yang dibahas dalam artikel.

Transformer modern bergantung pada matematika matriks Self-Attention, compiler GPU tingkat tinggi seperti Triton, dan kernel CUDA yang sadar memori seperti FlashAttention untuk memproses urutan token yang sangat besar pada kecepatan puncak perangkat keras.

1. Cara Kerja Self-Attention secara Matematis dalam Vision Transformer

Dalam Vision Transformer (ViT), sebuah gambar (mis. 224×224 piksel) dipotong menjadi grid patch 16×16, menghasilkan N = 196 patch yang diratakan. Setiap patch diproyeksikan secara linear menjadi vektor berdimensi D (mis. 768) dan digabung dengan positional embedding, membentuk matriks masukan X ∈ R^(N × D).

Membuat Query, Key, dan Value. Model mengalikan matriks patch X dengan tiga matriks bobot hasil pembelajaran (W_Q, W_K, W_V ∈ R^(D × d_k)) untuk membuat tiga representasi berbeda bagi setiap patch gambar:

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

Query (Q): Mewakili apa yang sedang dicari oleh sebuah patch gambar tertentu (mis. "Saya tekstur bulu, mencari mata atau telinga di dekat saya").
Key (K): Mengumumkan informasi apa yang dimiliki sebuah patch (mis. "Saya memuat tepi melengkung yang tajam").
Value (V): Menyimpan muatan fitur visual yang sebenarnya, yang akan diteruskan jika terjadi kecocokan.

Menghitung Skor Attention. Untuk mengukur seberapa kuat setiap patch berhubungan dengan setiap patch lainnya, model menghitung dot product antara semua Query dan Key:

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

Pembagian dengan sqrt(d_(k)) mencegah dot product yang besar mendorong fungsi softmax berikutnya ke wilayah dengan gradien sangat kecil. Ini menghasilkan matriks attention N × N, di mana entri (i, j) menilai seberapa besar patch i harus memperhatikan patch j.

Pembobotan Softmax dan Agregasi. Menerapkan Softmax per baris mengubah skor mentah menjadi probabilitas yang berjumlah 1. Mengalikan probabilitas ini dengan matriks Value V menghasilkan representasi campuran untuk setiap patch:

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

Pada Multi-Head Self-Attention, proses yang sama persis ini berjalan paralel di beberapa head (mis. 12 head), sehingga head yang berbeda dapat berfokus pada jarak spasial, warna, atau bagian objek yang berbeda secara bersamaan.

2. Perbandingan Kode Langsung: CUDA C++ Mentah vs. OpenAI Triton

Untuk melihat bagaimana pemrograman GPU berkembang, perhatikan penjumlahan vektor elemen demi elemen sederhana (C = A + B) untuk N elemen.

CUDA C++ Mentah (Pengelolaan Tingkat Thread)

Di CUDA C++, Anda memprogram dari sudut pandang satu thread skalar. Anda harus menghitung ID thread secara manual, memeriksa batas array, dan menulis konfigurasi peluncuran sisi host yang berulang-ulang.

#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 (Pemrograman Python Tingkat Blok)

Di Triton, Anda memprogram dari sudut pandang satu blok (tile) data langsung di Python. Compiler Triton secara otomatis menangani sinkronisasi thread, memory coalescing, dan vektorisasi di balik layar.

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: Mengatasi Bottleneck Memori Transformer

Self-Attention standar membentur tembok yang berat saat memproses gambar beresolusi tinggi atau urutan teks panjang karena matriks attention S = QK^T tumbuh secara kuadratik (N × N) terhadap panjang urutan N.

Bottleneck Bandwidth Memori (HBM vs. SRAM). GPU NVIDIA modern memiliki dua tingkat memori utama:

HBM (High Bandwidth Memory / VRAM): Besar (mis. 80 GB) tetapi relatif lambat untuk dibaca dan ditulis.
SRAM (Shared Memory On-Chip): Sangat kecil (~20 MB total di seluruh Streaming Multiprocessor) tetapi sekitar 10× hingga 20× lebih cepat daripada HBM.

Pada attention standar, GPU menghitung S = QK^T, menulis matriks N × N yang masif ke HBM yang lambat, membacanya kembali untuk menerapkan Softmax, menulis probabilitas kembali ke HBM, lalu membacanya lagi untuk dikalikan dengan V. Tensor Core menghabiskan sebagian besar waktunya menganggur menunggu transfer memori.

Tiling (Tidak Pernah Mematerialisasikan Matriks N × N). Diperkenalkan oleh Tri Dao dkk. (2022), FlashAttention memakai kernel fusion untuk menjaga seluruh perhitungan attention di dalam SRAM on-chip yang cepat. Ia memotong Q, K, dan V menjadi blok kecil yang muat di SRAM, menghitung attention tile demi tile, dan hanya menulis matriks keluaran akhir kembali ke HBM. Matriks perantara N × N yang besar tidak pernah disimpan ke VRAM.

Trik Online Softmax. Biasanya, Softmax membutuhkan nilai maksimum dan jumlah eksponensial di seluruh baris sebelum menormalisasi satu elemen pun. FlashAttention memakai algoritma online Softmax yang melacak nilai maksimum berjalan dan faktor normalisasi blok demi blok, menskalakan ulang keluaran parsial secara matematis saat setiap tile baru diproses.

Komputasi Ulang Saat Backpropagation. Karena FlashAttention tidak menyimpan matriks attention N × N selama forward pass, ia memakai kompromi cerdas saat pelatihan: ia hanya menyimpan statistik normalisasi kecil di HBM dan menghitung ulang tile attention secara langsung di dalam SRAM yang cepat saat backpropagation. Melakukan matematika tambahan di SRAM terbukti jauh lebih cepat daripada membaca matriks besar dari HBM, menghasilkan percepatan waktu nyata 2× hingga 4× sambil menyusutkan penggunaan memori dari kuadratik O(N²) menjadi linear O(N).

Berita lainnya: AI & Teknologi