[Paper] 엣지 GPU‑NDP 시스템에서 효율적인 MoE 추론을 위한 스케줄링 프레임워크

Source: arXiv - 2601.03992v1

(위에 제공된 내용만으로는 번역할 텍스트가 없습니다. 번역을 원하는 본문을 알려주시면 도와드리겠습니다.)

Overview

이 논문은 Mixture‑of‑Experts (MoE) 모델이 GPU‑NDP (Near‑Data Processing) 아키텍처를 갖춘 엣지 디바이스에서 효율적으로 실행되도록 하는 새로운 추론 프레임워크를 제시합니다. 로드 불균형, GPU 활용도 저하, 비용이 많이 드는 전문가 프리페칭 문제를 해결함으로써, 저자들은 지연 시간에서 2배 이상의 속도 향상을 달성했으며, 전력 제한이 있는 엣지 하드웨어에서 복잡한 AI 워크로드를 실행할 수 있는 길을 열었습니다.

핵심 기여

• Tensor‑parallel MoE inference: 사용되지 않은 텐서 병렬성을 활용하여 대규모 전문가 가중치를 여러 NDP 유닛에 분산시켜, 저배치 엣지 추론을 가능하게 함.
• Load‑balancing‑aware scheduler: GPU와 NDP 코어 간에 전문가 연산을 동적으로 할당하여, 불규칙한 전문가 활성화 패턴에도 불구하고 전체 하드웨어 활용도를 극대화함.
• Dataset‑free pre‑fetching: 가벼운 통계 기반 전략을 도입해 가장 가능성이 높은 전문가를 사전에 예측 및 로드함으로써 비용이 많이 드는 프로파일링 단계를 제거함.
• Comprehensive evaluation: 기존 최고의 엣지‑GPU‑NDP 베이스라인과 비교해 실제 MoE 모델에서 평균 2.41× (최대 2.56×) 엔드‑투‑엔드 지연 감소를 입증함.

Methodology

Tensor Parallelism for Experts

1. 기존 MoE 추론은 선택된 각 전문가를 단일 NDP 유닛에서 실행하므로, 소수의 전문가만 활성화될 때 리소스가 유휴 상태가 됩니다.
2. 저자들은 각 전문가의 가중치 행렬을 여러 NDP 코어에 걸쳐 분할합니다(대형 언어 모델에서의 모델 병렬화와 유사). 이를 통해 하나의 전문가에 대해 다수의 NDP 유닛이 협업할 수 있어, 배치 크기가 매우 작아도(엣지 디바이스에서 흔한 상황) 유닛들을 계속 바쁘게 유지할 수 있습니다.

Scheduling Algorithm

1. 스케줄러는 주어진 입력 배치에 대해 전문가 활성화 분포를 프로파일링합니다(저비용 런타임 히스토그램 사용).
2. 그런 다음 전문가 샤드를 NDP 유닛 및 GPU에 매핑하는 경량 bin‑packing 문제를 해결하여, 메모리 제약을 만족하면서 연산 부하를 균등하게 맞춥니다.
3. 활성화 패턴이 크게 변할 때만 스케줄을 재계산하므로 오버헤드가 낮습니다.

Dataset‑Free Pre‑fetching

1. 전체 데이터 기반 프로파일링 단계를 수행하는 대신, 시스템은 각 전문가가 선택된 횟수를 실시간으로 카운트합니다.
2. 자주 사용되는 전문가들은 추론이 시작되기 전에 미리 NDP의 로컬 메모리로 복사되어, 해당 전문가들의 “콜드‑스타트” 지연 시간을 감소시킵니다.

Implementation Details

• 각 NDP 코어에 작은 온‑칩 SRAM과 전용 연산 파이프라인을 갖춘 CUDA 호환 GPU‑NDP 프로토타입 위에 구축되었습니다.
• GPU 측에서는 표준 CUDA 커널을 사용하고, NDP 측에서는 맞춤형 마이크로‑커널을 사용하며, 경량 런타임 라이브러리를 통해 조정됩니다.

결과 및 발견

실험은 여러 MoE 구성(전문가 2–8명, 은닉 크기 1K–4K)과 현실적인 엣지 워크로드(음성 인식, 추천)를 포괄합니다. 이 프레임워크는 엣지 AI 서비스에서 일반적으로 병목 현상이 되는 꼬리 지연 시간을 지속적으로 감소시킵니다.

실용적 함의

• Edge AI 서비스 (예: 음성 비서, 디바이스 내 추천) 가 이제 응답 시간을 희생하지 않고 더 크고 정확한 MoE 모델을 실행할 수 있습니다.
• Developer tooling: 스케줄링 라이브러리를 기존 추론 스택(TensorRT, ONNX Runtime)에 통합하여 NDP 하드웨어를 자동으로 활용하고, 텐서 병렬성의 복잡성을 추상화할 수 있습니다.
• Hardware design guidance: NDP 유닛의 적당한 온칩 메모리와 스마트 스케줄러를 결합하면 비약적인 성능 향상을 얻을 수 있음을 보여주며, 차세대 엣지 GPU‑NDP 칩 설계에 정보를 제공합니다.
• Cost & power savings: 높은 하드웨어 활용도는 대기 전력 소모를 낮추어 AI 워크로드를 호스팅하는 모바일 및 IoT 디바이스의 배터리 수명을 연장합니다.

제한 사항 및 향후 작업

• 현재 스케줄러는 비교적 정적인 전문가 활성화 분포를 가정합니다; 입력 도메인의 급격한 변동은 더 빈번한 재스케줄링을 필요로 하며, 이는 오버헤드를 추가합니다.
• 실험은 프로토타입 NDP 플랫폼에 한정되어 있습니다; 다양한 메모리 계층을 가진 상용 엣지 GPU로 확장하면 새로운 병목 현상이 드러날 수 있습니다.
• 사전 가져오기 전략은 단순한 빈도 카운트에 의존합니다; 보다 정교한 예측 모델(예: 강화 학습)은 미스율을 더욱 감소시킬 수 있습니다.
• 향후 작업으로는 프레임워크를 확장하여 동적 MoE 라우팅(전문가를 실시간으로 선택) 을 지원하고, 멀티 엣지 시나리오를 위한 크로스‑디바이스 스케줄링을 탐구하는 것이 포함됩니다.

저자

• Qi Wu
• Chao Fang
• Jiayuan Chen
• Ye Lin
• Yueqi Zhang
• Yichuan Bai
• Yuan Du
• Li Du

논문 정보

• arXiv ID: 2601.03992v1
• 카테고리: cs.DC, cs.AI
• 출판일: 2026년 1월 7일
• PDF: PDF 다운로드