Apple Silicon向けLLM推論エンジンを開発して、MTPとメモリ帯域を検証した

Apple Silicon Mac上でローカルLLMを実行するための
オープンソース推論・サービングエンジン「ax-engine」を開発しています。

この記事では、ax-engineを作った背景、設計方針、MTPによる高速化、
mlx-lmとの比較、現時点での課題についてまとめます。

この記事は自作OSSの技術的な開発記録です。
実際のベンチマーク結果と、他のApple Silicon環境でのフィードバックを
共有していただくことを目的にしています。

ax-engineとは

ax-engineは、Apple Silicon Mac向けのLLM推論・モデルサービングエンジンです。

チャットUIやクラウドAPIサービスではなく、ローカルでモデルを実行し、
他のアプリケーションから利用できる推論ランタイムを目指しています。

主な特徴は以下です。

• Apple Silicon向けのローカル推論
• OpenAI互換API
• 1つのプロセスから複数モデルをサービング
• 対応モデルでのMTP（Multi-Token Prediction）
• Prefix reuse
• ローカル環境での推論
• 既存のOpenAI API対応ツールとの接続

GitHubリポジトリはこちらです。

defai-digital/ax-engine

開発の背景

Apple Silicon Macは、GPUとCPUがメモリを共有するUnified Memory Architectureを採用しています。

そのため、ローカルLLMを実行する際には、計算能力だけではなく、
モデルの読み出しやKV cacheの処理に必要なメモリ帯域が
ボトルネックになるケースがあります。

実際にローカルモデルを使っていると、モデルを読み込めても、
トークン生成速度が期待より遅く感じられることがあります。

そこで、Apple Silicon上でのモデル実行とサービングに焦点を当て、
推論速度、メモリ使用量、複数モデルの管理方法を検証しています。

MTPによるデコード高速化

ax-engineでは、対応モデルに対して
MTP（Multi-Token Prediction）を利用したデコードを試しています。

通常の自己回帰型デコードでは、1回のステップで次のトークンを1つずつ予測します。

MTPでは、補助的な予測ヘッドを利用して、
複数のトークン候補を先に予測し、その候補をメインモデルで検証します。

ただし、MTPの効果は常に同じではありません。

以下の要素によって結果が変わります。

• モデルのアーキテクチャ
• MTPに対応しているかどうか
• 量子化形式
• プロンプトの長さ
• コンテキスト長
• 生成するトークン数
• Apple Siliconの世代
• Unified Memoryの容量

そのため、単純なモデルサイズだけで速度を比較するのではなく、
同じモデル、同じ量子化、同じプロンプトで測定する必要があります。

ベンチマーク環境

今回の測定では、以下のような環境を使用しました。

項目	内容
チップ	M4 Pro / M5 Max
メモリ	XX GB
OS	macOS XX
モデル	Qwen XX
量子化	XX
コンテキスト長	XX
生成トークン数	XX
比較対象	mlx-lm XX
測定対象	Decode tok/s、Prefill tok/s、メモリ使用量

測定結果

特定のモデルと設定では、以下の結果が得られました。

Engine	Model	Quantization	Hardware	Decode
ax-engine	Qwen XX	XX	M4 Pro	XX tok/s
mlx-lm	Qwen XX	XX	M4 Pro	XX tok/s
ax-engine	Qwen XX	XX	M5 Max	XX tok/s
mlx-lm	Qwen XX	XX	M5 Max	XX tok/s

M5 Maxの特定の構成では、76 tok/s以上を測定しました。

また、条件によってはmlx-lmと比較して約2.4〜2.7倍の
デコード速度が得られました。

ただし、これはすべてのモデルや環境で再現される固定値ではありません。

モデル、量子化、コンテキスト長、プロンプト、MTPの設定によって
結果は大きく変わる可能性があります。

OpenAI互換API

ax-engineはOpenAI互換APIを提供しています。

そのため、OpenAI互換のBase URLを指定できるアプリケーションであれば、
クラウドAPIの代わりにローカルのax-engineへ接続できます。

curl http://localhost:XXXX/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{
    "model": "MODEL_NAME",
    "messages": [
      {
        "role": "user",
        "content": "Hello"
      }
    ]
  }'

既存のツールを大きく変更せずに、ローカル推論へ切り替えられることを
重要な設計方針の1つにしています。

他の環境からのベンチマークを募集

現在、Apple Silicon Macの世代やメモリ容量によって
どの程度の差が出るのかを調べています。

以下の環境で試した結果があれば、ぜひ共有してください。

• M1 / M2 / M3 / M4 / M5
• 16 GB / 24 GB / 32 GB / 48 GB / 64 GB以上
• Qwenなどのローカルモデル
• 異なる量子化形式
• mlx-lmやllama.cppとの比較

結果を共有する場合は、以下の情報があると比較しやすくなります。

• MacのチップとUnified Memory容量
• macOSのバージョン
• ax-engineのバージョン
• モデル名
• 量子化形式
• コンテキスト長
• Decode tok/s
• Prefill tok/s
• メモリ使用量
• エラーや互換性の問題

現時点での課題

まだ改善が必要な部分もあります。

• モデルによって対応状況が異なる
• MTPの効果がモデルと設定に依存する
• Apple Siliconの世代によって結果が異なる
• インストールやモデル変換で問題が発生する場合がある
• すべてのOpenAI互換クライアントとの互換性を保証できるわけではない

特に、ベンチマーク結果については、
同じ条件で再現できるかどうかを重視しています。

まとめ

ax-engineは、Apple Silicon Mac上でのローカルLLM推論と
モデルサービングに焦点を当てたオープンソースプロジェクトです。

MTP、Prefix reuse、複数モデルのサービング、OpenAI互換APIなどを通して、
ローカルLLMをより実用的に利用できる環境を目指しています。

まだ開発途中の部分も多いため、Apple Silicon環境で試した結果、
改善点、バグ、ベンチマーク結果などを共有していただけると助かります。

GitHub: