Codex를 활용한 자동 연구: QR 분해 커널 232배 빠르게 달성하기

실행 요약

OpenAI의 Codex를 사용해 자동화된 연구 루프를 구현한 개발자는 배치 처리된 정사각형 컴팩트-하우스홀더 QR 분해에 대해 torch.geqrf 기준 대비 232배의 성능 향상을 달성했다. 성공의 핵심은 일반적인 라이브러리 호출에서 블록화된 하우스홀더 알고리즘으로 전환하고, '후보자 빔(Beam)' 전략을 통해 국소 최적값에 갇히는 것을 방지하며, GPU 모드 popcorn CLI와 Modal 프로파일링이 제공하는 밀접한 피드백 루프를 활용한 점이었다.

최적화 과제: QR 분해

목표는 FP32 CUDA 행렬에 대해 배치 처리된 정사각형 컴팩트-하우스홀더 QR 분해를 구현하는 것이었다. 출력은 H 행렬(상삼각부는 R이고 하삼각부는 하우스홀더 벡터를 저장)과 반사계수를 담은 tau 벡터로 구성된 컴팩트 표현이 필요했다.

시리얼 성능 저하 요인

표준 하우스홀더 QR은 본질적으로 시리얼이다: 반사자 $j+1$는 반사자 $j$가 생성한 행렬에 의존한다. 이 순차적 의존성은 텐서 코어의 효과적인 사용을 방해한다. 작업이 행렬-벡터 형태를 유지하기 때문에 GPU의 가장 강력한 계산 유닛이 비활성 상태로 남아 있게 된다.

블록화된 하우스홀더 해법

시리얼 성능 저하를 극복하기 위해 블록화된 하우스홀더 알고리즘을 사용했다. 이 방법은 시리얼 작업을 $b$개의 열로 구성된 좁은 패널에 국한시킨다. 이후 $b$개의 반사자를 하나의 랭크-$b$ 업데이트(즉, WY 표현)로 압축하여, 행렬의 뒷부분을 세 개의 연속된 GEMM(일반 행렬 곱셈)으로 업데이트할 수 있게 한다. 이로써 계산의 대부분이 텐서 코어가 최대 효율로 실행할 수 있는 행렬-행렬 형태로 전환된다.

자동 연구 방법론: '루프 엔지니어링'

개발자는 14일 동안 1,500건 이상의 제출을 수행한 고주파 반복 루프를 활용했다. 이 과정은 LLM(Codex와 Claude)과 구조화된 하네스의 조합에 의존했다.

하네스 설정 및 방향 조정

• 도구 사용: 개발자는 벤치마킹 및 리더보드 제출을 위해 popcorn CLI를 사용했으며, 프로파일링 및 nsys/NCU 분석을 위해 Modal을 사용했다.
• 로그 기록: 매 제출마다 상태(수락/거절)와 크기별 시간을 기록하는 log.md 파일을 통해 중복 실험을 방지했다.
• 목표 중심 프롬프트: Codex의 /goal 명령어를 사용해 정량적 목표(예: "현재 최고 성능(n=512)보다 빠르게")를 설정하여 모델이 수시간 또는 수일 동안 자율적으로 반복할 수 있도록 했다.
• 감시 및 관리: /btw 및 /side 명령어를 통해 메인 최적화 루프를 중단하지 않고도 에이전트의 진행 상황과 현재 가설을 질의할 수 있었다.

국소 최적값 탈출을 위한 빔 서치

성능이 3,000 $\mu$s 수준에 도달하자, 모델은 자주 국소 최적값에 갇히며 구조적 혁신 대신 미세한 파라미터 조정만 반복했다. 이를 해결하기 위해 개발자는 후보자 빔(Beam) 전략을 도입했다:

• 단일 최고 성능 후보 대신, 에이전트는 3~5개의 활성 "아이디어 가족"을 유지했다.
• 이로 인해 현재 최고 성능보다 처음에는 성능이 낮게 나타나는 고위험 구조적 변경도 조기에 폐기되지 않게 되었다.
• 일반적으로 한 개의 "이득 추구" 빔(세부 조정), 한 개의 "근접 실패" 빔, 한 개의 "구조적/고위험" 빔이 포함되었다.

커널의 기술적 진화

커널은 최종 1,805 $\mu$s의 기하평균 성능에 도달하기 위해 10회의 주요 구조적 돌파구를 거쳤다:

단계	변경 사항	영향
1	torch.geqrf	기준 시작점 (>108.8k $\mu$s)
2	블록화된 WY QR (n=512)	패널 분해 및 뒷부분 업데이트 도입
3	블록화된 경로 (모든 크기)	모든 행렬 크기에 블록화 논리 적용
4	Triton 패널	커스텀 panel16/32 커널 구현
5	Cholesky-ORHR (n=4096)	가장 큰 행렬에 대해 그램-코레스키 사용
6	CUDA 그래프 재생	커널 실행 오버헤드 감소
7	융합된 V/T 레이아웃	슬라이스 복사 및 임시 메모리 제거
8	Split16 패널	그람-슈미트를 통한 꼬리 처리 최적화
9	크기 특화	행 수 하드코딩 및 축소 연산 융합
10	슈퍼패널	V256/T256 패킹 및 직접-H 반환 (1.80k $\mu$s)

핵심 통찰 및 트레이드오프

도메인 전문성 vs. 자동화

에이전트가 상당한 성능 향상을 이끌 수 있었지만, 도메인 전문성은 방향 조정에 필수적이었다. 개발자는 상위 10개 솔루션들이 다음과 같은 방식으로 추가 최적화했다고 지적했다:

• 데이터 탐지: 특정 입력 분포(예: 낮은 랭크 케이스)를 활용.
• 라이브러리 제거: 여전히 남아 있는 PyTorch 함수(예: 삼각형 해법)를 커스텀 CUDA/Triton 구현으로 대체.
• 정밀도 관리: 뒷부분 행렬을 FP16 상태로 유지하여 반복적인 형식 변환을 피함.

일반화 vs. 특정성

커뮤니티 토론에서는 이 '루프 엔지니어링' 접근 방식의 중요한 위험 요소를 지적했다: 벤치마크에 과도하게 최적화되는 것(overfitting).

"상위 10개 솔루션 중 8개는 경쟁용 입력 외에서는 완전히 작동하지 않았다. 유일하게 작동하지 않은 솔루션은... 합리적인 범위 내에서 자신의 솔루션을 조정한 전문가들이 만든 것이었다."

이는 LLM 기반 루프가 특정하고 명확하게 정의된 벤치마크를 해결하는 데 매우 강력하지만, 무거운 인간의 조정 없이는 일반적인 견고성을 유지하기 어려운 점을 시사한다.