Cơ chế tự chú ý, kernel CUDA và FlashAttention
Các Transformer hiện đại dựa vào phép toán ma trận tự chú ý, các trình biên dịch GPU cấp cao như Triton và các kernel CUDA có tính đến bộ nhớ như FlashAttention để xử lý những chuỗi token khổng lồ ở tốc độ phần cứng tối đa.
1. Cơ chế tự chú ý hoạt động về mặt toán học trong Vision Transformer như thế nào
Trong Vision Transformer (ViT), một ảnh (ví dụ: 224×224 điểm ảnh) được cắt thành lưới các mảnh 16×16, tạo ra N = 196 mảnh đã làm phẳng. Mỗi mảnh được chiếu tuyến tính thành một vectơ có số chiều D (ví dụ: 768) và kết hợp với một embedding vị trí, tạo thành ma trận đầu vào X ∈ R^(N × D).
- Tạo truy vấn, khóa và giá trị. Mô hình nhân ma trận mảnh ảnh X với ba ma trận trọng số đã học (W_Q, W_K, W_V ∈ R^(D × d_k)) để tạo ba biểu diễn khác nhau cho mỗi mảnh ảnh:
Q=XW_(Q), K=XW_(K), V=XW_(V)
- Truy vấn (Q): Biểu thị điều một mảnh ảnh cụ thể đang tìm kiếm (ví dụ: "Tôi là kết cấu lông, đang tìm mắt hoặc tai ở gần").
- Khóa (K): Cho biết một mảnh chứa thông tin gì (ví dụ: "Tôi chứa một cạnh cong sắc nét").
- Giá trị (V): Chứa dữ liệu đặc trưng thị giác thực tế sẽ được truyền tiếp nếu có sự phù hợp.
- Tính điểm chú ý. Để đo mức độ liên quan giữa mỗi mảnh với mọi mảnh khác, mô hình tính tích vô hướng giữa tất cả truy vấn và khóa:
Scores=(QK^(T))/(sqrt(d_(k)))
Chia cho sqrt(d_(k)) ngăn các tích vô hướng lớn đẩy hàm softmax ở bước sau vào những vùng có gradient rất nhỏ. Việc này tạo ra ma trận chú ý N × N, trong đó phần tử (i, j) chấm điểm mức độ mảnh i nên chú ý tới mảnh j.
- Gán trọng số Softmax và tổng hợp. Áp dụng Softmax theo từng hàng chuyển các điểm số thô thành xác suất có tổng bằng 1. Nhân các xác suất này với ma trận giá trị V tạo ra một biểu diễn pha trộn cho mỗi mảnh:
Attention(Q,K,V)=softmax((QK^(T))/(sqrt(d_(k))))V
Trong cơ chế tự chú ý đa đầu, chính quy trình này chạy song song trên nhiều đầu (ví dụ: 12 đầu), cho phép các đầu khác nhau đồng thời tập trung vào các khoảng cách không gian, màu sắc hoặc bộ phận vật thể khác nhau.
2. So sánh mã thực hành: CUDA C++ thuần và OpenAI Triton
Để thấy lập trình GPU đã phát triển như thế nào, hãy xét phép cộng vectơ đơn giản theo từng phần tử (C = A + B) với N phần tử.
CUDA C++ thuần (quản lý ở cấp luồng)
Trong CUDA C++, bạn lập trình từ góc nhìn của một luồng vô hướng đơn lẻ. Bạn phải tự tính ID luồng, kiểm tra giới hạn mảng và viết phần mã khung cấu hình khởi chạy phía máy chủ.
#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 (lập trình Python ở cấp khối)
Trong Triton, bạn lập trình từ góc nhìn của một khối (ô) dữ liệu trực tiếp bằng Python. Trình biên dịch của Triton tự động xử lý việc đồng bộ luồng, gộp truy cập bộ nhớ và vectơ hóa ở bên trong.
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: Giải quyết nút thắt bộ nhớ của Transformer
Cơ chế tự chú ý tiêu chuẩn gặp trở ngại nghiêm trọng khi xử lý ảnh độ phân giải cao hoặc chuỗi văn bản dài vì ma trận chú ý S = QK^T tăng theo bậc hai (N × N) với độ dài chuỗi N.
- Nút thắt băng thông bộ nhớ (HBM và SRAM). GPU NVIDIA hiện đại có hai tầng bộ nhớ chính:
- HBM (bộ nhớ băng thông cao / VRAM): Lớn (ví dụ: 80 GB) nhưng đọc và ghi tương đối chậm.
- SRAM (bộ nhớ dùng chung trên chip): Rất nhỏ (tổng cộng ~20 MB trên tất cả các Streaming Multiprocessor) nhưng nhanh hơn HBM khoảng 10× đến 20×.
Trong cơ chế chú ý tiêu chuẩn, GPU tính S = QK^T, ghi ma trận N × N khổng lồ ra HBM chậm, đọc lại để áp dụng Softmax, ghi các xác suất trở lại HBM rồi đọc chúng lần nữa để nhân với V. Tensor Cores dành phần lớn thời gian nhàn rỗi để chờ truyền dữ liệu bộ nhớ.
- Chia ô (không bao giờ tạo tường minh ma trận N × N). Do Tri Dao và cộng sự giới thiệu (2022), FlashAttention dùng hợp nhất kernel để giữ toàn bộ phép tính chú ý bên trong SRAM nhanh trên chip. Nó cắt Q, K và V thành các khối nhỏ vừa với SRAM, tính chú ý theo từng ô và chỉ ghi ma trận đầu ra cuối cùng trở lại HBM. Ma trận trung gian N × N khổng lồ không bao giờ được lưu vào VRAM.
- Thủ thuật Softmax trực tuyến. Thông thường, Softmax cần biết giá trị lớn nhất và tổng các hàm mũ trên toàn bộ một hàng trước khi chuẩn hóa bất kỳ phần tử đơn lẻ nào. FlashAttention dùng thuật toán Softmax trực tuyến để theo dõi các giá trị lớn nhất hiện tại và hệ số chuẩn hóa theo từng khối, điều chỉnh tỷ lệ các đầu ra từng phần bằng toán học khi mỗi ô mới được xử lý.
- Tính lại trong quá trình lan truyền ngược. Vì FlashAttention không lưu ma trận chú ý N × N trong lượt truyền xuôi, nó sử dụng một sự đánh đổi khéo léo khi huấn luyện: chỉ lưu các thống kê chuẩn hóa nhỏ trong HBM và tính lại các ô chú ý ngay trong SRAM nhanh trong quá trình lan truyền ngược. Thực hiện thêm phép tính trong SRAM tỏ ra nhanh hơn nhiều so với đọc các ma trận lớn từ HBM, giúp tăng tốc thời gian thực tế từ 2× đến 4× đồng thời giảm mức sử dụng bộ nhớ từ bậc hai O(N²) xuống tuyến tính O(N).