기술 요약: KForge – AI 가속기를 위한 LLM 기반 교차 플랫폼 커널 생성

문제 정의

에이전트형 AI 애플리케이션의 급증으로 인해 추론 워크로드가 단일 실행 구조에서 이기종 파이프라인으로 전환되었습니다. 이러한 파이프라인은 추론(reasoning), 도구 호출(tool calls), 다중 에이전트 조율 등을 교차하여 수행하며, 각 단계는 고유한 컴퓨팅 및 메모리 프로파일을 가집 모델을 가집니다. 결과적으로 최적의 배포를 위해서는 특정 파이프라인 단계를 가장 적합한 가속기에 매핑해야 합니다. 이러한 이기종성은 심각한 시스템적 과제를 야기합니다. 개발자는 다양한 하드웨어 백엔드(NVIDIA, AMD, Intel, Apple)와 프로그래밍 모델(CUDA, Triton, SYCL, Metal 등)에 걸쳐 고성능 커널을 생산해야 하기 때문입니다.

수동으로 이러한 커널을 작성하고 포팅하는 것은 시간이 많이 소요될 뿐만 아니라 깊은 저수준 전문 지식을 요구하며, 커널 복잡성이 증가함에 따라 확장하기 어렵습니다. 대규모 언어 모델(LLM)이 코드 생성 분야에서 가능성을 보여주었으나, 저수준 커널 최적화에 적용하는 것은 여전히 어렵습니다. 이는 성능 코드의 취약성(메모리 레이아웃이나 정밀도의 미세한 변화가 정확성을 깨뜨릴 수 있음)과 비-CUDA 플랫폼에 대한 학습 데이터의 부족 때문입니다. 또한, 기존 벤치마크는 커널을 개별적으로 평가하는 경우가 많아, 커널 최적화가 실제 프로덕션 런타임 내의 전체 모델 엔드 투 엔드(end-to-end) 실행에 미치는 영향을 포착하지 못합니다.

방법론: KForge 프레임워크

KForge는 커널 엔지니어의 반복적인 워크플로우를 모방하도록 설계된 교차 플랫폼, 다단계 자율 프로그램 합성 프레임워크입니다. 이 프레임워크는 두 개의 특화된 LLM 기반 에이전트 간의 협업 루프를 활용합니다.

1. 생성 에이전트(Generation Agent): 커널 후보를 생성하고 점진적으로 개선하는 역할을 담당합니다. 이는 다음을 포함하는 프롬프트 기반 루프 내에서 작동합니다:

   • 반복적 개선(Iterative Refinement): 이전 반복에서의 컴파일 오류, 런타임 오류, 수치적 불일치를 사용하여 코드를 수정합니다.
   • 교차 플랫폼 번역(Cross-Platform Translation): 기존의 기능적 구현체(예: CUDA)를 참조하여 저평가된 플랫폼(예: Metal 또는 SYCL)을 위한 코드를 생성합니다.
   • 후보 가드레일(Candidate Guardrails): 데이터 타입 다운캐스팅이나 정밀도 변경과 같은 원치 않는 동작을 방지하기 위해 제약 조건을 강제합니다.

2. 성능 분석 에이전트(Performance-Analysis Agent): 대규모 프로파일링 데이터셋 내의 희소한 최적화 신호를 처리하기 위해 생성 에이전트와 분리된 특화된 에이전트입니다. 이 에이전트는 다양한 프로파일링 입력(Nsight Systems와 같은 도구의 프로그래밍적 메트릭 또는 Xcode Instruments의 시각적 출력)을 해석하여 구체적인 코드 수준의 권장 사항을 생성합니다. 이러한 권장 사항은 메모리 대역폭, 워프 점유율(warp occupancy), 산술 집약도(arithmetic intensity)와 같은 하드웨어 활용도 메트릭을 목표로 합니다.

합성 루프(The Synthesis Loop):
시스템은 두 단계 사이를 교차합니다:

• 기능적 패스(Functional Passes): 생성 에이전트는 코드가 컴파일되고, 오류 없이 실행되며, 수치적으로 정확한 출력을 생성할 때까지 반복합니다.
• 최적화 패스(Optimization Passes): 정확성이 확립되면, 성능 분석 에이전트가 피드백을 제공하여 다음 합성 라운드가 핸드 튜닝된 베이스라인과의 성능 격차를 좁히도록 유도합니다.

KForge 프레임워크는 모델 레지스트리를 사용하여 기본 LLM에 무관하게 유지되면서, 4개의 벤더(NVIDIA, AMD, Intel, Apple)와 6개의 프로그래밍 모델에 대해 균일한 인터페이스를 지원합니다.

주요 기여

• KForge 프레임워크: 컴파일, 정확성 및 프로파일링 피드백을 사용하여 생성 및 성능 분석 에이전트가 협업하는 다단계 자율 합성 시스템을 도입했습니다.
• 교차 플랫폼 아키텍처: 알고리즘 아이디어를 플랫폼 간에 전송할 수 있도록 4개의 가속기 벤더와 6개의 프로그래밍 모델을 지원하는 통합 인터페이스를 갖춘 시스템 설계를 제공합니다.
• 이중 모드 평가: 성숙한 벤더 튜닝 런타임과 경쟁하거나, 참조 구현이 부족한 신규 하드웨어의 브링업(bring-up) 도구로서의 효능을 입증하는 케이스 스터디를 제시합니다.

실험 결과

저자들은 서로 다른 베이스라인 가용성을 가진 두 가지 백엔드에서 KForge를 평가했습니다.

1. NVIDIA B200 (벤더 튜닝 런타임과 경쟁)

• 컨텍스트: TensorRT-LLM v1.3.0rc9를 대상으로 gpt-oss-20b 모델의 디코드 경로(decode-path) 커널을 최적화함.
• 마이크로 벤치마크: KForge 생성 커널은 다양한 배치 크기에서 베이스라인보다 일관되게 우수한 성능을 보였습니다. 특히, "MoE finalize" 커널은 배치 크기 1의 1.06배에서 배치 크기 128의 1.43배까지 속도 향상을 보였습니다.
• 엔드 투 엔드 성능: 이러한 최적화된 커널을 전체 추론 파이프라인에 통합한 결과, 엔드 투 엔드 처리량에서 2.12% 개선(2656.61 vs 2601.55 tokens/sec) 및 월클락 타임(wall-clock time)에서 2.07% 감소를 달이트했습니다. 저자들은 이러한 이득이 작아 보일 수 있으나, 이미 벤더 베이스라인에 투입된 수년간의 엔지니어링 노력을 고려할 때 매우 유의미하다고 언급했습니다.

2. Intel Arc B580 (신규 하드웨어 브링업)

• 컨텍스트: 비교 가능한 핸드 튜닝 참조가 존재하지 않는 Intel Arc B580 (Battlemage)를 위한 Triton 커널 생성.
• 벤치마크: KernelBench Level 2의 37개 GEMM + tail-ops 워크로드에 대해 평가함.
• 결과: KForge는 PyTorch eager mode 및 torch.compile 중 더 빠른 쪽보다 5.13배의 기하 평균 속도 향상을 달성했습니다.
• 최적화 기법: 생성된 커널 분석 결과 다음의 성공적인 적용이 확인되었습니다:
  • 레이어 단위 퓨전(Layer-wise Fusion): 포스트-GEMM 연산을 단일 커널로 결합하여 전체 텐서 라운드트립을 제거함.
  • 혼합 정밀도 실행(Mixed-Precision Execution): XMX 유닛에서 FP32 누산을 사용하는 FP16 행렬 곱셈을 활용함.
  • 단일 패스 리덕션(Single-Pass Reductions): 멀티 패스 GroupNorm 및 LogSumExp를 스트리밍 온라인 알고리즘으로 대체함.

의의 및 주장

본 논문은 KForge를 프로덕션급 최적화와 신속한 하드웨어 브링업을 위한 실질적인 도구로 포지셔닝합니다. 저자들의 주장은 다음과 같습니다:

• 에이전틱 합성의 생존 가능성: 생성과 성능 분석의 분리는 LLM이 저수준 커널 최적화의 복잡한 트레이드오프를 효과적으로 탐색할 수 있게 합니다.
• 이득의 확장성: 성숙한 벤더 런타임에 대한 미미한 백분율 개선(예: ~2%)조차도 배포 규모에서는 효율성 증대로 결합되어 상당한 가치를 가집니다.
• 격차 해소: 프레임워크는 기존의 핸드 튜닝 참조 없이도 성능이 뛰어난 커널을 생성함으로써, 새로운 하드웨어(예: Intel Arc)의 "콜드 스타트" 문제를 성공적으로 해결하는 동시에 성숙한 하드웨어에서도 기존 런타임과 경쟁할 수 있습니다.

결론적으로, 본 연구는 에이전틱 커널 합성이 단순한 연구 개념을 넘어 현대의 이기종 추론 워크로드를 처리할 수 있는 AI 시스템 스택의 기능적 구성 요소로 전환되고 있음을 보여줍니다.