AI・テクノロジー

Self-Attention、CUDA カーネル、FlashAttention

2026-11-15

マイクロコントローラのメモリブロックのイメージ図。記事で扱う GPU ではありません。

現代の Transformer は、Self-Attention の行列演算、Triton のような高レベル GPU コンパイラ、FlashAttention のようなメモリを意識した CUDA カーネルに依存し、膨大なトークンシーケンスをピークのハードウェア速度で処理しています。

1. Vision Transformer における Self-Attention の数学的仕組み

Vision Transformer（ViT）では、画像（例：224×224 ピクセル）が 16×16 パッチのグリッドに切り分けられ、N = 196 個の平坦化されたパッチが生成されます。各パッチは次元 D（例：768）のベクトルに線形投影され、位置埋め込みと結合されて、入力行列 X ∈ R^(N × D) を形成します。

Query、Key、Value の作成。モデルはパッチ行列 X に 3 つの学習済み重み行列（W_Q、W_K、W_V ∈ R^(D × d_k)）を掛け、各画像パッチに対して 3 つの異なる表現を作成します：

Q=XW_(Q), K=XW_(K), V=XW_(V)

Query（Q）：特定の画像パッチが何を探しているかを表します（例：「私は毛皮のテクスチャで、近くの目や耳を探しています」）。
Key（K）：パッチがどんな情報を含んでいるかを示します（例：「私は鋭い曲線のエッジを含んでいます」）。
Value（V）：マッチが起きた場合に前方へ受け渡される実際の視覚特徴の内容を保持します。

Attention スコアの計算。すべてのパッチが他のすべてのパッチとどれほど強く関連しているかを測るため、モデルはすべての Query と Key の間のドット積を計算します：

Scores=(QK^(T))/(sqrt(d_(k)))

sqrt(d_(k)) で割ることで、大きなドット積がこの後の softmax 関数を勾配の極めて小さい領域へ押し込むのを防ぎます。これにより N × N の Attention 行列が生成され、要素 (i, j) がパッチ i がパッチ j にどの程度注意を向けるべきかをスコア化します。

Softmax の重み付けと集約。行ごとの Softmax を適用すると、生のスコアが合計 1 の確率に変換されます。これらの確率を Value 行列 V に掛けると、各パッチのブレンドされた表現が得られます：

Attention(Q,K,V)=softmax((QK^(T))/(sqrt(d_(k))))V

Multi-Head Self-Attention では、このまったく同じプロセスが複数のヘッド（例：12 ヘッド）で並列に実行され、異なるヘッドが異なる空間距離、色、物体のパーツに同時に注目できます。

2. 実践コード比較：生の CUDA C++ と OpenAI Triton

GPU プログラミングがどのように進化したかを見るために、N 要素に対する単純な要素ごとのベクトル加算（C = A + B）を考えてみましょう。

生の CUDA C++（スレッドレベルの管理）

CUDA C++ では、単一のスカラースレッドの視点でプログラミングします。スレッド ID の計算、配列境界のチェック、定型的なホスト側の起動設定の記述をすべて手動で行う必要があります。

#include <cuda_runtime.h>

// Device kernel: executed by thousands of individual threads
__global__ void vector_add_cuda(const float* A, const float* B, float* C, int N) {
    // Manually compute global thread index from block and thread IDs
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    
    // Guard against out-of-bounds memory access
    if (idx < N) {
        C[idx] = A[idx] + B[idx];
    }
}

// Host launch setup
void launch_cuda(const float* d_A, const float* d_B, float* d_C, int N) {
    int threads_per_block = 256;
    int blocks_per_grid = (N + threads_per_block - 1) / threads_per_block;
    vector_add_cuda<<<blocks_per_grid, threads_per_block>>>(d_A, d_B, d_C, N);
}

OpenAI Triton（ブロックレベルの Python プログラミング）

Triton では、Python で直接、データのブロック（タイル）の視点でプログラミングします。Triton のコンパイラがスレッド同期、メモリコアレッシング、ベクトル化を内部で自動的に処理します。

import triton
import triton.language as tl

@triton.jit
def vector_add_triton(A_ptr, B_ptr, C_ptr, N, BLOCK_SIZE: tl.constexpr):
    # Identify which block of data this program instance handles
    pid = tl.program_id(axis=0)
    
    # Create a vector of offsets for the entire block at once
    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = offsets < N
    
    # Load chunks of memory into fast SRAM, add them, and store back
    a = tl.load(A_ptr + offsets, mask=mask)
    b = tl.load(B_ptr + offsets, mask=mask)
    tl.store(C_ptr + offsets, a + b, mask=mask)

# Launch directly from PyTorch tensors
def launch_triton(A, B, C):
    N = A.numel()
    grid = lambda meta: (triton.cdiv(N, meta['BLOCK_SIZE']),)
    vector_add_triton[grid](A, B, C, N, BLOCK_SIZE=1024)

3. FlashAttention：Transformer のメモリボトルネックを解決する

標準的な Self-Attention は、高解像度画像や長いテキストシーケンスを処理する際に深刻な壁に直面します。Attention 行列 S = QK^T がシーケンス長 N に対して二次関数的（N × N）に増大するためです。

メモリ帯域幅のボトルネック（HBM と SRAM）。現代の NVIDIA GPU には 2 つの主要なメモリ階層があります：

HBM（高帯域幅メモリ / VRAM）：大容量（例：80 GB）ですが、読み書きは比較的低速です。
SRAM（オンチップ共有メモリ）：非常に小さい（全ストリーミングマルチプロセッサ合計で約 20 MB）ですが、HBM より約 10 倍から 20 倍高速です。

標準的な Attention では、GPU は S = QK^T を計算し、巨大な N × N 行列を低速な HBM に書き出し、それを読み戻して Softmax を適用し、確率を HBM に書き戻し、再び読み出して V に掛けます。Tensor Core は時間の大半をメモリ転送の待機に費やしています。

タイル化（N × N 行列を決して実体化しない）。Tri Dao ら（2022 年）によって導入された FlashAttention は、カーネル融合を使って Attention 計算全体を高速なオンチップ SRAM 内に収めます。Q、K、V を SRAM に収まる小さなブロックに切り分け、タイルごとに Attention を計算し、最終出力行列だけを HBM に書き戻します。巨大な N × N の中間行列が VRAM に保存されることはありません。

オンライン Softmax の工夫。通常、Softmax は単一の要素を正規化する前に、行全体の最大値と指数和を知る必要があります。FlashAttention はオンライン Softmax アルゴリズムを使い、ブロックごとに実行中の最大値と正規化係数を追跡し、新しいタイルが処理されるたびに部分的な出力を数学的に再スケーリングします。

逆伝播時の再計算。FlashAttention は順伝播で N × N の Attention 行列を保存しないため、訓練時に巧妙なトレードオフを使います。小さな正規化統計量だけを HBM に保存し、逆伝播時に高速な SRAM 内で Attention タイルをその場で再計算します。SRAM での追加計算は HBM から大きな行列を読むよりもはるかに高速であることが実証されており、実時間で 2 倍から 4 倍の高速化を実現し、メモリ使用量を二次の O(N²) から線形の O(N) へ削減します。

その他のニュース: AI・テクノロジー