기술 요약: 커널 포지 (Kernel Forge)

문제 정의

머신러닝 모델은 실행 시간의 대부분을 소수의 연산 커널(예: 행렬 곱셈, 컨볼루션, 정규화)에 의존합니다. 이러한 커널을 최적화하는 것은 지연 시간과 비용을 줄이는 직접적인 경로이지만, 전통적인 방식은 전문 엔지니어가 저수준 GPU 코드를 직접 작성해야 하므로 속도가 느리고 비용이 많이 들며 높은 전문성을 요구합니다.

최근 대규모 언어 모델(LLM) 기반의 에이전트 시스템은 인간의 노력을 줄이면서 커널을 생성하고 최적화할 수 있는 가능성을 보여주었습니다. 그러나 기존 도구들은 실제 시나리오에 적용할 때 다음과 같은 네 가지 결정적인 한계점을 가집니다:

1. 통합 격차(Integration Gap): 이들은 종종 독립적인 CUDA 코드를 생성하며, 이로 인해 개발자가 최적화된 커널을 수정되지 않은 PyTorch 모델에 수동으로 다시 통합해야 하는 번거로움이 발생합니다.
2. 워크로드의 협소함(Workload Narrowness): 많은 도구가 LLM 워크로드나 고립된 벤치마크 커널에만 집중하여, 비전, 확산(diffusion), 오디오 모델에서 발견되는 다양한 형태(shape), 레이아웃 및 인자(argument)를 다루지 못합니다.
3. 탐색의 한계(Search Limitations): 현재의 접근 방식은 종종 선형 정제(linear refinement)나 빔 탐색(beam-style search)에 의존하는데, 이는 초기에는 최적이 아닌 것처럼 보이거나 솔루션 공간을 효과적으로 탐색하지 못할 수 있는 유망한 최적화 경로를 제거할 위험이 있습니다.
4. 관측 가능성(Observability): 대부분의 시스템은 연구용 프로토타입이나 스크립트에 불과하여, 진행 상황을 모니터링하고, 후보를 검사하거나, 실패를 디버깅하기 위한 통합된 그래픽 인터페이스가 부족합니다.

또한, 기존 벤치마크는 무작위로 생성된 텐서와 고립된 커널 호출을 사용하여 정확성과 지연 시간을 평가하는 경우가 많습니다. 이는 텐서 형태, 활성화 분포, 인접 연산자의 동작과 같은 실제 모델 특유의 조건을 포착하지 못하며, 결과적으로 고립된 커널 성능과 실제 모델 실행 간의 간극을 만듭니다.

방법론: 커널 포지 (Kernel Forge)

커널 포지는 사용자가 모델 코드를 수정할 필요 없이 실제 PyTorch 모델 실행을 최적화하도록 설계된 오픈 소스 엔드 투 엔드 에이전트 하네스입니다. 이 시스템은 다음의 다단계 파이프라인을 통해 작동합니다.

1. 모델 섭취 및 연산 캡처 (Model Ingestion and Operator Capture)

커널 포지는 표준 PyTorch 모델과 사용자가 제공한 예시 입력을 받습니다. 시스템은 모델을 실행하여 추론 중에 실행되는 특정 연산자들을 캡처합니다. 연산자를 일반적인 패밀리로 취급하는 대신, 구체적인 호출 패턴(텐서 형태, 데이터 타입, 레이아웃, 인자, 장치 배치)을 기준으로 **런타임 호환 가능한 변체(runtime-compatible variants)**로 그룹화합니다. 이러한 변체들은 "연산자 카드(operator cards)"로 캡슐화되어 최적화 대상 역할을 수행하며, 정답 확인을 위한 참조로서 PyTorch 이저(eager) 실행을 보존합니다.

2. 에이전트 기반 커널 생성 (Agentic Kernel Generation)

시스템은 특화된 CUDA 커널을 생성하기 위해 LLM 기반 에이전트를 채택합니다.

• 입력: 에이전트는 연산자 카드와 하드웨어/런타임 구성을 전달받습니다.
• 과정: 첫 번째 후보의 경우, LLM은 CUDA 소스 코드와 런치 래퍼(launch wrapper)를 생성합니다. 이후 반복 과정에서 에이전트는 검색 컨트롤러로부터 부모 커널과 그 계보(lineage)를 전달받아, 동작을 유지하면서 지연 시간을 줄이기 위한 수정안을 제안합니다.
• 제약 사항: 에이전트는 기존 라이브러리(예: cuDNN, cuBLAS)에 위임하기보다 커스텀 CUDA 코드를 생성하도록 프롬프트됩니다. 다만, 생성된 코드가 때때로 이러한 백엔드를 효과적으로 래핑할 수 있음을 시스템은 인지하고 있습니다.

3. 검증 및 복구 (Validation and Repair)

생성된 후보들은 엄격한 검증 루프를 거칩니다.

• 컴파일: 소스는 nvcc로 컴파일됩니다. 컴파일 오류가 발생하면 제한된 범위 내의 복구 루프가 트리거됩니다.
• 실행: 성공적으로 컴파일된 후보는 캡처된 입력과 함께 실행됩니다. 런타임 실패(크래시, 잘못된 메모리 접근)는 복구 루프로 피드백됩니다.
• 수치적 정확성: 후보들은 섭취 단계에서 캡처된 PyTorch 이저 출력과 비교되며, 연산자별 수치 허용 오차를 사용합니다. 이 워크로드 특화 필터를 통과한 후보만이 검색 트리에 진입합니다.

4. 몬테카를로 트리 탐색(MCTS) 컨트롤러

선형 정제와 달리, 커널 포지는 여러 최적화 경로를 탐색하기 위해 MCTS를 사용합니다.

• 트리 구조: 컨트롤러는 후보 커널을 나타내는 노드로 구성된 지속적인 수정 트리를 유지하며, 각 노드는 소스 코드, 계보, 방문 횟수 및 측정된 지연 시간을 저장합니다.
• 선택 전략: 컨트롤러는 점진적 확장(progressive widening)과 결합된 UCT(Upper Confidence bounds applied to Trees)의 최소화 버전을 사용합니다. 이는 덜 방문된 브랜치를 탐색하는 것과 유망한 경로를 활용하는 것 사이의 균형을 맞추어, 당장은 이득이 보이지 않더라도 나중에 더 나은 솔루션으로 이어질 수 있는 후보를 다시 방문할 수 있게 합니다.
• 피드백: 지연 시간 측정값은 트리를 따라 상위로 전파되어 다음 정제를 위한 부모 커널을 선택하는 데 가이드 역할을 합니다.

5. 가드된 디스패치 및 패키징 (Guarded Dispatch and Packaging)

최종 출력은 최적화된 커널을 모델에 다시 통합하는 가드된 패키지(.cast)입니다.

• 가드 로직: 생성된 커널은 유효하고, 캡처된 호출 패턴과 호환되며, 수치적으로 정확하고, 해당 특정 변체에 대해 측정된 PyTorch 이저 지연 시간보다 빠를 경우에만 디스패치됩니다.
• 폴백(Fallback): 후보가 검증을 통과하지 못하거나 더 느린 경우, 시스템은 자동으로 원래의 PyTorch 이저 경로로 폴백합니다. 이를 통해 안전성을 보장하고 성능 저하를 방지합니다.
• 감사 가능성(Auditability): 시스템은 커스텀 생성 CUDA, 백엔드 래퍼, 그리고 폴백 경로를 구분하여 실제로 무엇이 최적화되었는지에 대한 명확한 감사 추적을 제공합니다.

사용자 인터페이스

커널 포지는 GUI와 CLI를 모두 제공합니다. 이러한 도구를 통해 사용자는 최적화 진행 상황을 모니터링하고, 후보 커널을 검사하며, MCTS 수정 트리를 확인하고, 실패를 디버깅하며, 서로 다른 모델 제공자 간의 결과를 비교할 수 있습니다.

주요 기여

1. 엔드 투 엔드 에이전트 시스템: 연산자를 캡처하고, MCTS를 통해 여러 경로를 탐색하며, 생성된 CUDA 커널을 자동으로 모델에 통합함으로써 실제 PyTorch 모델 실행을 최적화하는 도구인 커널 포지를 소개합니다.
2. 관측 가능성 및 사용성: 연구용 프로토타입과 실용적인 도구 사이의 간극을 메우기 위해 최적화 과정을 모니터링, 검사 및 디버깅할 수 있는 GUI와 CLI를 제공합니다.
3. 실제 환경 평가: 비전, 확산, LLM 워크로드를 대표하는 네 가지 다양한 Py de PyTorch 모델(ResNet-50, Stable Diffusion 3.5 Medium, Gemma 4 E2B, Qwen 3.5 35B-A3B)에 대한 평가를 수행했습니다.
4. 오픈 소스 공개: 에이전트 기반 커널 최적화의 재현 가능한 연구와 실질적인 도입을 지원하기 위해 도구를 공개합니다.

결과

NVIDIA DGX Spark (GB10 GPU 탑재)에서 평가된 커널 포지는 커널당 50회의 최적화 반복을 수행했습니다. 결과는 실제 워크로드에서의 LLM 기반 최적화의 잠재력과 과제를 동시에 보여줍니다.

• 성능 향상: 커널 포지는 14개의 커널이 PyTorch 이저 모드보다 우수한 성능을 내도록 성공적으로 최적화했습니다. 주목할 만한 속도 향상은 다음과 같습니다:
  • ResNet-50의 adaptive_avgpool2d에서 1.52×
  • Stable Diffusion 3.5 Medium의 group_norm에서 1.70×
  • Gemma 4 E2B의 softmax에서 2.83×
  • Qwen 3.5 35B-A3B의 softmax에서 1.54×
• 연산자 특이성: 상당한 속도 향상은 주로 오픈 소스/네이티브 PyTorch 경로(예: tensor_iadd, relu, softmax)에서 발견되었습니다. 반면, 성숙한 독점 또는 벤더 구현에 의해 지원되는 연산자(예: conv2d, linear, scaled_dot_product_attention)는 PyTorch 이저 성능과 비슷하거나 그 이하에 머무는 경우가 많았습니다.
• 런타임 영향: 가장 높은 로컬 속도 향상을 달出的 연산자들은 종종 전체 캡처된 런타임에서 아주 작은 부분만을 차지했습니다. 반대로, 지배적인 런타임 기여자는 개선하기 어려운 고도로 최적화된 벤더 지원 경로인 경우가 빈번했습니다.
• 탐색 비용: 평가 결과, 높은 반복 횟수(opt50)에서 탐색 비용(LLM API 사용량)이 크게 증가하지만, 이것이 반드시 전체 연산자 영역의 총체적인 성능 향상과 비례하여 나타나지는 않는다는 점이 관찰되었습니다. 특히 지배적인 연산자들의 경우 그러했습니다.

의의 및 주장

본 논문은 커널 포지가 고립된 벤치마크를 넘어 전체 PyTorch 모델에서 캡처된 연산자에 효과적으로 적용될 수 있음을 입증한다고 주장합니다.

저자들은 **배포 인식 평가(deployment-aware evaluation)**의 중요성을 강조합니다. 그들은 LLM이 빠른 독립형 커널을 생성할 수는 있지만, 이 커널들의 진정한 가치는 실제 모델 실행 내에서 어떻게 동작하느냐에 달려 있다고 주장합니다. "가드된 디스패치(guarded dispatch)" 메커니즘은 이점이 있다는 증거가 있는 곳에만 생성된 커널을 적용하면서 기존 프레임워크 구현의 신뢰성을 보존하는 핵심적인 혁신으로 제시됩니다.

결론적으로, 본 논문은 에이전트 기반 커널 최적화가 단순히 빠른 코드를 생성하는 능력이 아니라, 최적화 깊이, 탐색 비용, 그리고 실제 워크로드의 특정 특성(특히 지배적인 연산자들에 대한 성숙한 백엔드 구현의 존재) 사이의 트레이드오프를 탐색하는 능력을 기준으로 평가되어야 한다고 주장합니다.