arXiv論文メモ
新着一覧
cs.AR / cs.LG · 査読状況未確認

AMDのNPUで注意計算をどこまで融合すべきか

Programming AMD XDNA NPUs with Open-source Compiler Tools: A FlashAttention Case Study

Erwei Wang, Ephrem Wu, Victor J. B. Jung, Jiajie Li, Andre Rosti, Joseph Melber, Samuel Bayliss

この論文をやさしく読む

ひとことで言うと

AI向けNPUでFlashAttentionを速く動かすため、中間結果をどのメモリーに置き、どこまで演算を一体化するべきかを比較した研究です。新しい世代で有効な最適化が、前の世代ではほとんど効かない理由も調べます。

何に役立つ?

NPU向けの実装で、データ移動の削減に開発時間をどこまで使うべきかを判断する参考になります。メモリー階層ごとのルーフライン解析で、融合の追加効果を予測します。

この研究の面白いところ

QKᵀを共有メモリーに戻さず計算タイルに保持することで、XDNA 2ではIRON設計の2倍の性能を得ています。一方、XDNA 1ではすでに演算が律速となるため、同じ融合の価値が小さくなります。

どこまで分かった?

3.62 TFLOP/sは対象処理のエンドツーエンド実行の値です。5.3〜7.2倍のエネルギー効率は2Kトークン以上で同一チップの統合GPUと比較した結果で、全モデルの推論速度やあらゆるGPUとの比較を意味しません。

v1のアブストラクトに基づくAI解説。日本語訳とは別に、用途の解釈を含みます。

アブストラクトの日本語訳

AMD XDNAのような空間型NPUは、小さなローカルメモリーの隣に計算タイルを配置し、それらの間のデータ移動をソフトウェアに任せる。多段階の処理をこのような装置へ割り当てる際の主要な問題は、中間テンソルをどこに置くかである。本研究では、オープンソースのIRONとMLIR-AIRの処理系を使い、FlashAttentionのためにその選択を行う中で得た知見を報告する。 XDNA 1とXDNA 2上で四つの参照設計を比較する。一つは演算を個別に実行し、二つは演算間でオンチップのストリーミングを行い、残る一つは注意計算の三段階すべてを単一カーネルへ融合する。融合カーネルはQKᵀのスコアを計算タイルのローカルメモリーに保持し、カスケード接続を介して部分結果を集約するため、スコアを共有MemTileメモリーへ戻さない。 XDNA 2では、エンドツーエンドの全実行で3.62 TFLOP/sに達し、IRONの設計の2倍となった。2Kトークン以上では、同じチップにある統合GPUに対し、5.3〜7.2倍のエネルギー効率を示す。BERTからDeepSeekまで12種類のLLM構成を対象とし、最大128Kトークンに対応する。各メモリーレベルでのルーフライン解析が結果を説明し、最適化をどこで止めるべきかも示す。 XDNA 1は演算律速とメモリー律速の境目が低いため、オンチップのストリーミングだけで演算律速の領域に達する。XDNA 2でスループットを倍増させる同じ融合は、XDNA 1ではほぼ効果がない。実装の演算強度を各メモリーレベルの境目と比較すれば、コードを書く前にどちらの場合か予測できる。実装がその境目を越えるまで融合し、そこで止めるべきである。参照設計を、保守されるオープンソースとして公開する。

v1の要旨から自動生成。本文の精読・人による確認は未実施。

初稿
2026-09-18(UTC)
最新改訂
2026-09-18 · v1
査読・掲載
査読状況未確認
arXivで読むPDF

更新履歴

取得できた版を表示。版の更新は査読済みを意味しません。過去版の本文差分は未解析です。

原文の要旨

Spatial NPUs such as AMD XDNA place compute tiles beside small local memories and leave data movement between them to software. Mapping a multi-stage workload onto such a device is largely a question of where the intermediate tensors live. We report what we learned making those choices for FlashAttention with the open-source IRON and MLIR-AIR flows. We compare four reference designs on XDNA 1 and XDNA 2: one runs each operator separately, two stream between operators on chip, and one fuses all three attention stages into a single kernel. The fused kernel holds the $\boldsymbol{QK}^{\mathsf T}$ scores in compute-tile local memory and reduces partial results over the cascade interconnect, so the scores never return to shared MemTile memory. On XDNA 2, it reaches 3.62 TFLOP/s over complete end-to-end execution, twice the IRON design, with 5.3 to 7.2 times the energy efficiency of the integrated GPU on the same chip at 2K tokens and above. It covers twelve LLM configurations, from BERT to DeepSeek, up to 128K tokens. Roofline analysis at each memory level explains this result and shows when to stop. XDNA 1 has lower ridge points, so streaming on chip already reaches the compute-bound regime: the same fusion that doubles throughput on XDNA 2 is nearly wasted on XDNA 1. Comparing a mapping's operational intensity against each level's ridge point predicts which case applies before writing any code. Fuse until the mapping clears that ridge point, then stop. We release the reference designs as maintained open source.

arXiv ID: 2609.21264 / 要約の誤りについて