메모리‑인‑컴퓨팅으로 구현한 Mixture‑of‑Experts: 교차바 multiplex와 게이트‑출력 캐시로 면적·에너지 효율 2배 이상 향상

초록

본 논문은 MoE 변환기에서 전문가(Expert)들을 교차바 단위로 공유하도록 설계하고, 전문가 그룹화·동적 스케줄링으로 부하 불균형을 완화한다. 또한 자동 회귀 생성 시 전체 토큰을 다시 라우팅해야 하는 문제를 해결하기 위해 게이트‑출력(GO) 캐시를 도입해 연산·전송량을 크게 줄인다. 실험 결과, MoE 연산부 면적 효율이 최대 2.2배, 생성 단계 성능·에너지 효율이 각각 4.2배·10.1배 향상되었으며, 전체 성능 밀도 15.6 GOPS/W/mm²를 달성했다.

상세 분석

Mixture‑of‑Experts(MoE) 구조는 다수의 전문가와 라우팅 게이트로 구성되며, 입력 토큰당 활성화되는 전문가 수가 제한돼 연산량을 크게 줄일 수 있다. 이러한 sparsity는 프로세스‑인‑메모리(PIM) 아키텍처와 자연스럽게 맞물리는데, PIM은 메모리 내부에서 가중치를 직접 연산함으로써 데이터 이동을 최소화한다. 그러나 기존 PIM 설계는 각 교차바(crossbar)마다 별도의 ADC·DAC 등 주변 회로를 필요로 하며, 이들 주변 회로가 전체 면적의 60% 이상을 차지한다는 점이 큰 병목이었다.

논문은 이 문제를 해결하기 위해 “교차바‑레벨 multiplexing”을 제안한다. 전문가들을 교차바에 배치한 뒤, 여러 교차바가 동일한 ADC·DAC 세트를 공유하도록 설계함으로써 주변 회로의 복제 수를 크게 줄인다. 이때 전문가들의 활성화가 희소하므로 동시에 여러 교차바가 같은 주변 회로를 사용할 상황이 드물어, 공유에 따른 충돌을 최소화할 수 있다.

하지만 교차바 공유는 구조적 contention을 야기할 수 있다. 이를 완화하기 위해 저자는 두 단계의 부하 균형 기법을 도입한다. 첫 번째는 “전문가 그룹화”로, 사전에 작은 샘플 데이터로 각 전문가의 평균 부하를 측정하고, 부하가 높은 전문가와 낮은 전문가를 같은 그룹에 배치한다(로드‑소팅 그룹화). 이렇게 하면 그룹 내 평균 부하가 비슷해져 주변 회로에 대한 동시 접근이 고르게 분산된다. 두 번째는 “동적 스케줄링”이다. 프리필(pre‑fill) 단계에서 토큰이 순차적으로 들어오는 대신, 여러 토큰을 동시에 여러 그룹에 할당하고, 데이터 재사용이 가능한 시점에 idle 슬롯을 삽입해 불필요한 데이터 전송을 최소화한다. 알고리즘은 토큰 길이에 대해 선형 시간 복잡도를 가지며, 하드웨어 파이프라인으로 구현해 지연을 숨길 수 있다.

자동 회귀 생성 시에는 전문가‑선택 라우팅(expert‑choice routing) 때문에 매 디코딩 단계마다 전체 히든 상태를 게이트에 전달해야 하는 비용이 발생한다. 기존 연구는 이 라우팅을 토큰‑선택 방식으로 바꾸어 성능을 맞췄지만, 학습‑추론 불일치를 초래한다. 논문은 이를 해결하기 위해 “Gate‑Output(GO) 캐시”를 설계한다. GO 캐시는 각 전문가가 선택한 토큰에 대한 게이트 스코어와 해당 토큰에 대한 전문가 출력값을 DRAM에 저장한다. 이후 디코딩 단계에서는 새로운 토큰만을 입력으로 받아, 기존에 캐시된 결과를 바로 재사용함으로써 게이트 연산과 전체 토큰 라우팅을 회피한다. 캐시 용량은 k × 전문가 수 × 차원(d)으로 고정돼, 토큰 길이에 비례하지 않는다. 또한, 필요 시 캐시된 상위‑k 결과만을 업데이트해 일관성을 유지한다.

실험은 Llama‑MoE‑4/16(7 B 파라미터, 32 블록, 16 전문가, top‑k=4) 모델을 대상으로 진행되었다. PIM 코어는 256 × 256 교차바, 8‑bit I/O, 130 ns 지연·0.096 nW 전력, 면적 0.635 mm²인 HERMES 사양을 사용했다. 전체 MoE 연산에 1536개의 교차바가 필요했으며, 주변 회로 공유·그룹화·스케줄링을 적용한 설계는 면적 효율을 최대 2.2배 향상시켰다. 생성 단계에서 GO 캐시와 KV 캐시를 결합하면, 8 토큰 생성 시 지연이 4.2배, 에너지 소비가 10.1배 감소했다. 최종적으로 전체 시스템의 성능 밀도는 15.6 GOPS/W/mm²에 도달했다.

이러한 결과는 MoE와 같은 고도로 희소한 신경망 구조가 PIM 아키텍처와 결합될 때, 주변 회로 공유와 부하‑균형 스케줄링, 그리고 라우팅 캐시와 같은 소프트웨어‑하드웨어 공동 최적화가 면적·에너지 효율을 크게 끌어올릴 수 있음을 입증한다.