Hướng dẫn lập trình GPU NVIDIA theo kiến trúc Tile: Từ cuTile, Triton đến Flash Attention

Trong bài hướng dẫn này, chúng ta sẽ khám phá lập trình GPU với TileGym bằng cách xây dựng một quy trình làm việc (workflow) trên Colab có thể chạy trên các điều kiện phần cứng khác nhau. Chúng ta bắt đầu bằng việc thăm dò môi trường CUDA hiện có, kiểm tra xem NVIDIA cuTile có chạy trực tiếp được không, và chuyển sang dùng Triton khi các GPU tiêu chuẩn trên Colab thiếu bộ stack cuTile cần thiết. Thông qua thiết lập này, chúng ta sẽ học được ý tưởng cốt lõi của lập trình theo tile (tile-programming): thay vì viết mã cho từng luồng (thread) một, chúng ta thao tác trên toàn bộ các tile dữ liệu, tải chúng vào kernel, tính toán hiệu quả và lưu kết quả trở lại. Chúng ta sử dụng mô hình này để triển khai phép cộng vector, fused GELU, row-wise softmax, nhân ma trận theo tile (tiled matrix multiplication) và flash attention, đồng thời so sánh từng kết quả với PyTorch để kiểm tra tính chính xác và đo lường hiệu năng (benchmarking).

Thăm dò môi trường CUDA

Chúng ta bắt đầu bằng việc thiết lập môi trường, nhập các thư viện cần thiết và kiểm tra xem CUDA có khả dụng trên runtime hiện tại hay không. Chúng ta kiểm tra khả năng của GPU, phiên bản CUDA và cấu hình PyTorch để xác định xem backend cuTile thực thụ có thể sử dụng được hay không. Sau đó, chúng ta chọn backend thực thi đang hoạt động, giải thích mô hình lập trình theo tile và lưu trữ các chuỗi nguồn kernel cuTile tham chiếu để so sánh.

Định nghĩa các Triton Kernel

Chúng ta định nghĩa các triển khai Triton cho phép cộng vector, fused GELU, row softmax, nhân ma trận theo tile và flash attention. Chúng ta biểu diễn từng thao tác bằng cách sử dụng các lệnh tải, tính toán, rút gọn (reduction), tích vô hướng (dot product) và lưu trữ ở cấp độ tile, để GPU có thể xử lý các khối dữ liệu một cách hiệu quả. Chúng ta cũng cung cấp các hàm dự phòng (fallback) bằng PyTorch thuần túy để bài hướng dẫn vẫn có thể chạy được khi Triton hoặc backend GPU hỗ trợ không khả dụng.

Cộng Vector và GELU

Chúng ta xây dựng các tiện ích đo lường hiệu năng và kiểm tra tính chính xác để so sánh từng kernel tùy chỉnh với bản triển khai tham chiếu của PyTorch. Sau đó, chúng ta chạy kernel cộng vector và xác minh rằng đầu ra dựa trên tile khớp với phép cộng tiêu chuẩn của PyTorch. Chúng ta cũng kiểm tra kernel fused GELU, minh họa cách kết hợp phép nhân, cộng bias và kích hoạt GELU thành một lượt thực thi hiệu quả duy nhất.

Softmax và Nhân ma trận theo Tile

Chúng ta chạy kernel row-wise softmax và so sánh với softmax của PyTorch để xác minh tính chính xác về mặt số học. Sau đó, chúng ta thực hiện nhân ma trận theo tile, nhân các khối ma trận và tích lũy dọc theo chiều K. Chúng ta đo lường hiệu năng các kernel này so với PyTorch để quan sát cách thực thi dựa trên tile hoạt động trên backend đang sử dụng.

Kernel Flash Attention

Chúng ta kết thúc với kernel flash attention, áp dụng online softmax để tính toán attention mà không cần tạo ra toàn bộ ma trận attention. Chúng ta so sánh đầu ra của nó với scaled dot-product attention của PyTorch và đo lường hiệu năng thời gian chạy khi có sẵn backend GPU. Chúng ta khép lại bài hướng dẫn bằng cách tóm tắt backend đã sử dụng và các khái niệm lập trình theo tile chính đã học được.

Kết luận

Tóm lại, chúng ta đã hiểu cách các kernel dựa trên tile ánh xạ các thao tác toán học cấp cao lên các mô hình thực thi GPU hiệu quả. Chúng ta đã thấy cách fusion giảm lưu lượng bộ nhớ, cách các phép rút gọn trên tile giúp ổn định và tối ưu hóa softmax, cách nhân ma trận theo tile tích lũy trên các khối K, và cách flash attention sử dụng online softmax để tránh việc tạo ra toàn bộ ma trận attention. Chúng ta cũng có được lộ trình để thử nghiệm: chúng ta đã chạy các kernel Triton trên các GPU Colab phổ biến hiện nay, đồng thời thấy được cách các khái niệm tương tự chuyển đổi sang các kernel cuTile thực thụ trên các hệ thống Ampere, Ada hoặc Blackwell mới hơn với CUDA 13.1+.

Hãy xem toàn bộ mã nguồn cùng với Notebook. Ngoài ra, đừng ngần ngại theo dõi chúng tôi trên Twitter, tham gia SubReddit ML với hơn 150 nghìn thành viên và đăng ký nhận bản tin của chúng tôi. Khoan đã! Bạn có dùng Telegram không? Bây giờ bạn cũng có thể tham gia cùng chúng tôi trên Telegram.

Bạn cần hợp tác với chúng tôi để quảng bá GitHub Repo, trang Hugging Face, sản phẩm mới hoặc hội thảo trực tuyến, v.v.? Hãy kết nối với chúng tôi.

Sana Hassan, thực tập sinh tư vấn tại Marktechpost và là sinh viên chương trình bằng kép tại IIT Madras, có niềm đam mê áp dụng công nghệ và AI để giải quyết các thách thức thực tế. Với sự quan tâm sâu sắc đến việc giải quyết các vấn đề thực tiễn, anh mang đến một góc nhìn mới mẻ cho sự giao thoa giữa AI và các giải pháp trong đời sống.