AIチップ・アーキテクチャ

AIチップ・アーキテクチャ 2018年の国際コンピュータアーキテクチャシンポジウムで、ジョン・ヘネシーとデイビッド・パターソンは、コンピュータアーキテクチャの新しい黄金時代についてテュアリング講演を行いました。ヘネシーとパターソンがテュアリング賞を受賞した研究を行った1980年代には、シングルスレッドCPUの性能は年間52%成長していましたが、ムーアの法則とデナード・スケーリングの終焉により、2018年にはその成長率は3%になっていました。ドメイン固有アーキテクチャ（DSA）の必要性が生じたのです。彼らの実例は、すでに実用化されていたGoogleのTPU v1で、ニューラルネットワーク推論においてCPUの29倍のスループットと80倍のエネルギー効率を実現しました。彼らの最後の予測は、「次の10年間で、新しいコンピュータアーキテクチャのカンブリア爆発が見られるだろう」というものでした。この予測は現実のものとなりました。現在、数十ものアーキテクチャが真剣に開発されています。GPU、TPU、LPU、NPU、DPU、ASIC、ウェハー・スケール・エンジン、再構成可能なデータフロー、ニューロモルフィック、フォトニック、アナログなどです。特に、これらのアーキテクチャはAIのための計算に焦点を当てています。これまでに実際に展開に成功しているアーキテクチャは、GPU（NVIDIA、AMD）、 σύστημα συστολικής διαστολής (systolic-array) アクセラレータ（TPU、Trainium）、Cerebras Wafer-Scale Engine、Groq LPUです。NVIDIAが明確なフロントランナーであり、AMDがそれに続き、両社ともOpenAIとMetaから6ギガワットのコミットメントを得ています。TPUはGeminiをトレーニングし、最大100万チップでAnthropicにサービスを提供します。Anthropicはまた、100万チップ以上のTrainiumチップでClaudeを実行しています。Cerebrasは現在、OpenAIの推論にサービスを提供しており、Groq LPUは200億ドルの買収合併（acquihire）を通じてNVIDIAに統合されました。この記事では、これらの様々なアプローチについて、その哲学、アーキテクチャ、スケーリング方法（スケールアップとスケールアウト）、およびソフトウェアスタック（チップのプログラミング方法）を調査することを目的としています。 問題 AI計算は行列乗算に支配されています。Transformerは、一連の行列乗算（Q/K/V射影、アテンション、出力射影、FFN）であり、要素ごとの演算（正規化、活性化、残差加算）と交互に行われます。最先端モデルのトレーニングでは、10^25回の乗算・累積演算（行列乗算は乗算・累積演算の連なり）が行われます。これらの行列乗算の形状はワークロードによって異なります。トレーニングでは、バッチ内のシーケンスがすべてのレイヤーを順方向に通過し、損失が逆伝播され、重みが更新されます。数千ものトークンが同じ重み行列を一度に流れます。プリフィル（Prefill）は、推論のプロンプト取り込みフェーズです。最初の出力トークンが生成される前に、モデル全体を通して入力シーケンス全体を1回のパスで射影します。トレーニングとプリフィルの両方で、同じ重み行列に対して多くのトークンが積み重ねられるため、各レイヤーの計算は大きな行列と行列の乗算（GEMM）となり、高い算術強度（計算バウンド）を持ちます。デコード（Decode）は自己回帰的です。モデルは一度に1つのトークンを出力し、各トークンはそれ以前のすべてのトークンに条件付けられており、トークンN+1はトークンNが生成されるまで開始できません。1回のステップで生成されるトークンは1つだけなので、すべての行列乗算は行列とベクトルの乗算（GEMV）になります。1つのトークンを生成するには、モデル内のすべての重みを完全に読み取る必要があり、さらにアテンションのためのKVキャッシュ全体を読み取る必要があります。算術強度は、プリフィルと比較して桁違いに低下します。推論システムは、トークンをバッチ処理することで、これらのGEMVをGEMMに戻し、ある程度の強度を回復します。連続バッチ処理（continuous batching）は、多くのユーザーのデコードステップを積み重ね、投機的デコード（speculative decoding）は、リクエストごとにK個の下書きトークンを積み重ねて1回のパスで検証し、マルチトークン予測は、モデル自体の中で同じトリックを折り込みます。これにより、行列乗算ユニットの利用率が向上し、Ops/Bが押し上げられます。連続バッチ処理の場合、各ユーザーのリクエストは依然として独自のKVキャッシュを読み取るため、長コンテキストデコードは、重み帯域幅バウンドからKV帯域幅バウンドに移行します。ここでのアーキテクチャの問題は、数値を計算が発生する場所に十分に速く移動させることです。これはメモリの壁として知られています。計算は指数関数的にスケールしましたが、メモリ帯域幅はそうではありませんでした。各アーキテクチャは、データ移動ゲームに勝つための異なる戦略を提案しています。チップの理解は、4つの質問に集約されます。データはどこに存在するか、計算ユニットにどのように移動するか、計算ユニットはどのように見えるか、そしてチップはスケールでどのように互いに通信するかです。 NVIDIA GPU NVIDIA GPUは、大規模並列プロセッサです。その哲学は、ホストCPUによってオーケストレーションされ、CUDAを通じて公開される数千のスレッドを持つプログラム可能なチップが、並列化可能なワークロードを実行するための適切なマシンであるということです。各世代は、プログラミングモデルを変更することなく、プログラム可能なストリーミングマルチプロセッサ（SM）にアクセラレーションプリミティブを追加します。同じチップがTransformerをトレーニングし、推論を提供し、グラフィックスをレンダリングし、科学シミュレーションを実行します（アクセラレーテッド・コンピューティング）。 系譜 2006 Tesla G80: CUDA対応GPUの最初のもので、ユニファイドシェーダーとSIMT実行モデルを備えています。 2010 Fermi GF100: 真のコンピューティングアーキテクチャの最初のもので、ユニファイドL1/L2キャッシュ、デュアルワープスケジューラ、IEEE-754 FP64を備えています。 2012 Kepler K20, K40: SMX、ダイナミックパラレリズム、Hyper-Q。GPUは独自のワークを起動できます。 2014 Maxwell M40: Keplerの約2倍のワットあたりの性能を持つ再設計されたSM。 2016 Pascal P100: NVLink 1.0、HBM2、ネイティブFP16スループット。ディープラーニング専用に設計された最初のGPU。 2017 Volta V100: 最初のテンソルコア。独立したスレッドスケジューリング。 2018 Turing T4: INT8/INT4を備えた第2世代テンソルコア。最初のRTコア。 2020 Ampere A100: TF32と構造化スパース性を持つ第3世代テンソルコア。マルチインスタンスGPUパーティショニング。 2022 Hopper H100, H200, GH200: 第4世代テンソルコア、FP8、Transformer Engine。HBM3、TMA、スレッドブロッククラスター、非同期wgmma。 2024 Blackwell B100, B200, GB200: FP4、テンソルメモリ（TMEM）を備えた第5世代テンソルコア。2ダイチップレットGPU、NVLink 5。 2025 Blackwell Ultra B300, GB300: ミッドサイクルリフレッシュ。FP4スループットが約1.5倍、288GB HBM3e。長コンテキスト推論向けに調整。 2026 Rubin Rubin, VR200, Rubin CPX: HBM4、第3世代Transformer Engine、Vera CPUペアリング、Rubin CPXによる分散プリフィル。 2027 Rubin Ultra Rubin Ultra: 4ダイGPUパッケージ、パッケージあたり1TB HBM4e。600kW NVL576 Kyberラックに100ペタフロップスFP4/GPUで展開。 アーキテクチャ NVIDIA GPUは、スループット指向のコアのグループ、それらにデータを供給するための深いメモリ階層、そして数千ものスレッドを飛行状態に保つのに十分なスケジューリングシリコンで構成されています。コアはストリーミングマルチプロセッサ（SM）であり、パッケージあたり100個以上複製されています。V100では80個、A100では108個、H100では132個、B200では148個、B300では160個、Rubinでは224個です。各SMの内部には同じレシピがあります。4つのSMサブパーティションがあり、それぞれが独自のワープスケジューラ、ディスパッチユニット、16k×32ビットレジスタファイル、スカラーCUDAコアレーン、超越関数用の特殊関数ユニット、そしてSMのテンソルコアへのプライベートポートを備えています。4つのパーティションは、L1/共有メモリブロックとTMAを共有します。スレッドは32個のワープにグループ化され、SIMTロックステップで実行されます。パーティションあたりの数十個の常駐ワープにより、スケジューラはそれらを切り替えることでメモリ/算術遅延を隠すことができます。 計算 CUDAコアは元の計算スループットであり、AIにおいては、行列乗算以外のすべて（活性化、残差加算、正規化、アドレス計算）を担当します。しかし、Transformerブロックは約99%が行列乗算のFLOPであるため、圧倒的な計算スループットはテンソルコアから得られます。これらのコアは、小さな行列タイルの上で融合された行列乗算・累積（MMA）を実行します。D = A・B + C 完全な行列乗算は出力タイルのセットに分解されます。1つの出力タイルを生成するために、カーネルは左入力行列の行ストリップからAAAを、右入力行列の列ストリップからBBBを抽出し、内部次元KKKを歩き回り、各部分積を実行中のアキュムレータに折り込みます。CCCはこれまでの部分和を保持し、DDDは次のステップに引き継がれる更新された値です。内部ループが完了した後、DDDは完全な出力行列の1つの完成したタイルになります。行列乗算全体は、これらの多くのタイルMMAから構築されます。タイル形状はM×N×Kと記述され、M×Nは出力タイルサイズ、KKKは命令が1回の実行で収縮する内部次元の量です。行列乗算のKKK軸の残りは、カーネルの内部ループによって歩行されます。アキュムレータはループ全体でスティッキーです。各MMAの出力DDDは次のMMAの入力CCCとなり、したがって方程式は実際には C ← A・B + C となります。