Dari CUDA ke MLX: Bagaimana K-Search Membawa Puluhan Tahun Keahlian Kernel ke Apple Silicon

Peneliti dari UC Berkeley Sky Lab mengembangkan K-Search, kerangka kerja pencarian kernel evolusioner yang menggunakan AI untuk mengoptimalkan kernel GPU, dan memperluasnya dengan backend untuk framework MLX milik Apple. Mereka mengembangkan lapisan terjemahan CUDA-ke-MLX terstruktur yang baru yang mengubah kernel CUDA yang ada menjadi kernel berkualitas tinggi untuk Apple Silicon. Pendekatan ini mencapai peningkatan kecepatan hingga 0,97x dibandingkan kernel Atensi MLX asli dan peningkatan kecepatan prapengisian hingga 20x dibandingkan implementasi mlx-lm komunitas pada kernel Mamba SSM. K-Search menggunakan loop optimasi iteratif di mana LLM mempertimbangkan optimasi, menghasilkan kandidat kernel, mengompilasi, dan mengujinya pada perangkat keras nyata. Lapisan terjemahan mencakup tabel pemetaan konsep, petunjuk khusus MLX, dan pernyataan yang dapat digunakan kembali yang mempertahankan perilaku kernel ahli. Pekerjaan ini menunjukkan bahwa pencarian kernel yang digerakkan oleh AI dapat mentransfer pengetahuan optimasi di seluruh ekosistem perangkat keras tanpa memulai dari awal.