Dùng ~25 dòng SQL để tự động suy ra kernel FlashAttention

Một thử nghiệm mới trên cộng đồng Hacker News đang thu hút sự chú ý của giới nghiên cứu hệ thống ML: chỉ với khoảng 25 dòng SQL mô tả attention cơ bản, hệ thống có thể tự động suy ra kernel FlashAttention tối ưu. Thành quả này đến từ ba phần mở rộng SQL có thể mang đi tái sử dụng, mở ra hướng tiếp cận mới ở giao điểm giữa lý thuyết cơ sở dữ liệu và hệ thống học máy.

Ý tưởng: tách lớp logic khỏi lớp vật lý

Điểm cốt lõi của dự án là luận điểm cho rằng việc tách các lớp logic trong một chương trình tensor khỏi các lớp vật lý vừa đúng đắn hơn, vừa dễ mở rộng hơn.

Thay vì viết tay một kernel tối ưu cho từng loại phần cứng, người dùng chỉ cần mô tả attention nguyên bản bằng SQL — biểu thức softmax(Q·Kᵀ/√d)·V quen thuộc. Sau đó, một bộ tối ưu truy vấn sẽ tự động tìm ra kế hoạch thực thi hiệu quả, trong trường hợp này chính là FlashAttention.

Nếu "Bài học Cay đắng" (Bitter Lesson) là giáo lý trung tâm của học máy, thì tại sao không hình thức hóa nó bằng đại số quan hệ?

Ba phần mở rộng SQL làm nên điều này

Thành quả đạt được nhờ ba thành phần được nhóm phát triển xây dựng:

• xarray-sql: thư viện biểu diễn hiệu quả các tập dữ liệu mảng quy mô petabyte trên đám mây dưới dạng bảng (quan hệ).
• ddx: phần mở rộng cơ sở dữ liệu bổ sung khả năng tự động vi phân theo phong cách JAX cho SQL chuẩn, dựa trên mô hình dữ liệu mảng nói trên.
• einfold: bộ lập kế hoạch truy vấn dạng mô-đun giúp co rút tensor nhanh chóng trong SQL thuần.

Đáng chú ý, einfold tìm ra kế hoạch hợp nhất (fused plan) mà không cần bất kỳ quy tắc nào viết riêng cho attention. Quy tắc của nó hợp nhất mọi phép trung bình có trọng số softmax mà điểm số được tính từ một phép join.

Kiểm chứng nghiêm ngặt

Bản demo không chỉ dừng ở việc chạy được, mà còn kiểm tra chéo nhiều khía cạnh:

• Lan truyền xuôi (forward): mỗi cách viết model trên mỗi engine (DataFusion, DataFusion kết hợp einfold và DuckDB) được đối chiếu với cùng phép attention viết bằng JAX.
• Gradient: một hàm mất mát và đạo hàm của nó theo toàn bộ giá trị Q, K, V được diễn đạt trong một câu lệnh SQL duy nhất, rồi kiểm chứng bằng jax.grad.
• Mask là các hàng không tồn tại: một giá trị vô cực trong V chỉ lan tới những truy vấn được phép thấy nó, đúng như ngữ nghĩa SQL yêu cầu.
• Đối chứng không phải attention: bộ ước lượng Nadaraya–Watson (trung bình có trọng số kernel Gauss) cũng được tăng tốc — bất kỳ quy tắc nào tối ưu attention đều phải tối ưu được truy vấn này.

Ý nghĩa với cộng đồng lập trình viên

Với các kỹ sư AI tại Việt Nam đang huấn luyện mô hình lớn trên hạ tầng hạn chế, hướng tiếp cận này gợi mở một ý tưởng đáng quan tâm: để bộ tối ưu tự lo phần tối ưu phần cứng, còn người phát triển chỉ cần mô tả ý định ở mức logic. Đây cũng là lời nhắc rằng những nền tảng lâu đời như đại số quan hệ vẫn có thể đóng góp góc nhìn mới cho các bài toán hiện đại.

Dự án đang tìm kiếm cộng tác viên quan tâm đến giao điểm giữa lý thuyết cơ sở dữ liệu và hệ thống ML, với cộng đồng thảo luận đặt tại xql.systems.