鳴かぬなら、カスタムカーネル、ほととぎす

要約

MoE(Mixture of Experts):
モデル全体を動かすのではなく、必要な部分のみが動くことで 「賢さは大きいモデル並み、速さ・メモリは小さいモデル並み」
Apple の新しいオンデバイス基盤 Core AI(Core ML の後継、iOS 27 / macOS 27 beta)では、 MoE デコードが**「使わないエキスパートまで毎トークン全部読む」**せいで遅い。MoE の旨味（オンデバイスでこそ活かしたい）が消えている。
カスタム Metal カーネル(gather_qmm)を書いてMoe構造で動かした。
ルートされたエキスパートだけ読むようにしたら 2.1〜3.6× 高速化、品質は不変。

【図: 緑がMoeでエキスパートのみ動かした速度】

(M4 Max、release llm-benchmark、p=128 g=256)

1. MoE の何がそんなに嬉しいのか

ローカル LLM をやっていると、いつも 「賢さ ↔ 速さ・メモリ」 のトレードオフにぶつかります。賢いモデルは大きく、大きいモデルは遅く、メモリも食う。

MoE(Mixture of Experts)はこのトレードオフをズラすアーキです。
モデルのネットワークを「たくさんの小さなエキスパート」に分割し、トークンごとにルーターが上位 k 個だけを選んで動かす。

例:Qwen3.6-35B-A3B は

総パラメータ 35B(= 賢さの器は 35B 級)
でも1トークンで動くのは ~3B だけ(256 エキスパート中 8 個)

つまり 「35B の賢さを、3B の計算量で」。

デコードは基本メモリ帯域律速なので、「毎トークン読む重みが 3B 分で済む」なら、
35B 級の品質が小型モデル並みの速度で回る
——これがオンデバイス、特に帯域が限られる Mac/iPhone で MoE が刺さる理由です。

ポイント:MoE の旨味は「active なエキスパートの重みしか読まない」ことに全面的に依存している。

2. Apple新フレームワークだと遅い

Apple の Core AI に MoE を移植して llm-benchmark を回すと、LFM2.5-8B-A1B(active ~1.5B)が int8 で 39 tok/s しか出ない。

おかしい。

1.5B しか動いていないはずなのに、まるで 8B 全部を読んでいるような遅さ。

調べると、犯人は Core AI が MoE を落とす GatherMM というコンポジット op でした。

これは

ルートされたエキスパートを gather → dense matmul

という素直な実装なのですが、GPU 上に lower された matmul が
結局全エキスパートの重みを毎トークン読みに行く。
プロファイルすると
8.8 GB / トークンを読んで帯域飽和
していました。active は 1.5B なのに、です。

MoE の唯一にして最大の旨味(active だけ読む)が、フレームワーク側で消されていた。

これは Apple 自身の MoE モデルも同じ GatherMM を使うので移植ミスではなく、
Core AI の GPU MoE パスの成熟度の問題です。

3. 鳴かぬなら...

「殺してしまえ」

→ int4 に落として無理やり速くしてみた。

…非 QAT int4 は品質が壊れた(fp32 比 ~12 トークン flip / 41、文法が崩れる)。
MoE を殺すに等しい。却下。

「鳴くまで待とう」

→ Apple が GatherMM を最適化するのを待つ。

…いつになるか分からない。却下。

「鳴かせてみせよう」

→ 自分でカーネルを書いて鳴かせる。

Core AI には TorchMetalKernel という、生の Metal シェーダを書いて torch.export 経由で .aimodel の中にカスタム op として埋め込める公式 API があります。

これで「ルートされたエキスパートだけ読む matvec」を書きました。これが本記事の gather_qmm。

4. カスタムカーネル `gather_qmm`

肝は、ルーティングのインデックス自体をカーネルの入力にすること。

// QP は torch [E, N, K/4] (全エキスパート、packed int8)
// IDX は torch [k] (ルートされた expert id)

const uint slot = tgid.z;            // どの routed-expert スロットか
const uint e = uint(IDX[slot]);      // ← gather: このエキスパートだけ読む

...

uint packed = uint(QP[w0 + wi, n, e]);  // QP[w,n,e] = torch qp[e,n,w]

// e は実行時インデックス。他の E-k 個の重みはグローバルメモリから一切フェッチされない

e を実行時の添字としてグローバルメモリ読み出しに使うので、ルートされていないエキスパートの重みは物理的に読まれません。

これでデコードがようやく active-param 帯域で回る = MoE 本来の速度。

MetalSwitchGLU という drop-in を作って、metalize_moe(model, scheme="sym8") で各 MoE 層の switch_mlp を差し替えるだけ。

LFM2.5(feed_forward)/ Qwen・GLM(mlp)両方のレイアウトに対応。

5. 結果

冒頭の表とチャートの通り、2.1〜3.6×。over-read がひどいモデルほど効きます:

GLM は 4/64 = 16× の over-read を除去 → 2.6×
Qwen3.6 は 8/256 = 32× の over-read → 2.1×(MLA を全47層で回す重さが残るので倍率はやや控えめ)

特に Qwen3.6 では、これがちょうど MLX 4-bit との速度差の「expert-gather 分」を埋め、残差は int8-vs-int4 のバイト差だけになりました。

6. 品質検証(正直に)

gather_qmm 自体は dequant 後の参照と bit-exact。なので品質は重みの量子化方式だけで決まります。

最初に k-means int8 を試したら fp32 オラクル比で +5 flip/41 出てしまい、「fp16 と一致」は誤りでした(これは訂正)。

出荷済みと同じ対称線形 int8(sym8、per-K-block-32)を bit-exact gather で読むようにしたらクリーン:

判定は margin ルール:fp32 オラクルと top-1 が食い違っても、fp32 のロジット差が 0.1 未満なら統計的タイとしてカウントしない。

7. ハマりどころ(同じことをやる人へ)

rank-3 バッファ添字 + データ依存 gather は GPU で動く。W[e](e は実行時 expert id)を読める。先に小さな probe で確認するのが吉。
down 射影は per-slot。gate/up はトークン x を共有するが、down は各エキスパートが自分の gated activation を入力に取る。取り違えると出力が静かに壊れる(相対誤差 ~1.3)。
GPU 専用(カスタム MSL は ANE で動かない)。ただし代わりに「MoE グラフの raw GPU ロードが GatherMM→ANE コンパイルで abort する」既知問題も消える。
int8 の符号拡張は MSL の char() で。
正直な限界:これは Mac の話。同じ int4 ビルドは iPhone 17 Pro でも動いた(~32 tok/s、自分が phone で動かせた初の MoE)が、非 QAT int4 は品質が保てないので phone 版は「動くが劣化」として出している。

まとめ

MoE の旨味は「active だけ読む」。

新しいフレームワークがそれを殺していたなら、
カスタム Metal カーネル 1 本書けばいい。

3つの MoE が 2〜3.6× 速くなり、品質は据え置き。

zoo: https://github.com/john-rocky/coreai-model-zoo
カーネル解説: knowledge/compute-units-and-authoring.md(gather_qmm で検索)
HF: LFM2.5-8B-A1B-CoreAI / Qwen3.6-35B-A3B-CoreAI / GLM-4.7-Flash-CoreAI

🐣

フリーランスエンジニアです。
AIについて色々記事を書いていますのでよかったらプロフィールを見てみてください。

もし以下のようなご要望をお持ちでしたらお気軽にご相談ください。
AIサービスを開発したい、ビジネスにAIを組み込んで効率化したい、AIを使ったスマホアプリを開発したい、
ARを使ったアプリケーションを作りたい、スマホアプリを作りたいけどどこに相談したらいいかわからない…

いずれも中間コストを省いたリーズナブルな価格でお請けできます。

お仕事のご相談はこちらまで
rockyshikoku@gmail.com

機械学習やAR技術を使ったアプリケーションを作っています。
機械学習／AR関連の情報を発信しています。

X
Medium
GitHub