QwenのMoEルーターを観測してホットエキスパートを先読みしても速くならなかった

QwenのMoEルーターからエキスパートの選択ログが取れたので、これを使って推論前の事前ウォームアップ（ホットエキスパートの先読み）ができないか試してみた。ホットエキスパートは、MoEのルーターがトークンごとに256個から8個を選ぶ中で、特に頻繁に選ばれていたエキスパートを指す。

Kimi K3のレポートは「Top-16 / 896」の仕様記述止まりだったし、SwiftLMの検証に至ってはSSDストリーミング以前にOSのmmapページキャッシュだけで処理が閉じていた。

なら「ログから事前にホットエキスパートを特定してメモリに乗せる」ほうが、OS任せのmmapキャッシュより速いのか。同一条件でベンチマークを取ってみた。

検証環境

項目	内容
マシン	MacBook Pro M1 Max 64GB統合メモリ
OS	macOS 26.5 (Darwin, build 25F71)
Python	3.13 (miniconda base)
mlx	0.31.2
mlx-lm	0.31.3
対象モデル	`unsloth/Qwen3.6-35B-A3B-UD-MLX-4bit`（4bitベースの動的量子化、21.7GB・約20.2GiB）
比較対象として試したモデル	`mlx-community/Qwen3.5-122B-A10B-4bit`（4bit量子化、69.6GB・約64.8GiB）
MoEアーキ	`qwen3_5_moe`（Qwen3-Next系の`Qwen3NextSparseMoeBlock`を共有）
生成	`mlx_lm.stream_generate`、既定サンプリングパラメータ

35B-A3Bは256エキスパート中top-8ルーティング、40層全層がMoE。
122B-A10Bも同じqwen3_5_moe系だが層数・パラメータ規模が大きい。

ルーターをフックする

Qwen3NextSparseMoeBlockのルーター部分は次の3行に集約される。

gates = self.gate(x)
gates = mx.softmax(gates, axis=-1, precise=True)
inds = mx.argpartition(gates, kth=-k, axis=-1)[..., -k:]  # 選ばれたexpert index

mlx_lmのsite-packagesは書き換えず、モデルロード後にQwen3NextSparseMoeBlock.__call__をモンキーパッチする方式にした。元の実装をそのまま複製し、indsを層・トークンごとに記録する処理だけを足した。

最初はindsが出るたびにmx.eval()していたが、これは層数×トークン数ぶんの同期ポイントを作ってしまい、MLXの遅延評価によるメモリ再利用を妨げる。結局、各トークンが最終MoE層（35B-A3Bなら層39）に到達した時点で、そのトークンぶんのindsをまとめて1回だけ評価する形に変えて、同期ポイントを層数ぶん減らした。

ルーター観測・対照実験・安全ラッパーのソース一式はGitHubに置いた → LiltingChannelLabo

プロンプトはわざとジャンルをばらした5種類を用意した。

キー	内容
bst	二分探索木への挿入関数（Python）
bbs	簡易BBS実装（HTML）
kana_intro	キャラクターの自己紹介
math	フィボナッチ数列の行列累乗による高速化
cn	中国語での自己紹介

122B-A10BはMetalのメモリ不足で動かなかった

先に大きい方のモデルで試した。69.6GB（約64.8GiB）の4bit量子化モデルを64GB統合メモリの機体に載せる時点で無理があるのは分かっていたので、mlx_lm.load(model_id, lazy=True)で遅延ロードにした。

ロード自体は1〜2秒で終わった。mlx_lm.load()はデフォルトでlazy=Falseになっていて、これだと内部で全パラメータを即座に評価してしまう。ところがlazy=Trueにしても、最初のprefill（プロンプト全体をまとめて処理する段階）だけで空きメモリが数十GB単位で急減する現象が起きた。

外部からvm_statの空きメモリを基準に安全域を判定するアプローチは、しきい値を何度変えても安定しなかった。

空きメモリではなくスワップを基準にする

何度か走らせると、空きメモリは50GB近くから数十MBまで乱高下する一方、スワップ使用量は常に同じ値のままだった。

mmapされた読み取り専用のsafetensorsファイルは、OSがページキャッシュへ載せるぶんだけ空きメモリを消費するが、未変更ページは書き戻しが不要なので、必要になれば（スワップを経由せず）破棄してファイルから読み直せる。つまり空きメモリの減少だけでは、実際にメモリが逼迫しているかどうかは断定できない。MLX自身の内部計測（mx.get_active_memory()）もMLXが管理するバッファしか数えておらず、モデルロード直後はactive=0.00GBのままで、空きメモリの急減とは無関係に推移していた。

危険域の判定をスワップ使用量の増加だけを基準にするよう直したところ、以降のプロセスは外部から強制終了されず、最後まで走るようになった。

それでもMetal自身がメモリ不足を報告した

スワップ基準に直した状態で生成を走らせたところ、今度はmlx_lm側から明確なエラーが返ってきた。

libc++abi: terminating due to uncaught exception of type std::runtime_error:
[METAL] Command buffer execution failed: Insufficient Memory
(00000008:kIOGPUCommandBufferCallbackErrorOutOfMemory)

これはガードスクリプトが止めたのではなく、Metal自身がGPUコマンドバッファの実行に必要なメモリを確保できず、未捕捉のC++例外としてプロセスが終了した結果だった。システムはフリーズせず、スワップ使用量もこの時一度も変わらなかった。

mx.set_memory_limit()で明示的な上限を設定しても結果は変わらなかった。ドキュメント上、この関数は設定値を超え、かつスワップまで含めてメモリが尽きた時だけ例外を出す「目安」で、値を下げても実効的な挙動に違いは出なかった。

lazy=Trueが遅らせるのは重みを評価するタイミングで、forward時のピークメモリを下げる保証はない。実際、122B-A10Bの4bit量子化モデルは、mlx_lmの素のgenerate実装では64GB統合メモリの機体に収まらなかった。SwiftLMはこのモデルを自動的にSSD逐次読み出しの「SSD STREAMING」戦略に分類していた。

対象を、64GBに確実に収まる35B-A3Bに切り替えた。

35B-A3Bでのエキスパート観測

35B-A3Bでは問題なく動いた。ロード直後のMLX内部計測はactive=0.00GB、生成が進むとactiveは18.5GB前後で安定し、モデルサイズ相応の値に落ち着いた。

5プロンプトをmax_tokens=150で流し、層ごとに選ばれたユニークなエキスパート数（256個中）を数えた。

プロンプト	生成時間	層あたりユニークエキスパート数（256中）
bst	4.0秒	158.3
bbs	3.3秒	156.4
kana_intro	3.4秒	149.3
math	3.2秒	161.5
cn	3.2秒	147.9

どのプロンプトでも、層あたり6割前後（147.9〜161.5個）のエキスパートが使われていた。1トークンで選ばれるのはtop-8だけでも、150トークンぶん重ねると選択は半数以上のエキスパートに広がっていた。「ホットエキスパートの上位だけ先読みすれば足りる」という前提は、150トークン程度の生成ではあまり成り立っていなかった。

事前ウォームアップとmmapキャッシュを比べる

SwitchLinear/QuantizedSwitchLinear（mlx_lm/models/switch_layers.py）の実装を確認すると、各MoE層のエキスパート重みは(num_experts, output_dims, input_dims)の単一配列として格納されている。量子化モデルではQuantizedSwitchLinearが使われ、gather_qmmがrhs_indicesでエキスパートのインデックスを指定しながら行列積する。

x = mx.gather_qmm(x, self["weight"], self["scales"], ...,
                   rhs_indices=indices, transpose=True, ...)

これを利用して、観測で得たホットエキスパートを指定したgather_qmmをダミー入力で先に1回実行する処理をウォームアップとして実装した。

同一プロセス・同一プロンプトで2条件を比較した。条件Aは素のstream_generateでウォームアップなし、条件Bは観測ログの上位20エキスパート（層ごと）をこのウォームアップで先読みしてから、同じプロンプトを生成した。

プロンプト	条件A（素の生成）	条件B（上位20エキスパートを事前ウォームアップ）	差
bst	61.83 tok/s	61.80 tok/s	-0.05%
kana_intro	62.02 tok/s	61.92 tok/s	-0.16%

ウォームアップ自体は0.02秒で終わったが、生成速度のtok/sで比べるとA/Bの差は誤差の範囲に収まった。2プロンプトとも同じパターンで、条件Bがわずかに低いくらいで有意な改善はなかった。

条件Aと条件Bを同一プロセス内でA→Bの順に流しているため、初回prefillのコールドスタート分だけ後から走るBが有利になる偏りはある。ただし比較している生成速度は、トークンを1個ずつ出力するデコード側の指標なので、prefillの差はこの数字にほとんど影響しない。

35B-A3Bの4bit量子化は21.7GBで、素のmmapキャッシュの時点で61.8〜62.0 tok/sが出ていた。どのエキスパートが選ばれるかを事前に知っていても、ウォームアップで削れるコールドフォルトのコストは誤差の範囲しか残っていなかった。モデルがメモリに収まらずSSD I/Oがボトルネックになる規模、つまり今回動かなかった122B-A10Bでどうなるかは、mlx_lmの素のgenerateでは生成自体が動かず、検証できていない。