[소개] Meta는 Llama 3.1을 기반으로 미세 조정된 8B 모델인 KernelLLM을 출시합니다. 이 모델은 PyTorch 코드를 효율적인 Triton GPU 커널로 자동 변환할 수 있습니다. 실제 측정 데이터에 따르면, 단일 샷 추론 성능은 GPT-4o와 DeepSeek V3를 능가하며, 여러 번 생성하면 점수가 급등합니다.
AI 분야에서 매개변수 크기는 한때 "성능 한계"로 여겨졌습니다.
Meta의 최신 KernelLLM은 매개변수가 8B로 작아 GPU 커널 생성 작업에서 200B GPT-4o를 능가합니다.
이는 PyTorch 모듈 효율적인 Triton GPU 커널로 자동 변환하는 것을 목표로 하는 Llama 3.1 Instruct를 기반으로 미세 조정된 8B 매개변수 모델입니다.
KernelLLM은 GPU 커널 개발을 위한 마법같은 도구입니다. 매개변수가 적으면서도 더 강력한 성능을 발휘하며, 간단하고 사용하기 쉽습니다.
매개변수는 8B에 불과하지만, KernelBench-Triton Level 1에서는 단일 추론 성능이 GPT-4o와 DeepSeek V3를 능가합니다.
여러 추론을 통해 KernelLLM은 DeepSeek R1보다 우수한 성능을 보입니다.
이 모든 것은 경쟁 제품보다 매개변수 크기가 두 배나 작은 모델에서 비롯된 것입니다.
@Denis Kanonik은 "이게 또 테스트 세트를 사용해서 훈련된 건가요?"라고 불평했습니다.
KernelLLM은 커널 개발을 더 쉽게 만듭니다
KernelLLM은 Llama 3.1 Instruct를 기반으로 한 8B 모델로, Triton으로 GPU 커널을 작성하는 작업을 위해 특별히 훈련되었습니다.
GPU 프로그래밍을 더 쉽게 만들고 고성능 GPU 커널 생성을 자동화합니다.
KernelLLM은 효율적인 Triton 구현을 자동으로 생성하여 고성능 GPU 커널에 대한 증가하는 수요를 충족합니다.
작업 부하가 증가하고 가속기 아키텍처가 다양화됨에 따라 맞춤형 핵심 솔루션에 대한 수요가 크게 증가했습니다.
시중에 나와 있는 관련 도구 중 상당수는 테스트 중에만 최적화가 가능하거나 KernelBench 문제 조정에만 초점을 맞추기 때문에 다양한 시나리오에 대처하기 어렵습니다.
KernelLLM은 외부(PyTorch, Triton) 코드의 데이터에 따라 미세 조정된 최초의 LLM입니다.
Triton 커널 생성 워크플로
PyTorch 코드를 입력하면 KernelLLM이 Triton 커널 후보 코드를 생성합니다.
그런 다음 단위 테스트를 사용하여 코드를 검증하고, 무작위 입력으로 실행한 다음 출력이 올바른지 확인합니다. 여러 후보 코드를 생성하면 어느 코드가 가장 좋은지 비교하고 가장 좋은 코드를 선택할 수 있습니다.
KernelLLM의 Triton 커널 생성 프로세스 : KernelLLM을 사용하여 PyTorch 코드를 Triton 커널의 후보 코드로 변환합니다. 생성된 코드는 알려진 형태의 무작위 입력 데이터로 커널을 시험하는 단위 테스트를 통해 검증됩니다. 이 프로세스는 여러 후보 코드 생성(pass@k에서 평가), 품질 향상을 위한 후보 수 증가, 최종적으로 최상의 Triton 커널 구현 선택(녹색 부분)을 지원합니다.
이 모델을 훈련하기 위해 팀은 25,000개 이상의 (PyTorch, Triton) 코드 예제와 합성 샘플 쌍을 사용하여 많은 노력을 기울였습니다.
이 데이터 중 일부는 TheStack의 필터링 코드에서 나오고, 일부는 torch.compile() 및 힌팅 기술을 통해 생성됩니다.
데이터셋 KernelBook, 참조 링크: https://huggingface.co/datasets/GPUMODE/KernelBook.
Llama3.1-8B-Instruct 모델을 학습에 사용하였고, 사용자 지정 데이터 세트에 대한 지도 미세 조정(SFT)을 수행하여 KernelBench-Triton에서 올바른 Triton 커널과 호출 코드를 생성하는 능력을 테스트했습니다.
KernelBench-Triton은 KernelBench[Ouyang et al.]를 기반으로 개발된 변형입니다. 2025], Triton 커널 생성에 초점을 맞춥니다.
학습 및 평가 시 PyTorch 코드는 형식 예제를 지침으로 포함하는 프롬프트 템플릿으로 구성됩니다.
이 모델은 표준 SFT 방법을 사용하여 배치 크기 32로 10개 에포크 동안 학습되었으며, 하이퍼 매개변수는 검증 세트의 복잡도에 따라 선택되었습니다.
학습은 16개 GPU에서 12시간(GPU 시간 192시간)이 걸렸으며, 최상의 체크포인트에 대한 검증 결과가 보고되었습니다.
성과 평가
모델 크기는 작지만 성능은 최첨단 LLM과 비슷합니다.
KernelBench-Triton 테스트에서 8B 매개변수를 가진 KernelLLM은 20.2의 단일 추론 점수를 받았는데, 이는 671B 매개변수를 가진 DeepSeek V3(16점)과 200B 매개변수를 가진 GPT-4o(15점)보다 높습니다.
더 많은 후보 코드가 생성되면 점수가 더욱 높아질 수 있습니다. 후보코드 10개가 생성되면 점수는 51.8점에 도달할 수 있고, 후보코드 20개가 생성되면 점수는 57.1점에 도달할 수 있습니다.
KernelLLM 추론은 temperature=1.0, top_p=0.97로 실행되었습니다.
이 모델은 LLM으로 작성된 효율적인 GPU 커널의 성능을 평가하는 오픈 소스 벤치마크인 KernelBench에서 테스트되었습니다.
여기에는 Conv2D나 Swish(레벨 1)와 같은 간단한 단일 작업부터 완전한 모델 아키텍처(레벨 3)까지 작업 부하에 맞게 확장 가능한 신중하게 선택된 PyTorch 모듈 250개가 포함되어 있습니다.
다양한 난이도의 작업에서 안정적으로 수행되며, 간단한 단일 연산자와 복잡한 모델 아키텍처를 모두 처리할 수 있습니다.
테스트는 코드의 정확성(참조 PyTorch 출력과 비교)과 성능(기준 구현과 비교한 속도 향상을 통해)을 모두 측정합니다.
이 팀은 LLM이 Triton 커널을 생성하는 능력을 특별히 평가하는 새로운 KernelBench-Triton 변형을 개발했는데, 이는 KernelLLM 테스트에 적합합니다.
모든 테스트는 NVIDIA H100 GPU에서 수행되었습니다.
KernelLLM은 pass@k에서 대략 로그 선형 스케일링 동작을 보입니다.
KernelLLM을 어떻게 사용하나요?
먼저 여러 종속 패키지를 설치하세요.
- pip install 변압기 가속 토치 트리톤
이를 사용할 때는 먼저 라이브러리를 가져온 다음 generate_triton 함수를 호출하여 최적화된 Triton 코드를 생성합니다.
KernelLLM은 PyTorch 코드에서 Triton 커널을 생성하기 위한 간단한 인터페이스를 제공합니다.
스크립트를 작성하고 싶지 않다면 python kernelllm.py를 직접 실행하고, 내장된 REPL 인터페이스를 사용하고, 대화형 인터페이스를 열어 실시간으로 결과를 볼 수도 있습니다.
kernelllm.py는 모델과 상호작용하기 위한 다양한 방법을 제공합니다.
- 파이썬 kernelllm.py
KernelLLM은 빌드 프로세스를 사용자 정의하는 여러 가지 방법을 제공합니다.
때로는 잘못된 API 참조, 구문 오류 등 사소한 실수가 발생하고, 때로는 지침에 따라 이상적인 커널을 생성하지 못할 수도 있습니다.
생성된 코드의 구조는 컴파일러가 자동으로 내놓는 것과 비슷하며, 변수 명명, 텐서 형태, 유형 처리, 수치적 정밀도와 같은 세부 사항에서 문제가 발생할 수 있습니다.
참고문헌:
https://x.com/reach_vb/status/1924478755898085552
https://huggingface.co/facebook/KernelLLM
본 기사는 위챗 공개 계정 '신지위안' 에서 발췌하였으며, 영지가 편집하고 36Kr이 허가를 받아 게시하였습니다.
