125B MoE を 3090 3 枚で運用する vLLM 最適化編

前回の記事では、Qwen3.8-Flash-Next (125B-A6B MoE + 51B の n-gram テーブル) を RTX 3090 3 枚の vLLM で、コンテキスト 262,144・1 並列 80 tok/s で動作させた。本記事はその続きで、サービスとして運用を始めてから現在のモデルカードの数値に至るまでの 10 日間の作業をまとめる。

先に結果を示す。

※ 80.0 は prefix caching を無効にした構成の値。後述するが、prefix caching を有効にした本番構成では 59〜73 tok/s だった。

3x3090+MTP 列は配布設定 (--kv-cache-memory-bytes 550000000、262,144 x 2) での計測値である。2 並列は --max-num-seqs 2 (配布設定)、それ以外は変更前の --max-num-seqs 4 で計測した。7 節の表は開発時の構成 (--kv-cache-memory-bytes 783000000、2.85x) での別の計測回で、値が少し異なる。

現在の配布物には、前回の 3 本 (vllm.patch、decode-01、decode-02) に加えて 11 本のパッチがある。

このほかに、パッチではない設定変更が 2 つ、実装したが採用しなかった案が 3 つある。以下、時系列順に記述する。

1. 初回の TTFT が 100 秒かかる

前回の構成を本番に入れ、qwen-code から接続した。初回リクエストで最初のトークンが返るまでに 1 分以上かかった。

サーバログの出力は以下のとおり。

Avg prompt throughput: 4206.0 tokens/s

この値だけ見ると prefill は速い。しかしログのタイムスタンプを確認すると、prefill 開始 (08:43:44) から最初のトークン (08:45:23) まで、42k トークンで 100 秒 かかっていた。

Avg prompt throughput は TTFT を表さない。 vLLM v1 はプロンプトのトークン数を、最初のトークンが出力された 10 秒間の集計 window にまとめて加算する。したがってこの値は「プロンプト長 ÷ 10 秒」である。TTFT はクライアント側で計測する必要がある。

また、同じ会話の 2 回目以降 (55k〜73k トークン) は 10〜25 秒だった。初回のみが極端に遅い。

原因: PLE テーブルのページキャッシュミス

前回記述したとおり、95.37 GiB の PLE n-gram テーブルは NVMe 上のファイルを mmap し、np.take で行を gather している。1 トークンあたり 16 行 (bigram と trigram の 2 種類 x 8 head)、1 行 320 B で、各行は別々のページにある。

np.take は 1 スレッドで行を順に読むため、ページキャッシュに無いページは 1 ページずつ、queue depth 1 の同期読み出し になる。readahead は前回 MADV_RANDOM で無効化している (無効化しないと NVMe の読み出し量が 23 倍になる)。

コールドの場合、512 トークンに 1.9 秒かかり、prefill は 270 tok/s まで低下する。42k トークンのすべてのページがキャッシュミスになる場合、82 chunk x 1.9 s ≈ 156 秒となる。これは上限で、実測の 100 秒はこれより短い。テーブルの一部 (下記の 2.5 GiB) はキャッシュ済みで、プロンプト内で繰り返し出現する n-gram は同じ行を読むため、NVMe から実際に読むページはこの見積もりより少ない。差の内訳は計測していない。

ページキャッシュに載らない理由。 RAM は 125 GiB だが、QSA のホスト K/V プールが 12 layer x 3.96 GiB = 47.5 GiB を pinned (ページ固定、evict されない) で確保している。95.4 GiB のテーブルは RAM に収まらない。fincore で確認すると、キャッシュされていたのは 95.4 GiB 中 2.5 GiB だった。新しいテキストのたびに NVMe から読み出すことになる。

オフロードを使わない 122k の構成ではテーブルが RAM に収まる余地があったため、262k 構成に移行して初めて問題が表面化した。

修正: リクエスト受付時にプロンプト全体を prefetch する

テーブルのどの行を読むかは、トークン列のみで決まる。 リクエスト受付時点でプロンプト全体は確定しているため、prefill がその位置に到達する前にページを要求できる。

1. Qwen4ExpModelState.add_request にフックを追加し、新規リクエストのプロンプトのトークン列を prefetch スレッドに渡す。
2. prefetch スレッドは、GPU の Triton カーネル _ple_ngram_ids_kernel と同じハッシュを numpy で計算し、[トークン数, 16] の行番号を求める。int64 の乗算のオーバーフローは wrap、剰余は torch.remainder と同じ符号、直前の EOS より前のコンテキストは EOS として扱う、という点までカーネルと揃え、テストでビット一致を確認した。値が誤っていても誤ったページを prefetch するだけで、forward はこの値を使わない。
3. 1024 トークンごとにユニークなページを求め、隣接ページを範囲にまとめて process_madvise(pidfd, iovec[≤1024], MADV_WILLNEED) を発行する。カーネルが非同期に読み出しを発行するため、ドライブの queue depth が上がる。
4. gather 側 (_lookup) は 変更していない。gather 時点でページはキャッシュ済みか、読み出し中である。
5. リクエストの終了時または preempt 時に、残りの prefetch を破棄する。

注意点として、podman のデフォルトの seccomp プロファイルでは、process_madvise は CAP_SYS_PTRACE が無いと EPERM を返す。compose に cap_add: [SYS_PTRACE] を追加した。CAP_SYS_PTRACE を付与すると、コンテナ内の他のプロセスへの ptrace も可能になる。付与しない場合もサーバは起動し、起動ログに PLE page prefetch uses per-range madvise と出力して範囲ごとの madvise にフォールバックする。NVMe のスループットは同じで、prefetch スレッドの CPU 使用量が約 20 倍になる。prefetch 自体は VLLM_PLE_PREFETCH=0 で無効にできる。

同じマシンで、それぞれ未読の約 37k トークンのプロンプト (vLLM の Python ソース) を使って計測した。テーブルは 3.7 節で述べる P310 への移動前のドライブにあり、冒頭の 100 秒と同じ条件である。

ウォーム時の値が変わらないのは、キャッシュ済みページへの madvise が 1 ページあたり約 1 µs で済むためである。ウォーム時の処理経路には何も追加していない。

すべてのページがキャッシュミスになる場合 (1 トークンあたり 16 ページ)、prefill の上限はドライブのランダム読み出し性能 (4 万ページ/s 弱) で決まり、約 2,400 tok/s である。これ以上速くするにはテーブルを RAM に置く必要がある。上の表の prefetch 有りの 4,430 tok/s はこの上限を超えている。計測に使った Python ソースは共通の n-gram が多く、一部のページは計測前からキャッシュ済みだったため、NVMe から読んだページは 1 トークンあたり 16 ページより少ない。prefetch 無しの 44 秒が冒頭の 100 秒 (42k) より短いのも同じ理由である。

コールドの計測には未読のテキストが必要。 同じプロンプトを 2 回送ると、2 回目はページキャッシュも Triton のキャッシュもウォームになっている。A/B の各側で別々の未読ファイルを使う必要がある。

2. prefix caching を有効にすると KV 容量が 3 割減る

前回の構成は --no-enable-prefix-caching だった。qwen-code のようなエージェントは、毎ターン会話全体を送信する。prefix caching が無効だと、70k の会話を毎ターン最初から prefill することになる。

有効にしたところ、容量は以下のようになった。

prefix caching OFF:  GPU KV cache size: 1,067,300 tokens  (4.07x)
prefix caching ON:   GPU KV cache size:   734,862 tokens  (2.80x)

このとき GPU には約 500 MB の空きがあった。当初は VRAM の割り当ての問題と考えたが、空き VRAM は原因ではなかった。

容量を決めるのは block 数で、block のコストはホスト RAM

前回の 5.2 節で、QSA をオフロードする場合、attention の block size を mamba (gated delta-net の state) のページサイズに合わせて 12,144 トークンまで引き上げると述べた。このとき容量は以下の式で決まる。4 点の実測値でキャリブレーションした。

blocks    = floor(--kv-cache-memory-bytes / 3,207,168)
容量(tok) = blocks / R x max-model-len

• 3,207,168 B は gated delta-net の state 1 つ分 (ssm fp32 3,145,728 + conv bf16 61,440)。
• R は 262,144 トークンのリクエスト 1 本が使用する block 数で、prefix caching OFF で 42、ON で 61。

差の 19 は linear attention の KV cache group の数である。prefix caching には --mamba-cache-mode align が必要で、このモードは各 group の state を 1 ページ (現在の state) から 2 ページ (現在の state + checkpoint) に増やす。内訳は 19 x 2 + 22 (QSA、cdiv(262,144, 12,144)) + 1 (PLE の short-conv) = 61 となる。

つまり prefix caching は追加の VRAM を使用していない。550 MB / 171 block のまま、1 リクエストあたりの block 数が 42 から 61 に増えたため、4.07 本が 2.80 本になった。

したがって block 数を増やせばよい。block 1 個あたりのコストは以下のとおり。

VRAM 1 MiB あたり 1,405 トークン、RAM 1 GiB あたり 15,460 トークンとなる。--kv-cache-memory-bytes は、実質的にはホスト RAM の使用量を決めるパラメータである。 pinned されるホスト RAM はこの値のおよそ 93 倍になる。

pinned メモリの 2 の冪への切り上げを止める

--kv-cache-memory-bytes を 590 MB に上げたところ、ホストの RAM 使用量が 112 GiB、空きが 13 GiB になった。

原因は前回の 5.4 節で述べた PyTorch の caching host allocator で、pinned メモリの確保サイズを 2 の冪に切り上げる。1 layer 4.24 GiB のプールが 8 GiB になり、12 layer で 96 GiB になる。前回は 550 MB (1 layer 3.96 GiB) に抑えて 4 GiB 以下に収めることで回避していた。

同梱の PyTorch 2.13 のソース (CachingHostAllocator.h) によると、切り上げには閾値があり、PYTORCH_CUDA_ALLOC_CONF に pinned_max_round_threshold_mb:2048 を追加すると 2 GiB を超える確保は切り上げられない。

最終的には環境変数に依存せず、ホストプールを直接確保する実装にした。匿名 mmap で必要なサイズだけ確保し、cudaHostRegister で pin する。プールはプロセス終了まで解放しないため、allocator のキャッシュは不要である。また、pin する前に空き RAM を確認し、不足していれば必要サイズを表示して起動を中止する。PP の 3 rank が同時に確保するため、ファイルロックで直列化している。これが無いと、2 つの rank が同時にチェックを通過し、その後両方とも OOM になる可能性がある。

採用した値:

247,823 トークンのリクエスト 4 本を同時に保持した状態で、needle 4/4、preemption 0。prefix caching を有効にした状態で 262,144 x 4 本を確保できた。

3. decode 高速化 第 2 弾: 本番構成は 80 tok/s ではなかった

前回、safetensors のヘッダから decode 1 ステップの読み出し量を 5.83 GB/token と算出し、3090 の 936 GB/s から上限を 160 tok/s とした。PP=3 で 3 リクエストを同時に処理すれば、スループットの上限は 480 tok/s になる。80 tok/s はその半分であり、改善の余地があった。

3.0 計測方法

この段階の改善は 1 件あたり 0.3〜2 ms で、起動ごとのばらつきより小さい。そのため計測方法を先に決めた。

• 固定プロンプト・固定 seed で、ウォームアップ 1 回 + 計測 3 回。 ランダムなプロンプトは PLE のキャッシュミスを起こし、中央値も変動する。
• inter-token latency (ITL) の中央値と平均 tok/s の両方 を見る。平均には PLE のキャッシュミスや NVMe のストールなどの外れ値が含まれる。ユーザーの体感に近いのは平均である。
• A/B は 1 つのビルド内で環境変数により切り替える。 同じ構成でも起動し直すと数値が変わる。そのため今回のパッチはすべて環境変数 1 つで無効化できるようにした。
• torch profiler のカーネル時間は FULL cudagraph 内では実際より長く出る。 CUPTI がカーネルごとに数 µs を加算するため、カーネル数が多いほど誤差が大きい。合計 12.15 ms と出たが、実際の ITL は 10.66 ms だった。GPU 時間は graph replay の前後に CUDA event を記録して計測する。
• ホスト側のタイムラインは、.pth で全プロセスに tracer を注入して monotonic_ns を記録した。rank 間の対応付けは開始時刻ではなく完了時刻で行う。

3.1 prefix caching ON では 70 tok/s 前後、8k 以上で 60 tok/s

この方法で本番構成 (prefix caching ON) を計測し直すと、1 並列の ITL 中央値は 13.2〜13.7 ms、平均 59〜73 tok/s だった。前回の 80 tok/s は prefix caching OFF の値である。ON にすると、mamba align モードの mamba_get_block_table_tensor が KV cache group ごと (1 rank あたり 13 個) に 7 本の eager カーネルを実行する。1 ステップ・1 rank あたり約 95 本のカーネルが増える。

さらに、コンテキストが 2,048 トークンを超えるとステップ時間が 2 ms 増え、60 tok/s になる。QSA のオフロードを無効にした 32k 構成で計測すると、コンテキスト 0 / 8k / 24k のいずれも 12.5 ms で一定だった。したがってこの 2 ms は すべて QSA のホスト K/V を PCIe 経由で読む時間 である。

これは前回記事の記述の訂正になる。前回、「1 トークン 48 MiB を 80 tok/s で読んでも PCIe 帯域の 6 分の 1 であり、計算とオーバーラップできる」と書いたが、実際にはオーバーラップしていなかった。 indexer が選択した行は、その layer の indexer の実行後でなければ読み出しを開始できない。そのため他の処理とオーバーラップできない。45 tok/s の段階ではホスト側の待ち時間の方が大きく、この問題は見えていなかった。

3.2 ボトルネックはホスト側の直列処理 (decode-03)

3.3 節の hyper-connection の fusion を単独で適用すると、コンテキスト 0 で −0.4 ms、8k で −2 ms だった。GPU の処理量を減らしたのに、短いプロンプトでは効果が小さい。 短いプロンプトのステップは GPU 以外の要因で制限されており、その下限が約 12.5 ms にあった。

ホスト側のトレースで原因が判明した。PP の先頭以外の rank では、Worker.execute_model の冒頭で irecv_tensor_dict() を呼ぶ。この関数は最初に、テンソルのメタデータを gloo で受信する recv_object を実行し、前段の rank が forward を launch して isend するまでブロックする。 その後に、入力の準備と attention メタデータの構築 (このモデルでは Python で 3.2〜3.5 ms) が始まる。

旧:  rank0 [準備][launch]──→ rank1 [受信待ち][準備 3ms][launch]──→ rank2 [受信待ち][準備 3ms][launch]
新:  rank0 [準備][launch]──→ rank1 [準備 3ms][受信][launch]──→ rank2 [準備 3ms][受信][launch]
                               ↑ 前段の GPU 処理とオーバーラップする

rank あたりの GPU 時間が約 3.5 ms で、ホスト側の準備もほぼ同じ長さである。準備が前段の GPU 処理とオーバーラップせず、毎ステップ critical path に入っていた。GPU 側を高速化しても 12.5 ms で頭打ちになる理由はこれである。

修正として、受信処理を DeferredRecvIntermediateTensors というラッパーに入れ、model runner が forward の直前に .tensors を参照する時点まで遅延させた。rank 上のデバイス演算の順序 (受信 → forward → 送信) は変わらない。.tensors が参照されなかった場合も execute_model の後で必ず受信し、前段の送信と対応させる。

3.3 小さい GEMV の実効帯域が 2 割程度 (decode-04 / 05)

ホスト側のボトルネックを解消すると、GPU 時間が効いてくる。カーネルごとの実効帯域を算出すると、大きい projection (gated delta-net の qkvz 740 GB/s、QSA の qkv 750、INT8 の lm_head 900) は問題なく、小さいものだけが遅かった。

Marlin の VLLM_MARLIN_USE_ATOMIC_ADD=1 も試したが、sm8x + BF16 では Marlin 側で無視されるため効果は無かった。

hyper-connection (decode-04)

hyper-connection の GatedResidual 1 回は、inject の GEMV、down、silu、up、gate mix の 6 カーネルで約 32 µs かかる。decoder layer 1 つに 2 回あるため、1 トークンあたり 96 回実行される。

compressed-tensors の INT8 pack-quantized は、int32 1 ワードに 4 値を下位バイトから格納している。したがって weight_packed.view(uint8) がそのまま q[N, K] (値 + 128) になる。 repack も追加メモリも不要である。これを直接読む Triton カーネルを 2 本実装した。

• _hc_down_inject_silu_kernel: down (INT8) と inject (BF16) の GEMV を 1 カーネルで行い、silu を store に fuse
• _hc_up_gate_mix_kernel: up (INT8) の GEMV を行い、HC の 4 ストリームにわたる gate mix をレジスタ上で処理

BF16 への丸め位置は fuse 前と同じにしている。GatedResidual 1 回あたり 32 µs → 13.3 µs で、1 トークンあたり 96 回分、約 1.8 ms の削減になる。

4 トークンを超えるバッチ (prefill) では、同じ重みを読む tl.dot の W8A16 GEMM を使う。512 トークンでは Marlin より速い (down 102 vs 120 µs)。ただし タイルを 64x64 にすると 176 µs になり、prefill が 4% 遅くなった。 32x64 を採用した。decode 用に選んだパラメータを prefill にそのまま使うことはできない。

MoE ブロック (decode-05)

1 トークンの MoE ブロックでは、expert の GEMM よりも周辺の小さいカーネルの時間が長かった。router の GEMV (12 µs)、topk_softmax (1 CTA で 7 µs)、moe_align_block_size (2 カーネルで 7 µs)、shared expert の gate (4 カーネルで 7 µs)、moe_sum、最後の加算。expert の重みは 25 MB で、読み出しは 42 µs である。

• router と shared expert の gate を 1 本の GEMV にまとめた (512 行 + 1 行)。
• 最後に終了した CTA が top-k、renormalize、Marlin 用の block alignment を書き込む。 各 CTA は tl.debug_barrier() の後に atomic_add(sem="acq_rel") でカウンタを加算し、最後の CTA だけが後処理を行う。grid 同期用の別カーネルは不要になる。
• norm_topk_prob=True のため、全 expert にわたる softmax の分母は打ち消される。選択された 10 個だけで softmax を取れば同じ重みになる。
• top-k は、BF16 の logit を大小関係を保つ 16 bit 整数に変換し、下位に 65535 − expert id を格納した 32 bit のキーで計算する。1 回の max で値と id を同時に得られ、同値の場合は id の小さい方が選ばれる (vLLM の元のカーネルと同じ挙動)。argmax を 2 回行うより速い。
• routed expert の GEMM は vLLM の _fused_marlin_moe をそのまま呼ぶ。
• INT6 の shared expert は、compressed-tensors の packing (32 値を int32 6 ワードに隙間なく格納) のままだと GEMV での展開処理が複雑になるため、ロード時に 4 bit の bit plane と 2 bit の bit plane に並べ替えた。バイト数は変わらない。gate_up + SiluAndMul で 1 カーネル、down + 最終合算 (routed の和 + sigmoid(gate) x shared) で 1 カーネル。13.4 / 7.0 µs → 7.8 / 5.3 µs。

3.4 decode 中に生成されたトークンの PLE gather (decode-06)

1 節の prefetch はプロンプト部分にしか適用できない。decode でサンプリングされたトークンの n-gram 行は事前に分からない。ページキャッシュに無い場合、np.take が 16 行分のページを 1 ページずつ QD1 で読み、2〜4.6 ms が critical path に入る。p90 が p50 より 4〜5 ms 高い原因はこれだった。

gather の直前に、16 行分のページ (各行の先頭と末尾のページ) を 1 回の process_madvise で要求するようにした。これによりページフォルトが並列に処理される。prefetch スレッド用の汎用関数 (unique を取って範囲をまとめる処理) は 0.15 ms かかったため、iovec 配列を使い回す専用の関数を実装した。8k の p90 が 16 → 12.4〜13.0 ms、平均が 72〜74 → 80〜82 tok/s になった。

3.5 メモリ: allocated が同じまま reserved が 190 MiB 増加

最初の実装では、アイドル時に GPU1 のメモリ使用量が +190 MiB、prefill 後に 24,126 MiB (上限 24,124) になった。torch.cuda.memory_allocated とそのピークはベースラインと同一で、reserved のみが増えていた。

原因は load_model 中の確保パターンの変化だった。hyper-connection の重みを Marlin で repack しなくなった (新しく確保して古い領域を解放する処理が無くなった) ため、他の layer の repack が残す空き領域の位置が変わり、再利用できない断片が残った。この断片は expandable segments の 2 MiB 単位より小さく、torch.cuda.empty_cache() では解放されない。

対策として、確保パターンを元の実装に合わせた。 hyper-connection の重みと scale を一度 clone する (Marlin と同じく、新規確保してから旧領域を解放する)。INT6 の bit plane は 2 つを別々に確保すると 100〜190 MiB の断片が残ったため、元と同じサイズの 1 つの領域に書き込んでから元を解放する。結果として、4 x 247k 保持時の GPU ピークは [24072, 24100, 24090] → [24022, 24090, 24038] MiB と減少した。

3.6 結果

同じ計測スクリプト、本番と同じ compose 設定で計測した。

1 つのビルドで環境変数を切り替えたときの各パッチの寄与:

短いプロンプトでは、03 と 04 はそれぞれ単独では約 0.5 ms の改善だが、両方適用すると 2.5 ms 改善する。ホスト側の処理と GPU 時間が交互にボトルネックになっていたため、片方だけを改善してももう片方で制限される。

残りの 10 ms はほぼすべて GPU 時間で、rank あたり graph 約 3.2 ms と、lm_head およびサンプリングの時間である。

正しさの確認は以下のとおり。hyper-connection の fusion は 1 トークンで fp32 の参照実装と完全一致。MoE は実際の重み・実際の decode 192 回で、元の経路との最大相対差 9.7e-3 (BF16 の数 ULP) で、選択される expert の変化は無し。INT6 の bit plane はビット単位の往復テストで一致。サーバ経由の logprob の平均差は旧 vs 新が 0.001995、旧 vs 旧が 0.002010 で、run 間のばらつきと同程度である。

3.7 ドライブ側の遅延

decode 中に 100〜780 ms 停止することがあった。停止中のスレッドの wchan は folio_wait_bit_common (ページ I/O 待ち) だった。O_DIRECT で 4 KiB のランダム読み出しを継続的に計測すると、テーブルを置いていた NVMe (DRAM レスの QLC) では、数秒間すべての読み出しが約 140 ms になる期間がある (最大 970 ms)。同じ条件で別の NVMe (Crucial P310) は p50 0.106 ms、最大 22 ms だった。

ソフトウェア側では対処できないため、最終的にテーブルを P310 に移した。PLE テーブルを置くドライブは、ランダム読み出しのレイテンシのテール (p99 や最大値) で選ぶべきである。

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

実装・計測したうえで採用しなかった案を 3 つ記述する。

4.1 TP=3

前回、各次元が 3 で割り切れないため TP ではなく PP を採用したと書いた。パディングすれば TP=3 も実装可能である。実装する価値があるか判断するため、先に性能の上限を計測した。

重みは --load-format dummy とし、config を 3 で割り切れる形 (attention head 24 → 36、KV head 2 → 3、linear attention の head 16 → 18、shared expert 640 → 768、vocab を 192 の倍数) に書き換えて、パディング後と同じ計算量にした。routed expert は EP で 171 / 171 / 170 に分割した。

比較対象の PP=3 も同じ dummy 条件で計測した。両者とも、パッチは配布物の 4 本 (vllm.patch、decode-01、decode-02、ttft-01) のみで、3 節の decode-03 以降は含まない。QSA のオフロードは TP=1 専用のため無効にし、KV は VRAM に置いた (max-model-len 32k、prefix caching 無効、PLE 無し)。そのため下表の PP=3 の値は、3 節の本番構成の値とは直接比較できない。

重みを読むカーネルの時間は 1/3 になるが、1 トークンあたり約 2,000 本の小さいカーネルは 3 枚それぞれで全部実行されるため減らない。これに加えて all-reduce が 1 ステップあたり約 96 回発生する。vLLM の custom all-reduce は world size 3 に対応していない (Supported world sizes: [2, 4, 6, 8, 16]) ため、NCCL が使われる。prefill は、PCIe のみで接続された 3 枚では通信がボトルネックになる。

上限でも +10% であり、さらに QSA のホストオフロードを TP に対応させる改修も必要になるため、実装しなかった。

4.2 MTP の expert をホストメモリに置く

このモデルには 4B の MTP (speculative decoding 用の draft head) がある。前回は VRAM 不足のため無効にしていた。MTP の routed expert を INT4 に量子化してホストの pinned メモリに置き、UVA で GPU から直接読めば、3 枚でも載せられる可能性がある。

短いプロンプトでは速くなるが、それ以外の指標はすべて悪化した。

ホスト上の expert は、1 行あたり top-10 x 2.46 MB = 24.6 MB を PCIe 経由で読むため、1 行 0.9 ms かかる。verify は 1 リクエストあたり 2 行なので、decode で +1.8 ms。prefill では、draft head がプロンプトの全行で MTP layer を実行する (自身の QSA layer の KV を作るため) ため、512 行の chunk ごとに 512 expert すべてを PCIe 経由で読み、+45 ms になる。4 並列で遅くなるのは、並列時は PP の 3 ステージがすでに埋まっており、そこに verify の行数の倍増と、rank 2 のみにかかる UVA 読み出しが加わるためである。

重みをホストメモリに置けるのは、その重みを多数のトークンの計算に使う場合に限られる。 decode 時の expert の重みは 1〜4 トークンの計算にしか使われないため、PCIe の転送時間を償却できない。

4.3 PLE テーブルの RAM キャッシュ

テーブルのうち頻繁に使う行だけを RAM に保持すれば、コールド時の読み出しを減らせる可能性がある。公式ベンチの生成結果・日本語ドキュメント・vLLM のソースコードから 11.8M トークンのトレースを作成し、サーバと同じハッシュで行番号に変換して、LRU・静的な hot set・ページキャッシュの 3 方式を C のシミュレータで比較した。

アクセスされた行は 41.3M 行 = 12.3 GiB、ページ単位では 71.6 GiB だった。行単位のキャッシュはページキャッシュの 12.8 倍の密度になる。ただし trigram の行の 3 分の 1 は初めて出現する行 であり、どの方式でもキャッシュできない。静的な上位 8M 行 (2.4 GiB) の場合、コールド時の NVMe 読み出しは prefill で −32%、decode で −18% だった。decode で 1 行以上 NVMe から読むステップの割合は 0.46 → 0.40 にしか減らない。decode-06 により 16 行の読み出しは並列に発行されるため、1 行でも読み出しがあれば待ち時間はほぼ同じになる。

時間に換算すると、コールド TTFT が約 1 割、decode は 1% 未満の改善である。hot set のリストの作成方法と配布方法の問題もあり、採用しなかった。

5. decode 高速化 第 3 弾: 未使用の VRAM に K/V 行をキャッシュする

3.1 節で述べたとおり、長いコンテキストでの +2 ms は QSA の K/V を PCIe 経由で読む時間で、計算とオーバーラップできない。読み出し自体を減らす方針にした。

indexer の選択結果を計測すると、連続するステップの選択は 60〜85% が重複していた (前ステップとの重複率は平均 0.72、直近 4 ステップの和集合との重複率は 0.86)。過去に読んだ行を VRAM に保持すれば、読み出しの大部分を省ける。

問題は VRAM の確保である。1 リクエスト分でも 1 rank あたり約 20 MiB 必要だが、4 x 262k の最悪ケースで GPU1 の空きは 24 MiB しかない。

VRAM は prefill の staging arena を使う

前回の 5.3 節で述べたとおり、prefill ではコンテキストを区間ごとに VRAM 上の staging buffer (staging arena、QSA の 4 block 分 = 4 x 12,144 x 2 KiB ≈ 95 MiB/GPU) にコピーしてから計算する。staging arena は prefill の forward 中にしか使われず、decode 中は未使用である。 起動時に確保済みなので、これを使えば追加の VRAM は不要になる。

• arena を rank 上の QSA layer 数 (4) で分割し、layer ごとに 12,144 行の direct-mapped cache とする。1 スロットは 1 トークン分の K と V (KV head 2 本分) で 2 KiB。
• キーは論理位置ではなく 物理キャッシュスロット (block x block_size + offset) である。prefix caching で block を共有するリクエスト間では、行のキャッシュも共有される。
• ホスト側の K/V が変わるのは store 時のみなので、store の直後に該当スロットの行を無効化すれば整合性が保たれる。block の再割り当ても store を経由する。
• staged prefill は arena を上書きするため、その後に全タグを無効化する。

1 layer・1 ステップの処理

1. invalidate: 書き込まれたスロットの行を無効化し、rank の epoch を +1 する。
2. plan: ヒットした列は、そのスロットに hitmark = epoch を書き込む。ミスした列は claim = atomic_max(epoch << 32 | slot) でスロットの書き込み権を取得する。
3. attention: split-K カーネルをベースにしたカーネルで、ヒットは VRAM から、ミスはホストから読む。書き込み権を取得でき、かつそのステップでヒットが無かったスロットにのみ、読んだ行とタグを書き込む。

競合の回避は以下の 2 点による。ヒットがあったスロットには同じステップで書き込まない (hitmark == epoch のスロットは書き込み対象にならない)。書き込まれる可能性のあるスロットのタグは読まない (hitmark != epoch の場合はミスとして扱う)。これにより同じカーネル内でタグを読み書きしても不整合が起きず、grid 同期が必要なカーネルは plan の 1 本だけで済む。claim と hitmark は epoch を含むため、ステップごとのクリアも不要である。

また、ホストの pinned メモリと VRAM は UVA で同一のアドレス空間にあるため、以下のように 1 回のロードでヒット・ミスの両方を処理できる。

ptr = tl.where(hit, region_ptr, host_ptr)
row = tl.load(ptr)

ヒットとミスで別々にロードする実装では shared memory が 108 KB になり、3090 の上限 101 KB を超えて起動できなかった。

ヒット時はホストと同じバイト列を読み、タイル内の演算は変更していないため、出力は 近似ではなくビット一致 になる。合成データでの 1,000 ステップ (小さい領域での大量のハッシュ衝突、ステップ間でのホスト側の行の書き換え、arena の上書きなどを含む) と、実モデルの eager 実行で両方の経路を実行して torch.equal で比較する検査を rank あたり 8,000 回以上行い、不一致は 0 だった。

12 行以下に制限した理由

最初の実装では行数を制限しておらず、4 並列 8k の prefill 中に GPU1 が OOM になった。大きい eager バッチ用の Triton カーネルの variant をサービス中にコンパイル・ロードし、そのメモリが予算に含まれていなかったためである。GPU1 の空きは 34 MiB しかない。

12 行以下のバッチではカーネルの launch 設定が 1 通りしかなく、起動時の cudagraph capture の時点でロードされる。そのため decode 形状のバッチ (12 行以下) のみキャッシュを使い、それより大きいバッチは従来の経路を使う (store 時の無効化は常に行う)。

結果

ヒット率は実際の文章で 86〜91%、合成プロンプトで 88〜97%。前ステップとの重複率 (0.72) より高いのは、数ステップ前の行もキャッシュに残っているためである。カーネル単体では、全行をホストから読む場合 1 layer 174 µs、重複率 72% で 78 µs、100% で 38 µs。重複率 0 でも 177 µs で、オーバーヘッドはほぼ無い。

長いコンテキストでの decode 速度が、短いプロンプトとほぼ同じになった。

6. 画像入力への対応

ここまではテキストのみ (--language-model-only) で運用していた。このモデルには ViT (27 ブロック、幅 1152、BF16 で 0.836 GiB) が含まれている。

そのままでは VRAM に収まらない / 不要な重みの発見

--language-model-only を外すと、vLLM は visual を PP の 全 rank に作成する。0.84 GiB x 3 になる。rank 0 のピークはすでに 24,072 MiB で、KV の割り当てを全部削っても収まらない。

vLLM の runner が encoder を実行するのは最初の rank のみなので、他の rank の ViT は不要である。これを調べる過程で、embed_tokens (INT8 0.60 GiB) と lm_head (0.60 GiB) も 3 rank すべてに配置されている ことが分かった。embed を使うのは最初の rank、lm_head を使うのは最後の rank のみである。upstream の qwen4_exp の実装がこうなっており、他の vLLM のモデルのように PPMissingLayer を使っていない。

mem-01 はこの点のみを修正するパッチで、ロード後の重みは [21.58, 21.55, 21.55] → [20.98, 20.44, 21.05] GiB になった (削減量は rank 0 / 1 / 2 で 0.60 / 1.11 / 0.50 GiB)。3 枚合計で 2.2 GiB が不要な重みに使われていた。

ViT のブロックをホストから転送する

重複を削除しても、ViT を rank 0 に常駐させると 2 つの問題があった。

1. ロード中に OOM になる。 MoE の Marlin repack が一時的に 400 MiB を確保する。ViT が先に GPU 上にあると、ここで rank 0 が OOM になる。→ ViT は CPU 上で作成・ロードし、全 repack が完了した後の process_weights_after_loading() で GPU に移す。
2. 常駐させると KV の割り当てを 783 → 480 MB に下げる必要がある。 262,144 x 4.00 が 2.44 になる。

そこで、ViT の 27 ブロックはホストの pinned メモリに置いたままにし、GPU 上の 2 つのスロット (58 MiB) にブロック単位で prefetch しながら転送する 実装にした。ブロック i の計算中に、別ストリームでブロック i+1 をもう一方のスロットにコピーする。各ブロックの forward をラップし、実行直前にパラメータの .data をスロットの view に差し替える (ViT のブロックは @support_torch_compile により __call__ が直接 forward を呼ぶため、nn.Module の hook は呼ばれない)。

4.2 節で述べたとおり、ホストメモリに置けるのは多数のトークンの計算に使われる重みである。ViT の重みは 1 画像の全 patch (数千) の計算に使われるため、この条件を満たす。

1 ブロックの転送は PCIe 4.0 x16 で 1.16 ms で、patch 数に依存しない。計算時間は patch 数に比例し、1 MP で 1 ブロックあたり約 4 ms。計算時間が転送時間より長ければ、表に出るのは最初の 1 ブロックの転送時間のみになる。break-even point は約 0.4 MP で、それより小さい画像では転送時間が支配的になるが、差は最大でも約 12 ms である。TTFT の大部分は LM 側の prefill のため、2 MP 画像で 0.85 → 0.85 s と変化しなかった。

画像は --mm-processor-kwargs '{"max_pixels":2097152}' で 2 MP (2,048 トークン) に縮小している。上限を指定しないと、起動時の profiling が 16.7 MP (16,384 トークン) の画像を想定し、rank 0 の空きが不足する。

画像入力を有効にした状態で、テキストのみの構成より全 rank のピークが低くなった。4 x 247,823 保持の状態で needle 4/4、新規の長いプロンプトの prefill と並行して画像 45/45 を処理できた。

7. MTP を GPU に載せる

4.2 節の方式は採用しなかったが、mem-01 により VRAM の空きが増えたため再検討した。

rank ごとに 0.5〜1.1 GiB 空いたため、PP の layer 配分を均等な 16 / 16 / 16 から 16 / 17 / 15 に変更できる。09-15 に 17 layer を配置した rank が OOM になったのは、不要な embed と lm_head を保持していたためだった (この rank 1 は mem-01 で重みが 1.11 GiB 減る)。この配分で rank 2 のピーク時の空きは約 2.1 GiB になる。MTP の routed expert を INT4 g128 (1.17 GiB) にし、dense 部分 (0.17 GiB) と合わせて GPU に置けば収まる計算になる。

PP で MTP を動作させるための修正 (mtp-01)

stock の vLLM では、PP>1 で MTP を有効にすると起動しなかった。

• draft head の embed_tokens を量子化設定付きで作成する。 PP>1 では本体の embed が最初の rank にあり共有できないため、draft head は自身の embed を持つ。これが quant_config 無しで作成されていたため、INT8 の embed をロードできなかった。
• forward の分岐を修正する。 draft head は最後の rank にのみ存在し、本体の hidden state を直接受け取る。しかし分岐条件が「グローバルの PP rank が先頭かどうか」だったため、PP>1 では常に intermediate tensors の経路に入り、その値が設定されないため assert で停止していた。
• QSA のホスト RAM の見積もりに draft head の layer を加える。 draft head の attention も full attention で、ホストプールを確保する。見積もりは本体の layer_types のみを数えていた。
• attention の block size を、QSA layer 数が最も多い rank に合わせる。 draft head の QSA layer により、最後の rank の QSA group は 4 layer から 5 layer になる。block size は全 rank 共通のため、この group に合わせないとページサイズが mamba のページを超え、プール全体の block が大きくなる。

3 枚に収める (mtp-02)

さらに、最後の rank に置く必要のないものを 2 つ移動した。

• draft head の lm_head: draft head は量子化された lm_head のコピーを作成し、ロード直後に本体の lm_head (同じ rank 上にある) に置き換えていた。置き換えまでの間しか使われないコピーのために、ロード中に Marlin の repack で 0.6 GiB を確保し、rank 2 が OOM になっていた。最初から PPMissingLayer にした。
• draft head の embed_tokens (INT8 0.60 GiB): 1 ステップで数行しか読まないため、CPU 上で作成・ロードし、ロード後は pinned メモリの UVA view にした。数行分の PCIe 転送は無視できる。

MTP の expert の INT4 化は make_mtp_int4.py で行う。RTN の INT4 g128 対称量子化で、group ごとに二乗誤差が最小になる clip 値を探索する。ダウンロードしたチェックポイントは変更せず、3 ファイルだけを別ディレクトリに出力し、compose でマウントして上書きする。

結果

この表は開発時の構成 (--kv-cache-memory-bytes 783000000、--max-num-seqs 4) で、同じ日・同じ計測スクリプトで MTP の有無のみを切り替えて計測した。冒頭の表の 3x3090+MTP 列は配布設定での別の計測回である。

expert を GPU に置いたことで、4.2 節の問題 (prefill 1.5 倍、UVA 読み出し) はほぼ解消した。1 並列では 15〜43% 高速になる。4 並列で遅くなるのは speculative decoding の性質によるもので、GPU の稼働率が高い状態では verify の行数が倍になる分だけ遅くなる。

KV 容量が減るのは VRAM の問題ではなく、2 節の R (262,144 トークンのリクエスト 1 本が使用する block 数) が 61 から 85 に増えるためである。

• QSA: draft head の QSA layer により rank 2 の QSA group が 5 layer になり、block size が 12,144 → 9,776 トークンに小さくなる。
• linear attention: speculative decoding を有効にすると、vLLM は mamba の各 group に draft token 用のページを num_speculative_tokens (1) 個追加する。align モードの 2 ページと合わせて 1 group あたり 3 ページになる。
• block のサイズ: gated delta-net の conv state が draft token 1 個分 (20,480 B) 大きくなり、mamba のページが 3,227,648 B になる。block size 9,776 は、このページに 5 layer x 66 B/トークン (QSA の VRAM 上の部分) が収まる 16 の倍数の最大値である。

783 MB では 242 block / 85 = 2.85 本、550 MB では 170 block / 85 = 2.00 本になる。

decode-07 の行キャッシュも block size に合わせて分割が変わる。staging arena は「QSA layer 数が最も多い rank の layer 数 x block size」トークン分で、MTP 構成では 5 x 9,776 x 2 KiB ≈ 95 MiB になる。rank 2 は 5 layer x 9,776 行、rank 0 / 1 は同じ大きさの arena を 4 layer で分割し、12,220 行ずつ使う。

配布版では --kv-cache-memory-bytes 550000000 (262,144 x 2、pinned 41.2 GiB) と --max-num-seqs 2 にした。4 に設定した場合、250k のリクエスト 2 本を保持している間に画像リクエストが 3 本目として入り、preemption が 20 回発生して画像 1 件が 177 秒待たされた。

MTP の有無にはトレードオフがあるため、両方の構成を配布している。1 ユーザーでレイテンシを優先する場合は 3x3090+MTP、スループットとコンテキスト長を優先する場合は 3x3090 を使う。筆者の本番環境は後者である。

8. MTP + PP で structured output が 500 になる (mtp-03)

MTP 構成を qwen-code で使用中、応答自体は返るが、バックグラウンドで 500 エラーが発生していた。

ERROR [backend_xgrammar.py:168] Failed to advance FSM for request chatcmpl-... for tokens 198.
ERROR [scheduler.py:2073] Unexpected: grammar rejected tokens [198, 198] for request chatcmpl-.... Terminating request.
"POST /v1/chat/completions HTTP/1.1" 500 Internal Server Error

エラーになっていたのはメインの応答ではなく、qwen-code がメインのリクエストと並行して送信する response_format: json_schema 付きのリクエスト (入力候補の生成など) だった。

MTP が有効で、かつ 2 本以上のリクエストが同時に実行されているときのみ発生する。配布パッチは scheduler にも grammar の bitmask にも変更を加えていないため、upstream の vLLM の問題である。

原因は V2 model runner の DraftTokensHandler にある。structured output のリクエストがある場合、engine はサンプリングを遅延させ、take_draft_token_ids() で GPU が生成した draft token を受け取ってから grammar の bitmask を作成する。しかし DraftTokensHandler は 最後にサンプリングしたバッチの draft token しか返さない。

PP の場合、V2 runner は各リクエストを pp_size ステップごとにしか decode しない。リクエストが 2 本あると、交互に別のバッチに入る。そのため bitmask を作成する対象のリクエストの draft token が返されず、-1 のままになる。scheduler は -1 を受け取ると、bonus 位置 (draft token の次の位置) に「draft token を受理する前の grammar state」の mask を適用する。一方、GPU は実際の draft token を verify して受理する。その結果、bonus 位置では draft token 受理後の grammar state では許可されないトークンが選択されうる。accept_tokens がこれを拒否し、リクエストが終了する。拒否されたトークン列が [draft, bonus] の 2 個組であることとも一致する。

修正として、DraftTokensHandler がリクエストごとに最新の draft token を保持し、すべて返すようにした。次の verify で GPU が使うのはこの draft token なので、scheduler 側と一致する。終了したリクエストは dict から削除する。修正前は 2 並列で 16 本中 9 本が 500、修正後は 45 本すべて成功し、acceptance rate は変化しなかった。

発生条件は「V2 runner + async scheduling + PP>1 + speculative decoding + structured output + 2 並列以上」である。

9. 品質: 公式ベンチマークの再現

前回は KLD (BF16 に対して 0.0149 nats) のみを示した。KLD は出力分布の差を表す指標で、回答の正否は表さない。公式モデルカードの表のうち、エージェント環境も外部の judge モデルも不要な 3 つを、本番と同じイメージ・重みで計測した。サーバは本番の compose を元に、prefix caching のみ無効にした (--max-num-seqs 4、262,144 x 4)。

3 つとも公式値との差は 1 SE 前後で、1 回の計測では量子化による劣化は確認できない。サンプリングは generation_config.json のデフォルト (T=1.0 / top_p=0.95 / top_k=20)、thinking 有効、出力長の上限は指定していない。最長の応答は LiveCodeBench の 201k トークンで、上限を 81,920 にすると GPQA と LiveCodeBench の一部の応答が途中で打ち切られる。生成は合計 9.7M トークン、約 13 時間 (平均 約 205 tok/s) だった。

同じサーバで decode 速度を計測すると (出力 400 トークン)、1 並列 95〜100 tok/s、4 並列で合計 283〜298 tok/s だった。冒頭の 4 並列 243〜245 tok/s とは、prefix caching の有無 (有効時は 3.1 節の mamba align モードのカーネルが追加される) と計測スクリプトが異なるため、直接比較できない。

8 並列も試したが、合計 230 tok/s で 4 並列より遅かった。このときは cudagraph_capture_sizes に 8 を追加しており、8 トークンのステップも cudagraph で実行される。一方、decode-04 / 05 の fused カーネルは 4 トークンまでしか対応していないため、8 トークンのステップは fuse 前の経路で実行される。fused カーネルの有無を 8 並列で切り替えた計測はしていないため、この差への寄与は切り分けていない。ベンチマークは 4 並列 (cudagraph_capture_sizes [1,2,4]) で実行した。

LiveCodeBench では採点スクリプトに問題があった。2026-09 時点の lcb_runner の testing_util.py は、正しい解答を 8 件不正解と判定する。

Qwen3.8 系は sys.stdin.buffer.readline や sys.stdout.buffer.write をよく使うため、lcb_runner のままだと 86.3 になる。stdin を使う問題を実際のサブプロセスで採点し直すと 92.4 になった。lcb_runner で正解だった問題はすべてサブプロセスでも正解だったため、判定を緩めたことによる差ではない。他のモデルと LiveCodeBench のスコアを比較する場合は、採点スクリプトを揃える必要がある。

10. 現在の構成

--pipeline-parallel-size 3 --tensor-parallel-size 1
--max-model-len 262144 --max-num-seqs 4
--max-num-batched-tokens 512
--enable-prefix-caching --mamba-cache-mode align
--gpu-memory-utilization 0.97 --kv-cache-memory-bytes 783000000
--limit-mm-per-prompt '{"image":4,"video":0}' --mm-processor-kwargs '{"max_pixels":2097152}'
--compilation-config '{"mode":0,"cudagraph_mode":"FULL_DECODE_ONLY","cudagraph_capture_sizes":[1,2,4]}'

VLLM_PLE_MMAP_PATH=<テーブルファイル>   VLLM_QSA_KV_OFFLOAD=1
VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS=0
PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
cap_add: SYS_PTRACE   ulimits: memlock -1

前回の構成から VLLM_QSA_KV_OFFLOAD_MAX_GIB と VLLM_QSA_KVO_ARENA を削除した。前者はホストプールの確保処理が空き RAM を確認するようになったため、後者は推奨値をデフォルトにしたため不要になった。パッチはすべて Dockerfile でビルド時に適用しており、配布物と筆者の本番環境は同一のコードで動作している。

まとめ

• TTFT: PLE テーブルの読み出しが QD1 の同期読み出しになっていた。リクエスト受付時にプロンプト全体を process_madvise で prefetch し、コールド TTFT を 44 → 8.3 s にした。
• KV 容量: prefix caching の有無で 1 リクエストあたりの block 数が 42 → 61 に変わる。block のコストは主にホストの pinned RAM で、pinned メモリの確保を正確なサイズにしたうえで 262,144 x 4 を確保した。
• decode: PP のホスト側の直列処理、小さい GEMV の低い実効帯域、PLE のキャッシュミス、QSA の K/V の PCIe 読み出しを順に対処し、1 並列 59-73 → 95-99 tok/s、4 並列 120-154 → 243-245 tok/s にした。
• メモリ: prefill の staging arena を decode 中の K/V 行キャッシュに流用し、不要な embed / lm_head を削除して空いた領域に ViT と MTP を配置した。3 枚とも上限まで数十 MiB の状態のため、新規の VRAM 確保は行っていない。
• 採用しなかった案: TP=3 (上限 +10%、prefill −60%)、MTP の expert のホスト配置 (4 並列 −34%)、PLE テーブルの RAM キャッシュ (TTFT −1 割、decode −1% 未満)。

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