FORGE-UGC: FX 최적화 및 레지스터 그래프 엔진 — 범용 그래프 컴파일러

자율 AI 에이전트가 NPU, GPU, CPU가 혼합된 복잡한 하드웨어에서 효율적으로 동작하려면 각 장치에 최적화된 컴파일러가 필수적이다. 이 논문은 기존 프레임워크의 불투명한 구조를 개선하여 Intel NPU 환경에서 추론 속도와 에너지 효율을 동시에 극대화하는 투명한 컴파일 파이프라인을 제시한다.

기존의 AI 모델 배포 도구들은 모델을 하드웨어 명령어로 변환하는 과정을 '블랙박스'처럼 처리하여 개발자가 내부 최적화 과정을 파악하거나 수정하기 어려웠다. 특히 Transformer 모델의 핵심인 Attention 연산은 여러 개의 작은 연산으로 쪼개져 실행되는데, 이때마다 NPU와 CPU 사이를 오가며 발생하는 데이터 전송 지연이 전체 성능의 발목을 잡는 한계가 있었다.

FORGE-UGC는 이 문제를 해결하기 위해 PyTorch의 그래프 구조를 그대로 유지하면서 NPU에 특화된 '연산 통합(Fusion)' 기술을 적용한다. 흩어져 있는 연산들을 하나의 큰 덩어리로 묶어 NPU 내부에서 한 번에 처리하도록 유도함으로써 불필요한 데이터 이동을 원천적으로 차단하는 원리이다.

또한, 메모리 관리 측면에서도 단순히 데이터를 쌓아두는 것이 아니라 각 데이터가 실제로 살아있는 시간(Liveness)을 계산하여 더 이상 필요 없는 메모리 공간을 즉시 재사용한다. 이를 통해 메모리가 제한된 엣지 장치에서도 거대 언어 모델이 병목 현상 없이 매끄럽게 동작할 수 있는 기반을 마련했다.

FORGE-UGC는 총 4단계의 컴파일 프로세스를 거친다. 1단계(Frontend)에서는 torch.export.export()를 사용하여 모델을 FX 그래프 형태로 캡처한다. [PyTorch 모델 입력 → 심볼릭 트레이싱 → ATen 연산 기반 FX 그래프 출력] 과정을 통해 모델의 의미론적 정보를 온전히 보존한다.

2단계(Middle-end)에서는 6가지 최적화 패스를 순차적으로 적용한다. 특히 Attention Fusion 패스는 [Q, K, T, Scale, Mask, Softmax, V 입력 → 패턴 매칭 및 통합 연산 변환 → 단일 NPU 커널 출력] 과정을 수행하여 연산 효율을 높인다. 3단계(Lowering)에서는 최적화된 그래프를 NPUIR로 변환하며, 각 노드에 가상 레지스터 ID와 실행 장치(NPU/CPU)를 할당한다.

4단계(Backend)에서는 Liveness Analysis를 수행한다. [명령어 인덱스 입력 → 변수의 첫 쓰기와 마지막 읽기 시점 계산 → 변수별 생존 구간 도출] 과정을 거쳐 Linear-scan 알고리즘으로 물리적 버퍼를 할당한다. 마지막으로 Priority-based Topological Sort를 통해 동일 장치 연산들을 최대한 묶어 스케줄링함으로써 장치 전환 비용을 최소화한다.

Intel AI Boost NPU 환경에서 125M8B 파라미터 규모의 6개 모델군을 대상으로 실험한 결과, OpenVINO 및 ONNX Runtime 대비 6.99.2배 빠른 컴파일 속도를 기록했다. 추론 지연 시간(Latency)은 18.2~35.7% 단축되었으며, 특히 모델의 깊이가 깊어질수록 성능 향상 폭이 커지는 경향을 보였다.

에너지 효율성 측면에서는 추론당 에너지 소비량을 30.2~40.9% 절감했다. 이는 불필요한 장치 전환을 줄이고 메모리 할당 최적화를 통해 DRAM 전력 소모를 억제한 결과이다. 수치적 정확도(Numerical Fidelity) 분석에서도 KL Divergence가 8.4e-9 미만으로 유지되어 원본 모델과 거의 동일한 출력을 생성함을 입증했다.

FORGE-UGC는 하드웨어 불가지론적(Hardware-agnostic) 설계를 채택하여 프론트엔드와 미들엔드 최적화 로직을 백엔드 구현과 분리했다. NPUIR은 가상 레지스터 추상화를 통해 하드웨어 독립적인 최적화를 가능하게 하며, 실제 실행 시점에만 물리적 버퍼와 매핑된다.

최적화 패스 중 하나인 FXOperatorFusionPass는 Linear 연산 뒤에 오는 ReLU, GELU, SiLU 등의 활성화 함수를 감지하여 단일 NNFactory 그래프로 통합한다. 이는 중간 텐서의 구체화(Materialization)를 방지하여 메모리 대역폭 요구량을 획기적으로 줄인다.

스케줄러는 장치 친화적 휴리스틱을 사용하여 이전에 선택된 장치와 동일한 장치에서 실행 가능한 연산을 우선적으로 배치한다. Llama-3.1-8B 모델의 경우, 이 기법을 통해 장치 전환 횟수를 264회에서 93회로 줄여 추론당 50~130ms의 오버헤드를 제거했다.

현재 프로토타입은 단일 배치(Single-batch) 추론만 지원하며, 실제 서비스 환경에 필요한 배치 처리 기능은 향후 과제로 남아있다. 또한 NPUIR이 현재 고정된 연산자 세트만 지원하므로 커스텀 연산자 등록을 위한 확장성이 보완되어야 한다.