Flash-MoE: Running a 397B Parameter Model on a Laptop

Context

대규모 언어 모델(397B 파라미터)은 일반적으로 고가의 GPU 클러스터가 필요하며, 개인용 노트북에서 실행하기 어렵다는 제약이 있었다.

Technical Solution

• SSD Expert Streaming: 209GB 모델 가중치를 NVMe SSD에서 필요한 K=4개 활성화 전문가만 병렬 pread()로 온디맨드 로드, OS 페이지 캐시(~71% 히트율)에 의존
• FMA-최적화 역양자화 커널: (nibble × scale + bias) × x를 fma(nibble, scale×x, bias×x)로 재정렬하여 Fused Multiply-Add 명령 활용, GPU 연산 시간 12% 단축
• Metal 컴퓨트 셰이더: 4-bit/2-bit 역양자화 행렬-벡터 곱, SwiGLU 활성화, RMS 정규화, 배치 GPU 어텐션, RoPE, MoE 결합을 손-튜닝된 셰이더로 구현
• 레이어별 파이프라인: GPU 어텐션/선형 → SSD에서 전문가 로드 병렬화 → GPU 전문가 계산을 다음 레이어 준비 중에 지연 실행으로 오버래핑
• Linear Attention BLAS: Accelerate cblas_sscal/cblas_sgemv/cblas_sger 활용으로 GatedDeltaNet 상태 행렬 업데이트 64% 가속

Impact

• 기본 구성(4-bit 전문가, FMA 커널): 4.36 tokens/second 달성
• FMA 커널 최적화 전: 3.90 tokens/second → 최적화 후: 4.36 tokens/second (12% 향상)
• 2-bit 양자화(신뢰도 낮음): 5.74 tokens/second, 120GB 디스크 필요
• 개별 토큰 피크 성능(2-bit, 따뜻한 캐시): 7.05 tokens/second
• Linear Attention 어텐션 계산: 0.78ms → 0.28ms (64% 단축)
• C BPE 토큰화기: 180ms vs 기존 3500ms (20배 빠른 시작)
• 레이어별 평균 파이프라인 시간(4-bit): 4.28ms
• OS 페이지 캐시 자연 히트율: ~71%

Key Takeaway

Apple Silicon의 통합 메모리 아키텍처에서 SSD DMA와 GPU 연산이 메모리 컨트롤러를 놓고 경합하므로, 커스텀 캐싱보다 OS 페이지 캐시를 신뢰하고 직렬 파이프라인(GPU → SSD → GPU)으로 하드웨어 최적성을 달성할 수 있다는 원칙이 도출되었다.