최적화 시점을 학습하기: 전문가 GPU 커널 계보에서 얻은 검증된 최적화 기술

Learning When to Optimize: Verified Optimization Skills from Expert GPU-Kernel Lineages

LLM 기반 에이전트가 GPU 커널 생성을 위해 점점 더 많이 사용되고 있지만, 이러한 에이전트는 종종 어떤 최적화를 시도해야 하는지는 알지만, 그 최적화가 언제 유효한지 알지 못합니다. 본 연구에서는 KLineage를 소개하며, 이는 전문가 커널로부터 누락된 "언제"에 대한 지식을 학습합니다. KLineage는 순방향 실행을 기반으로 하는 대신, 검증 단계를 거친 단순화를 통해 전문가 구현을 역으로 추적하고, 각 수용된 단계를 재사용 가능한 최적화 기술로 변환합니다. 각 기술은 최적화 의도뿐만 아니라 코드 내 적용 위치, 유효성을 결정하는 조건, 효과 및 가정 위반으로 인한 실패 사례를 기록합니다. 이후 LLM은 이러한 기술을 동일한 컴파일/정확성/프로필링 단계를 거쳐 새로운 코드에 적용합니다. NVIDIA의 두 가지 아키텍처에서 사용된 다섯 가지 전문가 워크로드에서, 계보에서 파생된 이러한 기술들은 효과적인 최적화 교육 과정을 제공하며, 동일한 예산 하에서 최근의 메모리 기반 LLM 커널 기준을 뛰어넘는 최종 커널 품질과 최적화 효율성을 달성합니다. 또한, 원본 코드 암기 여부를 확인하기 위해 별도의 22개 테스트 케이스를 사용하여 추가적인 검증을 수행했습니다.

LLM-based agents are increasingly used to generate GPU kernels, but they often know what optimizations to try without knowing when those optimizations are sound. We introduce KLineage, which learns this missing "when" knowledge from expert kernels: instead of relying on forward rollouts, KLineage walks expert implementations backward through validation-gated simplifications and reverses each accepted step into a reusable optimization skill. Each skill records not only the optimization intent, but also where it applies in code, what conditions made it valid, what effect it had, and what failures its assumptions avoid. A downstream LLM materializes these skills on new code surfaces under the same compile/correctness/profile gate. On five expert workloads across two NVIDIA architectures, these lineage-derived skills serve as an effective optimization curriculum, exceeding recent memory-based LLM-kernel baselines in both final kernel quality and optimization efficiency under the same fixed budget. We additionally use a separate 22-instance held-out check as a sanity test against source-case memorization.