CUDA에서 MLX로: K-Search가 수십 년간의 커널 전문성을 Apple Silicon에 이식하는 방법

UC Berkeley Sky Lab의 연구진은 AI를 활용해 GPU 커널을 최적화하는 진화적 커널 검색 프레임워크인 K-Search를 기반으로, 여기에 Apple의 MLX 프레임워크용 백엔드를 확장하여 적용했습니다. 이들은 기존 CUDA 커널을 Apple Silicon용 고품질 커널로 변환하는 혁신적인 구조화된 CUDA-to-MLX 변환 계층을 개발했습니다. 이 접근 방식은 네이티브 MLX 어텐션 커널 대비 0.97배의 속도 향상을 달성했으며, Mamba SSM 커널에서는 커뮤니티 mlx-lm 구현에 비해 프리필 속도가 최대 20배 빨라졌습니다. K-Search는 반복 최적화 루프를 사용하는데, 여기서 대규모 언어 모델(LLM)이 최적화에 대해 추론하고 후보 커널을 생성한 후 실제 하드웨어에서 컴파일 및 벤치마킹을 수행합니다. 변환 계층에는 개념 매핑 테이블, MLX 특화 힌트, 전문가 커널 동작을 보존하는 재사용 가능한 어서션이 포함됩니다. 이 연구는 AI 기반 커널 검색이 처음부터 다시 시작하지 않고도 하드웨어 생태계 간 최적화 지식을 이전할 수 있음을 보여줍니다.