Jalapeño MLA 커널이 HotChips에서 발표되고 SemiAnalysis의 후속 논평이 나온 이후로 많은 논의가 있었습니다. OpenAI의 하드웨어 팀으로서, 우리는 이 작은 금덩어리를 아주 살짝만 건드렸습니다: AI가 우리의 커널을 작성하고 있으며, 그럴 때 우리가 커널이 하는 일을 한 줄 한 줄 이해할 필요가 정말로 없다는 사실입니다. 우리는 다음과 같은 점을 눈에 띄게 빠뜨렸습니다: 어떻게 이런 일이 가능한가? 전통적인 코드 생성 방식과 비교하여 이것을 올바르게 이해하는 방법은 무엇인가? 최적화된 커널이 최적화되지 않은 커널만큼 건전한가?
제 배경을 말씀드리자면, 저는 10년 넘게 가속기용 컴파일러를 연구해왔습니다. 저는 XLA를 시작했는데, 이는 훌륭한 컴파일러 인프라로, 뛰어난 회사 간 팀과 노력이 함께하고 있습니다. 지난 2년 넘게 OpenAI에서 저는 AI 시대에 컴파일러가 어떻게 작동해야 하는지에 대한 개념을 재정립하려고 노력해왔습니다. 새로운 컴파일러 공식은 기존의 강점을 활용하겠지만, 도구 상자에 활용할 강력한 새 도구가 있다는 사실을 부인할 수 없습니다.
이것은 다소 긴 여정이 될 수 있지만, 컴퓨터 프로그램 개선의 자동화, 즉 최적화 컴파일을 위해 AI가 어떻게 사용되고 있는지 조명하고자 합니다. 저는 AI를 통해 우리가 "컴파일러 2.0"이라고 부를 수 있는 것을 경험할 수 있다고 생각합니다. AI가 제안할 수 있는 것에는 근본적인 제약이 덜하며, 그 제안은 모델의 훈련과 맥락의 결과입니다. 이 때문에 저는 이를 "확률적 최적화 도구"로 분류합니다. 이는 어려움을 초래할 수 있지만, 앞으로 보게 될 것처럼 큰 강점의 원천이기도 합니다...
많은 학술 연구와 산업 응용이 이미 이 방향으로 나아가고 있으며, 최적화 컴파일러 영역에서 AI의 참여 가능성을 빠르게 발견하고 있습니다. 하지만 우리는 이에 대한 광범위한 설명이 필요한 시점에 있습니다.
배경
컴파일러는 프로그램을 입력받아 해당 프로그램의 번역되거나 개선된 버전을 출력합니다.
입력과 출력 양쪽의 프로그램에는 의미론이 있어 프로그램이 무엇을 의미하는지, 무엇을 할 수 있는지, 그리고 그렇게 할 수 있는 것들에 대해 어떻게 추론할지 알려줍니다.
컴파일러를 연구하는 사람들은 컴파일러를 순수 함수처럼 생각합니다. 데이터 구조를 입력받아 해당하는 의미론을 가져야 하는 데이터 구조를 출력합니다.
때때로 우리의 컴파일러는 "낮춤" 또는 "번역"에 초점을 맞춥니다. 예를 들어, C를 입력받아 x86-64 어셈블리를 출력할 수 있으며, 이를 종종 "더 낮은 수준"이라고 간주합니다. 하지만 실제로는 그 과정의 하위 부분으로 단순한 번역 이상을 수행하는 경우가 많습니다...
실제로 우리의 컴파일러는 "최적화"에 초점을 맞춥니다. 프로그램을 나타내는 데이터 구조(우리의 용어로 "중간 표현(IR)")를 입력받아 해당 프로그램의 더 나은 버전을 생성하려고 시도합니다. 때로는 "더 나은" 것이 실행 사이클이 더 적게 걸리는 것을 의미하고, 때로는 불필요한 코드가 더 적은 것을 의미하며, 때로는 프로그램에 대해 "반드시 참이어야 하는" 것들을 증명할 수 있는 것에 특화하는 것(부분 평가)을 의미합니다.
이제 잠시 생각해보면, LLM은 원래 인간의 텍스트를 한 언어에서 다른 언어로 번역하기 위해 만들어졌습니다. 분명히 번역은 그들의 전문 분야입니다. 그리고 일상적인 작업에서 LLM을 사용함으로써 그들이 새로운 해결책을 작성하고 기존 해결책을 개선할 수도 있다는 것을 알 수 있습니다. 많은 코더들 또한 LLM에게 "이 코드 조각을 최적화해줘"라고 요청한 경험이 있으며, 놀랍게도 그렇게 할 수 있습니다. (하지만 코드가 올바르게 최적화되었는지 확인해야 하며, 이에 대해서는 나중에 다루겠습니다!) 이는 LLM이 최적화 컴파일러에서 찾는 능력을 가지고 있음을 강조하기 위한 것입니다.
최적화와 최적성
최적화 컴파일러는 당연히 작업 중인 프로그램의 최적성을 높이기 위해 노력하며, 일반적으로 실행 시간이라는 목표를 기준으로 합니다. 일반적인 경우, 임의의 프로그램에 대해 이것을 수행하는 것은 매우 어려워서 컴파일러 엔지니어를 위한 완전 고용 정리라는 정리가 있습니다. (저는 컴파일러 엔지니어가 되기로 선택한 후에야 이것을 알게 되었지만, 여전히 위로가 되었습니다!)
"슈퍼옵티마이저"는 최적화 컴파일러의 놀라운 작은 하위 분야입니다. 주어진 프로그램이 있고 의미론을 통해 그것이 무엇을 하는지 말할 수 있다고 상상해보세요. 동일한 의미론을 가진 가장 최적의 프로그램은 무엇일까요? 이것이 슈퍼옵티마이저가 해결하려고 시도하는 것이며, 본질적으로 검색 문제입니다...
동일한 의미론을 가진 가장 짧은 프로그램을 찾으려고 하고, 후보 프로그램이 동일한 의미론을 가지고 있는지 물어볼 수 있는 방법이 있다고 상상해보세요. 가상적으로, 모든 프로그램을 목표 순서대로 열거하고 동일한 의미론을 가진 가장 작은 프로그램을 선택할 수 있습니다.
하지만 모든 프로그램을 목표 순서대로 열거하는 것은 상당히 다루기 어려워 보입니다. 2013년에 만들어진 "STOKE"(확률적 슈퍼최적화)라는 제가 가장 좋아하는 학술 논문 중 하나는 다음과 같이 질문했습니다: "음, 프로그램을 계속해서 무작위로 조정하면 결국 최고의 프로그램을 관찰할 수 있을까?" 그들은 무작위 보행(그리고 우리의 원조 머신러닝 친구인 Markov Chain Monte Carlo / Metropolis-Hastings)을 통해 결국 최적의 프로그램을 볼 수 있다고 제안했습니다.
Monte Carlo 조정은 일반적으로 어리석지만(무작위로 조정을 선택), 빠릅니다. LLM은 매우 똑똑하지만(많은 추론 토큰), 상대적으로 느립니다.
어리석고 빠른 Monte Carlo 조정 대신 LLM이 프로그램을 이끌어갈 방향을 알아내도록 하면 어떨까요? 우리는 프로그램을 최적화된 프로그램 공간으로 안내하는 매우 지능적인 확률적 최적화 도구를 갖게 될 것입니다.
최적화에 대한 직관
잠시 뒤로 물러서 봅시다. 여러분이 아는 사람 중 "코드 조각을 극한으로 최적화하는" 사람을 가장 잘 대표하는 사람을 생각해보세요. 줄여서 "최적화하는 올리"라고 부르겠습니다. 올리는 아마도 어떤 종류의 코드 조정이 효과를 볼 수 있을지에 대한 직관적인 감각을 가지고 있을 것입니다. 올리는 아마도 어떤 것들이 효과가 있는지 확인하기 위해 몇 가지를 시도해보고, 효과가 없으면 롤백하고 다른 것을 시도할 것입니다. 하지만 어떤 종류의 것들이 가능한지, 그리고 어떻게 컴파일러를 이길 수 있는지에 대한 직관을 가지고 있습니다.
올리가 가진 이러한 직관은 종종 컴파일러가 하는 것 이상입니다. 현대의 최적화 컴파일러는 결과가 상당히 인상적이지만, 상당히 단순한 규칙과 휴리스틱에 기반합니다. 기술 용어로, 고정점까지 실행되는 로컬 데이터 흐름 변환의 아이디어에 기반합니다. 또한 고려 사항을 단계적으로 순서화합니다. 즉, A를 고려한 다음 B를 고려하도록 컴파일러 파이프라인을 구축하지만, 복합적인 AB 문제는 고려하지 않습니다. 스케줄러와 레지스터 할당기는 이에 대한 악명 높은 예시이며, 복합 스케줄러-레지스터 할당기(단계 순서화를 붕괴시키는 이점을 얻기 위해)에 대한 많은 박사 학위 논문이 시도되었지만, 실제로 작동하게 만드는 것은 어려웠습니다.
이것이 올리의 전문성이 가치 있는 이유입니다. 올리는 종종 상황에 맞춤화된 휴리스틱으로 여러 NP-완전 문제를 균형 있게 조정하는 방법을 알고 있습니다. 따라서 더 맞춤화된 맥락 인식과 민감도가 있습니다. 올리는 또한 최적화 컴파일러가 수익성 있게 적용하지 못할 수 있는 기술, 특히 아웃라이닝, 사용자 정의 ABI 또는 벡터화를 가능하게 하는 변환, 또는 "컴파일러가 이것을 할 방법이 있었으면..."이라고 불평하게 만드는 다른 여러 가지 기술들을 조합하여 사용할 수 있습니다.
이제 AI가 가지고 있는 추론 능력을 통해 미니 올리 역할을 할 수 있다고 생각해보세요. AI는 무엇이 결실을 맺을지에 대한 완벽한 직관을 가지고 있지는 않을 수 있지만, 무엇이 수익성이 있을 수 있는지에 대한 감이 있고, 매우 많은 시도를 할 수 있습니다.
이 접근 방식에서는 STOKE 논문과 달리 시간이 무한대로 갈 때 최적 프로그램을 볼 수 있다고 보장할 수는 없지만, AI가 "더 인간과 같은" 추론 능력을 가지고 있기 때문에 단위 시간당 상당한 인간과 같은 추진력을 얻을 수 있습니다.
다시 연결하기: MLA 커널
먼저 말씀드리자면: 저는 AI가 우리의 Jalapeño MLA 커널을 위해 어떤 저수준 코드를 뱉어냈는지 알지 못합니다. 하지만 numpy로 MLA를 입력하는 방법은 알고 있습니다.
제가 이전에 작업했던 XLA 컴파일러에서는 이러한 numpy 연산들을 덩어리로 융합한 다음, "에미터"라는 메타프로그램을 사용하여 루프, 명령어 및 저수준 프리미티브로 낮췄습니다.
XLA 컴파일러/에미터 프로그램이 그렇게 했을 때, 저는 뒤에서 어떤 어셈블리가 나오는지 신경 쓸 필요가 없었습니다. 우리의 확률적 최적화 도구의 경우, AI는 개념적으로 에미터 메타프로그램의 자리를 차지합니다. 즉, 낮춤과 최적화를 모두 수행하며, 우리는 roofline을 향해 더욱 더 최적화하도록 요청할 수 있습니다.
이것이 AI가 어디에 끼어들고 기존 최적화 컴파일러 시스템의 구성 요소와 어떻게 유사한지 명확히 해주기를 바랍니다. 또한 "어셈블리 코드"로 간주하는 계층이 이제 위로 이동하고 있다고 생각하는 것도 도움이 됩니다. 일반 C++를 입력하고 -O3(가장 높은 일반적인 최적화 수준)으로 컴파일할 때, 입력한 C++를 이해하더라도 나오는 어셈블리를 이해할 것이라고 기대하지 않습니다. 우리는 여기서 유사한 작업을 수행하고 있지만, 더 높은 수준의 수학적 입력 사양을 사용합니다.
이제 핵심 질문은 AI에서 얻은 프로그램이 실제로 더 높은 수준의 설명/numpy와 동등한지 어떻게 확인하는가입니다. 그 확인 메커니즘은 AI 확률적 최적화 프로세스의 건전성을 확립합니다. 향후 블로그 게시물에서 이에 대해 더 자세히 다룰 것으로 예상되지만, 지금은 의미론적 동등성을 테스트하는 것이 가능하며 우리가 실제로 하고 있다고 말씀드리는 것으로 충분합니다. 가속기 프로그램은 "AI 최적화 프로그램이 정확히 수행하는 것"임을 검증할 수 있는 강력하고 완전한 계약에 특히 적합하며, 광범위한 맥락에서 상당히 수학적이고 데이터 흐름 지향적입니다.
이 분야의 많은 관련 기술은 프로그램 합성이라는 하위 분야의 노력에 의해 개척되었습니다. 최적화 컴파일러가 "여기 의미론을 가진 프로그램이 있습니다. 동등한 의미론으로 더 나게 만드세요!"라고 말하는 반면, 프로그램 합성은 "이러한 의미론을 가진 프로그램이 존재한다고 믿습니다. 찾을 수 있는 최고의 프로그램을 찾아주세요."라고 말합니다. 프로그램 합성은 최적화 컴파일보다 더 어려운 문제이지만, 근본적으로 덜 제약적입니다. 이는 본질적으로 올리와 같은 인간이 최적화 컴파일러보다 더 잘할 때 하는 일이며, AI가 이제 우리가 자동화하도록 도울 수 있는 것입니다. AI는 원래 프로그램에서 "영감"을 얻을 수 있지만, 단지 사소한 로컬 변환만 수행할 필요는 없습니다. 고전적인 최적화 컴파일러는 "아, 버블 정렬을 작성했군"을 보고 계약을 이해하여 퀵 정렬로 전환하지 않을 것이지만, 올리와 AI 모두 그렇게 할 수 있습니다. 이것이 우리를 고전적인 최적화 컴파일러 체제보다는 확률적 최적화를 통한 프로그램 합성 체제에 더 가깝게 만드는 것입니다.
이 모든 것은 "numpy에서 멀지 않은" 것으로 시작하여 48시간을 기다리면 동일한 의미론을 가진 최적화된 커널을 얻을 수 있다는 사실로 귀결됩니다. 이는 HotChips 발표에서 보여드린 바와 같습니다:
슬라이드에서도 언급했듯이, 우리 머신에서는 상당히 잘 튜닝되었다고 생각했던 커널에서도 AI가 인간 전문가의 성능을 능가하는 것을 종종 관찰할 수 있습니다. 탐색해야 할 수많은 조합/순열의 다양성으로 인해 달성 가능한 상당한 비율이 여전히 남아 있는 경우가 많습니다. 이러한 작업은 인간 성능 엔지니어에게는 종종 다루기 힘들 정도로 지루합니다.
요약 및 결론
컴파일러는 결국 하나의 함수일 뿐입니다. 우리는 프로그램을 그 함수에 제공하고, 프로그램의 더 나은 버전을 돌려받습니다. 우리가 얻는 프로그램과 입력한 프로그램은 동일한 의미론을 가질 것으로 기대합니다.
전통적인 최적화 컴파일러는 데이터 흐름 규칙과 휴리스틱을 통해 프로그램을 더 좋게 만듭니다. 이러한 규칙과 휴리스틱은 그 기원을 완전히 이해할 수 있지만, 수행할 수 있는 변환 측면에서 더 제한적일 수 있습니다.
대조적으로, AI는 확률적 최적화 도구로서 "열심히 생각"하고 무언가를 내뱉기만 하면 됩니다. 그 변환은 근본적으로 제한되지 않아 인간 전문가 최적화 도구와 더 유사합니다. AI가 내뱉는 프로그램이 건전하고 우리가 입력한 것과 동일한 의미론을 구현하는지 확인하는 방법이 필요하며, 우리는 이미 이를 마련했습니다. 그리고 이러한 종류의 AI 최적화는 매우 강력한 계약을 가진 수학적 연산에 특히 적합합니다. 계약은 커널이 하는 일을 한 줄 한 줄 이해할 필요를 없애줍니다.
그것이 바로 우리가 AI 생성 MLA 커널을 얻은 방법입니다!
