💻 computer science

Equivalence Checking of ML GPU Kernels

이 논문은 수작업, 컴파일러 또는 LLM에 의해 최적화된 머신러닝 연산의 정확성을 형식적으로 검증하는 최초의 사운드하고 완전한(sound and complete) GPU 커널 동등성 검사기인 Volta를 소개한다.

원저자: Benjamin Driscoll, Kshitij Dubey, Anjiang Wei, Neeraj Kayal, Rahul Sharma, Alex Aiken

게시일 2026-08-18

📖 4 분 읽기☕ 가벼운 읽기

원저자: Benjamin Driscoll, Kshitij Dubey, Anjiang Wei, Neeraj Kayal, Rahul Sharma, Alex Aiken

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

현대 인공지능의 거대하고 보이지 않는 기계 장치 속에서, 가장 중요한 작업은 클라우드가 아니라 GPU라고 불리는 특수 컴퓨터 칩에서 일어납니다. 이 칩들은 수백만 개의 작은 계산을 동시에 수행하도록 설계되었는데, 이는 현재 코드를 작성하고, 언어를 번역하며, 예술을 생성하는 대규모 언어 모델을 훈련하는 데 필수적인 요소입니다. 이러한 모델들을 유용할 만큼 빠르게 구동하기 위해, 엔지니어들은 GPU가 데이터를 어떻게 이동시키고 수학 연산을 수행해야 하는지 정확히 알려주는 '커널(kernel)'이라 불리는 매우 특화된 명령어를 작성해야 합니다. 지난 몇 년 동안 기업들은 인간 엔지니어보다 더 빠른 작업 방식을 찾기 위해 인공지능 자체를 사용하여 이러한 커널을 작성하기 시작했습니다. 그러나 이러한 속도에는 위험이 따릅니다. AI나 컴파일러가 코드를 더 빠르게 재작성할 때, 의도치 않게 미묘한 오류를 도입할 수 있기 때문입니다. 이러한 오류는 컴퓨터가 잘못된 답을 내놓게 하거나, 더 심각하게는 표준적인 테스트로는 찾아내기가 거의 불가능한 방식으로 소리 없이 시스템을 충돌시킬 수 있습니다. 핵심적인 문제는 이 칩들이 수천 개의 작업 스레드를 동시에 실행한다는 점이며, 만약 이들이 완벽하게 조율되지 않으면 서로의 영역을 침범하여, 결과가 사건이 발생하는 예측 불가능한 순서에 따라 달라지는 '경쟁 상태(race condition)'를 유발할 수 있다는 것입니다.

연구진은 이 문제를 해결하기 위해 '볼타(Volta)'라고 불리는 새로운 도구를 개발했습니다. 새로운 버전의 커널이 올바른지 추측하는 대신, 볼타는 두 버전이 동일한 결과를 생성한다는 것을 수학적으로 증명하는 형식 검증기(formal verifier) 역할을 합니다. 연구진은 참조 커널(신뢰할 수 있는 원본 버전)과 최적화된 커널(새롭고 더 빠른 버전)의 저수준 명령어를 입력받아 심볼릭 엔진(symbolic engine)으로 실행하는 시스템을 구축했습니다. 엔진은 코드에 특정 숫자를 입력하여 결과값을 확인하는 대신, 입력을 추상적인 기호로 취급합니다. 엔진은 코드가 취할 수 있는 모든 가능한 경로를 추적하며, 데이터가 수천 개의 병렬 스레드를 통해 어떻게 이동하고 스레드들이 서로 어떻게 동기화되는지를 추적합니다. 만약 코드가 충돌을 일으킬 수 있는 방식으로 메모리에 접근하려고 시도하거나, 스레드들이 서로를 기다리며 영원히 갇혀 버린다면, 이 도구는 즉시 오류를 표시합니다. 코드가 문제없이 실행되면, 도구는 두 커널의 최종 출력을 복잡한 수학적 표현식으로 변환하고, 입력되는 특정 숫자와 관계없이 그 표현식들이 근본적으로 동일한지 확인합니다.

연구진은 행렬 곱셈, 컨볼루션(convolutions), 그리고 대규모 언어 모델을 구동하는 어텐션 메커니즘을 포함한 다양한 실제 머신러닝 작업에 대해 볼타를 테스트했습니다. 그들은 이 도구가 수작업으로 최적화된 커널, 컴파일러에 의해 최적화된 커널, 심지어 대규모 언어 모델에 의해 생성된 커널까지 성공적으로 검증할 수 있음을 발견했습니다. 한 사례에서는 13차례의 자동 개선 과정을 거쳐 최적화된 AI 생성 커널을 조사했습니다. 볼타는 이 AI 생성 코드가 원래의 인간이 작성한 참조 모델과 수학적으로 동일함을 확인하여, 공격적인 최적화가 논리 구조를 깨뜨리지 않았음을 증명했습니다. 또한 이 도구는 다른 방법들이 놓친 오류를 잡아냄으로써 그 가치를 입증했습니다. 예를 들어, 수천 명의 개발자가 수년간 사용해 온 유명하고 널리 인용되는 GPU 프로그래밍 튜토리얼에서 데이터 경합(data races)을 감지했습니다. 이러한 오류는 표준 테스트로는 거의 포착되지 않는 매우 특정한 타이밍 조건에서만 나타나기 때문에 숨겨져 있었습니다. 도구는 또한 AI가 생성한 커널에서 존재하지 않는 메모리 위치의 데이터를 읽으려고 시도하는 버그를 식별했습니다. 현재의 하드웨어는 이 실수를 무시하고 지나갔지만, 연구진은 이 코드가 근본적으로 안전하지 않으며 미래의 기기에서는 실패할 수 있음을 보여주었습니다.

이 접근 방식의 강점은 수천 개의 스레드가 동작을 조율해야 하는 GPU 프로그래밍의 독특한 복잡성을 다룰 수 있는 능력에 있습니다. 기존의 도구들은 단일 스레드 프로그램이나 고수준의 수학적 연산은 확인할 수 있었지만, GPU의 거대한 병렬성을 관리 가능한 단위로 분해하는 데 어려움을 겪었습니다. 볼타는 머신러닝에서 흔히 나타나는 구조적 패턴, 즉 스레드의 수와 데이터의 크기가 사전에 알려져 있다는 가정을 통해 이를 극복합니다. 이 프레임워크 내에서, 도구는 스레드가 적절히 동기화되지 않을 경우 경쟁 상태가 존재함을 확실히 증명할 수 있으며, 두 프로그램이 동일한 심볼릭 결과를 생성한다면 서로 동등하다는 것을 증명할 수 있습니다. 연구진은 이 도구가 수십만 개의 명령어를 포함하는 커널에 대해서도 불과 몇 초 또는 몇 분 만에 이러한 속성을 검증할 수 있음을 보여주었습니다. 또한 그들은 도구의 수학적 논리가 건전(sound)하다는 것을 증명했는데, 이는 도구가 두 프로그램이 같다고 판정했다면, 모든 가능한 입력에 대해 실제로 같다는 것을 의미합니다.

이 연구는 인공지능 개발을 더욱 안전하고 신뢰할 수 있게 만드는 데 있어 중요한 진전을 의미합니다. 기업들이 모델을 구동하는 코드를 생성하기 위해 자동화된 시스템에 점점 더 의존함에 따라, 그 코드를 점검할 엄격한 방법이 필수적이 되었습니다. 연구진은 제한된 시나리오만을 확인할 수 있는 단순한 테스트를 넘어, 정당성을 보장하는 방법론으로 나아가는 것이 가능하다는 것을 보여주었습니다. 최적화된 커널의 동등성을 검증함으로써, 볼타는 개발자들이 조용한 버그를 도입할 두려움 없이 더 빠르고 공격적인 최적화를 사용할 수 있는 확신을 줍니다. 이 도구는 현재 사용 가능하며, 연구진은 코드와 그 배후의 증명 과정을 공개하여 다른 이들이 이 토대 위에 구축할 수 있도록 했습니다. 비록 이 도구가 아직 모든 유형의 GPU 코드를 다루지는 못하지만, 현대 머신러닝을 이끄는 대다수의 커널을 성공적으로 처리하며 고성능 컴퓨팅 코드의 자동 생성에 대한 새로운 신뢰의 기준을 제시하고 있습니다.

기술 요약: ML GPU 커널의 동등성 검사 (Equivalence Checking)

문제 정의

딥러닝과 대규모 언模型(LLM)의 급격한 발전은 GPU 커널 실행에 막대한 비용을 초래했으며, 이제 GPU 커널은 공격적인 최적화의 주요 대상이 되었습니다. 이러한 최적화는 수동, 컴파일러, 그리고 점점 더 LLM에 의해 수행됩니다. 비록 경험적 테스트와 프로그램 분석이 이러한 커널을 검증하는 데 사용되고 있으나, 이들은 형식적 보증(formal guarantees)을 제공하지 못합니다. 이는 다음과 같은 GPU 프로그래밍의 고유한 도전 과제들 때문에 특히 중요합니다:

동시성 오류 (Concurrency Errors): 수천 개의 스레드 간의 미세한 동기화는 테스트만으로는 탐지하기 어려운 미묘한 데이터 경합(data races)과 데드락(deadlocks)을 유발합니다. 이러한 오류는 드문 실행 스케줄에서만 나타나거나, 하드웨어 버전 간에 보장되지 않는 암시적 동기화에 의존할 수 있습니다.
부동 소수점 의미론 (Floating-Point Semantics): 최적화 과정에서 산술 연산의 순서가 재배치되면 참조 구현과 최적화된 구현 사이에 작은 수치적 차이가 발생할 수 있습니다. 테스트는 종종 이를 허용하지만, 이는 실제 의미론적 버그(예: 클리핑 연산 누락)를 은폐할 수 있습니다.
형식적 도구의 부재 (Lack of Formal Tools): 기존의 동등성 검사 기법은 단일 스레드 정수 프로그램이나 고수준 텐서 연산을 지원하지만, GPU 커널 특유의 병렬성, 동기화 및 비결정론적 스케줄링을 처리하는 데 실패합니다. GPU 프로그램을 위한 이전의 동등성 검사기는 존재하지 않습니다.

방법론

저자들은 실용적인 ML GPU 커널 클래스를 위해 특별히 설계된 최초의 동등성 검사기인 Volta를 제시합니다. Volta는 최적화 과정을 알 필요 없이 PTX 어셈블리(NVIDIA GPU 스택의 가장 낮은 문서화된 수준)를 분석하는 블랙박스 검증기로 작동합니다.

1. 대상 클래스: Structured-CTAs

Volta는 Structured-CTAs(Cooperative Thread Arrays)를 대상으로 합니다. 이 클래스는 다음을 가정합니다:

스레드 수, 텐서 크기, 포인터 타겟이 정적으로 알려져 있음.
브랜치 타겟과 메모리 액세스 주소가 스레드 ID 및 루프 인덱스에 기반하여 정적으로 해결 가능함.
제어 흐름이 런타임 데이터 값에 의존하지 않음 (효율적인 ML 커널은 데이터 의존적인 런타임 결정을 내리지 않음).
이러한 가정들은 JAX 및 XLA와 같은 많은 고성능 ML 커널에서 유효하며, Volta에 의해 체크됩니다. 만약 위반될 경우 예외가 발생합니다.

2. 심볼릭 실행 및 합류성 (Symbolic Execution and Confluence)

Volta의 핵심은 텐서 값을 실수로 취급하는 심볼릭 실행 엔진입니다 (최적화 도구인 --ffast-math에서 사용하는 관례를 모델링함).

심볼릭 상태 (Symbolic State): 엔진은 참조 커널과 최적화된 커널을 심볼릭하게 실행하여, 입력 텐서의 함수로서 출력 텐서 요소에 대한 식을 도출합니다.
스케줄링 (Scheduling): 라운드 로빈 스케줄링 전략을 채택합니다. 스레드는 배리어(예: syncthreads, syncwarp)에 도달할 때까지 실행됩니다.
합류성 (Confluence): 핵심적인 기술적 결과는 합류성의 증명입니다. 심볼릭 평가기는 모든 가능한 실행 스케줄이 동일한 심볼릭 식을 도출함을 보장합니다 (또는 경합/데드락을 보고합니다). 이 속성은 Agda 증명 보조기를 통해 공식적으로 검증되었으며, 이는 검사기의 결과가 분석 중에 선택된 특정 스케줄에 독립적임을 보장합니다.

3. 경합 및 데드락 탐지 (Race and Deadlock Detection)

Volta는 배리어 동기화를 활용하여 데이터 경합과 데드락을 탐지하도록 심볼릭 실행을 확장합니다:

컨텍스트 추적 (Context Tracking): 각 스레드에 대한 메모리 이벤트(읽기/쓰기) 및 동기화 상태를 추적하는 컨텍스트 X를 유지합니다.
경합 탐지 (Race Detection): 스레드가 마지막 작성자 또는 이전 읽기 작업과 동기화하지 않은 상태에서 메모리 위치를 읽거나 쓸 경우 경합이 감지됩니다.
데드락 탐지 (Deadlock Detection): 스레드가 만족될 수 없는 배리어를 기다리며 차단되는 경우 데드락이 보고됩니다.
보증 (Guarantees): 시스템은 사후 검증(sound)이 입증되었습니다 (경합이 보고되었다면, 구체적인 경합이 존재함). 또한 "진양성(true positives)" 속성을 만족합니다.

4. 결정 절차 (Decision Procedure)

심볼릭 식이 생성되면, Volta는 참조 출력과 최적화된 출력이 수학적으로 동등한지 결정해야 합니다.

식의 형태 (Expression Form): 식은 다변수 다항식, 덧셈, 곱셈 및 지수 함수를 포함합니다 (예: $p(\bar{x})e^{h(\bar{x})}$ ).
결정 가능성 (Decidability): 저자들은 $\sum p_i(\bar{x})e^{h_i(\bar{x})} = 0$ 형태의 등식에 대한 결정 가능성을 증명합니다. 이 증명은 지수 $h_i$ 가 서로 다른 다항식인 경우, 합이 항등적으로 0이 되기 위해 계수 $p_i$ 가 0이어야 한다는 사실에 기반합니다. 이는 무거운 수론적 기법(예: Lindemann–Weierstrass)을 피하면서 임의의 실수 계수를 처리합니다.
정규화 (Canonicalization): 구현체는 유리 함수(rational functions)와 지수 항을 정규화하기 위해 캐노니컬라이저(canonicalizer)를 사용하여 softmax 리스케일링과 같은 사례를 처리합니다.

주요 기여

선형 시간 경합/데드락 알고리즘: 배리어 동기화를 사용하는 Structured-CTA에서의 경합 및 데드락 탐지를 위한 선형 시간 알고리즘.
결정 가능성 증명: 실수 상에서의 다변수 다항식 및 지수 함수를 포함하는 등식 $\sum p_i(\bar{x})e^{h_i(\bar{x})} = 0$ 에 대한 결정 가능성 증명.
Volta 구현: 실용적인 GPU 커널 클래스를 위한 최초의 동등성 검사기. 컨볼루션, 행렬 곱셈, 리덕션(reduction), 어텐션 메커니즘의 정확성을 검증함.
형식적 검증: 핵심 합류성 레마(confluence lemma)는 Agda에서 공식적으로 검증되었으며, 시스템은 정의된 프로그램 클래스에 대해 건전성(soundness)과 완전성(completeness) 보증을 제공함.

결과

Volta는 다음과 같은 다양한 벤치마크를 통해 평가되었습니다:

인간 최적화 커널: 표준 튜토리얼의 행렬 곱셈 및 리덕션 커널에 대해 동등성을 검증했습니다. 또한, 폐기된 warp-synchronous 실행 패턴으로 인해 데이터 경합을 포함하고 있던 리덕션 커널 버전(Red-5, Red-6, Red-7)을 정확하게 거부했습니다.
LLM 생성 커널: LLM(예: 2D 컨볼루션 및 행렬 곱셈용)이 생성한 커널을 검증했습니다. 한 사례에서는 LLM 자체 테스트나 NVIDIA의 Compute Sanitizer가 잡아내지 못한 LLM 생성 커널의 공유 메모리 범위 외 읽기(out-of-bounds read)를 탐지했습니다.
컴파일러 생성 커널: Tensor Core를 사용하는 TileLang 컴파일러 생성 커널의 동등성을 검증했습니다.
어텐션 메커니즘: 지수 함수와 max 연산을 포함하는 복잡한 심볼릭 식에도 불구하고, 표준 소프트맥스 구현을 최적화된 온라인 소프트맥스 공식(FlashAttention에서 사용됨)에 대해 성공적으로 검증했습니다.
성능: Volta는 복잡한 어텐션 커널(예: FlashAttention-2)을 3분 이내에 검증했습니다. 결정 절차는 지수 함수가 포함된 벤치마크에서 일반 목적 SMT 솔버(예: Z3)보다 성능이 뛰어났으며, 특히 여러 출력 요소에 걸쳐 캐노니컬라이제이션 비용을 분담할 때 더욱 그러했습니다.

의의 및 주장

본 논문은 GPU 커널을 위한 최초의 형식적 동등성 검사기를 제공한다고 주장합니다. 그 의의는 다음과 같습니다:

형식적 보증: 버그가 비용이 큰 프로덕션 배포 환경에서 경험적 테스트보다 강력한 신뢰를 제공합니다.
동시성 처리: 기존의 동등성 검사기가 처리할 수 없었던 GPU 동기화(배리어, 워프 수준 동기화)의 고유한 문제를 해결합니다.
최적화에 대한 견고성: 값을 실수로 모델링하고 비트 단위 일치(bit-wise equality)가 아닌 수학적 항등성을 체크함으로써, 최적화가 부동 소수점 연산 순서를 변경하더라도 정확성을 검증합니다.
실용적 적용 가능성: "Structured-CTA" 가정이 실제 ML 워크로드(컨볼루션, 행렬 곱셈, 어텐션)의 광범위한 영역을 커버함을 입증합니다.

저자들은 본 연구가 NVIDIA GPU(PTX를 통해)에 집중되어 있지만, 해당 기술이 다른 런타임에도 적용 가능함을 언급했습니다. 또한, 동적 제어 흐름(예: top-k/argsort)이나 비동기 프리미티브(예: pipeline, tma)를 포함하는 커널은 현재 지원하지 않으며, 이는 향후 과제로 남겨두었다는 한계점을 명시했습니다.