KernelBand: Boosting LLM-based Kernel Optimization with a Hierarchical and Hardware-aware Multi-armed Bandit
Created by
Haebom
Category
Empty
저자
Dezhi Ran, Shuxiao Xie, Mingfang Ji, Ziyue Hua, Mengzhou Wu, Yuan Cao, Yuzhe Guo, Yu Hao, Linyi Li, Yitao Hu, Tao Xie
개요
KernelBand는 대규모 언어 모델(LLM)의 훈련 및 추론 비용을 줄이기 위해 고품질 커널을 최적화하는 새로운 프레임워크입니다. LLM 기반 코드 생성을 활용하여 커널 최적화를 수행하지만, 하드웨어 도메인 지식 부족으로 인한 탐색과 활용의 불균형 문제를 해결합니다. KernelBand는 커널 최적화를 계층적 다중 밴딧 문제로 공식화하여, LLM 에이전트가 커널 선택과 최적화 전략 적용을 순차적 의사 결정 프로세스로 처리하도록 합니다. 하드웨어 프로파일링 정보를 활용하여 유망한 최적화 전략을 식별하고, 런타임 동작 클러스터링을 통해 커널 후보 간 탐색 오버헤드를 줄입니다. TritonBench 실험에서 KernelBand는 기존 방법보다 우수한 성능을 보이며, 적은 토큰으로 더 나은 성능을 달성하고, 계산 자원 증가에 따라 지속적인 개선을 보입니다.
시사점, 한계점
•
시사점:
◦
LLM 기반 커널 최적화에서 하드웨어 지식 부족 문제를 해결하는 새로운 프레임워크 제시
◦
커널 선택 및 최적화 전략 적용을 다중 밴딧 문제로 공식화하여 최적화 공간 탐색 효율성 향상
