UC 버클리, CUDA 최적화 지식 MLX로 전환…K-Search로 애플 실리콘 성능 높여

UC Berkeley BAIR 연구자 원문에 따르면 연구진은 CUDA 커널에 축적된 최적화 지식을 애플 실리콘용 MLX 커널 개발에 활용하는 K-Search 확장 방식을 공개했다. 기존 코드를 그대로 번역하는 대신 하드웨어별 제약과 최적화 원리를 구조화해 AI 탐색 과정에 제공하는 것이 핵심이다.

K-Search는 대규모 언어모델이 커널 최적화 방법을 제안하고 코드를 작성하면 실제 하드웨어에서 컴파일·측정한 결과를 다시 탐색에 반영하는 진화형 최적화 프레임워크다. 연구진은 기존 CUDA용 K-Search에 MLX 백엔드를 추가하고, CUDA의 메모리·동기화·타일링 전략을 Metal과 MLX 환경에 맞게 해석하는 변환 계층을 만들었다.

두 아키텍처의 차이는 단순한 명령어 치환으로 해결하기 어렵다. 연구진은 CUDA의 공유 메모리와 동기화 원리를 Metal의 스레드그룹 메모리와 장벽으로 대응시키면서, 메모리 한도와 대역폭 차이까지 명시했다. NVIDIA H100과 애플 M3 Max처럼 메모리 구조와 대역폭이 다른 장치에서는 같은 최적화 기법의 우선순위가 달라지기 때문이다.

탐색에는 Gemini 3.5 Pro Preview가 최적화 전략을 관리하는 역할과 후보 커널 코드를 작성하는 역할을 함께 맡았다. 모델은 커널 유형과 데이터 배치, 병목 가능성을 분석한 뒤 한 번에 하나씩 최적화를 제안한다. 각 후보는 실제 애플 하드웨어에서 측정되며, 성능이 정체되면 다른 탐색 경로로 이동한다.

연구진이 공개한 Attention 커널 실험에서 변환 지식 없이 진화 탐색만 적용한 구현은 애플의 기본 MLX Attention 커널 대비 0.26배 성능에 머물렀다. 구조화한 CUDA-to-MLX 지식을 함께 제공하자 0.97배까지 높아졌다. 연구진은 이 과정에서 스레드그룹 메모리 타일링과 온라인 소프트맥스, 메모리 접근을 위한 K 전치, 애플 하드웨어의 빠른 exp2 명령 활용 같은 최적화가 선택됐다고 설명했다.

Mamba 상태공간모델 커널에서는 더 큰 차이가 보고됐다. M1 Max 64GB에서 mamba-370m을 f16으로 실행한 결과, 연구진 구현은 시퀀스 길이 4096의 사전 채우기에서 초당 6,743토큰을 처리했다. 비교한 커뮤니티 mlx-lm 구현은 초당 339토큰이었다. 약 20배 차이는 순차 상태 갱신을 병렬 접두 스캔으로 바꾼 데서 나왔으며, 단일 토큰 디코딩에서는 각각 초당 152토큰과 116토큰으로 격차가 작았다.

이번 결과는 이미 축적된 한 아키텍처의 최적화 경험을 새로운 하드웨어 생태계로 옮길 때 AI 탐색을 활용할 수 있음을 보여준다. 다만 MLX 확장 실험은 Attention과 Mamba 두 커널, 특정 M1 Max 구성에서 이뤄졌다. 연구진도 다른 커널과 장치에서 같은 방식이 얼마나 일반화되는지는 추가 검증이 필요하다고 밝혔다.

출처: UC Berkeley BAIR 연구자 원문 https://bair.berkeley.edu/blog/2026/07/29/cuda-to-mlx-k-search/
출처: K-Search 연구 논문 https://arxiv.org/abs/2602.19128