GPUのメモリ階層 — HBM・SRAM・レジスタと「移動が支配する」現実

GPUの速さを決めているのは演算器ではなく、レジスタ・共有メモリ・L2・HBMのどこにデータを置き、何回運ぶかです。階層の容量と帯域の桁、演算強度とタイル化、階層ごとに引くルーフライン、そして「FLOPsを増やしたのに速くなる」FlashAttentionまでを前提知識ゼロから追います。

FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness

比喩: 手元・作業台・棚・倉庫

腕のいい職人がひとり、部品を組み立てています。手のひらには2〜3個しか持てません。目の前の作業台には数十個、振り向いた先の棚に数百個、材料の全量は中庭を渡った先の倉庫です。

この人の出来高を決めているのは手の速さではなく、次に要る部品がいまどこにあるかです。手のひらから取れば一瞬、倉庫まで歩けば1分。倉庫に取りに行くたびに1分かかるなら、どれほど手が速くても1分に1個しか作れません。

GPUも同じ構造です。演算器は十分すぎるほど速く、実行時間を決めているのは「そのデータがいまどの引き出しにあるか」と「何往復するか」。メモリの壁を1からでは、なぜ運ぶほうが高いのかを配線の物理から見ました。ここではその続きとして、GPUという具体的な機械の引き出しを開けます。

GPUの引き出しは4段ある

数値はNVIDIA A100を例に取ります。

容量は上から下へ5桁以上増え、速さは同じだけ落ちます。FlashAttentionの論文はこの落差をA100の実測で示していて、オンチップSRAMは合計約20 MBで帯域は約19 TB/s、HBMは40 GBで約1.5 TB/s。容量は2000倍、帯域は10倍以上違う。 ホストまで出ればさらに二桁落ちます。

GPU固有の意外な点が2つあります。

ひとつ、レジスタファイルが異様に大きい。CPUのレジスタは1コア数百バイトですが、GPUは1 SMあたり256 KB、A100は108 SMなので合計27 MB前後——L2と同じ桁です。何万本ものスレッドの状態を全部保持したまま切り替え、メモリ待ちを他のスレッドの仕事で埋めるからで、レジスタは速い一時置き場であると同時に待ち時間を隠すための在庫でもあります（GPUはなぜ速いのかの占有率と表裏一体です）。

ふたつ、共有メモリはキャッシュではない。キャッシュは何が残るかをハードが決めますが、共有メモリ（__shared__）は載せるものをプログラマが決めるSRAMで、「置いたものは消すまでそこにある」と保証されます。この保証があるからこそ、後半のFlashAttentionのような運ぶ回数を設計する最適化が成立します。

「移動が支配する」を式にする

カーネル1本の実行時間には、動かしようのない下限があります。

$W$ は演算回数（FLOPs）、$\pi$ はピーク演算性能、$Q$ はある階層をまたいで運ぶバイト数、$\beta$ はその階層の帯域です。つまり計算し終わる時間と運び終わる時間の、遅いほうより速くは終われない。この2つの比が演算強度になります。

要するに「1バイト運ぶごとに何回の演算をするか」、運んだ材料を何回使い回せたかです。$I$ が機械のマシンバランス（ピーク性能 ÷ 帯域）を下回れば帯域律速、上回れば演算律速——ここまではメモリの壁と同じ話です。

この記事で足したいのは一点だけ。$Q$ は階層ごとに違う値になる。 同じカーネルでもHBMから見た $Q$、L2から見た $Q$、共有メモリから見た $Q$ は別物で、演算強度もルーフラインも階層の数だけ存在します。「このカーネルは帯域律速か」は問いとして不完全で、正しくは「どの階層の帯域で律速しているか」です。