Umakini wa kujitegemea, kerneli za CUDA na FlashAttention

Transformer za kisasa hutegemea hesabu za matriki za Umakini wa Kujitegemea, vikompaila vya GPU vya kiwango cha juu kama Triton, na kerneli za CUDA zinazozingatia kumbukumbu kama FlashAttention kuchakata mifuatano mikubwa sana ya tokeni kwa kasi ya juu kabisa ya maunzi.

1. Jinsi Umakini wa Kujitegemea Unavyofanya Kazi Kihesabu katika Vision Transformer

Katika Vision Transformer (ViT), picha (k.m., pikseli 224×224) hukatwa katika gridi ya vipande vya 16×16, ikitoa N = 196 vipande vilivyosawazishwa. Kila kipande hubadilishwa kwa makadirio ya mstari kuwa vekta yenye kipimo D (k.m., 768) na kuunganishwa na upachikaji wa nafasi, na kuunda matriki ya ingizo X ∈ R^(N × D).

• Kuunda Hoja, Funguo, na Thamani. Modeli huzidisha matriki ya vipande X kwa matriki tatu za uzani zilizojifunzwa (W_Q, W_K, W_V ∈ R^(D × d_k)) ili kuunda viwakilishi vitatu tofauti kwa kila kipande cha picha:

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

• Hoja (Q): Huwakilisha kile kipande mahususi cha picha kinachotafuta (k.m., "Mimi ni umbile la manyoya, natafuta macho au masikio karibu").
• Ufunguo (K): Hutangaza taarifa zilizomo katika kipande (k.m., "Nina ukingo mkali uliopinda").
• Thamani (V): Hushikilia data halisi ya sifa za picha itakayopitishwa mbele iwapo ulinganifu utapatikana.

• Kukokotoa Alama za Umakini. Kupima nguvu ya uhusiano wa kila kipande na kila kipande kingine, modeli hukokotoa zao la nukta kati ya Hoja zote na Funguo zote:

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

Kugawanya kwa sqrt(d_(k)) huzuia mazao makubwa ya nukta kusukuma kitendakazi cha softmax kinachofuata katika maeneo yenye gradienti ndogo sana. Hili hutoa matriki ya umakini ya N × N ambapo ingizo (i, j) hutoa alama ya kiasi ambacho kipande i kinapaswa kuzingatia kipande j.

• Uwekaji wa Uzani wa Softmax na Ukusanyaji. Kutumia Softmax kwa kila safu mlalo hubadilisha alama ghafi kuwa uwezekano ambao jumla yake ni 1. Kuzidisha uwezekano huu kwa matriki ya Thamani V hutoa uwakilishi uliochanganywa kwa kila kipande:

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

Katika Umakini wa Kujitegemea wa Vichwa Vingi, mchakato huuhuu huendeshwa sambamba katika vichwa vingi (k.m., vichwa 12), na kuruhusu vichwa tofauti kuzingatia umbali tofauti wa kieneo, rangi, au sehemu za vitu kwa wakati mmoja.

2. Ulinganisho wa Vitendo wa Msimbo: CUDA C++ ya Moja kwa Moja dhidi ya OpenAI Triton

Ili kuona jinsi upangaji wa GPU ulivyobadilika, fikiria uongezaji rahisi wa vekta kipengele kwa kipengele (C = A + B) kwa vipengele N.

CUDA C++ ya Moja kwa Moja (Usimamizi wa Kiwango cha Uzi)

Katika CUDA C++, unaandika programu kwa mtazamo wa uzi mmoja wa skala. Lazima ukokotoe mwenyewe vitambulisho vya nyuzi, ukague mipaka ya safu, na uandike mipangilio ya kawaida ya kuanzisha utekelezaji upande wa mwenyeji.

#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 (Upangaji wa Python wa Kiwango cha Kikundi)

Katika Triton, unaandika programu kwa mtazamo wa kikundi (kigae) cha data moja kwa moja kwa Python. Kikompaila cha Triton hushughulikia kiotomatiki usawazishaji wa nyuzi, kuunganisha ufikiaji wa kumbukumbu, na ubadilishaji kuwa operesheni za vekta ndani kwa ndani.

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: Kutatua Kikwazo cha Kumbukumbu cha Transformer

Umakini wa Kujitegemea wa kawaida hukumbana na kikomo kikubwa unapochakata picha zenye msongo wa juu au mifuatano mirefu ya maandishi kwa sababu matriki ya umakini S = QK^T hukua kwa mraba (N × N) kulingana na urefu wa mfuatano N.

• Kikwazo cha Upana wa Bendi wa Kumbukumbu (HBM dhidi ya SRAM). GPU za kisasa za NVIDIA zina viwango viwili vikuu vya kumbukumbu:

• HBM (Kumbukumbu Yenye Upana Mkubwa wa Bendi / VRAM): Kubwa (k.m., GB 80) lakini kusoma na kuandika ni kwa polepole kiasi.
• SRAM (Kumbukumbu Shiriki Iliyo Kwenye Chipu): Ndogo sana (~MB 20 kwa jumla katika Streaming Multiprocessor zote) lakini ina kasi ya takriban 10× hadi 20× kuliko HBM.

Katika umakini wa kawaida, GPU hukokotoa S = QK^T, huandika matriki kubwa sana ya N × N kwenye HBM ya polepole, huisoma tena kutumia Softmax, huandika uwezekano kurudi kwenye HBM, na kuusoma tena ili kuzidisha kwa V. Tensor Cores hutumia muda wao mwingi bila kufanya kazi zikisubiri uhamishaji wa kumbukumbu.

• Ugawaji katika Vigae (Bila Kuwahi Kuunda Matriki ya N × N katika Kumbukumbu). FlashAttention, iliyoanzishwa na Tri Dao na wenzake (2022), hutumia uunganishaji wa kerneli kuweka ukokotoaji wote wa umakini ndani ya SRAM ya haraka iliyo kwenye chipu. Hugawanya Q, K, na V katika vikundi vidogo vinavyotoshea ndani ya SRAM, hukokotoa umakini kigae kwa kigae, na huandika matriki ya matokeo ya mwisho pekee kurudi kwenye HBM. Matriki kubwa sana ya kati ya N × N haihifadhiwi kamwe kwenye VRAM.

• Mbinu ya Softmax ya Hatua kwa Hatua. Kwa kawaida, Softmax inahitaji kujua thamani ya juu kabisa na jumla ya thamani za kipeo katika safu mlalo nzima kabla ya kusanifisha kipengele chochote kimoja. FlashAttention hutumia algoriti ya Softmax ya hatua kwa hatua inayofuatilia thamani za juu na vigawo vya usanifishaji vinavyoendelea kikundi kwa kikundi, ikibadilisha vipimo vya matokeo ya sehemu kihesabu kila kigae kipya kinapochakatwa.

• Kukokotoa Tena Wakati wa Uenezaji wa Nyuma. Kwa kuwa FlashAttention haihifadhi matriki ya umakini ya N × N wakati wa upitishaji wa mbele, hutumia njia ya kubadilishana gharama kwa ujanja wakati wa mafunzo: huhifadhi takwimu ndogo tu za usanifishaji kwenye HBM na kukokotoa tena vigae vya umakini papo hapo ndani ya SRAM ya haraka wakati wa uenezaji wa nyuma. Kufanya hesabu za ziada kwenye SRAM huwa kwa kasi zaidi kuliko kusoma matriki kubwa kutoka HBM, na kutoa ongezeko la kasi la muda halisi la 2× hadi 4× huku ikipunguza matumizi ya kumbukumbu kutoka ya mraba O(N²) hadi ya mstari O(N).