The Other Half of the Memory Wall: Serving 35B MoEs from SSD with Trained Routing Prediction

Yu Lin; Yiming Wang; Runyuan Cai; Hanze Liu; Xiaodong Zeng

논문 · arXiv 2609.18063

메모리 벽의 또 다른 절반: 라우팅 예측 학습을 통해 SSD에서 35B MoE 서빙하기

The Other Half of the Memory Wall: Serving 35B MoEs from SSD with Trained Routing Prediction

Yu Lin, Yiming Wang, Runyuan Cai, Hanze Liu, Xiaodong Zeng

게시일 2026-09-16cs.AI arXiv PDF

XYZNOT로 읽기 (전문 번역과 AI 해설)

세 줄 요약

🎯 35B급 MoE는 int4에서도 약 19.5GB를 요구하여 소비자용 RAM을 초과하며, 레이어 라우팅이 이전 레이어의 출력에 의존하기 때문에 단순 SSD 오프로딩은 디스크 지연 시간으로 인해 지연을 유발합니다.
🔧 Edge0는 한 토큰 앞서 라우팅을 예측하는 레이어별 prerouter를 사용하여 SSD에서 MoE 전문가를 스트리밍하고 '라우팅으로서의 예측' 방식을 채택하며, 학생 경로에서 증류된 미병합 병렬 복구 LoRA와 결합합니다.
📈 24GB Mac mini M4 Pro에서 최대 2.9GiB 메모리 내로 35B MoE를 약 20.4tok/s 속도로 서빙하며, fp16 교사 모델 대비 평균 벤치마크 점수 차이를 3.9점 이내로 유지합니다.

초록

소비자용 하드웨어에서 Mixture-of-Experts (MoE) 추론은 가중치 메모리의 한계에 부딪힙니다. 35B급 모델은 4-bit에서 19.5GB에 달하며, 희소성은 토큰당 연산량을 줄여줄 뿐 유지해야 하는 바이트 수를 줄이지는 못합니다. SSD로의 단순 오프로딩만으로는 도움이 되지 않는데, 레이어 N의 출력이 나오기 전에 레이어 N+1의 전문가를 선택해야 하므로 연산 시간 뒤로 숨길 만큼 읽기 작업을 충분히 일찍 시작할 수 없기 때문입니다. 본 논문에서는 prerouter를 통해 이러한 격차를 해소하는 스트리밍 MoE 추론 엔진인 Edge0를 제시합니다. 레이어별 헤드가 다음 레이어의 라우팅을 한 토큰 앞서 예측하고, 이 예측을 라우팅 자체로 직접 사용하여 스테이징된 전문가 세트와 라우팅된 세트가 일치하도록 함으로써 누락을 방지합니다. 학생 경로에서 학습된 미병합 복구 LoRA는 int4 양자화 및 라우팅 대체로 인한 품질 저하를 보상합니다. 단일 24GB 머신에서 Edge0는 3GiB의 최대 활성 메모리 내에서 35B MoE를 20tok/s로 서빙하며, 5개 공개 벤치마크 평균에서 fp16 교사 모델 대비 수 포인트 이내의 성능을 유지합니다. 동일한 프레임워크에서 8B 티어도 실행되며, 프레임워크와 체크포인트 및 어댑터는 오픈 소스로 제공됩니다.

영어 원문 초록

Mixture-of-experts (MoE) inference on consumer hardware is bounded by weight memory: a 35B-class model is 19.5GB at 4-bit, and sparsity shrinks the compute per token, not the bytes that must be held. Naive offloading to SSD does not help on its own, because layer N+1's experts must be chosen before layer N's output exists, so the reads cannot start early enough to hide behind compute. We present Edge0, a streaming MoE inference engine that closes the gap with a prerouter: a per-layer head predicts the next layer's routing one token ahead, and the prediction is consumed as the routing itself, so the staged expert set equals the routed set and nothing is dropped. An unmerged recovery LoRA, trained on the student path, pays back the quality lost to int4 quantization and routing replacement. On a single 24GB machine, Edge0 serves a 35B MoE at 20tok/s inside 3GiB of peak active memory, within a few points of its fp16 teacher on average across five public benchmarks. An 8B tier runs on the same framework, and the framework, checkpoints, and adapters are open source.

논문 해설

문제 및 동기

Mixture-of-Experts (MoE) 모델은 토큰당 활성 파라미터 수를 적게 유지하면서(예: 35B 모델에서 토큰당 약 3B 파라미터 활성화) 총 파라미터를 확장합니다. 그러나 소비자용 하드웨어에서의 추론은 정적 가중치 메모리의 제약을 받습니다. 35B MoE는 int4로 양자화하더라도 약 19.5GB가 필요합니다. 24GB 머신에서는 이 크기가 운영체제와 KV 캐시 공간을 완전히 잠식하여 심각한 메모리 페이징을 유발하고 디코딩 처리량을 크게 떨어뜨립니다(예: M4 Pro Mac mini에서 상주 모델 실행 시 3.9 tok/s).

전문가 가중치를 SSD로 오프로드하면 메모리 점유율을 줄일 수 있지만, 단순한 온디맨드 로딩 방식은 직렬화 병목을 초래합니다. 레이어 $N + 1$ 의 전문가 선택이 레이어 $N$ 의 출력에 의존하기 때문입니다. 연산 시간 뒤로 숨길 만큼 SSD 읽기를 일찍 시작할 수 없어, 토큰마다 레이어당 한 번씩 디스크 지연 시간으로 인해 실행이 중단됩니다.

핵심 아이디어

가중치 티어로서의 SSD: MoE 전문가 가중치는 int4 형식으로 SSD에 유지되며 컴팩트한 스트리밍 풀로 메모리 매핑(mmap)됩니다. 최대 메모리는 전체 모델 파라미터 크기가 아니라 활성 워킹 세트에 의해 결정됩니다.
라우팅으로서의 예측 (Prediction as Routing): 레이어 $N$ 이 완료될 때까지 기다리는 대신, 레이어별 prerouter가 레이어 $N + 1$ 의 라우팅을 한 토큰 미리 예측합니다. 중요한 점은 모델이 디코드 시점에 이 예측을 라우팅 자체로 취급한다는 것입니다. 이로 인해 스테이징된 세트와 라우팅된 세트가 완전히 동일해지므로 토큰 드롭이나 폴백 로드가 구조적으로 제거됩니다.
미병합 복구 LoRA (Unmerged Recovery LoRA): 4-bit 양자화와 근사 라우팅은 품질 저하를 초래합니다. Edge0는 지식 증류를 통해 학생 라우팅 경로에서 저순위 어댑터를 학습시킵니다. LoRA를 int4 베이스에 병합하고 재양자화하면 델타의 대부분이 사라지기 때문에(로짓 수준 보존율이 18%에 불과함), Edge0는 거의 무시할 수 있는 오버헤드로 어댑터를 병렬 미병합 델타로 서빙합니다.

방법론

Prerouter 헤드 아키텍처

레이어 $N$ 에 속한 각 prerouter 헤드는 토큰 $t$ 의 어텐션 후 레이어 정규화 출력에서 실행되며, 토큰 $t + 1$ 에서의 레이어 $N + 1$ 라우팅을 예측합니다. 입력 피처는 은닉 상태와 두 개의 top- $k$ 원-핫 라우팅 벡터(현재 및 이전 토큰)를 결합합니다:

$dim = d_{model} + 2 E$

35B 티어의 경우 이는 $2048 + 2 \times 256 = 2560$ 입니다. 헤드는 다음을 계산합니다:

$head (x) = fc_{2} (erf - gelu (fc_{1} (x))) + ℓ (x)$

여기서 $ℓ (x)$ 는 다음 레이어의 라우터 가중치로 웜스타트된 선형 잔차 경로입니다.

미병합 어댑터 실행

서빙되는 MoE 레이어는 양자화된 베이스 바이트를 유지하면서 병렬 LoRA 델타를 실행합니다:

$y = W_{int4} (x) + \frac{α}{r} B A x$

3단계 학습

헤드 증류 (Head Distillation): 네트워크의 나머지 부분을 고정한 상태에서, 교사 모델의 다음 레이어 실제 라우터 결정을 모방하도록 prerouter 헤드를 학습시킵니다.
학생 경로 SFT (Student-Path SFT): 라우팅이 prerouter를 거치도록 강제하면서 교사가 생성한 텍스트를 사용하여 복구 LoRA 어댑터(어텐션, 선형 어텐션, 공유 전문가 프로젝션에 부착되며 스트리밍 전문가에는 부착되지 않음)를 학습시킵니다.
온폴리시 증류 (On-Policy Distillation): 학생 모델이 사전 라우팅된 실행 환경에서 텍스트를 생성하고, fp16 베이스 교사 모델이 생성 결과를 평가하며, 꼬리 인식 top- $k$ 확률 목적 함수를 갖는 역 쿨백-라이블러 발산을 통해 LoRA 파라미터를 최적화합니다.

결과

처리량 및 메모리 (Mac mini M4 Pro, 24 GB):
- edge0-35b ( $K = 4$ ): 19.5 GB int4 체크포인트에서 단 2.9 GiB의 최대 활성 메모리만을 사용하여 20.4 tok/s 디코드를 달성했습니다(완전 메모리 상주 시 3.9 tok/s 및 18.2 GiB와 비교).
- edge0-8b ( $K = 8$ ): 1.5 GiB 최대 메모리 내에서 28.0 tok/s 디코드를 달성했습니다.
프리페칭 속도 향상 (MacBook M2, 16 GB):
- 순수 반응형 온디맨드 스트리밍과 비교하여, 1토큰 prerouter 프리페칭은 라우팅 폭 전반에 걸쳐 디코드 속도를 +80% ~ +84% 향상시켰습니다( $K = 2$ : 4.8에서 8.6 tok/s, $K = 4$ : 3.5에서 6.4 tok/s, $K = 8$ : 1.8에서 3.3 tok/s).
- 메인 스레드의 블로킹된 로드 시간이 대폭 감소했습니다(예: $K = 4$ 에서 스텝당 244.0 ms에서 101.9 ms로 감소).
모델 품질 (OpenCompass):
- 5개 벤치마크(AIME 2026, HumanEval, GPQA-Diamond, MMLU-Pro, IFBench)에 걸쳐, edge0-35b는 평균 79.2점을 기록하여 fp16 Qwen3.6 베이스의 83.2점과 비교해 3.9점 차이를 보였습니다.
- edge0-8b는 69.9점을 기록하여 Ling 3.0 tiny fp16 베이스의 72.7점과 비교해 2.8점 차이를 달성했습니다.

한계점 및 열린 질문

서빙 동시성: Edge0는 단일 요청에 대해 FIFO 직렬화 방식으로 동작합니다. 배칭은 전문가 워킹 세트와 캐싱 역학을 변경하지만, 이는 단일 스트림 프로파일링에서 다루어지지 않았습니다.
CPU 그래프 빌드 오버헤드: 디코드 속도는 SSD 최적화로 완화할 수 없는 CPU 측 프레임워크 오버헤드(40개 레이어에 걸친 MLX 그래프 빌드에서 스텝당 약 44 ms 소요)에 의해 제한됩니다.
추론 품질 저하: 성능 저하는 복잡한 추론 벤치마크에 집중되어 나타났습니다(예: 35B의 경우 AIME에서 6.1점 하락, 8B의 경우 10.0점 하락).
하드웨어 의존성: prerouter의 상대적 처리량 이득은 디스크 지연 시간에 따라 달라지며, 가중치 캐시가 웜 상태이거나 초고속 스토리지 티어에서 실행될 때는 이득이 감소합니다.

핵심 용어

Edge0: 라우팅 예측을 통해 읽기 작업과 연산을 중첩하고 미병합 복구 LoRA로 품질 손실을 보상함으로써 SSD에서 모델을 서빙하는 스트리밍 MoE 추론 엔진입니다.
prerouter: 레이어 N의 은닉 상태와 최근 라우팅 기록을 사용하여 레이어 N+1의 라우팅 결정을 한 토큰 앞서 예측하는 경량 레이어별 신경망 헤드입니다.
prediction-as-routing: 런타임 토큰 드롭이나 미스 페널티를 없애기 위해 prerouter가 예측한 전문가 세트를 디코드 시점에 실제 레이어 라우팅으로 직접 실행하는 설계 원칙입니다.
recovery LoRA: 학생 모델의 사전 라우팅 경로에서 지식 증류를 통해 학습되고, int4 양자화 및 라우팅 근사 오차를 보상하기 위해 미병합 병렬 브랜치로 실행되는 저순위 어댑터입니다.
staged decode: 스텝별 메모리 할당을 방지하기 위해 GPU의 슬롯 테이블을 통해 전문가 인덱스를 매핑하는 고정 슬롯 더블 버퍼링 기반 전문가 실행 모드입니다.
incremental sticky-slot stack (incr_stack): 모든 전문가 텐서를 매번 다시 조합하는 대신, 영구 스택된 전문가 텐서를 제자리(in-place)에서 업데이트하여 스텝당 변경된 전문가 슬롯만 덮어쓰는 메커니즘입니다.
unmerged delta: LoRA 어댑터를 양자화된 베이스 가중치에 융합하지 않고 병렬로 실행하여, 재양자화 과정에서 발생하는 미세 적응 정보의 손실을 방지하는 방식입니다.
memory wall: 프로세서 처리량과 DRAM 대역폭/용량 간의 격차를 의미하며, 로컬 LLM 서빙에서는 동적 KV 캐시 상태와 정적 모델 가중치 저장 공간으로 나뉩니다.

자주 묻는 질문

Edge0는 왜 예측을 단순 힌트로 사용하는 대신 라우팅 자체로 사용할까요?

예측을 힌트로 사용하면 예측이 빗나갔을 때 런타임 폴백을 수행하거나 토큰을 드롭해야 합니다. 예측을 라우팅으로 직접 사용하면 스테이징된 전문가 세트가 항상 라우팅된 세트와 일치하게 되므로, 근사 비용을 오프라인 학습으로 넘겨 지식 증류를 통해 품질을 회복할 수 있습니다.

복구 LoRA를 4-bit 가중치에 병합하지 않고 미병합 상태로 서빙하는 이유는 무엇인가요?

LoRA 가중치 델타의 RMS는 약 10^-3으로, 4-bit 양자화의 스텝 크기보다 작습니다. 병합 후 재양자화하면 어댑터 효과의 대부분이 지워져 로짓 수준에서 18%만 남게 되지만, 병렬 미병합 델타 방식을 취하면 충실도를 100% 유지할 수 있습니다.

Edge0가 동일 토큰 내가 아니라 한 토큰 앞서 예측하는 이유는 무엇인가요?

동일한 토큰 내에서 예측하려면 레이어별로 GPU 파이프라인 플러시가 필요하여 스텝당 30~100 ms가 소요되며, 이는 숨기고자 하는 로드 시간보다 깁니다. 한 토큰 앞서 나가면 모든 스테이징 레이어를 단일 배치 플러시로 예측할 수 있어, 전체 디코드 스텝의 연산 시간 동안 SSD I/O를 충분히 숨길 수 있습니다.

Edge0는 어떻게 24GB 머신에서 완전 메모리 상주 베이스라인보다 더 빠르게 디코딩하나요?

19.5GB 모델이 메모리에 완전히 상주하면 18.2GiB의 RAM을 차지하여 운영체제와 KV 캐시를 위한 공간이 거의 남지 않으므로, 심각한 OS 메모리 페이징이 발생해 디코드 속도가 3.9 tok/s로 떨어집니다. Edge0는 약 2.9GiB의 최대 활성 메모리 내에서 SSD로부터 전문가를 스트리밍하여 OS 메모리 압박을 완전히 피하고 20.4 tok/s로 실행됩니다.

이 페이지의 해설은 AI가 논문을 바탕으로 작성했으며 오류가 있을 수 있습니다. 자세한 내용은 원문을 확인하세요.

XYZNOT로 읽기 (전문 번역과 AI 해설)