turboquant 및 (커스텀 커널)을 사용하여 MLX에서 실행되는 Gemma4 26b MoE

요약(TL;DR): 지난 한 주 동안 적절한 turbo quant 및 rotating KV cache 지원을 통해 Gemma4를 실행할 수 있는지 확인하며 미친 듯한 저녁 시간을 보냈습니다. 결과는 '예'였으며, 이제 제 MacBook Air M5에서 128k 컨텍스트와 4개의 동시 배치(concurrent batches)로 Gemma4 26b를 실행할 수 있습니다. 😄

mmap 없이 실행되는 8k 컨텍스트 환경에서, 이 방식은 프롬프트 처리(prompt processing), 생성 속도(generation speed), 그리고 런타임 메모리(runtime memory) 측면에서 llama.cpp를 능가합니다:

backend	model	bpw	pp tok/s	gen tok/s	runtime mem
llama.cpp	IQ4_XS + q4_0 KV + flash-attn	4.25	260.6	14.66	16.0 GB
MLX (ours)	nvfp4 + polar2	4.5	348.4	17.15	15.22 GB

이 속도에 도달하기 위해 많은 수동 튜닝(hand-tuning)이 필요했습니다. 여기에는 fp16 프롬프트 처리 속도에 근접하면서도 더 높은 배치 크기(batch sizes)를 가능하게 하는 실제 런타임 2bit 메모리 절감을 달성하기 위한 SWA 레이어용 커스텀 커널(custom kernel)이 포함되었습니다.

프롬프트 처리 속도는 배치 크기에 따라 비교적 잘 확장되지만, 주요 이득은 텍스트 생성(text generation)에서 나타납니다. 32 GB M5에서 512 토큰 길이의 프롬프트를 실행했을 때의 결과입니다:

B	pp tok/s	gen tok/s
1	353	16.0
...

직접 다운로드하여 서빙(serve)하고 싶다면, 리포지토리(repo)를 다운로드할 디렉토리에서 터미널을 열고 다음을 실행하세요:

git clone https://github.com/lovelacemadeline/gemma4-turboquant-mlx

그 다음 uv가 설치되어 있다면 (pip3로도 가능하지만 저는 uv를 선호합니다), 다음과 같이 진행하세요:

cd gemma4-turboquant-mlx
uv tool install --from . --reinstall gemma4-turboquant-mlx

설치가 완료되면, 다음 명령어로 백엔드(backend)를 구동할 수 있습니다:

mlx_lm.server --model mlx-community/gemma-4-26b-a4b-it-nvfp4

그러면 잘 작동할 것입니다. 😄

(참고: 16 GB RAM을 탑재한 Mac에서 실행하는 경우, Gemma MoE 모델의 대부분의 양자화(quants) 모델을 실행하기 위해 wired memory hack을 수행해야 합니다. 이에 대한 지침을 리포지토리에 포함해 두었습니다.)