はじめに
ここ2本、NVIDIA の Vera Rubin プラットフォームと、その7番目のチップである Groq 3 LPU / LPX ラックについて書きました。
書き終えてから、自分の理解に穴があることに気づきました。両方の記事で私は「attention は KV キャッシュという大容量データを読む処理なので HBM を持つ GPU 側、FFN は重みのストリーミング読み出しなので SRAM の帯域が活きる LPU 側」と書いています。NVIDIA の資料に書いてある通りの分担ですし、字面としては筋が通っています。しかし、なぜそう分けられるのかを、演算の中身から自分で導けたわけではありませんでした。attention だって重みを読むはずですし、FFN だって行列積です。両方とも帯域律速だと私は書いてきたのに、なぜ片方だけが SRAM に移せるのか。
これは要するに、Transformer の演算そのものを手で追ったことがない、ということです。ハードウェアの分担の話は、突き詰めれば「どの演算が、何バイト読んで、何 FLOP 使うか」という算数に還元されるはずで、そこを飛ばして LPU の看板だけ眺めても、腑に落ちないのは当然でした。
そこで本記事では、Transformer が1トークン生成するあいだに何をしているのかを、順番に一つずつ追います。そのうえで、各演算の FLOPs とメモリアクセス量を数え、ルーフラインのバランス点という物差しに載せます。結論を先に書いてしまうと、
- attention と FFN の違いは「行列積かどうか」ではなく、読むデータがユーザー間で共有できるかどうかにある
- FFN の重みは全ユーザー共有の定数なので、バッチで償却できる。attention が読む KV キャッシュはユーザー固有なので、償却できない
- さらに FFN は演算の形が文脈長に依存しないため、静的スケジューリングと相性が良い。attention は文脈長で形が変わる
という3点が、あの分担のすべてでした。attention と FFN の非対称性は、Transformer の数式の中に最初から書いてあったわけです。
なお本記事は LLM の入門記事ではなく、半導体側の人間が演算量とメモリアクセスの観点から Transformer を読み直したメモです。学習ではなく推論、それも自己回帰生成の話に絞ります。数値はすべて後述する題材モデルでの机上計算で、実測ではありません。実装は Flash Attention をはじめ多数の最適化が入るため、実機の挙動はここまで単純ではない点はあらかじめお断りしておきます。
題材にするモデル
具体的な数字がないと算数ができないので、Llama 3 8B 相当の構成を題材にします。以下、この数字を使い回します。
| 記号 | 意味 | 値 |
|---|---|---|
| d_model | 隠れ次元（1トークンを表すベクトルの長さ） | 4,096 |
| n_layers | Transformer ブロックの段数 | 32 |
| n_head | クエリのヘッド数 | 32 |
| d_head | 1ヘッドあたりの次元 | 128 |
| n_kv | Key/Value のヘッド数（GQA） | 8 |
| d_ff | FFN の中間次元 | 14,336 |
| V | 語彙数 | 128,256 |
| L | 文脈長（これまでのトークン数） | 可変 |
d_model 4,096 とは、1個のトークンが 4,096 個の数値からなるベクトルとして表現される、という意味です。このベクトルを本記事では隠れ状態（hidden state）と呼びます。推論中、1トークンぶんの隠れ状態は bf16 なら 4,096 × 2 = 8,192 バイト、つまり 8KB です。この 8KB という数字は後半で効いてくるので、頭の片隅に置いておいてください。
全体の流れ
まず俯瞰します。トークン ID が入って、次のトークンの確率分布が出てくるまでの流れです。
Transformer ブロック1段の中身は次の通りです。
Attention と FFN が交互に並び、それぞれを残差接続（＋）が迂回する。この2段構成が32回繰り返されるだけの、拍子抜けするほど規則的な構造です。ここで押さえておきたいのは、ブロックの入口も出口も同じ 4,096 次元だという点です。32段を通しても、1トークンを表すベクトルの長さは最初から最後まで変わりません。層を重ねるとは、この 4,096 次元のベクトルを少しずつ書き換えていく作業のことです。
演算を1つずつ追う
1. 埋め込み
トークン ID（0 から 128,255 の整数）を使って、128,256 行 × 4,096 列のテーブルから1行を引いてきます。行列積ではなく、ただのテーブル参照です。演算量はゼロ、読むバイト数も 4,096 × 2 = 8KB だけです。
テーブル自体は 128,256 × 4,096 × 2 バイト ≈ 1.05GB あってメモリ上では大きいのですが、1トークンにつき1行しか触らないので、帯域の話としてはほぼ無視できます。パラメータ数が多いことと、メモリ帯域を食うことは別だ、という最初の例になっています。
2. RMSNorm
隠れ状態を、二乗平均平方根で割って正規化し、学習された重み（4,096 個）を要素ごとに掛けます。
\mathrm{RMSNorm}(x) = \frac{x}{\sqrt{\frac{1}{d}\sum_{i=1}^{d} x_i^2 + \epsilon}} \odot g
演算量は 4,096 要素ぶんの積和が数回、パラメータも 4,096 個です。行列積が数千万〜数億 FLOP を要求するのに対して、こちらは4桁くらい小さい世界です。数値的な安定性のために必須ですが、性能の議論にはほとんど登場しません。以降、正規化と残差加算は「無視できるほど軽い」ものとして扱います。
3. Attention
ここが本記事の主役の片方です。5つの段階に分けて追います。
3-1. QKV 射影
正規化した隠れ状態 x（4,096 次元）に、3つの重み行列を掛けて Query、Key、Value を作ります。
q = W_Q x,\quad k = W_K x,\quad v = W_V x
行列のサイズは以下の通りです。GQA（Grouped Query Attention）を使っているので、Key と Value のヘッド数はクエリの4分の1しかありません。
| 射影 | 行列サイズ | 出力次元 | パラメータ数 |
|---|---|---|---|
| W_Q | 4,096 × 4,096 | 4,096（32ヘッド × 128） | 16.8M |
| W_K | 4,096 × 1,024 | 1,024（8ヘッド × 128） | 4.2M |
| W_V | 4,096 × 1,024 | 1,024（8ヘッド × 128） | 4.2M |
この段階は純粋な行列積で、しかも重み行列は全ユーザー共通の定数です。性質としては後述する FFN とまったく同じで、attention という名前がついているだけの GEMM だと思ってかまいません。
3-2. RoPE（位置エンコーディング）
q と k に、トークン位置に応じた回転を掛けます。次元を2つずつペアにして複素数とみなし、位置 m に比例した角度だけ回す操作です。
\tilde{q}_m = R_m q_m,\quad \tilde{k}_n = R_n k_n
回転行列は学習パラメータではなく、位置から決まる定数です。内積 q·k を取ったときに角度差 (m − n) だけが残るので、相対位置の情報が自然に入ります。演算量は軽微ですが、KV キャッシュに保存されるのは回転を掛けた後の k だという点は実装上重要です。
3-3. スコア計算
いま作った q（現在のトークン）と、過去すべてのトークンの k との内積を取ります。ここで初めて、過去のデータを読みに行きます。
s_n = \frac{\tilde{q} \cdot \tilde{k}_n}{\sqrt{d_{head}}} \quad (n = 1, \dots, L)
L 個の内積を、32個のクエリヘッドそれぞれについて計算します。GQA なので、4つのクエリヘッドが1組の KV ヘッドを共有します。
3-4. softmax
L 個のスコアを確率分布に変換します。
a_n = \frac{\exp(s_n)}{\sum_{j=1}^{L} \exp(s_j)}
非線形かつ全体の総和が必要なので、L 方向のリダクションが入ります。オンライン softmax（Flash Attention の中核）は、この総和を1パスで済ませることで、L × L のスコア行列を一度もメモリに書き出さずに済ませる工夫です。なお、この恩恵が大きいのは L × L の行列が実際に立つ prefill や長系列の学習で、decode ではスコアが L × 1 のベクトルにしかならないため、書き出し削減の効果は限定的です（decode 側で効いてくるのは、後述する KV の読み出し量そのものです）。
3-5. 加重和と出力射影
得られた重み a で、過去すべての v を加重平均します。
o = \sum_{n=1}^{L} a_n v_n
これを32ヘッド分つなげて 4,096 次元に戻し、最後に出力射影 W_O（4,096 × 4,096、16.8M パラメータ）を掛けます。
attention の本質は 3-3 と 3-5、つまり「過去の全 k と内積を取り、全 v を加重平均する」ことです。この2ステップだけが、モデルの重みではなく、そのユーザーのその会話に固有のデータを読みます。ここが後の議論のすべての起点になります。
図の青が重みを読む処理、赤が KV キャッシュを読む処理です。1トークン生成するたびに、新しい k と v を1組追記する一方で、過去すべての k と v を読み直しています。
4. KV キャッシュはどこで生まれるか
3-1 で作った k と v は、次のトークンを生成するときにも、その次のトークンのときにも、まったく同じ値が必要になります。位置 n のトークンの k と v は、位置 n の隠れ状態だけから決まり、後から来るトークンの影響を受けないからです。
したがって毎回作り直すのは無駄なので、一度作ったら取っておく。これが KV キャッシュです。裏を返せば、KV キャッシュとは「再計算を避けるために、演算をメモリと交換した結果」であり、Transformer の数式に元から存在する概念ではありません。
題材モデルでの1トークンあたりのサイズを数えます。
| 項目 | 計算 | 値 |
|---|---|---|
| 1層・1トークンの K | 8ヘッド × 128次元 × 2バイト | 2KB |
| 1層・1トークンの K と V | 上の2倍 | 4KB |
| 全32層・1トークン | 4KB × 32 | 128KB |
| 文脈32,768トークン分 | 128KB × 32,768 | 約 4.0GB |
8B のモデルの重みが FP8 で約 8GB であることを思うと、32K の文脈を持つユーザーが1人いるだけで、その半分に相当する KV が積み上がることになります。しかもこれはユーザーごとです。前々回の記事で「モデル重みは全ユーザー共有の定数で済むのに対し、KV キャッシュはユーザーごとに固有の状態である」と書いたのは、この構造のことでした。
なお GQA を使わない素の MHA（KV ヘッドも32個）なら、この 4.0GB が 16GB になります。GQA/MQA、そして DeepSeek の MLA といった一連の工夫は、すべてこの KV の量を削る方向の設計です。
5. FFN
もう一方の主役です。SwiGLU 型の FFN は3つの行列を使います。
\mathrm{FFN}(x) = W_{down}\big(\mathrm{SiLU}(W_{gate} x) \odot W_{up} x\big)
gate と up で 4,096 次元を 14,336 次元に持ち上げ、片方を非線形関数に通してもう片方と要素ごとに掛け合わせ、down で 4,096 次元に戻します。
| 行列 | サイズ | パラメータ数 |
|---|---|---|
| W_gate | 4,096 × 14,336 | 58.7M |
| W_up | 4,096 × 14,336 | 58.7M |
| W_down | 14,336 × 4,096 | 58.7M |
3つ合わせて 176.2M。1層の attention 側（QKVO 合計 41.9M）の4倍以上です。
FFN の演算はこれだけです。過去のトークンを一切参照せず、現在のトークンの 4,096 次元ベクトルを、固定の重み行列に通すだけ。トークン間の情報のやりとりは attention が全部やっているので、FFN は各トークンを独立に非線形変換する係になっています。この「過去を見ない」という性質が、後で決定的に効いてきます。
6. 出力層
32段を抜けたら、最後の RMSNorm を通し、LM ヘッド（4,096 × 128,256、525M パラメータ）を掛けて語彙数ぶんのスコアを出し、softmax して確率分布にし、サンプリングして次のトークンを決めます。決まったトークン ID は、また埋め込みテーブルの入口に戻ります。これが自己回帰です。
7. パラメータの内訳
ここまでを合計します。
| 部位 | 1層あたり | 全体（32層） | 比率 |
|---|---|---|---|
| Attention（W_Q, W_K, W_V, W_O） | 41.9M | 1.34B | 16.7% |
| FFN（gate, up, down） | 176.2M | 5.64B | 70.2% |
| 埋め込み ＋ LM ヘッド | ─ | 1.05B | 13.1% |
| 合計 | 218.1M | 8.03B | 100% |
Transformer ブロックの中だけで見れば、パラメータの約8割が FFN です。埋め込みを含めた全体でも7割。私たちが「モデルの重み」と呼んでいるものの正体は、その大部分が FFN の3つの行列だということになります。
この時点で、LPX に FFN を持っていく意味が少し見えてきます。デコード時に読み出される重みの大半を SRAM 側に引っ越しさせられる、ということですから。
prefill と decode で何が変わるか
同じ数式でも、フェーズによって計算機から見た顔がまったく変わります。
prefill では、入力プロンプトの L 個のトークンをまとめて処理します。QKV 射影も FFN も、[L × 4,096] の行列に重み行列を掛ける GEMM になります。L が数千あれば、行列積として十分に大きく、演算器を埋められます。演算律速です。
decode では、1トークンずつしか進めません。処理するのは [1 × 4,096] のベクトル1本で、行列積というより行列ベクトル積（GEMV）です。重み行列は全部読むのに、掛ける相手はベクトル1本しかない。重みの各要素は1回しか使われず、演算器は待ちぼうけです。帯域律速です。
この非対称性は、ローカル LLM を触っていれば体感でわかります。長いプロンプトの読み込みは GPU が唸って一気に終わるのに、生成は淡々と一定速度で進む。前者が演算律速、後者が帯域律速です。
算数: 1トークンあたりの FLOPs とバイト数
decode で1トークン生成するとき、何 FLOP 必要で、何バイト読むのかを数えます。行列積の FLOPs は、パラメータ1個あたり積和2回なので「2 × パラメータ数」で近似できます。
重みに由来する分（文脈長 L に依存しない）:
| 項目 | FLOPs | 読むバイト数（FP8） |
|---|---|---|
| Attention 射影（32層） | 2.7 GFLOP | 1.3 GB |
| FFN（32層） | 11.3 GFLOP | 5.6 GB |
| LM ヘッド | 1.1 GFLOP | 0.5 GB |
| 小計 | 約 15 GFLOP | 約 7.5 GB |
KV キャッシュに由来する分（L に比例）:
| 項目 | 1トークンあたり |
|---|---|
| FLOPs | 16,384 × L × 32層 ≒ 0.52 MFLOP × L |
| 読むバイト数（bf16） | 4KB × L × 32層 ＝ 128KB × L |
L = 32,768 を入れると、KV 由来は 17.2 GFLOP / 4.0GB になります。つまり 32K の文脈では、attention のスコア計算と加重和だけで、モデル全体の重み行列積を上回る FLOPs を使っている。文脈が長くなるほど attention が支配的になるというのは、こういうことです。
合計すると、32K 文脈で1トークン生成するのに約 32 GFLOP、約 11.5GB の読み出しが必要という計算になります。仮に 614 GB/s のメモリ帯域（手元の M5 Max です）なら 1 トークン約 19ms、53 tok/s が理論上限。実際にはもっと遅くなりますが、桁としてはこんなものです。なおこれは重み1バイト相当（FP8）・KV は bf16 という仮定で、書き込みや量子化の展開コストも無視した数字です。手元で動かしている MLX の Q4 モデルとは前提が違うので、そのまま実測と比べられる数字ではありません。
ルーフライン: バランス点という物差し
ここで、半導体屋にとって馴染みのある道具を持ち出します。演算強度（arithmetic intensity）です。
\text{演算強度} = \frac{\text{FLOPs}}{\text{読み書きしたバイト数}}
そしてチップ側にも、同じ次元の数字があります。ピーク演算性能をメモリ帯域で割った値で、これをバランス点と呼びます。ワークロードの演算強度がバランス点を下回れば帯域律速、上回れば演算律速になります。
| チップ | ピーク演算（FP8） | メモリ帯域 | バランス点 |
|---|---|---|---|
| Rubin GPU | 約 16 PFLOPS | HBM4 22 TB/s | 約 730 FLOP/byte |
| Groq 3 LPU | 1.2 PFLOPS | SRAM 150 TB/s | 約 8 FLOP/byte |
（いずれもベンダー公称値からの単純計算で、実機ベンチマークによる検証は本稿執筆時点で存在しません。Rubin の FP8 値は資料により 16〜17.5 PFLOPS と幅があります。また NVFP4 推論の 50 PFLOPS は前々回記事で触れた通りスパース性を織り込んだ数字で、これでバランス点を計算すると 2,270 になります。以下では条件の揃う FP8 の dense 値で比較します）
一方、decode でバッチサイズ1のときの演算強度は、重み読み出しに関して言えば、
\frac{2 \times N_{params}}{1 \times N_{params}} = 2 \ \mathrm{FLOP/byte}
です（FP8 なら重み1個が1バイトなので）。2 対 730。このルーフラインモデルの上では、達成できる演算性能は FP8 ピークの 0.3% 相当が上限になる、ということです。念のため補足すると、これは nvidia-smi の GPU 使用率が 0.3% になるという意味ではなく、ピーク演算性能に対して理論上ここまでしか引き出せないという話です。ローカル LLM で GPU が忙しそうに見えるのに生成が遅い、という現象の背景がこれです。
同じワークロードでも、バランス点が 8 しかない LPU なら、バッチ1の decode でも、ルーフライン上の上限としてピークの 25% まで届きます。SRAM の帯域を積むというのは、バランス点を左に引き寄せて、演算強度の低いワークロードでも演算器を埋められるようにする、という話でもあります。
ここまでは前2記事でも書いた話の再確認です。本題は次です。
なぜ attention と FFN で性質が違うのか
バッチサイズ B を大きくすると何が起きるか、を FFN と attention で別々に考えます。ここが本記事でいちばん書きたかった部分です。
FFN の場合。B 人ぶんのトークンをまとめて処理しても、読む重みは同じ 176.2M 個の行列1組だけです。重みを1回読んで B 回使えるので、
\text{演算強度}_{FFN} = \frac{2 \times N \times B}{N} = 2B \ \mathrm{FLOP/byte}
演算強度がバッチサイズに比例して上がります。B が 365 くらいになれば Rubin のバランス点 730 に届き、演算律速に転じる。これが「大きなバッチを組んで重み読み出しを償却する」という GPU の必勝法の正体です。
attention のスコア計算と加重和の場合。B 人ぶんまとめても、読む KV キャッシュは各ユーザー固有です。B 人ぶん読めば B 倍の FLOPs が出ますが、読むバイト数も B 倍になる。
KV の要素数を N、bf16 なので1要素2バイト、GQA のグループサイズを g とします。KV 1要素につき g 本のクエリがそれぞれ1回の積和（2 FLOP）を行うので、
\text{演算強度}_{attn} = \frac{2 g N}{2 N} = g = 4 \ \mathrm{FLOP/byte}
演算強度はグループサイズそのものになります。素の MHA（g = 1）なら 1 FLOP/byte です。
B をいくら増やしても 4 のままです。バッチで償却できません。attention の中核部分は、原理的に帯域律速から抜け出せない演算だということになります。
厳密には例外があります。共有プレフィックス、つまり複数のリクエストが同じシステムプロンプトや few-shot 例を先頭に持つ場合です。Hydragen は、この共有部分に対する複数のクエリをまとめて処理し、同じ KV の読み出しを計算上も共有することで、共有部分に限って演算強度を引き上げます。なお SGLang の RadixAttention はやや性質が違い、こちらは主に KV の保存の重複排除と prefill の再計算削減が狙いです。KV を1コピーにしても、通常の decode カーネルがクエリごとに同じ KV を読み直すのであれば、読み出し量が自動的に 1/B になるわけではありません。
もう1つの例外は GQA/MQA のグループサイズで、これは上の式の分子を定数倍します。ただしどちらの例外も、バッチサイズを増やしても演算強度が伸びないという性質そのものを覆すわけではありません。
同じ図の中で、右上がりの直線と水平な直線が並んでいます。この傾きの差が、本記事のいちばん言いたいことです。
整理します。
| FFN（および QKVO 射影） | attention の中核（スコアと加重和） | |
|---|---|---|
| 読むデータ | モデルの重み | KV キャッシュ |
| 誰のものか | 全ユーザー共有の定数 | ユーザー・会話ごとに固有 |
| サイズ | モデルサイズで固定 | 文脈長 × 同時ユーザー数に比例して増える |
| バッチによる償却 | 効く（演算強度は 2B） | 効かない（演算強度は一定） |
| 演算の形 | GEMM / GEMV | 内積とリダクション |
| 文脈長への依存 | なし | O(L) |
同じ「帯域律速」でも、中身がまるで違います。FFN の帯域律速は「大きくて動かない塊を毎回読み直す」問題で、置き場所を速いメモリにすれば解けます。attention の帯域律速は「ユーザーごとに増え続ける状態を読む」問題で、大容量メモリが要ります。
ここで一つ、注意深く読むと引っかかる点があります。attention の演算強度 4 は、LPU のバランス点 8 すら下回っています。つまり attention は LPU に載せても帯域律速のままで、演算強度の観点だけを見るなら、むしろ SRAM の 150 TB/s が最も効くワークロードだとすら言えます。それでも attention が GPU 側に残されるのは、演算強度ではなく容量の問題だからです。
ラック単位で容量を並べると、桁の違いがはっきりします。LPX ラックの SRAM は256基あわせて 128GB。対する Vera Rubin NVL72 は HBM4 だけで 20.7TB（288GB × 72基）あり、その下に Vera CPU 側の LPDDR5X がさらに大きな容量で控えています。32K 文脈のユーザー1人で KV 4GB を要求するのですから、何百人も同時に捌く前提では、SRAM 側には桁が足りません。正確に言えば LPX 側にも DRAM の拡張はありますが、超高帯域 SRAM だけで長文脈・高並列の KV を抱えようとすれば、FFN の重みと容量を奪い合うことになります。前回記事で書いた「SRAM 単独方式では KV とモデル重みが SRAM を食い合う」の中身は、この容量の桁差でした。
演算強度が分業の理由を説明してくれるのは FFN 側だけで、attention 側の理由は容量なのだ、という非対称は押さえておきたいところです。
そして両者に必要なメモリの性質は、正反対です。
| データ | 題材モデルでの大きさ | 必要な性質 |
|---|---|---|
| 隠れ状態（層をまたぐ活性） | 8KB / トークン | ほぼ何でもよい |
| KV キャッシュ | 128KB / トークン、32K 文脈で 4GB / ユーザー | 大容量。帯域も要る |
| モデルの重み | 8GB（うち FFN が 5.6GB） | 超高帯域。容量は固定 |
500MB の SRAM に KV キャッシュは置けません。32K 文脈のユーザー1人ぶんで 4GB あるのですから、SRAM 8個ぶんが1人に食われます。逆に、重みは全ユーザー共有なので、いったん SRAM に載せてしまえば何人でも共有できます。
つまり LPU の SRAM は「全員で共有できる大きな定数」を置く場所としては最高で、「ユーザーごとに増える状態」を置く場所としては最悪なわけです。FFN が前者、attention が後者。分担はここで決まっていました。
前回記事で「単体で完結させなかったのは、SRAM 単独方式では KV キャッシュがモデルの重みと SRAM を食い合うからだ」と書きましたが、その食い合いの正体は、この表の2行目と3行目です。
ついでに、GQA が何をしているのかもこの枠組みで読めます。KV ヘッドを 32 から 8 に減らすと、KV の容量が 4 分の 1 になるだけでなく、KV 1バイトあたりの積和回数が4倍になります。容量削減の手法として紹介されることが多いのですが、演算強度を引き上げる処方箋でもある、という見方ができます。MLA も同様に、圧縮した KV を読んでから展開することで、読むバイト数を減らしつつ演算強度を上げています。
静的スケジューリングとの相性
もう一つ、前回書ききれなかった論点があります。LPU はデータフロー型の静的スケジューリングで、データの経路と時刻をコンパイル時に確定させます。実行時間が決定論的になるのはそのおかげですが、裏を返せば、演算の形が実行時に変わってはならないということです。
ここで前節の表の最終行を見ます。
FFN の演算は、文脈長 L に一切依存しません。入力は常に 4,096 次元のベクトル1本、重みは 4,096 × 14,336 の固定行列。1トークン目だろうが10万トークン目だろうが、まったく同じ形の計算です。コンパイル時に形が完全に決まる。静的スケジューリングにとって、これ以上ないほど素直な相手です。
attention は違います。読む KV の数は L 個で、L は実行時にしかわかりません。しかもトークンを1つ生成するたびに1ずつ増えていきます。バッチ内のユーザーごとに L はバラバラですし、プレフィックスキャッシュのヒット状況で prefill の長さも変わる。形が実行時に決まる演算です。
ここは誤解のないように書いておきます。だから LPU で attention が動かない、という話ではありません。実際、GroqCloud は attention を含む Llama をそのまま LPU 上で提供していますし、131K の文脈長にも対応しています。最大長を固定してパディングする、長さ別にバケットを切って複数のプログラムをコンパイルしておく、といった方法で可変長は捌けます。NVIDIA 自身も、LPU の SRAM 上の作業セットとして重みと活性に加えて KV state を挙げています。
正確に言うならこうなります。FFN は形が固定なので、静的スケジューリングの実行モデルと素直に噛み合う。attention は形がリクエストごとに変わるので、実行できないわけではないが、可変長の扱いと KV の容量管理を抱え込むぶん、動的な実行とメモリ管理が得意な GPU 側に置いたほうが素直になる。分業の主因はあくまで前節の容量とバッチ償却で、実行モデルの相性はそれを後押しする副次的な理由、というのが妥当な位置づけだと思います。
つまり FFN と attention の分担は、メモリの性質（共有できる定数か、ユーザー固有の状態か）だけでなく、実行モデルの性質（形が静的か動的か）でも、同じ線で切れています。強さの異なる2つの理由が同じ結論を指しているわけで、この分担がかなり自然なものだということがわかります。
Groq のアーキテクチャがもともと静的スケジューリングだったこと、そしてその Groq のチップに FFN を割り当てたことは、偶然の組み合わせではなかったのだと思います。
層をまたぐデータは 8KB しかない
分業のもう一つの疑問が、これで解けます。FFN を LPU、attention を GPU に分けるなら、1トークンを生成するあいだに32層ぶん、つまり32往復ぶん、GPU と LPU のあいだをデータが行き来することになります。しかもその配線は NVLink ではなく Spectrum-X Ethernet です（LPU 同士のスケールアップ接続には専用のファブリックが使われますが、GPU との接続はスケールアウト側の Ethernet 経由だと説明されています）。なぜそれで足りるのか。
ちなみにこの分離は、NVIDIA では AFD（Attention-FFN Disaggregation）と呼ばれています。学術側にも、attention とそれ以外を異種デバイスに分離する近縁の研究（Model-Attention Disaggregation など。ハードウェアの割り当ては NVIDIA の構成と同一ではありません）があり、NVIDIA の発明というより、業界が同じ結論に収束しつつある設計パターンだと見るのが正確だと思います。
答えは、層の境界を流れるのが隠れ状態、つまり 4,096 次元のベクトル1本＝8KB だけだからです。
1トークン・1層あたり往復 16KB、32層で 512KB。一般形で書けば、1ステップあたりの通信量は
V_{comm} = 2 \times B \times n_{layers} \times d_{model} \times s
（B はバッチサイズ、s は1要素のバイト数）となり、同時処理数と d_model と、投機デコードを使うならその先読みトークン数に比例して増えます。それでも、仮に1ストリーム 1,000 tok/s で回しても 0.5 GB/s（机上計算）という桁で、帯域としては Ethernet で十分に通ります。
むしろ本質的なのはレイテンシのほうだと思います。往復は層の数だけ発生し、そのたびに同期が入るからです。NVIDIA も AFD を、トークンごとに中間活性を交換する反復ループとして説明しています。ここで LPU の決定論的な実行時間が効いてくるはずですが、具体的にどう隠蔽しているのかは公開情報からは読み切れませんでした。
一方、もし KV キャッシュのほうを動かそうとしたら、ユーザー1人で 4GB です。到底運べません。だから「KV のある側に attention を置き、細い活性だけを配線に流す」。分業の切り口は、通信量の観点からもここしかなかったことになります。
先ほどの3つのデータの大きさの表を、置き場所と結びつけると次のようになります。
| データ | 大きさ | 置き場所 | 動くか |
|---|---|---|---|
| 隠れ状態 | 8KB / トークン | ─ | 層ごとに GPU と LPU を往復する |
| KV キャッシュ | GB / ユーザー | HBM4 →（アイドル時）LPDDR5X → フラッシュ | ユーザー単位で階層を退避する |
| 重み（FFN） | GB / モデル | 各 LPU に割り当てられた作業セットを SRAM に置く | 動かない |
Vera Rubin の設計を、この3行に圧縮できると思っています。前2記事で書いた KV キャッシュの3階層退避も、FFN の SRAM 移設も、Ethernet 接続で足りる理由も、すべてこの表の帰結です。
MoE はこの絵をどう変えるか
最後に MoE に触れておきます。MoE は FFN 層を N 個のエキスパートに分割し、トークンごとにルーターが上位 k 個だけを選んで通す構造です。attention 側は何も変わりません。
ここで、私が最初に誤解していた点を書いておきます。MoE は演算強度を上げる最適化ではありません。エキスパートを一部しか通さないと読むバイト数は減りますが、FLOPs も同じだけ減ります。バッチ1・FP8 のエキスパート GEMV の演算強度は、
\frac{2 N_{active}}{N_{active}} = 2 \ \mathrm{FLOP/byte}
で、dense の FFN とまったく同じ 2 のままです。分母だけ減らして分子を維持できるわけではない、ということです。
MoE がやっているのは、総パラメータの容量と、1トークンあたりの計算量・重み読み出し量を切り離すことです。巨大な総容量を持ちながら、各トークンでは一部しか使わない。Switch Transformer や Mixtral が狙っているのも「計算量をほぼ据え置いたまま総パラメータを増やす」ことであって、演算強度の改善ではありませんでした。
なお DeepSeek のように MoE と MLA を併用するモデルもありますが、あれは FFN 側を MoE で、attention 側を MLA で、それぞれ別々に殴っているのであって、MoE が attention に手を入れているわけではありません。本記事の分類で言えば、MoE は重み側の最適化、MLA は KV 側の最適化です。
そして容量の問題は残ります。どのエキスパートが選ばれるかはトークンごとに変わるので、全エキスパートをメモリに置いておく必要がある。速度はアクティブパラメータ数で決まり、容量は総パラメータ数で決まる、という非対称がここから来ます。前回記事で「SRAM 常駐方式なら MoE でも総パラメータ分の容量が要る」と書いたのも同じ理屈です。
ここで、本記事の題材モデルについて一つ注意書きを足しておきます。Llama 3 8B 相当の dense モデルは算数を追う教材としては最適ですが、GPU と LPU を分けて運用する構成が実際に想定されているのは1兆パラメータ級の MoE です。FFN の重みが十分に大きく、かつ attention 側の KV が支配的になる規模でこそ、分離のオーバーヘッドを払う価値が出てくるからで、8B の dense を32層ぶん Ethernet で往復させる構成に意味があるという話ではありません。本記事の数字は、あくまで比率と桁を掴むためのものとして読んでください。
最後に、私が手元のベンチマークで観測してきた「同じ保存容量なら dense のほうが SV/UVM のプロトコル精度で MoE を上回る」という結果にも触れておきます。gpt-oss-120b が Gemma 4 31B dense に負けた、というあれです。
アクティブパラメータ数の少なさが効いている可能性は高いと思っていますが、これをもって「精度に効くのはアクティブパラメータだけだ」と一般化するのは行き過ぎでしょう。Mixtral が 13B 級のアクティブパラメータでより大きな dense に匹敵する結果を出していることを思えば、総パラメータが担う知識容量、エキスパートの専門化、ルーティングの質、学習データと学習計算量など、効いている要素は他にもあるはずです。私の観測は「私のタスクで、私が試した数モデルでは、そうだった」という以上のものではありません。MoE が計算量と重み読み出し量を節約する代わりに何を失っているのかは、アーキテクチャの一言では片づかない、というのが今の理解です。
所感
書き終えて思うのは、前2記事で私が引用符つきで書いていた分担が、Transformer の数式から素直に導けるものだった、ということです。
- FFN が読むのは全ユーザー共有の定数だから、SRAM に常駐させて共有できる
- attention が読むのはユーザー固有の状態で、文脈長に比例して膨らむから、大容量の HBM とその下の階層が要る（演算強度だけを見れば attention こそ SRAM 向きですが、容量がそれを許しません）
- FFN は形が実行時に変わらないから、静的スケジューリングのチップと素直に噛み合う
- attention は形がリクエストごとに変わるから、可変長と容量管理が得意な GPU 側に置いたほうが素直になる（LPU で動かないという意味ではありません）
- 層をまたぐ活性は 8KB しかないから、両者を Ethernet で結んでも成立する
これらは NVIDIA の設計判断というより、Transformer という計算構造がハードウェアに突きつけている要求です。NVIDIA が偉いのではなく、Transformer がそういう形をしている。200億ドルの取引も、7チップのプラットフォームも、この非対称性への応答だと考えると、腑に落ちる感じがあります。
エンジニアとしての感想も一つ。Transformer の演算を追ってみると、正規化と残差加算という「軽い」処理を除けば、やっていることは行列積とリダクションと softmax だけです。この単純さこそが、専用チップを起こす動機になっている。RTL で書けと言われたら、演算そのものは難しくありません。難しいのはメモリ階層とスケジューリングのほうで、つまり Transformer 用チップの設計とは、演算器の設計ではなくデータ移動の設計なのだと、算数を追ってようやく実感しました。SystemVerilog で検証すべき対象も、乗算器の正しさではなく、可変長のデータが階層をまたいで正しく流れるかどうか、というところに寄っていくはずです。
そして、手元の PC でトークンがぽつぽつ出てくるときも、Vera Rubin のラックが 1,000 tok/s を叩き出すときも、動いている数式は同じです。読むべきバイトを読み終えるまで、次のトークンは出てこない。この身も蓋もない事実が、200億ドルの買収から私の 128GB の PC まで、同じ形で貫いているのは、やはり面白いと思います。
用語集
前2記事の用語集も併せて参照してください。本記事で新しく使った用語を中心に載せます。
| 用語 | 分類 | 説明 |
|---|---|---|
| 隠れ状態 | 演算 | 1トークンを表す d_model 次元のベクトル。層をまたいで流れる実体。題材モデルでは 8KB |
| 埋め込み | 演算 | トークン ID からベクトルを引くテーブル参照。行列積ではない |
| RMSNorm | 演算 | 二乗平均平方根による正規化。演算量は軽微 |
| 残差接続 | 演算 | サブ層の出力を入力に足し戻す配線。勾配の通り道 |
| QKV 射影 | 演算 | 隠れ状態から Query / Key / Value を作る行列積。性質は FFN と同じ |
| RoPE | 演算 | 位置に応じて q と k を回転させる位置エンコーディング。学習パラメータではない |
| softmax | 演算 | スコアを確率分布に変換する非線形処理。L 方向のリダクションを伴う |
| SwiGLU | 演算 | gate / up / down の3行列を使う FFN の形。現代の主流 |
| GQA | 演算 | Grouped Query Attention。複数のクエリヘッドで KV ヘッドを共有し、KV 容量を減らして演算強度を上げる |
| MLA | 演算 | Multi-head Latent Attention。KV を圧縮して保持する DeepSeek の方式 |
| GEMM / GEMV | 演算 | 行列積 / 行列ベクトル積。prefill が前者、decode が後者になる |
| 演算強度 | 性能 | FLOPs ÷ 読み書きバイト数。ワークロード側の性質 |
| バランス点 | 性能 | ピーク演算 ÷ メモリ帯域。チップ側の性質。演算強度がこれを下回れば帯域律速 |
| ルーフライン | 性能 | 演算強度を横軸に、達成可能性能の上限を描く性能モデル |
| アクティブパラメータ | MoE | 1トークンの処理で実際に使われるパラメータ数。速度と精度に効く。容量は総パラメータで決まる |
| AFD | 分業 | Attention-FFN Disaggregation。attention と FFN を別のハードウェアに分離して実行する構成。NVIDIA の GPU＋LPU 分業もこれにあたる |
| 共有プレフィックス | 推論 | 複数リクエストが先頭に持つ共通の文脈。この範囲の KV は1コピーで済み、attention の演算強度を上げられる数少ない例外 |
参考リンク
- Vaswani et al., Attention Is All You Need（原論文） https://arxiv.org/abs/1706.03762
- Su et al., RoFormer: Enhanced Transformer with Rotary Position Embedding（RoPE） https://arxiv.org/abs/2104.09864
- Ainslie et al., GQA: Training Generalized Multi-Query Transformer Models https://arxiv.org/abs/2305.13245
- Shazeer, GLU Variants Improve Transformer（SwiGLU） https://arxiv.org/abs/2002.05202
- Dao et al., FlashAttention https://arxiv.org/abs/2205.14135
- Williams et al., Roofline: An Insightful Visual Performance Model https://dl.acm.org/doi/10.1145/1498765.1498785
- Efficient Heterogeneous Large Language Model Decoding with Model-Attention Disaggregation（attention の演算強度がバッチサイズに依存しないことの定式化） https://arxiv.org/abs/2405.01814
- Juravsky et al., Hydragen: High-Throughput LLM Inference with Shared Prefixes https://arxiv.org/abs/2402.05099
- NVIDIA Technical Blog: Inside NVIDIA Groq 3 LPX https://developer.nvidia.com/blog/inside-nvidia-groq-3-lpx-the-low-latency-inference-accelerator-for-the-nvidia-vera-rubin-platform/
- Kipply, Transformer Inference Arithmetic https://kipp.ly/transformer-inference-arithmetic/
- 前々回記事: NVIDIA Vera Rubin のアーキテクチャを理解する https://qiita.com/sukimaengineer/items/76ee0281a53bb293292a
- 前回記事: Groq 3 LPX はなぜ「エージェント AI 向け」と言われるのか？ https://qiita.com/sukimaengineer/items/1d090d090a36d5034085
本記事の調査・執筆・図版作成には Claude（Anthropic）とその Research 機能を利用しています。