体系 › 計算機アーキテクチャ
·無料·論文·12分で読めます

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

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

要点

  • GPUのメモリはレジスタ・共有メモリ/L1・L2・HBMの4段で、容量は5桁以上、帯域は桁で違う。
  • 共有メモリはキャッシュではなく、載せるものをプログラマが決めるSRAMで、置いたものは消すまでそこにある。
  • 演算強度 I = W/Q の Q は階層ごとに違うので、どの階層の帯域で律速しているかまで問う必要がある。
  • I を上げる方法は再利用だけで、タイル化では I = b/s となり、タイルの一辺に比例する。
  • FlashAttentionはスコア行列をHBMに書かず逆伝播で再計算し、FLOPsを増やして運搬を減らす。
対象textタスクhardware

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

一次資料 — この記事の根拠

この解説の公開 2026-08-22

FlashAttention: Fast and Memory-Efficient Exact Attention with IO-AwarenessarXiv:2205.14135論文ページ·PDF

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

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

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

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

GPUの引き出しは4段ある

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

階層 実体 置き場所を決めるのは 容量の目安
レジスタ SRAM コンパイラ 256 KB / SM(32ビット×65,536本)
共有メモリ・L1 SRAM 共有メモリはプログラマ、L1はハード 192 KB / SM(共有メモリは最大164 KB/ブロック)
L2キャッシュ SRAM ハード 40 MB(チップ全体で共有)
グローバルメモリ 積層DRAM(HBM) プログラマ(cudaMalloc) 40 GB
ホストのDRAM DRAM 明示的な転送のみ 数百GB

容量は上から下へ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本の実行時間には、動かしようのない下限があります。

t≥max⁡ ⁣(Wπ, Qβ)t \ge \max\!\left(\frac{W}{\pi},\ \frac{Q}{\beta}\right)
(1)

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

I=WQI = \frac{W}{Q}
(2)

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

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

FIG 1行列積では計算量が n³、データ量が n² で伸びる。伸び方が違えば差は必ず桁になり、大きな行列ほど「1バイトあたりの仕事」=演算強度が上がる。縦軸を線形に切り替えると下の曲線が潰れて読めなくなる——メモリ階層を対数で語るのも同じ理由

演算強度は再利用でしか上がらない

WW は解くべき問題が決めてしまうので、II を上げるには QQ を減らすしかありません。その方法は事実上ひとつ、一度運んだデータを速い階層に置いたまま何度も使う——タイル化です。

n×nn \times n の行列積を b×bb \times b のタイルに切ると、タイルの組は (n/b)3(n/b)^3 通り、各組でタイル2枚(2b22b^2 要素)を読んで 2b32b^3 回の演算をします。1要素 ss バイトなら

I=2b32b2s=bsI = \frac{2b^3}{2b^2 s} = \frac{b}{s}
(3)

つまり演算強度はタイルの一辺に比例する。fp16(s=2s=2)で b=128b=128 なら64 FLOP/byteです。速い引き出しが大きいほど大きなタイルを置け、大きなタイルほど運搬が減ります。

逆に言えば、要素ごとに1回しか触らない計算——足し算、ReLU、正規化、埋め込みの引き当て——はタイルを大きくしても II が上がりません。構造的に永久に帯域律速です。

階層ごとにルーフラインを引く

紙の上で計算してみます。n=8192n = 8192、fp16、b=128b = 128 のGEMM。HBMから運ぶのは 2n3/b2n^3/b 要素×2バイト ≈ 17 GB、演算量は 2n3≈1.12n^3 \approx 1.1 TFLOP。A100のfp16テンソルコアのピークを312 TFLOP/s、HBM帯域を1.5 TB/sとすると、演算側3.5 msに対して転送側11 ms。「大きなGEMMでも帯域律速」という結論になってしまいます。

ところが実際の高性能GEMMはこの見積もりより速く回ります。間違っているのは前提のほうで、上の計算は「タイルを読むたびに毎回HBMまで行く」と仮定していました。現実にはL2が40 MBあり、隣り合うSMがほぼ同じ時刻に同じタイルを要求するので2度目以降はL2で返ります。CUTLASSなどが行うスレッドブロックのスウィズル(実行順の並べ替え)は、このL2ヒットを狙って巡回順を組み替える最適化です。

これが階層ルーフラインの実用的な意味です。紙の上でHBMトラフィックを見積もったら、必ず実測の転送量と突き合わせる。実測が見積もりを大きく下回っていればL2が効いており、一致していればL2が素通りされていて改善の余地がある。Nsight Computeの dram__bytes.sum が、実際にHBMを渡ったバイト数です。

コードで数える: 何往復しているか

もっと日常的な例で往復を数えます。

y = torch.relu(x)   # x を読み、y を書く   → 2往復
z = y * scale       # y を読み、z を書く   → 2往復
w = z + bias        # z を読み、w を書く   → 2往復

3行がそれぞれ独立したカーネルなら、要素1個につきHBMを6回渡ります。演算は3回。I=3/(6×2)=0.25I = 3/(6 \times 2) = 0.25 FLOP/byte で、A100のfp16のマシンバランス(約200 FLOP/byte)から3桁下です。演算器は99%以上の時間を待って過ごします。

1本のカーネルに融合すれば、xx を1回読んで ww を1回書くだけ、往復は6回から2回へ。FLOPsは1回も減っていないのに3倍速くなるわけです。torch.compile の仕事のかなりの部分がこの融合で、Tritonを手書きするときも最初に確認するのはここです。

FlashAttentionがなぜ効くのか

同じ論理をattentionに当てます。Attention機構を1から理解するで見たとおり、自己注意は QK⊤QK^\top で N×NN \times N のスコア行列を作り、softmaxをかけ、VV を掛けます。素直に実装すると、この行列がHBM上に実体として書き出され、softmaxのために読み直され、VV との積でまた読まれます。

大きさを数字にすると効きます。系列長 N=4096N = 4096、fp16なら1ヘッドあたり 40962×2=33.54096^2 \times 2 = 33.5 MB。32ヘッド・バッチ8なら8.6 GBが、層ごとに何度も往復する。しかもこれは最終出力ですらない、捨てられる中間結果です。

FlashAttentionの発想は、この行列を一度もHBMに書かないことです。QQ・KK・VV をブロックに切って共有メモリに載せ、ブロックごとに部分スコアを計算し、softmaxの最大値と分母を走らせながら更新していく(オンラインsoftmax)。N×NN \times N の実体はどこにも存在せず、SRAMの中で生まれて消えます。論文の解析では、HBMアクセスは標準実装の Ω(Nd+N2)\Omega(Nd + N^2) に対して O(N2d2/M)O(N^2 d^2 / M)(MM はSRAM容量)。要するに速い引き出しが大きいほど運搬が減るという、さきほどのタイル化の式と同じことを言っています。

いちばん大事な点は逆伝播にあります。FlashAttentionはスコア行列を保存しない代わりに、逆伝播でもう一度計算し直します。演算回数は増える。それでも速い。帯域律速の世界では、FLOPsは通貨ではありません。 通貨はバイトで、計算をやり直してでも運搬を減らすほうが安い、という交換レートが成立しています。

現場ではこう使う

誰がいつ触るか。 推論基盤の担当が「同じGPUで何リクエスト捌けるか」を上げるとき、学習側が1ステップの時間を削るとき、カスタムカーネル(Triton/CUDA)を書くとき。モデル構造を一行も変えずに残っている改善の余地は、たいていこの層にあります。

最初に見る指標。 Nsight Compute の Memory Workload Analysis(ncu --set full)で、dram__bytes.sum(HBM実転送量)、L1/L2のヒット率、smsp__sass_average_data_bytes_per_sector_mem_global_op_ld(1トランザクションあたりの有効バイト=コアレス度)。PyTorchなら先に torch.profiler でカーネル別の時間を取り、上位数本にだけ降ります。紙の上の QQ を出してから実測を見る順番を守ること。 見当なしに数字を眺めても、良いのか悪いのか判断できません。

触るパラメータ。 ブロックあたりの共有メモリ量(164 KBを超えるには cudaFuncAttributeMaxDynamicSharedMemorySize の明示設定が要る)、__launch_bounds__ や -maxrregcount によるレジスタ数の上限、タイル形状(Tritonの BLOCK_M/BLOCK_N/BLOCK_K)、Ampere以降の非同期コピー cp.async、HopperのTMA。どれも「速い引き出しに何をどれだけ置き、いつ運び始めるか」を決めるつまみです。

知らないと事故になる落とし穴。

設計レビューの定番。 「系列長4096のattentionで、スコア行列をfp16でHBMに置くと1ヘッド何バイトか。それは何を意味するか」——33.5 MB、ヘッド数とバッチを掛ければ即座にGB級、しかも捨てる中間結果。この見積もりを即答できることは、FlashAttentionを使う理由を自分の言葉で説明できることと同義です。

まとめ

次は、そのバイト自体を小さくする方向——数値表現を削る量子化を、精度との交換レートとして見ていきます。

参考文献

  1. FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. arXiv:2205.14135論文ページ·PDF

本記事は上記論文の本文にもとづいて執筆しています。数値・主張は原典を優先してください。

運営: TechAthletes · 編集方針 · 出典の付け方 · 訂正の記録

コメント

コメントにはログインが必要です