Công cụ AI tự động tối ưu kernel CUDA bằng LangGraph

Một lập trình viên vừa chia sẻ dự án Agentic CUDA Kernel Optimizer trên Hacker News — một công cụ dùng tác tử AI để tự động biến mô tả bài toán thành các phiên bản kernel CUDA tối ưu. Hệ thống chạy theo vòng lặp gồm sinh mã, kiểm tra tính đúng đắn, đo hiệu năng và tinh chỉnh, nhờ đó tìm ra cấu hình chạy nhanh nhất mà vẫn đảm bảo kết quả chính xác.

Điểm đáng chú ý là công cụ không chỉ đo thời gian chạy, mà còn có thể tra cứu tài liệu NVIDIA để tìm gợi ý tối ưu, đồng thời đọc các chỉ số từ Nsight Compute để định hướng cho lần thử nghiệm tiếp theo — điều mà trước đây vẫn là công việc thủ công của kỹ sư hiệu năng.

Cơ chế hoạt động

Quy trình của công cụ gồm các bước chính:

• Nạp hoặc sinh ra chữ ký hàm, các bộ dữ liệu đầu vào, kernel tham chiếu và kernel khởi tạo.
• Chạy kernel tham chiếu và đánh giá phiên bản ban đầu.
• Đề xuất thay đổi, biên dịch, so sánh kết quả bằng NumPy và đo độ trễ của kernel.
• Đưa kết quả trở lại cho lần thử kế tiếp; tự sửa các ứng viên không hợp lệ trong giới hạn số vòng lặp.
• Lưu lại ứng viên nhanh nhất đã được kiểm chứng, lịch sử thực thi và bản đồ nhiệt thời gian chạy.

Sơ đồ quy trình tác tử tối ưu kernel CUDA

Mô hình có thể thay đổi cả mã kernel lẫn cấu hình khởi chạy cho từng trường hợp. Một khung kiểm thử C++ độc lập chịu trách nhiệm biên dịch kernel bằng NVRTC, khởi chạy qua CUDA Driver API và lưu kết quả đầu ra. Phần Python đảm nhận việc so sánh và chọn ứng viên.

Mọi trường hợp đều phải vượt qua kiểm tra tính đúng đắn. Xếp hạng dựa trên trung bình nhân độ trễ trên các trường hợp hiệu năng, còn các trường hợp kiểm tra nhỏ không ảnh hưởng đến điểm số. Thời gian đo mặc định gồm 10 lần khởi chạy làm nóng và 100 lần đo thực tế bằng CUDA events; thời gian biên dịch và phát lại của profiler không được tính vào xếp hạng.

Kết quả thử nghiệm trên RTX 3060

Dự án được phát triển trên Windows với GPU RTX 3060 Laptop. Ví dụ minh họa là bài toán GEMM số thực dấu phẩy động 32-bit (Float32) trên ma trận chữ nhật, bật cả tính năng nghiên cứu tài liệu NVIDIA lẫn Nsight Compute.

Bản đồ nhiệt thời gian chạy của lần chạy 022

Bản đồ nhiệt thời gian chạy của lần chạy 023

Lần chạy 023 tiếp tục từ kernel tốt nhất của lần chạy 022 với cùng dữ liệu đầu vào và kernel tham chiếu, cho thấy khả năng tinh chỉnh lặp lại qua nhiều phiên làm việc.

Mỗi phiên làm việc tạo một thư mục trong results/run-NNN/, chứa mã nguồn kernel, yêu cầu gửi tới mô hình, dữ liệu vào/ra, phản hồi của mô hình và công cụ, cùng các tệp history.json và summary.json. Những lần chạy thành công sẽ xuất ra best.cu, các yêu cầu phát lại theo từng trường hợp và tệp heatmap.png/heatmap.svg.

Giới hạn cần lưu ý

Đây là công cụ thử nghiệm cho từng kernel riêng lẻ, không phải một giải pháp tối ưu tổng quát. Tác giả cũng thẳng thắn nêu rõ một số hạn chế quan trọng:

• Vượt qua các trường hợp được cung cấp không chứng minh kernel đúng trong mọi tình huống, và kernel tham chiếu do AI sinh ra không phải là một "chuẩn đúng" độc lập.
• Mức cải thiện phụ thuộc nhiều vào bài toán cụ thể; hiện chưa có so sánh với cuBLAS hay các thư viện của nhà cung cấp khác.
• Nên để GPU rảnh trong lúc đo để tránh sai lệch kết quả.
• Các script đầu vào do AI sinh ra được chạy cục bộ dưới dạng tiến trình con Python không có sandbox, và kernel CUDA sinh ra cũng chạy trực tiếp trên GPU của máy — đây là rủi ro bảo mật cần cân nhắc.

Yêu cầu cài đặt

Dự án cần Python 3.12 trở lên, GPU NVIDIA cùng CUDA Toolkit/driver tương thích, CMake 3.24+, trình biên dịch C++17 và khóa API OpenAI. Lệnh chạy cơ bản:

.venv\Scripts\python optimizer_agent/optimizer_agent.py --description "Single-precision GEMM with rectangular matrices." --max-iterations 7

Mô hình mặc định là gpt-5-mini với mức suy luận trung bình, và chi phí API sẽ tính vào tài khoản người dùng. Có thể dùng cờ --use-nsight để bật profiling và --nvidia-research để tra cứu hướng dẫn từ NVIDIA trước khi sinh kernel.

Với các kỹ sư hiệu năng GPU tại Việt Nam đang làm việc với CUDA, đây là một hướng tiếp cận đáng theo dõi: để AI đảm nhận phần thử-sai lặp đi lặp lại, còn con người tập trung vào thiết kế thuật toán và đánh giá tính đúng đắn ở mức hệ thống.

Dù còn nhiều giới hạn, dự án cho thấy tiềm năng rõ rệt của tác tử AI trong các bài toán tối ưu hiệu năng phần cứng — lĩnh vực vốn đòi hỏi kiến thức chuyên sâu và rất nhiều thời gian thử nghiệm thủ công.