Skip to content
Published on

cuDNN完全解剖:なぜ深層学習演算はGPU上で光速なのか

Authors

はじめに:cuDNNがなかったら何が起きるか

PyTorchでtorch.nn.Conv2dを呼び出すとき、内部で何が起きているか知っているだろうか。単に「GPUで畳み込みを実行する」だけではない。cuDNNが自動的にアルゴリズムを選択し、メモリレイアウトを調整し、ハードウェアに最適化されたカーネルを実行している。

生のCUDAカーネルで直接畳み込みを実装すると、cuDNNと比べて10〜100倍遅くなる。その理由を正確に理解することがこの記事の目標だ。


1. なぜcuDNNが必要なのか:Raw CUDA vs 最適化ライブラリ

CUDAは汎用の並列プログラミングモデルだ。あらゆる演算を並列化できるが、深層学習で支配的な特定の演算(畳み込み、バッチ正規化、アテンション)のための専門的な最適化は含まれていない。

cuDNN(CUDA Deep Neural Network library)はこれらの演算のための手作業でチューニングされたカーネルライブラリだ。同じ演算でも:

Naive CUDAの畳み込み:        約5 TFLOPS達成(H100理論値の0.5%cuDNN最適化畳み込み:          約900 TFLOPS達成(H100理論値の90%差:約180

なぜこれほどの差があるのか:

  1. アルゴリズム選択:入力サイズに応じてDirect、im2col+GEMM、Winograd、FFTから最適を自動選択
  2. メモリレイアウト最適化:Tensor Coreを有効にするNHWCレイアウトを使用
  3. カーネル融合:複数の演算を1つのカーネルにまとめてメモリ往復回数を削減
  4. Warpレベル最適化:メモリコアレッシング、レジスタ再利用などのHW特性を完全活用

PyTorchとTensorFlowはcuDNNを基盤ライブラリとして使用している:

import torch

# このコードは内部的にcuDNNを呼び出す
conv = torch.nn.Conv2d(64, 128, kernel_size=3, padding=1).cuda()
x = torch.randn(32, 64, 56, 56, device='cuda')
y = conv(x)  # → cudnnConvolutionForward() を呼び出す

# cuDNNの使用状況を確認
print(torch.backends.cudnn.enabled)    # True
print(torch.backends.cudnn.version())  # 例: 8906

2. 4種類の畳み込みアルゴリズム:cuDNNの核心

畳み込みは深層学習で最も演算集約的な処理だ。cuDNNは状況に応じて4つのアルゴリズムから選択する。

アルゴリズムA:Direct Convolution(Naive)

最も直感的な実装。出力の各要素を独立して計算する:

各出力ピクセル(i, j)について:
  for c_out in range(C_out):
    for c_in in range(C_in):
      for kh in range(K):
        for kw in range(K):
          output[c_out, i, j] += input[c_in, i+kh, j+kw] * filter[c_out, c_in, kh, kw]

総演算数:N × C_out × C_in × K^2 × H_out × W_out

長所:シンプル。短所:メモリアクセスパターンが非効率。大きな入力では非常に遅い。

アルゴリズムB:im2col + GEMM(cuDNNデフォルト)

畳み込みを行列乗算に変換する核心トリック。cuBLASの高度に最適化されたGEMMを活用できる。

im2col変換:

入力特徴マップ(3x3画像、3x3カーネル):

元の入力:           im2col変換結果:
┌─────────┐          ┌────────────────────────────┐
1  2  3 │  im2col  │ 1  2  4  5  (パッチ0, 位置0,0)4  5  6 │ ────────→│ 2  3  5  6  (パッチ1, 位置0,1)7  8  9 │          │ 4  5  7  8  (パッチ2, 位置1,0)└─────────┘          │ 5  6  8  9  (パッチ3, 位置1,1)                     └────────────────────────────┘
各行 = 1つの受容野(receptive field)内のすべての画素値

その後:output = filter_matrix × im2col_matrix
        (標準GEMMTensor Coreをフル活用!)

im2colの欠点:入力データを再配置するため追加メモリが必要(元の入力のK^2倍)。

# im2colの動作をPythonで確認する例
import numpy as np

def im2col(input, kernel_h, kernel_w, stride=1, pad=0):
    N, C, H, W = input.shape
    out_h = (H + 2*pad - kernel_h) // stride + 1
    out_w = (W + 2*pad - kernel_w) // stride + 1

    # パディングを適用
    img = np.pad(input, [(0,0),(0,0),(pad,pad),(pad,pad)], mode='constant')

    # im2col行列を構築:(C*kernel_h*kernel_w) × (N*out_h*out_w)
    col = np.zeros((N, C, kernel_h, kernel_w, out_h, out_w))
    for j in range(kernel_h):
        jj = j + stride * np.arange(out_h)
        for i in range(kernel_w):
            ii = i + stride * np.arange(out_w)
            col[:, :, j, i, :, :] = img[:, :, jj[:, None], ii[None, :]]

    col = col.transpose(0, 4, 5, 1, 2, 3).reshape(N*out_h*out_w, -1)
    return col
# cuDNNの実際のim2colはこれよりはるかに最適化されている

アルゴリズムC:Winogradアルゴリズム(小型カーネルの最強)

3x3畳み込みでcuDNNがデフォルトで選択するアルゴリズム。Winogradの最小フィルタリングアルゴリズム(1980年)を深層学習に適用したものだ。

核心的なアイデア: 線形代数変換で乗算回数を大幅に削減する。

通常の3x3畳み込み(2x2出力):
- 入力パッチ:4x4 = 16要素
- カーネル:3x3 = 9要素
- 出力:2x2 = 4要素
- 必要な乗算回数(naive):4 × 9 = 36
Winograd F(2x2, 3x3)変換後:
- 変換済み入力:4x4 = 16要素(線形変換、乗算なし)
- 要素ごとの乗算:16回(4x4)
- 逆変換:4x4 → 2x2(線形変換、乗算なし)
- 必要な乗算回数:16- 削減率:3616 = 2.25倍の削減

数式で表すと:

Y = A^T [(G × g × G^T)  (B^T × d × B)] A

ここで:
d = 入力タイル(4x4)
g = 3x3カーネル
B, G, A = 固定された変換行列(事前計算済み定数)
= 要素ごとの乗算(アダマール積)

重要な点:G × g × G^T はカーネルに対して1回だけ計算される
           (推論時に事前計算可能)
           B^T × d × B は各入力タイルに対して計算
           要素ごとの乗算により36回ではなく16回の乗算で同じ結果

cuDNNは3x3、stride=1の畳み込みで自動的にWinogradを選択する。ResNet、VGGなどほとんどのCNNが3x3畳み込みを主に使うため、実際のところ非常に重要だ。

アルゴリズムD:FFTベースの畳み込み(大型カーネル)

7x7や11x11のような大型カーネルには周波数領域での畳み込みが効率的だ。

空間領域の畳み込み:    O(N × K^2)N = 出力サイズ、K = カーネルサイズ)
周波数領域の畳み込み:  O(N × log N)FFT後に要素ごとの乗算)

Kが大きいほどFFTが有利:
K=3:  9 ops vs log(N)7  → 大した差なし
K=11: 121 ops vs log(N)7FFT17倍少ない乗算

3. cuDNN Auto-Tuner:benchmarkモードの真実

# この1行が実際に何をするのか?
torch.backends.cudnn.benchmark = True

benchmark=False(デフォルト):

  • cuDNNが入力サイズに基づいてヒューリスティックでアルゴリズムを選択
  • 初回から高速だが最適でない可能性がある

benchmark=True:

  • 最初のforward passでcudnnFindConvolutionForwardAlgorithm()を呼び出す
  • 現在の入力サイズで利用可能なすべてのアルゴリズムを実際に実行しベンチマーク
  • 最速のものを選択してキャッシュに保存
  • 以降、同じ入力サイズではキャッシュされたアルゴリズムを使用
# benchmarkの効果を直接確認
import torch
import time

torch.backends.cudnn.benchmark = False
model = resnet50().cuda()
x = torch.randn(32, 3, 224, 224).cuda()

t0 = time.time()
for _ in range(10): y = model(x)
torch.cuda.synchronize()
print(f"benchmark=False: {(time.time()-t0)*1000:.0f}ms")

torch.backends.cudnn.benchmark = True
# 最初の実行にはベンチマーキングが含まれる(遅い!)
t0 = time.time()
y = model(x)  # この1回の呼び出しに数秒かかる
torch.cuda.synchronize()
print(f"最初の実行(ベンチマーキング含む): {(time.time()-t0)*1000:.0f}ms")

# 以降の実行(キャッシュされた最適アルゴリズムを使用)
t0 = time.time()
for _ in range(10): y = model(x)
torch.cuda.synchronize()
print(f"benchmark=True(ウォームアップ後): {(time.time()-t0)*1000:.0f}ms")
# 通常20〜40%高速

注意: バッチごとに入力サイズが変わる場合、benchmark=Trueはむしろ遅くなる可能性がある。毎回新しいベンチマーキングが必要になるためだ。学習・推論の入力サイズが固定の場合にのみ使用すること。


4. バッチ正規化とカーネル融合:メモリトラフィック削減の魔法

バッチ正規化(Batch Normalization)は見た目はシンプルだ:

y = gamma * (x - mean(x)) / sqrt(var(x) + eps) + beta

しかしこれを別々のカーネルとして実装すると致命的に非効率になる:

Naive実装のメモリ往復:
1. HBMからxをロード
2. 平均値を計算 → HBMに保存
3. x + 平均値をロード(HBM読み取り2回)
4. 分散を計算 → HBMに保存
5. x + 平均値 + 分散をロード(HBM読み取り3回)
6. 正規化を計算 → HBMに保存
7. ReLU:HBMから読み取り → 計算 → HBMに書き込み

合計:約10回のHBMアクセス

cuDNNはBN + ReLUを単一カーネルに融合する:

融合BN+ReLUカーネル:
1. xをレジスタ/共有メモリにロード(HBM読み取り1回)
2. レジスタで平均値、分散を計算(HBMアクセスなし)
3. レジスタで正規化(HBMアクセスなし)
4. レジスタでReLU(HBMアクセスなし)
5. 結果をHBMに保存(HBM書き込み1回)

合計:HBMアクセス2回(5倍削減!)
# PyTorchで融合BN+ReLUを使用
import torch.nn as nn

# Naive:Conv → BN → ReLU(3つの別々のカーネル起動)
class NaiveBlock(nn.Module):
    def __init__(self, c):
        super().__init__()
        self.conv = nn.Conv2d(c, c, 3, padding=1)
        self.bn = nn.BatchNorm2d(c)
        self.relu = nn.ReLU()

    def forward(self, x):
        x = self.conv(x)
        x = self.bn(x)    # 別のcuDNNカーネル
        x = self.relu(x)  # 別のelementwiseカーネル
        return x

# torch.compile()を使うと、コンパイラが自動的に
# cuDNN Graph APIとnvFuserを通じてこれらを融合
model = torch.compile(model)  # conv+bn+reluを自動的に1つのカーネルに融合

最新のcuDNN(v8以降)はグラフAPIを通じてより複雑な融合もサポートしており、torch.compile()が自動的にこれを活用する。


5. アテンションとFlashAttention:cuDNNを超えて

Transformerのアテンション演算を標準的に実装すると:

# Standard Attention — メモリ非効率
def standard_attention(Q, K, V, scale):
    # S = Q × K^T : N×Nのアテンションスコア行列を生成
    S = torch.matmul(Q, K.transpose(-2, -1)) * scale  # O(N^2)のメモリ!
    # ソフトマックス
    P = torch.softmax(S, dim=-1)
    # Vとの積
    O = torch.matmul(P, V)
    return O

# 問題:SはO(N^2)のメモリを消費する
# N=8192(文書長)、FP16の場合:
# Sのサイズ = 8192 * 8192 * 2バイト = ヘッド1つあたり128MB!

FlashAttentionはcuDNNを完全に迂回する新しいCUDAカーネルでこの問題を解決する。

FlashAttentionの核心的なアイデア:N×N行列をHBMに書き出さない

Standard Attentionのメモリフロー:
Q, K, VHBMから読み取り
S = QK^THBMに書き込み(N^2サイズ!)
softmax(S)HBMに書き込み
P × VHBMに書き込み
HBMアクセス:O(N^2)

FlashAttentionのメモリフロー:
Q, K, Vをタイル単位でロード → 共有メモリへ
共有メモリでアテンションの部分和を計算
  (完全なN^2行列を生成しない)
最終出力のみHBMに書き込み
HBMアクセス:O(N)N^2からNへ!

タイリングトリックの詳細:

FlashAttentionのタイリング:

QをブロックQ_1, Q_2, ... Q_Tcに分割
K, VをブロックK_1, V_1, ... K_Tr, V_Trに分割

for i in range(Tc):
  Q_iをSRAMにロード
  O_i = 0, l_i = 0, m_i = -inf  (ソフトマックス統計を初期化)

  for j in range(Tr):
    K_j, V_jをSRAMにロード

    # SRAMのみで計算(HBMアクセスなし)
    S_ij = Q_i × K_j^T  (タイルサイズのアテンションスコア)
    m_ij = max(m_i, rowmax(S_ij))  (数値安定性)
    P_ij = exp(S_ij - m_ij)

    # オンラインソフトマックス更新(数値的に安定!)
    O_i = diag(exp(m_i - m_ij)) × O_i + P_ij × V_j
    l_i = exp(m_i - m_ij) × l_i + rowsum(P_ij)
    m_i = m_ij

  # 最終正規化後、HBMに保存(書き込み1回)
  O_i = diag(l_i)^(-1) × O_iHBMに保存

結果:

  • メモリ使用量:O(N^2) → O(N)(N=8192で数百MBの節約)
  • 速度:2〜4倍高速(HBMトラフィック削減のおかげ)

6. メモリレイアウト:NCHWとNHWC

深層学習テンソルのメモリレイアウトは性能に決定的な影響を与える。

NCHWレイアウト(バッチ × チャネル × 高さ × 幅):
バッチ=1、チャネル=3RGB)、4x4画像:

メモリ配置:
[R00 R01 R02 R03 | R10 R11 ... | R30 R31 R32 R33 |
 G00 G01 G02 G03 | G10 G11 ... | G30 G31 G32 G33 |
 B00 B01 B02 B03 | B10 B11 ... | B30 B31 B32 B33]

→ 同じチャネルのピクセルが連続して配置

NHWCレイアウト(バッチ × 高さ × 幅 × チャネル):
メモリ配置:
[R00 G00 B00 | R01 G01 B01 | R02 G02 B02 | R03 G03 B03 |
 R10 G10 B10 | R11 G11 B11 | ...
 R30 G30 B30 | R31 G31 B31 | R32 G32 B32 | R33 G33 B33]

→ 同じ空間位置のすべてのチャネルが連続して配置

Tensor CoreがNHWCを好む理由:

Tensor Coreは16x16の行列タイルを処理する。NCHWで畳み込み用の16x16タイルを構成すると、チャネル方向で不連続なメモリアクセスが発生する。NHWCではチャネルが連続しているため、タイルのロードが連続メモリ読み取りになる。

# PyTorchでNHWC(channels_last)を使用
x_nchw = torch.randn(32, 64, 56, 56, device='cuda')

# NHWCに変換(channels_lastフォーマット)
x_nhwc = x_nchw.to(memory_format=torch.channels_last)

# モデルをchannels_lastフォーマットに変換
model = model.to(memory_format=torch.channels_last)

# このforward passではcuDNNがNHWCカーネルを自動選択
output = model(x_nhwc)

# ベンチマーク比較(H100での実測値):
# NCHW: 約12ms/バッチ
# NHWC: 約9ms/バッチ  (約25%高速)

7. TensorRT:推論のためのcuDNNの次の段階

TensorRTはcuDNNの上に構築された推論最適化エンジンだ。学習済みモデルをデプロイ環境で最大性能で実行する。

TensorRTの最適化パイプライン:

元のモデル(ONNX/PyTorch)
  グラフ解析
  レイヤー融合
  ┌─────────────────────────────────────────────┐
ConvBNReLUConvBNReLU  │           ↓ 融合後                          │
  │     単一の最適化カーネル                     │
  └─────────────────────────────────────────────┘
  精度選択(FP32FP16INT8  カーネル自動選択(入力サイズごとにベンチマーク)
  最適化された実行エンジン

精度段階別の性能比較(ResNet-50、バッチ=32、H100):

精度レイテンシスループット精度低下
FP324.2ms7,600 img/s基準値
FP161.8ms17,800 img/s無視できるレベル
INT80.9ms35,500 img/s約0.1% Top-1

INT8量子化のキャリブレーション:

import tensorrt as trt

# INT8キャリブレーターの設定
class MyCalibrator(trt.IInt8EntropyCalibrator2):
    def __init__(self, data_loader, cache_file):
        super().__init__()
        self.data_loader = iter(data_loader)
        self.cache_file = cache_file
        self.batch_allocation = None

    def get_batch_size(self):
        return 32

    def get_batch(self, names):
        try:
            batch = next(self.data_loader)[0].numpy()
            # 代表データセットを使ってアクティベーション範囲をキャリブレーション
            # TensorRTが各レイヤーのFP32範囲をINT8にマッピング
            if self.batch_allocation is None:
                self.batch_allocation = cuda.mem_alloc(batch.nbytes)
            cuda.memcpy_htod(self.batch_allocation, batch)
            return [int(self.batch_allocation)]
        except StopIteration:
            return None

# ビルダーの設定
builder = trt.Builder(logger)
config = builder.create_builder_config()
config.set_flag(trt.BuilderFlag.INT8)  # INT8を有効化
config.int8_calibrator = MyCalibrator(calibration_loader, 'cache.bin')

8. 実際のLLMでのcuDNN呼び出しを追跡する

GPT-2やLLaMAを実行するとき、内部でどのカーネルが呼び出されるかを追跡してみよう:

# PyTorch ProfilerでLLMの演算を追跡
import torch
from torch.profiler import profile, ProfilerActivity

model = GPT2Model.from_pretrained('gpt2').cuda()
input_ids = torch.randint(0, 50257, (1, 512)).cuda()

with profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
    with_stack=True
) as prof:
    output = model(input_ids)

# 上位CUDAカーネルを出力
print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))

実際に実行されるカーネル:

GPT-2 Forward Pass内部CUDAカーネル(512トークン、バッチ=1):

演算              カーネル                        時間   割合
──────────────────────────────────────────────────────────
Linear(QKV)    cublasSgemm(Tensor Core)      2.1ms  35%
Attention        flash_attn_fwd_kernel           1.4ms  23%
Linear(出力)   cublasSgemm(Tensor Core)      0.9ms  15%
LayerNorm        layer_norm_kernel               0.3ms   5%
GELU             vectorized_elementwise_kernel   0.2ms   3%
Residual Add     vectorized_elementwise_kernel   0.1ms   2%
Embedding        Embedding_cuda                  0.1ms   2%
その他                                           0.9ms  15%

合計時間:約6ms(H100での計測)

重要な観察点:

  • Linearレイヤー = cuBLAS GEMM:LLM演算の約50%が行列乗算
  • Attention = FlashAttention:cuDNNではなくカスタムCUDAカーネル
  • LayerNorm = カスタム融合カーネル:平均 + 分散 + 正規化が1つのカーネル
  • 活性化関数 = elementwiseカーネル:非常に高速(メモリ帯域幅制限)

9. まとめ:cuDNNが生み出す性能差の源泉

cuDNNがRaw CUDAより10〜100倍速い理由をまとめると:

  1. アルゴリズム最適化:Winograd、im2col+GEMMなど数学的により効率的なアルゴリズムを状況に応じて選択
  2. カーネル融合:複数の演算を1つのカーネルにまとめてHBM往復回数を最小化
  3. メモリレイアウト最適化:NHWC + Tensor Coreの有効化でハードウェア効率を最大化
  4. Auto-Tuner:実際のハードウェアでベンチマーキングして入力サイズごとに最適な実装を選択
  5. FlashAttentionのような革新:O(N^2) → O(N)のメモリアクセスでアテンション演算を革新

PyTorchの1行の裏には、このような何十年もの最適化研究が隠されている。この内部動作を理解することが、LLMサービング最適化、カスタムCUDAカーネルの作成、ハードウェア選択の出発点となる。