K-Search: Chuyển đổi tối ưu hóa CUDA sang Apple Silicon MLX với hiệu năng vượt trội

Hình 1: Bản đồ chuyển đổi tối ưu hóa từ CUDA sang MLX. Kiến thức tối ưu hóa CUDA có thể được chuyển đổi thành các chiến lược MLX gốc của kiến trúc thay vì sao chép từng lệnh một.

Chúng ta đang đối mặt với một kỷ nguyên mới trong lĩnh vực điện toán. Phần cứng đang thay đổi nhanh chóng — không chỉ là các GPU nhanh hơn, mà còn là sự gia tăng của hàng loạt chip từ các nhà cung cấp khác nhau, mỗi loại đều có kiến trúc riêng và thường được tùy chỉnh cho các khối lượng công việc AI cụ thể. Phần mềm cũng đang thay đổi với tốc độ tương đương, và các công cụ lập trình AI hiện nay có thể tạo ra trong vài phút những thứ mà vài năm trước phải mất hàng tháng trời mới hoàn thành.

Với việc phần lớn hoạt động điện toán hiện nay tập trung vào AI, các GPU kernel trở thành thành phần then chốt cho sự thành công của nó. Đây là các chương trình cấp thấp chạy bên trong GPU, và việc viết ra những kernel hiệu quả không hề đơn giản — cần nhiều năm kinh nghiệm chuyên môn mới có thể thực hiện đúng. Việc chuyển một kernel từ phần cứng của nhà cung cấp này sang nhà cung cấp khác thậm chí còn khó khăn hơn, và thường đồng nghĩa với việc phải tìm ra các phương pháp tối ưu hóa từ con số không. Ví dụ, hệ sinh thái CUDA đã tích lũy hàng thập kỷ kinh nghiệm quý báu về kernel: các triển khai được tinh chỉnh thủ công cho attention, state space models và các thao tác quan trọng khác, đại diện cho hàng ngàn giờ kỹ thuật. Các hệ sinh thái phần cứng mới hơn (Apple Silicon, các bộ tăng tốc AI tùy chỉnh và những loại khác) đang phát triển nhanh chóng nhưng lại thiếu chiều sâu này.

Trong công trình này, chúng tôi đặt câu hỏi liệu chuyên môn đó có thể được chuyển giao tự động hay không. Chúng tôi xây dựng dựa trên K-Search, một khung tìm kiếm kernel tiến hóa do Cao và cộng sự tại Berkeley Sky Lab giới thiệu, sử dụng AI để tối ưu hóa các GPU kernel, và mở rộng nó với một backend cho MLX — khung học máy của Apple dành cho các chip Apple Silicon của riêng họ. Chúng tôi đã phát triển một lớp chuyển đổi CUDA-sang-MLX có cấu trúc mới, cho phép K-Search lấy các CUDA kernel hiện có làm cơ sở tri thức và điều chỉnh chúng thành các GPU kernel chất lượng cao cho Apple Silicon, thay vì phải xây dựng lại từ đầu.

Chúng tôi chứng minh rằng phương pháp của mình đạt được hiệu suất gần mức chuyên gia trên Apple Silicon với tốc độ nhanh gấp 0,97 lần so với kernel Attention gốc của MLX, và tốc độ prefill nhanh gấp 20 lần so với bản triển khai mlx-lm của cộng đồng trên kernel Mamba SSM; chúng tôi báo cáo các con số này và mức độ đóng góp của lớp chuyển đổi vào hiệu suất đạt được trong các phần dưới đây. Mặc dù chúng tôi tập trung vào các MLX kernel cho Apple Silicon, phương pháp này không chỉ dành riêng cho MLX mà còn áp dụng cho bất kỳ hệ sinh thái nào có thể chuyển giao chuyên môn CUDA.

Tại sao lại là MLX?

Khung MLX của Apple đã chứng kiến sự đón nhận đáng kể kể từ cuối năm 2023. Với việc Apple Silicon có mặt trong hàng trăm triệu chiếc MacBook và Mac Studio, MLX cho phép thực hiện suy luận AI cục bộ mà không tốn chi phí đám mây. Kiến trúc bộ nhớ thống nhất (unified memory architecture) làm cho nó trở nên đặc biệt hấp dẫn đối với các mô hình cỡ trung bình (7B–70B tham số trên các chip dòng M).

Tuy nhiên, ẩn sau đà phát triển này là một khoảng cách đáng kể: nhiều kernel quan trọng về hiệu suất mà hệ sinh thái NVIDIA coi là hiển nhiên: paged attention, các kernel quét SSM được tối ưu hóa, fused MoE routing, hoặc là không tồn tại, hoặc là ở dạng sơ khai nếu không được tinh chỉnh theo phần cứng cụ thể. MLX chạy các mô hình một cách chính xác nhưng thường bỏ lỡ đáng kể tiềm năng hiệu suất.

Khoảng cách này chính là động lực cho phần còn lại của bài viết này.

K-Search là gì?

K-Search là một khung tối ưu hóa kernel tiến hóa ban đầu được phát triển bởi tác giả chính của chúng tôi, Shiyi Cao, tại UC Berkeley Sky Lab. Với một kernel sơ khai và thông số kỹ thuật phần cứng, nó chạy một vòng lặp tối ưu hóa lặp đi lặp lại: một LLM suy luận về các tối ưu hóa nên thử tiếp theo, một mô hình viết mã tạo ra các kernel ứng viên, và các ứng viên đó được biên dịch và đánh giá hiệu năng trên phần cứng thực tế.

Các phép đo được phản hồi ngược lại vào quá trình tìm kiếm, giúp liên tục tinh chỉnh, theo đuổi các hướng đi hứa hẹn và loại bỏ các ngõ cụt cho đến khi hiệu suất hội tụ.

Thuật toán 1: K-Search thông qua các mô hình thế giới đồng tiến hóa. Quá trình tìm kiếm luân phiên giữa việc chọn hành động hứa hẹn nhất, khởi tạo và đánh giá mã cho đến khi sự cải thiện bị đình trệ, và tiến hóa mô hình thế giới thông qua các thao tác chèn, cập nhật và cắt tỉa. Được chuyển thể từ Cao và cộng sự (2026).

Quá trình tìm kiếm được dựa trên một Spec (Thông số kỹ thuật): một tài liệu chuyên biệt mã hóa các quy tắc phần cứng, các mẫu tối ưu hóa và các ràng buộc toán học, giúp ngăn mã được tạo ra không bị "ảo giác" các nguyên hàm không hợp lệ và đảm bảo các ứng viên sẽ thực sự biên dịch và chạy hiệu quả.

Trong các lần chạy của chúng tôi, một mô hình duy nhất (Gemini 3.5 Pro Preview) đóng cả hai vai trò: duy trì trạng thái suy luận và viết các kernel. Phần suy luận được nhắc nhở đóng vai trò là một "kỹ sư hiệu năng GPU kernel" và được yêu cầu thực hiện một phân tích cố định trước khi đề xuất bất cứ điều gì: phân loại kernel (reduction, scan, attention/softmax, …), viết lại phép tính tham chiếu ở dạng chuẩn, lập bản đồ bố cục dữ liệu và các mẫu truy cập, đồng thời đưa ra giả thuyết về nút thắt cổ chai có khả năng xảy ra (băng thông, độ trễ, tính toán hoặc đồng bộ hóa) trong mỗi chế độ runtime. Chỉ sau đó, nó mới đưa ra các tối ưu hóa ứng viên, mỗi cái là một thay đổi đơn lẻ có thể thực hiện trong một lần lặp.

Chúng tôi gọi trạng thái suy luận bền vững này là mô hình thế giới (world model). Thay vì một danh sách phẳng các thứ cần thử, nó là một cây quyết định (tiền tố): mỗi đường dẫn từ gốc đến lá tạo thành một kế hoạch tối ưu hóa hoàn chỉnh, và các nhánh anh em là các lựa chọn thay thế cạnh tranh. Mỗi nút đều được chấm điểm — một overall_rating trong khoảng [0, 10], độ tin cậy trong khoảng [0, 1], và tác động trên mỗi nút đối với băng thông bộ nhớ, áp lực thanh ghi và sự phù hợp với tính toán/phần cứng — để quá trình tìm kiếm có thể xếp hạng các kế hoạch một phần và mở rộng các kế hoạch hứa hẹn nhất. Cây này tồn tại và phát triển qua các vòng: việc tinh chỉnh một ý tưởng sẽ thêm một nút con thay vì ghi đè lên nút cha của nó, và nếu điểm số tốt nhất không cải thiện trong vài vòng (cửa sổ đình trệ), quá trình tìm kiếm sẽ quay lại để khám phá một nhánh thay thế. Một nút đơn lẻ, như nó xuất hiện giữa quá trình chạy trên kernel attention, trông như thế này:

Danh sách 1: Ví dụ về nút mô hình thế giới của K-Search. Mỗi tối ưu hóa ứng viên ghi lại một hành động cụ thể, các tác động phần cứng ước tính, xếp hạng ưu tiên tổng thể và độ tin cậy của mô hình.

Hình 2: Tổng quan về K-Search. Khung này hoạt động trên một Trạng thái Tìm kiếm $S_t$ được cấu trúc dưới dạng một cây tìm kiếm. Cây bao gồm các nút Đóng (màu xanh, các trạng thái đã truy cập kèm theo chương trình như $x_{12}$) và một Biên của các nút Mở (màu cam, các giả thuyết đang chờ xử lý như $u_{13}$). Quy trình làm việc lặp qua ba giai đoạn: (1) Chọn Hành động, nơi nút hành động hứa hẹn nhất được truy xuất từ biên dựa trên điểm ưu tiên ước tính $V$ của mô hình thế giới; (2) Tinh chỉnh Cục bộ, nơi một chính sách ngẫu nhiên $\pi_{\mathrm{code}}$ lấy mẫu các triển khai cụ thể cho đến khi đình trệ; và (3) Cập nhật Mô hình Thế giới, nơi LLM suy luận về quỹ đạo để cập nhật cây tìm kiếm thông qua Chèn (thêm hành động mới), Cập nhật (điều chỉnh $V$, ví dụ: $u_{11}$ giảm từ 0,9 xuống 0,6), và Cắt tỉa (loại bỏ các nút ít hứa hẹn hơn như $u_{10}$).

Bài báo K-Search gốc đã đánh giá chiến lược tìm kiếm này trên các CUDA kernel từ FlashInfer. Trên khắp các tác vụ GQA decode, MLA decode, MLA prefill và MoE, K-Search đã cải thiện nhất quán hơn so với OpenEvolve và ShinkaEvolve trong cùng ngân sách 120 lần lặp. Những kết quả này thiết lập khung tìm kiếm mà chúng tôi xây dựng dựa trên đó; phần còn lại của bài viết này đặt câu hỏi liệu kiến thức tối ưu hóa của nó có thể chuyển giao ra ngoài CUDA hay không.

Hình 3: Các kết quả chính từ bài báo K-Search gốc. Qua ba lần chạy, K-Search đạt được điểm số tìm kiếm tốt nhất, hiệu suất kernel trên mỗi khối lượng công việc và phân phối tốc độ mạnh mẽ hơn so với OpenEvolve và ShinkaEvolve trên bốn CUDA kernel của FlashInfer. Được tái tạo chính xác từ Cao và cộng sự (2026).

Xây dựng một backend MLX

Để mang K-Search đến với Apple Silicon, trước tiên chúng tôi đã xây dựng một backend MLX gốc. Chúng tôi đã triển khai một bộ điều hợp tác vụ dành riêng cho MLX cho K-Search, bao gồm:

Chuyển đổi chuyên môn CUDA sang MLX

Tuy nhiên, thách thức thú vị hơn không chỉ đơn thuần là chạy K-Search trên MLX. Thông tin chi tiết chính là các CUDA kernel chuyên gia mã hóa hàng thập kỷ kiến thức tối ưu hóa có thể chuyển giao sang GPU Apple nếu bạn có thể thu hẹp khoảng cách về khái niệm. Chỉ đơn giản là đưa cho LLM một CUDA kernel và yêu cầu nó chuyển đổi là không đủ: nếu không có bối cảnh phần cứng sâu sắc, nó sẽ tạo ra mã hợp lệ về cú pháp nhưng sai về kiến trúc (kích thước tile sai, nguyên hàm không hợp lệ, giả định bộ nhớ không khớp).

Lớp chuyển đổi của chúng tôi bao gồm:

Đạt hiệu suất kernel chuyên gia: kernel Attention

Chúng tôi đánh giá ba cấu hình của một kernel attention MLX cho Apple Silicon: (1) một baseline sơ khai, (2) tiến hóa thuần túy không có thêm bối cảnh được cung cấp, và (3) một lớp chuyển đổi bối cảnh đầy đủ, cung cấp cho bộ tối ưu hóa kiến thức triển khai cụ thể theo kiến trúc được trích xuất từ các kernel hiệu năng cao (ví dụ: FlashAttention-2), cho phép tìm kiếm tiến hóa suy luận về các chiến lược triển khai thay vì bắt đầu từ một kernel sơ khai. Cùng với nhau, ba cấu hình này cho phép chúng tôi cô lập tác động chính xác của lớp chuyển đổi.

Hình 4: Mở rộng hiệu suất của Attention Kernel thông qua các tối ưu hóa xếp chồng. Cấu hình "Full Context" (Bối cảnh đầy đủ) khám phá và triển khai thành công các chiến lược nâng cao như double buffering và loop unrolling, đạt được hiệu suất gần mức chuyên gia.

Bước nhảy từ 0,26× lên 0,97× tốc độ của kernel attention hiện đại nhất của Apple minh họa tầm quan trọng của lớp chuyển đổi. Với bối cảnh đầy đủ, kernel đã tiến hóa tự khám phá các tối ưu hóa chính trong FlashAttention 2: threadgroup memory tiling, online softmax, K-transposition để truy cập bộ nhớ và thủ thuật exp2. Thủ thuật cuối cùng này thay thế mọi hàm mũ softmax bằng hàm mũ cơ số 2,

điều này là chính xác và cho phép kernel sử dụng trực tiếp lệnh phần cứng fast::exp2 của Apple thay vì phải trả phí cho việc chuyển đổi cơ số tại thời điểm chạy.

Prefill nhanh gấp 20 lần: kernel Mamba SSM

Để đánh giá liệu K-Search có khái quát hóa ngoài các kernel attention hay không, chúng tôi đã áp dụng nó cho kernel state-space model (SSM) được Mamba sử dụng. Không giống như attention, nút thắt cổ chai tính toán là một bản cập nhật trạng thái đệ quy thay vì softmax, mang lại một thách thức tối ưu hóa khác biệt đáng kể. Chúng tôi so sánh bản triển khai đã tiến hóa với bản triển khai MLX của cộng đồng (mlx-lm) và bản triển khai tham chiếu PyTorch (mamba.py) trên M1 Max.

Được đánh giá trên mamba-370m f16, M1 Max 64GB:

Bảng 1: Thông lượng prefill và decode trên mamba-370m (f16, M1 Max 64GB). mlx-mamba (của chúng tôi) đạt thông lượng prefill cao hơn khoảng 20 lần so với baseline mlx-lm của cộng đồng, trong khi decode vẫn tương đương.

Tốc độ prefill nhanh hơn khoảng 20 lần so với mlx-lm xuất phát từ một sự khác biệt: mlx-lm không triển khai quét song song (parallel scan) cho SSM. Sự đệ quy trạng thái

trông có vẻ tuần tự về bản chất, nhưng mỗi bước có thể được viết dưới dạng một cặp $(\bar{a}_t, \bar{b}_t)$ theo phép kết hợp kết hợp (associative combine)

điều này tái tạo chính xác sự đệ quy. Vì toán tử có tính kết hợp, toàn bộ chuỗi có thể được đánh giá bằng một phép quét (tiền tố) song song trong $O(\log N)$ bước phụ thuộc thay vì $O(N)$. mlx-lm bỏ qua điều này và xử lý từng token một, khiến phần lớn khả năng tính toán của Apple Silicon bị nhàn rỗi; kernel Metal đã tiến hóa của chúng tôi áp dụng phép quét và tận dụng đầy đủ hơn thông lượng của GPU. Mức tăng này xuất hiện ở prefill, nơi toàn bộ chuỗi có sẵn để quét song song, và không xuất hiện ở decode từng token, nơi chỉ có một token mới mỗi bước và không có phép quét để song song hóa — đó là lý do tại sao hàng decode gần như phẳng trong khi prefill nhanh hơn khoảng 20 lần.

mamba.py chậm ở cả prefill và decode vì nó là bản triển khai tham chiếu PyTorch, phải quay lại sử dụng CPU hoặc MPS trên Apple Silicon, từ bỏ các tối ưu hóa dành riêng cho phần cứng mà backend Metal của MLX có thể thực hiện được.

Tiếp theo là gì?