6月26日、NVIDIAが「Scaling AI Inference Across Multiple GPUs Using NVIDIA TensorRT with Multi-Device Inference Support」と題した記事を公開した。この記事では、TensorRT 11.0で新たに導入されたマルチGPU推論機能の仕組みと実装方法について詳しく紹介されている。
シングルGPUの限界を破る：TensorRT 11.0のマルチデバイス推論
生成AIの推論ワークロードは、単一GPUのメモリと演算能力を急速に超えつつある。特に動画生成・高解像度画像生成パイプラインでは、Attentionブロックが扱うトークン列が数万トークン規模に達し、シーケンス長に対して二乗スケールするAttentionの計算コストがボトルネックになる。
TensorRT 11.0では、この課題に対応するため「マルチデバイス推論サポート（Multi-Device Inference Support）」がネイティブ機能として追加された。従来はモデル並列化を自前で実装する必要があったが、今回からTensorRTランタイムに直接組み込まれた分散通信レイヤーを使って、複数GPU推論を構成できる。
従来のTensorRTは、グラフ最適化エンジンがシングルデバイス向けに設計されており、NCCLなどの集合通信ライブラリとの統合がランタイムレベルで提供されていなかった。そのため、マルチGPU推論を実現するには、TensorRTエンジンの外側でユーザーが独自に通信ロジックを記述し、テンソルの分割・結合・同期を手動で管理する必要があった。TensorRT 11.0ではこの統合がランタイムに組み込まれ、グラフ内に集合通信オペレーションを直接表現できるようになった。
通信基盤はNCCL、戦略の選択がカギ
マルチGPU間の通信には**NVIDIA NCCL（Collective Communications Library）**が使われる。NCCLはNVLink、NVSwitch、PCIe、InfiniBandといったトポロジーを自動判別して最適なトランスポートを選択するライブラリで、大規模モデル訓練では実績のある基盤技術だ。TensorRT 11.0はこのNCCLと直接統合することで、推論ワークロードにも同様のトランスポート最適化を適用する。
対応する集合通信オペレーションは、AllReduce、Broadcast、Reduce、AllGather、ReduceScatter、AlltoAll、Gather、Scatterの全8種類。
並列化戦略の選び方
分散推論には主に2つの戦略がある。
テンソル並列（Tensor Parallelism）は、単一層の重みをGPU間で分割する手法だ。一つの層の重みが1台のGPUメモリに収まらない場合の「唯一の選択肢」になることも多い。Transformerブロックでは列並列・行並列の投影を組み合わせることで、1ブロックあたりのAllReduceを1回に抑えられる。
コンテキスト並列（Context Parallelism）は、入力シーケンスをシーケンス次元でGPU間に分割する手法だ。Attentionの二乗スケーリング問題を直接解決するため、長シーケンスワークロードに特に有効で、因果マスクなしの双方向Attentionを使う拡散モデル（Diffusion/DiTモデル）とも相性が良い。
TensorRT 11.0ではIDistCollectiveLayerプリミティブを導入してこれらの戦略を表現できる。以下ではコンテキスト並列の3実装を比較したベンチマーク結果を紹介する。
コンテキスト並列の3実装：AllGather KV / Ring Attention / Ulysses
コンテキスト並列には実装アプローチが複数あり、それぞれトレードオフが異なる。
AllGather KV：各ランクがQ（クエリ）の自分のスライスを処理しつつ、K・VをAllGatherで全GPU間に共有してからローカルAttentionを計算する。シンプルで実装しやすいが、K・V全体をGPUメモリ上に展開する必要がある。
**Ring Attention**：K・Vをリングトポロジーでストリーミングしながら、各GPUがローカルQの計算をオーバーラップさせる手法。オンラインsoftmaxを使うことでK・VのフルサイズテンソルをどのGPU上にも展開せずに済み、メモリフットプリントを削減できる。
**DeepSpeed Ulysses**：シーケンスをGPU間で分割した後、Attention計算前にQ・K・VにAll-to-Allを適用。各GPUは全シーケンスを受け取るが、担当するAttentionヘッドのサブセットのみを処理する。2回目のAll-to-AllでAttentionヘッド次元とシーケンス次元を再配置して結果を結合する。
ベンチマーク結果：NVIDIA Cosmos 3とFLUX.1で比較
実際のメディア生成パイプラインを対象に、8GPU（シングルノード）構成でのエンドツーエンドレイテンシを計測している。
動画生成（NVIDIA Cosmos3-Nano）：数万トークン規模の長コンテキストを扱うワークロードでは、元記事のベンチマーク結果においてUlyssesが最も低いレイテンシを示した。
画像生成（FLUX.1-dev）：画像生成でもUlyssesが最速だったが、Ring Attentionも4GPU構成でのスケーリングは良好だった。
長コンテキスト（数万トークン規模）においては、Ulyssesがコンテキスト並列の主力戦略として機能する傾向が示されている。具体的な数値については元記事のグラフを参照されたい。
実装ガイド：C++での最小構成
基本的なワークフローはシングルデバイスのTensorRTと同様で、ネットワーク定義に分散通信レイヤーを追加する点が異なる。INetworkDefinition::addDistCollectiveでAllReduceなどの集合通信オペレーションをネットワークグラフに直接組み込める。
auto* collectiveLayer = network->addDistCollective(
inputTensor,
CollectiveOperation::kALL_REDUCE,
ReduceOperation::kSUM,
-1, // root: -1 for collectives without a root rank
nullptr, // groups: nullptr means all ranks participate
0 // groupSize
);
collectiveLayer->setNbRanks(8);
推論実行時はNCCLコミュニケータを実行コンテキストに渡し、OpenMPIで起動する。起動コマンドの詳細はTensorRTリポジトリのサンプルに用意されている。NCCLコミュニケータは、それを使う実行コンテキストの生存期間中は有効でなければならない点に注意が必要だ。
PyTorchモデルを本番C++環境にデプロイする場合は、Torch-TensorRTを使ってPyTorchモデルをTensorRTエンジンに変換するワークフローも利用できる。PyTorchをモデル開発環境として維持しながら、最適化済みのTensorRTエンジンを本番システムにデプロイするアプローチだ。
詳細はScaling AI Inference Across Multiple GPUs Using NVIDIA TensorRT with Multi-Device Inference Supportを参照していただきたい。
