노트북에서 3,970억 파라미터 모델을 실행하는 Flash-MoE

4 days ago 5

Flash-MoE는 3,970억 파라미터의 Mixture-of-Experts 모델을 MacBook Pro(48GB RAM) 에서 초당 4.4토큰 이상으로 실행하는 C/Metal 기반 추론 엔진
전체 209GB 모델을 SSD에서 스트리밍하며, Python이나 프레임워크 없이 C와 Metal 셰이더만으로 구현
SSD Expert Streaming, FMA 최적화 커널, Deferred GPU Compute 등으로 GPU·SSD·CPU 병렬 효율을 극대화
4-bit 양자화 구성이 품질과 속도의 균형을 이루며, 도구 호출 기능을 포함한 프로덕션 수준 출력 생성
노트북 환경에서도 초대형 MoE 모델을 실시간 추론 가능하게 만든 경량화·최적화 사례

성능 결과

4-bit 전문가(FMA 커널) 구성에서 4.36 tok/s, 품질 우수, 전체 209GB 디스크 사용
4-bit 기본 구성은 3.90 tok/s로 FMA 최적화 전 단계
2-bit 전문가(trust OS) 구성은 5.74 tok/s로 속도는 빠르지만 JSON 출력 오류로 도구 호출 불가
2-bit 피크 단일 토큰은 7.05 tok/s까지 도달하나 실사용에는 부적합
4-bit 양자화가 실제 운용에 적합한 구성

하드웨어 환경

MacBook Pro (Apple M3 Max), 16코어 CPU(12P+4E), 40코어 GPU, 16코어 ANE
48GB 통합 메모리, 대역폭 약 400GB/s
SSD는 1TB Apple Fabric, 17.5GB/s 순차 읽기 속도
macOS 26.2 (Darwin 25.2.0) 환경

모델 아키텍처

총 60개 트랜스포머 레이어: 45개 GatedDeltaNet(선형 어텐션) + 15개 풀 어텐션
각 레이어는 512명의 전문가를 가지며, K=4명이 토큰당 활성화됨 (공유 전문가 1명 포함)
히든 차원 4096
핵심 기술
- SSD Expert Streaming
  - 전문가 가중치(4-bit 기준 209GB)를 NVMe SSD에서 병렬 pread() 로 필요 시 로드
  - 각 레이어에서 활성화된 4명의 전문가만 로드(약 6.75MB씩)
  - OS 페이지 캐시가 자동으로 캐싱을 관리하며, 별도 캐시 불필요
  - Apple의 “LLM in a Flash” 논문에서 영감을 받은 구조
- FMA 최적화 디퀀트 커널
  - (nibble * scale + bias) * x 연산을 fma(nibble, scale*x, bias*x) 형태로 재배열
  - scale*x와 bias*x를 사전 계산하여 GPU FMA 유닛이 한 번의 명령으로 수행
  - 단순 구현 대비 12% 속도 향상
- Metal Compute Shaders
  - 4-bit/2-bit 디퀀트 행렬-벡터 곱, SwiGLU 활성화, RMS 정규화, GPU 어텐션(Q@Kᵀ, softmax, scores@V), RoPE, MoE 결합+잔차+게이트 등을 핸드코드 Metal 커널로 구현
- Deferred GPU Expert Compute
  - CMD3(전문가 순전파) 명령을 비동기 제출하여 GPU가 실행 중일 때 CPU가 다음 레이어 준비
  - 결합+정규화+잔차 연산도 GPU에서 수행되어 다음 레이어로 직접 전달
- Accelerate BLAS 활용
  - GatedDeltaNet의 순환 계산에 cblas_sscal, cblas_sgemv, cblas_sger 사용
  - 스칼라 코드 대비 64% 빠른 성능
- Trust the OS
  - 커스텀 캐시 제거, OS 페이지 캐시(LRU 기반, 약 35GB)가 전문가 데이터 캐싱 담당
  - 자체 Metal LRU, malloc 캐시, LZ4 압축 캐시보다 모두 느림
  - 자연스러운 71% 캐시 적중률 달성
레이어별 파이프라인 (4-bit 기준 평균 4.28ms)
CMD3(prev) → CMD1: attention projections + delta-net [1.22ms GPU] → CPU: flush results [0.01ms CPU] → CMD2: o_proj + norm + routing + shared [0.55ms GPU] → CPU: softmax + topK routing [0.003ms] → I/O: parallel pread K=4 experts [2.41ms SSD] → CMD3: expert forward + combine + norm [0.04ms encode, DEFERRED]
통합 메모리 제약
- Apple Silicon에서 SSD DMA와 GPU 연산이 동일 메모리 컨트롤러를 공유
- 병렬 실행 시 GPU 대역폭 포화로 지연 급증
- GPU → SSD → GPU 순차 파이프라인이 하드웨어 최적 형태

프로젝트 구조

metal_infer/ 디렉터리에 주요 구성 포함
- infer.m: 약 7000라인의 완전한 추론 엔진
- shaders.metal: 약 1200라인의 Metal 커널
- chat.m: 도구 호출 지원 대화형 TUI
- tokenizer.h: 단일 헤더 C BPE 토크나이저(449라인)
- extract_weights.py, repack_experts_2bit.py: 가중치 변환 및 재양자화 스크립트
- train_predictor.py: 전문가 라우팅 예측 분석
- model_weights.bin/json, vocab.bin, tokenizer.bin: 모델 및 토크나이저 데이터
루트에는 repack_experts.py, progress.py, results.tsv 등 실험 기록 및 시각화 도구 포함

실험 결과

유지된 접근법
- FMA 디퀀트 커널**: GPU 연산 -12%,** 토큰 속도 +12%
  - OS 페이지 캐시 신뢰: Metal LRU 제거로 +38% 향상, 핵심 기반
  - GPU combine+norm 통합: CPU 왕복 제거, 파이프라인 효율 향상
- BLAS Delta-Net**: CPU 어텐션 0.78→0.28ms,**+64% 향상
- F_NOCACHE (2-bit): 페이지 스래시 방지로+3% 향상
- GPU RoPE 통합 어텐션**: 풀 어텐션 레이어에서**+2% 향상
- C BPE 토크나이저**: 시작 시간 3500ms→180ms,** 20배 개선
  - CMD3 지연 실행: GPU/CPU 오버랩, 파이프라인 효율 향상
폐기된 접근법
- LZ4 전문가 압축: -13%, 압축 해제 오버헤드가 캐시 이득 상쇄
- F_RDADVISE 프리페치: 효과 없음, SSD DMA가 GPU 속도 -73%
- Temporal 전문가 예측: -18%, 25% 적중률로 SSD 대역폭 낭비
- MLP 라우팅 예측기: 31% 정확도, 기존보다 낮음
- GPU LUT 디퀀트 커널: -2%, 간접 레지스터 접근 병렬화 저하
- GPU 프라이빗 버퍼 압축: -20%, 블릿 비용이 절감분 초과
- Spin-poll GPU 대기: -23%, CPU 발열로 GPU 성능 저하
- 전문가 파일 클러스터링: 효과 없음, NVMe가 7MB 단위로 무시
- dispatch_io: -70%, 관리 오버헤드 과다
- mmap 전문가 파일: -5배 느림, 콜드 데이터 페이지 폴트 과다
- Speculative early routing: -38%, 캐시 오염 및 오버헤드
- MTP speculative decoding: 효과 없음, MoE I/O가 토큰 단위로 확장

안전성 및 메모리 관리

비전문가 가중치: 5.5GB (mmap, 읽기 전용)
Metal 임시 버퍼: 약 200MB
총 사용량 약 6GB, 42GB는 OS 및 페이지 캐시에 사용 가능
OOM(Out of Memory) 위험 없음, 전문가 데이터는 SSD에서 필요 시 스트리밍
커스텀 캐시 없음, OS 캐시 신뢰 원칙 유지

실행 예시

4-bit 추론: ./infer --prompt "Explain quantum computing" --tokens 100
2-bit 추론(도구 호출 비활성): ./infer --prompt "Explain quantum computing" --tokens 100 --2bit
대화형 모드: ./chat
레이어별 타이밍 분석: ./infer --prompt "Hello" --tokens 20 --timing

Read Entire Article