dflash-mlx — Apple Silicon DFlash 추론 (MLX)
Apple Silicon에서 DFlash(block-diffusion speculative decoding)를 MLX로 구현한 고성능 추론 엔진. 260+ star.
Key Points
- DFlash(block-diffusion speculative decoding)의 MLX 포트
- Apple Silicon에서 3~4.6x 추론 속도 향상 (Qwen3.5 기준)
- Lossless 추론: 모든 토큰 검증 후 출력 — 토큰 생성 품질 유지
- MLX의 speculative-decoding primitive 부재 → Metal 커널부터 재구현
- OpenAI-compatible server mode (
dflash-serve) 지원
What It Does
DFlash는 작은 block-diffusion 모델(드래프터)이 여러 토큰을 동시에 제안하고, 대상 모델(타겟)이 이를 한 번의 순전파로 검증하는 speculative decoding 기법입니다. 원본 DFlash는 CUDA 대상이며, dflash-mlx는 Apple Silicon용 MLX 포트입니다.
동작 방식
- 드래프터: 작은 block-diffusion 모델이 16개 토큰을 병렬로 생성
- 타겟 검증: 대상 모델이 16개 토큰을 single forward pass로 검증
- 최장 정렬 접두사 수용: 가장 긴 올바른 토큰 접두사를 채택
- Lossless: 모든 출력 토큰이 타겟 모델로 검증됨
기술적 도전과 해결책
MLX에는 speculative-decoding primitive가 없으므로, draft/verify 루프의 모든 요소를 Metal 위에서 재구현해야 했습니다:
- Hidden-state 추출: 드래프터는 로깅이 아닌 중간 계층 활성화에 조건부. 특정 Qwen 레이어에 훅을 연결해 KV 캐시를 해치지 않고 텐서 노출
- 병렬 블록 제안: draft 모델이 block-diffusion denoising 루프로 여러 토큰 동시에 제안
- Single-pass batched 검증: 모든 제안 블록을 단일 타겟 순전파로 검증
- Tape-replay 롤백: 전체 GatedDeltaNet 상태 스냅샷 대신 tape 기록으로 롤백
- JIT SDPA 2-pass: long-context 검증(
N >= 1024)용 커스텀 Metal attention 커널
Benchmarks
M4 Max 36GB (Qwen3 BF16)
| Max new tokens | MLX-LM BF16 tok/s | dflash-mlx BF16 tok/s | 속도 향상 |
|---|---|---|---|
| 512 | 42.3 | 133.1 | 3.1x |
| 1024 | 42.0 | 144.6 | 3.4x |
| 2048 | 41.3 | 174.4 | 4.2x |
| 4028 | 40.6 | 186.4 | 4.6x |
M5 Max 64GB (Qwen3.5 BF16)
| 모델 | 토큰 | Baseline | DFlash | 속도 향상 | 수용률 |
|---|---|---|---|---|---|
| Qwen3.5-4B | 1024 | 53.48 tok/s | 197.49 tok/s | 3.69x | 88.67% |
| Qwen3.5-4B | 2048 | 53.74 tok/s | 219.83 tok/s | 4.10x | 89.26% |
| Qwen3.5-9B | 1024 | 31.09 tok/s | 127.47 tok/s | 4.10x | 88.96% |
| Qwen3.5-35B-A3B-4bit | 4096 | 140.73 tok/s | 189.62 tok/s | 1.35x | 86.96% |
Supported Models
| 타겟 | 드래프터 |
|---|---|
mlx-community/Qwen3-4B-bf16 (기본) | z-lab/Qwen3-4B-DFlash-b16 |
mlx-community/Qwen3.5-4B-MLX-bf16 | z-lab/Qwen3.5-4B-DFlash |
mlx-community/Qwen3.5-27B-4bit | z-lab/Qwen3.5-27B-DFlash |
mlx-community/Qwen3.5-35B-A3B-4bit | z-lab/Qwen3.5-35B-A3B-DFlash |
Qwen3.5 지원은 작동하지만 아직 불완전 — hybrid GatedDeltaNet + attention 아키텍처 때문에 exact partial-block acceptance가 더 복잡함.
Quick Start
git clone https://github.com/aryagm/dflash-mlx.git && cd dflash-mlx
uv sync
uv run dflash-mlx --max-new-tokens 128기본값은 Qwen3-4B BF16. 첫 실행 시 타겟/드래프트 체크포인트가 Hugging Face 캐시에 다운로드됨 (~12GB).
# 타겟/드래프트 모델 지정
uv run dflash-mlx --target-model mlx-community/Qwen3.5-4B-MLX-bf16 --draft-model z-lab/Qwen3.5-4B-DFlash
# Chat mode
uv run dflash-mlx-chat
# OpenAI-compatible server
uv run dflash-serveFeatures
- Auto draft resolution: 수동
--draft플래그 불필요 - Streaming: CLI와 서버 모두 토큰 단위 출력
- Chat templates: 기본 활성화
- Open WebUI / Continue 호환: OpenAI-compatible API
- Streaming SSE 지원
Sources
- Aryagm/dflash-mlx — DFlash MLX 포트, MIT License
- X: Tom Doerr — 소개 트윗