LLM Wiki

Flash-MoE — 순수 C/Metal로 M3 Max에서 397B MoE 모델 실행

6 min read

Flash-MoE — 순수 C/Metal로 M3 Max에서 397B MoE 모델 실행

Python 제로. 순수 C + Objective-C + Metal만으로 MacBook Pro M3 Max (48GB)에서 Qwen3.5-397B-A17B를 4.36 tok/s로 실행. 3.4K★.

개요

danveloper/flash-moe는 Apple Silicon에서 대규모 MoE 모델을 실행하는 순수 C/Metal 추론 엔진이다. Python 기반 프레임워크(vLLM, HuggingFace 등)에 대한 의존성을 완전히 배제하고, 손으로 작성한 Metal 컴퓨트 파이프라인만으로 209GB 모델 가중치를 SSD에서 스트리밍 방식으로 처리한다.

  • 스타: 3,401
  • 포크: 399
  • 작성자: danveloper (1인, 24시간 만에 AI와 협업하여 구현)
  • 생성: 2026-03-18
  • 라이선스: MIT
  • 언어: Objective-C (59.5%), C (13.6%), TeX (9.7%), Python (8.7%), Metal (7.4%)

핵심 성과

48GB RAM 노트북에서 397B 파라미터 모델 실행

설정tok/s품질비고
4-bit + FMA 커널4.36우수Tool calling 완전 지원. 209GB
4-bit 기본3.90우수FMA 최적화 전
2-bit5.74양호*120GB. JSON 불안정
2-bit 피크7.05양호*따뜻한 캐시 순간 최대

대상 모델: Qwen3.5-397B-A17B

  • 총 60개 트랜스포머 레이어
    • 45개: GatedDeltaNet (선형 어텐션)
    • 15개: 표준 전체 어텐션
  • MoE 구조: 512개 전문가 중 토큰당 K=4 활성화 + 공유 전문가 1개
  • Hidden dimension: 4096
  • 4-bit 양자화 시 전체 가중치: 209GB

핵심 최적화 기술

1. SSD 전문가 스트리밍 (“Trust the OS” 원칙)

  • 전문가 가중치(209GB)를 SSD에서 pread() + GCD로 온디맨드 로드
  • K=4개 활성 전문가만 로드 (각 ~6.75MB)
  • OS 페이지 캐시가 자연스럽게 71% 캐시 히트율 달성
  • Apple의 “LLM in a Flash” 논문에서 영감
  • 커스텀 캐시(Metal LRU, malloc, LZ4 압축)는 모두 OS 페이지 캐시보다 성능이 낮았음

2. FMA 역양자화 커널

  • (nibble * scale + bias) * x → fma(nibble, scale*x, bias*x) 수학적 재배열
  • GPU FMA 유닛이 역양자화+곱셈을 단일 명령어로 처리
  • 12% 속도 향상

3. 손으로 작성한 Metal 커널

  • 4-bit/2-bit 역양자화 행렬-벡터 곱셈 (FMA 최적화 포함)
  • 퓨즈된 SwiGLU 활성화 함수
  • RMS 정규화 (2단계: 제곱합 리덕션 + 적용)
  • 배치 GPU 어텐션 (Q@K^T → softmax → scores@V)
  • GPU RoPE (Q 역인터리브 + K 정규화와 퓨즈)
  • MoE 결합 + 잔차 + 시그모이드 게이트 (단일 커널)
  • 지연 GPU 전문가 실행 (Deferred CMD3)

4. Accelerate BLAS 선형 어텐션 가속

  • GatedDeltaNet 점화식을 cblas_sscal/sgemv/sger로 처리
  • 64 헤드 × 128×128 상태 행렬 업데이트
  • 스칼라 대비 64% 빠름

레이어당 파이프라인 (4-bit, 평균 4.28ms/layer)

이전 CMD3 (GPU 비동기)
  → CMD1: 어텐션 프로젝션 + GatedDeltaNet  [1.22ms, GPU]
  → CPU: 결과 플러시                        [0.01ms]
  → CMD2: o_proj + 정규화 + 라우팅 + 공유   [0.55ms, GPU]
  → CPU: softmax + topK 라우팅              [0.003ms]
  → I/O: 병렬 pread K=4 전문가 로드         [2.41ms, SSD]
  → CMD3: 전문가 순방향 + 결합 + 정규화     [0.04ms 인코딩, 지연 실행]

핵심 인사이트: Apple Silicon에서 SSD DMA와 GPU 컴퓨트는 동일 메모리 컨트롤러를 공유 → 병렬 오버랩 불가. 직렬 파이프라인(GPU → SSD → GPU)이 최적.

58가지 실패한 최적화 (귀중한 연구 자료)

시도결과이유
LZ4 전문가 압축-13%압축 해제 오버헤드 > 캐시 절약
F_RDADVISE 프리페치변화 없음SSD DMA가 GPU 성능 73% 저하
시간적 전문가 예측-18%25% 적중률로 대역폭만 낭비
MLP 라우팅 예측기31% 정확도시간적 기준선보다 나쁨
mmap 전문가 파일5배 느림콜드 데이터 페이지 폴트 오버헤드
스핀 폴 GPU 대기-23%CPU 열 발생이 GPU와 경쟁
투기적 조기 라우팅-38%캐시 오염 + 오버헤드
투기적 디코딩변화 없음MoE에서 I/O는 토큰별 스케일

메모리 안전

  • 비전문가 가중치: 5.5GB (mmap’d, read-only)
  • Metal 스크래치 버퍼: ~200MB
  • 총 사용: ~6GB → 42GB를 OS + 페이지 캐시에 확보
  • OOM 리스크 없음. 전문가 데이터는 SSD에서 온디맨드 스트리밍

그래프 뷰

백링크