ハードウェアとテンソル 入門

テンソルとはポインタと形状とストライドであり、それが一次元しかないメモリの上に載る。GPU がうるさく言うこと —— 連続性、コアレッシング、タイル化、次元にこうあってほしいという形 —— はすべてこの一事から出てくる。このページの数値はすべて隣の図が計算しているので、どこまで動かしても主張は崩れない。

01

テンソルはポインタと形状とストライド

メモリは一次元しかない。深層学習の矩形はすべて、線の上に載せた算術だ。

A[2, 1] と書いても、機械は「2 行目」を探しに行かない。 2 行目など存在しない。あるのはひとつのバッファの先頭を指すポインタと、それをどう読むと決めたかを述べる形状と、軸ごとにどれだけ跨ぐかを述べるストライドだけだ。バッファに沿ってスライダーを引き、矩形のどのマスに当たるかを見てほしい:

スロット 0 には A[0, 0] が入る

走査が一直線であることに注目してほしい。バッファは最後までスロット順のままで、折れているのは矩形のほうだ。1 行は 4 スロット。行 1 がから始まる理由はただ一つ、その前に 4 列あるからだ。

この「前に 4 列」が添字の式のすべてである。行ストライドが 4 なのは行が 4 要素だから、列ストライドが 1 なのは列が隣り合うから、そしてオフセットは掛けて足した結果にすぎない。i と j を決めて、その和を線の上で読んでほしい:

A[0, 0] はスロット 0 になる

青緑の括弧が i × 4 を測り、琥珀の括弧が j を足す。これが offset = i·s₀ + j·s₁ の全部で、しかも一般化する —— k 階テンソルは k 本のストライドを持ち、オフセットは添字とのドット積になる。PyTorch が保持しているのはまさにこの三つだ —— storage()、stride()、storage_offset()。以下の操作はすべて、このどれか一つを書き換える。

ストライドが (4, 1) でなければならない理由はない。 NumPy と PyTorch はそれを既定にし、Fortran、MATLAB、そして BLAS の子孫はどれも (1, 3) を既定にする。両者を切り替えて、ビューの最初の行がどの 4 スロットに落ちるかを見てほしい:

ストライドは 4 と 1

行優先は同じ行の列を隣り合わせに置くので、行はバッファの連続した区間、列は跨いだ区間になる。列優先はその逆だ。これは好みの問題ではない。転置された被演算子に対して torch.matmul が先にコピーを取ることがあるのはこのためであり、列優先の cuBLAS にわざと転置した B を渡すのもこのためだ。

アドレスにはまだ足りないものがある。スロット一個の幅だ。ストライドは要素を数え、ハードウェアはバイトを数える。dtype はその為替レートだ。 2026 年のモデルが実際に使う三つの幅を辿り、同じ十二要素が縮む様子を見てほしい:

fp32 なら 700 億パラメータが 260.8 GiB

バッファが短くなる一方で形状は動かない。三行とも (3, 4) のままで、変わったのは element_size() だけだ。これに本物のパラメータ数を掛けると、もう細部ではなくなる —— 700 億の重みは fp32 で 260.8 GiB、bf16 で 130.4 GiB、。H100 が積むのは 80 GB だ。

02

ビューはストライドの算術にすぎない

転置、スライス、ブロードキャスト。どれもバイトを触らない —— そのうちの一つが触らざるをえなくなり、コピーが黙って現れるまでは。

アドレスが式になった時点で、書く操作のほとんどはデータではなく式を書き換える。転置は何も動かさない。形状の二項とストライド組の二項を入れ替えるだけだ。A と Aᵀ を切り替え、ストレージがぴくりとも動かないのを見てほしい:

ストライドは 4 と 1 になり、何もコピーされていない

配線に注目してほしい。A ではビューが読み順のまま真下に降り、Aᵀ では編み込まれる。転置をその順で読むと、一歩ごとに 4 スロット跳ぶからだ。バイトは同一で、走査は同一ではない —— この差がこのページの残り全部の性能の話になる。A がどれだけ大きくても A.t() は O(1) だ。

スライスは同じ手品にノブが二つ増えただけだ。ある列から step 列おきに残すことは、オフセットに start × s₁ を足し、s₁ に step を掛けることに等しい。二つのスライダーを動かし、残された列が線のどこに落ちるかを見てほしい:

オフセット 0、列ストライド 1

どちらの書き換えも算術なので、A[:, 1::2] の生成には代償がない —— そして A とストレージを共有するので、スライスに書けば A が変わる。この別名づけが無料のビューの代価であり、一群のバグがそこに住む。コピーだと思い込んだスライスと、そこへ黙って書き続けるループだ。

三つ目の書き換えが変わり種だ。ストライド 0 は、すべての添字が同じアドレスに落ちることで、一行の記憶に好きなだけの行を演じさせる。行数を上げて、ストレージが大きくなるのを拒む様子を見てほしい:

4 スロットの記憶から 1 行を読む

配線が一点に集まるのを見てほしい。バイアスを 4,096 行にブロードキャストする代価は 16,384 スロットではなく 4 スロットだ。(4096, 768) の活性化に x + b と書いても b のために何も確保されないのはこのためである。ストライドは 0、メモリは一行、 4,096 回の読み直しはキャッシュがほぼ無料にする。

ではビューはいつ無料でなくなるのか。バッファに強いる走査が毎歩 +1 でなくなった瞬間だ。それが連続の定義であり、.view() が形状の読み替えに応じる前に確かめる条件でもある。 step を広げ、走査が壊れるところまで行ってほしい:

連続なので view はビューを返す

step = 1 なら走査は +1 十二回で、.view(-1) は無料だ。では隙間ができ、.view() は例外を投げ、.reshape() は黙って確保してコピーする。その静けさが罠だ。reshape は寛容なので皆がそれを掴み、 O(1) の呼び出しが百万回まわるループで O(n) の確保に変わる。x.contiguous().view(-1) ならコピーは自分が置いた場所で起きる。

03

走査が実際に払う代価

メモリシステムは要求したバイトを売らない。売るのは32 バイトのセクタで、それより小さい単位はない。

ここまでストライド走査の代価は掛け算一回だけだった。ハードウェアは同意しない。 4 バイトを届けるメモリシステムなど地上に存在しないからだ。NVIDIA の GPU はグローバルメモリ要求を32 バイトのセクタ単位で処理する。バッファに沿ってスライダーを動かし、fp32 を一個要求して、実際に届くものを見てほしい:

要素 0 はセクタ 0 にある

ついでに運ばれた七要素に注目してほしい。メモリバスを渡り、L2 に着き、そして捨てられた。それ自体は些事にすぎない —— 次に欲しいのが要素 1 なら、そのセクタはもう手元にある。代価が姿を現すのは、要求するものがどれも隣り合っていないときだけだ。

ここで機械のもう半分が入ってくる。GPU はロードを一本ずつ出さない。warp 単位で出す —— 32 レーン、1 命令、32 アドレスだ。アドレスが連続なら warp 全体で 4 セクタ。跨いでいれば 32 セクタになりうる。ストライドを開き、セクタが増えていくのを見てほしい:

ストライド 1 では 4 セクタを運び、100% しか使わない

レーンが散り、その背後でセクタが埋まっていくのを見てほしい。ストライド 1 では warp が 128 バイトを欲しがり、ハードウェアが 128 バイトを運ぶ。でも欲しいのは 128 バイトのままなのに、運ばれるのは 1,024 バイト —— レーンごとに専用のセクタで、流量の八分の七が捨てられる。これがコアレッシングであり、メモリ律速のカーネルにおける最大の梃子だ。

この関係は厳密で、二つの姿勢ではなく一本の曲線として見る価値がある。全レーンが自分のセクタを持つまで効率は 1/ストライドで、そこから先はもう失うものがない —— 同じストライドのスライダーを引き、曲線の上で読んでほしい:

ストライド 1 は流量の 100% を使う

効率の落ち方が双曲線なので、最初の倍化がいちばん高くつく。ストライド 1 から 2 で帯域の半分が消え、4 から 8 ではさらに八分の一しか減らない。 H100 では実効帯域 3.35 TB/s との有効帯域の差になる。算術は同一だ。カーネルは何も変わっていない —— 変わったのは要求した順序だけである。

だから同じ総和を二通りに書けば、二つの別のプログラムになる。行優先の行列を行に沿って足せば連続アドレスを読み、列に沿って足せばセクタごとに一要素しか読まない。向きを切り替え、セクタが点くのを見てほしい:

1 セクタから 8 要素を読んだ

この行列の一列は 32 バイト間隔の八要素なので、読むと八セクタすべてに触れ、運ばれた 256 バイトのうち 32 バイトしか使わない。行版は八要素ごとに一セクタだ。同じループ、同じ FLOP、八倍の流量 —— しかも 4096 × 4096 では、列版は L2 でも毎回外す。

04

行列積はメモリの工程表である

算術はほぼ無料だ。A の行と B の列をそこまで運ぶことが、工学のすべてになる。

行列積は Transformer で最も計算密度の高い操作だが、素朴に書けばやはりメモリ律速になる。出力の各要素はA の一行とB の一列のドット積だ。出力の上をドラッグして、一要素が何を読まねばならないかを見てほしい:

C[0, 0]. ドラッグで出力要素を選ぶ。矢印キーは一マスずつ動かし、Home は C[0, 0] に戻す
C[0, 0] には 6 回の積和が要る

隣の出力がまたA の同じ行を要ることに注目してほしい。再利用なしで書けば、n × n の行列積は 2n³ FLOP を行い2n³ 個の被演算子を読む —— 1 個 2 バイトなら 1 バイトあたり 0.5 FLOP だ。H100 の演算器が休まずに済むには 295 が要る。

この比には名前がある。演算強度、すなわち FLOP を DRAM から運んだバイト数で割った値だ。それを上げる手段が再利用である。スレッドブロックが出力の T × T ブロックを保持し、 A と B の対応する帯をそこへ流せば、各被演算子は n 回ではなくn/T 回しか読まれない。タイル辺長を上げ、曲線から強度を読んでほしい:

辺長 1 のタイルは 1 バイト 0.5 FLOP を買う

両軸とも対数で、そう明記してある。だからここでの直線は一定の傾きではなくべき則だ。bf16 では強度は T/2 —— 辺長が倍になるたび流量が半分になる。算術がようやく限界になるリッジに届くには T = 591 が要り、SM 一基が抱えられる量をはるかに超える。なぜ 128 で済むのかはすぐ後で戻ってくる。

リッジはカーネルではなく機械の性質だ。ピーク演算をピーク帯域で割れば、それより下ではメモリが天井、それより上ではテンソルコアが天井になる強度が出る。強度を動かし、どちらの天井の下にいるかを見てほしい:

1 バイト 0.5 FLOP では天井は 1.7 TFLOP/s

1 バイト 295 FLOP を下回れば屋根の斜面の上にいて、スループットはダイに載るテンソルコアの数によらず帯域 × 強度になる。素朴な行列積はで 1.7 TFLOP/s に頭打ちする。定格は 989 だ(H100 SXM5、NVIDIA の 2023 年データシート)。softmax、LayerNorm、そしてあらゆる要素ごとの演算は、恒久的にその斜面に住む。

出力をブロックとして描けば再利用は一目瞭然になる。各ブロックはA の横帯を一本とB の縦帯を一本読むので、行列積全体では各被演算子を n/T 回読む。辺長を変え、DRAM の流量を見てほしい:

辺長 128 のタイルは 2.0 GiB を運ぶ

倍になるたび流量が半分になるのを見てほしい。T = 128 で 2.0 GiB、 256 で 1.0 GiB、512 で 512 MiB —— 三つの行列が実際に抱えるのは 96 MiB だ。その差を吸うのが L2 である。H100 は 50 MB 積んでおり、だから実際の 128 タイルのカーネルは、この模型の予想よりずっと理想に近い数値を出す。

タイルはいくらでも大きくはできず、その壁は共有メモリだ。T × T の被演算子タイル二枚を 1 要素 2 バイトで持ち、次のロードが今の計算と重なるよう二重緩衝すると 8T² バイトになり、 SM が確保できるのは 228 KiB だ。辺長を壁の向こうへ押してほしい:

辺長 128 のタイルは共有メモリに収まる

T = 128 ならタイル一対は 128 KiB で、累算器の余地を残して収まる。では 288 KiB になり、カーネルはそもそも起動しない —— 遅いカーネルではなく、起動時のエラーだ。実際に収まる最大の正方タイルは 168 で、それでも cuBLAS が 128 を選ぶのは、2 の冪がテンソルコアの MMA 形状を割り切り、累算器にレジスタを残すからである。

05

GPU が欲しがる形には理由がある

何かの倍数でない次元は、代金を払ってから捨てる算術になる。

タイルは幅 128 で、ハードウェアはその一部だけを計算できない。幅 512 列の出力はぴたりと覆われるが、520 列なら 640 列を計算して 120 列を捨てることになる。出力を広げ、パディングが現れるのを見てほしい:

N = 512 では最後のタイルの 0% がパディング

無駄が境界をどれだけ越えたかに比例しないことに注目してほしい。では最後のタイルの 94% がパディングで、N = 632 では 6% だ。これがタイル量子化であり、実運用の Transformer の次元が 64 や 128 の倍数になっている理由でもある —— GPT-2 の 50,257 語彙を 50,304 に詰め、ヘッド次元を 64 にし、FFN 幅をモデル幅の四倍にする。

その上にもう一段、より粗い量子化がある。タイルは SM に配られ、H100 には 132 基ある。132 タイルはちょうど一波、133 タイルは二波で、二波目は SM を一基だけ占め、残り 131 基は遊ぶ。タイル数を波の境界の向こうへ歩かせてほしい:

132 タイルでは機械の 100% しか動かない

崖を見てほしい。仕事がタイル一つ増えただけで稼働率は 100% からへ落ち、264 タイルまで戻らない。これがウェーブ量子化で、測る形が測るカーネルと同じくらい効くという話でもある。波数が増えれば影響は縮み、八波では同じ一タイルの代価は半分ではなく 12% になる。

最後の形の問題はカーネルではなくデータから来る。系列の長さはばらばらで、テンソルは矩形だ。だからバッチは最長の一本に合わせて詰められ、あとからマスクが余りを捨てる。バッチを長さのバケットに切り、無駄な算術が落ちるのを見てほしい:

1 バケットでは計算の 55% がパディング

バケットが一つなら、このバッチの FLOP の 55% は誰も使わないトークンを計算している。長さで並べ替えてだけで 12% になる。これが長さでまとめるバッチングだ。可変長アテンションが矩形を捨てて cu_seqlens を取るのも同じ理由 —— つまりはストライドだ。

06

持ち帰るもの

不変条件が一つ、定数が三つ。

(i, j) の要素とは base + (offset + i·s₀ + j·s₁) × itemsize のバイトであり、コピーしない操作はこの四数を書き換えるだけだ。定数は 32 バイトのセクタ、幅 128 のタイル、132 基の SM。