TensorRT-LLM

Definition:

TensorRT-LLMとは, NVIDIAが開発した, 大規模言語モデル(LLM)の推論をGPU上で高速かつ効率的に実行するための最適化・実行基盤であり, その本質は, 第1にGPUアーキテクチャに密接に結びついたカーネルレベルの最適化を行う点, 第2にトークン逐次生成という動的過程そのものを最適化対象とする点, 第3にモデル変換から実行, 分散処理に至るまでを統合したエンドツーエンドの推論基盤である点にある.従来のTensorRTを基盤としつつ, トランスフォーマー系モデル特有の計算構造と生成過程に適応するよう拡張されたものである.

TensorRT-LLMの開発の背景には, トランスフォーマーの普及とモデル巨大化がある.特にGPT-3以降, 数十億から数千億パラメータ規模のモデルが実用化されるにつれ, 推論時の計算コスト, メモリ消費, レイテンシが重大な制約となった.従来のTensorRTは主に畳み込みニューラルネットワークを対象として設計されていたため, 自己回帰生成に伴う逐次依存性やKVキャッシュ管理といった課題には十分対応できなかった.このため, LLM特化の推論基盤が求められ, TensorRT-LLMが登場したのである.

同時期には, vLLMやDeepSpeedといったフレームワークが登場し, 推論効率化を巡る競争が激化した.NVIDIA自身も従来はFasterTransformerを提供していたが, 同プロジェクトは現在アーカイブ化[メンテナンス終了]されており, TensorRT-LLMはその後継として, より統合的かつ実運用志向の設計を採用している点に特徴がある.

TensorRT-LLMのアーキテクチャは, モデル変換層, カーネル最適化層, 実行エンジン層から構成される.モデル変換層では, PyTorchなどで記述されたモデルをTensorRTエンジンへと変換し, 演算グラフの最適化や演算子融合が行われる.カーネル最適化層では, AttentionやMLP, Layer Normalizationといった主要演算に対してGPUアーキテクチャに最適化された専用カーネルが適用される.実行エンジン層では, CUDA Graphや非同期実行, ストリーム管理により, 低レイテンシかつ高スループットな推論が実現される.

機能面では, 量子化, KVキャッシュ管理, バッチング戦略, 分散推論が中核を成す.量子化ではFP16やBF16に加え, INT8やFP8といった低精度形式に対応し, メモリ使用量と計算コストを削減する.KVキャッシュ管理では, 自己回帰生成において過去トークンの情報を効率的に保持・再利用するためのメモリ戦略が実装されている.特に, KVキャッシュをページ単位で動的に管理するPagedKVキャッシュ機構を採用しており, これによりメモリフラグメンテーションが抑制され, より多くのリクエストの同時処理が可能となる.バッチングでは, リクエストを逐次的に取り込みながら処理するContinuous Batching[In-flight Batching]によりGPU利用率を最大化する.さらに, Tensor ParallelismおよびPipeline Parallelismを組み合わせた分散推論により, 大規模モデルの実運用が可能となる.

加えて, スループット向上のための重要な手法としてSpeculative Decoding[投機的デコーディング]が実装されている.これは, 軽量なドラフトモデルが複数トークンを先行生成し, メインモデルがその正否を並列に検証することで, 自己回帰生成の逐次的ボトルネックを緩和するものである.また, LoRA[Low-Rank Adaptation]のランタイムサポートも提供されており, 複数のLoRAアダプタを実行時に動的に切り替えるマルチLoRA推論が可能である.これにより, 単一のベースモデルに対して複数のタスク特化アダプタを効率的に運用できる.

ここで重要なのが, TensorRT-LLMにおける「コンパイル」という概念である.TensorRT系の実行モデルでは, 推論に先立ってモデルをGPU向けの最適化済み実行エンジンへと変換する工程が必要となる.このコンパイルは単なる形式変換ではなく, モデル構造および実行条件を固定したうえで, カーネル融合, メモリアロケーション, レジスタ使用といった低レベル最適化を適用する工程である.

したがって, TensorRT-LLMではモデルごとにコンパイルが必要であるだけでなく, 正確には「モデル構造 × 実行設定」の組み合わせごとにエンジンを生成する必要がある.例えば, 同一のLLaMAであっても, 最大シーケンス長やバッチサイズ, あるいは精度設定[FP16かINT8か]が異なれば, 別のエンジンとして再コンパイルが必要となる.なお, モデルの重みについては, エンジンへの埋め込み[ベイクイン]に加え, weight streaming機能を用いて実行時に外部からロードする方式も選択可能であり, 特に大規模モデルの運用においてはこの分離が有効に機能する.この制約は柔軟性を制限する一方で, 極めて高い実行効率を実現するための前提条件である.

また, 生成途中のリクエストを動的にスケジューリングするIn-flight Batchingにより, トークン単位での実行制御が可能となっている.この点は, 従来の固定入力に対する一括推論とは異なり, 逐次生成過程そのものを最適化対象とする設計思想を具体化したものである.さらに, プラグイン機構によりカスタム演算やモデル固有の最適化を組み込むことも可能であり, 柔軟性と性能の両立が図られている.

TensorRT-LLMは, LLaMAをはじめとするGPT系モデル, T5やWhisperといったEncoder-Decoder型モデル, MixtralなどのMixture-of-Expertsモデルに対応し, Triton Inference Serverなどのサービング基盤とも統合されている.このように, 単なる推論高速化ライブラリではなく, 実運用を前提とした包括的なエコシステムを形成している.

特徴と設計思想の比較

比較項目 vLLM TensorRT-LLM
開発元 UCバークレー(オープンソース) NVIDIA
強み PagedAttentionによる効率的なメモリ管理 NVIDIA GPUへの極限の最適化
導入の容易さ 非常に簡単(Pythonフレンドリー、pip installで完結) 複雑(モデルのコンパイル作業が必要)
対応ハード NVIDIA、AMD、CPUなど幅広く対応 NVIDIA製GPU専用(特にH100/A100に最適化)
モデル対応 Hugging Face上の最新モデルに即座に対応 主要モデルに対応するが、変換・最適化の手間がある

参考文献

Mathematics is the language with which God has written the universe.





















追記型アーキテクチャ TriAttention Deep Agents LiteRT-LM Transformerの万能近似能力 LightMem