Ollama 0.31: Multi-Token Prediction을 통한 Gemma 4 성능 향상

Ollama 0.31은 Apple Silicon에서 Gemma 4를 위한 multi-token prediction (MTP)을 구현하여, 코딩 에이전트 벤치마크 전반에서 평균적으로 거의 90% 더 빠른 토큰 생성 속도를 달성했습니다. 이 성능 향상은 기본적으로 활성화되며 모델의 출력 결과는 변경되지 않습니다.

Multi-Token Prediction (MTP) 메커니즘

Gemma 4는 메인 모델과 함께 실행되는 작고 빠른 draft model을 사용하여 다음 여러 토큰을 제안합니다. 그런 다음 메인 모델은 이러한 제안들을 단 한 번의 패스(pass)로 검증하여, draft model의 예측이 정확할 때마다 하나의 비용으로 여러 개의 토큰을 확정합니다.

이 접근 방식은 닫는 괄호, 반복되는 식별자, 보일러플레이트와 같이 예측 가능한 패턴이 자주 포함되는 코드 생성에 특히 효과적입니다. 결과적으로, 파일을 읽고 도구를 실행하기 위해 모델을 지속적으로 호출하는 코딩 에이전트의 응답성이 현저히 향상됩니다.

기술적 구현

Ollama 0.31의 속도 향상은 auto-tuning draft length, engine-level speculative decoding, 그리고 GPU kernel 최적화를 통한 세 가지 주요 기술적 최적화를 통해 이루어집니다.

Auto-Tuning Draft Length

초안을 작성할 토큰의 이상적인 수는 모델, 양자화(quantization), 하드웨어 및 텍스트의 예측 가능성에 따라 달라지기 때문에, Ollama는 런타임에 draft length를 결정합니다. 시스템은 제안 수락률과 검증 패스 지속 시간을 추적하여 초당 토큰 수를 최대화하는 길이를 선택합니다. 만약 제안이 지속적으로 거부된다면, 시스템은 추측(speculation)이 성능을 저하시키지 않도록 자동으로 표준적인 one-at-a-time decoding 방식으로 되돌아갑니다.

Speculative Decoding Engine

생성 프로세스는 CPU round-trips를 피하기 위해 GPU에서 완전히 실행되는 특정 시퀀스를 따릅니다:

Drafting: draft model이 토큰 시퀀스를 예측합니다.
Sampling and Verification: 메인 모델이 단 한 번의 패스로 제안된 전체 시퀀스를 검증합니다.
Commitment: 수락된 토큰은 유지됩니다. 거부된 토큰의 경우, 엔진은 각 제안 전에 기록된 rollback point를 사용하여 이전 상태를 재계산하지 않고 마지막으로 수락된 토큰으로 되돌아갑니다.

MLX Kernel 최적화

검증은 프로세스 중 계산 비용이 가장 많이 드는 부분입니다. draft batch는 일반적으로 크기가 작기 때문에(2~8개 토큰), 단일 토큰 디코딩과 대규모 배치 prefill 사이의 크기에 해당하며, 이로 인해 표준 행렬 곱셈 커널이 비효율적일 수 있습니다.

Ollama는 각 가중치 블록을 한 번만 읽고 언팩(unpack)하여 전체 배치에 걸쳐 재사용하는 특수 커널을 MLX에 기 contributed했습니다. M5 Max에서 nvfp4를 사용할 경우, 이 최적화는 동일한 계산을 유지하면서 불필요한 작업을 제거하여 Gemma 4의 가장 큰 행렬 곱셈을 2배에서 2.5배 더 빠르게 만듭니다.

벤치마크 및 사용 가능 여부

성능은 코딩 에이전트와 함께 실제 프로그래밍 작업을 시뮬레이션하는 Aider polyglot 벤치마크를 사용하여 측정되었습니다. 결과에 따르면 M5 Max에서 Gemma 4 12B (nvfp4)의 생성 속도가 거의 90% 향상되었습니다.

이러한 개선 사항을 사용하려면 macOS용 Ollama 0.31 이상 버전을 다운로드해야 합니다. 이전에 Gemma 4를 다운로드했던 사용자는 다음 명령어를 사용하여 모델을 다시 pull 해야 합니다:

ollama pull gemma4:12b-mlx