LLM 기반 Triton 커널 최적화를 위한 컴파일러 기반 계층적 진단
요약
LLM 기반의 Triton 커널 최적화를 위해 컴파일러 피드백과 IR 구조를 연결하는 계층적 진단 프레임워크를 제안합니다. Ascend NPU 환경에서 실험한 결과, 벤치마크 커널들에 대해 평균 4.35배의 속도 향상을 달성했습니다.
핵심 포인트
- 표면적 프로파일링을 넘어 IR 및 컴파일러 동작을 분석하는 교차 계층 진단 방식 도입
- 경량 패턴 분류부터 심층 컴파일러 분석까지 단계별로 진행되는 계층적 최적화 프레임워크
- Ascend NPU 기반 Triton 커널에서 기하 평균 4.35배의 성능 향상 입증
- 런타임 증상을 소스 수준의 재작성으로 연결하는 자동화된 최적화 프로세스
최근 대규모 언어 모델 (LLMs)의 발전으로 자동화된 커널 생성 및 최적화가 가능해졌으나, 기존의 대부분의 접근 방식은 컴파일 피드백(compilation feedback) 및 프로파일링 지표(profiling metrics)와 같은 표면적인 신호에 의존합니다. 이러한 신호들은 커널이 느리다는 점은 알려주지만, 특히 NPU와 같은 신흥 가속기(accelerators)에서 백엔드 컴파일러가 왜 수익성 있는 최적화를 실현하지 못하는지에 대한 이유는 알려주지 않습니다. 따라서 우리는 소스 코드를 다시 작성하기 전에 런타임 증상(runtime symptoms)을 IR 구조 및 컴파일러 동작과 연결하는 점진적인 교차 계층 진단(cross-layer diagnosis) 문제로 커널 최적화를 정식화합니다. 이러한 통찰을 바탕으로, 우리는 Triton 커널을 위한 컴파일러 기반의 계층적 최적화 프레임워크인 우리 시스템을 제시합니다. 이 시스템은 더 깊은 증거가 필요할 때만 경량 패턴 분류(pattern triage) 및 프로파일링 진단에서 IR 귀속(IR attribution) 및 컴파일러 기반 분석으로 단계를 격상하며, 그 후 증거에 기반한 소스 수준의 재작성(source-level rewrites)을 제안합니다. 우리는 Ascend NPU를 위한 Triton에 이 시스템을 구현하였으며, 표준화된 NPUKernelBench에서 파생된 Ascend 950 벤치마크의 성공적으로 변환된 37개 항목에 대해 평가를 수행했습니다. 이 항목들 전반에 걸쳐, 시스템은 초기 커널 대비 최적화된 Triton 커널로부터 4.35$ imes$의 기하 평균 속도 향상(geometric-mean speedup)과 2.73$ imes$의 중앙값 속도 향상(median speedup)을 달성했습니다. 37개 중 22개는 2$ imes$를 초과했고, 13개는 5$ imes$를 초과했습니다. 전체 분포는 베이스라인에 가까운 항목부터 큰 성능 향상을 보인 항목까지 다양하며, 이는 현재 시스템의 범위와 한계에 대한 투명한 보고의 필요성을 시사합니다.
AI 자동 생성 콘텐츠
본 콘텐츠는 arXiv cs.PL (Programming Languages)의 원문을 AI가 자동으로 요약·번역·분석한 것입니다. 원 저작권은 원저작자에게 있으며, 정확한 내용은 반드시 원문을 확인해 주세요.
원문 바로가기