[Paper] TritonForge: 프로파일링 기반 자동 Triton 커널 최적화 프레임워크

Source: arXiv - 2512.09196v1

Overview

TritonForge는 Triton이라는 고성능 딥러닝 프리미티브를 위한 인기 DSL로 작성된 GPU 커널을 자동으로 튜닝하는 새로운 프레임워크입니다. 정적 커널 분석과 실시간 프로파일링 피드백을 결합하여 코드를 반복적으로 재작성함으로써 병목 현상을 제거하고, 개발자가 GPU 아키텍처 전문가가 아니어도 5배까지 속도 향상을 달성할 수 있습니다.

Key Contributions

• 프로파일링 기반 최적화 루프: 런타임 메트릭을 소스 레벨 변환에 직접 연결합니다.
• 모듈형 코드 생성 파이프라인: 어떤 추론 엔진도 플러그인할 수 있습니다(프로토타입은 대형 언어 모델을 사용하지만 설계는 모델에 구애받지 않습니다).
• 자동 병목 탐지: 서브 최적 블록 크기, 메모리 레이아웃 불일치, 부족한 명령 수준 병렬성 등 흔히 발생하는 Triton 함정을 자동으로 감지합니다.
• 실증적 검증: 다양한 GPU 세대에서 매트릭스 곱셈, 컨볼루션, 어텐션 등 여러 커널을 대상으로 평균 1.76× 성능 향상과 최고 5× 향상을 보여줍니다.
• 오픈소스 프로토타입: 개발자가 기존 CI 파이프라인에 TritonForge를 통합하여 지속적인 성능 회귀 테스트를 수행할 수 있는 방법을 제공합니다.

Methodology

1. Static Analysis – TritonForge는 Triton 소스를 파싱해 고수준 IR(루프 네스트, 메모리 접근, 스레드/블록 구성)을 추출합니다.
2. Initial Profiling – 커널을 컴파일하고 대상 GPU에서 실행하면서 가벼운 프로파일러가 점유율, 메모리 대역폭, 스톨 원인 등의 메트릭을 수집합니다.
3. Bottleneck Classification – 규칙 기반 휴리스틱(예: 낮은 점유율 → 블록 크기 증가)과 선택적인 LLM 기반 추론을 사용해 성능을 저해하는 코드 패턴을 pinpoint합니다.
4. Transformation Generation – 타일링 조정, 공유 메모리 버퍼링, 루프 언롤링, 데이터 타입 변경 등 후보 재작성안을 생성합니다.
5. Iterative Evaluation – 변환된 각 커널을 다시 컴파일·프로파일링하고 현재까지 최고의 결과와 비교합니다. 더 이상의 개선이 없거나 타임아웃에 도달하면 루프를 종료합니다.

루프가 실제 런타임 데이터를 기반으로 구동되기 때문에, 옵티마이저는 Ampere와 Hopper와 같은 서로 다른 GPU 마이크로아키텍처의 특성에 맞게 자동으로 조정될 수 있으며, 아키텍처‑특정 규칙을 하드코딩할 필요가 없습니다.

Results & Findings

• Average improvement: 모든 테스트 커널에서 평균 1.76× 향상.
• Optimization time: 단일 GPU당 커널당 약 2–5 분 소요, CI 파이프라인에 실용적.
• Model‑agnosticity: LLM을 더 단순한 규칙 엔진으로 교체해도 파이프라인은 동작하지만 성공률이 약 12 % 감소, 이는 핵심 프로파일링‑피드백 루프가 성능 향상의 주요 동인임을 확인합니다.

Practical Implications

• Developer productivity: 개발자는 간단한 Triton 커널을 작성하고 TritonForge가 저수준 튜닝을 담당하므로 알고리즘 혁신에 집중할 수 있습니다.
• Performance portability: 동일한 Triton 소스를 새로운 GPU 세대에 맞게 자동으로 재튜닝할 수 있어 하드웨어 업그레이드 시 유지보수 부담이 감소합니다.
• CI/CD integration: 최적화 루프가 결정적이고 비교적 빠르기 때문에 지속적 통합 단계에 추가해 성능 회귀를 조기에 감지할 수 있습니다.
• Cost savings: 더 빠른 커널은 대규모 학습·추론 워크로드에서 클라우드 GPU 비용 절감으로 직결됩니다.
• Foundation for higher‑level tools: TritonForge의 모듈형 설계는 다른 DSL(CUDA‑Python, JAX XLA 등)이나 에너지·지연 시간과 같은 도메인‑특화 비용 모델을 통합하는 도구로 확장될 수 있습니다.

Limitations & Future Work

• Reliance on profiling accuracy: 프로파일링 카운터가 제한된 GPU에서는 병목 분류기가 스톨을 오해해 비최적 재작성으로 이어질 수 있습니다.
• Search space explosion: 현재 휴리스틱 탐색은 매우 불규칙한 커널에 대해 전역 최적 구성을 놓칠 수 있으며, 베이지안 최적화와 같은 고급 탐색 알고리즘을 도입할 계획입니다.
• LLM dependence: 프레임워크 자체는 모델에 구애받지 않지만, 프로토타입의 최고 성능은 여전히 LLM에 의존합니다. 추론 비용을 낮추기 위해 경량 정적 분석 대안을 연구할 예정입니다.
• Broader benchmark coverage: 현재 평가는 선별된 커널 집합에 국한되어 있으므로, 전체 트랜스포머 학습 루프와 같은 엔드‑투‑엔드 모델에 확장해 실사용 영향을 검증할 필요가 있습니다.

TritonForge는 데이터‑드리븐 자동 튜닝이 전문가 수준의 GPU 성능을 일상 개발자에게도 제공할 수 있음을 보여주며, 보다 확장 가능하고 유지보수하기 쉬운 고성능 ML 코드베이스 구축의 길을 열어줍니다.

Authors

• Haonan Li
• Keyu Man
• Partha Kanuparthy
• Hanning Chen
• Wei Sun
• Sreen Tallam
• Chenguang Zhu
• Kevin Zhu
• Zhiyun Qian

Paper Information

• arXiv ID: 2512.09196v1
• Categories: cs.SE
• Published: December 9, 2025
• PDF: Download PDF