『GPUとLLM』第2章。文章がトークンと T × d の行列になり、線形層・Attention・MLP のどれもが行列積として計算される様子を、図が動く形でたどります。計算量 2·m·n·k、推論 ≈ 2N・学習 ≈ 6N FLOP、prefill と decode の形の違い、KV キャッシュと重みのメモリまで、Llama 3 70B の寸法を例に概算します。
ChatGPT のような大規模言語モデル（LLM）は、文章を読んで次に来る言葉を選んでいます。けれども GPU の中で実際に起きているのは、言葉の処理というより、大きな数の表どうしの掛け算、つまり行列積のくり返しです。図の X × W ＝ Y が、この章でずっと追いかける形です。
この章では、文章が数の表になるところから始めて、線形層の計算量、Transformer の1層の中身、推論の2つの段階、そして動かすのに必要なメモリまでを順にたどります。具体的な数には、公開されている代表例として Llama 3 70B の寸法を使います。
LLM は文章を1文字ずつでも1単語ずつでもなく、トークンという単位に区切って扱います。図では「GPUは行列が得意です」が6つのトークンに分かれました。どこで区切るかはトークナイザという前処理が決めていて、よく出てくる文字の並びほど1つのトークンにまとまりやすくなります（図の区切り方と番号は説明用の例です）。
区切られたトークンには、語彙の中での番号が付きます。Llama 3 の語彙は 128,256 種類です。ここから先、モデルが受け取るのはこの番号の列だけです。トークンの数を T と書くことにします。図の例では T ＝ 6 です。
番号のままでは計算に使えないので、各トークンを d 個の数の並び（ベクトル） に置き換えます。右の紫の表が埋め込み表で、語彙の数だけ行があり、トークンの番号の行をそのまま取り出します。Llama 3 70B では d ＝ 8,192 です。
取り出した行を上から積むと、左の青い表になります。これが入力の行列 X で、大きさは T 行 × d 列です。図の例なら 6 × 8,192 ＝ 49,152 個の数が並びます。この先モデルの中を流れていくのはずっとこの T × d の形の表で、各層はこれを同じ形のまま少しずつ書き換えていきます。
X の1行、つまり1トークン分のベクトル x（1 × k）に、重みの行列 W（k × n）を掛けると、長さ n の新しいベクトル y ができます。これが線形層で、LLM の部品のほとんどはこの形をしています。k は入力の長さで、最初の層なら d ＝ 8,192 です。
図は掛け算を十字の形に並べたものです。濃い青の x の行と、濃い紫の W の1列が交わる濃いオレンジのマスが、y の1つの値です。中身は k 個の数を順に掛けて足す計算（内積）で、掛け算 k 回と足し算 k 回、約 2k 回の計算（FLOP）が要ります。これを W の n 列すべてについて行うと、薄いオレンジの y の1行が埋まります。
1行分の計算を T 個のトークン全部について行うと、オレンジの表 Y が全部埋まります。式で書けば Y ＝ X W で、（m × k）×（k × n）→（m × n）の行列積です。m はトークンの数 T にあたります。紫の W は1回用意すれば、すべての行で使い回せることにも注目してください。
計算量は数えるだけで分かります。1行に約 2·k·n 回、それが m 行なので、合計は約 2·m·n·k FLOP です。たとえば m ＝ 1,000、k ＝ n ＝ 8,192 なら約 1.3 × 10¹¹ 回。H100 の BF16 の理論ピーク（毎秒約 10¹⁵ 回）で割れば約 0.14 ミリ秒ですが、実際にどこまで近づけるかは第3章の話題です。
LLM の本体は、同じ形の層を何十枚も重ねたものです（Llama 3 70B は80層）。1層は大きく2つの部品でできています。トークンどうしで情報をやりとりする Attention と、トークンを1つずつ加工する MLP です。どちらの出力も、上の線（残差）を通ってきた元の X に足し戻され、形は T × d のまま次の層へ渡ります。
図の部品の色を見てください。濃い紫は重みの行列との掛け算、薄い紫は Q・K・V どうしの行列積、灰色は正規化や softmax のような1マスずつの軽い計算です。ほとんどが紫、つまり行列積だと分かります。ここから Attention、MLP の順に中を見ていきます。
Attention ではまず、X から3種類の表を作ります。探しているものを表す Q（クエリ）、自分の見出しにあたる K（キー）、相手に渡す中身の V（バリュー） です。作り方は3つとも同じで、Q ＝ X W_Q、K ＝ X W_K、V ＝ X W_V。前のページの紫の部品が、そのまま3枚の重み行列になりました。
Q はヘッドという小さな単位に分かれていて、Llama 3 70B では 64 ヘッド × 128 次元 ＝ 8,192 です。一方、K と V は8ヘッド分（1,024）だけです。8つの Q ヘッドで1組の K・V を共有する工夫（GQA）で、図の K・V が細いのはそのためです。この差は、あとで KV キャッシュの大きさに効いてきます。
次に、各トークンの Q と、すべてのトークンの K の内積をとります。1つのヘッドについて、Q（T × 128）と、K を転置した Kᵀ（128 × T）の行列積で、結果は T × T の点数表 S になります。図で K が横に倒れたのが転置です。S の行が「見る側」、列が「見られる側」のトークンです。
点数は √128 ≈ 11.3 で割って大きくなりすぎないようにし、softmax で各行を合計1の割合に直します。ただし先の文章を見てはいけないので、自分より後ろのトークン（右上の灰色）は0にします。まとめると S ＝ softmax(Q Kᵀ ÷ √d_head) です。計算量は1ヘッドで約 2·T·T·128 回で、T の2乗で増えるのが特徴です。
点数表 S ができたら、それを V に掛けます。O ＝ S × V で、（T × T）×（T × 128）→（T × 128）。各トークンは、点数の割合に応じて他のトークンの V を混ぜたものを受け取ります。図では緑の S が左へ動き、V と掛け合わされて O になりました。これも行列積です。
64 ヘッド分の O を横に並べると T × 8,192 に戻り、最後にもう1枚の重み W_O（8,192 × 8,192）を掛けて Attention の出力になります。数えると、Attention の中身は重みとの行列積が4回と、Q・K・V どうしの行列積が2回でした。softmax のような軽い計算は、量としてはわずかです。
Attention の次は MLP です。こちらはトークンどうしのやりとりはなく、各行を別々に加工します。図のように、T × 8,192 の X に重み W_up を掛けて T × 28,672 へと約3.5倍に広げ、活性化関数で1マスずつ形を変えてから、W_down を掛けて T × 8,192 に戻します。
Llama 3 の MLP はもう少し工夫されていて、W_gate と W_up の2枚で広げた結果を、片方に活性化（SiLU）をかけてから1マスずつ掛け合わせ、W_down で戻します（SwiGLU と呼ばれる形）。重みは 8,192 × 28,672 の大きさが3枚で、1層あたり約7億個です。ここでも主役は行列積で、活性化は量としてはわずかです。
1層の重みを数えると、Attention が約1.5億個、MLP が約7億個で、MLP が8割以上を占めます（上のバー）。これを80層分まとめ、入口の埋め込み表と出口の語彙の表（合わせて約21億）を足すと、N ≈ 706億、つまり約 70B になります。
重み1個は、1トークンを通すたびに掛け算1回と足し算1回に使われます。だから1トークンの推論（順伝播）は約 2N FLOP、70B なら約 1.4 × 10¹¹ 回です。学習では誤差を逆向きに伝える計算がその約2倍かかるので、約 6N FLOP／トークンが目安です。Llama 3 は15兆トークン以上で学習したと発表されていて、6N を掛けると約 6 × 10²⁴ 回になります。
文章を生成するとき、推論は2つの段階に分かれます。最初の prefill では、プロンプトのトークン全部をまとめて一度に通します。図のように X は T 行あり、計算は（T × k）×（k × n）の行列積です。紫の W を1回読み込めば、T 行分すべてに使い回せます。
このとき各層で作った K と V は、右の緑の KV キャッシュ に保存しておきます。次の段階で、過去のトークンの K と V を計算し直さずに済ませるためです。prefill の最後には、最後のトークンの出力から、生成する最初の1トークンが選ばれます。
2つ目の decode では、1トークンずつ生成します。新しく入るのは直前に生成した1トークンだけなので、図の X は下の1行だけになりました（薄い行は過去のトークンで、もう計算し直しません）。計算は（1 × k）×（k × n）、つまり行列 × ベクトルです。紫の W は prefill と同じ大きさのまま、それを丸ごと読んで1行分にしか使いません。
新しいトークンの k と v は KV キャッシュの一番下に1行足され、q はキャッシュに並んだ過去全部の K・V と Attention を計算します。同じ重み、同じ式でも、prefill と decode では計算の形がまったく違います。この違いが GPU の速さにどう効くかは、第3章で見ていきます。
KV キャッシュの大きさは、掛け算で見積もれます。上の式のとおり、K と V の2つ × 層数 × KV ヘッド数 × ヘッド次元 × トークン数 × 1つの数のバイト数です。Llama 3 70B を BF16（1つ2バイト）で動かすと、1トークンあたり 2 × 80 × 8 × 128 × 2 ＝ 327,680 バイト、約 0.33MB になります。
トークン数に比例して増えるので、8,192 トークンで約 2.7GB、Llama 3.1 が扱える最大の 131,072 トークンなら約 43GB です（下のバー）。これは会話1つ分なので、何人もの会話を同時にさばけば、その人数分だけ増えます。もし K・V を Q と同じ 64 ヘッドで持っていたら、さらに8倍でした。
動かすのに必要なメモリを並べてみます。まず重みで、容量は「パラメータ数 × 1個のバイト数」です。706億 × 2バイト（BF16）で約 141GB。これに KV キャッシュ（128k トークンの会話1つで約 43GB）と、計算途中の値（アクティベーション）が加わります。途中の値の量は、まとめて処理する数や長さで変わります。
縦線が H100 1枚分の 80GB です。重みだけで線を大きくはみ出し、2枚分の 160GB でも KV キャッシュまでは入りません。1枚に載らないモデルを動かすには、1個のバイト数を減らす量子化と、複数の GPU に分ける並列化が必要で、これが第4章のテーマです。
この章の要点は次の4つです。
計算のほとんどが行列積なので、行列積を速くできれば LLM 全体が速くなります。第1章で見た Tensor Core は、まさに小さな行列積をまとめて計算する回路でした。次の第3章では、同じ行列積でも速く回るときと回らないときがあるのはなぜかを、算術強度とルーフラインで見ていきます。
