Cleave: 분리된 대수적 탐색과 연산자 스케줄링을 통한 텐서 프로그램 최적화 확장

FlashAttention이나 FlashDecoding 같은 최적화된 커널은 오늘날의 거대 모델(large models) 가속에 매우 중요합니다. 이들 대부분은 기존 ML 컴파일러로는 효율성을 따라잡을 수 없기 때문에 전문가들이 직접 작성한 것입니다. 이러한 커널을 생성하려면 여러 감소(reduction)를 포함하는 계산들을 융합해야 하며, 이는 계산 그래프의 대수적 변환과 변환된 그래프의 연산자 스케줄링 둘 다를 필요로 합니다. 불행하게도 이 두 가지를 함께 탐색하면 너무 거대한 공간이 되어 탐색하기 어렵습니다. 우리는 심볼릭 디커플링(symbolic decoupling)을 기반으로 구축된 ML 컴파일러 Cleave를 제안합니다: Cleave는 심볼릭 형태(symbolic shapes)를 가진 그래프에 대한 슈퍼최적화(superoptimization)를 수행하여 변환을 발견하고, 그런 다음 결과로 나온 각 그래프를 구체적인 형태(concrete shapes)에서 스케줄링합니다. 형태를 심볼로 표현하는 것은 동등성 검사(equivalence checking)를 저렴하게 만들고, 심볼릭 분할 횟수(symbolic split count)를 가진 새로운 Split 연산자를 통해 감소 차원(reduction dimension)을 따라 병렬화할 수 있게 합니다. Cleave의 스케줄러는 반복적인 타일링(iterative tiling)과 수평 융합(horizontal fusion)을 통해 여러 감소를 포함하는 그래프들을 융합합니다. 일반적인 LLM 서브그래프에 대한 평가 결과, Cleave가 생성한 커널은 최고 기준선보다 최대 2.8배 빠르며 (평균 1.6배), Mirage 대비 컴파일 시간은 평균 5.9배 단축되는 것으로 나타났습니다. 프로덕션 서비스 추적(production serving traces)에서 포착된 동적 워크로드의 경우, Cleave는 각 연산자를 한 번만 컴파일하여 FlashInfer가 수작업으로 작성한 FA2 및 FA3 백엔드 대비 기하 평균 속도 향상률 1.4배와 1.7배를 달성했습니다. Cleave 코드는 다음에서 이용 가능합니다: https://github.com/nyu-systems/cleave