실시간 예측 프리패칭으로 MoE 추론 효율 극대화

초록

PROBE는 Mixture‑of‑Experts( MoE ) 모델의 추론 단계에서 발생하는 공간적 부하 불균형과 시간적 급변을 실시간으로 예측·계획·프리패치하여 계산과 통신을 동시에 최적화한다. 라우터 기반 Lookahead Predictor, 하드웨어 인식 Balance Planner, Phase‑Locked Co‑Scheduling 세 가지 핵심 모듈을 통해 프리필 단계에서 1.32배, 디코딩 단계에서 1.26배의 지연 감소와 처리량 향상을 달성한다.

상세 분석

본 논문은 대규모 언어 모델에서 파라미터는 유지하면서 실제 연산량을 줄이는 Mixture‑of‑Experts 구조가 추론 시 ‘이중 페널티(double penalty)’라는 새로운 병목 현상을 만든다는 점을 지적한다. 전문가 병렬화(Expert Parallelism, EP)는 메모리 효율성을 크게 높이지만, 토큰‑전문가 매핑이 입력 의미에 따라 급격히 변동하면서 특정 GPU에 과부하가 집중되고, 이와 동시에 All‑to‑All 통신에서 네트워크 혼잡이 겹쳐 전체 레이턴시가 가장 느린 스트래거러에 의해 제한된다. 이러한 공간적 불균형과, 연속 배치와 요청이 실시간으로 들어오면서 발생하는 전문가 핫스팟의 급격한 이동(시간적 변동)은 기존의 정적 복제나 히스토리 기반 로드 밸런싱으로는 해결할 수 없다.

PROBE는 세 가지 혁신적인 설계를 제시한다. 첫째, Gate‑Initialized Lookahead Predictor는 현재 레이어의 라우터 파라미터를 고정된 prior 로 사용하고, 이전 레이어의 히든 상태를 입력으로 받아 다음 레이어의 전문가 활성화를 90% 이상의 정확도로 예측한다. 이는 라우터 연산을 별도 수행하지 않으면서도 거의 실시간에 가까운 예측을 가능하게 한다. 둘째, Hardware‑Aware Balance Planning은 ‘숨김 윈도우(hiding‑window)’라는 제약 하에 전문가 복제와 토큰 할당을 동시 최적화한다. 즉, 복제에 소요되는 전송량이 메인 연산 파이프라인의 오버랩 가능한 시간보다 크지 않도록 보장함으로써 프리패치가 크리티컬 패스를 차단하지 않게 만든다. 셋째, Phase‑Locked Co‑Scheduling은 예측·계획·프리패치 단계를 메인 스트림과 완전히 분리된 듀얼 트랙으로 운영한다. 전송을 두 단계(split‑phase)로 나누어 All‑to‑All 집합 연산과 겹치지 않게 스케줄링함으로써 네트워크 대역폭 경쟁을 완전히 제거한다.

실험에서는 GPT‑OSS‑120B와 Qwen3‑235B 같은 최신 MoE 모델을 8× H800 GPU 클러스터에 배치해, 프리필 단계에서 최대 1.32배, 디코딩 단계에서 1.26배의 레이턴시 감소와 처리량 향상을 입증했다. 특히 워크로드 변동성이 큰 상황에서도 스트래거러를 효과적으로 억제해 기존 Grace‑MoE, Libra, FasterMoE 등과 비교해 일관된 성능 우위를 보였다. 논문은 또한 메모리 사용량을 최소화하면서 KV‑캐시와 경쟁하지 않도록 복제량을 동적으로 조절하는 방법을 제시하고, 하드웨어 특성(컴퓨팅 파워 vs. 인터커넥트 대역폭)에 따라 최적의 숨김 윈도우 크기를 자동 튜닝하는 메커니즘을 포함한다.

요약하면 PROBE는 “예측 → 계획 → 프리패치”라는 연속 파이프라인을 도입해 MoE 추론의 핵심 병목인 공간·시간 불균형을 실시간으로 해소하고, 계산·통신을 동시에 최적화함으로써 레이턴시 민감 서비스에 적합한 새로운 시스템 설계를 제시한다.