125B MoE を 32GB 2 枚で運用する vLLM 最適化編

前回の記事では、Qwen3.8-Flash-Next (125B-A6B MoE + 51B の n-gram テーブル) を RTX 3090 3 枚の vLLM で 1 並列 95〜99 tok/s まで高速化した。その後、公開している量子化モデルについて、RTX 5090 を 2 枚持つ利用者から「32 GB 2 枚で動かせるか。動かす場合、expert をどれだけ RAM に置くことになり、decode はどれだけ遅くなるか」という趣旨の質問を受けた。

本記事は、その回答として作成した配布構成 2x32GB の内容をまとめる。32 GB 2 枚には重みが収まらないため、使用頻度の低い routed expert をホスト RAM に置き、選ばれたときだけ PCIe 経由で読む。

筆者は 32 GB のカードを持っていない。本記事の数値はすべて RTX 3090 での計測である。 計測の条件は次のとおり。

• 主な計測は 3090 3 枚 (PP=3) で、2 枚構成用に選んだ expert (1 枚あたり 3 GiB、計 6 GiB) をホストに置いて行った。1 並列の decode では PP の各段が順に実行されるため、1 トークンの時間は全段の合計になる。どの expert がホストにあるかは層の分け方に依存しないので、1 トークンあたりのホスト読み出しの回数と時間は 2 枚構成と同じになる。
• 3090 は 3 枚とも PCIe 4.0 x16 で接続している。帯域は 5090 を x8 のスロットに挿した場合の PCIe 5.0 x8 と同じである。
• 2 枚での動作は 3090 2 枚 (PP=2) でも確認した。ただし 24 GB に収めるため、ホストに置く量は 1 枚あたり 10.5 GiB になる (6 節)。
• 5090 自体の演算性能・メモリ帯域、PP が 2 段になることの影響、32 GB に実際に収まるかどうかは計測していない (8 節)。

先に結果を示す。

条件はすべて 3090 3 枚 (PP=3)、--max-num-batched-tokens 2048、--kv-cache-memory-bytes 550000000、--max-num-seqs 2、prefix caching 有効である。初版の列は、現在のイメージで VLLM_EXPERT_ONE_TABLE=0 として offload-02 の経路で動かした値である (5 節の修正は適用済み)。decode はどれも temperature 0・ignore_eos・thinking 無効・出力 512 トークンで、「新規プロンプト 16 本」は分野 (コード・日本語・英語の物語・数学・その他) を交互に替えたプロンプトを起動直後のサーバに 1 回ずつ流した平均、「コード / 英語の物語 / 日本語」は同じプロンプトを 2 回続けて流した値である。

前回の配布物の 12 本に、以下の 3 本を追加した。

以下、作業の順に記述する。

1. 何が収まらないか

前回の mem-01 適用後の重みは 62.27 GiB である。PP=2 で 24 層ずつに分けると、1 枚あたり 31.36 / 31.21 GiB になる。

重み以外に必要な VRAM は、3090 での実測で 1 枚あたり約 1.6 GiB (KV の 0.34 GiB を含む) である。prefill の chunk を 2048 にすると、さらに 0.4〜0.7 GiB 増える (3.5 節)。5090 の VRAM は 32,607 MiB で、使える量を 1 枚あたり約 31.4 GiB と見積もると、1 枚あたり最低約 2.2 GiB を VRAM の外に出す必要がある。

出す対象は routed expert にした。

• routed expert は 48 層 x 512 個で 58.01 GiB あり、重みの 93% を占める。
• 1 個は INT4 g128 で 2.417 MiB (w13 [1280, 2560] + w2 [2560, 640] + BF16 の scale)。
• 1 トークンで選ばれる expert は 48 層 x top-10 = 480 個で、全体の 2% である。

選ばれない expert はそのトークンの計算に使われない。使用頻度の低い expert をホストに置けば、PCIe の読み出しは選ばれたときだけで済む。

KV キャッシュは前回までと同じく QSA の K/V をホストに置く。PP=2 では rank あたりの QSA 層が 6 層になり、attention の block size は 8,096 トークンになる (3 枚構成では 12,144)。配布設定の --kv-cache-memory-bytes 360000000 で 112 block、262,144 x 1.04 本、ホストの pinned メモリは 20.8 GiB である。

2. 層単位でホストに置く (offload-01)

最初に、指定した層の routed expert をすべてホストに置く方式を実装した。VLLM_EXPERT_HOST_LAYERS に層番号を指定する。

• 構築時に expert の重みをホストの pinned メモリに確保し、GPU からは UVA の view で読む。MoE のカーネルは変更せず、選ばれた expert の分だけが PCIe で転送される。
• ロード後の Marlin の repack は新しい GPU テンソルを返すため、repack の直後にホストへ戻す。GPU に同時に存在するのは 1 層分だけで、重みが VRAM に収まらないモデルでもロードできる。

1 層あたり VRAM 1.2 GiB を空けられるが、decode では 1 層あたり 1 トークン約 24 MB (top-10 の expert) を PCIe で読むため、約 1 ms/token かかる。prefill では chunk ごとにほぼ層全体 (1.2 GiB) を読む。

6 層で decode が約 4 割低下した。3090 2 枚で 18 層をホストに置いた場合は 35 tok/s だった。この方式は配布物に残しているが、既定では使わない。

3. 使用頻度の低い expert だけをホストに置く (offload-02)

層ではなく expert 単位で選べば、選ばれにくい expert だけをホストに置ける。ホストに置く expert の集合を、以下 cold set と呼ぶ (配布物の VLLM_EXPERT_COLD_FILE / cold-experts/)。

3.1 expert の選択の偏り

gate の top-10 を層ごとに数える計数用のフックを vLLM に入れ、以下のデータを流した。ある位置で選ばれる expert は、その位置までのコンテキストだけで決まる。そのため生成済みのトークン列を prefill で流せば、decode のときと同じ選択が得られる。

3.5 GiB (1,482 個、全体の 6.0%) をホストに置く場合に、評価用の集合で選択のうちホスト側の expert が占める割合 h を求めた。

1 種類のデータだけで選ぶと、別の種類の作業にはほとんど効果が無い。競技プログラミングの Python と思考文、実リポジトリのコードでは、よく使われる expert が異なる。2 種類を混ぜると、どちらでも使用頻度の低い expert だけが残り、選択に使っていない G2 でも h は一様にランダムな場合の約 1/11 になった。

偏り自体は大きい。LCB-B では、各層の下位 10% の expert が受ける選択は全体の 0.29% (層の中央値) だった。

3.2 ホスト読み出し 1 回のコスト

3090 3 枚の本番構成で、fused decode の Marlin MoE を GPU 側とホスト側の 2 回に分けて呼ぶ実験用の実装を作り、1 回あたりのコストを計測した。重みは GPU にすべて残したまま、ホスト側の expert は pinned メモリの UVA view から読むようにした。条件は LiveCodeBench の 5 問 (チャット) と 30k のコードの続き 3 本、temperature 0・ignore_eos・出力 768 トークン。

対照との ITL の差を読み出し回数で割ると、ホスト読み出し 1 回あたり 111〜136 µs (実効 19〜23 GB/s) になる。一様な cold set では読み出しが各トークンに均等に発生し、p90 は p50 + 0.8 ms 程度に収まる。選んだ cold set では読み出しが少数のトークンに集中し、p90 が最大 15.8 ms まで伸びた。

2 回に分けて呼ぶだけで +0.8 ms かかっている。この実験では並べ替えと Marlin の呼び出しを 48 層それぞれで 2 回ずつ行っていたためで、4 節でこの分を削減する。

3.3 分野を混ぜて選ぶ

3.1 節の cold set はコードのデータだけで選んでいた。これを使うと、英語の物語の decode が大きく低下した。

そこで、コード以外のデータも数えて cold set を選び直した。追加したのは、日本語 (量子化のキャリブレーションデータ 160 本、62 万トークン)、arXiv の論文と科学系の文書 (各 32 本、26 万トークン)、英語と日本語の雑多なチャットに対する生成 C1 (16 本、decode 1.8 万トークン) である。LCB-A・code・G1・日本語・arXiv・科学・C1 の 7 集合をそれぞれ合計 1 に正規化して足し、その値の小さい順に選ぶ。PP=2 の 1 枚目 (層 0〜23) と 2 枚目 (層 24〜47) で、それぞれの予算まで選ぶ。

評価には選択に使っていない G2 (コーディングの decode) と C2 (英語と日本語の雑多なチャットの decode) を使った。1 トークンあたりのホスト読み出し回数は以下のとおり。

コードのみで選んだ cold set は、C2 では一様にランダムな場合より読み出しが多い。コードで使われない expert には、チャットでよく使われる expert が含まれている。汎用の cold set は、コードではコード専用のものより読み出しが多いが、どちらの評価でも一様にランダムな場合の約 1/4〜1/10 になる。配布物にはこの汎用の cold set を、1 枚あたり 1.5 / 2 / 2.5 / 3 / 4 GiB の 5 種類入れた。

3.4 実装

• GPU 側の RoutedExperts は行数を減らし、expert map でホスト側の expert を -1 にする (EP と同じ形)。
• ホスト側用に同じクラスの 2 つ目のモジュールを作る。model の子モジュールにはしないため、パラメータ名とロード済みの集計は変わらない。loader は weight_loader をラップし、ホスト側の expert id を 2 つ目のモジュールへ渡す。
• decode の fused 経路では、router のカーネルが GPU 側とホスト側の 2 本のブロック表を書き、Marlin を 2 回呼ぶ。

注意点として、Marlin MoE は expert id が -1 のブロックを飛ばさず、範囲外を読む。 単体テストで illegal address を再現した。最初の実装は -1 を含む 1 本のブロック表で Marlin を 2 回呼んでいたため、ここで停止した。ブロック表には、そのモジュールが持たない expert を入れてはいけない。

出力は、同じ生成結果を採点させたときの logprob の差 (生成部分) が 0.005〜0.024 で、退避なしのサーバ同士の差 (0.004〜0.021) と同程度だった。

3.5 prefill と chunk サイズ

prefill では、chunk ごとにその chunk が使うホスト側の expert を PCIe で読む。chunk を大きくすると、プロンプト 1 トークンあたりの読み出し量が減る。

chunk 2048 で、退避なし (chunk 512) と同程度になった。代わりに VRAM が 1 枚あたり 0.4〜0.7 GiB 増える。1 節の見積もりはこの増加分を含む。配布設定は chunk 2048、1 枚あたり 3 GiB にした。

3.6 offload-02 の結果

3090 3 枚、汎用 3 GiB/枚、chunk 2048 で、コード / 英語の物語 / 日本語の decode は 80〜83 / 83〜84 / 78〜80 tok/s だった。これで最初の配布物とした。同じ条件の新規プロンプト 16 本では、退避なしの 97.7 tok/s に対して 81.9 tok/s で、約 16% 低い。4 節でこれを削減する。

4. decode 高速化: 1 本のテーブルと実行時の入れ替え (offload-03)

4.1 offload-02 で遅かった箇所

新規プロンプト 16 本の計測で、退避なしが 97.7 tok/s、offload-02 が 81.9 tok/s だった。差は 1 トークンあたり約 2 ms で、以下の 2 つに分けられる。

• 呼び出しの追加。 offload-02 は層ごとに Marlin の w13、activation、w2 をホスト側用にもう 1 組実行する。ホスト読み出しが無いステップでも、1 トークンあたり 144 カーネルが追加される。
• ホスト読み出し。 汎用の cold set は分野を混ぜて選んだため、どの会話でも一定の頻度で読み出しが発生する。

4.2 VRAM とホスト RAM にまたがる 1 本のテーブル

1 本の expert テンソルの前半の行を VRAM、後半の行をホストの pinned RAM に置けば、Marlin は 1 回で済む。CUDA の仮想メモリ管理 (VMM) API を ctypes で libcuda から直接呼んで実装した。

1. cuMemAddressReserve で連続した仮想アドレス範囲を確保する。
2. cuMemCreate で、デバイス上の物理メモリと、CU_MEM_LOCATION_TYPE_HOST_NUMA のホスト上の物理メモリを作る。
3. 両者を隣り合わせに cuMemMap する。

GPU からホスト側を読む帯域は 25.1 GB/s で、従来の pinned メモリの UVA と同じだった。

• GPU 側の expert を前の行、ホスト側の expert を後ろの行に並べ、router の gate の行も同じ順に並べ替える。router が出力する id がそのままテーブルの行番号になるため、vLLM の MoE 経路は変更不要で、expert map も不要、Marlin は 1 回になる。
• ロードは offload-02 の処理 (2 つのモジュール) をそのまま使い、repack の後に 1 本にまとめる。moe_quant_config が scale への参照を持っているため、_setup_kernel を呼び直す。
• VRAM 側は 2 MiB 単位で切り捨てるため、VRAM の使用量は offload-02 から増えない (nvidia-smi で +40 MiB 以内、KV 容量は同じ)。
• VLLM_EXPERT_ONE_TABLE=0 を指定した場合と、ドライバがホスト側のマップを拒否した場合 (起動時に 4 MiB で試す) は、offload-02 の経路で動作する。
• gate の行を並べ替えるため、BF16 の logit が同値の場合に選ばれる expert が元と変わる可能性がある (同値の場合は id の小さい方が選ばれる)。この影響は vLLM の実行ごとの非決定性より小さい。

注意点が 2 つある。vLLM の起動ログの Model loading took は VMM で確保したメモリを数えないため、実際より小さい値が出る。また、VMM で確保したホスト側のメモリは CPU の仮想アドレスにマップしないため、プロセスの RssShmem には現れない (free の used には現れる)。KV 容量の自動決定は空き VRAM の実測で行われるため、前者の影響は受けない。

この段階で 81.9 → 86.6 tok/s になった。約 0.66 ms/token が呼び出しの追加分だった。

4.3 実行時の入れ替え: シミュレーション

残りはホスト読み出しである。汎用の cold set は、どの分野でも一定の頻度で読み出しが発生するように選ばれている。一方、1 つの会話の中では使われる expert に偏りがある。会話で使われている expert を実行時に GPU 側へ移し、使われていない expert をホスト側へ移せば、読み出しを減らせる。

実装の前にシミュレーションで効果を確認した。26 セッション (コード、既存コードの解説、日本語、英語の物語、数学、その他、各約 2k トークン) の生成結果と LiveCodeBench の実際の生成 8 本、計 13.7 万トークンを prefill で流し、層ごとの top-10 をトークン順に記録した。

• 固定の汎用 cold set (3 GiB/枚): 1 トークンあたりのホスト読み出し 8.08 回 (日本語 13.9、その他 17.1、コード 10.5)。
• 入れ替えあり (8 ステップごと、使用回数を半減期 512 ステップで減衰、差が 2 回を超えたら入れ替え、1 層 1 組、prefill は数えない): セッションの順序どおりで読み出し 1.71 回 + 入れ替え 0.36 回、セッションの順序をシャッフルして 2.12 回 + 0.45 回。
• prefill のトークンも数えた場合は 2.20 → 2.18 回で差が無かった。そのため decode の router カーネルだけで数える。
• 閾値・半減期・間隔を変えても、入れ替え 1 回を読み出し 2 回と換算した合計は 2.9〜3.7 回相当で、変化は小さかった。上記の既定値は最小値の付近にある。

4.4 実行時の入れ替え: 実装

• decode の fused 経路の router カーネル (前回の decode-05) が、選んだ行ごとに atomic_add で使用回数を数える。CUDA graph の中で実行される。
• model_state.prepare_inputs から毎ステップ expert_table_step() を呼び、8 ステップごとに 2 本のカーネルを forward と同じ stream に積む。判断のカーネルは、層ごとに使用回数を減衰させ、ホスト側で最も使われている expert と GPU 側で最も使われていない expert を選ぶ。交換のカーネルは、4 本のテンソルと gate の行を、ポインタ表を使って全層まとめて交換する。
• ホストへの読み戻しは無い。prepare_inputs は CUDA graph の capture 中には呼ばれない (capture 時は prepare_dummy_inputs が呼ばれる)。

正しさの確認として、起動直後に行ごとの指紋 (5 本のテンソルの和) を取り、数千行を入れ替えた後も全 48 層で行の組が保たれていることを確認した (この検査は配布物には含めていない)。判断のカーネルだけを実行して交換しない構成は 86.6 tok/s で、1 本化のみの構成と同じだった。判断のカーネルのコストは計測できない程度である。

4.5 結果

3090 3 枚、汎用 3 GiB/枚、chunk 2048、新規プロンプト 16 本の decode:

回数は 3 rank の合計で、router カーネルのカウンタから取った。退避なしとの ITL の差から逆算すると、ホスト読み出し 1 回は約 0.108 ms (1 本のテーブルのみの行で 1.31 ms / 12.1 回)、入れ替え 1 回は約 0.2 ms になる。読み出し 1 回のコストは 3.2 節の実験 (111〜136 µs) と一致する。

入れ替えの効果は、同じ会話が続くほど大きくなる。開発中のビルド (3090 3 枚) で、コード / 英語の物語 / 日本語のプロンプトをそれぞれ 2 回ずつ流し、これを 2 巡繰り返した。

各セルはコード、英語の物語、日本語の順である。入れ替えの無い offload-02 では 1 回目と 2 回目の差が 1 tok/s 未満だが、入れ替えありでは 2 回目以降が 94〜98 tok/s になり、退避なしとほぼ同じになった。

配布イメージでの最終確認 (3090 3 枚、配布設定の cold set と chunk):

logprob の差は 0.006〜0.024 で、退避なしのサーバ同士の差 (0.004〜0.021) と同程度だった。90k の needle は 1 並列・2 並列とも正答し、画像入力も正常に動作した。

5. offload-02 のホスト RAM の二重確保

3090 2 枚 (1 枚あたり 10.5 GiB をホストに配置) で起動すると、起動に 27 分かかり、worker のホストメモリ (RssAnon + RssShmem) が 48.8 / 42.4 GiB あった。

pinned メモリの allocator の統計を層ごとに出力して調べたところ、ホスト側の expert の、ロード時の形式のコピーが repack 後も解放されていなかった。 vLLM は構築時に record_metadata_for_reloading → capture_layer_to_meta でパラメータの属性をコピーして保持する。その中に offload-02 の weight_loader (クロージャ) が含まれ、このクロージャが引数 _p=cold_p でロード時の形式のテンソルへの参照を持っていた。1 本化の有無に関係なく、ホスト側の expert は RAM を 2 倍使っていた (配布設定で +6 GiB、10.5 GiB/枚で +21 GiB)。

修正として、クロージャはテンソルを持たず、名前で引く (getattr(cold, _n)) ようにした。さらに層ごとに torch._C._host_emptyCache() を呼ぶ。offload-02 のパッチ自体を修正し、offload-03 はその上に作り直した。

6. 3090 2 枚での計測 (PP=2)

2 枚構成の動作確認として、3090 2 枚で PP=2 を動かした。24 GB に収めるため、ホストに置く量は汎用の cold set で 1 枚あたり 10.5 GiB、chunk は 512 である。

offload-02 の列は --kv-cache-memory-bytes 360000000、offload-03 の列は 550000000 で計測した。コード / 英語の物語 / 日本語の offload-03 の値は、1 回目と 2 回目の差が大きい。4.5 節のとおり、入れ替えが進むと速くなるためである。

90k の needle は 1 並列・2 並列とも正答した。ただし 1 並列の 89,716 トークンのプロンプトは、応答まで 92.3 s かかった。3090 3 枚・3 GiB/枚・chunk 2048 では同じプロンプトが 16.5 s である。ホストに置く量が多く chunk が小さいため、prefill のホスト読み出しが多い。32 GB のカードでは 1 枚あたり 3 GiB・chunk 2048 で動作する見積もりなので (1 節)、この構成は 5090 の性能の目安ではなく、2 枚構成で正しく動作することの確認である。

VLLM_EXPERT_ONE_TABLE=0 (offload-02 の経路) は起動に 1,831 s かかり、計測スクリプトの待ち時間を超えたため数値は無い。

7. 試して採用しなかった案

7.1 ホスト側の expert を CPU で計算する

ホスト側の expert を PCIe で読む代わりに、CPU で計算する案を検討した。

1 トークンでは CPU の方がわずかに速いが、以下の理由で採用しなかった。

• GPU → CPU → GPU の同期の往復に 32 µs かかる。decode は FULL cudagraph で実行しているため、各層に host node を置いて CPU を待つことになり、ホスト側の expert が選ばれなかった層も含めて 48 層分の往復が毎トークン加わる。これは汎用 3 GiB/枚でのホスト読み出しのコスト (7〜13 回 x 0.12 ms) と同程度である。
• prefill では、PCIe は 1 chunk につき 1 回読むだけだが、CPU はトークン数に比例して遅くなる。
• 5090 は PCIe 5.0 のため、PCIe 側がさらに有利になる。CPU の DRAM 帯域 (実測 79 GB/s) も PCIe の 3〜4 倍にとどまり、その帯域を計算にも使う。

7.2 入れ替えを copy engine で非同期に行う

入れ替えを別 stream の copy engine で非同期に行えば、入れ替えの時間を隠せる。ただし予備の行と expert map 経由の router が必要になり、prefill の経路に expert map を持ち込むことになる。得られるのは入れ替え 0.9 回 x 0.2 ms ≈ 0.18 ms/token (約 1.5 tok/s) なので、実装しなかった。

ホスト読み出しを GPU 側の計算とオーバーラップさせる効果は、1 本化により同じカーネルの中で既に得られている。読み出し 1 回のコストは PCIe の帯域で決まるため、約 0.11 ms のままである。

8. 配布設定と未確認の点

配布物の vllm-patch/2x32GB/ の主な設定:

--pipeline-parallel-size 2 --tensor-parallel-size 1
--max-model-len 262144 --max-num-seqs 2
--max-num-batched-tokens 2048
--kv-cache-memory-bytes 360000000        # 262,144 x 1.04、QSA の pinned 20.8 GiB
--enable-prefix-caching --mamba-cache-mode align

VLLM_EXPERT_COLD_FILE=/cold/3GiB-per-card.json   # ホストに 1 枚あたり 3 GiB、計 6 GiB

起動時に GPU メモリが不足する場合は、1 段大きい cold set (4 GiB) を使う。起動後に 1 枚あたり約 1.5 GiB 以上空いている場合は、小さい cold set を使えば decode が速くなる。

32 GB のカードで確認していない点:

• 32 GB に実際に収まるか。1 節の見積もりでは 3 GiB/枚で約 0.8 GiB の余裕がある。
• 5090 での退避なしの decode 速度。メモリ帯域は 3090 の約 1.9 倍で、PP は 3 段から 2 段になる。
• ホスト読み出し 1 回のコスト。PCIe 5.0 x8 では 3090 と同じ約 0.11 ms、PCIe 5.0 x16 ではその約半分と見積もっている。
• 前回までの Triton のパッチが sm_120 で動作するか。

まとめ

• 重み 62.3 GiB は 32 GB 2 枚に収まらない。重みの 93% を占める routed expert のうち、使用頻度の低いものを 1 枚あたり 3 GiB ホスト RAM に置き、選ばれたときだけ PCIe で読む。
• ホストに置く expert は、1 種類のデータだけで選ぶと別の分野で読み出しが増える (コードのみで選ぶと英語の物語の decode が 95〜96 → 57〜61 tok/s)。7 種類のデータを正規化して混ぜた汎用の cold set を配布した。
• ホスト読み出しは 1 回あたり約 0.11 ms (PCIe 4.0 x16)。層単位でホストに置く方式は 6 層で decode が約 4 割低下するため、既定では使わない。
• offload-03: CUDA VMM で VRAM とホスト RAM にまたがる 1 本の expert テーブルを作り、Marlin を 1 回にした (81.9 → 86.6 tok/s)。router が数えた使用回数で実行時に expert を入れ替え、ホスト読み出しを 12.1 → 4.3 回/token に減らした (→ 90.0〜91.8 tok/s)。
• offload-02 の初版は、vLLM がコピーして保持する weight_loader のクロージャ経由で、ホスト側の expert を RAM に 2 重に保持していた。修正後、2 枚構成の起動は約 27 分 → 275 s になった。
• すべて RTX 3090 での計測で、32 GB のカードでは未検証である。3090 3 枚で同じ cold set を使った decode は、退避なしの 97.7 tok/s に対して 91.5 tok/s だった。

パッチと compose ファイルは以下の vllm-patch/2x32GB/ に置いている。対象は vllm/vllm-openai の 0.29.1rc1.dev47+gdc36fcce9。 https://huggingface.co/Minachist/Qwen3.8-Flash-Next-INT4-Mixed-AutoRound