KernelArc: A Multi-Agent Framework for GPU Kernel Optimization
KernelArc는 공유 메모리와 결정론적 가드(deterministic guards)를 통해 협력하는 전략 특화형 에이전트들을 통해 GPU 커널을 자율적으로 최적화하며, NVIDIA H100 및 B200 GPU의 다양한 워크로드에 대해 SOL-ExecBench 리더보드에서 최상위 순위를 달성한 멀티 에이전트 프레임워크입니다.
원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기
현대적인 컴퓨터는 인공지능에 필요한 방대한 계산을 처리하기 위해 그래픽 처리 장치(GPU)라고 불리는 특화된 칩에 의존합니다. 이 칩들은 매우 강력하지만, 동시에 많은 다양한 부품들이 완벽하게 동기화되어 작동해야 하는 복잡한 기계이기도 합니다. 이 칩들의 성능을 최대한으로 끌어올리기 위해, 엔지니어들은 칩에게 데이터를 어떻게 이동시키고 수학 연산을 수행할지 정확히 알려주는 커널(kernel)이라 불리는 아주 작고 고도로 특화된 프로그램을 작성해야 합니다. 수년 동안 이것은 인간 전문가들이 코드의 미세한 부분을 조정하고, 메모리 사용량을 조절하며, 연산 타이밍을 맞추어 단 1비트의 속도라도 더 짜내기 위해 수없이 많은 시간을 소비하는 작업이었습니다. 칩이 점점 더 정교해짐에 따라, 이러한 수동 작업은 더 어려워졌으며, 하드웨어가 할 수 있는 능력과 소프트웨어가 실제로 달성하는 성과 사이의 간극은 더욱 벌어졌습니다.
연구자들은 이 코드를 작성하고 개선하는 데 도움을 주기 위해 방대한 양의 텍스트를 학습한 컴퓨터 프로그램인 대규모 언어 모델을 사용하기 시작했습니다. 이러한 모델은 변경 사항을 제안하고, 이를 테스트하며, 결과로부터 학습함으로써 자동화된 엔지니어처럼 역할을 수행할 수 있습니다. 그러나 단일 자동화 에이전트는 종-종 한계에 부딪힙니다. 모델은 좋은 해결책을 찾아내면 그것을 계속 다듬는 데만 집중하여, 다른 방향에 있는 더 나은 접근 방식들을 놓칠 수 있습니다. 이는 마치 등산객이 산의 멋진 길 하나를 발견한 뒤 그 길을 따라 계속 올라가기만 하느라, 근처의 다른 경로가 훨씬 더 높은 정점으로 이어진다는 사실을 깨닫지 못하는 것과 같습니다. 이를 해결하기 위해, IMEC의 연구진은 여러 개의 자동화된 에이전트가 동시에 다양한 경로를 탐색하도록 하는 'KernelArc'라는 새로운 시스템을 개발했습니다.
KernelArc 시스템은 병렬로 실행되는 여러 에이전트에게 서로 다른 전략을 할당함으로써 작동합니다. 한 명의 에이전트가 문제를 혼자 해결하려고 노력하는 대신, 그룹 내의 각 에이전트는 데이터 저장 방식 변경, 수학적 정밀도 수정, 또는 여러 단계를 하나로 결합하는 것과 같이 서로 다른 공격 각도에 집중합니다. 이 에이전트들은 너무 무질서하고 혼란스러워질 수 있는 전체 사고 과정이나 모든 단계를 공유하지 않습니다. 대신, 그들은 오직 최종 결론, 즉 무엇이 성공했고 무엇이 실패했는지만을 공유합니다. 그들은 이 결과를 게시판 역할을 하는 공유 메모리 공간에 기록합니다. 만약 한 에이전트가 계산 속도를 높이는 비결을 발견하면, 그 결과를 게시합니다. 다른 에이전트들은 이 게시물을 읽고 그 통찰력을 자신의 작업에 활용하여, 막다른 길을 피하고 성공 사례를 바탕으로 발전할 수 있습니다.
에이전트들이 고장 난 코드로 시간을 낭비하지 않도록, 시스템에는 엄격한 자동 심판이 포함되어 있습니다. 이 심판은 모든 새로운 코드 제안을 일련의 테스트에 통과시켜 정확성을 확인하고 속도를 측정합니다. 만약 제안이 테스트를 통과하지 못하면 즉시 폐기됩니다. 만약 코드가 정상적으로 작동하더라도 현재의 최고 버전보다 빠르지 않다면, 백업용으로는 유지되지만 리더를 대체하지는 않습니다. 새로운 버전이 정확하면서도 더 빠를 때만 비로소 새로운 표준이 됩니다. 이 과정을 통해 에이전트 팀은 모든 시도의 세부 사항에 매몰되지 않고도 광범위한 가능성을 탐색할 수 있습니다. 또한 시스템에는 에이전트가 더 이상의 개선을 찾을 수 없는 정체기(plateau)에 빠졌을 때 작동하는 안전 장치가 있어, 시스템이 완전히 다른 접근 방식을 시도하도록 강제함으로써 탐색이 계속 전진할 수 있도록 보장합니다.
연구진은 이 시스템을 사용 가능한 가장 진보된 GPU 중 두 가지인 NVIDIA H100과 B200에서 테스트했습니다. 그들은 행렬 곱셈, 언어 모델에 사용되는 어텐션 메커니즘, 그리고 다양한 형태의 데이터 융합과 같은 작업을 포함하여 성능을 측정하는 데 사용되는 표준 작업 세트에 집중했습니다. 상세한 가이드와 함께 단일 에이전트가 작업하는 특정 테스트에서, 시스템은 해당 작업에 대한 기존 최고의 라이브러리보다 약 3.2% 빠른 766 테라플롭스의 속도에 도달했습니다. 이는 단일 에이전트가 명확한 경로가 주어졌을 때 특정 문제에 매우 깊게 파고들 수 있음을 보여주었습니다. 그러나 연구진이 더 넓은 범위의 작업을 다루기 위해 다중 에이전트 KernelArc 시스템으로 전환했을 때, 결과는 더욱 인상적이었습니다. 이 시스템은 특화된 어텐션 메커니즘과 대규모 언어 모델을 위한 융합 레이어 등 다양한 복잡한 연산을 위한 맞춤형 구현을 만들어냈습니다.
다양한 크기와 형태의 데이터에 대한 연산 속도를 순위 매기는 공개 리더보드에서, KernelArc의 제출물은 여러 카테고리에서 1위를 차지했습니다. 어텐션과 잔차 더하기(residual addition)가 포함된 한 작업에서, 단일 에이전트는 0.441의 성능 점수에서 멈춰 섰습니다. 통찰력을 공유하고 다양한 방향을 탐색하는 다중 에이전트 시스템은 이 장벽을 깨고 0.481에 도달했습니다. 또 다른 복잡한 어텐션 작업 테스트에서, 다중 에이전트 시스템은 표준 참조 구현 대비 거의 291배, 고도로 최적화된 베이스라인 대비 143배 이상의 속도 향상을 달상했습니다. 이러한 결과는 여러 에이전트가 자신들의 발견을 공유하는 것이, 특히 문제 공간이 크고 다양할 때 단일 에이전트가 혼자 작업하는 것보다 더 나은 해결책을 더 빠르게 찾을 수 있게 해준다는 것을 시사합니다.
연구는 또한 에이전트의 수와 공유 메모리의 양이 변할 때 시스템이 어떻게 작동하는지 조사했습니다. 연구진은 에이전트들이 성공과 실패를 게시할 수 있는 공유 메모리를 제공하는 것이 정해진 시도 횟수 내에서 더 강력한 결과를 도출하는 데 도움이 된다는 것을 발견했습니다. 메모리가 무제한일 때 시스템은 가장 우수한 성능을 보였으며, 시작점 대비 290배 이상의 속도 향상을 달성했습니다. 이는 집단의 집단적 경험으로부터 배우는 능력이 성공의 핵심 요소임을 나타냅니다. 연구진은 공유 메모리나 전략 전문화와 같은 각 기능의 가치가 특정 작업과 탐색 단계에 따라 달라진다고 언급했습니다. 때때로 에이전트들은 폭넓게 탐색해야 하고, 다른 때에는 특정 솔루션을 정교하게 다듬어야 합니다.
이 연구는 컴퓨터 최적화의 모든 문제를 해결했다고 주장하거나, 인간 엔지니어가 더 이상 필요하지 않다고 제안하는 것이 아닙니다. 결과는 테스트된 작업과 하드웨어에 국한되며, 시스템은 여전히 에이전트를 안내하기 위해 인간이 설계한 프레임워크에 의존합니다. 그러나 이번 연구 결과는 조정된 그룹의 자동화된 에이전트들이 단일 에이전트보다 더 넓은 범위의 솔루션을 탐색할 수 있음을 입증합니다. 가장 가치 있는 결론만을 공유함으로써, 이 에이전트들은 실수를 반복하지 않고 서로의 발견을 토대로 발전할 수 있습니다. 이러한 접근 방식은 하드웨어의 잠재력과 소프트웨어 성능 사이의 간극이 계속 커지고 있는 현대 컴퓨팅의 복잡성을 다루는 유망한 방법을 제시합니다. 이 시스템은 적절한 협업이 있다면, 자동화된 도구들이 차세대 인공지능을 구동하는 칩의 완전한 힘을 끌어내는 데 도움을 줄 수 있음을 보여줍니다.
연구 분야의 논문에 파묻히고 계신가요?
연구 키워드에 맞는 최신 논문의 일일 다이제스트를 받아보세요 — 기술 요약 포함, 당신의 언어로.