특화된 GPU 커널 생성을 통한 극도의 효율성 달성

기존의 프로덕션 추론 시스템은 다양한 모델과 워크로드를 처리하기 위해 범용 커널에 의존합니다. 이는 GPU 연산 형태가 정적 모델 매개변수와 동적 요청 시간 요소의 조합에 의해 결정되기 때문에 최적의 방법이 아닙니다. 예를 들어, 모델이 행렬 곱셈의 차원 중 하나를 정의하는 동안 다른 차원은 각 요청의 특정 토큰 수에 따라 변동됩니다. 에이전트 기반 GPU 커널 생성에 대한 관심이 커지고 있으며, 최근의 노력들이 가능성을 보여주었습니다. 저희는 핵심적인 질문을 탐구했습니다. 커널 생성을 자동화할 수 있다면 왜 크기가 크게 다른 모델(10억 개에서 1조 개의 매개변수)이 동일한 커널에 의존해야 할까요? 런타임에 마주하는 특정 형태에 맞게 커널을 특화함으로써 극도로 높은 효율성을 달성할 수 있습니다.

이 블로그에서는 에이전트를 사용하여 GPU 커널을 생성하면서 얻은 성공 사례와 인사이트를 공유합니다. 저희는 극도의 특성화를 달성하도록 설계된 시스템인 Proteus를 구축했으며, 이를 위해서는 엄격한 최적화, 검증 및 컨텍스트 관리에 맞춤화된 하네스가 필요합니다.

기존의 코딩 하네스는 에이전트가 '보상 해킹(reward-hack)'을 하는 경향이 있기 때문에 여기서 실패하는 경우가 많습니다. 즉, 본질적인 의도보다는 규칙의 문자 그대로만 따르는 것입니다. 에이전트에게 벤치마크를 제공하면 의도한 연산이 아니라 벤치마크 자체만 최적화할 수 있습니다.

이를 해결하기 위해 Proteus는 커널을 제안하고, 통제된 참조 구현과 비교하여 검증하고, 성공한 커널의 실행 시간을 측정하며, 최상의 결과를 바탕으로 반복적으로 개선합니다. 이 프로세스는 간단해 보이지만, 성공 여부는 두 가지 근본적인 과제를 해결하는 데 전적으로 달려 있습니다. 그림 1은 저희 디자인의 간소화된 아키텍처를 보여줍니다. Proteus 하네스를 사용하여 vLLM에서 제공하는 최상의 커널보다 1.8~5.2배 빠른 Qwen 3.5 122B 커널을 생성했습니다.

검증

원래 저희는 커널 탐색을 어려운 부분으로 생각했습니다. 개선 없이 정체 상태(plateau)에 빠지지 않고 어떻게 방대한 프로그램 공간을 탐색할 것인가? 하지만 실제로 첫 번째 질문은 더 기본적이었습니다. '우리가 측정하고 있다고 생각하는 것을 실제로 측정하고 있는가?'였습니다.

모델은 제공된 점수를 최적화합니다. 특별한 편법을 쓸 필요도 없이, 단순히 평가 과정에서 특정 가정이 개입되었기 때문일 수 있습니다. 한 가지 예로 어텐션 레이어의 일반적인 단계인 RoPE(rotary position embeddings)용 커널이 있었습니다. 후보 커널이 이전 시도에서 남은 컴파일된 코드를 재사용하여, 처음부터 공정하게 다시 빌드하는 것보다 비용이 적게 드는 것처럼 보일 수 있었습니다. 또 다른 후보는 그래프(예: CUDA 그래프)에 일련의 GPU 실행을 기록하고 이를 하나의 단위로 재생하는 반면, 비교 대상인 베이스라인은 여전히 각 부분을 개별적으로 실행하여 두 대상이 동일한 작업을 수행하지 않는 경우도 있었습니다. 또한 공개된 테스트 세트에 포함된 입력 크기에는 강했지만, 보여주지 않은 크기에는 약한 경우도 있었습니다.

그래서 저희는 초기 설계 작업을 프롬프트가 아닌 검사기(checker)에 집중했습니다. 교차 검증이 필요할 때는 여러 타이머(예: CUDA 이벤트 타이머, 실제 경과 시간, CUPTI 타이머)를 사용하여 양쪽의 시간을 동일한 방식으로 측정합니다. 유지되어서는 안 되는 컴파일된 잔여 상태를 지우고, 설정(setup) 및 해제(teardown) 순서를 일관되게 유지하여 한쪽이 다른 쪽이 감당해야 하는 작업을 건너뛰지 못하도록 합니다. 다음 라운드의 시작점으로 사용하기 전에 우승한 커널의 시간을 다시 측정합니다. 후보 커널이 볼 수 없는 일부 테스트를 유지하여 시험 문제에만 맞추는 현상(오버피팅)을 방지합니다. 인위적으로 부풀려진 성능으로 평가를 "속이는" 행위를 방지하기 위해, 물리적인 GPU 대역폭과 연산 한계를 초과하는 이론적으로 불가능한 속도 향상(예: 100배 이상)을 감지하는 자동 일관성 검사를 구현합니다. 이는 에이전트가 실제 성능 향상이 아닌 하네스 메트릭에만 최적화되었던 과거 업계 사례에서 나타난 보상 해킹 함정을 방지합니다. 이러한 제약 조건이 없었다면 더 많은 커널을 생성하는 것은 대부분 더 많은 노이즈를 생성하는 것에 불과했을 것입니다.

검사기를 강조하는 것은 에이전트 기반 커널 생성의 병목 구간도 변화시킵니다. 단순히 프로그램 탐색(즉, 시스템이 변형을 반복적으로 생성하여 프로그램을 탐색하는 반복적 최적화) 작업에서는 좋은 후보가 드물기 때문에 후보를 작성하는 데 대부분의 비용이 발생합니다. 여러 초안을 병렬로 생성할 수는 있지만 검증을 건너뛸 수는 없습니다. 검증 과정을 신중하게 설계해야 하며, 검사는 실제 GPU에서 격리된 상태로 여러 번 실행되어야 합니다. 시스템은 커널을 작성하는 속도가 아니라 커널을 신뢰할 수 있는 속도만큼 빠르게 움직입니다.

컨텍스트 관리

또 다른 과제는 커널 생성 모델이 볼 수 있는 정보를 결정하는 것입니다. 이는 트레이드오프 관계에 있습니다. 모델에 더 큰 프롬프트를 제공하면 현재 최상의 커널, 최근의 실패 사례, 프로파일러 힌트, 이전 실행의 기록 등 더 많은 정보를 얻을 수 있어 도움이 될 수 있습니다. 하지만 모델이 읽는 모든 토큰에 비용을 지불해야 하므로 비용도 더 많이 듭니다. 또한 프롬프트가 커질수록 다음 시도가 빗나가기 쉬워집니다. 유용한 신호가 오래된 조언, 상충되는 팁, 다른 입력 크기나 다른 연산에 적용되는 세부 정보와 섞이게 됩니다. 모델은 어떤 문장을 신뢰해야 할지 항상 알지 못하므로 가장 두드러진 문장을 따르거나 모든 문장을 조금씩 따르게 됩니다.

정보를 너무 적게 주면 반대 현상이 발생합니다. 모든 시도가 처음부터 시작됩니다. 동일한 막다른 골목으로 다시 돌아가게 됩니다. 지난 실행이나 관련 연산에서 이어지는 내용이 없으므로 루프가 진행되지 않습니다.

저희는 이를 돕기 위해 지식 레이어(knowledge layer)를 원했습니다. 무엇이 효과가 있었는지 기억하고, 나중에 이를 재사용하며, 사람의 개입 없이 이 작업을 수행하는 것입니다. 이 레이어에는 저장된 교훈이 얼마나 세부적인지, 그리고 얼마나 광범위하게 적용되는지 사이의 두 번째 트레이드오프가 존재합니다.

매우 구체적인 메모("이 커널에서, 이 입력 크기로, 이 루프를 언롤링할 것")는 다음 시도에 정확히 필요한 것일 수 있습니다. 하지만 다음 연산, 다음 GPU 또는 다른 입력 크기에서 잘못 사용되기도 쉽습니다. 매우 일반적인 메모("온칩 메모리를 더 잘 활용할 것")는 거의 모든 곳에 적용되지만 모델에 구체적으로 무엇을 해야 하는지 알려주지 못합니다. 저희는 두 가지 실패 모드를 모두 목격했습니다. 교훈이 너무 일반적일 때는 조치 없이 실패만 재진술했습니다. 더 자세한 내용을 저장했을 때는 종종 하나의 실행에 너무 종속되어 다음 실행에 도움이 되지 않았습니다. 한 번의 긴 실행에서 모델이 읽고 쓴 내용의 대부분은 커널을 작성하는 대신 해당 메모리를 가져오고 라우팅하는 데 소비되었습니다. 메모리 레이어는 많은 작업을 수행하고 있었지만 다음 후보를 더 개선하지는 못했습니다. 그림 2는 이러한 시스템의 토큰 비용 분석을 보여줍니다. 비용의 대부분을 지식 레이어가 차지하고 있습니다.

보존할 가치가 있는 지식 버전은 더 작고 일반성과 구체성 사이의 균형을 이룹니다. 모델이 커널을 작성하려고 할 때, 프롬프트에는 신뢰도가 높은 컨텍스트만 포함되어야 합니다. 즉, 특정 상황과 조치를 쌍으로 연결한 실행 가능한 시사점(과거의 수정 사항 대 영향 매핑에서 추출됨)과 밀접하게 관련된 상위 실행의 간결한 실패 기록입니다. 계층적 태그 필터링과 하이브리드(키워드 + 시맨틱) 검색을 통해 검색되는 교훈은 실행에 옮길 수 있을 만큼 구체적이어야 하며, 적용되지 않는 부분을 명확히 할 수 있을 만큼 범위가 한정되어야 합니다. 교훈 저장소를 재구성하고 추가로 추출하는 것과 같은 더 깊은 연산은 매 시도마다 과거 실행에 대해 동기식 멀티홉 탐색을 수행하는 대신 백그라운드 작업에 속해야 합니다. 시사점이 상황과 조치를 명시할 수 없다면 프롬프트에 넣을 가치가 없습니다. 그림 3은 지식 레이어를 수정한 후의 토큰 비용 분석을 보여주며, 이 수정 이후 대부분의 토큰이 후보 생성에 사용됩니다.

사례 연구: Gated DeltaNet packed decode

구체적인 예로 Qwen 3.5 122B의 Gated DeltaNet 경로에 있는 packed decode 커널이 있습니다. 이 연산은 순환 상태(recurrent state)를 업데이트하고 패킹된 QKV 입력, 게이트 매개변수 및 상태 인덱스로부터 디코드 출력을 작성합니다. 저희는 Triton 백엔드가 탑재된 NVIDIA B200 GPU에서 전체 Proteus 루프를 실행하기 위해 이 작업을 사용했습니다. 작업 계약 검증, 참조 구현 측정, 에이전트에 후보 커널 요청, 정적 검사 및 빌드 실행, 통제된 참조와 비교하여 정확성 검증, 검증된 후보만 벤치마킹한 다음 최상의 후보를 다시 측정하는 과정을 거쳤습니다.

그림 4는 왼쪽에서 오른쪽으로 읽습니다. 베이스라인 노드는 벤치마크 기준을 0.025ms로 설정합니다. 후보 0000은 안전한 시드(seed)입니다. packed-decode 구조를 재현하고 검증을 통과했지만 참조보다 느렸기 때문에 Proteus는 이를 성공으로 처리하는 대신 측정된 상위 노드로 유지했습니다. 거기서부터 Proteus는 모든 형태에 대해 하나의 범용 커널을 최적화하는 것을 중단하고 탐색을 형태별(shape-specific) 경로로 분할했습니다.

Batch-1 복구 경로는 후보 012에서 형태별 커널을 생성하여 단일 배치 디코드 형태에서 1.5배의 속도 향상을 달성했습니다. 가장 강력한 결과는 serving-decode 경로에서 나왔습니다. 후보 030은 0.018ms로 가장 낮은 측정 커널 지연 시간을 기록했고, 후보 036은 1.6배로 가장 우수한 형태 속도 향상을 보여주었습니다. 우승한 해당 서빙 프래그먼트는 Batch=4, Key=128, Value=128 레이아웃에 특화되어 값 차원을 64개 폭의 청크로 처리하므로, 범용 대체재가 아닌 해당 특정 형태에 안전한 커널입니다.

타임라인의 마지막 우회로는 왜 트레이스가 중요한지 보여줍니다. 이후 (Triton 대신) C++ 생성 시도에서 빌드 및 생성 실패가 발생했고, 해당 분기의 시도 횟수를 모두 소진한 후 긴 실행이 종료되었습니다. 따라서 유용한 결과물은 단순히 가장 빠른 후보에 그치지 않습니다. 그것은 그림에 표시된 전체 경로입니다. 즉, 의미론적 오류는 거부되었고, 올바르지만 더 느린 커널이 측정되었으며, 실제 성능 향상은 프로덕션 커널로 안전하게 구성할 수 있는 형태로 유지되었습니다.

다음으로 작업할 내용

우리가 구축한 진화 루프는 작성자로서 엄격합니다. 이 루프는 종종 고정된 패턴으로 모델을 호출합니다. 즉, 현재 가장 좋은 커널을 가져와서 약간 수정하고, 확인하고, 반복하는 방식입니다. 이는 에이전트에게 필요한 자율성을 빼앗아 갑니다. 구조를 쉽게 변경하거나, 언어를 전환하거나, 막다른 디자인을 포기할 수 없습니다.

그럼에도 루프는 여전히 필요합니다. 커널을 직접 작성하기 위해서가 아니라, 에이전트에게 신뢰할 수 있는 다음 힌트를 제공하기 위해서입니다. 그 힌트는 두 가지 소스에서 나와야 합니다.

첫째, 지식 레이어와의 통신입니다. 즉, 조치를 취할 수 있을 만큼 구체적이고 적용되지 않는 범위를 명확히 알 수 있는 몇 가지 시사점을 제공하는 것입니다. 이것이 없다면 모든 시도는 제로에서 시작해야 합니다.

둘째, 신뢰할 수 있는 검증기(checker)의 결과입니다. 즉, 에이전트가 스스로 측정하지 않은 정확성과 타이밍 정보입니다. 이러한 수치는 다음 시도를 위한 힌트가 됩니다. 또한 우리가 신뢰해야 할 유일한 점수이기도 합니다. 에이전트가 자체적으로 작업 시간을 측정하면 남은 캐시, 일치하지 않는 비교, 에이전트가 볼 수 있는 테스트 등의 문제로 다시 돌아가게 됩니다.

따라서 우리가 원하는 구분은 "에이전트 대 루프"보다 더 세분화되어 있습니다. 커널이 작성되는 방식에 대해서는 에이전트에게 자율성을 부여하십시오. 루프는 메모리와 평가를 위한 채널로 유지하십시오. 에이전트가 제안하면, 루프는 에이전트가 볼 수 있도록 허용된 내용과 마지막 제안이 실제로 성공했는지 여부를 반환합니다.

Proteus는 NVIDIA B200 GPU에서 실행되는 Qwen 3.5 122B의 Gated DeltaNet 경로(선형 어텐션 스타일 블록)의 일부를 위한 특화된 커널을 구축했습니다. 개별 커널의 속도 향상은 1.8배에서 5.2배에 달했습니다.

여기서 얻은 교훈은 생성이 비용이 적게 드는 단계라는 점입니다. 검증과 컨텍스트 관리가 까다로운 부분입니다. 바로 이 부분에서 신중한 설계 시간과 혁신이 필요합니다.

에이전트 기반 GPU 커널 생성은 극단적인 특화의 놀라운 잠재력을 열어주었지만, 신뢰할 수 있고 프로덕션에 즉시 적용 가능한 하네스(harness)를 구축하는 것은 여전히 도전적인 과제입니다. 우리는 AI와 시스템의 교차점에서 가장 까다로운 문제를 해결하고 있으며, 효율적인 추론의 미래를 함께 만들어갈 대담한 엔지니어를 찾고 있습니다. 가능성의 한계를 뛰어넘는 데 열정이 있으시다면, 지금 지원하세요!

(이 글은 AI의 도움을 받아 번역되었습니다. 원문이 궁금하시다면 여기를 클릭해 주세요)