vLLM の MoE モデル向け Multi-LoRA サービング

vLLM は、Multi-Low-Rank Adaptation(Multi-LoRA)サービングの効率的なソリューションを導入し、組織が単一の GPU 上で数十個のファインチューニングされたモデルをホストできるようにしました。これにより、各カスタムモデルごとに専用のコンピュートエンドポイントを用意する必要がなくなります。個々のモデルが十分なトラフィックを受けずハードウェアを飽和させない場合、GPU がアイドル状態になることがよくあります。

MoE モデル向け Multi-LoRA サービング

Multi-LoRA は、元のベースモデルの重みを凍結したまま、リクエストごとに小さな学習可能アダプタを入れ替えることで、複数のカスタムモデルが同じ GPU を共有できるようにします。これは、ルーターが入力トークンを専門化されたニューラルネットワーク(エキスパート)に振り分ける Mixture of Experts(MoE)モデルに特に有益です。

MoE アーキテクチャでは、各エキスパートが「拡張後圧縮」パターンを使用します。gate_up プロジェクションは隠れ状態をより大きな中間空間に拡張し、down プロジェクションはそれを元に圧縮します。LoRA を適用すると、各エキスパートは gate_updown の両方のプロジェクションに対して、縮小と拡張の操作をそれぞれ行う、合計4つの追加カーネル操作が必要になります。

これに対処するため、vLLM は fused_moe_lora カーネルを実装しました。このカーネルは LoRA 操作を fused_moe カーネルに直接統合し、エキスパートルーティング(トークンが異なるエキスパートに割り当てられる)とアダプタ選択(リクエストが異なる LoRA アダプタを使用する)によって生じる複合的なスパース性を管理します。

パフォーマンス向上のための技術的最適化

MoE モデル向けのマルチ LoRA の初期実装は高レイテンシに悩まされていました。vLLM と AWS は NVIDIA Nsys と NCU を活用し、3 つの主要な最適化層を通じてボトルネックを特定・解消しました。

実行最適化

プロファイリングにより、Triton コンパイラが入力長に依存する変数をコンパイル時定数として扱い、fused_moe_lora カーネルが新しいコンテキスト長ごとに再コンパイルされていることが判明しました。その結果、ベースモデルに比べて Time To First Token(TTFT)が 10 倍に悪化しました。do_not_specialize コンパイラヒントを追加することで、カーネルが一度だけコンパイルされ、すべてのコンテキスト長で再利用されるように解決しました。

カーネルレベルの最適化

LoRA 行列は「スキニー」(ランク r が隠れ状態次元に比べて大幅に小さい)ため、標準的な GEMM カーネルの性能が低下します。以下の最適化が実装されました:

  • Split-K Work Decomposition: この戦略は内部次元 K の総和を複数のスレッドグループに分割し、部分和を並列に計算してスキニー行列のロードバランシングを改善します。アトミック加算は Triton コンパイラで sem="relaxed" を使用して最適化されました。
  • CTA Swizzling: Cooperative Thread Array(CTA)スウィズリングは、GPU のスケジュールを再配置し、近接した列で作業するスレッドグループが同時に実行されるようにして、L2 キャッシュの再利用を高めます。
  • Masking Reduction: 行列次元がブロックサイズで均等に割り切れる場合に条件付きマスクチェックを省略するため、EVEN_K パラメータが導入され、各ロード操作のオーバーヘッドが削減されました。
  • Kernel Fusion: LoRA 重みとベースモデル重みの加算が LoRA 拡張カーネルに統合され、カーネル起動オーバーヘッドが削減されました。

Amazon 固有のチューニング

標準的な fused MoE 用に最適化されたデフォルトの Triton カーネル設定は、マルチ LoRA の複合スパース性を考慮していませんでした。AWS は 4 つの fused_moe_lora 操作(gate_up_shrinkgate_up_expanddown_shrinkdown_expand)向けにカスタムチューニングされた設定を開発しました。これらの設定は、Amazon SageMaker AI および Amazon Bedrock を利用する顧客に自動的にロードされます。

パフォーマンスベンチマークと利用可能性

これらの最適化により、GPT-OSS 20B、Qwen3-MoE、DeepSeek、Llama MoE といった MoE モデルや、Llama 3.3 70B、Qwen3 32B といった密結合モデルのパフォーマンスが大幅に向上しました。

GPT-OSS 20B において、vLLM 0.11.1rc3 から vLLM 0.15.0 への移行により、Output Tokens Per Second(OTPS)が 454% 向上し、TTFT が 87% 短縮されました。Amazon SageMaker AI と Amazon Bedrock で利用可能なさらなる最適化により、vLLM 0.15.0 と比較して OTPS がさらに 19% 増加し、TTFT が 8% 短縮されます。

GPT-OSS 20B の結果概要:

バージョン/プラットフォーム OTPS TTFT
vLLM 0.11.1rc3 32 (Initial) ~1.2s (Estimated)
vLLM 0.15.0 144 135 ms
Amazon SageMaker/Bedrock 171 124 ms

これらの改善は vLLM バージョン 0.15.0 以降で利用可能です。

Sources