기술 요약: GPU 포캐스터(Forecasters): 커널 런타임 최적화를 위한 선택적 대리 모델로서의 언어 모델

문제 정의

GPU 커널은 현대 딥러닝의 핵심적인 작업 동력이며, 여기서 발생하는 미세한 효율성 개선은 훈련 및 추론 비용의 상당한 절감으로 이어집니다. 이러한 커널을 최적화하는 과정은 일반적으로 대규모 언어 모델(LLM)이 후보 커널을 제안하고, 이를 컴파일하여 물리적 하드웨어에서 실행해 성능을 측정하는 생성-평가 사이클을 포함합니다. 그러나 이 평가 과정은 심각한 병목 구간입니다. 각 평가는 컴파일(CUDA의 경우 60~120초 소요)과 장치 내 실행을 필요로 하므로, 평가 처리량은 생성 속도보다 수십 배 더 느립니다. 탐색 예산이 확장됨에 따라, 반복적인 물리적 측정 비용은 잠재적인 커널 구현의 방대한 공간을 탐색하는 능력을 제한합니다.

기존 연구들이 더 나은 커널을 '생성'하기 위해 LLM을 사용하는 데 집중했다면, 본 논문은 이와 상보적인 문제인 '평가' 비용을 줄이는 문제를 다룹니다. 저자들은 LLM을 단순한 생성기가 아니라, 물리적 GPU를 위한 **선택적 대리 모델(selective surrogates)**로 사용하는 것을 제안합니다. 목표는 후보 커널의 상대적 성능을 참조 구현(reference implementation)과 비교하여 예측하고, 모델이 불확incity(불확실성)를 느끼는 시점을 식별하여 불확실한 경우에만 비싼 물리적 GPU를 사용하도록 미루는 것입니다.

방법론

과업 정식화: 선택적 예측(Selective Forecasting)

저자들은 이 문제를 하드웨어 사양을 바탕으로 후보 커널의 상대적 가속도(Speedup, $S_i = T_{ref} / T_i$)를 예측하는 문제로 정식화했습니다. 하드웨어의 마이크로 아키텍처 세부 사항과 컴로 컴파일러 동작의 복잡성 때문에 절대적인 실행 시간(초 단위)을 예측하는 것은 실행 불가능하다고 판단했습니다. 대신, 타겟을 8개의 서수적 가속도 빈(bin)으로 이산화했습니다( "심각한 성능 저하"부터 "극단적인 가속"까지의 범위).

대리 모델(LLM)은 참조 커널, 후보 커널, 그리고 하드웨어 사양을 입력으로 받습니다. 모델은 추론 체인(chain of reasoning)을 생성하고 8개의 빈에 대한 확률 분포를 출력합니다. 모델이 명시하는 신뢰도는 단일 빈에 할당된 최대 확률로 정의됩니다.

선택적 예측 전략

시스템은 선택적 예측 프레임워크 하에서 작동합니다. 대리 모델은 후보들을 순위 매기는 데 사용됩니다. 만약 대리 모델의 신뢰도가 특정 임계값을 초과하면, 해당 예측을 수용하고 후보를 GPU에서 측정하지 않습니다. 만약 신뢰도가 낮으면(불확실성을 나타냄), 해당 후보는 실제 성능 측정을 위해 물리적 GPU로 전달됩니다. 이를 통해 고정된 GPU 측정 예산 내에서 훨씬 더 많은 후보를 고려할 수 있습니다.

보정(Calibration)을 위한 강화 학습(RL)

대리 모델의 정확도, 특히 결정적으로 보정(명시된 신뢰도와 실제 정확도 사이의 일치)을 개선하기 위해, 저자들은 오픈 웨이트 모델(GPT-OSS-20B)에 강화 학습을 적용합니다. 이들은 다음을 기준으로 보상을 정의합니다:

1. 정확성(Correctness): 예측된 빈이 실제 측정된 빈과 일치하는지 여부.
2. 범주형 보정(Categorical Calibration, Brier Score): 실제 원-핫(one-hot) 타겟과의 거리에 따라 분포에 페널티를 부여함.
3. 거리 인식 보정(Distance-Aware Calibration, CRPS): 누적 분포에 페널티를 부여하며, 가속도의 서수적 특성을 고려함 (예: 실제 가속도가 1.5배일 때 2배라고 예측하는 것은 4배라고 예측하는 것보다 덜 심각한 오류임).

학습 데이터는 이전 검색 실행에서 성공한 후보들을 학습 예제로 사용하여, 커널 검색 과정의 부산물로 구축되었습니다.

주요 기여

1. GPU 대리 모델로서의 LLM: 본 논문은 LLM이 물리적 GPU가 후보 커널에 어떻게 반응하는지를 예측하는 데 있어 유의미한 정확도를 달al 수 있음을 입증하며, 측정 대신 예측을 기반으로 후보의 순위를 매김으로써 가용한 가속도의 상당 부분을 회복할 수 있음을 보여줍니다.
2. 선택적 예측 프레임워크: 저자들은 GPU 예측을 선택적 예측 문제로 규정하고, 대리 모델이 보정된 불확실성을 보고할 수 있음을 보여주었습니다. 이는 대리 모델이 높은 신뢰도의 예측은 신뢰하고, 불확실한 경우에만 물리적 측량을 예약하는 탐색 전략을 가능하게 합니다.
3. 대리 모델 행동에 미치는 RL의 영향: 실험을 통해 RL 학습이 대리 모델의 특성을 어떻게 변화시키는지 규명했습니다. 구체적으로, 보정 보상(Brier 또는 CRPS)을 추가하면 신뢰도 추정의 신뢰성이 향상되는 반면, 정확성만을 학습시키면 신뢰도와 정확도 사이의 관계가 악화될 수 있음을 보여주었습니다.
4. 엔드 투 엔드(End-to-End) 탐색 개선: 대리 모델을 커널 검색 루프(특히 PUCT 기반 탐색)에 통합하면, 물리적 GPU 측정 없이도 훨씬 더 많은 후보를 탐색할 수 있음을 보여줍니다. 대부분의 작업에서, 이는 동일한 GPU 예산 하에서 모든 제안을 측정하는 베이스라인보다 더 빠른 커널을 찾아냅니다.
5. 데이터셋 공개: 저자들은 118개의 문제, 다양한 GPU 유형(A100, L40S, H100) 및 다양한 생성 방법을 아우르는, 측정된 실행 시간이 포함된 12,388개의 LLM 생성 GPU 커널 컬렉션을 공개합니다. 이 데이터셋은 GPU 성능 예측 및 커널 탐색에 관한 광범위한 연구를 지원합니다.

결과

• 예측 정확도: 기성 LLM(예: GPT-OSS-20B)은 상위 1%의 후보만을 GPU에서 측정할 때 가용한 최적 가속도의 약 83%를 회복합니다. 이 성능은 상위 50%를 측정할 때 94%까지 향상됩니다.
• 보정 및 RL: RL 학습은 대리 모델의 행동을 크게 변화시킵니다.
  • 정확성 중심 학습은 정확한 빈(bin) 적중률은 높이지만, 신뢰도와 오차 사이의 상관관계를 저하시킵니다 (모델이 틀린 예측에 대해 과잉 확신을 갖게 됨).
  • **보정 학습(Brier/CRPS)**은 예측 오차와 신뢰도 간의 단조성(monotonicity)을 개선합니다. 특히, Brier로 보정된 학습은 '회복된 가속도' 지표를 베이스 모델보다 약 2.1 퍼센트 포인트 향상시키는데, 이는 더 잘 보정된 분포가 탐색 알고리즘에 더 나은 순위 신호를 제공함을 시사합니다.
• 탐색 효율성: 엔드 투 엔드 커널 탐색에서, 대리 모델을 활용한 탐색은 6개 작업 중 4개에서 베이스라인보다 더 빠른 커널을 찾아내며, 종종 베이스라인의 최고 가속도에 도달하면서도 실질적으로 더 적은 GPU 측정을 사용합니다. 베이스라인이 승리하는 작업의 경우, 탐색이 조기에 포화 상태에 도달하는 것으로 보아 대리 모델은 탐색 공간이 크고 미개척 상태일 때 가장 유용합니다.
• 발견 모멘트(Discovery Moments): 대리 모델은 무작위 선택보다 더 잘 "발견 모멘트"(자식 커널이 부모 커널보다 현저히 빨라지는 변이)를 식별할 수 있지만, 정밀도 측면에서 여전히 과제가 남아 있어 고위험 결정에는 GPU 확인이 필수적입니다.

의의 및 주장

본 논문은 LLM이 단순한 코드 생성을 넘어 커널 최적화에서 더 넓은 역할을 수행할 수 있다고 주장합니다. 가상 GPU 모델로서 작용함으로써, LLM은 물리적 하드웨어 자원의 희소성을 관리하는 고처리량, 저비용 필터 역할을 할 수 있습니다.

저자들은 이 접근 방식의 주요 가치가 단순히 속도에 있는 것이 아니라, 매 단계마다 원래의 시스템을 호출하지 않고도 모델 기반 최적화를 수행할 수 있다는 점에 있다고 강조합니다. 대리 모델은 물리적 GPU보다 높은 처리량과 낮은 쿼리당 비용을 제공하며, 이는 보정된 불확실성 추정을 통해 관리 가능한 수준의 예측 노이즈를 트레이드오프(trade-off)하는 것입니다. 이 연구는 커널 최적화를 복잡한 물리 시스템을 위한 신경망 시뮬레이터를 구축하는 광범위한 분야와 연결하며, 하드웨어 측정을 선택적으로 사용해야 할 희소한 자원으로 취급하는 옵티마이저로 나아가는 경로를 제시합니다.