cuBLASは、バッチ化されたFP32ワークロードごとに非効率なカーネルをディスパッチします。サイズは256×256から8192×8192×8まで。RTX GPU上で利用可能な演算のうち、およそ40%しか使いません。RTX 5090でテストしましたが、おそらくすべてのRTX非Pro GPUの影響を受けています。
CUDA 13.2.51、cuBLAS 13.3.0、ドライバ595.58.03の最新バージョンでテストしました。以前のバージョンはさらに悪いです。
私は単純ですが効率的なカーネルを書き、さまざまなワークロードに対してcuBLASと比較しました。
5090でのバッチ性能 vs cuBLAS（>100%なら私のカーネルの方が速いことを意味します）：
| Size | B=4 | B=8 | B=16 |
|---|---|---|---|
| 256 | 91% | 80% | 90% |
| 512 | 120% | 153% | 135% |
| 1024 | 137% | 142% | 142% |
| 2048 | 158% | 155% | 157% |
| 4096 | 157% | 162% | 170% |
| 8192 | 158% | 152% | 148% |
cuBLASは他のGPUでは適切なカーネルを使用します。RTX GPUは明らかにNVIDIAからの優遇が少ないようです：
- Pro 6000：3種類のタイルサイズを順に拡張し、73%のFMA（Fused Multiply-Addパイプ）に到達
- H200：最良の実装で、CUTLASSとxmmaファミリを混在させ、82%のFMAに到達
3つのGPUすべてにわたる詳細な分析と完全なNCUプロファイリングデータ、さらにSASSスケジューリングを掘り下げて、私のカーネルと適切なcuBLAS SGEMMの間に残る5%のシングルモード差を説明する内容、そして再現スクリプトは、以下のリンク先の記事に用意されています。
バグに加えて、この記事ではシンプルなTMA（tensor memory accelerator）のダブルバッファカーネルも扱っており、5090上のバッチモードでcuBLASを46〜65%上回ります。また、適切に選択されたカーネルの性能の80〜120%を達成し、シンプルながら非常に高性能なカーネルを書くための良いテクニックになっています。
VS 適切なPro6000カーネル：
| Size | B=4 | B=8 | B=16 |
|---|---|---|---|
| 256 | 87% | 95% | 77% |
| 512 | 102% | 124% | 101% |
| 1024 | 101% | 104% | 96% |
| 2048 | 90% | 102% | 93% |
| 4096 | 93% | 93% | 93% |
| 8192 | 94% | 95% | 95% |
VS 適切なH200カーネル：
| Size | B=4 | B=8 | B=16 |
|---|---|---|---|
| 256 | 85% | 104% | 77% |
| 512 | 105% | 97% | 88% |
| 1024 | 87% | 89% | 89% |
| 2048 | 89% | 90% | 92% |
| 4096 | 91% | 89% | 90% |
| 8192 | 88% | 87% | 87% |
ダブルバッファパイプラインの可視化：
Tile 0: [load buf0] [wait] [compute buf0 + load buf1] Tile 1: [wait buf1] [compute buf1 + load buf0] Tile 2: [wait buf0] [compute buf0 + load buf1] ... 簡略化したカーネルソース：
__global__ __launch_bounds__(256) void fused_matmul( const __grid_constant__ CUtensorMap A_tma, const __grid_constant__ CUtensorMap B_tma, float* C) { extern __shared__ __align__(128) char dsmem[]; float* smem = (float*)dsmem; // ダブルバッファ同期のためのmbarrierを2つ uint64_t* mbar = (uint64_t*)(dsmem + 2 * STAGE * 4); // TMAターゲットのための共有メモリアドレス const int as0 = __cvta_generic_to_shared(&smem[0]); const int bs0 = __cvta_generic_to_shared(&smem[A_SIZE]); const int as1 = __cvta_generic_to_shared(&smem[STAGE]); const int bs1 = __cvta_generic_to_shared(&smem[STAGE + A_SIZE]); // スレッドID int tid = threadIdx.y * 32 + threadIdx.x; int tr = threadIdx.y * TM, tc = threadIdx.x * 4; int bm = blockIdx.y * BM, bn = blockIdx.x * BN; // mbarrierの初期化（スレッド0のみ） if (tid == 0) { mbarrier_init(mbar[0]); mbarrier_init(mbar[1]); } __syncthreads(); float c[TM][4] = {}; // アキュムレータ // 最初のタイルを事前ロード if (tid == 0) { mbarrier_expect_tx(mbar[0], BYTES); tma_load_2d(as0, &A_tma, /*k=*/0, bm, mbar[0]); tma_load_2d(bs0, &B_tma, bn, /*k=*/0, mbar[0]); } for (int t = 0; t < K/BK; t++) { int s = t % 2; // 現在のバッファ // 現在のタイルのTMA完了を待機 mbarrier_wait(mbar[s], phase[s]); // 次のタイルのロードを開始（計算と重ねる） if (tid == 0 && t + 1 < nt) { tma_load_2d(next_buf_a, &A_tma, next_k, bm, next_mbar); tma_load_2d(next_buf_b, &B_tma, bn, next_k, next_mbar); } // 計算：全256スレッドが共有メモリからFMA float* As = &smem[s * STAGE]; float* Bs = &smem[s * STAGE + A_SIZE]; #pragma unroll for (int kk = 0; kk < BK; kk++) { float b0 = Bs[kk*BN+tc], b1 = Bs[kk*BN+tc+1], ...; for (int i = 0; i < TM; i++) { float a = As[(tr+i)*BK+kk]; c[i][0] += a * b0; c[i][1] += a * b1; // ... 行あたり4つのFMA } } __syncthreads(); } // 結果をグローバルメモリへ書き込み for (int i = 0; i < TM; i++) store_row(C, bm+tr+i, bn+tc, c[i]); 全文の記事はこちら
再現スクリプト付きのRepoと、ベンチマークデータ
[link] [comments]
