0
0

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?

3値27B(Ternary-Bonsai-2-27B)を RTX 5070×2 の vLLM TP=2 で本番運用した記録 — 最後の壁は PCIe x4 と WSL2 の通信

0
Posted at

RTX 5070(12GB)2枚の WSL2 環境で、3値量子化の 27B モデル Ternary-Bonsai-2-27B を vLLM のテンソル並列(TP=2)で本番運用するまでの記録です。移植で踏んだ穴、効いた高速化と効かなかった高速化、そして最後まで残った「長文の読み込み中は他の人の生成が 4 tok/s まで落ちる」問題の真因を、全部実測の数字で書きます。

結論を先に書くと、

  • 1本 71 tok/s・16本同時で合計 422 tok/s、KV プール約 40 万トークンで動いた(以前の sglang 構成はプール 21 万)
  • 長文 prefill 中に他人が遅い真因は 2枚目の GPU が PCIe Gen4 x4 接続だったこと。1ステップの半分が GPU 間通信で、NCCL は既に物理上限
  • WSL2 では NCCL の all-reduce が計算カーネルと一切重ならない。通信を計算の裏に隠す系の最適化は、自作の通信まで作って試したが全滅

です。

前提

  • GPU: RTX 5070 12GB × 2(sm_120)。WSL2 上。GeForce なので P2P 無し
  • マザーボードは B760 チップセット。1枚目は CPU 直結 x16、2枚目はチップセット経由の x4
  • vLLM 0.30.0(V1)、torch 2.13 + CUDA 13.0
  • モデル: fraserprice/Ternary-Bonsai-2-27B-vllm。重みは {-1, 0, +1} の 2bit 表現+128要素ごとのスケールで、入力側にブロック単位のアダマール回転が掛かっている。構造は Qwen3.5 系のハイブリッド(線形アテンション GDN 48層+フルアテンション 16層)
  • 利用者は数人。長文(5〜12万トークン)を投げるエージェントと、短い質問のチャット UI が同居

1. TP=2 への移植: アダマールブロックの真ん中で割れる

配布されているプラグインは TP 非対応でした。問題は2つあります。

  1. 回転の符号表が全幅をキーにしているので、行並列の層で分割後の幅を引くと KeyError
  2. down_proj の入力幅 17408 を2分割すると 8704 で、1024 のアダマールブロックのちょうど真ん中で割れる

2 は回転が層をまたいでいるので、単純に各ランクで回すと答えが変わります。アダマール行列の再帰構造(Sylvester 型)を使うと、

H1024 [a; b] = [ H512·a + H512·b ;  H512·a − H512·b ]

なので、各ランクが自分の半分 512 列だけ回して all_gather で交換し、ランク0は和、ランク1は差を取れば元と同じになります。カーネルのバタフライ最終段がこの形になっていることを確認したうえで実装しました。CPU 上の模擬テストで TP=1 比の相対誤差 2.6e-3(bf16 の丸め並み)、わざとランクを入れ替えると 1.44 になって検出できることも確かめています。

2. KV を 3bit にして MTP と両立させる

12GB 2枚で KV プールを稼ぐため、vLLM 標準の TurboQuant(turboquant_k3v4_nc、キー3bit・バリュー4bit)を使いました。bf16 KV ではプール 8.8 万トークンでしたが、これで 37〜40 万トークンになります。

ところが TurboQuant と MTP(投機デコード)を併用すると出力が崩壊しました(17×23 に 393 と答える、1 1 1 … を繰り返す)。これは上流の既知バグで、未マージの修正 PR #53410 を手で当てて解決しました。

MTP の予測数は実測で 1 にしています。

予測数 1本 4本合計 8本合計 16本合計 プール
なし 53 176 302 487 47.4万
1 71 224 370 422 39.6万
2 85 258 315 372 35.3万
3 77 211 256 222 —

(単位 tok/s)予測2は少人数なら速いのですが、8本以上で負けます。2個目の当たり率が 28% しかなく、外れの計算が同時利用時の GPU を食うためです。複数人運用なので 1 を選びました。

3. 最初に付け忘れて精度テストを無効にした話

新しいサーバを立てるとき、--reasoning-parser qwen3 --tool-call-parser qwen3_coder --enable-auto-tool-choice を付け忘れていました。思考が本文に混ざったまま採点したので、その回の精度テストは全部無効です。比較対象と同じパーサーは最初の起動から付けるのが教訓です。

付け直した後の自前の長文集計テスト(96K トークン文脈で複数箇所を数える問題)では、以前の sglang 構成が 10/14 だったのに対し 14/14 でした。

4. 効いた高速化

prefill の線形層を int8 GEMM に(242 → 114 ms/ステップ)

3値重みは各行・各128要素群に fp16 スケールを持っています。このスケールを「行ごとの最大値 × 8bit の比率」に分解すると、

q[n, g] = round(127 * s[n, g] / smax[n])
W_int8[n, k] = code[n, k] * q[n, k/128]     (厳密に int8 に収まる)

となり、prefill の行列積を丸ごと cuBLAS の int8 GEMM(torch._int_mm)1回+行・トークン単位の再スケールにできます。活性化はアダマール回転後にトークン単位で int8 化します(回転が外れ値を均してくれる)。512トークンで線形層が 242 → 114 ms、層ごとの相対誤差は 2.4e-3 → 9.5e-3 でしたが、精度テストでは差が出ませんでした。

TurboQuant のデコードをまとめ読みに(長文の生成 +29%)

標準のデコードカーネルは Q ヘッドごとに1プログラムなので、GQA(1つの KV ヘッドを6本の Q ヘッドが共有)だと同じ KV を6回読んで6回解凍します。KV ヘッド単位で1プログラムにして6本の Q をまとめて計算するよう書き換えました(考え方は上流 PR #40792 と同じで、MSE キー向けに自作)。8万トークンで 1層あたり 3.82 → 1.91 ms、実運用の 4.6万トークン文脈で生成が 37.9 → 48.9 tok/s になりました。

mamba 状態は fp16 ではなく bf16

GDN の状態を fp16 にすると block サイズが小さくなってプールは増えるのですが、融合 GDN デコードカーネルが使えず Triton 実装に落ちて 14% 遅くなりました。bf16 なら融合カーネルのままです。

prefix キャッシュの照合単位

ハイブリッドモデルの align モードでは、キャッシュ命中がブロック(数千トークン)刻みになり、命中率が 28% しかありませんでした。--prefix-match-unit 16 --prefix-caching-hash-algo xxhash で 16 トークン刻みにすると、会話の2ターン目以降の命中が 99% になりました(xxhash は venv に別途入れる必要があり、入れ忘れると起動に失敗します)。

長文 prefill を細かく刻む

長文の読み込みと他人の生成は同じステップに相乗りするので、1ステップに混ぜる prefill 量が大きいほど他人が遅くなります。生成中の人がいるときは1ステップの prefill 合計を 256 トークンまで、誰もいないときは1件あたり 512 まで、とスケジューラに手を入れて切り替えています。

5. 効かなかったもの

試したこと 結果
dp4a による decode GEMV(bonsai-turbo の手法を移植) カーネル単体 2.1倍、本番では +2% のみ。誤差が増えるので不採用
Triton で 2bit 解読と行列積を融合 cuBLAS に負ける(bf16 で 0.62倍、int8 でも 11% 短縮止まり)
CPU への KV 退避(--kv-offloading-backend native) CPU 側が GPU プールより小さく、あふれた会話は CPU からも先に消える
LMCache ハイブリッドモデルで同梱コネクタが RuntimeError
パイプライン並列(PP=2) prefill +50% だが decode 半減
NCCL_PROTO=LL 単発 +9%、prefill -25%
1ステップの prefill 上限を 128 に 他人の生成は速くならず、長文が遅くなるだけ

dp4a の件は教訓になりました。カーネル単体の倍率は、CUDA グラフの中で動く本番のステップ全体の改善を保証しません。

6. 最後に残った問題: 長文の読み込み中、他人の生成が 4.5 tok/s

エージェントが数万トークンの文脈を読み込んでいる間、チャット側の生成が 4.2〜4.8 tok/s まで落ちます。GPU の温度も電力も上がらない(全力 204W に対して 115〜130W)ので、どこかで GPU が遊んでいるはずです。

py-spy は「CPU のコピー待ち」に見えた

ワーカーを py-spy で見ると、40.8% が入力バッファの同期コピー、38.8% が torch.cuda.synchronize でした。WSL2 では vLLM がピン留めメモリを使えないと判定して、同期コピーのフォールバックに落ちているのが見えます。

vLLM には VLLM_WSL2_ENABLE_PIN_MEMORY=1 という公式の環境変数があり、手元の WSL カーネル(6.18)ではピン留めメモリも GPU からの直接読み取りも正常に動きました。コピー1回は 27.9 → 4.5 µs になります。ところが本番に入れても速度は変わりませんでした。 py-spy に見えていたのは CPU の無駄ではなく、GPU の処理が終わるのを待っている時間だったわけです。

部品ごとに測る

そこで、1ステップ(prefill 256 トークン)を部品ごとに単体で測りました。

部品 時間
GPU 間 all-reduce(64層 × 2回 = 128回 × 0.74 ms) 95 ms
3値の線形層(うち重みの int8 展開 31 ms) 83 ms
その他(GDN・アテンション等) 約 20〜30 ms

合計が実測の約 0.2 秒と一致します。ステップの半分が通信で、その間 GPU は計算していません。電力が低いのもこれで説明できます。

2枚目が x4 だった

ホストとのコピー速度を GPU ごとに測ると、はっきり差が出ました。

PCIe D2H (2.6MB) H2D (2.6MB)
GPU0 Gen4 x16 25.6 GB/s 23.2 GB/s
GPU1 Gen4 x4 7.1 GB/s 6.9 GB/s

P2P が無いので all-reduce はホストメモリ経由です。GPU1 の 7 GB/s で 2.6MB を送って受けると 0.37 + 0.37 = 0.74 ms で、NCCL の実測値そのものです。B760 は x8/x8 の分割に対応していないので、スロットの差し替えでは改善できません。

念のため NCCL の設定(NCCL_PROTO 3種、チャネル数、バッファサイズ、Tree、SHM の memcpy モード)を9通り振りましたが、どれも同じか悪化でした。

通信を減らす・隠す、は全部だめだった

通信の int8 圧縮。 128要素ごとのスケール付きで int8 にして all_gather し、自分の分は bf16 のまま足す方式です。バイト数は半分でも 1ステップ 100 → 79 ms(-21%)止まりで、1回あたり 1.5% の誤差が 128回積み重なります。見送りました。

通信と計算を重ねる。 本命でした。リクエストの束を2つに割り、片方の all-reduce 中にもう片方の線形層を計算します(どちらの束もモデル全体を通るので、答えは変わりません)。vLLM 内蔵の Dual Batch Overlap は MoE とデータ並列向けで、この構成では使えないため自作しました。ところが、1層分の試験で重ねた版のほうが遅くなりました(164 → 203 ms)。

原因を切り分けると、こうなりました。

試験(WSL2) 時間
計算カーネルだけ 7.0 ms
NCCL all-reduce だけ 7.6 ms
別ストリームで同時に投げる 14.6 ms(ただの和)
カーネル同士 / カーネル+cudaMemcpyAsync 7.0〜7.1 ms(重なる)

WSL2 では NCCL の all-reduce が、別ストリームの計算カーネルと一切重なりません。 NCCL のプロトコルやチャネル数、SHM の memcpy モードを変えても同じでした。

コピーエンジンなら重なるので、最後にコピーだけで all-reduce を作りました。/dev/shm を両プロセスで mmap して cudaHostRegister し、各ランクが自分の部分和を D2H、CPU のワーカースレッドが共有メモリのフラグで相手と同期して、相手の部分和を H2D します。結果は NCCL とビット単位で一致しました。しかし、

方式(256トークン × 64層) 時間
NCCL・順番に流す 164 ms
自作・順番に流す 181 ms
自作・2分割で重ねる 190 ms
自作・4分割 261 ms
自作・8分割 413 ms

分割するほど悪化します。WSL2 では CPU と GPU の間の同期(イベント待ち・コピーの投入)に1回あたり約 0.3 ms の固定遅延があり、受け渡しの回数が増えるほど積み上がるためです。そもそも GPU1 は1層あたりの通信(約 1.5 ms)が計算(約 1.1 ms)より長いので、完全に隠せても下限は約 96 ms でした。

途中で見つけた WSL2 の癖も書いておきます。

  • 長い単一カーネルの実行中に後から投げたコピーは、そのカーネルが終わるまで待たされる(短いカーネルが並んでいる場合は間に入れる)。torch.cuda._sleep のような長いカーネルで重なりを試すと誤判定する
  • cudaHostRegister した共有メモリからの DMA 自体は問題なく動く(以前、同じメモリに GPU カーネルが書けない問題を踏んでいるが、コピーエンジン経由なら可)

まとめ

  • 3値 27B は RTX 5070 × 2 の vLLM TP=2 で実用速度になる。1本 71 tok/s、16本同時で合計 422 tok/s、プール約 40 万トークン
  • 移植のポイントは、アダマールブロックを跨ぐ分割の組み直し、TurboQuant×MTP の上流修正、int8 prefill、KV まとめ読みのデコード
  • コンシューマ向けマザーボードで GPU を2枚挿すときは、2本目のレーン数を確認する。 TP=2 の混載ステップの半分は GPU 間通信で、x4 だとそこが物理上限になる
  • WSL2 では通信の隠蔽が効かない。NCCL は計算と重ならず、コピーエンジンで自作しても同期の固定遅延が勝つ
  • プロファイラの「待ち」は、何を待っているかまで部品単位で測って初めて意味がある
0
0
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
0
0

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?