Keyboard shortcuts

Press or to navigate between chapters

Press S or / to search in the book

Press ? to show this help

Press Esc to hide this help

LLMシステムの基礎

GPUアーキテクチャ — シリコンからLLM学習まで

現代の大規模言語モデルは、ほぼ例外なくGPU(Graphics Processing Unit、画像処理装置)上で学習・提供されます。GPUアーキテクチャを理解することは、並列化戦略、メモリ管理、カーネル最適化、インフラ規模のサイジングについて、情報に基づく判断を行うために不可欠です。ここでは、LLMワークロードに関係するGPUハードウェアを包括的に紹介します。

なぜ深層学習にGPUを使うのか?

GPUとCPUは、根本的に異なるハードウェア設計思想を持っています。この違いを理解すると、LLMの学習がGPU上では100〜1000\(\times\)高速になる理由が分かります。

Tip

CPUとGPU — 根本的な設計思想

  • CPUレイテンシ向けに最適化されています。大容量キャッシュ、分岐予測、アウト・オブ・オーダー実行を備え、少数のスレッドを可能な限り高速に実行します。現代のCPUは8〜96コアです。

  • GPUスループット向けに最適化されています。それぞれが単純な処理を行う数千のスレッドを並列に実行します。現代のGPUには数千の「コア」(実行ユニット)があり、ストリーミング・マルチプロセッサ(SM)にグループ化されています。

深層学習のワークロードは、行列乗算(\(O(n^3)\)の演算を\(O(n^2)\)のデータに対して行う)に支配されます。行列乗算は非常に並列化しやすい処理です。70BモデルのTransformerの1回のフォワードパスは、トークンあたり約\(\sim\)140 TFLOPの計算を必要とするため、GPUのスループットに最適です。

NVIDIA GPUマイクロアーキテクチャの世代

NVIDIAはGPUアーキテクチャを継続的にリリースしており、それぞれが深層学習に重要な革新をもたらしています。

アーキテクチャ代表製品深層学習における主な革新
Pascal2016P100初のHBM GPU、FP16対応、NVLink 1
Volta2017V100Tensor Core (第1世代)、混合精度学習
Turing2018T4INT8推論、RTコア(ML向けではない)
Ampere2020A100BF16 Tensor Core、TF32、第3世代NVLink、MIG
Hopper2022H100FP8 Tensor Core、TMA、Transformer Engine、NVLink 4
Blackwell2024B200第2世代Transformer Engine、NVLink 5(1.8 TB/s)、FP4

深層学習向けNVIDIA GPUマイクロアーキテクチャのタイムライン。

LLMの学習と推論に使われる主なGPU

GPUアーキテクチャHBMBF16 TFHBM帯域幅NVLinkLLMでの役割
V100-32GBVolta32 GB125 TF*900 GB/s300 GB/sレガシー、小規模モデルのファインチューニング
A100-40GBAmpere40 GB312 TF1.5 TB/s600 GB/s低予算の学習・推論
A100-80GBAmpere80 GB312 TF2.0 TB/s600 GB/s標準的なRLHF(70Bで8〜64枚)
H100 SXMHopper80 GB990 TF3.35 TB/s900 GB/s学習が3\(\times\)高速
H200 SXMHopper141 GB990 TF4.8 TB/s900 GB/sより少ないGPUで70Bの方策+参照モデルを収容
B200 SXMBlackwell192 GB2250 TF8.0 TB/s1800 GB/s次世代、H100の2\(\times\)
MI300XCDNA3192 GB1300 TF5.3 TB/sN/A最大のメモリ容量、ROCm
TPU v5eGoogle16 GB197 TF1.6 TB/sICI 1.6 TB/sクラウド専用、JAX/XLA

Warning

どのGPUを選ぶべきか?

  • 70B以上のモデルを学習する場合 :NVLinkを備えたH100/B200ノード(テンソル並列化には高速なインターコネクトが必要)。1インスタンスあたり最低8\(\times\)H100。

  • 推論(レイテンシ重視) :高帯域幅が必要ならH100/H200、メモリ帯域律速(巨大なKVキャッシュ)ならMI300X。

  • 7B〜13Bのファインチューニング :A100-80GBが費用対効果に優れています。LoRAなら単一GPUで可能です。

  • 低予算 :A100-40GB、あるいは7BモデルのLoRAならA10(24GB)でも可能です。

GPU内部アーキテクチャ — ストリーミング・マルチプロセッサ(SM)

GPUは ストリーミング・マルチプロセッサ(SM) の配列として構成されます。各SMは、専用のレジスタファイル、共有メモリ、実行ユニットを持つ独立したプロセッサです。GPUの性能を理解するには、SMを理解することが重要です。

Important

SMの主要コンポーネント

  • CUDAコア :FP32/INT32演算用のスカラーALU。A100ではSMあたり64個。要素単位の演算、リダクション、行列以外の演算に使われます。

  • Tensor Core :行列積和(MMA)に特化したユニット。各ユニットは1サイクルで\(4{\times}4{\times}4\)の融合積和を実行します。A100ではSMあたり4個で、対応する精度ではCUDAコアの\(16\times\)のスループットを実現します。

  • レジスタファイル :最速の記憶領域(レイテンシ1サイクル)。アクティブな全スレッドで共有されます。L1へのレジスタ・スピルは大幅な低速化を引き起こします。

  • 共有メモリ/L1 :プログラマーが明示的に管理するオンチップSRAM。Flash Attentionの性能の鍵です(タイル全体を共有メモリに収められます)。

  • ワープ・スケジューラ :各SMにはワープ・スケジューラが4個あります(A100)。ワープはロックステップで実行される32スレッドの単位です(SIMTモデル)。スケジューラはワープを切り替えてメモリレイテンシを隠蔽します。

Tip

SIMT実行モデル

GPUはSingle Instruction, Multiple Threads(SIMT)実行を使います。ワープ(32スレッド)の内部では、全スレッドが同じ命令を異なるデータに対して実行します。スレッドが分岐すると(例:if/else)、両方の経路が直列化されます。これをワープ分岐と呼びます。そのため、GPUカーネルでは分岐を最小限に抑える必要があります。

LLMワークロードの主要な演算(GEMM、アテンション、softmax)はスレッド間で制御フローが均一であり、SIMT実行に適しています。

世代をまたぐGPUチップのスケーリング

NVIDIA GPUアーキテクチャの進化では、計算密度、オンチップメモリ、深層学習用の専用ユニットが一貫して拡大してきました。

アーキテクチャSM数TC/SMSRAM/SML2主な変化
Volta (V100)808128 KB6 MBTensor Coreを導入
Ampere (A100)1084192 KB40 MBBF16/TF32、L2の大容量化
Hopper (H100)1324256 KB50 MBTMA、FP8、Thread Block Cluster
Blackwell (B200)1484256 KB128 MB2\(\times\)ダイ、FP4、TMEM、NVLink 5

NVIDIAアーキテクチャにおけるSMレベルのスケーリング。

GPUメモリ階層と帯域幅

現代のGPUにおける学習・推論性能は、メモリ階層をまたぐデータ移動をどれだけうまく管理できるかによって、ほぼ完全に決まります。メモリ階層を理解することは任意ではありません。後の各最適化手法の基盤です。

Important

GPUメモリ階層 — A100 80GBの基準値

レベル容量帯域幅レイテンシ場所
レジスタ\(\sim\)256 KB/SM\(>\)100 TB/s1サイクルオンチップ、スレッド単位
SRAM(共有)164 KB/SM\(\sim\)19 TB/s\(\sim\)20サイクルオンチップ、SM単位
L2キャッシュ合計40 MB\(\sim\)5 TB/s\(\sim\)200サイクルオンチップ、共有
HBM2e(VRAM)80 GB2 TB/s\(\sim\)200 nsパッケージ上(5スタック)
CPU DRAM512 GB以上\(\sim\)25 GB/s\(\sim\)10 \(\mu\)sホスト(PCIe 4)
NVMe SSDTB単位7 GB/s\(\sim\)100 \(\mu\)sホストストレージ

Tip

なぜ差がこれほど大きいのか

階層の各レベルは、その1つ上のレベルに比べておおむね 10\(\times\)遅く、100〜1000\(\times\)大きい ものです。A100はBF16 Tensor Coreで312 TFLOP/sのスループットを持ちますが、HBM帯域幅は2 TB/sしかありません。つまり、HBMからロードした1バイトごとに、次のバイトが到着するまでに\(312 \times 10^{12} / (2 \times 10^{12}) \approx 156\)回の浮動小数点演算を実行できます。カーネルが1バイトあたり156 FLOP未満しか実行しない場合、そのカーネルはメモリ律速です。計算ユニットはデータを待ってアイドル状態になります。

レジスタ

各CUDAスレッドはプライベートなレジスタファイルを利用できます。レジスタはチップ上で最速の記憶領域です。調停なしに1クロックサイクルで読み書きできます。A100にはSMあたり65,536個の32ビットレジスタがあります。レジスタをローカルメモリ(L1/L2)へスピルさせることは、大きな性能上の危険要因です。

SRAM — 共有メモリ/L1

各SMには、A100で192 KB(H100では256 KB)のL1/共有メモリ統合プールがあり、A100では最大164 KBを共有メモリとして構成できます。共有メモリはプログラマー(または新しいCUDAバージョンのコンパイラー)が明示的に管理します。たとえばFlash Attentionは、アテンションのタイル計算をSRAMに収められるという洞察を中心に構築されています。

L2キャッシュ

A100の40 MBのL2は、108個すべてのSMで共有されます。SRAMとHBMの間のステージング領域として機能します。空間的局所性が高いワークロード(たとえばバッチ内で重み行列に繰り返しアクセスする場合)では、L2のヒット率が実効的なHBMトラフィックを大幅に削減します。

HBM — High Bandwidth Memory

HBMはGPUパッケージに直接搭載された積層DRAMで、広いインターポーザーを介して接続されています。A100 SXMは、2 TB/sで80 GBのHBM2eを備えています。H100 SXM5は、3.35 TB/sで80 GBのHBM3を備えています。これは、モデルの重み、KVキャッシュ、アクティベーション、オプティマイザー状態の主要な作業メモリです。

PCIe経由のCPU DRAM

GPUのHBMとCPUのDRAM間のデータ転送は、PCIeバスを通過します。PCIe Gen4 \(\times\)16は、方向あたり約\(\sim\)32 GB/s(双方向で64 GB/s)を提供し、Gen5ではその倍になります。これはHBM(片方向)と比べて約\(\sim\)60\(\times\)帯域幅が小さい値です。CPUオフロード(ZeRO-Infinity、DeepSpeed)はこのリンクを利用しますが、ボトルネックにならないよう注意して使う必要があります。

NVMe

NVMe SSD(Samsung 990 Proなど)は、シーケンシャル読み出しで約\(\sim\)7 GB/sに達します。ZeRO-Infinityはオプティマイザー状態をNVMeへオフロードできますが、これは計算とI/Oの比率が非常に高い場合(大きなバッチサイズ、遅い学習ステップ)に限って実用的です。

算術強度とRooflineモデル

Important

算術強度

\[ I = \frac{\text{FLOPs}}{\text{Bytes accessed from HBM}} \quad \text{(FLOPs / Byte)} \] カーネルは、\(I < I_{\text{ridge}}\)のとき メモリ律速 、\(I > I_{\text{ridge}}\)のとき 計算律速 です。ここで \[ I_{\text{ridge}} = \frac{\text{Peak FLOP/s}}{\text{Peak Bandwidth}} = \frac{312 \times 10^{12}}{2 \times 10^{12}} = 156 \text{ FLOP/Byte (A100 BF16)} \]

Note

アテンションの算術強度

シーケンス長\(n=4096\)、ヘッド次元\(d=128\)の単一アテンションヘッドを考えます。

FLOP :\(QK^T\)には\(2n^2d\)、softmaxには\(O(n^2)\)、\(\text{Attn} \times V\)には\(2n^2d\)のコストがかかります。合計は\(\approx 4n^2 d = 4 \times 4096^2 \times 128 \approx 8.6\) GFLOPです。

メモリトラフィック (標準的な非Flash実装):

  • \(Q, K\)を読み出す:\(2 \times n \times d \times 2 = 2\) MB

  • アテンションスコア\(S = QK^T\)を書き出す:\(n^2 \times 2 = 33.5\) MB

  • softmaxのために\(S\)を読み出す:\(n^2 \times 2 = 33.5\) MB

  • softmax出力\(P\)を書き出す:\(n^2 \times 2 = 33.5\) MB

  • 最終行列乗算のために\(P\)と\(V\)を読み出す:\(n^2 \times 2 + n \times d \times 2 = 34.5\) MB

  • 出力\(O\)を書き出す:\(n \times d \times 2 = 1\) MB

総メモリ量 :\(\approx 138\) MB(\(n^2\)のアテンション行列を4回走査することが支配的です)。

算術強度 : \[ I = \frac{8.6 \times 10^9}{138 \times 10^6} \approx 62 \text{ FLOP/Byte} \]

これはA100のリッジポイントの\(62/156 = 40%\)であり、 明確にメモリ律速 です。GPUはメモリを待って60%の時間アイドルになります。

Flash Attentionによる解決 :\(n \times n\)行列を実体化せず(\(Q, K, V\)をSRAMでタイル化する)、Flash AttentionはHBMトラフィックを\(Q, K, V\)の読み出しと\(O\)の書き出しだけに削減します:\(4 \times n \times d \times 2 = 4\) MB。ロードした各バイトは\(O(n)\)回の計算で再利用されます(各クエリがすべてのキーにアテンションするため)。したがって、 \[ I = \frac{4n^2 d}{4 \cdot n \cdot d \cdot 2} = \frac{n}{2} = \frac{4096}{2} = 2048 \text{ FLOP/Byte} \]

これはリッジポイント(156)の\(13\times\)上であり、 深く計算律速 です。GPUはピークの312 TFLOPSに達し、必要な帯域幅は\(312\text{T}/2048 \approx 152\) GB/s(HBM容量の7.6%)にすぎません。もはやメモリはボトルネックではありません。

アテンションはメモリ律速、FFNは計算律速

Tip

Transformerにおける2つの領域

Transformerブロックには、算術強度が大きく異なる2つの主要コンポーネントがあります。

  • アテンション :\(n \times d\)テンソルを処理します。\(QK^T\)積は\(O(n^2 d)\) FLOPを必要としますが、アテンションスコアのために\(O(n^2)\)のメモリを必要とします。長いシーケンスではメモリトラフィックが支配的になり、アテンションはメモリ律速です。

  • FFN(MLP) :形状\([d_{\text{model}}, 4d_{\text{model}}]\)の重み行列を持つ2つの大規模線形層です。これらは算術強度の高い大規模GEMMであり、FFNは計算律速です。

そのため、Flash Attention(メモリ最適化)はアテンションには有効ですがFFNには効きにくく、量子化(重みサイズの削減)はアテンションよりFFNに大きな効果があります。

Tensor Core

Important

Tensor Coreとは何か?

Tensor Coreは、Volta(2017年)で導入された行列積和(MMA)専用ユニットです。各Tensor Coreは、1クロックサイクルで\(4\times4\times4\)の行列乗算を実行します。 \[ D = A \times B + C \quad (4\times4 \text{ matrices}) \] A100には、108個のSM全体で 432個のTensor Core があります(SMあたり4個、各サブパーティションに1個)。BF16精度では312 TFLOP/sを実現し、FP32 CUDAコアのスループットの約\(16\times\)です。

  • 対応精度 :FP64、TF32、BF16、FP16、INT8、FP8(H100以降)。

  • 累積 :BF16入力であっても、内部では常にFP32で累積します。これにより、ドット積での壊滅的な桁落ちを防ぎます。

  • 要件 :Tensor Coreは、行列の次元が8(BF16)または16(FP8)の倍数のときに最も効率的です。こうした倍数に合わせたパディングは、しばしば実施する価値があります。

  • WGMMA(H100) :Hopperでは、より大きなタイル(64\(\times\)256\(\times\)16)を処理するワープグループ単位のMMA命令が導入され、TMA(Tensor Memory Accelerator)によるデータ移動とパイプライン化できます。

Warning

Tensor Coreの落とし穴

Tensor Coreが役立つのは、カーネルが計算律速の場合だけです。小さなバッチ(バッチサイズ1の推論など)ではGEMMのタイルが小さく、Tensor Coreの利用率が低くなり、再びメモリ律速の領域に戻ります。そのため推論エンジンはリクエストを積極的にバッチ化します。

通信アーキテクチャ — NVLink、InfiniBand、PCIe

分散LLMの学習と推論では、GPU、ノード、ストレージ間で膨大なデータを移動する必要があります。大規模学習では、この通信ファブリックがボトルネックになることがよくあります。

PCIe — ホストとデバイスを結ぶリンク

Important

PCIeの世代

世代x16帯域幅(各方向)双方向注記
PCIe Gen316 GB/s32 GB/s古いサーバーで一般的
PCIe Gen432 GB/s64 GB/sA100 PCIe、現在の大半のサーバー
PCIe Gen564 GB/s128 GB/sH100 PCIe、普及途上

PCIeは次の用途に使われます。

  • CPU \(\leftrightarrow\) GPU間のデータ転送(モデルの読み込み、CPUオフロード)

  • NVLinkが利用できない場合のノード間GPU通信(まれで非常に低速)

  • NVMeストレージへのアクセス(CPU経由)

Warning

PCIeはGPU間通信向けではない

NVLinkが利用できる場合は、GPU間通信をPCIe経由にしてはいけません。PCIeの帯域幅(32 GB/s)はNVLink 4(900 GB/s)より28\(\times\)低速です。NVLinkのないマルチGPUサーバー(一般向けGPUなど)ではGPU間帯域幅がPCIeに制限されるため、テンソル並列化は極めて遅くなります。

Important

NVLinkの世代

世代リンク数総帯域幅GPU
NVLink 26300 GB/sV100
NVLink 312600 GB/sA100
NVLink 418900 GB/sH100
NVLink 5181800 GB/sB200 (Blackwell)

NVLinkは同じノード上のGPU間を結ぶポイント・ツー・ポイントのインターコネクトです。各リンクは双方向です。H100 SXM5にはNVLink 4のリンクが18本あり、1本あたり双方向で50 GB/s、合計900 GB/sを提供します。

NVSwitch

DGX H100システムでは、8枚すべてのGPUがNVSwitchを介して接続されます。NVSwitchは完全二分帯域幅を提供する専用スイッチングチップです。これにより、リング内の隣接GPUだけでなく、任意のGPUが他の任意のGPUと、同時にNVLinkの最大速度で通信できます。

Tip

リングと完全二分帯域幅

リングトポロジ(8 GPU)では、AllReduceのデータはリングを一周する必要があります。各リンクは総データ量の\(\frac{2(N-1)}{N}\)を運ぶため、アルゴリズム上の帯域幅は\(B_{\text{link}} \times \frac{N}{2(N-1)}\)(\(N=8\)では約\(0.57 \times B_{\text{link}}\))です。NVSwitchの完全二分帯域幅では、AllReduceがツリーベースのアルゴリズムで全リンクを同時に利用でき、ピークに近い帯域幅を達成します。実際のDGX H100では、リングは約\(\sim\)700 GB/sのバス帯域幅、NVSwitchは約\(\sim\)900 GB/sを実現します。

InfiniBand — ノード間通信

ノード(サーバー)間の通信には、InfiniBandがGPUメモリへのダイレクトアクセスを備えた、高帯域幅・低レイテンシのネットワークを提供します。

Important

InfiniBand NDR

  • NDR 400Gb/s = ポートあたり50 GB/s(単方向)

  • HDR 200Gb/s = ポートあたり25 GB/s(前世代)

  • RDMA :Remote Direct Memory Access(リモート・ダイレクト・メモリアクセス)。GPUはリモートCPUを介さずに、リモートGPUのメモリを読み書きできます。

  • GPUDirect RDMA :データはCPUとシステムDRAMを完全にバイパスし、HBM \(\to\) NIC \(\to\) ネットワーク \(\to\) NIC \(\to\) HBMへ直接流れます。

  • レイテンシ :小さなメッセージでは約\(\sim\)1〜2 \(\mu\)s(TCP/IPでは約\(\sim\)100 \(\mu\)s)。

Fat-Treeトポロジ

大規模GPUクラスターではFat-Treeネットワークトポロジが使われます。\(k\)ポートのスイッチを使う3層Fat-Treeは、完全二分帯域幅で\(k^3/4\)ノードをサポートします。\(k=64\)ポートの400Gb/s NDRスイッチでは、\(64^3/4 = 65{,}536\)ノードです。

レール最適化トポロジ

実際のクラスターでは、各ノード内の各GPUを異なるトップ・オブ・ラックスイッチに接続するレール最適化トポロジが使われます。これにより、すべてのGPUが関係するAllReduce操作で、すべてのネットワークリンクを同時に使い、帯域幅を最大化できます。

分散LLM学習における通信パターン

分散学習は、集合通信プリミティブに依存します。どのプリミティブを選ぶかによって、必要な帯域幅とスケーリングの挙動が決まります。

Important

通信プリミティブ

プリミティブ用途通信量
AllReduce勾配同期(DDP、FSDP)\(2(N-1)/N \times\)パラメータサイズ
AllGatherシャード化された重みを収集(FSDP)\((N-1)/N \times\)パラメータサイズ
ReduceScatter勾配を分散(FSDP)\((N-1)/N \times\)パラメータサイズ
AllGatherテンソル並列化のアクティベーションアクティベーションサイズ
Point-to-Pointパイプライン並列化(send/recv)マイクロバッチのアクティベーション
Broadcast重みの同期(新しいワーカー)モデル全体のサイズ

Note

帯域幅の計算 — 70Bモデルの勾配AllReduce

設定 :70Bパラメータモデル、BF16勾配、8ノード\(\times\)8 GPU = 64 GPU。データ並列度 = 64。

勾配サイズ :\(70 \times 10^9 \times 2\)バイト\(= 140\) GB。

GPUあたりのAllReduce通信量 (リング):\(2 \times (64-1)/64 \times 140 \approx 275\) GB。

利用可能なノード間帯域幅 :ノードあたり8 GPU\(\times\) GPUあたり50 GB/s \(= 400\) GB/s(レール最適化トポロジで8個すべてのNICがアクティブ)。

AllReduce時間 :\(275 / 400 \approx 0.69\)秒/ステップ。

意味 :計算ステップが1秒の場合、通信によって0.69秒(ステップ時間全体の41%)が追加されます。これが、勾配圧縮、混合精度、FSDP(通信と計算をオーバーラップさせる)が重要な理由です。

ネットワークトポロジ図

次の図は、ノード内(NVLink)とノード間(InfiniBand)の通信経路を示す、典型的な2ノードGPUクラスターのトポロジを表しています。

Tip

帯域幅に基づく並列化の選択

  • テンソル並列化(TP) :各層でAllReduceが必要です。ノード内でNVLinkを使う場合に限って使用します。TP=8はH100 DGXノードで標準的です。

  • パイプライン並列化(PP) :ステージ間のポイント・ツー・ポイント通信です。ノードをまたげますが、パイプライン・バブルのオーバーヘッドが追加されます。TPだけではモデルを収められない場合に使います。

  • データ並列化(DP) :勾配のAllReduceです。IB経由でノードをまたげます。高速なIBがあれば適切にスケールします。

  • FSDP/ZeRO :AllGather+ReduceScatterです。DPと似ていますが、オプティマイザー状態をシャード化します。大規模モデルではDPより推奨されます。

vLLM — PagedAttentionと高スループット推論

vLLM (Kwon et al. 2023)は、PagedAttentionを導入しました。これは、オペレーティングシステムがRAMに使うページング抽象化を借用し、GPUのKVキャッシュに適用したものです。LLM推論では、過去のすべてのトークンのキー・テンソルとバリュー・テンソルを保存するKVキャッシュが、メモリ消費の大部分を占めます。これを効率的に管理することが、高スループット推論の中心的な課題です。

KVキャッシュの断片化問題

Important

KVキャッシュのメモリ公式

\(L\)層、\(H\)ヘッド、ヘッド次元\(d\)、\(n\)トークンのシーケンスを持つモデルの場合: \[ \text{KV cache size} = 2 \times L \times H \times d \times n \times \text{bytes_per_element} \] Llama-3 70B(BF16)の場合:\(L=80\)、\(H=8\)(GQA)、\(d=128\): \[ = 2 \times 80 \times 8 \times 128 \times n \times 2 = 327{,}680 \times n \text{ bytes} \] \(n=4096\)トークンの場合:シーケンスあたり\(\approx 1.3\) GB。

Tip

内部断片化と外部断片化

従来の推論システムは、各シーケンスのKVキャッシュ用に、最大シーケンス長に合わせた連続メモリ領域をあらかじめ確保します。これにより、2種類の無駄が生じます。

内部断片化 :500トークンしか生成しないシーケンスでも、4096トークン用に予約された領域を保持します。未使用の3596トークン分が無駄になります。

外部断片化 :多数のシーケンスが完了すると、空きメモリは小さく不連続な隙間に分かれます。空き容量の合計が十分でも、単一の連続領域が十分に大きくないため、新しい長いシーケンスを割り当てられません。

実際には、単純な割り当てでのGPUメモリ利用率は20〜40%にすぎないことがよくあります。

PagedAttention — KVキャッシュの仮想メモリ

PagedAttention(Kwon et al., 2023)は、オペレーティングシステムのページング抽象化を借用します。シーケンスごとに1つの連続領域を持つ代わりに、KVキャッシュを固定サイズの ページ (ブロック)に分割します。そしてCPUのページテーブルに相当する間接参照テーブルが、各シーケンスの論理的なトークン位置を、GPUメモリ上に分散した物理アドレスへ変換します。

Important

PagedAttentionの中核概念

  • ブロックサイズ :通常は1ブロックあたり16トークン(調整可能)。各ブロックには\(16 \times 2 \times L \times H \times d\)個の要素が格納されます。

  • ブロックテーブル :シーケンスごとに、論理ブロックのインデックスをGPUメモリプール内の物理ブロックのインデックスへ対応付けます。

  • 物理ブロックプール :固定サイズのブロックを事前に確保したプール。割り当ては\(O(1)\)で、空きリストから取り出すだけです。

  • アテンションカーネル :アテンション計算中にブロックテーブルを使い、不連続な物理位置からKVブロックを収集するよう変更されています。

Note

ブロックテーブルの例

ブロックサイズを4トークンとし、2つのシーケンスがあるとします。

  • シーケンスA(7トークン):論理ブロック [0,1] \(\to\) 物理ブロック [3, 7]

  • シーケンスB(5トークン):論理ブロック [0,1] \(\to\) 物理ブロック [1, 5]

物理ブロック3にはシーケンスAのトークン0〜3が格納されます。物理ブロック7にはシーケンスAのトークン4〜6が格納されます(部分的に使用)。シーケンスAのアテンションカーネルは、間接参照層としてブロックテーブルを使い、物理ブロック3、7の順に読み出します。

PagedAttentionの利点

ほぼゼロの無駄

内部断片化は、シーケンスごとに最大でも部分的に埋まった最後のブロック1つ分に抑えられます。ブロックサイズが16なら、最悪の場合の無駄はシーケンスあたり15トークンであり、無視できます。ブロックは固定サイズで交換可能なので、外部断片化はなくなります。

動的割り当て

ブロックはシーケンスの成長に応じてオンデマンドで割り当てられます。最終的なシーケンス長を事前に知る必要はありません。出力長が未知である生成処理では、これは非常に重要です。

プレフィックス共有(コピーオンライト)

共通のプレフィックス(システムプロンプトなど)を持つ複数のシーケンスは、そのプレフィックスに対して 同じ物理ブロック を共有できます。ブロックテーブルは、複数のシーケンスを同じ物理ブロックへ指し示すだけです。シーケンスが共有ブロックへの書き込みを必要とするとき(プレフィックスから分岐するとき)、コピーオンライトが発生します。

Tip

プレフィックス共有による節約

1000トークンのシステムプロンプトを128人の同時ユーザーに提供するチャットボットを考えます。

  • プレフィックス共有なし:システムプロンプトのKVキャッシュだけで\(128 \times 1000 \times 327{,}680 / 10^9 \approx 42\) GB

  • プレフィックス共有あり:\(1 \times 1000 \times 327{,}680 / 10^9 \approx 0.33\) GB

  • 節約量:共有プレフィックス部分で約\(\sim\)128\(\times\)

スワップによるプリエンプション

GPUメモリが枯渇すると、vLLMはKVブロックをCPU DRAMへスワップすることで、シーケンスをプリエンプトできます(またはブロックを単純に破棄し、後で再計算できます)。ブロックが自己完結していて不連続であるからこそ可能な方法です。連続領域の割り当てをスワップするには、バッファ全体をコピーしなければなりません。

継続的バッチ処理

従来のバッチ処理(「静的バッチ処理」)は、すべてのシーケンスが完了するまで新しいシーケンスを開始しません。あるシーケンスが500トークン、別のシーケンスが10トークンを生成する場合、短いシーケンスについてGPUは490ステップの間アイドルになります。これは非常に非効率です。

Important

継続的バッチ処理

継続的バッチ処理(イテレーションレベル・スケジューリングとも呼ばれます)は、一度に1つのデコードステップを処理します。各ステップの後に次を行います。

  1. 生成されたEOSトークンを確認し、完了したシーケンスを調べる

  2. 完了したシーケンスをバッチから削除し、KVブロックを解放する

  3. 待機中の新しいシーケンスを追加し、空いたスロットを埋める

  4. 更新されたバッチで次のデコードステップを実行する

バッチの構成は各ステップで変わり、シーケンスが動的に参加・離脱します。これによりGPU利用率をほぼ100%に保ち、スループットを大幅に改善できます(静的バッチ処理に比べて1.5〜3\(\times\))。ここではPagedAttentionが不可欠です。バッチ途中でのシーケンス追加・削除にはKVブロックの動的な割り当て・解放が必要であり、ページ化メモリでのみ効率的に実行できるからです。

vLLMにおける投機的デコーディング

投機的デコーディングは、小さなドラフトモデル(例:1Bパラメータ)を使って\(k\)個の候補トークンを素早く提案し、大きなターゲットモデルが1回のフォワードパスで検証します。最初に拒否されるまでのすべてのトークンが受理されます(検証ステップあたりの期待受理数は3〜5トークン)。これにより、品質を損なわず、レイテンシ重視の単一シーケンス生成を2〜3\(\times\)高速化できます。

vLLMは投機的デコーディングをPagedAttentionと統合しています。

  • ドラフトトークンには投機的KVブロックが割り当てられる

  • 拒否された場合、投機的ブロックは解放される(ページ化割り当てなら低コスト)

  • 受理された場合、投機的ブロックはメインシーケンスへ昇格される

  • ブロックテーブルの更新は\(O(k)\)で、数個のテーブルエントリを更新するだけ

大規模環境での具体的なメモリ節約 — 70Bモデル

Note

メモリ予算 — 70B BF16推論

設定 :Llama-3 70B、BF16、単一のA100 80GBノード(8 GPU、テンソル並列化)。

モデルの重み :\(70 \times 10^9 \times 2\)バイト\(= 140\) GB \(\div\) 8 GPU \(= 17.5\) GB/GPU。

KVキャッシュ用の残り :\(80 - 17.5 - 3\)(オーバーヘッド)\(= 59.5\) GB/GPU。

GPUあたりのトークンごとのKVキャッシュ (TP=8で、各GPUはヘッドの\(1/8\)を保持):\(2 \times 80 \times 1 \times 128 \times 2 = 40{,}960\)バイト\(\approx 40\) KB/トークン。

KVキャッシュ内の最大トークン数 :\(59.5 \times 10^9 / 40{,}960 \approx 1.45\)百万トークン。

4096トークンのシーケンスを128本同時実行する場合 :\(128 \times 4096 = 524{,}288\)トークンで、予算内に十分収まります。

PagedAttentionなし (各シーケンスに最大長4096を事前割り当て):計算は同じですが、平均で約\(\sim\)50%が断片化によって無駄になります \(\to\) 収容できるのは64シーケンスだけです。

Warning

ブロックサイズのトレードオフ

ブロックサイズを大きくすると、ブロックテーブルのオーバーヘッドが減り、メモリアクセスの局所性が改善します(分散した読み出しが減ります)。一方、ブロックサイズを小さくすると、内部断片化が減り、より細粒度のプレフィックス共有が可能になります。vLLMのデフォルトは16トークン/ブロックで、バランスのよい設定です。非常に長いシーケンス(10万トークン以上)では、より大きなブロック(32〜64)が適している場合があります。

vLLM:エンドツーエンドシステム

vLLMは、PagedAttentionを完全なサービングスタックの内部に組み込みます。継続的バッチ処理、プレフィックスキャッシュ、投機的デコーディング、テンソル並列モデルのシャーディングが連携し、GPU 1ドルあたりのスループットを最大化します。

アーキテクチャの概要

中核コンポーネント

  • API Server :OpenAI互換のリクエスト(completion、chat)を受け付けます。入力をトークン化し、ビームサーチや複数サンプルのための「シーケンスグループ」を作成します。

  • Scheduler :vLLMの司令塔です。次の3つのキューを管理します。

    • waiting:まだ開始されていない新しいリクエスト(プリフィル待ち)

    • running:トークンをアクティブに生成しているリクエスト(デコード段階)

    • swapped:KVキャッシュがCPUへオフロードされたプリエンプト済みリクエスト

    各イテレーションで、Schedulerは利用可能なGPUメモリブロックに基づき、どのリクエストを実行するか決定します。

  • Block Manager :KVキャッシュの仮想メモリ抽象化を実装します。論理ブロック(シーケンスごと)を物理ブロック(GPUメモリプール内)へ対応付けます。次を処理します。

    • 割り当て(新しいトークンが生成される \(\rightarrow\) 新しいブロックが必要)

    • コピーオンライト(ビームサーチ用:複数のビームがプレフィックスブロックを共有し、分岐時だけコピー)

    • スワップ(プリエンプト/再開時のGPU \(\leftrightarrow\) CPU移行)

    • プレフィックスキャッシュ(プロンプトが共通プレフィックスを持つとき、キャッシュ済みブロックを再利用)

  • Model Executor :実際のLLMフォワードパスを実行します。GPU間のテンソル並列化を管理し、ページ化KVキャッシュブロックを読み取るアテンションカーネルをディスパッチします。

  • KV Cache Pool :固定サイズのブロックに分割された事前割り当てGPUメモリ(デフォルト:16トークン \(\times\) num_heads \(\times\) head_dim \(\times\) 2バイト/ブロック)。実行時に動的割り当てを行わない \(\rightarrow\) 断片化はゼロです。

リクエストのライフサイクル(エンドツーエンドの流れ)

  1. 到着 :クライアントがプロンプトを送信します。APIサーバーはトークン化し、SequenceGroupを作成してwaitingキューに入れます。

  2. スケジューリング :各ステップでSchedulerは次を実行します。

    1. swappedのシーケンスを再開できるか確認する(空きブロックが十分か)。

    2. waitingのシーケンスがプリフィルを開始できるか確認する(プロンプト全体に十分なブロックがあるか)。

    3. runningのシーケンスに残りのブロックを割り当てる(現在のブロックが満杯なら、ステップごとにシーケンスあたり1つの新しいブロックが必要)。

    4. 予算を超える場合、優先度の低いrunningのシーケンスをプリエンプトする(KVをCPUへスワップするか、後で再計算する)。

  3. プリフィル (リクエストの最初のイテレーション):プロンプト全体を1回のフォワードパスで処理します。すべてのプロンプトトークンのKVキャッシュを計算し、割り当てられたブロックに格納します。これは計算律速です(大量のトークンをまとめて処理するため)。

  4. デコード (後続のイテレーション):各シーケンスについて、ステップごとに新しいトークンを1つ生成します。実行中の全シーケンスをまとめてバッチ化します(継続的バッチ処理)。これはメモリ律速です(KVキャッシュ全体を読み出し、1トークンを生成するため)。

  5. ブロック割り当て :各デコードステップの後、シーケンスの最後のブロックが満杯なら、Block Managerが新しい物理ブロックを割り当て、次の論理ブロックへ対応付けます。

  6. 完了 :シーケンスがEOSまたは最大長に達すると、runningから削除されます。その物理ブロックは直ちに解放されます \(\rightarrow\) 他のシーケンスが利用できるようになります。レスポンスはクライアントへストリーミングされます。

プレフィックスキャッシュ(自動プロンプトキャッシュ)

複数のリクエストが共通プレフィックス(システムプロンプト、few-shot例など)を持つ場合:

  1. 各論理ブロックのトークン内容をハッシュする。

  2. 新しいリクエストの到着時に、プレフィックスブロックがすでにキャッシュにあるか確認する。

  3. ヒットした場合、そのトークンのプリフィルをスキップし、物理KVブロックを直接再利用する。Time-to-first-tokenは大幅に短縮される。

  4. 追い出しにはLRUポリシーを使う。メモリ圧迫が必要になるまで、キャッシュ済みブロックは解放しない。

効果 :長いシステムプロンプト(全ユーザーで共有される2000トークン以上)を持つチャットアプリケーションでは、プレフィックスキャッシュによってTTFTが60〜80%短縮されます。

vLLMにおけるガイド付き(制約付き)デコーディング

vLLMは、プラグイン可能なバックエンドを通じて制約付きデコーディング(セクション1.12.11)をネイティブにサポートしています。これにより、性能オーバーヘッドを最小限に抑えながら、サービング時に 構造化出力を保証 できます。

サポートされる制約の種類

OpenAI互換APIは、guided_*パラメータまたはresponse_formatフィールドを通じて制約を受け付けます。

from openai import OpenAI
client = OpenAI(base_url="http://localhost:8000/v1")

# --- JSON Schema制約 ---
response = client.chat.completions.create(
    model="meta-llama/Llama-3-70B-Instruct",
    messages=[{"role": "user",
               "content": "Extract: name, age, city from: "
                          "'John is 30 and lives in NYC'"}],
    extra_body={
        "guided_json": {
            "type": "object",
            "properties": {
                "name": {"type": "string"},
                "age": {"type": "integer"},
                "city": {"type": "string"}
            },
            "required": ["name", "age", "city"]
        }
    }
)
# スキーマに適合する有効なJSONが保証される

# --- 正規表現制約 ---
response = client.completions.create(
    model="meta-llama/Llama-3-70B-Instruct",
    prompt="Generate an IPv4 address: ",
    extra_body={
        "guided_regex": r"\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}"
    }
)

# --- 選択肢制約 ---
response = client.completions.create(
    model="meta-llama/Llama-3-70B-Instruct",
    prompt="Sentiment: ",
    extra_body={"guided_choice": ["positive", "negative", "neutral"]}
)

バックエンドアーキテクチャ

vLLMはマスク計算をバックエンドエンジンへ委譲します。

  • XGrammar (v0.7以降のデフォルト):JSONスキーマ、正規表現、任意のEBNF文法をサポートするプッシュダウンオートマトンエンジン。効率的なC++コアにより、複雑なスキーマでは最速です。

  • Outlines (Willard and Louf 2023):FSMベースで、JSONと正規表現をサポートします。XGrammarが利用できない場合のフォールバックとして使われます。

マスクは、モデルのフォワードパスがロジットを生成した後、サンプリングの前に適用されます。FSM/PDAの状態遷移と事前計算済みインデックスのルックアップは\(O(1)\)なので、実際のオーバーヘッドは1ステップあたり\(<\)1 msです。

性能への影響

制約はロジットをマスクするだけで、アテンションやFFNの再計算を行わないため、スループットの低下は無視できます(ベンチマークで\(<\)2%)。主なコストは、スキーマをFSM/PDAインデックスへコンパイルすることです。スキーマの複雑さに応じて0.5〜5秒かかります。vLLMはリクエスト間でコンパイル済みスキーマをキャッシュするため、このコストは一意なスキーマごとに1回だけ発生します。

Warning

構造化出力 $\neq$ 正しい出力

制約付きデコーディングは、出力が構文的に有効であること(JSONとして解析でき、スキーマの型に一致すること)を保証します。しかし、意味的な正しさを保証するものではありません。モデルは、正しく解析できるものの事実としては誤った値を、依然としてハルシネーションする可能性があります。下流で必ずビジネスロジックを検証してください。

指標vLLMHF Generate理由
スループット(tok/s)2,500〜4,000300〜600継続的バッチ処理+PagedAttention
メモリ利用率90〜95%50〜60%断片化ゼロ、動的ブロック割り当て
最大同時実行シーケンス数200〜50016〜32ページ化KVによりシーケンスごとの予約が不要
Time-to-first-token100〜300 ms500〜2000 ms繰り返し使うシステムプロンプトのプレフィックスキャッシュ

vLLMと代替手法の性能比較(70Bモデル、A100 \(\times\) 4、TP=4)。