vLLMがLibTorch Stable ABIへ移行、PyTorch更新時の再ビルドを削減

導入

vLLMはC++およびCUDAバックエンドでPyTorchのLibTorchインターフェースを広く利用しています。従来は、ソースコード上のAPIが変わっていなくても、LibTorchのバイナリインターフェースに変更があると、vLLMの再コンパイルと再リンクが必要になる場合がありました。これにより、新しいPyTorchへの対応が遅れ、wheelの保守や配布にも負担がかかります。

vLLMは今回、LibTorch Stable ABIへの移行過程を公開しました。v0.33.0以降、Stable ABIを利用したCUDA wheelを提供しています。ほかの互換条件を満たしていれば、対応するPyTorchの更新時にvLLMを作り直すのではなく、torch wheelを交換して既存のvLLMパッケージから読み込める構成を目指せます。

Stable ABIの役割

ABIは、コンパイル済みコードが関数を呼び出す方法、データを配置する方法、シンボルを解決する方法を定めます。そのため、APIが同じに見えても、既にコンパイルされた拡張機能が動作し続けるとは限りません。Stable ABIは、互換性のあるバージョン間で、サポート対象のバイナリインターフェースを維持する仕組みです。

PyTorchのLibTorch Stable ABIは、主に三つの層で構成されます。選ばれたLibTorch機能を公開する安定したC shim、利用しやすくするC++ラッパー、そしてlibtorch.soへリンクしないtorch::headeronlyのユーティリティです。つまり、LibTorch全体を安定化するものではなく、対象を限定した互換レイヤーです。

vLLMの移行手順

• 依存関係を棚卸しする。 PyTorch APIの呼び出しを調べ、未対応の機能と必要な安定APIを整理しました。
• 未対応の機構をPythonへ移す。 Stable ABIでは安定版のSymIntやPyTorchのautograd機構が利用できないため、一部のmeta functionとAutogradFunctionをPython実装へ移しました。
• 安定版ライブラリを分離する。 CUDAビルドだけでも175超のC++/CUDAソースと120超のヘッダーがあるため、全体を一度に変更せず、独立したライブラリへカーネルを段階的に移しました。TORCH_TARGET_VERSIONは不安定なLibTorchヘッダーを取り込むとビルドを失敗させます。
• 演算子登録を統一する。 Stable ABIの演算子はSTABLE_TORCH_LIBRARYで登録し、実装をTORCH_BOXで包んでC++型をStableIValueへ変換します。
• ファイル単位でAPIを置き換える。 例えばtorch::Tensorをtorch::stable::Tensorに変更し、安定ラッパーやheader-only型を使います。
• 動作とシンボルを検証する。 単体テスト、性能テスト、差分の行単位レビューに加え、torch-abi-auditで共有ライブラリ内のPyTorchシンボルを調査します。

残る制約とトレードオフ

Stable ABIが解決するのはLibTorchのバイナリ互換性だけです。Python API、最低対応CUDAバージョンなどのGPU要件、サードパーティ依存関係は別途確認が必要です。移行中、is_pinned()には安定版の等価APIがなかったため、vLLMはtorch_call_dispatcherを利用してATen演算子をABI安定な方法で呼び出しました。その後、PyTorch 2.14には安定版のC++ is_pinned()メソッドが追加されています。

C shimは型変換、演算子検索、boxingによるホスト側のオーバーヘッドを加える可能性があります。通常は小さく、CUDA Graphのウォームアップ後には見えにくくなる一方、eager実行で非常に高速なカーネルを繰り返し起動する場合は影響し得ます。vLLMは移行中に目立った実行時性能低下を確認していません。

意義

今回の移行は、推論エンジンとPyTorchのリリースを結び付ける要因を一つ減らします。ユーザーは互換性のあるPyTorch更新を早く利用できる可能性があり、開発者や配布者は再ビルドと再リンクの回数を減らせます。また、大規模CUDAプロジェクトが実際に利用するAPIが明らかになり、PyTorch側のStable ABI拡充にもフィードバックを与えます。

ただし、これは任意のPyTorchを無条件に差し替えられるという保証ではありません。段階的移行、ビルド時の制約、シンボル監査、実行時テストを組み合わせ、バイナリ互換性を検証可能な工程に変えた点が、vLLMの取り組みの核心です。

出典：vLLM Blog