Attention Through Arithmetic Intensity

目次 Su Jianlin氏の投稿を読んでいたところ、ある一文が私の足を止めました。「KVキャッシュに加え、デコーディングにはもう一つの変数がある — MTP、つまり投機的デコーディングだ。その考え方は、計算量と引き換えに速度を得るというものだ。しかし、MLAはデコーディング中にMQA（head_dims=512+）のように振る舞い、すでに計算量の大部分を前もって消費しているため、『MLA+MTP』は不利になりがちだ。」 私の最初の反応は、「それはどういう意味だろう？なぜMLAはデコード中に『計算量を前もって消費する』のだろうか？そして、なぜMTPと衝突し、そのデコードFLOPsがhead-dim-512+ MHAと同等になるのだろうか？」 この疑問を掘り下げていくと、古典的でありながら、Attentionに適用すると異常に美しいものに行き着きました。それは算術強度 — バイトあたりのFLOPsです。そして、その答えは驚くほど明快であることがわかりました（以下のすべてはBF16 KVキャッシュを想定しています）。 MHAを完全に削減すると、そのAIはちょうど1になります。GQAとMQAも同様に明快です — それらはヘッド数にのみ依存します。 コンテキスト長とヘッド次元は完全に相殺されます。MLAも同じ形状を持ち、潜在次元にも依存しません。定数は2弱です。 並べると、単一トークンデコードのAttentionコアのAIは以下のようになります。 Attention キャッシュが保持するもの AI、おおよそ MHA 各クエリヘッドが独自のKとVを持つ 1 GQA クエリヘッドのグループが1つのKとVを共有する クエリヘッド数 / KVヘッド数 MQA すべてのクエリヘッドが1つのKとVを共有する クエリヘッド数 MLA 1つの潜在空間からKとVの両方が拡張される クエリヘッド数の約2倍 セクション2と4では、これら4行を導出します。最初の3つはKVヘッド数が異なる同じ公式です。最後の行の定数は全く別の場所から来ています。 そして、この2弱という定数は、多くの現行GPUでAttentionデコードを、明らかにメモリバウンドからルーフラインの膝近くに座る状態にするのに十分です。 MTPをその上にスタックすると、ワークロードはコンピュートバウンドに傾きます — これがまさにSu氏がMLAがMTPに不向きだと言う理由です。 この記事では、完全な導出を行います。これは詳細に記述されています — 行列乗算でFLOPsを数える方法から始まり、すべての行列形状が明記されています。 Attentionの構造とデコードが何を計算するかをすでに知っている場合は、セクション2を結果の2.5までスキップし、セクション3にジャンプできます。 1. 算術強度とは何か 定義は簡単です。 つまり、HBMから読み込まれるバイトごとに、どれだけの浮動小数点演算が得られるかということです。 低いAIは、GPUがほとんどの時間データを移動に費やしていることを意味します。高いAIは、データが到着すると、そのデータが多くの算術演算のために再利用されることを意味します。 ハードウェアには対応するしきい値があります。 理想的なルーフラインモデルでは： そして BF16テンソルコアのスループットを使用すると、いくつかの一般的なカードの理論的なバランスポイントは以下のようになります。 GPU Dense BF16 Peak HBM Bandwidth Theoretical Balance Point H100 SXM ~989.5 TFLOP/s 3.35 TB/s ~295 FLOP/B H200 SXM ~989.5 TFLOP/s 4.8 TB/s ~206 FLOP/B B200 (HGX) ~2.25 PFLOP/s ~8 TB/s ~281 FLOP/B 出典: NVIDIA H100, NVIDIA H200, NVIDIA HGX B200, NVIDIA DGX B200。 実際のカーネルがピークFLOPsとピーク帯域幅を同時に飽和させることはないため、これらはプロファイラで見られる線ではなく、直感を構築するためのルーフラインの上限として扱ってください。 「数百FLOP/B」というのが覚えておくべきオーダーです。これは後で比較されます。 最初に述べておくべき注意点があります：AIは比率です。それは「ルーフラインのどちら側にいるか」を答えるものであり、「どちらのアプローチが速いか」を答えるものではありません。分子と分母の両方が成長し、すべてが遅くなる間にAIは変更されないままになる可能性があります。 セクション5には具体的なケースがあります：同じものを計算する2つのアルゴリズムで、AIが高い方が120倍のFLOPsを実行します。 以下は、AttentionのKV関連部分のみをカウントします。 1つのデコードされたトークンについて：隠れ状態から、Q、K、Vを計算し、KVキャッシュを読み込み、このレイヤーのAttention出力まで処理します。 Softmaxは2つの大きな行列乗算に比べて小さいため、除外されています。 目標は、レイテンシ全体を推定することではありません。それは、1つの質問を分離することです：Attention構造の変更が、その区間の算術強度に何をするのか？ セクション2から4はデコードのみ（一度に1トークン、履歴はキャッシュから読み込む）を扱います。プリフィル — シーケンス全体を一度に計算する — はセクション5まで待つことになります。そこでは、同じモデルが2つのフェーズで逆のアプローチを望むことがわかります。 PS：なぜ最終的なWOをカウントしないのか？それは、選択したAttention構造とは無関係だからです。WOは常にヘッドごとの連結された出力を受け取ります。その幅はHqとdvのみに依存します — KV側がどのように編成されているかはWOには見えず、MHA、GQA、MQA、MLAはすべて同じ幅のものを渡します。 Q/K/V射影と同様に、これは重み×ベクトルであり、KVキャッシュには触れず、Lと共に増加せず、バッチ=1ではそのAIは2/bに固定されます。 それを含めると、すべての構造に同じ定数が追加され、比較が薄まります。 MLPや通信についても同様です。 2. MHA / GQA / MQAの1つの公式 2.1 FLOPsの数え方 (m × k)行列と(k × n)行列の乗算は、m·n·k回の乗算・加算（MAC）を必要とします。1回のMACは乗算と加算の1回ずつなので、2 FLOPsです。 2.2 MHA、簡潔に 最も基本的なAttentionはMHAです：隠れ状態をヘッドに分割し、各ヘッドが独自のQuery、Key、Valueを計算し、Attentionを独立に実行し、連結します。 形状から始めます。 レイヤーは現在のトークンの隠れ状態を受け取ります。 3つの射影がそれをQuery、Key、Valueに変換します。 MHAでは、3つすべてが同じヘッド数ですが、GQAとMQAは後でK/Vヘッド数を減らすため、独自のシンボルが与えられます。 Hq クエリヘッド、Hkv KVヘッド。MHAはHkv = Hqの場合です。 ここで1つ決めることがあります：Keyヘッド数とValueヘッド数は、数学的に一致する必要はありません — 8グループのKと4グループのVでも紙の上ではうまく機能します。しかし、実際のモデルはそのように構築されていません。KとVはペアでキャッシュされ、位置のkを格納することは常にvを格納することを意味します。この投稿ではその慣例に従い、両方をHkvと呼びます。 ヘッド次元については：クエリとキーは一致する必要があります。そうでなければ、qとkをドットできません — dkと呼びます。Valueの次元は異なる場合があります。dvと呼びます。 ヘッドあたり： 新しいkとvはKVキャッシュに追加されます。履歴長Lで、レイヤーのキャッシュは以下のようになります。 次に、各クエリヘッドは3つのステップを実行します。 g(h)をクエリヘッドhが使用するKVヘッド（MHAではg(h) = h）とします。 1つ — スコア。現在のクエリを、キャッシュされたすべてのL個のキーとドット積します。これは(1 × dk) × (dk × L)の積です。 2つ — 正規化。Softmaxがスコアを重みに変換します。形状は変わりません。 3つ — 重み付き和。重みをValueに掛けます。これは(1 × L) × (L × dv)の積です。 Hq個のヘッドごとの出力が連結されて、レイヤーのAttention出力になります。 2.3 FLOPs 3つの部分があり、それぞれに2mnkの形状が代入されます。 1つ — このトークンのQ、K、Vを計算します。3つの(1 × dmodel) × (dmodel × ·)の積。m=1です。 2つ — スコア。クエリヘッドごとに、1つの(1 × dk) × (dk × L)の積。Hqヘッド分です。 3つ — 重み付き和。クエリヘッドごとに、1つの(1 × L) × (L × dv)の積です。 これら3つすべてを合わせると、レイヤーのAttention演算になります。 2つの項は完全に異なる振る舞いをします。Fprojは、現在のトークンのみを処理するため、履歴長Lに依存しません。一方、2番目の項は、履歴全体をスキャンするため、Lに比例します。 それらの比率は、オーダーとしては以下のようになります。 （MHAのHkv = Hqおよびdk = dvを代入すると、ちょうど1.5 dmodel/Lになります。） その比率は、聞くほど小さくはありません。dmodel = 7168の場合、10752/Lとなります。L = 8Kでは、射影はAttentionよりもコストが高く（1.3倍）、32Kでは33%になり、10万トークンを超えると10%を下回ります。したがって、Fprojを小さいためという理由で省略することは正当化されません。 それを正当化するのは、それが比較とは無関係であるということです。Fprojは重み×ベクトルであり、KVキャッシュではなく重みを読み込むため、どのAttention構造を選択してもAIは2/bに固定されます。FattnはLと共に増加する唯一の項であり、構造が変更する唯一の項です。 以下のAIは、それだけを保持します。 とはいえ、Fprojパス自体を忘れないでください。セクション3では、MLAが行うことは、Lで乗算されていた作業を、Lに依存しないパスに移動させることであると示しています。