Skip to content

Flash-MoE

노트북에서 3,970억 파라미터 모델을 실행

About

  • 397B 파라미터의 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, scalex, biasx) 형태로 재배열
      • scalex와 biasx를 사전 계산하여 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% 캐시 적중률 달성
  • 통합 메모리 제약
    • Apple Silicon에서 SSD DMA와 GPU 연산이 동일 메모리 컨트롤러를 공유
    • 병렬 실행 시 GPU 대역폭 포화로 지연 급증
    • GPU → SSD → GPU 순차 파이프라인이 하드웨어 최적 형태

See also

Favorite site