AscendCraft 자동 Ascend NPU 커널 생성

초록

AscendCraft는 경량 DSL을 도입해 LLM이 AscendC 커널을 단계적으로 생성하도록 설계된 시스템이다. DSL 단계에서 연산 패턴·타일링·데이터 흐름을 기술하고, 이후 구조화된 LLM 하향 패스를 통해 AscendC 코드로 변환한다. MultiKernelBench 7개 연산군에서 98.1% 컴파일 성공, 90.4% 기능 정확도를 달성했으며, 46.2% 커널이 PyTorch eager 실행보다 동등하거나 더 빠른 성능을 보였다. 또한 신규 mHC 아키텍처 커널 두 개를 자동 생성해 기존 구현 대비 6.6×·3.0× 가속을 기록했다.

상세 분석

본 논문은 NPU 전용 프로그래밍 모델인 AscendC의 복잡성을 완화하고, 대규모 언어 모델(LLM)의 코드 생성 능력을 효과적으로 활용하기 위해 두 단계 파이프라인을 제안한다. 첫 번째 단계는 “경량 DSL”을 정의하는데, 이는 기존 AscendC가 요구하는 저수준 메모리 정렬, 버퍼 관리, 큐 동기화와 같은 세부 사항을 추상화하면서도, 연산의 핵심인 타일링 전략, 데이터 흐름, 온칩 메모리 할당 등을 명시적으로 표현한다. DSL은 Triton 스타일의 간결한 구문을 채택해 LLM이 구조적 오류 없이 핵심 알고리즘을 기술하도록 돕는다.

두 번째 단계는 DSL → AscendC 변환을 위한 “구조화된 LLM 하향 패스”이다. 각 패스는 DSL의 특정 요소(예: copy‑in, compute, copy‑out 블록, 버퍼 선언, 큐 연산)를 담당하도록 설계되어, 변환 과정에서 발생할 수 있는 하드웨어 제약 위반을 사전에 차단한다. 특히, copy‑in·compute·copy‑out 순서를 AI‑Core 함수에 매핑하고, 버퍼 크기·정렬을 자동 검증함으로써 컴파일 오류를 크게 감소시켰다.

실험에서는 MultiKernelBench에 포함된 7개 연산 카테고리(벡터, 스칼라, 큐브 등) 전반에 걸쳐 98.1%의 컴파일 성공률을 기록했으며, 기능 정확도는 90.4%에 달했다. 성능 측면에서는 46.2%의 커널이 PyTorch eager 구현을 능가하거나 동등한 수준을 보였고, 82.7%는 최소 20% 수준, 57.7%는 80% 수준을 달성했다. 이는 기존 LLM 기반 직접 AscendC 생성이 5% 이하의 정확도를 보였던 것과 큰 차이를 만든다.

또한, 논문은 새롭게 제안된 mHC 아키텍처의 두 실제 커널을 자동 생성해 각각 6.6배·3.0배의 가속을 달성했으며, 이후 인간 개발자가 LLM과 협업해 추가 최적화를 수행해 최종적으로 15.9배·7.2배까지 성능을 끌어올렸다. 이는 DSL이 단순 코드 스켈레톤 제공을 넘어, 실무에서의 성능 튜닝 기반을 제공한다는 점을 시사한다.

핵심 인사이트는 다음과 같다. (1) 고수준 DSL이 LLM에게 “무엇을 해야 하는가”를 명확히 제시하면, LLM은 저수준 세부 사항을 안정적으로 채워 넣을 수 있다. (2) 하향 패스를 구조화하고 제약을 명시적으로 삽입하면, 하드웨어 특화 언어의 문법·시맨틱 오류를 사전에 차단해 컴파일 성공률을 크게 높인다. (3) 카테고리별 전문가 DSL 예시를 제공함으로써, LLM이 동일 카테고리 내 새로운 연산에도 일반화된 최적화 전략을 적용하도록 유도한다. 이러한 설계 원칙은 Ascend뿐 아니라 다른 전용 가속기(예: TPU, Habana)에도 적용 가능할 것으로 기대된다.