AIハードウェアとソフトウェアスタック

FlashAttentionなどの高速なアテンションの実装の主な考え方として正しいものはどれか。

  1. A計算をタイルに分けてチップ上のメモリで進め、n×nの大きな中間の表をHBMに書き出さないことで、メモリの読み書きを減らす
  2. Bアテンションの計算を省略し、近似の答えを出す
  3. CアテンションをCPUで計算する
答えと解説を見る

正解：A. 計算をタイルに分けてチップ上のメモリで進め、n×nの大きな中間の表をHBMに書き出さないことで、メモリの読み書きを減らす

素直な実装では、n×nの表をHBMに書き、ソフトマックスのために読み直し、さらにValueとの積のために読み直すので、メモリの読み書きが性能を決めてしまいます。FlashAttentionは、QueryとKey・Valueをブロックに分けて共有メモリに読み込み、ソフトマックスを少しずつ更新しながら計算を一度に進めて、表をHBMに書き出しません。計算の量は同じでも、メモリの読み書きが大きく減って速くなります。

  • A. 正しい。結果は通常のアテンションと同じです。
  • B. 近似ではなく、同じ結果を正確に出します。
  • C. GPUのチップ上のメモリを活用します。
この分野を解く 分野ページへ

AIハードウェアとソフトウェアスタックのほかの問題

  1. AIの計算で使われる「演算の融合(オペレータフュージョン)」の主な目的はどれか。
  2. 行列の積の結果にバイアスを足し、ReLUを通す処理を、3つの別々のカーネルから1つに融合した。主に減るものはどれか。
  3. XLAやTVM、PyTorchのコンパイル機能のようなAIコンパイラが行う主な処理として正しいものはどれか。
  4. AIのフレームワークやコンパイラで使われる「計算グラフ」の説明として正しいものはどれか。
  5. ONNXの説明として正しいものはどれか。