Self-attention, CUDA kernel และ FlashAttention

Transformer สมัยใหม่พึ่งพาคณิตศาสตร์เมทริกซ์ของ Self-Attention คอมไพเลอร์ GPU ระดับสูงอย่าง Triton และเคอร์เนล CUDA ที่คำนึงถึงหน่วยความจำอย่าง FlashAttention เพื่อประมวลผลลำดับโทเค็นขนาดมหาศาลด้วยความเร็วสูงสุดของฮาร์ดแวร์

1. Self-Attention ทำงานอย่างไรในเชิงคณิตศาสตร์ใน Vision Transformer

ใน Vision Transformer (ViT) ภาพ (เช่น 224×224 พิกเซล) ถูกหั่นเป็นกริดแพตช์ 16×16 ได้แพตช์ที่แผ่เรียบแล้ว N = 196 แพตช์ แต่ละแพตช์ถูกฉายเชิงเส้นเป็นเวกเตอร์ขนาดมิติ D (เช่น 768) และรวมกับ positional embedding เกิดเป็นเมทริกซ์อินพุต X ∈ R^(N × D)

• การสร้าง Query, Key และ Value โมเดลคูณเมทริกซ์แพตช์ 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): เก็บเพย์โหลดฟีเจอร์ภาพจริงที่จะถูกส่งต่อไปข้างหน้าเมื่อเกิดการจับคู่

• การคำนวณคะแนน Attention เพื่อวัดว่าแต่ละแพตช์สัมพันธ์กับทุกแพตช์อื่นมากเพียงใด โมเดลคำนวณดอตโปรดักต์ระหว่าง Query ทั้งหมดกับ Key ทั้งหมด:

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

การหารด้วย sqrt(d_(k)) ป้องกันไม่ให้ดอตโปรดักต์ค่าสูงผลักฟังก์ชัน softmax ที่ตามมาไปยังช่วงที่ gradient เล็กมาก ได้เมทริกซ์ attention ขนาด 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 พัฒนาไปอย่างไร ลองพิจารณาการบวกเวกเตอร์ทีละองค์ประกอบอย่างง่าย (C = A + B) สำหรับ N องค์ประกอบ

CUDA C++ ดิบ (การจัดการระดับเธรด)

ใน CUDA C++ คุณเขียนโปรแกรมจากมุมมองของเธรดสเกลาร์เพียงตัวเดียว คุณต้องคำนวณ ID ของเธรดเอง ตรวจสอบขอบเขตของอาร์เรย์ และเขียนโค้ดสำเร็จรูปสำหรับการตั้งค่าการเรียกใช้งานฝั่งโฮสต์

#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 มาตรฐานชนกำแพงรุนแรงเมื่อประมวลผลภาพความละเอียดสูงหรือลำดับข้อความยาว เพราะเมทริกซ์ attention S = QK^T โตขึ้นเป็นกำลังสอง (N × N) ตามความยาวลำดับ N

• คอขวดแบนด์วิดท์หน่วยความจำ (HBM เทียบกับ SRAM) GPU ของ NVIDIA สมัยใหม่มีหน่วยความจำหลักสองระดับ:

• HBM (High Bandwidth Memory / VRAM): ความจุใหญ่ (เช่น 80 GB) แต่อ่านและเขียนค่อนข้างช้า
• SRAM (Shared Memory บนชิป): เล็กมาก (ประมาณ 20 MB รวมทุก Streaming Multiprocessor) แต่เร็วกว่า HBM ประมาณ 10 ถึง 20 เท่า

ใน attention มาตรฐาน GPU คำนวณ S = QK^T เขียนเมทริกซ์ N × N ขนาดมหึมาออกไปยัง HBM ที่ช้า อ่านกลับมาเพื่อใช้ Softmax เขียนความน่าจะเป็นกลับไปยัง HBM แล้วอ่านอีกครั้งเพื่อคูณกับ V Tensor Core ใช้เวลาส่วนใหญ่ไปกับการรอการถ่ายโอนหน่วยความจำโดยไม่ได้ทำงาน

• Tiling (ไม่สร้างเมทริกซ์ N × N ขึ้นมาจริง) FlashAttention ซึ่งนำเสนอโดย Tri Dao และคณะ (2022) ใช้ kernel fusion เพื่อให้การคำนวณ attention ทั้งหมดอยู่ภายใน SRAM บนชิปที่เร็ว โดยหั่น Q, K และ V เป็นบล็อกเล็กที่พอดีกับ SRAM คำนวณ attention ทีละไทล์ และเขียนเฉพาะเมทริกซ์เอาต์พุตสุดท้ายกลับไปยัง HBM เมทริกซ์ตัวกลางขนาดมหึมา N × N จะไม่ถูกบันทึกลง VRAM เลย

• เทคนิค Online Softmax โดยปกติ Softmax ต้องรู้ค่าสูงสุดและผลรวมของเอกซ์โปเนนเชียลทั้งแถวก่อนจึงจะทำนอร์มัลไลซ์องค์ประกอบใดได้ FlashAttention ใช้อัลกอริทึม Softmax แบบออนไลน์ที่ติดตามค่าสูงสุดสะสมและตัวประกอบนอร์มัลไลซ์ทีละบล็อก และปรับสเกลเอาต์พุตบางส่วนใหม่ทางคณิตศาสตร์ทุกครั้งที่ประมวลผลไทล์ใหม่

• การคำนวณซ้ำระหว่าง Backpropagation เนื่องจาก FlashAttention ไม่บันทึกเมทริกซ์ attention N × N ระหว่าง forward pass จึงใช้การแลกเปลี่ยนที่ชาญฉลาดระหว่างการฝึก คือบันทึกเฉพาะสถิตินอร์มัลไลซ์ขนาดเล็กไว้ใน HBM และคำนวณไทล์ attention ซ้ำทันทีภายใน SRAM ที่เร็วระหว่าง backpropagation การคำนวณเพิ่มใน SRAM เร็วกว่าการอ่านเมทริกซ์ขนาดใหญ่จาก HBM มาก ให้ความเร็วเพิ่มขึ้น 2 ถึง 4 เท่าตามเวลาจริง พร้อมลดการใช้หน่วยความจำจากกำลังสอง O(N²) เป็นเชิงเส้น O(N)