Apple Silicon에서 DFlash speculative decoding 구현: Qwen3.5-9B에서 85 tok/s, 3.3배 성능 향상 (MLX, M5 Max)
DFlash speculative decoding on Apple Silicon : 85 tok/s, 3.3x on Qwen3.5-9B (MLX, M5 Max)
핵심 요약
Apple Silicon 환경에서 DFlash를 활용해 Qwen 모델의 추론 속도를 최대 3.3배까지 가속화한 구현 사례.
- 성능 가속 — DFlash speculative decoding을 MLX로 구현하여 Qwen3.5-9B 기준 85 tok/s 달성함.
- 하드웨어 최적화 — Apple Silicon의 통합 메모리 구조에 맞춰 커널 최적화 및 동기화 오버헤드를 줄임.
- 기술적 한계 — 양자화 모델에서는 draft 모델의 속도가 병목이 되는 구조적 한계를 발견함.
- 향후 계획 — 27B 모델을 위한 draft 압축 및 MoE 모델 지원을 준비 중임.
Apple Silicon을 위한 DFlash(논문)의 네이티브 MLX 구현을 만들고 있음. 작은 draft 모델이 block diffusion을 통해 16개의 토큰을 병렬로 생성하고, 타겟 모델이 한 번의 forward pass로 검증함. 출력은 베이스라인과 비트 단위로 동일함(greedy exact argmax 일치).
설정: M5 Max, 64GB, MLX, CUDA 미사용.
결과
Qwen3.5-9B bf16
|Gen length|DFlash|Baseline|Speedup|
|:-|:-|:-|:-|
|1024 tokens|85 tok/s|26 tok/s|3.3x|
|2048 tokens|80 tok/s|26 tok/s|3.1x|
Qwen3.5-4B bf16
|Gen length|DFlash|Baseline|Speedup|
|:-|:-|:-|:-|
|1024 tokens|109 tok/s|41 tok/s|2.7x|
|2048 tokens|133 tok/s|42 tok/s|3.2x|
4B 모델은 생성 길이가 길어질수록 실제로 더 빨라짐. 모델이 작아서 컨텍스트가 커져도 draft/verify 균형이 잘 유지되기 때문임.
Qwen3.5-27B 양자화
|Quant|Gen length|DFlash|Baseline|Speedup|
|:-|:-|:-|:-|:-|
|8bit|1024 tokens|35 tok/s|14 tok/s|2.5x|
|8bit|2048 tokens|26 tok/s|11 tok/s|2.3x|
|4bit|1024 tokens|44 tok/s|24 tok/s|1.9x|
|4bit|2048 tokens|40 tok/s|23 tok/s|1.7x|
8bit가 4bit보다 더 나은 속도 향상 비율을 보임. int4는 검증 속도가 너무 빨라서 bf16 draft가 병목이 됨. int8을 사용하면 draft/verify 균형이 더 건강해짐.
모든 수치는 생성 전용(첫 토큰부터 마지막 토큰까지, prefill 제외)임. 모든 모델에서 수락률은 80-87% 수준임.
내가 만든 것
기존에 DFlash MLX 구현은 없었음. 런타임부터 직접 작성함. 실제로 수치를 개선한 부분:
head_dim=256 패치. Qwen3.5-9B는 head_dim=256을 사용하는데, MLX의 steel_attention이 이를 지원하지 않았음. 2줄의 패치로 빠른 SDPA 경로를 활성화함.
Sync elision. 파이프라인을 사이클당 2번의 GPU→CPU 동기화에서 1번으로 재구성함. 80+ tok/s에서 각 동기화는 약 0.5ms의 비용이 발생함.
3개의 matmul을 1개의 matmul + split으로 변경함. 레이어당 커널 디스패치 횟수를 줄임.
