エンジニアリング調査：この MLX 移植はさらに 10× 速くなれるか？

日付：2026-09-19。マシン：Apple M3 Max、GPU コア 40、統合メモリ 128 GiB、 macOS 27.2、MLX / MLX Metal 0.32.2、FP16 推論。このレポートには、自作の Metal カーネルを含む実際のローカル実験が含まれます。本番ランタイムを 変えず、 量子化した重みも公開しません。

テストしたエンジニアリング上の変更は 10× を達成しません。 交互に並べた計測は、 コンパイルと最終 decision-head 層の未使用出力の枝刈りによる、控えめで形状依存の 改善を支持します。選んだケースは、ラウンドごとのペア中央値を使って約 3–8% 改善しました。 より大きなバッチの一部の区間は改善を含みません。自作の厳密な erf の GELU/gate カーネルは 数値的には成功しましたが、MLX のコンパイルに対する一貫した追加のエンドツーエンドの利得は 得られませんでした。素朴な 8 ビットと 4 ビットのバックボーン量子化はストレージを減らし、 より大きなパイロットワークロードを加速できず、予測や較正済み確率を変えました。

数学的な限界と近似のトレードオフは MATH_10X_RESEARCH.md で別に調べています。 元の実装レビューは PERFORMANCE_RESEARCH.md にあり、リリースされた チェックポイントのベンチマークは引き続き BENCHMARKS.md です。

実験の統制と限界

すべての研究用 GPU 作業は直列に実行しました。他のエージェントの作業は CPU／ファイルシステム／ ネットワークのみを使いました。マシンは AC 電源で、記録された pmset の熱／性能の警告はなく、 実験中にスワップの使用も報告されませんでした。通常のデスクトップ操作は続いていました。 これは統制された恒温槽でも、それ以外の点でアイドルな専用ベンチマークマシンでもありません。

最初のスクリーニング実行は、各候補を新しいプロセスで、4–5 回のウォームアップと 12–16 個の サンプルで実行しました。それらは実行間の大きなドリフトを明らかにしました。たとえば英語の 単一質問のパイロットは 1.24× のコンパイル改善を示唆しましたが、その後の交互実験では 約 1.03× しか見つかりませんでした。したがって逐次パイロットのレイテンシはスクリーニングの 根拠であり、主要な因果的な高速化の主張ではありません。

確認用スクリプト paired.py：

• 各ラウンド内で候補の順序を回転させ、そのラウンドのすべての候補に同じ入力を 使います。
• ラウンド間で実際の state テキストを変えます。最大 16 個の state 変種を生成し、 同じテンソル形状の変種を残します。多言語の短いケースには 10 個、その他の報告した ケースには 16 個の変種があります。
• 異なる自然言語の質問を使い、最大の短いワークロードには 50 個の異なる instruction を 含めます。入力ハッシュを確認し、答えをキャッシュせず、質問を重複排除せず、文脈的な エンコーダ状態を再利用しません。
• 各タイマーを止める前に結果を評価し GPU を同期します。準備済みのフォワード呼び出しと、 トークナイズと出力整形を含む公開の予測経路の両方を計測します。モデルのロードは除きます。
• ウォームアップ後、英語の head／コンパイル実験には 32 ラウンド、多言語と自作 Metal の 実験には 16 ラウンドを計測します。各候補は同じラウンド数と同じ入力系列を見ます。

これらの入力は公開されたベースラインフィクスチャとは異なります。以下の比較は 研究実験の内部でのものであり、無関係な表を割って得た前後比較ではありません。 研究のバッチ上限は 64 で、リリースされた API の既定は 16 です。繰り返した state 変種は 意図的な反復計測であり、結果キャッシュはありません。

analyze.py は、ラウンドごとの eager_time / candidate_time 比と、その中央値の探索的なパーセンタイルブートストラップ区間を、 ラウンド索引の 2,000 回のリサンプルを使って計算します。これらの区間は OS のノイズや系列相関の あらゆる源を考慮せず、複数セッションでの反復の代わりにはなりません。独立に計算した p50 値の比は、 ペア比の中央値と異なることがあります。

生の JSON には、すべてのタイミング、入力ハッシュ、環境メタデータ、同等性の指標、計測時に 記録したソースのフィンガープリントが含まれます。実験スクリプトはその後、整形され、 互いに素な任意候補で拡張されました。以前のフィンガープリントはそれらの以前のスクリプト版を 記述しています。

コンパイルと最終ヘッドの厳密な枝刈り

4 つの経路を比較しました：

1. Eager： リリースされた FP16 の DecisionModel。
2. コンパイル： ロード、評価、凍結したモデルを囲む mx.compile で、通常の形状 特殊化を使います。
3. 選択 Q + コンパイル： 最後の head 層で、完全長の QKV 射影と K/V は保ちつつ、 CLS と選択肢マーカーの注意クエリだけを発行します。出力射影と FFN は、これらの 選択された出力にだけ実行します。
4. 完全注意 + 選択出力 + コンパイル： 元の完全長の QKV と SDPA の呼び出しを保ち、 その後、出力射影と FFN の前に CLS／選択肢の出力を集めます。これにより、元の注意 カーネルの形状を保ちつつ、未使用の最終ヘッドの密な作業の大半を取り除きます。

どちらの枝刈りプロトタイプも、モデルの数学的な依存関係を保ちます。それらは依然として すべての QKV 射影を計算し、数学的な上限計算にある追加の Q のみの射影の節約を実現は しません。GEMM と SDPA の形状を変えると浮動小数点の丸めが変わることがあります。 どちらのプロトタイプも decoder キャッシュでも、早期終了でも、以前の transformer 層を 落とす近似でもありません。

エンドツーエンドの p50 レイテンシ（ミリ秒）：

出典：英語のペアデータと 多言語のペアデータ。

完全注意／選択出力の経路について、ペア中央値の高速化と探索的な 95% 区間には次が含まれます：

英語の 50 問と多言語の長バッチの区間は 1 を含みます。それらは再現可能な改善を確立しません。 選択 Q の経路は多言語の短 16 と長 1 のケースでやや良かったものの、すべての形状を 支配する単一の枝刈り経路はありません。すべての候補の区間、フォワード計測、生のラウンドごとの 比は paired_analysis.json にあります。

この head／コンパイル実験の 2 つのモデル系列にわたる 1,530 の変化入力の質問比較で、コンパイルは eager のロジット、action ロジット、較正済み確率に正確に一致しました。両方の枝刈り経路は すべての 1,530 の argmax 判断で一致し、較正済み確率の最大差は 0.0001883 でした。 完全注意の枝刈り経路はさらに、各モデルの別途の 63 問フィクスチャスイートにも合格しました： 126/126 の一致で、最大確率差は英語で 4.31e-5、多言語で 6.48e-6 でした。これらは回帰チェックであり、 1,530 個の独立にラベル付けした例でのタスク精度の主張ではありません。

モデル全体とブロックごとのコンパイルは両方ともスクリーニングしました。ブロック実験も 63 個の英語フィクスチャ出力をすべて保ちましたが、モデル全体のコンパイルに対する重要な 利点は確立できませんでした。形状の特殊化はサービス内で上限づけねばなりません。モデルは Python の形状依存の reshape とマスクを使うため、shapeless=True を無差別に適用するのは 安全ではありません。公式コンパイルガイドが 形状の特殊化と状態のキャプチャを文書化しています。

最初の英語のモデル全体候補の呼び出しは 2,166.7 ms かかり、そのパイロットでウォームな フォワードの p50 が約 12.75 ms 続きました。新しい B16 形状の初回呼び出しは 272.4 ms でした。 JSON のフィールドは cold_forward と名付けられていますが、意味するのは eager の参照推論の後の 最初の候補呼び出しであり、完全にコールドなアプリケーションや新しく初期化した Metal ドライバ ではありません。後続の候補は以前にコンパイルした Metal カーネルを再利用したため、それらの 初回呼び出し時間はコールドスタート費用の統制された順位づけではありません。コンパイル済みの 英語短 1 のアクティブ／ピーク MLX メモリはパイロットで約 803.6/918.6 MiB、多言語は約 614.1/676.9 MiB でした。これらのアロケータの計測はホスト側のコンパイラ割り当てをすべては 含まず、形状の無制限な変動下でのメモリ限界を確立しません。英語のコンパイルパイロットと 多言語のコンパイルパイロットを参照してください。

選択的量子化：ストレージ削減には有用だが、速度の主張には不適

プロトタイプは、密な FP16 モデルをロードした 後 に nn.quantize を呼びます。 encoder.layers.* の線形モジュールだけを選び、アフィン群サイズ 64 とし、その結果のモデルを コンパイルします。埋め込み、norm、decision head、scorer、action head は FP16 のままです。 これにより、パックされた整数重みを現在の密なローダーでキャストすることを避け、action head の 割り切れない 1028/772 の入力幅も避けます。量子化チェックポイント形式やロード契約は 出荷していません。公式の MLX 量子化層実装が この選択の仕組みを提供します。

多言語 4 ビットの結果は、小さなフィクスチャスイートだけでは不十分な理由を示します。その 63 個のフィクスチャ argmax は同じままでしたが、26 の異なるワークロード判断のうち 7 つが 変わりました。これらは FP16 に対する一致の計測であり、正解を基準にした精度の計測では ありません。0.8008 の絶対確率変化は 80.08 パーセントポイントです。

英語短 16 のパイロット入力で、FP16 の eager／コンパイルのエンドツーエンド p50 は 91.26/87.94 ms、8 ビット／4 ビットのコンパイルは 96.66/93.20 ms でした。短 1 の量子化は そのスクリーニング実行でやや速く見え、より大きな形状はそうではありませんでした。 多言語の大きな形状のスクリーニングも速度の利得を示せませんでしたが、その逐次実行には 大きなドリフトがありました。これらの観察は、無条件の高速化やリリースの主張を退ける ことを正当化するものであり、交互に並べた量子化の反復なしに正確な低速化係数を割り当てる ことではありません。さらなる量子化の作業には、活性化を考慮した較正やファインチューニング、 代表的なラベル付き品質スイートが必要です。

生のソース：英語 8 ビット、 英語 4 ビット、 多言語 8 ビット、 多言語 4 ビット。

自作 Metal：厳密な GELU/gate 融合を実装しテストした

kernels.py は、連結された 2 つの MLP 分岐を読み、 同じ erf ベースの GELU を計算し、gate を掛け、単一の出力を書き出す、実際の自作 Metal カーネルを実装します。tanh-GELU や sigmoid 近似で代用はしません。カーネルは MLX v0.32.2 自身の erf と expm1 のヘルパーを使い、そのライセンスと通知を vendor/README.md に保っています。FP16 のみを 明示的にサポートし、安全な Metal 数学モードを使います。公式の自作カーネルガイド がこの API とその数学モードの制御を説明しています。

8 つの代表的な活性化形状にわたって、27,958,016 個のランダム生成した FP16 出力要素が、 元の操作と厳密に等しい値を持ちました。マイクロベンチマークは数値の等価性を比べるもので、 ゼロの符号ビットではありません。モデル全体の変化入力テストも厳密に一致しました：2 つの モデル系列にわたる 474/474 の質問比較、さらに両方の 63 問フィクスチャスイートで、ロジット、 action ロジット、較正済み確率の差はゼロでした。

この正しさの結果は、MLX の融合したコンパイル済み式に対する一貫した速度の利点には つながりませんでした。たとえば 1,312 トークン、中間幅 2,624 で、呼び出しごとに同期した 活性化の時間は、eager の GELU してから gate で 0.378 ms、mx.compile で 0.268 ms、 自作カーネルで 0.280 ms でした。8,192 トークン、幅 1,152 では、対応する値は 0.846/0.764/0.714 ms でした。これらのマイクロベンチマークはディスパッチと同期の オーバーヘッドを含み、スクリーニングのプローブです。分離したデバイス実行時間の 計測ではありません。完全な入力、生のタイミング、等価性チェックは microbench.json にあります。

自作カーネルはその後、すべてのエンコーダ MLP に組み込まれ、候補の順序を回転させ 入力を変える完全なモデルで計測されました：

自作カーネルの完全なペア実行は、同じ重みを持つ 2 つ目のモデルインスタンスを使うので、 未変更の実装と自作の実装が、変異や古いコンパイル済みキャプチャなしに共存します。 それらの絶対タイミングを以前の head 枝刈りの実行と比べてはいけません。控えめで混在した 結果は、自作カーネルを一般的な性能改善として公開することを支持しません。出典： 英語 Metal のペアデータと 多言語 Metal のペアデータ。

自作エンジニアリングがさらなる調査に値する箇所

モデルはすでに mx.fast.scaled_dot_product_attention、mx.fast.rope、最適化された層正規化を 呼びます。その D64 の真偽値マスク SDPA 経路は融合されており、10× の差を説明する 見逃した Flash Attention のスイッチはありません。ローカル注意は依然として密な key/value タイルを走査します。真の双方向ウィンドウカーネルは、距離の包含 <=64 とパディングの 意味を保ちつつそれらのタイルを飛ばせますが、モデル全体の演算上の機会は短い入力では 小さく、公開された長い形状では上限があります。既存のソースレビューと 数学レポートがこの区別を定量化しています。

有用な次のプロジェクトと、その根拠の要件は：

• 長い入力のウィンドウ注意： D64、実際の双方向ウィンドウ、パディングされたバッチに 対してタイル境界を特殊化する。512/1024 トークンで融合した密な SDPA と比べ、その後 完全なモデルで比べる。このカーネルは本レポートでは構築もベンチマークもしていない。
• 密なカーネルのエピローグとスケジューリング： ゲート付き MLP のエピローグを GEMM へ 融合することや、短い M の行列スケジューリングの改善を調べる。MLX はすでに特殊化した Metal GEMM 実装を使うため、それらを置き換えるには、実際のディスパッチ／カーネル プロファイルと、正確な M/N/K 形状での計測済みの利得が必要です。単体の活性化の結果は、 もう 1 つの要素単位カーネルだけでは不十分な理由を示しています。
• 長さを考慮したバッチングと共有 CPU 準備： 正確な入力 ID を保ちつつ、共有 state テキストを各質問系列を構築する前に 1 回トークナイズする。小さな項目を無関係な長い項目へ パディングしない。多言語の長 8 のパイロットは入力の準備に約 13.1 ms を費やし、それは エンドツーエンドの数百ミリ秒と比べられます。その準備を完全になくしても、この ワークロードで 10× にはなりません。待ち行列のレイテンシと一意な推論数は、あらゆる バッチングの主張に含まれねばなりません。
• 共同で答えるより小さな student： 10× が製品要件なら、密な作業の大半を取り除くよう、 または 1 つの文脈エンコードで多くの固定した質問に答えるよう、モデルを蒸留または再設計 する。これは学習済みモデルを変え、代表的なラベル付きの訓練／評価を必要とし、厳密な 移植の最適化ではありません。現在の双方向エンコーダで任意の文脈的 state/KV を質問間で 再利用するのは無効です。

8 つの単体の FP16 エンコーダ入力射影 GEMM プローブは、呼び出しごとの同期を含めて 0.66–11.55 TFLOP/s を達成しました。大きな英語の M=4096, N=5248, K=1024 プローブは 11.55 TFLOP/s、多言語の M=8192, N=2304, K=768 プローブは 7.96 TFLOP/s を達成しました。 これらは 観測されたスループット値であり、ハードウェアのピーク仕様や、グラフ全体の スループットの上限ではありません。小さな M の計測は提出と同期のコストに特に支配されます。 ストリーム化したグラフはそれらを別様に償却します。それらはどの形状がプロファイリングに 値するかを示すもので、より良いカーネルが存在しえない証明ではありません。数学レポートの 10× 同一作業スループット予算は、計測されたデバイス能力ではなく、理論上の要件のままです。

再現とリリース判断

スクリプトは既存の .venv とローカルの固定したチェックポイントを使います。GPU コマンドは 直列に実行し、正式なベンチマークと並べて実行しないでください：

# Screening: repeat for eager, compiled, blocks, q8, q4, selected-compiled.
.venv/bin/python -m experiments.engineering.run_variants \
  --model laya --variant compiled --iterations 12 --warmup 4 --quality \
  --output experiments/engineering/reproduced-compiled.json

# Primary confirmation, including 50 genuinely different questions.
.venv/bin/python -m experiments.engineering.paired \
  --model laya --iterations 32 \
  --output experiments/engineering/reproduced-laya-paired.json
.venv/bin/python -m experiments.engineering.paired \
  --model laya-multilingual --iterations 16 --cases short1,short16,long1,long8 \
  --output experiments/engineering/reproduced-multilingual-paired.json

# Hand-written kernel microbench and complete-model comparison.
.venv/bin/python -m experiments.engineering.microbench
.venv/bin/python -m experiments.engineering.paired \
  --model laya --iterations 16 --cases short1,short16,long1 --metal \
  --output experiments/engineering/reproduced-metal-paired.json
.venv/bin/python -m experiments.engineering.run_variants \
  --model laya --variant metal-compiled --iterations 5 --warmup 3 \
  --cases short1 --quality --output experiments/engineering/reproduced-metal-quality.json

# CPU-only paired analysis.
.venv/bin/python -m experiments.engineering.analyze

実験用のすべての Python ファイルは、Ruff の整形と lint チェックに合格します。安定した ランタイム、元のベンチマーク結果、公開された FP16 チェックポイントがリリース成果物の ままです。コンパイルと最終ヘッドの厳密な枝刈りは、コールド形状／キャッシュポリシーと より広い品質検証の後で、信頼できる 将来の任意の最適化です。計測された利得は、 コンパイルのレイテンシや自作カーネルを黙って既定の経路へ加えることを正当化しません。 10× の高速化、本番対応の量子化チェックポイント、計測されたローカルウィンドウカーネルの 利得はいずれも主張しません。