Compiler-Grounded Hierarchical Diagnosis for LLM-Based Triton Kernel Optimization

기술 요약: LLM 기반 Triton 커널 최적화를 위한 컴파일러 기반 계층적 진단

1. 문제 정의

대규모 언어 모델(LLM)이 자동화된 커널 생성 및 최적화 분야를 발전시켜 왔으나, 기존 방식은 주로 컴파일 에러, 런타임 메트릭, 프로파일링 데이터와 같은 "표면적 신호(surface signals)"에 의존합니다. 이러한 신호들은 커널이 느리다는 사실은 알려주지만, 왜 백엔드 컴파일러가 수익성 있는 최적화를 실현하지 못하는지에 대한 이유는 거의 설명하지 못합니다. 이는 특히 컴파일 스택이 진화 중이고 타겟 특화적이며 개발자에게 불투명한 NPU(예: Ascend)와 같은 신흥 가속기에서 더욱 두드지게 나타납니다.

식별된 핵심 과제는 연산자 수준의 소스 코드와 백엔드 성능 동작 사이의 **의미론적 간극(semantic gap)**입니다. 성능 실패는 종종 코드가 중간 표현(IR) 단계, 스케줄링 결정, 그리고 백엔드 특화 컴파일 패스를 통과한 후에야 나타납니다. 결과적으로, GPU 중심 패턴에 사전 학습된 LLM 에이전트들은 시행착오식 최적화 도구로 작동하게 됩니다. 즉, 그들은 그럴듯한 소스 코드 수정을 제안하지만, 런타임 증상을 특정 컴파일 수준의 원인으로 매핑할 수 있는 근거가 부족합니다.

2. 방법론

본 논문은 커널 최적화를 **점진적인 교차 계층 진단 문제(progressive cross-layer diagnosis problem)**로 취급하는 컴파일러 기반의 계층적 최적화 프레임워크를 제 제안합니다. 시스템은 컴파일러 기반의 근거(grounding)를 항상 켜두는 기능이 아니라, 필요할 때만 경량 증거에서 심층 컴파일 분석으로 격상시키는 방식으로 작동합니다.

계층적 워크플로우:
시스템은 네 가지 단계로 구성된 폐쇄 루프 파이프라인을 통해 작동합니다.

• L1: 패턴 트리아지(Pattern Triage): 시스템은 먼저 큐레이션된 최적화 패턴, 휴리스틱, 전문가 규칙(이전 실행에서 추출된 내부 휴리스틱 포함)을 사용하여 패턴 기반의 강력한 가설을 확인합니다. 저위험 재작성이 기존 증거에 의해 정당화되면 프로세스가 여기서 중단됩니다.
• L2: 프로파일링 진단(Profiling Diagnosis): 패턴이 실패할 경우, 시스템은 런타임 증거(지연 시간, 처리량, 하드웨어 카운터)를 수집하여 병목 지점(예: 메모리 이동, 스칼라 활동, 동기화)을 국소화합니다. 이는 증상은 식별하지만 메커니즘까지는 식별하지 못할 수 있습니다.
• L3: IR 귀속(IR Attribution): 프로파일링이 불충분할 경우, 시스템은 중간 표현(IR)을 검사하여 런타임 증상을 구조적 진단(예: 비단위 스트라이드 로드, 퓨전 누락, 레이아웃 유발 벡터화 실패)으로 변환합니다. 이는 연산자의 로워링(lowering) 궤적을 재구성하여 단계별 변환을 식별합니다.
• L4: 컴파일러-소스 격상(Compiler-Source Escalation): 최후의 수단으로, 시스템은 진단된 패턴이 왜 백엔드 최적화를 차단하는지를 설명하기 위해 컴파일러 기반 분석을 호출합니다. 시스템은 컴파일러 지식(패스 동작, 로워링 규칙, 제약 조건)을 참조하여 성공적인 재작성이 충족해야 하는 소스 수준의 조건을 복구합니다.

시스템 아키텍처 및 상태 관리:

• 에이전트 중심 설계: 이를 맞춤형 컴파일러 서비스가 아닌, 에이전트 중심 환경(CannBot 스킬 저장소에 통합됨)으로 구현했습니다. 시스템은 검증, 정확성 및 증거 수집을 관리하기 위한 얇은 실행 하네스를 사용하며, 도메인 특화 "스킬"이 최적화 로직을 처리합니다.
• 상태 기반 라운드(Stateful Rounds): 최적화 프로세스는 명시적이고 감사 가능한 라운드로 구성됩니다. 각 라운드는 가설, 증거 출처 및 검증 결과를 기록합니다. 이러한 상태 기반 설계 덕분에 시스템은 재개 가능하며, 심층 분석(IR/컴파일러)이 일시적인 관찰이 아닌 재사용 가능한 증거로서 부착되도록 보장합니다.
• 회고적 합성(Retrospective Synthesis): 온라인 루프 외에도, 시스템은 더 느린 교차 연산자 합성 패스를 수행합니다. 검증된 최적화 이력을 기존 패턴 패밀리에 매핑하거나 새로운 재사용 가능한 패턴을 생성함으로써, 가장 "저렴한" 단계(Pattern Triage)가 시간이 지남에 따라 개선되도록 합니다.

3. 주요 기여

1. 문제 정식화: 저자들은 LLM 기반 커널 최적화를 점진적인 교차 계층 진단 문제로 정식화하였으며, 프로파일링 전용 에이전트가 백엔드 특화 최적화에서 실패하는 주요 이유로 컴파일러 기반 근거의 부재를 지목했습니다.
2. 계층적 워크플로우: 패턴 트리아지에서 프로파일링, IR 귀속, 컴파일러-소스 격상으로 이어지는 워크플로우를 설명하며, 더 깊은 증거가 필요할 때만 증거에 기반한 소스 수준 재작성을 제안합니다.
3. 평가 증거: Ascend 950 벤치마크(NPUKernelBench 유래)를 사용하여 종합적인 이득, 최적 라운드 계정 및 대표적인 교차 계층 트레이스를 보고하며 표준화된 분포 인식 평가 증거를 제시합니다.

4. 실험 결과

시스템은 Ascend NPU용 Triton으로 구현되었으며, 표준화된 Ascend 950 벤치마크(NPUKernelBench 유래)에서 추출한 37개의 성공적으로 변환된 연산자에 대해 평가되었습니다. 각 연산자는 5개의 케이스(총 185개 케이스)에 대해 테스트되었습니다.

• 성능 이득:
  • 기하 평균 속도 향상: 4.35× (초기 Triton 커널 대비 최적화된 커널).
  • 중앙값 속도 향상: 2.73×.
  • 분포: 37개 항목 중 22개가 2배 이상의 속도 향상을 보였고, 13개는 5배 이상을 기록했습니다. 결과는 긴 꼬리 분포를 보이며, 상당한 이득과 베이스라인에 근접한 항목이 공존함을 보여줍니다.
• 프로세스 계정:
  • 중앙값 최적 라운드 인덱스는 10으로, 최적의 결과에 도달하기 위해 여러 번의 반복이 필요했음을 나타냅니다.
  • 증거 깊이: 선택된 "최적" 라운드 중 33개는 패턴 트리아드에 의한 것이었으나, 시스템은 더 깊은 계층의 유용성을 입증했습니다. 32개의 궤적이 프로파일링 진단에 도달했고, 3개는 IR 귀속에, 4개는 컴파일러-소스 격상에 도달했습니다.
  • 사례 연구: 대표적인 트레이스(Hyena FFT-size padding)는 프로파일링이 어떻게 스칼라 병목을 식별하고, IR 귀속이 어떻게 로워링에서의 특정 동기화 비용을 식별하며, 컴파일러-소스 분석이 어떻게 제약 조건을 설명하여 견고한 2.08× 속도 향상을 이끌어냈는지 보여주었습니다.

5. 의의 및 주장

본 논문은 컴파일러 근거(compiler grounding)가 항상 켜져 있기보다는 선택적으로 적용되어야 한다고 주장합니다. 최적화를 점진적인 진단으로 처리함으로써, 시스템은 심층 컴파일 추론이 가벼운 무게의 증거로 충분한 근거를 확보하지 못할 때만 호출되도록 하여 심층 컴파일 추론을 실용적으로 만듭니다.

저자들은 결과가 일부 연산자에서 **상당한 헤드룸(headroom)**을 보여준다는 점을 강조하지만, 모든 케이스에서 균일한 개선을 주장하지는 않습니다. 저자들은 현재의 평가가 인과적 격리(예: 고정된 툴체인 하에서 점진적 격상 대 프로파일링 전용 또는 항상 심층 분석을 사용하는 전략 간의 비교)를 제공하지 않는다는 점을 명시적으로 밝힙니다. 본 논문은 투명한 벤치마크 및 프로세스 증거를 뒷받침하는 "시스템적 논증"으로 자리매김하며, 계층 구조의 인과적 분해는 향후 통제된 실험을 위해 남겨둡니다.

이 연구는 소스 코드와 런타임 성능 사이의 관계가 보이지 않는 컴파일러 변환에 의해 매개된다는 점을 강조하며, 이러한 간극을 계층적이고 증거 중심적인 접근 방식으로 메우는 것이 NPU와 같은 신흥 가속기에서 커널을 최적화하는 데 필수적임을 보여줍니다.