Kernel-Smith: 진화적 커널 최적화를 위한 통합 레시피

고성능 GPU 커널 작성은 AI 시스템 효율성의 핵심이지만, 사람이 직접 최적화하기 매우 어렵고 기존 LLM은 단발성 생성에 그치는 한계가 있다. 이 논문은 LLM이 스스로 코드를 실행하고 성능 피드백을 받아 반복적으로 개선하는 '진화적 루프'를 도입하여, 실제 상용 추론 엔진(SGLang, LMDeploy)에 채택될 수준의 최적화 성능을 자동화했다.

GPU 커널 최적화는 Tiling, Fusion 등 수많은 선택지 중 최적의 조합을 찾는 과정인데, 기존 LLM은 한 번에 완벽한 코드를 짜려다 보니 복잡한 최적화에 실패한다. 이는 마치 미로 찾기에서 한 번의 선택으로 출구까지 가려는 것과 같아, 조금만 경로가 복잡해져도 오류가 발생하거나 성능이 낮은 코드가 생성된다.

Kernel-Smith는 이를 해결하기 위해 '진화(Evolution)' 개념을 도입한다. 여러 개의 코드 후보(Population)를 만들고, 실제로 실행해본 뒤 "컴파일이 되는가?", "결과가 정확한가?", "얼마나 빠른가?"라는 구체적인 성적표를 LLM에게 준다. LLM은 이 성적표를 보고 기존 코드의 단점을 보완해 더 나은 다음 세대 코드를 만든다.

특히 모델 학습 시 단순히 최종 결과물만 흉내 내게 하는 것이 아니라, "A 코드에서 B 코드로 바꿨을 때 성능이 이만큼 좋아졌다"는 단계별 변화 과정을 학습시킨다. 이를 통해 모델은 현재 코드의 병목 지점을 찾아내고 이를 수정하는 '전문 수리공'과 같은 능력을 갖추게 되어, 반복적인 시도 끝에 인간 전문가 수준의 고성능 커널을 완성한다.

에이전트 프레임워크는 OpenEvolve를 기반으로 하며, MAP-Elites 알고리즘을 활용해 코드의 복잡도와 성능이라는 다차원 공간에서 다양한 고성능 후보군을 유지한다. 각 단계에서 에이전트는 아카이브된 상위 코드와 실행 피드백(컴파일 로그, 실행 시간, 하드웨어 메타데이터)을 입력받아 새로운 변이 코드를 제안한다.

데이터 합성 단계에서는 GitHub에서 59,000개의 고품질 PyTorch 모듈을 수집하고, 이를 Triton 커널로 변환하는 과정을 거친다. 특히 'Cold-start 데이터'와 'Cluster-Seeded 전문가 데이터'를 구분하여 생성하는데, HDBSCAN 클러스터링을 통해 대표적인 연산 패턴을 식별하고 전문가 모델(Gemini-3.0-pro 등)로 최적화 궤적을 생성한다.

학습 레시피는 SFT와 RL(GRPO)로 구성된다. SFT 단계에서는 정확성과 성능 향상이 확인된 단계별 수정 사항만 필터링하여 학습시킨다. RL 단계에서는 GRPO 알고리즘을 사용하며, 부모 코드 대비 속도 향상 비율(Speedup ratio)을 보상 신호로 활용한다. [부모 코드 실행 시간 / 생성 코드 실행 시간 → 속도 향상 배수 계산 → 1보다 크면 양의 보상 부여] 방식으로 모델이 성능 개선 방향으로 최적화되도록 유도한다.

NVIDIA Triton 백엔드 기반 KernelBench 평가에서 Kernel-Smith-235B-RL 모델은 평균 속도 향상 비율(Avg AMSR) 3.70을 기록하며 SOTA를 달성했다. 이는 독점 모델인 Claude-4.6-opus(3.33)와 Gemini-3.0-pro(2.83)를 크게 앞지르는 수치이다. 특히 난이도가 높은 Level 2 작업에서는 7.77의 속도 향상을 보여 Claude-4.6-opus(5.83) 대비 압도적인 효율성을 입증했다.

MetaX MACA 백엔드에서도 Kernel-Smith-MACA-235B는 평균 14.26의 속도 향상을 기록하며 DeepSeek-v3.2-think(8.01) 등 대규모 모델들을 능가했다. 이는 프레임워크가 특정 하드웨어에 종속되지 않고 이기종 플랫폼으로의 확장이 용이함을 보여준다.

실제 서비스 적용 사례로 SGLang의 메타데이터 업데이트 루틴에 적용되어 4.78배의 커널 속도 향상을 이끌어냈으며, LMDeploy의 MoE 라우팅 모듈에서는 1.36배의 속도 향상을 달성하여 실제 추론 엔진의 처리량(Throughput)을 최대 3.00% 개선하는 성과를 거두었다.

Kernel-Smith는 커널 생성을 '일회성 예측'이 아닌 '구조화된 탐색 문제'로 정의한다. 에이전트는 진화 아카이브에서 샘플링된 예시들과 현재 코드의 실행 피드백을 결합한 프롬프트를 통해 다음 변이를 생성한다. 이때 피드백은 단순 스칼라 보상이 아닌 컴파일 에러 메시지, 런타임 수치, 하드웨어 성능 카운터 등을 포함하는 다차원 정보를 제공한다.

사후 학습의 핵심은 '단계 중심(Step-centric)' 구성이다. 전체 진화 궤적을 그대로 모방하는 대신, 성능이 확실히 개선된 '원자적 수정(Atomic revisions)' 단위로 데이터를 분해하여 학습시킨다. 이는 긴 문맥에서 발생할 수 있는 정보 누출(Information leakage)을 방지하고 모델이 현재 상태에서 최선의 다음 단계를 찾는 능력을 극대화한다.

평가 시스템은 GPU 실행 시간의 변동성(Noise)을 줄이기 위해 Warm-up 실행, 다회 측정 후 이상치 제거, CUDAGraph 활용 등의 기법을 도입하여 측정 오차를 1% 이내로 억제했다. 또한 LLM이 기존 PyTorch 연산자를 그대로 호출하여 성능을 속이는 'Hacking'을 방지하기 위해 런타임 탐지 메커니즘을 구축했다.

현재 프레임워크는 Triton과 MACA 백엔드에 집중되어 있으며, 더 다양한 하드웨어 백엔드로의 확장이 필요하다. 또한 풀 리퀘스트(PR) 생성 및 통합 과정의 완전 자동화는 아직 향후 과제로 남아 있다.