PTXBench: Benchmark and Adapt LLMs for GPU Kernel Optimization with Architecture-specific PTX

현대적인 컴퓨터 내부에서 그래픽 처리 장치(GPU)는 방대한 양의 데이터를 동시에 처리하도록 설계된 강력한 엔진 역할을 합니다. 이 엔진을 최대한 활용하기 위해 소프트웨어 개발자들은 하드웨어가 데이터를 어떻게 이동시키고 계산을 수행해야 하는지 정확하게 지시하는 특화된 명령어를 작성합니다. 이러한 명령어는 PTX라고 불리는 저수준 언어로 작성되며, 이는 프로그래머에게 기계의 내부 작동을 직접 제어할 수 있는 권한을 부여합니다. 이 과정을 더 쉽게 만들기 위한 상위 수준의 도구들도 존재하지만, 때때로 이들은 최신 하드웨어 세대의 고유한 특징들을 놓치기도 합니다. 이러한 기계들이 급격히 진화함에 따라, 소프트웨어를 특정 성능에 맞춰 완벽하게 조정하는 것은 어려운 과제가 되고 있습니다. 여기서 질문이 생깁니다. 인공지능, 특히 거대 언어 모델이 스스로 이러한 정밀한 저수준 명령어를 작성하는 법을 배울 수 있을까요, 아니면 여전히 성능을 제대로 끌어내지 못하는 오래되고 일반적인 방식에 의존하고 있을까요?

연구진은 PTXBench라는 새로운 테스트 환경을 구축하여 이 질문에 답하고자 했습니다. 그들은 이 AI 모델들이 단순히 작동하는 코드를 쓰는 것을 넘어, 최신 GPU의 특정하고 진보된 기능들을 능동적으로 사용하는 코드를 작성할 수 있는지 확인하고 싶었습니다. 연구진은 현재 사용 가능한 가장 강력한 그래픽 카드인 H100과 B200 두 가지에 집중했으며, 행렬 곱셈 및 어텐션 메커니즘과 같이 인공지능에 사용되는 복잡한 수학적 연산 작업들을 대상으로 모델들을 테스트했습니다. 목표는 세 가지를 측정하는 것이었습니다. 코드가 오류 없이 실행되는지, 연구진이 요청한 특정 하드웨어 명령어를 실제로 트리거하는지, 그리고 기존의 최고 성능 소프트웨어 라이브러리보다 빠르게 실행되는지였습니다.

결과는 모델이 할 수 있는 능력과 최상위 성능을 위해 요구되는 수준 사이에 상당한 격차가 있음을 보여주었습니다. AI 모델들은 단순한 작업에 대해서는 올-바르게 작동하는 코드를 작성하는 데 종종 성공했지만, AI 시스템 학습에 필수적인 어텐션 메커니كن의 역전파(backward pass)와 같이 더 복잡한 연산에서는 상당히 어려움을 겪었습니다. 모델이 요청된 하드웨어 명령어를 실행하는 코드를 작성해냈을 때조차도, 그 코드가 선두적인 전문 라이브러리의 속도를 따라잡는 경우는 드물었습니다. 많은 경우 AI가 생성한 코드는 더 느렸으며, 이는 단순히 명령어를 실행하게 만드는 것만으로는 진정한 최적화를 달eric하기에 충분하지 않다는 점을 시사했습니다. 연구에 따르면, 모델들은 해당 기능들이 명시적으로 설명되었음에도 불구하고, 새로운 칩의 고유한 아키텍처 특화 기능을 활용하기보다는 빈번하게 일반적인 코딩 패턴으로 회귀하는 경ו를 보였습니다.

이러한 모델들을 어떻게 개선할지 이해하기 위해, 연구진은 지도 미세 조정(supervised fine-tuning)이라는 방법을 사용하여 특정 AI 모델을 교육하는 통제된 실험을 수행했습니다. 연구진은 모델에게 모델 자신의 실수 사례와 함께, 수정된 버전의 코드 및 왜 그 수정이 효과적이었는지에 대한 설명을 함께 보여주었습니다. 'Fixit'이라고 명명된 이 접근 방식은 모델이 여러 작업, 특히 기능적으로 올바른 코드를 생성하는 데 있어 개선되는 데 도움을 주었습니다. 그러나 개선이 균일하게 나타나지는 않았습니다. 모델은 학습 과정에서 보지 못한 다른 크기의 문제나 변형된 작업으로 새로운 기술을 일반화하는 데 여전히 어려움을 겪었습니다. 연구진은 수정의 질과 학습 데이터의 균형이 데이터의 양만큼이나 중요하다는 것을 발견했습니다. 다양하지 않은 더 많은 양의 데이터로 학습된 모델보다, 다양한 문제 세트로 학습된 모델이 더 나은 성능을 보였는데, 이는 학습 경험의 유형이 데이터의 양보다 더 중요하다는 것을 나타냅니다.

또한 이 연구는 AI에게 적절한 문맥(context)을 제공하는 것의 중요성을 강조했습니다. 연구진이 하드웨어의 구체적인 능력에 대한 상세한 정보를 모델에게 제공했을 때, 모델은 요청된 명령어를 사용할 가능성이 훨씬 높았습니다. 이러한 구체적인 지식이 없으면 모델은 요청된 PTX를 거의 사용하지 않습니다. 이는 AI가 이러한 저수준 프로그래밍을 배울 수는 있지만, 효과적으로 수행하기 위해서는 정밀한 가이드와 고품질의 예시가 필요함을 시사합니다. 연구진은 현재의 AI 모델들이 유망함을 보여주기는 하지만, 최신 하드웨어에 맞춰 코드를 최적화하는 전문 인간 엔지니어를 대체할 준비는 아직 되지 않았다고 결론지었습니다. 앞으로 나아갈 길은 더 나은 학습 데이터, 더 균형 잡힌 학습 경험, 그리고 진화하는 컴퓨터 아키텍처의 특정 제약 사항에 대해 모델이 추론할 수 있도록 가르치는 방법에 대한 더 깊은 이해를 포함합니다.