DeepSeekによって開発された行列計算ライブラリDeepGEMMの実装を解説していく。 なるべくGPGPUプログラミングやHopperアーキテクチャ、ハイパフォーマンスな行列計算に関する前提知識を問わないように心がけて書きます。
DeepSeekによって開発されたNVIDIA Hopper GPU向けの行列積計算ライブラリ。
https://github.com/deepseek-ai/DeepGEMM
2025年2月末くらいに公開された。BLASよりもシンプルなインタフェースで、300行程度の小さなカーネル関数にも関わらずPytorchの行列積よりも速いとのこと (README参照)。
インタフェースはPythonで、中身はC++ & CUDAだが、C++/CUDA部分は実行時にJITコンパイル・共有ライブラリを動的ロードした上で呼び出されるため、ビルドプロセスに影響を与えない。
行列積 とするとき、 と はe4m3のfp8で、 はbfloat16 (いずれもPytorchによる実装) を想定している。 READMEでも触れられているが、用途としてはDeepSeek-V3やR1などのLLMにおける量子化済みパラメータのハイパフォーマンスな推論のようだ。
行列積はコンピュータで計算すると計算量がかなり多くなる。 例えば行列積 を考え、 の形が , Bの形が とすると、 の形は になる (これは行列積の定義) 。 具体的には、Cのすべての要素について を計算する。
これだけだとまあどうってことはないのだが、これをシンプルにC言語で実装しようとすると下記のようになる。
for (int i = 0; i < M; i++) { // Cの行
for (int j = 0; j < N; j++) { // Cの列
for (int k = 0; k < K; k++) { // 実際の各要素の計算
C[i][j] += A[i][k] * B[k][j];
}
}
}
見ての通りの三重ループで、計算量は (正方行列だとすると) の計算量になる。これが行列積はナイーブに実装すると計算量が多いと言われる理由だ。 行列積はディープニューラルネットをはじめ、割と様々なジャンルで使われる計算なのでここが遅いと困る。ここについては色々と最適化テクニックがあるのでまずそれを見ていこう。
まだ全然DeepGEMMの話は出てこないのだが、とりあえずCPU向けの最適化から解説していく。いくつかの前提:
[[a, b], [c, d]] のような2*2行列 (二次元配列) は、メモリ上で [a, b, c, d] と並ぶ。これが [a, c, b, d] とならぶのがColumn-major。CやnumpyはRow-majorで、FortranやMATLABはColumn-majorである。詳しくはRow- and column-major orderなどを参照CPU向けの行列積演算最適化の基本的な戦略としては、シングルスレッドでの性能を限界まで上げて、その後マルチコアに対応する、という流れがよくあると思う。
行列積の場合、 の計算量自体はアルゴリズム的に避けられない。なので、シングルスレッドの性能を上げるために施す最適化はほぼ「メモリアクセスを減らす」ことが目的となる。メモリアクセスを減らしながら行列の積を計算するためには、次のようなことをすれば良い。
これでメモリアクセスが減らせる。また実際は、SIMD命令で可能な限りまとめて計算することで、さらに速くできる。 この、「メモリアクセスを減らす」と「SIMDをなるべくいっぱい使う」の両方を、プログラムのメンテナビリティをあまりにも損ねない範囲で実装するのが目標になる。
コードを再掲する。
for (int i = 0; i < M; i++) { // 行列Cの行
for (int j = 0; j < N; j++) { // 行列Cの列
for (int k = 0; k < K; k++) { // 実際の各要素の計算
C[i][j] += A[i][k] * B[k][j];
}
}
}
ループがi -> j -> kという順番になっている。一番内側がkのループになっていて、これは行列Aのi行目の列数であり、行列Bのj列目の行数である。当然だが、Cにセットする値を計算するとき、Aの各行・Bの各列はこのようにkに依存する形になる。
仮にA, Bを次元数4の正方行列として、i=0, j=0のときに最も内側のループはなにを計算するかというと:
// 初期状態
C[0][0] == 0
// k = 0
C[0][0] += A[0][0] * B[0][0]
// k = 1
C[0][0] += A[0][1] * B[1][0]
// k = 2
C[0][0] += A[0][2] * B[2][0]
// k = 3
C[0][0] += A[0][3] * B[3][0]
このように計算する。こう見るとわかりやすいが、Cについては同一メモリアクセス (= レジスタアクセス = 最速) 、AについてはRow-majorなのでメモリへのシーケンシャルアクセス、Bについては次元数4だとするとB[0], B[4], B[8], B[12]のような飛び飛びのメモリアクセス (ストライドアクセス) となる。Aのシーケンシャルアクセスは実際はキャッシュラインへのアクセスとなるのでメモリアクセスは発生していない。Bは、このような小さな行列だと十分キャッシュラインに乗る可能性が高いが、行列のサイズが大きくなるとキャッシュに乗らずメモリアクセスが発生してしまう。なので、このコードはBへのアクセスが遅いと言える。
CPUでの行列積の最適化は基本的にこのような、メモリアクセスを減らしていくことを積み重ねていく。
この場合は、ループ交換で次のように最適化できる:
for (int i = 0; i < M; i++) { // 行列Cの行
for (int k = 0; k < K; k++) { // 実際の各要素の計算
for (int j = 0; j < N; j++) { // 行列Cの列
C[i][j] += A[i][k] * B[k][j];
}
}
}
こうすると、Aは固定値なのでレジスタアクセス、B・Cについてキャッシュラインアクセスになる。
おさらいすると、ロードやストア命令は、レジスタとメモリ間でのデータの移動の命令であり、これが内部的にはメモリではなくキャッシュへのアクセスで済むことがある。レジスタへのアクセスはこの世で一番速いので、まず可能な限りデータをレジスタに乗せておいて、それを使い倒すようなコードにする。ただ汎用レジスタは16個くらいしかないのでそれを考慮してプログラムを書く必要がある。
例えば次のようなコードになる。
// 行のループ
for (int i = 0; i < M; ++i) {
// 列のループ
for (int k = 0; k < K; k += 4) {
// Aの連続する4つの要素をレジスタに保持 これがレジスタブロッキング
// 4で割り切れない次元数を考慮している
float a0 = (k < K) ? A[i*K + k] : 0.0f;
float a1 = (k+1 < K) ? A[i*K + k+1] : 0.0f;
float a2 = (k+2 < K) ? A[i*K + k+2] : 0.0f;
float a3 = (k+3 < K) ? A[i*K + k+3] : 0.0f;
// 列のループ
for (int j = 0; j < N; ++j) {
// Bへのアクセス
float b0 = (k < K) ? B[k*N + j] : 0.0f;
float b1 = (k+1 < K) ? B[(k+1)*N + j] : 0.0f;
float b2 = (k+2 < K) ? B[(k+2)*N + j] : 0.0f;
float b3 = (k+3 < K) ? B[(k+3)*N + j] : 0.0f;
// Cにセット
C[i*N + j] += a0 * b0;
C[i*N + j] += a1 * b1;
C[i*N + j] += a2 * b2;
C[i*N + j] += a3 * b3;
}
}
}
ループを4つずつにアンロールしながら、Aをレジスタブロッキングしつつ処理している。これはCの各行を左から右へ処理しているが、元のコードとは異なりAの各要素を保持するメモリ/キャッシュ参照が不要で、レジスタ参照ができるようになっているイメージ。
その時使いたいデータをなるべくレジスタに置きつつ、次に使いそうなデータがL1, L2, L3キャッシュに置かれるようにするのがキャッシュブロッキングだ。 プログラムとしては、レジスタブロッキングでは明示的に変数にセットすることで値がレジスタに保持されるようにしたが、キャッシュブロッキングではそういうことはしない。読み出された値は基本的にCPUによって自動でキャッシュに保持されるためだ。したがって、キャッシュとしてブロッキングしたいメモリサイズだけを読み取るようなループを重ねていくことで、内側のループアクセスはすべてキャッシュヒットし、それを使い終わったら外側のループがどかっとずれるようなコードを書くことになる。
ブロッキングサイズは、演算対象のデータの大きさと、使いたいCPUのキャッシュメモリのサイズ、そして計算したい行列の大きさによって決めることになる。具体的なコードはこんな感じになる。
// L3キャッシュブロッキング
const int BL3_M = 128;
const int BL3_K = 128;
const int BL3_N = 256;
// L2キャッシュブロッキング
const int BL2_M = 64;
const int BL2_K = 64;
const int BL2_N = 128;
// L1キャッシュブロッキング
const int BL1_M = 16;
const int BL1_K = 16;
const int BL1_N = 32;
// レジスタブロッキング
const int REG_K = 4;
// L3キャッシュブロッキング
for (int i3 = 0; i3 < M; i3 += BL3_M) {
for (int k3 = 0; k3 < K; k3 += BL3_K) {
for (int j3 = 0; j3 < N; j3 += BL3_N) {
// L3ブロックのサイズ制限
int i3_end = std::min(i3 + BL3_M, M);
int k3_end = std::min(k3 + BL3_K, K);
int j3_end = std::min(j3 + BL3_N, N);
// L2キャッシュブロッキング
for (int i2 = i3; i2 < i3_end; i2 += BL2_M) {
for (int k2 = k3; k2 < k3_end; k2 += BL2_K) {
for (int j2 = j3; j2 < j3_end; j2 += BL2_N) {
// L2ブロックのサイズ制限
int i2_end = std::min(i2 + BL2_M, i3_end);
int k2_end = std::min(k2 + BL2_K, k3_end);
int j2_end = std::min(j2 + BL2_N, j3_end);
// L1キャッシュブロッキング
for (int i1 = i2; i1 < i2_end; i1 += BL1_M) {
for (int k1 = k2; k1 < k2_end; k1 += BL1_K) {
for (int j1 = j2; j1 < j2_end; j1 += BL1_N) {
// L1ブロックのサイズ制限
int i1_end = std::min(i1 + BL1_M, i2_end);
int k1_end = std::min(k1 + BL1_K, k2_end);
int j1_end = std::min(j1 + BL1_N, j2_end);
// 実際の計算
for (int i = i1; i < i1_end; ++i) {
// k方向のレジスタブロッキング
for (int k = k1; k < k1_end; k += REG_K) {
int k_limit = std::min(k + REG_K, k1_end);
// Aをレジスタに保持
float a_reg[REG_K];
for (int kr = 0; kr < REG_K && k + kr < k_limit; ++kr) {
a_reg[kr] = A[i*K + (k + kr)];
}
// B, Cにアクセスするループ
for (int j = j1; j < j1_end; ++j) {
// 4つの乗算を同時に行い、結果を累積
for (int kr = 0; kr < REG_K && k + kr < k_limit; ++kr) {
C[i*N + j] += a_reg[kr] * B[(k + kr)*N + j];
}
}
}
}
}
}
}
}
}
}
}
}
}
i3, k3, j3のループがL3キャッシュのブロッキング、i2, k2, j2のループがL2キャッシュのブロッキング、i1, j1, k1のループがL1キャッシュのブロッキングをそれぞれするためのストリップマイニングになる。だいぶややこしいコードになっているが行列積だと割とよく見る見た目のコードではある。
最も内側のループで C[i*N + j] += a_reg[kr] * B[(k + kr)*N + j]; のようなことをしているのは、積和命令 (fma, fused multiply add) で置き換えられる。
の計算は、 を計算して、その後それに を足すが、この2つの命令をfmaという一つの命令で完了できる。これはCPUにFMA用の回路が実装されているためだ。また、計算を2回に分けると丸めが2回発生するが、FMAではまとめて1回の丸めにできるので精度的にも有利になる。
こうすればできる。
C[i*N + j] = std::fma(a_reg[kr], B[(k + kr)*N + j], C[i*N + j]);
メインメモリからキャッシュメモリにデータを乗せておくような命令がプリフェッチ。ハードウェアプリフェッチとソフトウェアプリフェッチがある。 IntelのCPUにはハードウェアプリフェッチャーが10種類くらいあり、L1キャッシュ用、ストライドパターンを認識して自動でプリフェッチするもの、現在のキャッシュラインの次のラインをプリフェッチするものなど色々ある。
ソフトウェアプリフェッチャーはprefetch0やprefetchw命令などで、強制的にハードウェアプリフェッチャーを起動できるものだ。
プリフェッチしなくても必要なタイミングでデータは読み込まれるのだから意味ないのでは、と思うかもしれないが、CPUはパイプライン処理しているのでバックグラウンドでプリフェッチができたり、キャッシュコントローラーが暇な時間を有効活用できたり、そもそもプリフェッチ専用回路が最近のCPUにはあったりするので、明示的にやっておくとお得らしい。
キャッシュブロッキングを行っている前提では、連続するキャッシュラインの読み込みはハードウェアプリフェッチが自動でやってくれる可能性が高いので、初回のデータ読込をソフトウェアプリフェッチでやっておくようなパターンがあるらしい。
これだけ色々やるとだいぶメモリへのアクセスが抑えられるので、ベクトル化したくなる。 ベクトル化はSIMDと呼ばれているやつとたぶん同じ意味で、配列などの連続したメモリをSIMD用レジスタに乗せ、各データに同じ演算をまとめてバッチのように実行してくれるものだ。
例えばある程度最近のIntelプロセッサだとAVX2というベクトル命令郡がサポートされていて、これはymmレジスタという256bitのレジスタに乗ったデータにまとめて同じ演算ができる。256bitあると単精度浮動小数点数であれば8個乗るのでスループットはここだけ取り出せばなんと8倍に向上する。
上記コードの、最内ループのこの箇所:
// 4つの乗算を同時に行い、結果を累積
for (int kr = 0; kr < REG_K && k + kr < k_limit; ++kr) {
C[i*N + j] += a_reg[kr] * B[(k + kr)*N + j];
}
ここがレジスタにあるAと、B, Cについては連続アクセスになっているのでベクトル化が可能だ。 このケースでは、コンパイラが自動ベクトル化をしてくれる可能性もある。実際ベクトル化の際はアラインメントしなければいけなかったりと色々やることはあるが使ったほうが大体速くなる。ベクトル化のコードはintrinsicを使うなりコンパイラに任せるなり色々あり、詳しくは省略。
ここまでやったコードをOpenMPなどでマルチコア対応させることが簡単にできる。行列を適当なブロックに分割してそれぞれ担当させることができる。
注意点としては、CPUはL1, L2キャッシュはコアごとにあるが、L3キャッシュは全コアで共有するようになっている。つまりシングルコアではL3キャッシュに乗せられていたデータを、複数コアで共有するようになるとあふれるということが起こる。これを前提としてL3キャッシュはブロッキングする必要がある。
行列積では、Cへのアクセスは完全に独立していてコアごとに依存関係がなく、AとBについても読み込みしかしないため、基本的にキャッシュコヒーレンシのためのコストがかからない (と理解している) 。なので、ブロッキングサイズの調整は必要だが基本的にはマルチコア化すると (他に大量のプロセスが動いたりしていなければ) 常に速くなる (と理解している) 。
上で行列積は は避けられないと書いたが、シュトラッセンのアルゴリズムでは になるらしい。
アルゴリズムについてあまり詳しいことは理解していないのだが、これが広く使われるライブラリとして実装されて使われていることは知る限りない。おそらく計算が複雑になって最適化がかえってやりにくくなるのではないかと思う。シュトラッセンだけでなくWinogradとか他にもいくつかアルゴリズムはある。
次にGPGPUでの行列積の最適化について。AMDのGPUについてはあまり良く知らなくて勉強中なので、基本的にNVIDIAのGPUについて書く。
まずGPUというのはGraphics Processing Unitの略で、元々はコンピューター・グラフィックスの処理に特化した計算機だった。ここで言うコンピューター・グラフィックスとは、コンピュータ上の仮想的な三次元空間の処理になる。代表的なのはゲームだ。
3D空間をゲームの中で表現するということはそもそも、我々の現実世界でものを眺めているような見た目を画面の中で実現することだ。このためには、3Dモデルを作ってそれを3D空間に配置するのだが、3Dモデルはたくさんの三角形を組み合わせて作られる (ポケモンのポリゴンをイメージするとわかりやすい) 。
このような3Dモデルを構成する三角形をポリゴンと呼ぶのだが、このとき3D空間ではモデルごとに独立なローカル座標と、空間全体を示すグローバル座標がある。
ここで、グローバル座標の中にある3Dモデルを移動させたいときにどうするかというと、3Dモデルを構成するすべての三角形を移動させればよいわけで、それはつまり、三角形の頂点をすべて移動させることを意味する。各頂点は (x, y, z) を持つので、これをX方向・Y方向・Z方向に移動させるためにはどのような計算をすればよいかというと、平行移動行列を掛けることで移動後の座標が得られる。平行移動行列は
こういう掛け算をすることになる。これを3Dモデルのすべての三角形のすべての頂点に対して計算しなければいけないので、けっこう大量の行列積を計算することになる。
なんとなく想像がつくと思うが、3Dモデルを構成するポリゴンが多くなればなるほどその3Dモデルは滑らかになる。なので、きれいでリアリティのある3D空間にしたければ、それだけ計算が重くなっていく。
また、3Dモデルの回転や、また最終的に画面に表示するために2Dに落とす処理、2Dの視点移動の際の座標再計算、光の反射の計算など、あらゆるところで行列計算が出てくる。
CPUは前述の理由で大量の行列積を高速に行えないので、行列計算に特化した計算機としてGPUが生まれた。
GPUは簡単に言うと、超大量の演算器を実装していて、CPUとは比べ物にならない量の並列計算ができる代わりに、ひとつひとつの演算はCPUより遅い、というものだ。行列積を例に取ると、大量の計算が必要な代わりに、 の各要素については、他の要素の計算結果への依存がないので並列に計算できる。この特性を利用して、でかい行列でも全部並列に計算すればいいじゃんというのがGPUだ。
つまりGPUは行列計算をするための専用計算機として生まれたと言ってもたぶん過言でない。そしてこの、行列計算を高速にできるという特性がディープラーニングやその他の科学技術計算にも使えるのでは？と目をつけられ、そっちの方向に進化して現在に至る。
ここからはGPUでの行列計算最適化と、そのための前提知識を見ていく。
GPUのプログラミングについて知る前に、ハードウェア的な仕組みについて知っておいたほうがわかりやすくなると思うので、まずハードウェアの話からする。
GPUはSIMT (Single Instruction, Multiple Threads) と言って、同じ命令を複数のスレッドで実行するようなパラダイムで動く。これ、SIMDと何が違うの？という話なのだが、(CPUの) SIMDは演算器を横に並べた実行ユニットで同じ命令を動かすもので、SIMTはもっと汎用的かつ軽量なプロセッサを並べている感じだと思う。実行する命令も、SIMD演算器は一個の機械語 (addとか) を実行するのに対し、SIMTではもっと大きな、関数とかプログラムと呼ばれるレベルのものを動かせる。
つまり、GPUは「同じ関数を複数のスレッドで動かすもの」であり、この点は単なるマルチコアCPUと大きく異なる。実際はパラメータを変えながらマルチスレッドしたいわけだが、この辺は、各スレッドで、自分自身のスレッド全体の中でのインデックスが取得できるので、それを使って処理を分ける感じになる。
例えば、長さ3のベクトル と を足したいとき、CPUでは
void func(int* c) {
for (int i = 0; i < 3; i++) {
c[i] = a[i] + b[i];
}
}
int c[3] = {0, 0, 0}
func(&c)
とするわけだが、GPUでは
void kernel(int* c) {
int i = get_my_thread_index(); // 各スレッドが自分のスレッドインデックスを得られる
c[i] = a[i] + b[i];
}
int c[3] = {0, 0, 0}
kernel<threads = 3>(&c)
となる。これは疑似コードであり、実際にコンパイルできるコードではないことに注意。 呼び出し側では、スレッド数を指定してカーネル関数を呼び出し、カーネル側では、各スレッドが自分のスレッドのインデックスを取得できる。各スレッドは完全にパラレルで動く。コードからforループは消える。
前述のように、SIMTでは各スレッドが同じ命令を実行するので、上で言うと kernel を各スレッドが実行することになる。ここで、異なるカーネル関数を複数書いて、それを1スレッドずつ実行すればSIMTっぽくない異なる命令を並列に動かす動きもできるのか？というと、やってみればできるし、GPUのリソースが空いている限りはちゃんと並列になる。ただし、後述するがかなりの数のコアが無駄になってしまいスループットがめちゃくちゃ下がり、やる意味があまりない、という感じになる。
ここからは、まずNVIDIA GPUの計算機ハードウェアとしての理解を深めていく。
NVIDIAのGPUにはブランドがあり、ゲーミング用のGeforce RTX、3DCGなど産業用のNVIDIA RTX / Quadro、データセンターやサーバ用のNVIDIA Tesla (今は単にNVIDIAと言う) などがある。NVIDIAは他にも普通にラップトップ作ったりや、組み込み用のJetson、車載コンピュータ用のNVIDIA Drive、医療機器用のClara、ネットワークインタフェースなど色々やっているのだが、どれも粒度が微妙に異なっており、いわゆるGPUといったらGeforce RTX, NVIDIA RTX / Quadro, NVIDIA (Tesla) があるという理解でいいはず。
このブランドの違いは、単に用途の違いを示している。用途が違うということは、計算の精度や備えている機能などが異なってくる。
一方、アーキテクチャという概念もある。ここでいうアーキテクチャはいわゆるコンピュータアーキテクチャで、計算機としての設計が異なる。例えば製造プロセスや電力効率、スケジューラやディスパッチの最適化など、アーキテクチャが違えば根本的に計算機として異なっている。
NVIDIAのGPUはこの、ブランド、アーキテクチャ、そしてその中のスペックによって一意に特定できる (と理解している) 。例えばみんな大好きGeforce RTX 4090は、ブランドはGeforce RTX、アーキテクチャはAda Lovelaceで、これだけだと他に4090 D、4080、4080 SUPER、4070、4070 Tiなど色々あり、その中の4090、という感じになる。
アーキテクチャには次のようなものがある。
| アーキテクチャ | 発表年 |
|---|---|
| Tesla | 2006 |
| Fermi | 2010 |
| Kepler | 2012 |
| Maxwell | 2014 |
| Pascal | 2016 |
| Volta | 2017 |
| Turing | 2018 |
| Ampere | 2020 |
| Ada Lovelace | 2022 |
| Hopper | 2022 |
| Blackwell | 2024 |
この中には、Geforce RTX系のプロダクトしかないものや、Quadro・Tesla系のプロダクトしかないものもある。例えばGeforce RTXにHopperアーキテクチャを採用しているGPUはない。
次に、NVIDIAのGPUの基本構造と概念について見ていく。コンピュータの仕組みを知らないとCがわからないのと同じで、GPUハードウェアのことを知らないとGPUプログラミングのことはちゃんとわからないと思っている。
図1: NVIDIA GPUの基本構造
いきなりこの世界から色が失われたみたいな図を乗っけてしまったが、これは別にふざけているわけではなく、まずGPUというもののメチャクチャ簡単な概念図がこれだ (これに相当する図がNVIDIAのドキュメントになかった) 。GPUというのはグラフィックボード上の計算機部分を指し、グラフィックボードにはGPUが使えるメモリ (device memoryとあるやつ) も載っている。永続ディスクというものはボード上にはない。
GPU側のボードはCPU側のボードとPCIeもしくはInfiniBandというNVIDIAのGPU用の接続インタフェースで接続される (民生用のGPUだとほぼPCIeを使う) 。 GPUでプログラムを動かすときの流れは、ホストでプロセス用に確保されたメモリ空間から、GPUのメモリにデータを転送、カーネルと呼ばれるGPUで動く関数を起動する。そしてGPU上でカーネルを計算して、GPUのメモリに計算結果を書き戻し、GPUのメモリからホストのメモリにデータを戻して、ホスト側で使う、という感じになる。
じゃあホストとGPUはどう通信しているの？という疑問が湧くわけだが、これは通常のデバイスプログラミングと同じように、カーネルモジュールを介して通信している。このカーネルモジュールにあたるのが「NVIDIA Driver」 というソフトウェアだ。ただ、これはカーネルツリーに取り込まれていたりはしなくて、クローズドソースのプロプライエタリソフトウェアである。
NVIDIAはNVIDIA Driverにあたるソフトウェアについては情報をほぼ出しておらず、これをインストールすればGPUが使えるよ、としか言ってないので、ここに書いてある内容は間違いを含んでいる可能性がけっこうある (他の箇所がそうではないわけではないが) 。
まずNVIDIA Driverをインストールすると insmod 的なコマンドが実行され、システムコールハンドラやカーネルモジュールのバイナリがカーネルにロードされる。またデバイスファイルが作られる。 /dev/nvidia0 がGPU本体で、 /dev/nvidiactl がその管理インタフェースみたいなものだと思うのだが詳しいことはわからない。GPU複数枚挿しをしていると、 /dev/nvidia1, /dev/nvidia2 とどんどんデバイスファイルができる。
NVIDIA GPUをプログラムから扱うには、それ専用の、Cとほぼ互換のプログラミング言語を書いてそれをNVIDIAが提供する nvcc というコンパイラでコンパイルする。このとき cuda.h という、これもNVIDIAが提供するヘッダファイルとライブラリを使ってオブジェクトファイルができるのだが、この cuda.h は内部ではioctlで /dev/nvidia0 や /dev/nvidiactl とやり取りするような命令をコンパイルしてくれる。つまり cuda.h はプログラマから見たGPUの抽象化を提供し、NVIDIA Driverは実行ファイルから見たGPUの抽象化を提供している。
次に、上の雑な図におけるGPU部分を見ていく。
まず、みんな大好きGeforce RTX 4090などで使われているAda lovelaceアーキテクチャの構造を見てみよう。 Geforce RTX 4090というのはゲーミング用のGPUとして最近までは最高スペックだった (現在は少し前に50系が出た) 。20~30万するはず。ディープラーニング用でも実際けっこう使われている。
これはまずデバイスメモリ部分を除いた、雑な図で言う「GPU」の部分の図である。中身を一つ一つ見ていこう。
ど真ん中にGPU全体で共有するL2キャッシュがある。これはデバイスメモリのデータをキャッシュするもので、CPUのL2キャッシュと役割としては似ているが、CPUのL2キャッシュはコアごとにあるのに対し、GPUでは全コアで共有する。なのでCPUでいうL3キャッシュのほうが近い。
CPUのキャッシュと大きく異なるのはその容量で、Ada lovelaceではなんと96MBもある。この一つ前の世代では6MBくらいだったのでここから大きく増えている。
GPUでもメモリIOが遅いのはCPUと同じで、最適化の観点ではL1キャッシュやL2キャッシュに以下にデータを載せるかが重要になる。
これは単にホスト側のマザーボードと通信するインタフェースである。ゲーミングPCを自作したことがある人はわかると思うがGPUとマザーボードを黒くて太い謎の線で接続するアレのことだ。
これはGPUとデバイスメモリを接続して通信するためのコントローラモジュールなはず。
これはいわゆるハードウェアアクセラレータで、ソフトウェア的に重たい処理をハードウェアで高速に実行するための専用モジュールである。 Optical flow acceleratorは私の理解だとDLSS 3に使われていて、つまり3D CGのより滑らかで美しいフレーム補完を実現するためのアクセラレータだ。
NVENC、NVDECは動画のエンコード・デコードのためのハードウェアアクセラレータだ。 Ada lovelaceはゲーミングやクリエイター向けのグラフィックス用途の製品として採用されているので、これらのハードウェアアクセラレータを搭載しているのだと思われる。ただ後述するが、テンソルコアなど科学技術計算向けのモジュールもあって、割とバランス良く作られている印象。
GPCはGraphics Processor Clusterといって、Raster Engine、TPC (Texture Processing Clusters) 、ROP (Raster Operations) をまとめたものだ。
GPCの中にTPCがあり、TPCの中にSM (Streaming Multiprocessor) がある。このSMの中に実際にスレッドを実行するコアなど色々あるのだがそれは後述する。
以前のアーキテクチャでは、GPC・TPCの概念はなく、デバイスメモリとたくさんのSMでGPUはできていたらしい。GPC、TPC、SMと階層構造になっていて、GPCの数を増減させたり、その中のTPCの数を増減させたり、その中のSMの数を増減させたりなどしてスペックの違いを表現させたりしている。階層構造の主な目的はモジュール化による歩留まり向上や電力効率のためだと思うのだが正確には理解していない。
Geforce RTX 4090には11個のGPCがある。図中には12個あるが、これはアーキテクチャの図であり、RTX 4090としては実際は11個だ。これは歩留まりを考慮しているためなんだと思う。
ラスタライズ処理を行うためのモジュール。つまりGPUの文脈では、3Dのデータを画面に移すために2Dのピクセルの集合変換する処理をするためのもの。それ以上のことを知らない。今後もよくわからないままだろう。
TPCはSMをまとめたもの、以上！という理解をしているのだが正しいのかわからない。 図中では1GPCに6TPCある。Geforce RTX 4090だと11GPCなので66TPC、ではなく、64TPCだ。これも歩留まりを考慮してのものだろう。
ついにSMに来た。これが相当大事。
SMはストリーミング・マルチプロセッサの略で、実際に計算をする部分と言って良い。例えばGeforce RTX4090では、64個のTPC中にそれぞれ2個、合計128個のSMがある。
SMの図中には4つの白枠で囲まれた領域があると思うのだが、これはワープと言う。ワープが実際の最低並列処理単位だ。ワープについては詳しくこの後で説明する。
まず、図中の「128 KB L1 Data Cache / Shared Memory」という部分について。これは物理的に同じメモリユニットをL1データキャッシュとシェアードメモリという2つの目的で共有しているメモリになる。メモリといってもいわゆるオンチップのメモリで、デバイスメモリよりも数桁高速である。このメモリ領域は128KBある (各SMで) 。
L1キャッシュはCPUのL1キャッシュと同じで、メモリ上のデータをキャッシュしておく領域のうち、L2よりも容量が小さく、かつコアに近いため速い、というものだ。
シェアードメモリはCPUにはない概念だ。シェアードメモリは、SM内で動くすべてのスレッドが共有するメモリという点でL1キャッシュと同じだが、これはプログラムから明示的に操作できる。L1キャッシュはプログラマが明示的に制御することは不可能だが、シェアードメモリはプログラマが自分で開放しない限り、確保したデータはずっと残り続ける。
逆に言うと、各スレッドが使いたいデータを適切にシェアードメモリに載せてからカーネルの実際の処理をするようにプログラミングすることで、メモリアクセスをかなり減らすことができる。この特性によって、GPUプログラミングの最適化においてシェアードメモリはかなり重要な位置を占めている。
L1キャッシュとシェアードメモリは同じメモリユニットを共有しているので、シェアードメモリでいっぱいデータを確保すると、L1キャッシュとして使える領域は少なくなる。
図中にある「Tex」のこと。 グラフィックスパイプラインにおいて、テクスチャマッピング処理を行うためのユニットである。128のSMごとにそれぞれ4つあるので全体では結構あることになる。これもGeforce系のグラフィックス処理に使われるものだがあまり詳しくないのでこのくらいで。
レイトレーシングコアはレイトレーシングという、光の物理的な挙動をシミュレートする処理を専門で行うモジュールだ。これ以上特に書くことがない。
さていよいよワープに入る。図中で4つの白枠で囲まれた領域があるがこれをワープという。これは実際にSM内に4つのワープしかないことを示している。 ワープの中にも色々あるが、まずワープというものの概念について話す。
ワープはGPUによって実行される最小のスケジューリング単位であり、32のスレッドで構成される。この32のスレッドは同じ命令を同期的に実行する。 最小のスケジューリング単位というのは、例えばカーネルAを16スレッド起動し、カーネルBを10スレッド起動したいとしても、カーネルAとBは同じワープに配置されないということだ。ワープ内では同じ命令を実行するので、カーネルAのスレッドとBのスレッドは異なるワープに配置される。つまり、カーネルAを実行するワープでは32ひく16で16コア、Bのワープでは22コアが遊んでいることになる。
では、例えばカーネルAを64スレッドで動かしたいとする。このとき、64のスレッドは32ずつの2グループに別れて、それぞれが別のワープで実行される。このとき、64のスレッドを分けた2つのワープ内のスレッドは、カーネル実行の最初から最後まで同じワープにいることになる。例えばスレッド番号1 ~ 32がワープ1、33 ~ 64がワープ2だとすると、このグルーピングはずっと固定になる。また、ワープ内のスレッド実行は同期的だが、各ワープは非同期に動くので、呼び出し元ではシンクロナイズ的な処理をする必要がある。
また上の例で、ワープ1とワープ2は、最初に割り当てられたSMで最初から最後まで実行されることが保証される。これは何を意味するかと言うと、途中でスケジューラの都合で、異なるSMに配置されたりすることがないので、シェアードメモリに確保したデータが急に見えなくなったりすることがない。
では、ワープ1とワープ2がどのSMに配置されるかというのは、これは「ブロック」によって決まる。これはここまで出てきてない言葉で、後でソフトウェアの説明をするときにも書くが、簡単に説明する。ブロックはスレッドをまとめる単位なのだが、例えばトータルで64のスレッドを動かしたいとき、64のスレッドを1ブロックで動かすのか、32のスレッドを2ブロックで動かすのか、16スレッドを4ブロックで動かすのか、1スレッドを64ブロックで動かすのか、、など、色々パターンがある。
このとき、あるブロックの中のスレッドは、必ず同じSMで動くことが保証される。例えば64スレッドを1ブロック64スレッドで動かすとき、この64スレッドは必ず同じSMで動く。ただし、ワープは32スレッドまでなので、64スレッドは2ワープに分かれることになる。
これが例えば1024スレッドを動かしたいとき。1ブロックだと32ワープ必要になる。このすべてのワープが同じSMで動くのだが、1つのSMには4ワープしかないので、32ワープを同時に動かすことはできない。そうなると、ワープスケジューラが32のワープを切り替えながらSM内で頑張ってスレッドを実行するわけだが、このときGPU全体で128個あるSMのうち127個は遊んでいるわけで、ちょっと効率が悪い感じがするだろう。 じゃあ1024を8ブロックにして、各ブロックは128スレッドにして、これを8SMに割り振れば遊びがないと思うのだが、これだとシェアードメモリも8SMに分かれるので、若干プログラムが複雑になる。こういうのをGPUプログラミングの最適化では考える必要がある。
ワープごとに必ず同じ命令を同期的に実行するということは、プログラムカウンタはワープごとにあればいいことになるのだが、実際はそうでもない。このへんちょっと正確な情報が見つけられなかったのだが、たぶんVolta以降のアーキテクチャでは、プログラムカウンタはスレッドごとに持つようになっている。
後で詳しく説明するが、ワープごとに同じ命令を実行するとなると、ワープ内の一部のスレッドだけが実行したい処理 (条件分岐などによる) があったときに、処理の実行が不要なスレッドはマスクと言って寝かせておく必要があった。
スレッドごとにプログラムカウンタを持つとこのようなケースで、マスクされていないスレッドがIOなどで待っているときにマスクされているスレッドからマスクを外して実行したりできるようになって効率が良くなる。
次はSMにおけるレジスタについて。 レジスタの役割自体はCPUと同じで、実際に演算する際のデータの置き場所となる。 違いは色々あり、まずCPUのレジスタはコアごとだが、GPUではSMごとになる。次に数が異なり、CPUでは数え方にもよるが数十個程度なところ、Ada lovelaceではなんとSMごとに16384個ある。
SMには4ワープあるので128スレッド、なので各スレッドが128個使えることになる。すべてのスレッドがこの個数制限を超えていない前提だと、ワープを切り替える際もレジスタはそのままにしておける。つまりCPUのように、コンテキストスイッチの際のレジスタからメモリへのデータ退避が発生しない。これはGPUのスイッチング時のオーバーヘッドを削減しスループットを向上させる上で重要である。
仮にレジスタを使いすぎるとシェアードメモリへの退避 (レジスタスピル) が発生する。レジスタから見るとシェアードメモリは数桁遅いので、パフォーマンスの大きなオーバーヘッドになる。
レジスタファイルという名前がついているが、レジスタとレジスタファイルは違うものだ。レジスタは物理的なメモリなのに対し、レジスタファイルはワープとスレッドから対応するレジスタのデータをIOするモジュールのことである。
ワープ内のスレッドは同じ命令を実行する。このときCPUのように、メモリIO待ちなどを考慮したワープの実行可能かどうかの状態という概念がある。ワープスケジューラはこのような状態を考慮して、実行可能なワープを取り出して実行する。 このような、待ち状態のワープと実行可能な状態のワープを切り替えながら実行して無駄をなくすことを「レイテンシの隠蔽」と呼ぶ。
このワープスケジューリングなのだが、調べてもあまり詳しい情報は出てこない。実際にはワープの状態管理やスケジューリングアルゴリズムなどもっと複雑だと思うのだが、ホワイトペーパーなどであまり公開されていないようだ。このスケジューリングは、GPUのスループットについてかなり重要な部分なので、あまり積極的にオープンにされないのかなと想像している。
ディスパッチはワープスケジューラによって実行をスケジューリングされたワープ内のスレッドを、実際に処理を行うコアに転送的なことをするモジュール。転送先はここから説明するCUDAコア、テンソルコア、LD/ST、SFUになる。
さていよいよCUDAコアに、、
SMの中にCUDAコアがある。 CUDAコアとは何かというと、スレッドを実際に実行する部分で、デバイスコードと呼ばれる命令 (CPUで言う機械語みたいなもの) を実際に実行する演算装置のことだ。なのでCPUでいうコアと考えてよい。
ワープには32スレッドと言っていたが、これに対応するのが32のCUDAコアだ。ただし、SMの図を見ると「FP32 / INT32」のコアが16個、「FP32」のコアが16個ある。ワープでの32スレッドが単精度浮動小数点数演算の場合はこれらのコア全てで演算ができるが、32ビット整数の場合はINT32のコアが必要なため16コアしか使えない。このときスループットは半分になってしまうのだが、普通のGPUのワークロードでは浮動小数点演算や、浮動小数点数と整数の混合演算がメインとなるので特に問題とならないことが多い。
次はテンソルコアだ。これは名前の通りテンソルの処理をハイパフォーマンスにできるコアである。これはGeforce RTX 4090にはSMごとに4コア、GPU全体で512コアある。AIやHPCで使われる行列演算に特化したハードウェアだ。
Ada lovelaceのテンソルコアはFP32、FP16、BF16、そしてe5m2とe4m3両方のFP8に対応しているので、学習時や推論時の量子化にも利用できる。また混合精度演算をハードウェアでサポートしている。 行列積及び行列積和が非常に高速に行えるので、特にディープラーニング系の計算に有利だ。
実際に計算するときは行列をブロックに分けて複数SMの複数テンソルコアに分配して演算するのだが、これはNVIDIA製のライブラリで実装されるので自分でプログラミングすることはあまりない。
これはロード・ストアユニットで、メモリとデータをやり取りしたいときに使われる。
SFUはSpecial Function Unitの略で、sin, cos, expなどの計算をハードウェア実装したもの。
最後に、図には出てこなかったがデバイスメモリについて。VRAMなどとも言われる。これはCPUのメモリのようなもので、揮発性かつディスクよりは早く、レジスタやキャッシュよりは遅いというものだ。GPUの箱にメモリは24GBとか10GBとか書いてある時、それはこのデバイスメモリのことである。
デバイスメモリはCPUのメモリとは違って、内部にいくつかの構成要素がある。このへんがハードウェアとしてどうなっているのかはよくわかってないが、プログラマ的にはこれら全部がデバイスメモリとして抽象化されているという理解で良いはず。
グローバルメモリは一番普通のメモリというか、「メモリ」と言うときこれを意味していると思って良いと思う。全スレッドからアクセスが可能で、普通のメモリなのでキャッシュやレジスタと比べると遅い。
これは定数を置いておく領域で、ソフトウェア側で定数として宣言したメモリにホストからデータをコピーすることで値が確保可能。容量制限があり、かつ定数なので読み取り専用で値の変更ができないが、SM側で必ずキャッシュされるので、定数をいろんなスレッドから参照するのであればここに置くのが最も速い。
テクスチャメモリはCUDAのtexture型もしくはテクスチャオブジェクトを使って宣言し、グローバルメモリの領域とバインドすることで使用可能になる。これはGPU側からは読み取り専用になっている。 テクスチャメモリのメリットは2次元・3次元のアクセスパターンにハードウェア的に最適化されていることだ。行列や三次元テンソルでは、例えばストライドアクセスしながら列を読み取るとか、三次元で言うと奥行き方向に移動したりとか、メモリ的にシーケンシャルでないアクセスをしなきゃいけないことがよくあって、そのときにうまくアクセスしたメモリの周囲をキャッシュしてくれたり、プリフェッチうまいことやってくれるたりする機能がハードウェアとして備わっている。またハードウェアフィルタリング機能で画像処理やテンソル処理をアクセラレーションできるらしいのだが、この辺は詳しく理解していない。
サーフェスメモリは、テクスチャメモリと大体同じだがカーネル側からも書き込みができるというもの。surface型もしくはサーフェスオブジェクトとして利用できる。ただハードウェアフィルタリング機能はないらしい (よくわかってない) 。
ローカルメモリはけっこう特殊で、まずCUDAから確保したりアクセスしたりするものではなく、CUDAコンパイラが自動で使うものだ。レジスタが不足したり、でかいサイズの配列や動的にサイズが決まる配列などが置かれる。あるデータがローカルメモリにあるかグローバルメモリにあるかは、プログラマは意識しなくて良いようになっている。また、このデータはグローバルメモリとは違って、そのデータを使っているスレッドからしかアクセスできない。基本的にはあるスレッドがSM側で持ちきれないデータを置くところで、CPUでいうスタックやヒープみたいなものなのだが、マルチスレッド前提のGPUではスレッド専用のデータ領域が必要で、その役割を担っている。
Ada lovelaceについて色々と説明してきた。次に、DeepGEMMが対応しているHopperアーキテクチャについて見ていく。Ada lovelace編を読んでいる前提で既に書いたことは省略して書くので、先に読んでおくほうが良い。
Hopper NVIDIA H100、H200など完全にサーバ・データセンター用のアーキテクチャで、OpenAIが使っているGPUもこれだ。今の製品の中で最も高級なもののひとつ。
NVIDIA H100 Tensor Core GPU Architecture Overviewより
Ada lovelaceと見比べるといくつか違いがあるのがわかる。
まずPCIeのバージョンが4から5になっている。PCIeは各バージョンで色々異なるものの、4.0と5.0ではレーンごとに帯域幅が増えていると理解している。 HopperのGPUはLLMをはじめとした超大量のデータを計算しなければいけないワークロードを見越して作られていて、基本的にはIOが律速になる。なのでPCIeの帯域幅にも余裕を持たせているのだと思われる。
Optical flow Accelerator、NVENC、NVDECはなくなっている。これらはグラフィックス用のものだが、Hopperは科学技術計算に使われることを想定しているので不要なのだと思われる。
NVLinkというものが生えている。これはGPU間通信のための高速インタフェースだ。 LLMなどのワークロードでは、1台のGPUではデータは乗り切らないため複数台 (数十 ~ 数百台) のGPUを使ってモデルを学習させるのだが、このときGPU間でデータ通信を行うケースがあり、それを前提としたIOがくっついている。
High Bandwidth Memory3というサブシステムによって、デバイスメモリとの通信時の帯域幅が大幅に向上しているらしい。LLM学習での帯域幅ボトルネックを解消するための改良。
L2キャッシュが2つあるように見えるが、パーティション化されて、それぞれが異なるGPCと接続されるようになっている。これはキャッシュすべきデータの局所性を高めることを目的にしているらしい。これによってキャッシュ利用が最適化される。
H100ではL2キャッシュ容量は2パーティション合計で50MBである。
H100ではGPCは8ある。GPCにはラスターエンジンがなくなっている。これもグラフィックス用途でないためだろう。
H100ではTPCは66ある。TPC内にSMが2つあることは同じ。特にここはAda lovelaceと変わってなさそうだった。
SMは色々変わっている。
NVIDIA H100 Tensor Core GPU Architecture Overviewより
なくなった。これもワークロードの違いから来るものだろう。全体的に、コンピュータグラフィックス系のハードウェアはごっそり削除されている。
命令キャッシュはワープ間共有のL1と、ワープごとのL0が存在する。
H100では、FP32専用コアが32個、INT32専用コアが16個に増えている。また、FP64用のコアも16個ある。
テンソルコア自体は4th GenerationでAda lovelaceと同じだ。テンソルコアの数もSMごとに4つで変わってない。
数がAda Lovelaceと比較して、4から8に増えている。
容量が128KBから256KBに倍増しているだけでなく、Hopperでは分散シェアードメモリが導入された。これは、複数のSMにまたがるシェアードメモリについて、デバイスメモリを経由せず直接SM間通信ができるようになったものだ。これも行列計算のユースケースをよく理解していて良い機能だと思う。
図では表現されていないのだが、先程スレッド、スレッドをまとめるブロックについて書いたが、Hopperではなんとブロックをまとめるクラスターという階層が追加されている。これはプログラミングモデル自体が拡張されていることを意味する。ブロック内のスレッドは固定のSMで動くことは変わらないのだが、ブロックをまとめる概念ができたことで、異なるSMの動作を同期させる制御ができるようになり、デバイスメモリを介さないレベルでSM間のコミュニケーションができるようになっている。
新たに増えたのがこれだ。TMAはTensor Memory Acceleratorの略。 TMAは大きなテンソルをメモリからシェアードメモリにロードしたり、メモリに書き戻したりするのをこうそくにやってくれるやつである。
通常、行列やテンソルは、実態としてはメモリ上に1次元で保存されている。これにアクセスするためにはストライドを考慮してアドレス計算をする必要がある。 TMAでは、テンソルの次元数やその中のコピー対象のブロック座標を指定することで、コピーしたいアドレス指定のオーバーヘッドを大きく削減する。
NVIDIA H100 Tensor Core GPU Architecture Overviewより
さらにわかりやすいのがこの図だ。
NVIDIA H100 Tensor Core GPU Architecture Overviewより
A100という一個前のGPUでは、各スレッドがアドレス生成して読み出しするところを、H100ではテンソルの次元数やブロックサイズなどを指定することでTMAがハードウェア側でアドレス計算してデバイスメモリからシェアードメモリに載せてくれる。
図中にはpaddingも考慮されているので、おそらくテンソルの転置やインデクシングなどでテンソルの実態がストライド違いのビューだったとしてもうまくやってくれるのだろう。これはかなり便利だ。
もう一個増えたのがTransformer Engineだ。BERTやGPTなどの学習にTransformerをHopperで動かしたいわけだが、セルフアテンション機構は計算量が非常に多い。 私の理解だと、Transformer Engineは学習・推論時のパラメータの精度を自動でアップ・ダウンスケーリングすることで精度を犠牲にせず高速化してくれるものだ。テンソルコアであるレイヤの計算を実行した後、その出力値を分析し、次のレイヤの計算を把握した上で、精度をFP8までの可能な範囲で落とすことで通信や計算を高速化する。動的精度スケーリングとでも言えばよいか。
NVIDIA H100 Tensor Core GPU Architecture Overviewより
これめちゃくちゃ頭良くてすごいなと思う。あまり詳しく説明されている資料は見つけられなかったのだが、、
さて、これまででGPUのハードウェアについては完全理解したので、ソフトウェアの話をしていく。まずGPUの一般的なプログラミングの話をして、その後いよいよDeepGEMMがやってる最適化を解説していく。
NVIDIAのGPUでコードを動かすためには、CUDAというCを拡張した言語を書き、それを「nvcc」というCUDAコンパイラでコンパイルすることで実行ファイルが出来上がる。 CUDAの内部では上にも書いたが、NVIDIA Driver経由でioctlでデバイスファイルを操作することでデータをGPUに送ったりカーネルを起動したりできる。
NVIDIAのGPUプログラミングで最も特徴的なのはスレッド・ブロック・グリッドの概念だ。スレッド・ブロックは既に出てきているものと同じで、グリッドはブロックをまとめるものだが、グリッド = GPUと捉えて良いので、重要なのはスレッドとブロックである。
行列積のカーネルはこんな感じで書ける。
#include <stdio.h>
#include <cuda_runtime.h>
#define N 1024
__global__ void matmul(float *A, float *B, float *C) {
// スレッドのインデックスを計算
int row = blockIdx.y * blockDim.y + threadIdx.y;
int col = blockIdx.x * blockDim.x + threadIdx.x;
float sum = 0.0f;
// 行列積を計算
for (int i = 0; i < N; i++) {
sum += A[row * N + i] * B[i * N + col];
}
// 結果を格納
C[row * N + col] = sum;
}
このカーネルは、無関係な部分を省くと次のように呼び出す。
// ブロックとグリッドの次元を設定
dim3 threadsPerBlock(16, 16);
dim3 numBlocks(64, 64);
// カーネルの実行
matmul<<<numBlocks, threadsPerBlock>>>(A, B, C);
行列の次元数は1024の平方行列を想定していると思って見てほしい。このカーネルの方針は、1024 x 1024の行列Cの各要素を、1つのスレッドが担当して計算する、というものだ。
matmulがカーネルだ。関数定義に __global__ とついているが、これが、この関数はGPU上で動くよ (正確には「この関数はGPUで動くのだが、この関数を呼び出すのはCPUからでもGPUからでもいけるよ」という意味) になる。
今回、呼び出し時にはブロックを (x=64、y=64) 、スレッドを (x=16、y=16) で指定している。 このように、スレッドやブロックはx, yが指定でき、zまでの最大3次元で指定が可能。
このとき、スレッドは のような形で、合計が256スレッドになる。 これを (x=256) と、yを指定しない形で呼び出した場合、 のようにスレッドが作られる。
この256スレッドが、同様の計算で64 * 64の合計4096ブロック分作られる。
これらの情報を使って、各スレッドが自分の処理すべきCのインデックスを知る必要がある。 その処理が以下の部分だ。
int row = blockIdx.y * blockDim.y + threadIdx.y;
int col = blockIdx.x * blockDim.x + threadIdx.x;
blockIdx はブロックのインデックス、 blockDim はブロックの次元数、 threadIdx はブロックの中におけるスレッドのインデックスだ。今回は2次元で指定しているので、インデックスも2次元になる。
例えば、ブロック (x=0, y=0) のなかのスレッド (x=0, y=0) は、 になり、col も同様に0になる。
ブロック (x=63, y=63) のなかのスレッド (x=15, y=15) は C[1023][1023] の計算を担当することになる。これが行列の一番右下になる。
ブロック (x=0, y=0) を考えたとき、スレッド (x=0, y=0) は C[0][0]を担当し、スレッド (x=15, y=15)はC[15][15]を担当する。
つまり、各ブロックは 1024 * 1024 の大きな行列を左上から 16 * 16 に切り刻んだものをそれぞれ担当し、そして各スレッドはその中のひとつをそれぞれ担当するのだ。
呼び出しのコードの完全版は次のような感じだ。
int main() {
// ホスト側の変数
float *h_A, *h_B, *h_C;
// デバイス側の変数
float *d_A, *d_B, *d_C;
size_t size = N * N * sizeof(float);
// ホストメモリの割り当て
h_A = (float*)malloc(size);
h_B = (float*)malloc(size);
h_C = (float*)malloc(size);
// 行列Aと行列Bの初期化（単純化のため、ここでは単純な値で初期化）
for (int i = 0; i < N * N; i++) {
h_A[i] = 1.0f;
h_B[i] = 2.0f;
}
// GPUメモリの割り当て
cudaMalloc((void**)&d_A, size);
cudaMalloc((void**)&d_B, size);
cudaMalloc((void**)&d_C, size);
// データをホストからデバイスへコピー
cudaMemcpy(d_A, h_A, size, cudaMemcpyHostToDevice);
cudaMemcpy(d_B, h_B, size, cudaMemcpyHostToDevice);
// ブロックとグリッドの次元を設定
dim3 threadsPerBlock(16, 16);
dim3 numBlocks((N + threadsPerBlock.x - 1) / threadsPerBlock.x,
(N + threadsPerBlock.y - 1) / threadsPerBlock.y);
// カーネルの実行
matrixMul<<<numBlocks, threadsPerBlock>>>(d_A, d_B, d_C);
// 結果をデバイスからホストへコピー
cudaMemcpy(h_C, d_C, size, cudaMemcpyDeviceToHost);
// 結果の検証（最初の要素だけ表示）
printf("Result C[0][0] = %.2f\n", h_C[0]);
// メモリの解放
free(h_A);
free(h_B);
free(h_C);
cudaFree(d_A);
cudaFree(d_B);
cudaFree(d_C);
return 0;
}
cudaMalloc でGPU側のデバイスメモリ (グローバルメモリ) を呼び出し、 cudaMemcpy でホストからデバイスへ、そしてデバイスからホストへデータ転送ができる。 cudaFree でデバイス側のメモリを開放できる。
このコードは速いのか？というと、CPUよりは全然速いだろうけど、共有メモリなどを全く考慮していないので、L1L2キャッシュに偶然アクセスすることを祈るだけのコードになっており、最適化の余地は多分にある。でもシンプルで、ベースラインとしては悪くないというか、至って普通のコードという感じのものだ。
いくつか最適化に関する前提知識を書いていく。
コアレスドアクセス (Coalesced Access) は有名だ。これは、メモリアクセス時に連続したメモリにアクセスしたほうがパフォーマンスが出るというもの。 ワープは同じ命令を実行するので、ワープごとにメモリアクセスも同じタイミングでやるわけだが、このときワープ内のスレッドがアクセスする連続していると、メモリリクエストが一つのトランザクションにまとまってくれる。これがバラバラだと複数回のメモリアクセスになるので遅くなる。
「コアレスドアクセス」とは「まとめられたアクセス」を意味する。
上のコードでは、ブロック内の16*16のスレッドは、8ワープに分けられる。このとき、ひとつのワープは (x=0, y=0)から (x=1, y=15)までの32スレッドを持っているので、 (x=0, y=15) のスレッドと (x=1, y=0)のスレッドの間で次の行に飛んでしまうためコアレスでなくなる。なのでここは改善の余地がある。
バンクコンフリクトは共有メモリに対する悪いアクセスパターンを意味する言葉だ。 まず、シェアードメモリやレジスタはバンクと呼ばれ、物理的に領域が分割されている。かつ、各バンクは同時に1つのメモリリクエストしか処理できない。
このとき、連続したメモリアドレスは異なるバンクになるようになっている。例えば32個のバンクがあるとき、アドレス0番地がバンク1、アドレス1番地がバンク2、、、アドレス31がバンク32となって、アドレス32がバンク1、という感じになっている。
例えばあるワープが共有メモリにアクセスするとき、ワープ内のスレッド1が0番地、スレッド2が1番地、、とアクセスすると、バンクコンフリクトが発生しない。 しかし、スレッド1が0番地、スレッド2が32番地、スレッド3が64番地、、とアクセスすると、各スレッドが同じバンクにアクセスしてしまい、共有メモリアクセスがものすごく遅くなる。
原因は違うのだが、これもコアレスドアクセスと似ていて、つまり連続してアクセスすることで解決が可能だ。
例外は、複数のスレッドが同じアドレスにアクセスする場合はバンクコンフリクトは起きないらしい。
後はワープダイバージェンスについて。ワープ内の32のスレッドは基本的に同じ命令を実行するのだが、異なる命令を実行したくなったとする。たとえば、スレッドのインデックスが15以下なら1を足す、そうでなければ1を引く、のようなカーネルを考える。このとき、足し算の命令と引き算の命令は同時に実行できない（ワープ自体の制約）ので次のような順序でカーネルが実行される。
このように、ワープ内で異なる命令を実行しようとするとマスク状態のスレッドが発生して並列実行が妨げられることを「ワープダイバージェンス」という。これがGPUは条件分岐に弱いと言われる理由だ。
ただ、例えばマスクされていないスレッドがIO待ちなどをしているときは、マスク状態のスレッドを先に進めることができるらしい。これはプログラムカウンタのところで少し触れたが、プログラムカウンタをスレッドごとに持つとこのような最適化ができるようになる。これはハードウェアで実装されているらしい。
さてようやくDeepGEMMの話をする。ここまでの話を前提に、DeepGEMMがどのように行列積を最適化しているかを見ていこう。
DeepGemmのgemm_fp8_fp8_bf16_nt関数は、 torch.float8_e4m3fn の行列を受け取って、積を計算し、bf16にして返すものだ。
DeepGemmはJIT型になっていて、渡された行列のサイズを元に最適な設定を生成し、それを元にCUDAのソースコードを生成し、それを共有ライブラリとしてnvccコマンドでビルド、得たsoファイルをdlopenでメモリにロードして実行している。
実行時にコンパイルするので最初は遅いのだが、soはキャッシュされるので、同じサイズの行列であれば2回目以降は全くコンパイルは発生しない。これはパラメータだけ変えて同じ計算を何度も繰り返すディープラーニングの学習や推論と相性が良い。
また、実行したカーネルは内部でベンチマークしており、カーネルが更新されると速い方を選んで使うようになっている。
jit_kernels/gemm.pyではカーネルを呼び出すコードを生成し、それを実行している。
まず呼び出されるコードはこんな感じになる。
// DeepGEMM auto-generated JIT CUDA source file
#include <cuda.h>
#include <cuda_fp8.h>
#include <cuda_runtime.h>
#include <iostream>
#include "cutlass/cutlass.h"
#include "deep_gemm/fp8_gemm.cuh"
extern "C" void launch(void* __raw_lhs, void* __raw_lhs_scales, void* __raw_rhs, void* __raw_rhs_scales, void* __raw_out, int m, void* __raw_stream, int num_sms, int smem_size, int& __return_code) {
// Cast raw types (if needed)
auto lhs = reinterpret_cast<__nv_fp8_e4m3*>(__raw_lhs);
auto lhs_scales = reinterpret_cast<float*>(__raw_lhs_scales);
auto rhs = reinterpret_cast<__nv_fp8_e4m3*>(__raw_rhs);
auto rhs_scales = reinterpret_cast<float*>(__raw_rhs_scales);
auto out = reinterpret_cast<__nv_bfloat16*>(__raw_out);
auto stream = reinterpret_cast<cudaStream_t>(__raw_stream);
using namespace deep_gemm;
// Templated args from Python JIT call
constexpr auto N = 128, K = 512;
constexpr auto BLOCK_M = 128;
constexpr auto BLOCK_N = 16;
constexpr auto kNumStages = 8;
constexpr auto kNumTMAMulticast = 1;
// Make a templated GEMM
using GemmType = Gemm<N, K, BLOCK_M, BLOCK_N, 128, 1, kNumStages, kNumTMAMulticast, GemmType::Normal>;
// Launch kernel
auto tma_a_desc = GemmType::make_2d_tma_a_desc(lhs, m);
auto tma_b_desc = GemmType::make_2d_tma_b_desc(rhs);
auto tma_scales_a_desc = GemmType::make_2d_tma_scales_a_desc(lhs_scales, m);
auto tma_d_desc = GemmType::make_2d_tma_d_desc(out, m);
GemmType::run(out, rhs_scales, nullptr,
m,
tma_a_desc, tma_b_desc, tma_scales_a_desc, tma_d_desc,
stream, num_sms, smem_size);
}
GemmTypeというのはcudaで書かれたソースコード内で、テンプレートを使って生成されるカーネルの更にラッパーだ。GemmType::runが内部でカーネルを呼び出すようになっていて、このへんはそれにパラメータを渡すだけになっている。
変数的には、lhsが行列積の左側、rhsが右側、lhs_scalesがlhsのスケーリング係数、rhs_scalesがrhsのスケーリング係数、outが結果を入れる行列だ。
実際にカーネルの処理と最適化を見ていく。
一番大掛かりな最適化はこのパイプライン化かなと思った。 READMEにもあるのだが、パイプライン化のイメージはこんな感じだ。
そもそも通常の行列計算では、共有メモリを使った処理は次のような流れになる。
パイプライン化では、次のような感じになる。
このように、データを (TMAで) ロードすることを担当するワープと、計算をする担当のワープを完全に分けてしまうことで、処理がオーバーラップするようにする。これはCPUのパイプライニングとよく似ていることに気付くだろう。
実際には、ここにあるifの中身がロード担当のワープ、elseの中身が計算担当のワープのパスだ。
if (threadIdx.x >= kNumMathThreads) {
// TMA warp-group for loading data
cutlass::arch::warpgroup_reg_dealloc<kNumTMARegisters>();
// NOTES: only one thread (or warp) will be used
if (threadIdx.x == kNumMathThreads) {
// Persistently schedule over blocks
while (scheduler.get_next_block(m_block_idx, n_block_idx)) {
launch_k_iterations([&](int k_iter, auto type) {
constexpr bool kHasDivisibleStages = std::is_same_v<decltype(type), DivisibleK>;
constexpr int kNumInnerStages = kHasDivisibleStages ? kNumStages : (SHAPE_K % kFullKOfAllStages) / BLOCK_K;
DG_STATIC_ASSERT(kNumInnerStages != 0, "Invalid number of inner stages");
// NOTES: unrolling and `kNumInnerStages` are vital for performance, NVCC will try to eliminate all
// shared memory pointers, e.g. `full_barriers` registers, if all the access indices are constant
#pragma unroll
for (uint32_t s = 0; s < kNumInnerStages; ++ s) {
// Wait consumer release
empty_barriers[s]->wait((scheduler.current_iter * kNumIterations + k_iter + 1) & 1);
// Issue TMA A with broadcasting
auto& full_barrier = *full_barriers[s];
int k_idx = k_iter * kFullKOfAllStages + s * BLOCK_K;
tma_copy<kNumTMAMulticast>(&tensor_map_a, reinterpret_cast<uint64_t*>(&full_barrier),
smem_a[s], k_idx, scheduler.get_global_idx(shape_m, BLOCK_M, m_block_idx));
tma_copy<kNumTMAMulticast>(&tensor_map_scales_a, reinterpret_cast<uint64_t*>(&full_barrier),
smem_scales_a[s], m_block_idx * BLOCK_M,
scheduler.get_global_idx(SHAPE_K_SCALES, 1, k_idx / BLOCK_K));
// Issue TMA B without broadcasting
tma_copy(&tensor_map_b, reinterpret_cast<uint64_t*>(&full_barrier),
smem_b[s], k_idx, scheduler.get_global_idx<false>(SHAPE_N, BLOCK_N, n_block_idx, m_block_idx));
full_barrier.arrive_and_expect_tx(SMEM_A_SIZE_PER_STAGE + SMEM_B_SIZE_PER_STAGE + SMEM_SCALES_A_SIZE_PER_STAGE);
}
// Wait unaligned cases
#pragma unroll
for (uint32_t s = kNumInnerStages; s < kNumStages; ++ s) {
empty_barriers[s]->wait((scheduler.current_iter * kNumIterations + k_iter + 1) & 1);
full_barriers[s]->arrive();
}
});
}
// To safely deconstruct distributed shared barriers, we need another round of empty waits
if constexpr (kNumTMAMulticast > 1) {
#pragma unroll
for (uint32_t s = 0; s < kNumStages; ++ s)
empty_barriers[s]->wait((scheduler.current_iter * kNumIterations + 1) & 1);
}
}
} else {
...
tma_copy を使ってTMAでグローバルメモリから、確保済みの共有メモリにデータをロードしている。このとき、 tensor_map_ とつく変数を渡しているのだが、これは確保済みのTMAディスクリプタだ。
// Prefetch TMA descriptors at very beginning
if (threadIdx.x == kNumMathThreads) {
cute::prefetch_tma_descriptor(reinterpret_cast<cute::TmaDescriptor const*>(&tensor_map_a));
cute::prefetch_tma_descriptor(reinterpret_cast<cute::TmaDescriptor const*>(&tensor_map_b));
cute::prefetch_tma_descriptor(reinterpret_cast<cute::TmaDescriptor const*>(&tensor_map_scales_a));
cute::prefetch_tma_descriptor(reinterpret_cast<cute::TmaDescriptor const*>(&tensor_map_d));
}
TMAディスクリプタはデータレイアウトなどテンソルの情報を記述するデータ構造なのだが、これは上の方で既にL1キャッシュにプリフェッチされている。
さて、上のコードを読むと、smem_a (lhs) 、smem_scales_a (lhsのスケーリング係数) 、smem_b (rhs) がTMAでロードされていることに気付く。smem_scales_b (rhsのスケーリング係数) は何故かここでTMAロードされず、計算担当のワープによってロードされている。ここは後で出てくる。
次に、上のコードの以下について。
full_barrier.arrive_and_expect_tx(SMEM_A_SIZE_PER_STAGE + SMEM_B_SIZE_PER_STAGE + SMEM_SCALES_A_SIZE_PER_STAGE);
full_barrierと、少し下にempty_barrierがある。これはキューみたいなもので、TMAで得たデータはfull_barrierに通知され、計算担当によって使い終わったデータはempty_barrierに通知されるようになっている。これらのバリアはシェアードメモリに確保されており、限られたシェアードメモリを使いまわしながらTMAと計算をうまくオーバーラップさせている、といった感じになる。
計算するのはelseの中で、ここだ。
まず以下のコード:
// Load B scales with math warp-groups
// NOTES: except the first warp, we want to overlap loading B scales with TMA stores between tasks
if (threadIdx.x >= 32) {
auto num_previous_lines = scheduler.get_global_idx<false>(ceil_div(SHAPE_N, BLOCK_K), 0, 0, m_block_idx);
auto local_scales_b = scales_b + (num_previous_lines + ((n_block_idx * BLOCK_N) / BLOCK_K)) * SHAPE_K_SCALES;
#pragma unroll
for (uint32_t i = threadIdx.x - 32; i < num_scales_b; i += kNumMathThreads - 32)
st_shared(smem_scales_b + i, __ldg(local_scales_b + i));
}
cutlass::arch::NamedBarrier(kNumMathThreads).sync();
上で少し触れたが、Bのスケーリング係数はTMAでロードせずに、計算担当のワープが自前で共有メモリから取り出している。 これはコメントにもあるが、TMAストアとBのスケーリング係数のロードをオーバーラップさせたいかららしいのだが、ちょっとよくわかってない。TMAユニットに個数制限があるのだろうか？
ワープグループという言葉がソースコードに出てきていると思うのでこれについて。ワープグループはCUDAのドキュメントによれば4つの連続したワープのことで、一定の条件を満たすことで連続したワープたちを強調して動かせるので命令発行やディスパッチのオーバーヘッドを削減できるものだと理解している。 以下のwgmmaは、実際にワープグループ単位で行列積を実際に計算する処理だ。
// Commit WGMMA instructions
#pragma unroll
for (int i = 0; i < WGMMA::kNumAccum; ++ i)
warpgroup_fence_operand(accum[i]);
warpgroup_arrive();
#pragma unroll
for (int k = 0; k < BLOCK_K / WGMMA::K; ++ k) {
auto desc_a = make_smem_desc(smem_a[s] + math_wg_idx * WGMMA::M * BLOCK_K + k * WGMMA::K, 1);
auto desc_b = make_smem_desc(smem_b[s] + k * WGMMA::K, 1);
WGMMA::wgmma(desc_a, desc_b, accum, k);
wmmaという、テンソルコアを直接使うAPIがあるのだが、wgmmaはそのワープグループ版になる。
その後、以下で共有メモリに書き戻し、その後TMAでグローバルメモリに書き戻す、という流れだ。
// Write back to shared memory using STSM
DG_STATIC_ASSERT(WGMMA::kNumAccum % 4 == 0, "Invalid STSM x2 vectorization");
#pragma unroll
for (auto i = 0; i < WGMMA::kNumAccum / 8; ++ i) {
SM90_U32x4_STSM_N<nv_bfloat162>::copy(
__float22bfloat162_rn({final_accum[i * 8 + 0], final_accum[i * 8 + 1]}),
__float22bfloat162_rn({final_accum[i * 8 + 2], final_accum[i * 8 + 3]}),
__float22bfloat162_rn({final_accum[i * 8 + 4], final_accum[i * 8 + 5]}),
__float22bfloat162_rn({final_accum[i * 8 + 6], final_accum[i * 8 + 7]}),
smem_d + (warp_idx * 16 + lane_idx % 16) * BLOCK_N + i * 16 + 8 * (lane_idx / 16)
);
}
if constexpr (WGMMA::kNumAccum % 8 != 0) {
SM90_U32x2_STSM_N<nv_bfloat162>::copy(
__float22bfloat162_rn({final_accum[WGMMA::kNumAccum / 8 * 8 + 0], final_accum[WGMMA::kNumAccum / 8 * 8 + 1]}),
__float22bfloat162_rn({final_accum[WGMMA::kNumAccum / 8 * 8 + 2], final_accum[WGMMA::kNumAccum / 8 * 8 + 3]}),
smem_d + (warp_idx * 16 + lane_idx % 16) * BLOCK_N + WGMMA::kNumAccum / 8 * 16
);
}
cute::tma_store_fence();
cutlass::arch::NamedBarrier(kNumMathThreads).sync();
// Use TMA store to write back to global memory
if (threadIdx.x == 0) {
cute::SM90_TMA_STORE_2D::copy(&tensor_map_d, smem_d, n_block_idx * BLOCK_N,
scheduler.get_global_idx(shape_m, BLOCK_M, m_block_idx));
cute::tma_store_arrive();
cute::tma_store_wait<0>();
}
__syncwarp();
このように、パイプラインを組んでTMAするワープグループと計算するワープグループに分かれることで、CPUのようなオーバーラッピングを実現している。
Schedulerというのがある。 このスケジューラが何をしているのかはわかりにくいのだが、たぶん以下のような感じだと思う。
スケジューラがやっているのはキャッシュブロッキングに近く、大きな行列を小さなブロックに分割して、うまくキャッシュに載せながらSMに割り当てることだ。
これを実際にやっているのがget_next_block関数である。
これは結局はget_swizzled_block_idxを呼び出してそれを返している。
__device__ __forceinline__ void get_swizzled_block_idx(const uint32_t num_m_blocks, int block_idx, uint32_t& m_block_idx, uint32_t& n_block_idx) {
DG_STATIC_ASSERT(kNumNBlocksPerGroup % kNumTMAMulticast == 0, "Invalid group size");
// Swizzle for better L2 usages
auto num_blocks_per_group = num_m_blocks * kNumNBlocksPerGroup;
auto group_idx = block_idx / num_blocks_per_group;
auto first_n_block_idx = group_idx * kNumNBlocksPerGroup;
auto num_n_blocks_in_group = min(kNumNBlocksPerGroup, kNumNBlocks - first_n_block_idx);
auto in_group_idx = block_idx % num_blocks_per_group;
m_block_idx = in_group_idx / num_n_blocks_in_group;
n_block_idx = first_n_block_idx + in_group_idx % num_n_blocks_in_group;
}
ここでは、m_block_idxとn_block_idxのポインタに適切な値をセットするのが目的なのだが、block_idxを受け取ってそれで行列上のどのブロックを計算させたいかを決定している。このとき、近いブロックが近いメモリになるようにしつつ、mとnからなる行列が、縦長でも横長でもなくいい感じの二次元になるようにしている。これをすると、近いブロックが近いデータを扱うようになるので、L2キャッシュにデータが残っている確率が高くなる。つまり空間的局所性を向上している。
また、この処理によってL2キャッシュのバンクコンフリクト発生を抑えている。近いブロックが近いデータを扱うようになると、同時にアクセスされるL2キャッシュのバンクも隣接しやすいと考えられるのでコンフリクトが起きにくくなる。
一番意味がわからなかったやつがこれなのだが、処理的にはここになる。
なんかバイナリをいじっているのがわかると思う。これはコードを読んでもわかるはずがないのでREADMEから引用する:
We observe a performance improvement in the CUTLASS FP8 kernel between NVCC 12.2 and 12.3. By comparing the compiled SASS, we discover that one bit in a series of
FADDinstructions is flipped in an interleaving pattern. After referencing some open-source CUDA assembler implementations, we identified that this bit controlsyield, which may enhance warp-level parallelism (just a guess, yielding the current warp and let other warps work). To leverage this, we develop a similar script to modify theFFMAinstructions in the compiled binary. Besides simply modifying theyieldbit, we also flip thereusebit (registers cannot be reused if the warp is yielded). This adjustment improves performance (10%+ in some cases) for fine-grained scaling FP8 GEMMs by creating more opportunities to overlap MMA instructions with promotionFFMAinstructions.
つまり、CUTLASSというNVIDIAの行列計算ライブラリで、nvccのバージョンを変えたらパフォーマンスが上がったので、SASS (アセンブラみたいなやつ) を見たところ、FADD命令のyieldビットが反転していたことがわかった。その後、reuseビットも反転させてみたら最大10%パフォーマンスが上がったとのこと。なんじゃそれという感じだけど早くなってるならヨシという感じなのだろう。
レジスタ数制御とかいくつか細かいのがあるようだが、あまり理解できなかったので省略する、、
DeepGEMMの場合、CUTLASSをかなりベースとしてTMA、パイプライニング、スケジューラは作っているというか、ほとんど真似している感じだった。たぶんそれで動かして、問題が合ったケースをひとつひとつ潰したり、バイナリをいじったりしている感じなんじゃないだろうか。
JITにしたのはすごく良いアイディアだと思っていて、ディープラーニングはパラメータだけ変えて同じ計算を何度も繰り返すことが多く、その割に実行環境がかなり人それぞれ違うので、JITでその都度最適化してベンチマーク取りながら実行するというのはワークロードに合ってるなと思う。
このコードを読むだけならなんとなくいけるが、書けと言われると無理だ。NVIDIAのエンジニアが凄いし、DeepSeekにもこれを改造して自分たちのワークロードに合わせられるようなエンジニアがいるんだなと思うと、技術力ヤバいなという気持ちになる。
糸冬了！！