[Giới thiệu] Meta ra mắt KernelLLM, một mô hình 8B được tinh chỉnh dựa trên Llama 3.1, có thể tự động chuyển đổi mã PyTorch thành hạt nhân GPU Triton hiệu quả. Dữ liệu đo lường thực tế cho thấy hiệu suất suy luận một lần của nó vượt qua GPT-4o và DeepSeek V3, và điểm số của nó tăng vọt khi được tạo lần.
Trong lĩnh vực AI, kích thước tham số từng được coi là "giới hạn hiệu suất".
KernelLLM mới nhất của Meta, với tham số 8B nhỏ, đánh bại GPT-4o 200B trong nhiệm vụ tạo hạt nhân GPU.
Đây là mô hình tham số 8B được tinh chỉnh dựa trên Llama 3.1 Instruct, nhằm mục đích tự động chuyển đổi mô-đun PyTorch thành hạt nhân GPU Triton hiệu quả.
KernelLLM chỉ đơn giản là một công cụ kỳ diệu để phát triển nhân GPU. Nó đạt được hiệu suất mạnh mẽ hơn với ít thông số hơn, đơn giản và dễ sử dụng.
Nó chỉ có 8B tham số, nhưng trong KernelBench-Triton Cấp độ 1, hiệu suất suy luận một lần vượt quá GPT-4o và DeepSeek V3.
Qua lần suy luận, KernelLLM hoạt động tốt hơn DeepSeek R1.
Tất cả những điều này đều xuất phát từ một mô hình có kích thước tham số nhỏ hơn hai bậc so với đối thủ cạnh tranh.
@Denis Kanonik phàn nàn, "Liệu điều này có được đào tạo lại bằng bộ kiểm tra không?"
KernelLLM giúp phát triển kernel dễ dàng hơn
KernelLLM là mô hình 8B dựa trên Llama 3.1 Instruct, được đào tạo chuyên biệt cho nhiệm vụ viết hạt nhân GPU bằng Triton.
Nó giúp lập trình GPU dễ dàng hơn và tự động tạo ra các hạt nhân GPU hiệu suất cao.
KernelLLM đáp ứng nhu cầu tăng trưởng về hạt nhân GPU hiệu suất cao bằng cách tự động tạo ra các triển khai Triton hiệu quả.
Khi khối lượng công việc tăng lên và kiến trúc accelerator đa dạng hơn, nhu cầu về các giải pháp cốt lõi tùy chỉnh cũng tăng lên đáng kể.
Nhiều công cụ liên quan trên thị trường chỉ có thể được tối ưu hóa trong quá trình thử nghiệm hoặc chỉ tập trung vào việc điều chỉnh các vấn đề của KernelBench, khiến việc xử lý nhiều tình huống khác nhau trở nên khó khăn.
KernelLLM là LLM đầu tiên được tinh chỉnh dựa trên dữ liệu trong mã bên ngoài (PyTorch, Triton).
Quy trình tạo hạt nhân Triton
Nhập mã PyTorch và KernelLLM sẽ tạo mã ứng viên hạt nhân Triton.
Sau đó sử dụng các bài kiểm tra đơn vị để xác minh mã, chạy mã với các đầu vào ngẫu nhiên và xem đầu ra có chính xác không. Nếu bạn tạo ra nhiều mã ứng viên, bạn có thể so sánh mã nào tốt nhất và chọn mã tốt nhất.
Quy trình tạo hạt nhân Triton của KernelLLM : Sử dụng KernelLLM để dịch mã PyTorch thành mã ứng viên cho hạt nhân Triton. Mã được tạo ra sẽ được xác minh bằng các bài kiểm tra đơn vị thực hành hạt nhân với dữ liệu đầu vào ngẫu nhiên có hình dạng đã biết. Quá trình này hỗ trợ tạo nhiều mã ứng viên (đánh giá bởi pass@k), tăng số lượng ứng viên để cải thiện chất lượng và cuối cùng chọn triển khai hạt nhân Triton tốt nhất làm đầu ra (phần màu xanh lá cây)
Để đào tạo mô hình này, đội ngũ đã nỗ lực rất nhiều, sử dụng hơn 25.000 cặp ví dụ mã (PyTorch, Triton) và các mẫu tổng hợp.
Một số dữ liệu này đến từ mã lọc của TheStack và một số được tạo thông qua torch.compile() và các kỹ thuật gợi ý.
Dữ liệu KernelBook, liên kết tham khảo: https://huggingface.co/datasets/GPUMODE/KernelBook.
Mô hình Llama3.1-8B-Instruct được sử dụng để đào tạo và điều chỉnh có giám sát (SFT) được thực hiện trên dữ liệu tùy chỉnh để kiểm tra khả năng tạo hạt nhân Triton chính xác và mã gọi trên KernelBench-Triton.
KernelBench-Triton là một biến thể được phát triển dựa trên KernelBench [Ouyang et al. 2025], tập trung vào việc tạo ra hạt nhân Triton.
Khi đào tạo và đánh giá, mã PyTorch được cấu hình bằng mẫu nhắc nhở có chứa các ví dụ về định dạng làm hướng dẫn.
Mô hình được đào tạo trong 10 kỷ nguyên với quy mô lô là 32 bằng phương pháp SFT tiêu chuẩn và các siêu tham số được chọn dựa trên mức độ phức tạp của tập xác thực.
Quá trình đào tạo mất 12 giờ trên 16 GPU (192 giờ GPU) và kết quả xác thực báo cáo để tìm ra điểm kiểm tra tốt nhất.
Đánh giá hiệu suất
Mặc dù có kích thước mô hình nhỏ hơn, hiệu suất của nó có thể so sánh với LLM hiện đại.
Trong thử nghiệm KernelBench-Triton, KernelLLM có 8 tham số có điểm suy luận lần là 20,2, cao hơn DeepSeek V3 có 671 tham số (16 điểm) và GPT-4o có 200 tham số (15 điểm).
Nếu tạo ra nhiều mã ứng viên hơn, điểm số có thể tăng cao tăng. Khi tạo ra 10 mã ứng viên, điểm có thể đạt tới 51,8 điểm và khi tạo ra 20 mã ứng viên, điểm có thể đạt tới 57,1 điểm.
Suy luận KernelLLM được chạy với nhiệt độ = 1,0 và top_p = 0,97.
Các mô hình đã được thử nghiệm trên KernelBench, một chuẩn mực mã mã nguồn mở để đánh giá khả năng của các nhân GPU hiệu quả được viết bằng LLM.
Nó chứa 250 mô-đun PyTorch được lựa chọn cẩn thận, được điều chỉnh theo khối lượng công việc, từ các thao tác đơn giản (như Conv2D hoặc Swish, Cấp độ 1) cho đến kiến trúc mô hình hoàn chỉnh (Cấp độ 3).
Nó hoạt động rất ổn định trong nhiệm vụ có nhiều mức độ khó khác nhau và có thể xử lý cả những toán tử đơn giản và những kiến trúc mô hình phức tạp.
Các bài kiểm tra đo lường cả tính chính xác của mã của bạn (thông qua so sánh với đầu ra PyTorch tham khảo) và hiệu suất (thông qua tốc độ so với triển khai cơ sở).
Đội ngũ đã phát triển một biến thể KernelBench-Triton mới để đánh giá cụ thể khả năng tạo hạt nhân Triton của LLM, rất phù hợp để thử nghiệm KernelLLM.
Tất cả các thử nghiệm đều được thực hiện trên GPU NVIDIA H100.
KernelLLM thể hiện hành vi mở rộng tính gần đúng trong pass@k.
Làm thế nào để sử dụng KernelLLM?
Đầu tiên hãy cài đặt một số gói phụ thuộc:
- pip cài đặt máy biến áp tăng tốc ngọn đuốc triton
Khi sử dụng, trước tiên hãy import thư viện và gọi hàm generate_triton để tạo mã Triton đã tối ưu hóa.
KernelLLM cung cấp giao diện đơn giản để tạo hạt nhân Triton từ mã PyTorch.
Nếu bạn không muốn viết một tập lệnh, bạn cũng có thể chạy trực tiếp python kernelllm.py, sử dụng giao diện REPL tích hợp, mở giao diện tương tác và xem kết quả theo thời gian thực.
kernelllm.py cung cấp nhiều phương pháp khác nhau để tương tác với mô hình.
- python kernelllm.py
KernelLLM cung cấp một số phương pháp để tùy chỉnh quy trình xây dựng:
Đôi khi nó mắc những lỗi nhỏ, chẳng hạn như tham chiếu API không đúng, lỗi cú pháp và đôi khi nó không thể tạo ra hạt nhân lý tưởng theo hướng dẫn.
Cấu trúc của mã được tạo ra cũng giống như những gì trình biên dịch tự động đưa ra và đôi khi có thể xảy ra sự cố ở các chi tiết như đặt tên biến, hình dạng tenxơ, xử lý kiểu và độ chính xác về số.
Tham khảo:
https://x.com/reach_vb/status/1924478755898085552
https://huggingface.co/facebook/KernelLLM
Bài viết này trích từ tài khoản công khai WeChat "Xinzhiyuan" , do Yingzhi biên tập và được 36Kr xuất bản với sự cho phép.
