▲

Stanford CRFM: AI로 생성된 CUDA 커널, PyTorch 최적화 코드 성능을 넘다

 (crfm.stanford.edu)
7P by GN⁺ 5일전 | ★ favorite | 댓글 2개
  • AI가 생성한 CUDA-C 커널들이 PyTorch의 전문가 최적화 커널과 비슷하거나 더 나은 성능을 보임
  • 단일 LLM(대형언어모델)이 자연어 최적화 아이디어 생성과 다양한 코드 브랜칭을 반복, 기존 방법보다 최적화 다양성과 병렬 탐색에서 뛰어난 성능 달성
  • 대표 벤치마크 결과, Matmul(101%), Conv2D(179.9%), Softmax(111.8%), LayerNorm(484.4%), Conv2D+ReLU+MaxPool(290.1%) 등에서 PyTorch 대비 압도적
  • 기존 “순차적 커널 개선”의 한계를 넘기 위해 자연어 추론 + 브랜칭 구조 적용, 빠른 속도로 고성능 커널을 생성
  • FP16, Flash Attention 등 최신 ML 워크로드에서도 진보 중이며, 미래에는 AI가 자체적으로 더 빠른 커널을 탐색·개선하는 패러다임이 주류가 될 가능성을 보여줌

TL;DR 주요 성과

  • Stanford CRFM 연구팀은 AI가 생성한 고성능 CUDA-C 커널이 기존 PyTorch의 전문가 설계 커널과 유사하거나 더 나은 속도를 내는 예를 발견함
  • 원래는 합성 데이터를 통해 더 나은 커널 자동생성 모델을 학습시키려고 했으나, 합성 데이터 생성만으로도 놀랄 만한 수준의 빠른 커널이 자동 생성되는 현상을 관찰함
  • 이에 따라 방법, 성능 벤치마크, 최적화 전략, 앞으로의 방향 등을 조기에 공개함
  • 벤치마크: Nvidia L40S GPU 기준. 성능(%): PyTorch 기준 실행 시간 ÷ 생성 커널 실행 시간
    • Matmul (FP32): PyTorch 대비 101.3% (4096x4096 행렬)
    • Conv2D: 179.9% (입력: 100, 3, 224, 224; AlexNet Conv1 규격)
    • Softmax: 111.8% (4096x65536 텐서)
    • LayerNorm: 484.4% (16, 64, 256, 256 텐서)
    • Conv2D + ReLU + MaxPool: PyTorch reference 대비 290.1%, torch.compile() 대비 189.0%

연구 동기 및 방법

  • 원래 목적은 커널 생성 모델 학습용 합성 데이터 생성이었으나, 실험 과정에서 자체적으로 생성된 커널이 예상을 뛰어넘는 고성능을 달성
  • KernelBench(Stanford 공개 벤치마크, 2024년 12월 공개) 활용
    • 주어진 torch 코드에 대해, LLM이 최적 속도 커널을 새로 작성
    • 입력/출력 결과의 수치 동치 여부로 정확성 검증
  • 기존 방식: 커널을 단계별로 조금씩 고치며 점차 개선하는 순차적 수정 루프
    • 단점: 아이디어 다양성 부족, 같은 최적화 반복, 지역 최적점 수렴
  • 개선 아이디어
    1. 최적화 아이디어를 자연어로 발상 및 기록한 뒤, 그 아이디어들의 코드 구현분기를 여러 개 동시에 생성
    2. 각 라운드마다 여러 최적화 시도 병렬 진행 → 최고 성능 커널로 다음 라운드 씨앗(시드) 설정
    • 이렇게 하면 한정된 탐색 반복 안에서 다양한 최적화 전략 탐구와 병렬 탐색이 가능

최적화 아이디어 예시

  • 메모리 접근 최적화: global/shared/register 메모리 계층 효율 개선
  • 비동기 처리 및 레이턴시 은닉: 연산과 데이터 이동을 오버랩
  • 데이터타입/정밀도 최적화: FP16/BF16 활용 및 하드웨어 특화 연산
  • 계산 및 명령어 최적화: 연산량, 명령어 수, 하드웨어 특수 명령 최적 활용
  • 병렬성 및 오큐펀시: SM(Streaming Multiprocessors) 활용 극대화
  • 루프/분기/인덱싱 최적화: 루프 최소화, 불필요한 인덱스 연산 제거

커널 최적화 실제 과정(Conv2D 예시)

  • 라운드별 최적화 아이디어 및 성능 개선 흐름
    • 최초(0라운드): 단순 CUDA 포팅(PyTorch 대비 20% 속도)
    • 다음 라운드: → 읽기 캐시 활용 → FP16 Tensor Core 연산(GEMM 변환) → 이중 버퍼링/파이프라인 → 인덱스 사전계산/공유메모리 → 벡터화 → K-루프 동시버퍼링 등 고도의 GPU 아키텍처 활용
    • 최종(13라운드): PyTorch 대비 179.9% 성능 확보
  • 실제 코드는 고급 CUDA 프로그래밍 기법 대거 활용
    • 각 라운드마다 새로운 아이디어를 자연어로 생성하고, 여러 구현을 병렬 시도하여 최적 코드를 선택

Takeaways 및 시사점

  • AI 기반 커널 생성이 강력한 탐색 루프와 자연어 기반 다양한 최적화 아이디어 조합으로 인간 전문가 수준을 뛰어넘을 수 있음
  • 일부 최신 연산자(FP16 matmul, Flash Attention)는 현 시점에서 PyTorch 대비 아직 성능 낮음
  • FP32 계산이 최근 하드웨어에서는 FP16/BF16에 비해 상대적으로 최적화 덜 됨 → 해당 영역에서 성능 우위 가능
  • 제한된 탐색 토큰(입출력 합 700만 토큰) 상황에서도 지속적인 성능 개선 확인
  • AlphaEvolve, Gemini 2.5 Pro 등 최근 연구도 대량 브랜칭+검증 기반 탐색이 재학습 없이도 획기적 성능 개선 가능함을 시사
  • 앞으로는 이런 방식으로 합성 데이터 및 고성능 커널을 대량 생성하며, AI가 자체 개선하는 루프(자기개선형 AI)로 발전할 것임

결론

  • AI 기반 커널 자동 생성·최적화는 이미 전문가 핸드코딩 수준에 도달했으며, 빠른 미래에 모델+브랜칭 탐색+검증 조합으로 자가개선 AI 시스템이 가능할 전망
  • PyTorch·TensorFlow와 같은 프레임워크의 성능 한계를 AI가 자체적으로 넘어서게 될 가능성 대두

부록: Conv2D 커널 전체 코드

  • 실제 함수와 커널 전체 소스코드가 포함됨(상세 코드 생략)
  • 코드 내 주요 특징:
    • 공유 메모리 벡터화, K-dim 계층적 double-buffering, CUDA WMMA, warp-level output buffer 등
    • 실시간 동적 인덱스 계산, bias 처리, K 루프 안의 지연 데이터 동시 로드 등
  • 완전한 샘플 코드 및 예시 커널은 깃허브 저장소에서 확인 가능
▲
aer0700 3일전  [-]

일종의 앙상블 기법이라고 해야할까요. 굉장하네요

답변달기
▲
GN⁺ 5일전  [-]
Hacker News 의견

  • 나는 이 글의 저자들이 "AI 에이전트"를 바라보는 방식이 꽤 흥미롭다고 생각함
    대부분 사람들은 에이전트를 인간 직원처럼 생각하는 경향이 있음
    몇 명 안 되는 에이전트를 병렬로(종종 한 명만) 설정해 놓고, 각각이 루프를 돌면서 한 번에 하나의 일만 처리하게 함
    여전히 고정된 인원 수와, 한 번에 한 가지 일만 가능한 직원, 그리고 작업 전환이 느린 세상에 머물러 있음
    하지만 LLM은 그렇지 않음
    사실상 무한한 에이전트를 언제든지 마음껏 만들어낼 수 있음
    LLM 요청을 일렬로 처리하나 병렬로 처리하나 비용 차이가 없음
    이 사실을 인지하면, 각 에이전트가 필요에 따라 자신을 여러 하위 에이전트로 분기시키는 패턴이 자연스럽게 떠오름
    이것이 바로 저자들이 시도한 방식임
    에이전트를 "작업(task)"이나 "잡(job)"로 간주하고, Celery나 Sidekiq에서 배운 점을 적용하는 관점이 더 적합하다고 느낌

  • "FP32는 최신 ML 워크로드에서 부쩍 덜 사용되고, 최신 하드웨어에서는 FP16이나 BF16에 비해 최적화도 덜 되어 있음. 그래서 FP32 커널에서 PyTorch보다 성능을 쉽게 개선할 수 있었던 이유 중 하나가 될 수 있음"
    개발자들이 수년에 걸쳐 fp32 버전 커널을 최적화하는 데 시간을 거의 투자하지 않았음
    정말 재미있는 건, 실제로 집중적으로 개발된, 사람들이 진짜 사용하는 커널 쪽 성능을 올릴 수 있을 때라고 생각함

    • NVIDIA가 GPU에 대해 충분히 자세한 문서를 제공하지 않기 때문에 이런 좋은 결과가 부분적으로 설명된다고 생각함
      마이크로아키텍처가 잘 문서화된 프로세서라면, 프로그래머나 컴파일러가 결정적으로 최적 프로그램을 작성할 수 있기 때문에, 이미 알려진 해법을 찾는 수준을 넘어서서 ML/AI를 적용해 대단한 개선을 이루기는 쉽지 않음
      반면, NVIDIA GPU처럼 문서화가 덜 된 경우에는, 최적의 프로그램을 찾으려면 이전에 최적화된 프로그램의 예를 참고하면서 무작위 탐색을 하거나, 어떤 상황에서 실제 하드웨어가 어떻게 동작하는지 역공학적 분석이 필요한 경우가 많음
      이런 상황에서는 ML/AI가 가능성을 보일 수 있고, 잘 만들어진 프로그램들을 데이터로 학습함으로써 인간 프로그래머도 찾기 힘든 undocumented behavior를 포착할 수 있음

    • 혹시 fp16/bf16 커널의 이미 알려진 개선사항이 단순히 fp32에 옮겨진 결과일 수도 있지 않을까 궁금함

    • "사람들이 수년간 fp32 커널을 전혀 손대지 않았다는 이야긴가?"
      와, 만약 그렇다면 AI가 사전 솔루션이 없던 영역에서 새로운 알고리즘을 만들어낸 셈이라니 정말 멋진 일이라고 생각함

  • 내 결론은 이 아티클, Google의 AlphaEvolve(여기), 그리고 최근 o3가 리눅스 커널에서 zero day를 찾았다는 소식(여기)에서 볼 수 있듯이
    특히 Gemini Pro 2.5와 o3가 기존 모델로는 안되던 아이디어들도 갑자기 해내는 새로운 능력 레벨에 도달했다고 느꼈음

    • 나는 갑자기 기존에 안되던 게 갑자기 되는 게 아니라
      사실상 인간보다 훨씬 빠른 속도로 반복과 테스트가 가능해졌고
      즉각적으로 대량의 정보를 활용할 수 있게 되면서
      방대한 정보, 발전, 그리고 지능적으로 적용된 브루트포스 조합이 일부 응용 분야에서 성공을 거두고 있는 시점에 이른 것이라고 봄

    • 네가 언급한 사례들과 이번 결과가 실제로 유사점이 많다고 생각함
      본문에도 "테스트 시 반복 루프는 차례로 코드 수정 결과를 확인하는 대화식 챗봇이 아니라, 명확한 최적화 가설에 따라 적극적으로 병렬 평가를 진행하는 구조화된 탐색 방식에 가깝다"라는 언급이 있음
      내 결론은 이제 LLM의 능력을 바탕으로
      평가 함수가 명확하거나 유사한 패턴의 솔루션 공간을 현저하게 줄이는 방법을 터득했다는 점임
      어떤 모델이 다른 모델을 추월한다, 혹은 특정 모델만 풀 수 있다... 이런 모델 간 비교가 논점이 아니라
      이런 식의 응용이 충분히 드러나고 있다는 현실이 더 의미 있다고 느껴짐

    • Gemini Pro 2.5는 사람이 쓸 수 있는 첫 번째 AI라고 느꼈지만, 엄밀히 보면 겨우 그 임계점을 넘긴 수준임
      성공률이 20% 아래로 떨어지는 경우도 종종 있음
      3.0이 나오면... 이제는 살짝 무서울 수 있겠다는 생각이 듦

    • 잠깐, 무슨 뜻이지? 이거 리눅스 커널이랑 상관없고, GPU 프로그래밍에서 말하는 "커널"임
      혹시 본문 전체를 잘못 이해한 거 아님?

    • 흥미롭긴 한데, 좀 더 강한 증거가 있음? 한 번만의 결과로는 설득력이 충분하지 않음

  • "기준 코드는 디폴트 FP32고, 허용 오차는 1e-02임"
    이 정도 오차면 쉽게 fp16 연산으로 "fp32" 커널을 대체할 수 있음

    • 한 단계만 더 생각해 보면 실제로 이번 작업의 핵심이 해당 커널에서 가능한 많은 fp32 연산을 fp16으로 바꾸라는 게 아니었나 싶음
      이런 이식 작업이 실제로 얼마나 어려운지는 확인이 필요하지만
      직관적으로는 그리 인상적이지는 않다고 느껴짐
      내 생각이 혹시 틀렸다면, 저자들이 이 대목을 명확하게 다뤄줬으면 좋겠음

    • 이렇게 되면 결과는 무의미하게 됨
      상대 오차를 체크는 했을지 모르겠음
      float32를 float16으로 대체하는 건 의미 없음
      정밀도가 float32를 선택하는 유일한 이유라 할 만한데
      이 특징 자체를 잃으면 버전별로 차별화도 없음

  • 혹시 나만, 이 글 제목 보고 AI가 OS 커널을 만들어냈다고 착각해서 읽었던 사람임?

    • 나도 그랬음

  • 400% 속도 향상도 대단하지만
    무엇보다 흥미로운 점은: 반복마다 단순 연산 최적화가 아니라, 언어 추론 단계를 강제로 넣어 탐색 다양성을 이끌어낸 방법론임
    이 방식이 실제로 잘 작동했기에 매우 흥미로움

    • 나는 map-elites나 island 기법 같은 게 들어갔나 했는데, 이 부분을 놓친 듯함
      그냥 평범한 미메틱 진화라고 생각함

  • 정말 흥미로운 결과임
    이 사람들이 너무 신나서 블로그에 공개한 듯함
    사실 출판 전에 피드백 받으려는 의도가 있었을지도
    "자체 개선(self improvement)"의 전설적인 길이 이거일지는 모르겠지만
    이런 결과가 바로 그 길에서 볼 법한 사례라고 느껴짐

    • "진짜 자체 개선의 길인지는 모르겠다"
      이 방법들은 극도로 명확한 평가 함수가 있을 때만 효과가 있음

  • 내 경험 공유: replication 시도에서 LayerNorm 커널은 수치적으로 안정적이지 않아 검증 불가라고 봄
    평균 0, 표준편차 1로만 테스트해서 numerical cancellation 현상이 바로 드러나지 않았음
    덧붙여, 이후에 수치적으로 안정적인 버전을 새로 만들었다는 사실을 확인
    이 점은 훌륭하다고 생각함

  • 정말 멋진 결과임
    o3와 Gemini 2.5 Pro를 사용했지만
    어떤 쪽이 더 좋은 커널을 만들어냈는지는 아쉽게도 언급이 없음

  • AI가 생성한 코드가 분할 커널(fused kernel)처럼 넓은 영역을 어떻게 다루는지 관찰하는 게 흥미로운 포인트임
    예를 들어 gemm + relu + gemm + normalization 등이 있을 텐데
    튜너로 쭉 훑거나 사람이 일일이 작성하려면 상당히 고생스러운 작업임

    • 근데 여기서 말하는 "커널"이 AI 맥락에서 정확히 뭘 의미하는지 잘 모르겠음
      OS 커널은 아니라는 점만은 확실함
답변달기