B200アテンションカーネルをゼロから約SOTAレベルまで60の図で解説

FlashAttention-4 をゼロから約SOTAレベルまで60の図で解説（CUDA/PTX）

［動画: 動画をダウンロード］

以下のB200アテンションカーネルを構築することで、これらの動画を生成します。

この記事で実装されたBlackwell B200アテンションカーネルを使用したビデオモデルによって生成されました。以下の「キャップストーンプロジェクト」を参照してください。

［画像: 14のカーネル進行と終盤でカバーされるマイナー最適化における、Stock FA4の同一実行率。各実装は、最適な連続レイアウトでタイミングが測定されています。ベンチマークキャリブレーションを参照してください。］

14のカーネル進行と終盤でカバーされるマイナー最適化における、Stock FA4の同一実行率。各実装は、最適な連続レイアウトでタイミングが測定されています。ベンチマークキャリブレーションを参照してください。

このブログについて

このブログでは、CUDAと少量のPTXを用いて、ゼロから密なB200アテンションカーネルを構築します。ベースラインから、FA4論文で使用された4K、8K、16Kの形状でFlashAttention-4のパフォーマンスの94.4%に達します。

主な貢献は、ビジュアルガイドです。60個の図を中心に構築された、初心者フレンドリーな進行です。まず、ナイーブカーネルの動作原理を直感的に理解し、次に詳細な図、簡潔な説明、コードとともに、最適化を一つずつ追加していきます。

キャップストーンプロジェクトでは、最終的なカーネルをビデオ生成モデルに組み込みます。

ここでは、パフォーマンスの最後の数パーセントを絞り出すことは目的としていません。それは私の次のブログの焦点となります。

これは、最新のハードウェアで動作する、最も難しいカーネルの一つなので、楽しいでしょう。

---

リサーチ: これは、最新ハードウェアでの独自のGPUカーネルリサーチのための基盤を提供することを目的としています。B200アテンションに焦点を当てていますが、多くの概念やメンタルモデルはこれ以外のカーネルにも適用できます。最終的には、このカーネルを理解し、独自の最適化アイデアを思いつくことができるようになるでしょう。

前提条件: 初心者向けオンランプを含めました。もしあなたが完全な初心者であるか、CUDAの基礎が shakyだと感じる場合は、続行する前にこの脚注1に従ってください。基本的なCUDAの知識はあると仮定しますが、Blackwellに関する予備知識は不要です。新しい概念は視覚的に、一つずつ導入され、必要になったときにのみ提示されます。したがって、Blackwellが初めての場合でも、そのまま読み進めてください。進行についていけるはずです。

アテンションとは何かを知っていると仮定します。また、オンラインソフトマックスの背後にあるアイデアにもある程度精通していると仮定します。この部分はBlackwell固有ではなく、それに関するリソースは plentyあります。

追従するためにB200を自宅に所有する必要はありません。私も持っていません。

コード: すべてのコードは私のB200 Attention repoで利用可能です。ベースラインの後、各章で1つの主要な最適化が追加されます。ソースコードは、隣接するカーネルがほとんど簡単にdiffできるように整理されています。したがって、最適化が一度に何を変更したかを確認できます。

キャップストーン: 最後に、あなたが理解したカーネルをビデオ生成モデルに組み込み、美しいビデオを生成します。メインチャプターの後にあるキャップストーンプロジェクトを参照してください。

DSL vs CUDA: 元のFA4はCuTeで書かれていますが、私は個人的に生のCUDA + 少量のPTXの方が理解しやすい（抽象化レイヤーが少ない）と考えているため、CUDAでカーネルを実装します。FA4のCuTe実装をCUDA構文に翻訳するのではなく、クリーンで直感的な進行を通じて、一般的な高速B200アテンションカーネルを理解します。それでも、FA4はこの作業の主要な参照の一つであり、ほとんどの最適化アイデアはそこから適応されています（そしてFA4自体はcutlassとcute-dslカーネルに基づいています。謝辞を参照してください）。

既存リソース: H200およびB200での最適化されたmatmulsを説明する優れたリソースがあります（謝辞を参照）。しかし、B200アテンションについては、私が望むような深い解説を見つけられませんでした。一部のリソースは、最終的なカーネル結果をカバーしていますが、進行やほとんどの最適化の背後にある低レベルの動機を説明していません。他のものは、パイプラインやワープの役割を要約するようなハイレベルで表面的にとどまり、実際にカーネルを理解するために必要なほとんどの作業と手助けをスキップしています。どれも私が望むような深い解説を提供してくれませんでした。

スコープ: 最適化するカーネルは、密、ヘッド次元128、非因果、BF16です。

---

もしあなたにAIしかないなら、私たちはあなたと同じAIを持っており、おそらくそれを使うのが得意です

– tomcr00se

Part I — 基本

Chapter 1 – Baseline Blackwell B200 Attention

i. このチャプターのロードマップ

まず、各CTAがどのような作業を行うのか、そしてB200の異なるCTAに作業がどのように割り当てられるのかを調べます（ワーク並列化セクション）。

次に、初心者向けに、論理的表現と物理的表現、ポインタ、タイルド行列乗算についてオプションの寄り道をします。

その後、b200アテンションカーネルにズームインし、各CTA内で何が起こるかを議論します。

これらはすべて後で視覚的に詳細に説明されます。今は全体像を示すだけです。その後、コードにリンクします（図と直接対応します）。

ベースラインアテンションカーネル

最初のカーネルの理解から始めましょう。後続のカーネルは、ここで紹介された数学とtcgen05の概念のほとんどを同じように使用します。したがって、このチャプターでは、私たちが全体を通して使用する基礎をカバーします。だからこそ、最初のチャプターは後続のチャプターよりも長くなっています。

最初のカーネルは出発点として存在し、正しい数値結果を生成しますが、後続のカーネルほど効率的ではありません。この最初のカーネルはすでに非自明であり、多くのBlackwell固有の機能を使用しています。これらは後で徐々にカバーしていきます。

主なボトルネック

アテンションは基本的にソフトマックスを挟んだ2つの行列乗算です。ソフトマックスはFLOPs（浮動小数点演算数）ははるかに少ないですが、ALUとMUFUユニットで実行され、テンソルコアでは実行されません。そしてBlackwell B200では、テンソルコアのスループットは約2倍になりましたが、expユニットはほとんど変更されませんでした（FA4論文ではこれを非対称ハードウェアスケーリングと呼んでいます）。したがって、私たちのタイルサイズでは、ソフトマックスは非常に強力なMMAと同じくらいのサイクル数を消費するため、このカーネルの主なボトルネックとなります。私たちの最適化のほとんどは、この2つの側面からこれを攻撃します。ソフトマックスの作業をより安価にし、それを行列乗算とオーバーラップさせる（つまり、MMAの影に隠す）ことです。

ii. ワーク並列化

B200固有の詳細を議論する前に、まず作業がどのように分割されるかを見てみましょう。

［画像］

Q、K、V、および出力テンソルはすべて同じ形状 B, num_heads, seq_len, head_dim を持つと考えてください。

テンソルをタイルに分割し、高速だが小さいオンチップメモリに収まるようにします。上記の図に示すように、私たちのタイルサイズは [128, 128] です。

［画像］

出力（O）タイルの数だけCTAをスケジュールします。各CTAは単一のOタイルを生成します。そして、すべてのCTAが集合的に出力テンソル全体（すべてのタイル）を生成します。CTAは並列に実行されます（紫色の矢印）。

各CTA内では、K/Vループの作業はシーケンシャルです（オレンジ色の矢印）。1つのCTAにズームインしてみましょう（上記参照）。そのOタイルを生成するために、CTAは対応するQタイルを1回ロードし、次にすべてのK/Vタイルペアをループします。ループの各イテレーションで、S_tile = Q_tile @ K_tile.T を計算し、オンラインソフトマックスを適用してP_tileを生成し、P_tile @ V_tileをこのCTAのプライベートバッファO-tileに累積します。CTA内では、そのCTAプライベートのOタイルは「実行中のアキュムレータ」として使用されます（CTAはK/Vループの各イテレーションでそれを更新します）。K/Vループの最終イテレーションの後、CTAはそのOアキュムレータを正規化し、完了した出力タイルをグローバルメモリに格納します。

これらの一部がまだ理解できなくても心配しないでください。後で各ステップを詳細に説明します。

基本的に、この分割はタイルド行列乗算に似ています。異なるCTAが異なる出力タイルを生成し、1つの出力タイルのループはそのCTA内に残ります。しかし、行列乗算とは異なり、アテンションは追加でオンラインソフトマックスの状態をK/Vループ全体で保持します。

擬似コードでは、ハイレベルなフローは次のようになります。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22

# 紫色の矢印: 独立した作業項目がCTA全体にマッピングされる
parallel_for batch, head, q_tile:

    # ここから1つのCTAが開始
    Q = load_q_tile(batch, head, q_tile)

    # オレンジ色の矢印: このCTA内のシーケンシャルループ
    for K, V in kv_tiles(batch, head):

        # セクションA: Q @ K.T -> S
        S = Q @ K.T

        # セクションB: S -> P, rowmax/rowsumの更新、および古いOの補正
        P, softmax_state, O = online_softmax_update(
            S, softmax_state, O
        )

        # セクションC: 現在のP @ Vの寄与を追加
        O += P @ V

    # この単一のOタイルの最終正規化
    output[batch, head, q_tile] = O / softmax_state.rowsum

Qタイルの外側のループがCUDAグリッドになります。KVタイルの内側のループは各CTA内に残ります。

上記の擬似コードでラベル付けされたセクションは、私たちの最初のカーネルのB200コードに直接マッピングされ、このチャプターで詳細に扱われます。

• セクションA) Scores = Query @ Key.T (マイクロタイルでのtcgen05、Kメジャーレイアウト)
• セクションB) Scores -> Probabilities (TMEMからの読み取り)
• セクションC) O += Probabilities @ Values (マイクロタイルでのtcgen05、MNメジャーレイアウト)

Chapter 1全体を通して、このハイレベルなピクチャを念頭に置いてください。擬似コードの各ステップがカーネルによってどのように実装されるかにズームインしていきます。

iii. 初心者向けオンランプ — オプション

このセクションはオプションです。 タイルド行列乗算、メモリレイアウトの概念、データの論理ビューと物理ビューの区別、およびポインタを使用してメモリにインデックスを付ける方法にすでに精通している場合はスキップしてください。

タイルド行列乗算との類似性:

まず、後で利用する基本的なタイルド行列乗算の直感を構築しましょう。

［画像］

通常のタイルド行列乗算では、両方のオペランド行列をタイルに分割し、小さいタイルが高速だが小さいオンチップメモリに収まるようにします。これにより、各ロードされたタイルをオンチップメモリから複数回直接再利用できます。各イテレーションで、削減次元（図ではK）に沿って進み、対応するAとBのタイルを乗算し、結果の部分和を同じCタイルに累積します。すべてのKタイルを処理した後、これらの部分和を追加することは、Kを分割せずに完全な行列乗算を行うことと数学的に同等です。アテンションは形状は似ていますが、正確ではありません。

• 一般的なタイルド行列乗算は、1つのCタイルを固定し、その削減次元に沿って歩行します。
• アテンションは、1つのQタイルを固定し、K/Vシーケンス次元に沿って歩行します。

レイアウトの導入:

したがって、グローバルメモリにデータのタイルがあります。

AタイルとBタイルの行列乗算を行いたいとします。ソフトウェアでこの行列乗算を実装したくありません（個々の要素をループしてスカラー命令でドット積を計算する）、これは非常に遅くなります。代わりにテンソルコアを使用したいです。Blackwellでテンソルコアを使用するには、tcgen05.mmaファミリーの命令を使用する必要があります。

上記のtile_aとtile_bは2D行列のように見えますが、もちろんこれらは格納されたデータの論理ビューにすぎず、実際のメモリは1Dです。

［画像］

したがって、データが物理メモリにどのように「配置」されているかについての組織化が常に存在します。

AオペランドがGMEMで行メジャーであるとしましょう。その場合、物理的には行が1Dメモリに連続して格納されます（上記参照）。

線形メモリ内の要素を見つけるには、[row, col]の要素について、ベースポインタから開始し、ストライドを使用してオフセットを計算します。

小さな[2,2]行列があり、要素my_matrix[1,1]を取得したいとします。

［画像］

したがって、物理メモリに「配置」されているこれらの2Dタイルがあります。レイアウトとは、[row, col]のような論理座標を物理メモリアドレスにマッピングするルールです。行メジャーはそのようなレイアウトの1つです。

しかし、呼び出したいtcgen05.mma(tile_A, tile_B)は、その単純な行メジャー物理レイアウトでは動作しません。それはごく少数の非常に特定のレイアウト（単純な行メジャーや列メジャーだけでなく、mmaが期待するレイアウトが具体的に何であるかは、今はそれほど重要ではありません。後で議論します）しかサポートしていません。

基本的に、tcgen05はデータが特定のレイアウトのいずれかに配置されていることを期待します（つまり、データがメモリにどのように配置されているか）。さらに、tcgen05はGMEMオペランドをサポートしません（GMEM帯域幅が低すぎるため、オペランドはまず高速なオンチップメモリに書き込まれる必要があります）。

このセクションの要点は、これらのAとBのオペランドをSMEMに、mmaがサポートする特定のレイアウトのいずれかに配置する必要があるということです。

---

メインフローに戻る（オンランプから出る）：

最初の2つの図と擬似コードで見たように、各CTA内で、Qは固定され、K/Vループの各イテレーションは以下を計算します。

• A) Q @ K.T -> S
• B) S -> P
• C) P @ V -> O

これらは、KVループの各イテレーション（SEQ_LENに沿って）が行う3つの主要な計算です。K/Vループの単一イテレーションにズームインし、これらの3つのステップを1つずつカバーしましょう。

A) Q@K.T (Kメジャーの説明)

1つのK/Vループイテレーション内 — ステップ1/3。

QKを実行する前に、SMEMにQと現在のKタイルが必要です。QはCTAごとに1回ロードされ、新しいKタイルはK/Vイテレーションごとにロードされます。QとKは両方とも同じレイアウトを使用するため、以下ではQのコピーのみを描き、Kのコピーは省略します。

［画像］

注意：QK行列乗算の図とコードスニペットでは、従来の行列乗算の次元名を使用します：Q[M,K] @ K[N,K].T = S[M,N]。これらはカーネルの入力グローバルテンソルの次元名ではなく、ローカルな行列乗算の軸名です。

前述のように、Blackwellのテンソルコアはtcgen05ファミリーの命令で表されます。そしてtcgen05.mmaは、オペランドが特定のレイアウトで配置されていることを期待します（任意のレイアウトでグローバルメモリから共有メモリにタイルをコピーして、その上でMMAを実行することはできません）。その制約があるため、データをこれらのMMAがサポートするレイアウトの1つに「配置」する必要があります。

SMEMにデータをどのように配置すれば、MMAがサポートするレイアウトの1つになるのでしょうか？簡単にするために、QタイルのGMEMからSMEMへのコピープロセスのみを説明します（同じロジックがKタイルにも個別に適用されます）。論理的には、そのQタイルを垂直に分割し（図に示すように - ここでは視覚的に合計4つのスライス）、次にTMA2で各スライスを個別にGMEMからSMEMにコピーします（上記の図を参照）。

なぜタイル全体を一度にコピーするのではなく、タイルをこのようにスライスしてコピーする必要があるのですか？ 良い質問です。この考えをしばらく保留してください。一度に多くのことを紹介したくないので、それが重要になったときにこの点に戻ります。今はスライスでコピーすることを覚えておいてください。

すべてスライスをコピーした後、GMEMからSMEMへの移動で何も変わらなかったように見えます（視覚化されたように、論理的な[BLOCK_M, BLOCK_K]行列は同じですが、今度はGMEMではなくSMEMにあります）。しかし実際には、物理的には、表現は変更されました。物理ビューの緑色の注釈を参照してください。したがって、SMEM内の結果データは、GMEMにあった行メジャーレイアウトではなくなります。物理SMEMレイアウトでは、1つのスライスのすべての行、次に別のスライスのすべての行などが配置されます。

したがって、これは各スライス内では行メジャーですが、全体としては（SMEM内の[BLOCK_M, BLOCK_K]全体のレイアウト）単純な行メジャーまたは列メジャーではありません。むしろ、これはtcgen05.mmaがサポートするいわゆる「標準レイアウト」の1つです。この用語を覚えておいてください。後で頻繁に目にします。

［画像］

同じタイルのデータ（Qタイル）がGMEMとSMEMでどのように異なるかを示す別の視覚化（上記）。ピンクの矢印は要素の物理的な順序を示します。GMEMは通常の行メジャーでしたが、SMEMはtcgen05が処理する方法を知っている標準レイアウトです（この特定のレイアウトはKメジャー・ノースウィズルと呼ばれます - 後でこれに戻ります）。

Kメジャーの「K」はKeyテンソルを指しません。「K」は標準的な行列乗算の命名規則における削減次元を意味し、この場合は論理タイルの水平方向に視覚化されます。上記の各垂直スライスは、そのK次元の小さな部分をカバーします。

今のところ、この部分を無視して、簡単のためにこのレイアウトをtcgen05がサポートする唯一のレイアウトと考えてください。繰り返しますが、一度に多くのことを紹介したくありません。後で重要になったときに他のレイアウトについて議論します。それは、TMAで垂直スライスをコピーした結果として物理データの配置がどのように変更されたかを強調するためだけです。

物理レイアウトはなぜこのようになり、他の任意のようにならないのですか？これは直感的ではないように思えますが、tcgen05がサポートする標準レイアウトの1つであると言う以上の深い直感はないと信じています（NVIDIAのドキュメントによると）。最終的には、ハードウェアの契約に頼るしかありません。

---

ここでは、TMAの部分、つまりGMEMからSMEMへのデータのコピーをカバーしました。次に、MMAの部分、つまり準備したSMEMデータで実際にMMAを実行する方法についても議論しましょう。

私たちはTMAを使用してデータをグローバルメモリからSMEMにコピーしましたが、それは単なるコピーメカニズムであり、これから呼び出すtcgen05.MMAを本質的に認識していません。そして別途、tcgen05.MMAは特定のレイアウトのセットしか解釈する方法を知らず、読み込むデータがどのようにSMEMに入ったかは認識していません。したがって、次のように考えることができます。ロード部分（前述）とMMA部分（後述）。これらは2つの独立したステップであり、互いを認識していないため、それらが同じレイアウトに同意するようにするのは私たちの仕事です。

［画像］

図の簡略化：図は読みやすくするためにBLOCK_K = HEAD_DIM = 32 BF16要素（64B）を使用しています。コードと以下のテキストはBLOCK_K = HEAD_DIM = 128 BF16要素（256B）を使用しています。以下で説明するすべては同じように機能します。実際のコードにはスライスの4倍が含まれています。

Q_tileとK_tileの両方をSMEMにコピーしました（Kに沿って16Bスライス）。今、この(block_m, block_k) Qタイルを取得し、(block_n, block_k).T Kタイルと行列乗算する必要があります。今のところ、Qタイル（tcgen05.mmaの最初の入力）にのみ焦点を当て、Kタイル（tcgen05.mmaの2番目の入力）は一時的に省略します。

私たちのタイルはBLOCK_K = 128要素（BF16 dtypeの場合、256バイト）を持っています。しかし、1つのtcgen05.mma命令はKに沿ってMMA_K=16 BF16要素（32B）を消費します。3したがって、mma_K幅はタイルK幅よりも小さいため、QとKタイルのBLOCK_K次元を32バイトステップで歩行するforループが必要です（SMEM内のデータに対するソフトウェアループ、上記黄色で表示）。

私たちのQスライスの各々はKに沿って16B幅ですが、1つのtcgen05.mmaはKに沿って32Bを消費します。したがって、隣接する2つのQスライスが1つのQマイクロタイルを形成します。上記の図では、同じ色のスライスのペアが1つのマイクロタイルを形成します。

「マイクロタイル」と言うとき、それはtcgen05が直接消費するBLOCK_M=128, MMA_K=16のチャンクを意味します。「マイクロステップ」は、1つのQマイクロタイルと1つの対応するKマイクロタイルを消費する1回のtcgen05.mma呼び出しです。BLOCK_K=128およびMMA_K=16なので、128 / 16 = 8マイクロステップが完全な削減K次元をカバーします。簡略化された図は、それらの8ステップのうち2つのみを示しています。

［画像］

これをさらに視覚化するために、K_tileも描きました。これはTMAで独立して16B幅のスライスに分割されコピーされ、次にそれらのスライスがペアになってマイクロタイルとして扱われます（以前カバーしたQ_tileのコピーの表示を省略したのと同じです）。

次に、Q_tileとK_tileを一緒に乗算するには、対応するマイクロタイル（上記の図の赤いボックスに示されています）を乗算する必要があります。

• 最初のQマイクロタイル @ 最初のKマイクロタイル.T
• 2番目のQマイクロタイル @ 2番目のKマイクロタイル.T

上記のoutput_accumは私たちのScoresタイル（S）です。Blackwellでは、tcgen05.mmaはテンソルメモリ（TMEM）、SMEMやレジスタとは別の専用オンチップメモリにその出力を累積します。（以前のGPUでは、テンソルコアのアキュムレータは代わりにレジスタにありました）。すべてのQKマイクロステップは、TMEM内の同じSタイルに累積します。今のところ、MMAの出力はTMEMに存在することを覚えておいてください。セクションBで、ソフトマックスがSを読み取る必要があるときに、TMEMの割り当てとアクセス方法を説明します。

上記の直線的なMMA呼び出しを、MMA_K=16マイクロステップでQ_tileとK_tileのBLOCK_K次元を反復するループに変換します。擬似コードでは：

［画像］

これは、実際のカーネルコードからの「黄色い」ループであり、少し簡略化されています（単一のスレッドによって実行されます）。

// tcgen05 mma_k幅（K次元に沿って）
int MMA_K_BYTES = 32; // バイト単位
int MMA_K = 16; // BF16値単位

// 今のところ、各KメジャーSMEM "ディスクリプタ"を
// 現在のマイクロタイルを選択するポインタのような値として考えてください
// （これにはレイアウトメタデータも含まれています。後でカバーします）
uint64_t q_desc = noswizzle::desc_kmajor(Q_smem, BLOCK_M);
uint64_t k_desc = noswizzle::desc_kmajor(K_smem, BLOCK_N);

// 次のQとKマイクロタイルに進むためのディスクリプタをアドバンスするオフセット
// （マイナーな詳細ですが、ディスクリプタアドレスは16B単位でエンコードされているため16で割ります）
uint64_t q_desc_step = (BLOCK_M * MMA_K_BYTES) / 16;
uint64_t k_desc_step = (BLOCK_N * MMA_K_BYTES) / 16;

// BLOCK_K = 128, MMA_K = 16 なので、8マイクロステップ
for (int k = 0; k < BLOCK_K / MMA_K; ++k) {

    // 1マイクロステップ
    tcgen05::mma_f16(
        taddr_s, q_desc, k_desc, 
        /*MMA形状とデータ型=*/i_desc_qk,
        /*入力dを有効にする=*/(k == 0 ? 0 : 1));

    // 次のマイクロタイルのペアに進むためにディスクリプタをアドバンス
    q_desc += q_desc_step;
    k_desc += k_desc_step;
}

// すべての8マイクロステップの完了を追跡するためにmbarrierを設定
// （待機と必要な同期はここでは省略）
tcgen05::commit_arrive(mbar_addr);
// 待機と必要な同期はここでは省略

各イテレーションは、次のMMA_K=16チャンクに対してtcgen05.mmaを呼び出し、同じSタイルに累積します。最初のMMAマイクロステップはSを初期化し、残りの7つのMMAマイクロステップはその同じSタイルに累積します。すべてのマイクロタイルに対するループと完了待機の後、そのS_tileはQ_tile @ K_tile.Tの結果を含みます。

明示的なK_tileの転置は不要です。 QとKは両方ともKメジャーレイアウトを使用します。MMA AオペランドはQ[m,k]を直接読み取ります。Bオペランドの場合、tcgen05.mmaは格納されたKタイルをQ @ K.Tに必要な[k,n]オペランドとして解釈します。Kメジャー自体は「このオペランドを転置する」という意味ではありません。AとBは異なる論理形状[M,K]と[K,N]を持つ異なるMMAオペランドスロットであり、したがってMMAによって異なるように解釈されます。

---

QKマイクロステップループをコードで見てきました。そこでは、q_descとk_descをマイクロタイルを選択するポインタのような値として扱いました。次に、これらのディスクリプタが含むレイアウト情報と、MMAが各マイクロタイル内のデータを見つけるためにそれを使用する方法を解明しましょう。

［画像］

ここにさらにレベルがあります。Qマイクロタイルにズームインしましょう（BLOCK_M=128、MMA_K=16）。ただし、同じロジックがKマイクロタイルにも適用されます。

これらのマイクロタイル自体は、低レベルのユニットで構成されています。アトム（またはCUDAドキュメントの「スウィズルアトム」）と呼ばれます（ここではスウィズリングは行っていません。NVIDIAはスウィズリングモードがNoneの場合でも「スウィズルレイアウトアトム」を包括的な用語として使用しています。テーブル53）。

これらのアトムのサイズは、使用している特定の標準レイアウトによって異なります。この場合（使用しているレイアウトはKメジャー・ノースウィズルと呼ばれますが、今はこれらの詳細を無視して、tcgen05.mmaがサポートするレイアウトが1つだけであると仮定します）、これらのアトムは(8, 8)です。4

これまでのトップダウンストーリー（図の凡例を参照）：マイクロタイルは2つのスライスで構成され、それぞれがアトムで構成されます。アトムがSMEMにどのように配置されているかを理解する必要があります。しかし、それらは独立してプログラム可能なMMAユニットではありません。MMAを発行するとき、私たちはまだマイクロタイルのレベルで操作しており、個々のアトムのレベルではありません。

さて、以前の質問に答えることができます：なぜQとKをTMAで8要素幅の垂直スライスでコピーしたのですか？タイル全体を一度にコピーするのではなく？

前に使用したスライス幅は任意ではありませんでした。実際、それはアトム幅1つでした。Qの場合、使用した各(128,8)スライスは、(8,8)アトムの垂直スタックです。したがって、タイルをこれらのスライスに分割し、TMAを使用して完全なスタックを一度にコピーし、スタックをSMEMに次々と配置します。これにより、tcgen05.mmaがサポートするKメジャーの物理アトム順序が生成されます。

［画像］

上記で議論したように、各K/Vループイテレーションの「黄色い」ループの各イテレーションで、2つのマイクロタイルデータ（1つのQマイクロタイルと1つのKマイクロタイル）をtcgen05.mmaに供給します。

単一のマイクロタイルについて、すべてのアトムのアドレスをmmaに提供する必要はありません。代わりに、最初のアトムのベースポインタを提供する必要があります。図では、これが緑色でハイライトされています。

黄色いループが進むにつれて、ディスクリプタのベースアドレスをマイクロタイルから次のマイクロタイルに移動させます。最初のイテレーションでは、マイクロタイル0の開始であるシアンアトムを指します。2番目のイテレーションでは、マイクロタイル1の開始である紫アトムを指します。K次元全体で同様です。Kオペランドディスクリプタについても同様です。

しかし、ベースが最初のアトムのみを指しているため、直感的には、そのベースアドレスに加えて、MMAは、指定されたマイクロタイル内の他のアトムを見つける方法を知る必要があるということです。

［画像］

そのために、SBOとLBOという概念があります（図で示されています）。これらは追加のソフトウェアループイテレーションを作成するのではなく、ベースポインタに対する同じマイクロタイル内の他のアトムの位置を記述します。SBOは、同じスライス内の次アトムにステップする方法をMMAに伝え、LBOは、次のスライス内のアトム（同じインデックス）に到達する方法をMMAに伝えます。

この場合、議論しているSMEMメモリレイアウト（Kメジャー・ノースウィズル）では、各アトム、したがって各垂直スライスは、幅8 BF16要素（16B）なので、次のようになります。

SBO = 8行 * 16B/行 = 128B   // 1アトムをスキップ
LBO = BLOCK_M行 * 16B/行 = 2048B    // アトムの完全なスライスをスキップ

ここで、LBOはQコピーループで使用したSLICE_BYTESと同じです。SMEM内の隣接する垂直スライス間のオフセットです。各Qマイクロタイルには2つのスライスが含まれます。したがって、ループはディスクリプタのベースアドレスを2 * SLICE_BYTESバイトだけ進めて、次のマイクロタイルに到達します。ディスクリプタアドレスは16バイト単位でエンコードされているため、16で割るとコードで使用されるq_desc_stepが得られます。

これらをまとめると、ソフトウェアループはベースポインタをマイクロタイルから次のマイクロタイルに進め、SBOとLBOは現在のマイクロタイル内の他のアトムへのオフセットを記述します。

B) オンラインソフトマックス (S->Pの説明)

1つのK/Vループイテレーション内 — ステップ2/3。

この時点で、最初の行列乗算（Q@K.T）を終え、カーネルのオンラインソフトマックス部分（S -> P）に近づいています。

オンラインソフトマックスの数学は標準的で、B200固有のものではなく、他の多くの説明があるので、ここでは再導出しません。必要な状態更新を要約し、B200カーネルがそれをどのように実装するかを説明します。

［画像］

TMEMをセクションAで既に簡単に見てきました。Q @ K.TマイクロステップはSをそこに累積しました。ここでTMEM自体にズームインし、オンラインソフトマックスがSをどのように読み取り、Pに変換するかを見てみましょう。

TMEMは（128行、512列）であり、明示的に割り当てる必要があります。K/Vループの前に、カーネル1はS用に128列、O用にさらに128列を割り当てます。（割り当ては列単位で、各列には128要素が含まれるため、128列を割り当てることはTMEMに[128,128] fp32要素のチャンクを割り当てることになります）。

Sは現在のK/Vタイルのスコアを保持し、K/Vループの各イテレーションで上書きされます。Oは実行中の出力アキュムレータであり、後述のセクションCのP@V行列乗算は各K/Vイテレーションでそれに加算されます。tcgen05.allocは割り当てのベースTMEMアドレスを共有メモリのスロットtmem_addrに書き込みます。次に、その割り当てをSとOのアドレスに分割します。

int TMEM_S_COLS = 128;
int TMEM_O_COLS = 128;

__shared__ int tmem_addr[1];

// 簡単化されたPTX表記
// ワープ選択と同期は省略。
tcgen05.alloc(tmem_addr, TMEM_S_COLS + TMEM_O_COLS);

// S用の最初の128列、O用の次の128列。
int taddr_s = tmem_addr[0];
int taddr_o = taddr_s + TMEM_S_COLS;

前のセクションで見た抽象化を開いてみましょう。

# セクションB: S -> P, rowmax/rowsumの更新、および古いOの補正
P, softmax_state, O = online_softmax_update(
    S, softmax_state, O
)

1つのK/Vイテレーションについて、状態は次のように変化します（まず擬似コードで見てみましょう）。

// セクションA、前述
// S = Q @ K.T;

// セクションB:

tile_rowmax = row_max(S);
// 新しい基底
new_rowmax  = max(running_rowmax, tile_rowmax);

// 古い状態を新しい基底に再スケーリング
rescale     = exp(running_rowmax - new_rowmax);
running_rowsum *= rescale;
O              *= rescale;

// Pを生成し、分母を更新
P               = exp(S - new_rowmax);
running_rowsum  += row_sum(P);
running_rowmax   = new_rowmax;

// セクションC、後述:
// O += P @ V;
// 完全なK/Vループの後:
// output = O / running_rowsum;

これは一般的なオンラインソフトマックス更新であり、B200固有ではありません。そのため、ここでは詳しく説明しません。私たちの主な質問は、これをB200カーネルにどのようにマッピングするかです。

上記の擬似コードには多くの個別の状態更新が含まれていますが、私たちのカーネルでは、このオンラインソフトマックスの作業は基本的に3回のTMEMの走査を中心に組織されています。そして、他のすべての状態更新は、これらの3回の走査にほとんど折り込まれています。したがって、オンラインソフトマックス更新については、「データタイルの3回の走査」の粒度で考えてください。

3回の走査は次のようにグループ化できます。

• Sタイルの2回の走査 - Pを生成するために使用され、
• Oタイルの1回の走査 - Oを再スケーリングするために使用されます。

ソースでは、O-リスケール走査は2つのS走査の間に行われます。しかし、ここでは2つのS走査を一緒に説明し、その後Oリ スケーリングに戻ります。

これらを順番に見ていきましょう。コードウォークスルーとともに。

---

i) Pの生成

まず、SタイルをP（未正規化）に変換するために使用される2回の走査について説明します（このPは後続のP@V行列乗算で使用されます）。これを行うために、カーネル1はTMEM内のSの各行を2回走査します。

1. 最初に、行最大値を計算するため（2番目のステップでの数値的安定性のために使用されます）。
2. 次に、行最大値を減算し、結果を指数化して、正規化されていないPを生成するため exp(S - new_rowmax)。

Sの各行の最大値を計算し、その行の値から減算してから指数化することは、マックストリックと呼ばれ、ソフトマックス実装で数値的安定性のために一般的に使用されます。Sには大きな正のスコアが含まれる可能性があり、直接指数化するとオーバーフローする可能性があります。行最大値を減算すると、expへの各入力はゼロまたは負になり、ソフトマックスの結果を変更せずにオーバーフローを防ぎます。したがって、行最大値を最初に減算せずにexp(S)を計算することは、一般的に数値的に安定していません。

私たちの実装が2回の走査を使用する理由は、Sのすべてをレジスタに保持しないためです。最初の走査では、TMEMからチャンクごとにSをロードし、各行の最大値を更新し、次に同じレジスタを次のチャンクのロードに使用します。したがって、最初の走査で完全なSを走査し終える頃には、以前に読み取ったチャンクはレジスタになくなります。Pを生成するには、TMEMからSを再度読み取り、更新された行最大値を減算し、結果の値を指数化するために2回目の走査が必要です。

最初のS走査：行最大値の計算

最初の走査がSをどのように走査して行最大値を計算するかを見てみましょう。まず図、次にコードです。

［画像］

TMEMにアクセスできるのはtcgen05.ファミリーの命令のみであり、通常のロード/ストア命令はできません。したがって、行最大値を計算するために、まずTMEMから値をロードする必要があります（最大値を計算する前に）。そのため、TMEMから値をレジスタにロードするためにtcgen05.ld（ロードの略）命令を使用します。

この図を左から右に読みましょう。CTAレベル、次にワープ、最後にワープ内のスレッドという3つのレベルでTMEMの走査を追います。TMEM内のSタイルは（128行、128列）です。

CTAの観点から、4つのワープがSの128行を分割します。各ワープは32, 128の行チャンクを所有します。そのため、カーネル1には4つのワープがあります（Sの全128行を32行チャンクでカバーするため、全128行を並列に処理できます）。

しかし、各ワープは32, 128チャンク全体を一度に読み取るわけではありません。これはレジスタスピルを引き起こします（これら4つの32, 128チャンクすべてをレジスタに保持するのに十分なレジスタがないため、コンパイラは一部をグローバルメモリに書き込みます。これは信じられないほど遅いです）。これを避けるために、4つのワープのそれぞれは、(32, 8)の小さいステップでその32, 128チャンクを反復します。6 したがって、TMEM列に対するforループがあります。

1つのワープにズームイン： 各ループイテレーションで、ワープはtcgen05.ldを使用してTMEMから(32 rows, 8 cols)チャンクをロードします。32スレッドすべてが1つのロード命令（tcgen05.ldは「ワープコレクティブ」です）に参加します。このロードは、ロードしている32行のチャンクをワープの32スレッドに分散させます。各スレッドが1行あたり8つのFP32値を受け取ります。そして、ワープ内の32スレッド全体で、(32, 8)チャンク全体を保持します。

1つのスレッドの観点から、ループの各イテレーションで、同じ行をたどり、一度に8つの値を受け取ります。そして、8つの値のそれぞれを反復処理し、それまでに見た最大値を保持します。

言い換えれば、ロードはワープコレクティブですが、行の計算はスレッドローカルです。各スレッドは、その行のすべての128スコアを、一度に8つずつ見ます。したがって、各スレッドは最初のS走査でその行の最大値を独立して計算します。そのため、各スレッドはtile_rowmaxスカラーをレジスタに保持できます。

各スレッドは、K/Vループ全体で同じクエリ行を処理し続けます。その行最大値と行合計はイテレーション間で引き継がれますが、tile_rowmaxとtile_rowsumは各タイルごとに再計算されます。

4つのワープ全体で、1回のイテレーションはしたがってSの(128, 8)チャンクを処理します。次に、ループは8列右に進み、繰り返します。128/8=16回のイテレーションの後、4つのワープはSの全128列を走査しました。ループの終わりに、ワープ内の各スレッドは、指定されたTMEM行の最大値を保持します。

---

コードでは、tcgen05::ld_32x32b_x8ヘルパーを呼び出します。.32x32bは、1つの32ビット値あたり1行で32個のTMEM行をロードすることを意味します。.x8は、単一のtcgen05.ld命令内で、8列全体にわたってその形状を繰り返します。したがって、各スレッドは8つの32ビット行値を受け取り、8つのレジスタに格納します。それが前の図に示された(32,8)チャンクです。

ここに最初のS走査をコードで示します。各スレッドは、チャンクを移動しながら、独自のtile_rowmaxを更新します。

float tile_rowmax = -FLT_MAX;

for (int score_n8 = 0; score_n8 < BLOCK_N / 8; ++score_n8) {
    // レーンあたり8つのFP32一時変数
    float s8[8];

    int addr = ...; // 以下でカバー
    tcgen05::ld_32x32b_x8(addr, s8);

    for (int i = 0; i < 8; ++i) {
        tile_rowmax = fmaxf(tile_rowmax, s8[i]);
    }

}

では、その命令にどのチャンクを読み取るかをどのように指示しますか？

Sタイルは[BLOCK_M, BLOCK_N] = [128,128]であり、それぞれTMEMの行と列にマッピングされます。前述のように、4つのワープは別々の32行領域をカバーし、各ワープのループは列を横断します。

TMEMアドレスは、通常のバイトアドレスではなく、2つの座標を含む32ビット値です。

[ 上位16ビット: 行 ][ 下位16ビット: 列 ]

PTXはTMEM行を「レーン」と呼びます。このロードでは、これらの座標はワープチャンクの開始行と列を指定します。各スレッドの行と列を個別に指定するわけではありません。

次のようにアドレスを構築します。

// このワープ領域の最初の行：
// warp_idによって0、32、64、または96のいずれか
int row_base = warp_id * 32;
// 上位16ビットにシフト
int trow = row_base << 16;

// S列のループ内
int taddr =
    trow                // 開始行、上位ビットにエンコード
    + taddr_s           // Sの開始列
    + score_n8 * 8;     // S内の現在のチャンクの列オフセット

<< 16はrow_baseを上位16ビットに移動させ、下位半分をゼロのままにします。列番号は既に下位16ビットを占めています。その数値が小さいため、上位半分はゼロです（シフト不要）。少し視覚化すると（各括弧は16ビットを表します）。

row_base << 16:  [row_base][  0   ]
column:          [   0    ][column]
sum:             [row_base][column]

(ここで、column = taddr_s + score_n8 * 8)

2つのフィールドが重複しないため、通常の加算により、行と列が1つのTMEMアドレスに結合されます。

注意：すべての32スレッドが同じtaddrを渡します。このロード形状では、ハードウェアは各スレッドのレーンIDを使用して、行row_base + lane_idからの値を提供します。通常のld.globalのようなスレッドごとのアドレス指定とは異なり、アドレスにlane_idを自分で追加することはありません。例えば、row_base = 32の場合：

• レーン0は行32からの値を受け取ります。
• レーン1は行33からの値を受け取ります。
• …
• レーン31は行63からの値を受け取ります。

これはワープコレクティブロードの意味です。指定されたワープ内のすべての32スレッドが同じロード命令に参加します。ここでは、すべて同じtaddrを使用しますが、それぞれが自身の行からの8つの値を受け取ります。

ループが進むにつれて、trowは固定され、列オフセットは8ずつ増加します。

---

Sの2回目の走査の前：実行中の行最大値の更新

このSタイルの行ごとの最大値を計算したので、オンラインソフトマックスの定義に従って、実行中の行最大値と行合計を更新する必要があります。最初のS走査はtile_rowmax（現在のSタイル内のこの行の最大値）を見つけました。しかし、2回目の走査は実行中の最大値を使用します。これは以前のK/Vタイルも考慮に入れています。2回目の走査の前には、tile_rowmaxを以前の行最大値と組み合わせ、古い行合計も再スケーリングします。これらはレジスタ内のスレッドごとのスカラー更新であり、Sの別の走査ではありません。

2つの最大値のうち大きい方を取り、スケーリング係数を計算します。

float new_rowmax = fmaxf(rowmax, tile_rowmax);
float rescale = __expf(rowmax - new_rowmax);

rowmax = new_rowmax;
rowsum *= rescale;

現在のタイルが実行中の行最大値を増加させた場合、古い行合計は以前の最大値を使用して計算されました。したがって、それらの以前の寄与が更新された最大値と一致するように、スケーリング係数を掛けています。

2回目のS走査は、Pを生成するためにこの更新された行最大値を使用します。同じスケーリング係数は古いOにも適用する必要があります。これは後述のサブセクションiiでカバーします。

---

2回目のS走査：Pの生成と行合計の更新

更新された実行中の行最大値が得られたので、Sに対する2回目の走査を行います。Sをチャンクごとに再度読み取り、それらの行最大値をスコアから減算して、結果を指数化して正規化されていないPを生成します。

最初の走査と同じTMEM走査を繰り返します。したがって、走査図は別途ありません。変更されるのは、各スレッドがロードした8つの値に対して行うことです。更新された実行中の行最大値を減算し、指数化し、結果を行合計に追加し、PをSMEMに書き込みます。

この2回目の走査がPをSMEMにどのように書き込むかを詳しく見てみましょう。

［画像］

読み取ったSの各チャンクについて、対応するPのチャンクを生成し、KメジャーSMEMレイアウトに書き込みます。これは後続のP@V MMAが期待するレイアウトです。

2回目のS走査の各イテレーションは、Pの8列チャンクを生成します。隣接する2つのチャンクが1つの16列Pマイクロタイルを形成し、後続のPV MMAマイクロステップで使用されます。

S自体はTMEM内で変更されません。Pは以前Kが保持していたのと同じSMEMバッファに格納します。KはQKの後には不要になるため、そのストレージをPに再利用するのは安全です。

ここにコードでの2回目のS走査を示します。TMEMからSを読み取り、SMEMにPを書き込みます。

int row = row_base + lane_id;
float tile_rowsum = 0.0f;

for (int score_n8 = 0; score_n8 < BLOCK_N / 8; ++score_n8) {
    float s8[8];

    int addr = ...; // 最初の走査と同じTMEMアドレス計算
    tcgen05::ld_32x32b_x8(addr, s8);

    for (int i = 0; i < 8; ++i) {
        float p = __expf(s8[i] - rowmax);
        tile_rowsum += p;

        // PをKメジャーSMEMに書き込む
        int col = score_n8 * 8 + i;
        int p_idx = idx_kmajor_w8(BLOCK_M, row, col);
        P_smem_ptr[p_idx] = __float2bfloat16_rn(p);
    }
}

rowsum += tile_rowsum;

PをSMEMに書き込む

idx_kmajor_w8は、後続のPV行列乗算に必要な、私たちの馴染みのあるKメジャーSMEMレイアウトの各P値がどこに配置されるかを計算します。ヘルパーの実装自体は理解する上でそれほど重要ではありません。ここではrow変数が必要です。なぜなら、ワープコレクティブTMEMロードとは異なり、これらは通常の個々のスレッドSMEMストアだからです。

---

ii) Oのリスケール

ここで、2つのS走査の間でカーネルコード内で行われるO-リスケール走査に戻りましょう。以前に新しい最大基底を計算しました。しかし、古いOタイルはまだ古い基底にあります。したがって、O += P @ Vを実行する前に、Oをリ スケールする必要があります。オンラインソフトマックスアルゴリズムでは、古いOが現在の最大基底にあるまで、PVをOに累積するのは一般的に安全ではありません。

以前に実行中の行最大値を更新して、現在のSタイルを含めました。行最大値の変更は、正規化されていないP値のスケールを変更するため、Oに累積されたP@Vの寄与にも影響します。古いOをリ スケールして、現在のP@Vの寄与を追加する前に、それらの以前の寄与が更新された最大値と一致するようにします。

これは標準的なオンラインソフトマックスの定式化です。ここでの唯一の変更点は、これをBlackwell B200ハードウェアにマッピングしていることです。

前述のように、行合計と行最大値はスレッドごとのスカラーです。Oは異なります。これはTMEMに格納された完全な[BLOCK_M, HEAD_DIM]の部分的なアキュムレータタイルです。

したがって、これらの値をリ スケールするには、Oを保持するTMEMの列を読み取り、それらをリ スケールし（古い行最大値と新しい行最大値から計算したスケーリング係数で）、リ スケールされた値をTMEMに書き戻す必要があります。これは実行中の部分的なOであるため、最終的な出力ではないため、Oを書き戻す必要があります。これにより、後続のPV行列乗算が更新された基底にその部分を累積できるようになります。

各スレッドは、O行のすべての値に対して独自のリスケール係数を使用します。最初のK/Vイテレーションでは、以前のOがないため、この走査をスキップします。最初のPVマイクロステップがOを初期化します。

float o8[8];

for (int out_n8 = 0; out_n8 < HEAD_DIM / 8; ++out_n8) {

    // 以前と同じTMEMアドレス計算、今度はO用
    int taddr =
        trow                // 開始行、上位ビットにエンコード
        + taddr_o           // Oの開始列
        + out_n8 * 8;       // O内の現在のチャンクの列オフセット
    tcgen05::ld_32x32b_x8(taddr, o8);

    for (int i = 0; i < 8; ++i) {
        o8[i] *= rescale;
    }

    tcgen05::st_32x32b_x8(taddr, o8);
}

これは、S走査でカバーされたものと同じ4つのワープ、行の所有権、および8列チャンクを使用します。新しい部分は、レジスタ内の値をリ スケールした後、tcgen05.stを使用して、修正されたOをTMEMに書き戻すことです。

ストア完了待機と同期はここでは省略されています。完全なカーネルは、PVがアキュムレータを使用する前にOストアを完了します。

PがSMEMで準備され、OがTMEMでリ スケールされたので、セクションCはP @ VをOに加算できます。

C) P@V (MNメジャーの説明)

1つのK/Vループイテレーション内 — ステップ3/3。

Pタイルが生成された方法をカバーしましたが、今度はVタイルと行列乗算する必要があります。判明したことですが、Q、K、P（すべて同じKメジャーSMEMレイアウトを使用、前述）とは異なり、Vは異なるSMEMレイアウトを使用します。それが次の議論のトピックです。

カーネル1では、P_tileはソフトマックスによってSMEMに生成されるため、以下のグローバルメモリビューには表示されません。

カーネルでは、VはKと一緒にQKの前にロードされます。VのSMEMレイアウトをそれを使用するPV行列乗算と一緒に議論できるように、ここでVのコピーを説明します。

［画像］

注意：行列乗算の図とコードスニペットは、行列乗算中心の次元名を使用します。実際のカーネルソースは、固定のアテンションタイル名を使用します。

以下のPV行列乗算の図では、次元を従来の行列乗算の次元名で命名します：P[M,K] @ V[K,N] = O[M,N]

M、N、Kという軸の名前を付ける理由を少し立ち止まって説明させてください。各行列乗算で独立してM/N/Kという名前を使用するのは、各行列乗算が単独で見た場合、その入力形状を読みやすくするためです。したがって、これらのM/N/Kは行列乗算中心の命名選択であり、グローバルテンソル（Q、K、P、またはV）の固有のプロパティではありません。

そのため、同じソース次元がQKとPVで異なるローカル名を持つことがあります。したがって、K次元は現在の行列乗算の削減次元を意味します。QKの場合は特徴量、PVの場合はK/Vタイルの行です。それでは、私たちのより大きなPV行列乗算の議論に戻りましょう。

Q、K、Vは同じグローバルメモリレイアウトで始まりますが、カーネルでは、各MMAがどのように消費するかによって、それらをSMEMに異なるように配置します。QタイルにはKメジャーレイアウトを使用し、PタイルにはKメジャーレイアウトを使用するのに、Vタイルには異なるレイアウト（以下で説明）を使用するのはなぜですか？

QKはKを転置された状態で必要とし、PVでは内側の次元が既に一致するため、Vは転置を必要としません。より具体的には：

QK: Q[M,K] @ K[N,K].T
PV: P[M,K] @ V[K,N]

QKの場合、Kは[N,K]として格納され、MMAはBオペランドを論理[K,N]として消費します。そのため、KにはKメジャーを使用します。これにより、MMAは、個別の転置Kタイルを最初に構築することなく、格納されたKタイルを論理的な転置オペランドとして解釈できます。

PVの場合、Vは既にMMA Bオペランドに必要な[K,N]の向きになっています。そのため、VにはMNメジャーを使用します。これにより、MMAはVタイルを転置することなく、直接Vを消費できます。Pは[M,K] Aオペランドなので、Kメジャーのままです。

言い換えれば、これらの図で使用される格納タイルビューでは、私のショートカットは次のとおりです。

• Kメジャー：削減Kは右側にあります。Q[M,K]、K[N,K]、またはP[M,K]のように。
• MNメジャーB：削減Kは左側にあります。V[K,N]のように。

これらはKメジャーとMNメジャーレイアウトの正式な定義ではなく、これらのオペランドビューを読むための単なるショートカットです。

［画像］

これらは同じPV MMAの2つのオペランドです。tcgen05.mmaはAとBのメジャーモードを独立して選択できるため、PはKメジャー、VはMNメジャーにすることができます。Pは左側にあり、セクションAのQと同じKメジャーアトム順序を使用します。数字はSMEMアドレスの増加を示しています。両方のレイアウトは(8,8)アトムを使用します。変更されるのは、どの原子がバックツーバックに格納されるかであり、それはTMA呼び出しごとにコピーするデータの量によって決まります。

Qの場合、1回のTMA呼び出しで完全な垂直スライスをコピーし、そのアトムを正しい順序で配置したことを思い出してください。Vは異なるSMEMレイアウトと異なるアトム順序を使用するため、各TMA呼び出しがコピーすべきデータの量を再考する必要があります。

なぜQの場合のように、Vをより大きなスライスでコピーしないのですか？

Qの場合、1つの垂直スライスのアトムはSMEMにバックツーバックに格納されます。したがって、1回のTMA呼び出しでそのスライス全体をコピーできます。

Vの場合、図のアトム順序を見てください。最初の8行ストライプ全体にアトムを配置し、次に次のストライプに移動します。したがって、1つの垂直スライスのアトムは、SMEM内の他のアトムによって分離されます。したがって、1回のTMA呼び出しで完全な垂直スライスをコピーすると、そのアトムはSMEM内の間違った場所に配置されます。

しかし、水平ストライプ全体をコピーしないのはなぜですか？私たちの単純な、スウィズルされていない2D TMAコピーは、その長方形を行ごとに格納します。SMEMに行メジャー（値の完全な行、次に次の完全な行）が得られます。しかし、私たちのレイアウトでは、次のアトムの前に1つの(8,8)アトムのすべての8行が必要です。

そのため、1度に1つのアトムをコピーし、各アトムがSMEM内のどこに着地するかを明示的に選択します。

ここではVコピーのループを少し簡略化して示します。Qコピーとは異なり、各TMA呼び出しは、完全な高さの垂直スライスではなく、1つの(8,8) BF16アトムをコピーします。

V_tmapは、以前使用したQ_tmapと同じ種類のメタデータ（グローバルメモリアドレス、dtype、dims、strides）を含みますが、コピー形状は now (8 rows, 8 columns) です。kv_row0は現在のVタイルの最初のグローバル行です。図の名前を使用すると、このVタイルは(BLOCK_K rows, BLOCK_N columns)です。

int ATOM_BYTES = (8 * 8) * 2;          // 1つの8x8 BF16アトム
int STRIPE_BYTES = (8 * BLOCK_N) * 2;  // 1つの完全な8行ストライプ

// 各スレッドがこれらのTMAコピーすべてを発行します。
// バリア設定と完了待機は省略

// Vを下に移動、8行ずつ
for (int row8 = 0; row8 < BLOCK_K / 8; ++row8) {

    // 各イテレーションで1つの(8,8)アトムをコピー、左から右へ移動
    for (int col8 = 0; col8 < BLOCK_N / 8; ++col8) {

        // グローバルV内のこの(8,8)アトムの開始行と開始列；
        // 要素座標、バイトオフセットではない
        int src_row = kv_row0 + row8 * 8;
        int src_col = col8 * 8;

        // V_smemはSMEM内の開始バイトアドレスです
        int dst_smem = V_smem
                     + row8 * STRIPE_BYTES
                     + col8 * ATOM_BYTES;

        tma::copy_2d_gmem_to_smem(
            dst_smem, &V_tmap, src_col, src_row, mbar_addr);
    }
}

内側のループは、SMEM内に完全なアトムを隣接して配置します。次に、外側のループは次の8行ストライプに移動し、繰り返します。これらのオフセットは、上記のMNメジャー物理レイアウトを構築します。

私たちの(128,128) Vタイルでは、ループは(16コピー下、16コピー横) = 256コピーを発行します。個々の要素をループするのではなく、各呼び出しはアトムの64値すべてをコピーします。

---

これで、PVオペランドの両方がSMEMに準備される方法がわかりました。Pはソフトマックスによって（前のセクションB）、Vは別途SMEMにコピーする方法を説明しました。MMAがそれらをどのように消費するかを見てみましょう。

［画像］

QKの場合と同じ基本的なアイデアです（以前のセクションAで説明しました）。私たちの削減次元（この図ではBLOCK_Kとラベル付け）には128個のBF16要素がありますが、1つのtcgen05.mma発行はMMA_K=16 BF16要素しか消費しません。したがって、PVは128 / 16 = 8 MMAマイクロステップに分割されます。簡単にするために、ここではそれらのマイクロステップのうち2つだけを示します。

このVビューでは、削減Kは垂直です。したがって、1つのマイクロステップはVの隣接する8行スライス2つを消費します。

BLOCK_Kに沿って、各マイクロステップは16 BF16要素を消費します。各BF16要素は2バイトなので、それは図に示されている32バイトです。

---

PVループはQKと同じ構造に従います。PはKメジャーディスクリプタを使用しますが、異なるのはVがMNメジャーディスクリプタを使用することです。命令ディスクリプタi_desc_pvは、MMA形状、データ型、およびそれらのオペランドモードを指定します。

各PVマイクロステップは、Pの16列とVの対応する16行を消費し、同じOタイルに寄与します。

// Pは以前Kが使用していたSMEMストレージを使用します
int P_smem = K_smem;

uint64_t p_desc = noswizzle::desc_kmajor(P_smem, BLOCK_M);
uint64_t v_desc = noswizzle::desc_mnmajor(V_smem, BLOCK_N);

// MMA_K = 16 BF16値; MMA_K_BYTES = 32。
// ディスクリプタアドレスは16バイト単位でエンコードされています
uint64_t p_desc_step = (BLOCK_M * MMA_K_BYTES) / 16;
uint64_t v_desc_step = (BLOCK_N * MMA_K_BYTES) / 16;

// 単一スレッドがすべてのMMAマイクロステップを発行します
for (int k = 0; k < BLOCK_K / MMA_K; ++k) {

    // Oは最初のPVマイクロステップでのみ初期化され、
    // それ以降のすべてのマイクロステップは既存のOに累積します。
    // (kv_tileは現在のK/Vタイルのゼロベースインデックスです)。
    int enable_input_d = (kv_tile == 0 && k == 0) ? 0 : 1;

    tcgen05::mma_f16(
        taddr_o, p_desc, v_desc, i_desc_pv, enable_input_d);

    p_desc += p_desc_step;  // Pの次の16列
    v_desc += v_desc_step;  // Vの次の16行
}

// すべての128/16=8 PVマイクロステップの完了を追跡します
tcgen05::commit_arrive(mbar_addr);

QKとの1つの違いは初期化条件です。QKはK/Vイテレーションごとに新しいSタイルを開始します。しかし、Oは前のK/Vタイルの寄与を保持する必要があります（既にセクションBでリ スケールされています）。したがって、最初のK/Vイテレーションの最初のマイクロステップのみがOを初期化し、それ以降のすべてのマイクロステップはそれに累積します。最初のK/Vイテレーションでは、カーネルはOリ スケールもスキップします（リ スケールする前のOはありません）。

ループは両方のディスクリプタを次のマイクロタイルのペアに進めます。では、Vディスクリプタにズームインしましょう。そのベースは何を指し、MMAはそのマイクロタイル内の他のアトムをどのように見つけるのでしょうか？

---

［画像］

各PVマイクロステップについて、Vオペランドディスクリプタのベースは、その16行チャンク内の最初の(8,8) Vアトムを指します。

マイクロステップ0の場合、Vディスクリプタのベースは緑色でハイライトされたアトムです。次に、マイクロタイルをループすると、ディスクリプタは次の16行チャンクに進みます。マイクロステップ1の場合、Vディスクリプタのベースは紫色でハイライトされたアトムになります。

次の図は、LBOとSBOが同じマイクロタイル内の他のアトムをそのディスクリプタベースからどのようにカバーするかを示しています。

［画像］

各PVマイクロステップについて、ディスクリプタベースは最初の8行Vスライスの最初のアトムを指します。

LBOは右側の次アトムに到達し、SBOはNに沿って幅8列のアトムを横断します。

これらはVコピーループで使用したものと同じオフセットです。私たちの128列Vタイルについて：

SBO = ATOM_BYTES   = 8 * 8 * 2       = 128 バイト
LBO = STRIPE_BYTES = 8 * BLOCK_N * 2 = 2048 バイト

したがって、コピーループとMMAディスクリプタはアトムがどこにあるかについて一致しています。TMAはそれらをこれらのオフセットに配置し、MMAはそれらを見つけるために同じオフセットを使用します。

その開始アドレス、2つのオフセット、およびMMA形状を使用して、1つのtcgen05.mma発行は完全な16行Vマイクロタイルをカバーします。

次に、ループはPディスクリプタを次の16列に、Vディスクリプタを対応する次の16行に進めます。Vの場合、次のマイクロタイルは16行下から始まります - 2つの8行ストライプ下です。したがって、ベースアドレスを2 * STRIPE_BYTES = 4096 バイト進めます。ディスクリプタアドレスは16バイト単位でエンコードされているため、16で割るとコードで使用されるエンコードされたv_desc_stepが得られます。

---

ズームバックアウト：1つのK/Vイテレーション

Chapter 1のセクションA、B、およびCでは、1つのK/Vループイテレーション内の個々のステップにズームインしました。今度は、完全なKVループイテレーションのレベルにズームバックします。

上記の図が省略している1つのことは同期です。

TMAロードとtcgen05.mmaは非同期です。しかし、カーネル1では、すべての主要なハンドオフで即座に待機します。現在のK/Vタイルをロードして待機し、QKを実行して待機し、行処理を行い、次にPVを実行して待機します。PVが完了した後のみ、次のK/Vタイルのロードを開始します。

したがって、このベースラインは依然として1つのセマンティックステージを次々と実行します。1つのK/Vタイルのみがアクティブであり、K/Vロード、QK、行処理、およびPVの間に有用なオーバーラップはありません。後続のカーネルはこれらのステージのオーバーラップを開始します。SMEMのバリアは、非同期TMAおよびMMA作業が完了したときに私たちを通知します。__syncthreads()は、すべてのCTAスレッドがそのポイントに到達するまで待機するだけです。非同期TMAまたはMMA作業の完了を待機しません。

エピローグ: KVループ全体を終了した後、TMEMからOを読み取ります。これは、4つのワープ、8列ずつです。各スレッドは、その行の値を除算し、BF16に変換し、グローバルメモリに書き込みます。以前のOリ スケール走査とは異なり、ここでは最終的な出力をグローバルメモリに書き込みます（TMEMに戻すのではなく）。

カーネルチェックポイント

このブログは図だけではありません。コードも同様に重要です。

1_baseline.cuを参照してください。

それは私たちが研究したのと同じ正確なシーケンスに従います。

• 各CTAは単一のQタイルを保持します。すべてのK、Vタイルをループします。
• S=Q@K.T (マイクロタイルでのtcgen05、Kメジャーレイアウト)
• S->P (TMEMにアクセス)
• O+=P@V (マイクロタイルでのtcgen05、MNメジャーレイアウト)

今、コードに慣れてください。

これまでのパフォーマンス

このテーブルは、各チャプターの後に1行ずつ増えていきます。Stock FA4 (cute DSL) は100%です。各値は、6回の同一実行率の中央値です。各ベンチマーク呼び出しは、このカーネルとStock FA4を同じB200で測定しました。

⭐ リポジトリをスターすると、後続のチャプターを進む際にすべてのカーネルを手元に置くことができます。

さらに多くのMLシステムの記事、ビデオ、コードが登場します：LinkedIn · X · YouTube

Chapter 2 – PをSMEMからTMEMへ移動

カーネル1では、2番目のソフトマックス走査はTMEMからSを読み取り、それを正規化されていない確率Pに変換し、Pを以前Kが保持していたSMEMバッファに書き込みます（QKの後、Kバッファは不要になるため、Pに安全に再利用しました）。

カーネル2はPの場所を変更します。PのフラグメントをSMEMに書き込む代わりに、TMEMに書き戻します。それがカーネル2で導入された主な変更点です。

この変更はSMEMの使用量を削減しませんが、後続の最適化の準備をします。PをKのデッドSMEMスロットに格納することで、カーネル1はそのライフタイムをカップリングしました。QKの後、Kはもはや必要ありませんが、PVがそれを消費するまでそのスロットはPによって使用され続けます。これは最適ではありません。Pに別のSMEMバッファを割り当てることもできますが、後続の最適化にはほとんどすべてのSMEM予算が必要になります。代わりに、カーネルはPをTMEMに移動させ、ライフタイムを分離します。これにより、SMEMをさらに割り当てることなく、後続のチャプターでこれらの分離されたPとKのライフタイムを使用して、さらなる最適化が可能になります。

割り当てられたSMEMは減少しません。カーネルは依然としてQ、K、Vのストレージを必要とします。

［画像］

最適化として、Sが占めていた割り当てに直接Pを書き込みます。これにより、Pに追加のTMEM割り当ては不要になります（2番目のTMEM走査が既に消費したSの要素を安全に上書きできるため）。

したがって、カーネル2では、同じTMEMスロットが時間とともに意味を変えます。

• QKの後：スロットはFP32 Sを保持します。
• ソフトマックス走査2：消費されたSチャンクは、パックされたBF16 Pで上書きされます。
• PV：スロットはPとして消費されます。

これによりSとPのライフタイムのカップリングが発生しますが、後で見るように、PとKのライフタイムのカップリングほど問題ではありません。7

TMEMは32ビットメモリセルとして編成されていますが、私たちのP値はbf16であるため、私たちのP値の2つは1つのTMEMセルにパックできます。同様に、SはFP32値を含んでいるため、スコア1つ（S行列の値1つ）は1つのb32 TMEMセルを占有します。したがって、PはSのTMEM列の半分を使用します。

図は、論理的な16スコアステップを1つ示しており、異なる色の矢印で示されています。コードでは、そのステップは2つのtcgen05.ld.x8ロード（オレンジ色の矢印で示されています）、それに続く1つのtcgen05.st.x8ストア（緑色の矢印で示されています）を使用して実装されます。カーネル1は幅8チャンクでPを処理します。カーネル2はこれらのチャンクのうち2つを1つのイテレーションにグループ化し、行ごとに、16個のFP32 S値をロードし、16個のBF16 P値を生成します。幅16のグループ化は根本的なものではありません。8

カーネル1から覚えているように、Pを生成する際にTMEMをチャンク単位で反復します（前のチャプターの関連図を参照）。これはチャプター1と非常に似ています。チャプター2で新しいのは、結果のPチャンクをTMEMに格納することです（16個のP値を8個のb32 TMEMセルにパックして）。したがって、Sポインタは16列進みますが、パックされたP書き込みポインタは8列しか進みません。したがって、各P書き込みは読み取りポインタの後ろに残り、既にこの2回目のTMEM走査で消費されたスコアのみを上書きします。したがって、Pを生成してTMEMに（チャンクで）書き込む際に、未読のS値を上書きするリスクはありません。図にも示されているように。9

私たちのSは128個のTMEM列を占め、言及したように16列チャンクで読み取るため、全TMEM S列をカバーするには128/16=8回のイテレーションが必要です。8回のステップすべてが完了した後、下位64個のb32列には完全な128列のBF16 Pタイルが含まれます。したがって、PはSバッファ全体を上書きするのではなく、半分だけを上書きします。残りの半分には古いS値が含まれており、二度と読み取られることはありません。次のQK（KV forループの次のイテレーションで）は、S/Pスロット全体をS(i+1)で完全に上書きします。

［画像］

PV行列乗算は変更されません：P[BLOCK_M, BLOCK_K] @ V[BLOCK_K, BLOCK_N] -> O[BLOCK_M, BLOCK_N]

変更されるのはPのソース（Aオペランド）だけです。カーネル1では、PとVの両方がSMEMから消費されました。カーネル2では、PはTMEMから、VはSMEMから消費されます。

カーネル1では、PはKメジャーSMEMレイアウトとSMEMディスクリプタを必要としました。カーネル2では、PはTMEMアドレスによって直接識別され、PのSMEMディスクリプタは使用しなくなりました。

KメジャーとMNメジャーはSMEM専用レイアウトであり、TMEMオペランドには使用されません。VはMNメジャーSMEM Bオペランドのままです。

---

これで、Pがどのように生成されTMEMに書き込まれるか（チャプター1のようにSMEMに書き込まれるのではなく）を議論しました。では、P@V行列乗算を再度見てみましょう。ただし、オペランドの1つ（P）はTMEMから来ます。

［画像］

図のすべてのMMA関連次元は従来のローカル行列乗算表記を使用していることを覚えておいてください：A[M,K] @ B[K,N] = C[M,N]

したがって、PVはBLOCK_K = 128で削減されます。チャプター1と同様に、各MMA発行はPとVの対応するMMA_K=16マイクロタイルを消費します：P[:, 16k : 16(k+1)] @ V[16k : 16(k+1), :]。8つのMMA_K=16発行が完全なBLOCK_K=128削減をカバーします。煩雑さを避けるため、図は最初の2つのみを示しています。

［画像］

ループは、PをTMEMアドレス経由で選択する代わりにSMEMディスクリプタを進めるのと同じ論理MMA_K=16マイクロタイル順序に従いますが、Pの選択はTMEMアドレス経由で行います。taddr_p + k * 8は、次の128x16 Pマイクロタイルのベースを指します。各行内では、その16 BF16値は8個のb32 TMEM列を占めます。Vディスクリプタは対応する16行に進み、各発行は同じOタイルに累積します。

Pはtcgen05.stで生成されるため、行スレッドはこれらのTMEMストアを待機し、PVが開始する前にフェンスする必要があります。これは発行メカニズムへの単なるマイナーな変更です。

したがって、カーネル2はソフトマックス数学、PV削減順序、およびハイレベルスケジュールを変更しません。その唯一のセマンティックな変更は次のとおりです。QKはTMEMでSを生成し、ソフトマックスはそのストレージをPに変換し、PVはTMEMから直接Pを消費します。

カーネルチェックポイント

2_p_to_tmem.cuを参照し、カーネル1と比較してください。

それは私たちが研究したのと同じ正確なシーケンスに従います。

• ソフトマックスは、消費されたSをパックされたBF16 Pで上書きします。
• PVはTMEMからPを、MNメジャーSMEMからVを読み取ります。
• PはKのSMEMスロットを再利用しません。
• 数学とスケジュールは変更されていません。

今、コードに慣れてください。

これまでのパフォーマンス

これはパフォーマンスの向上にはつながりませんでしたが、これらの変更によりKとPのライフタイムが分離され、後続の主要な最適化が可能になります。

私はこれをカーネル1に折り込むのではなく、別個のステップとして維持しました。なぜなら、PをTMEMに書き込むことは、追加のアドレス指定、パッキング、および同期の詳細を導入するからです。カーネル1は既に十分な新しい概念を導入しているため、不必要に複雑にしたくありませんでした。そして、初期ベースラインをできるだけシンプルに保つようにしました。

Chapter 3 – Swizzling

これまでのカーネルでは、各オペランドのMMAでの役割に基づいて、Q/KをKメジャー、VをMNメジャーとして選択しました（チャプター1でカバーしました。必要であれば復習してください）。カーネル3は同じスケジュールとメジャーモードを維持し、スウィズルモードのみを変更します（各タイルを独立して）。Q/K/VはノースウィズルからSMEMのSW128に移動し、SMEMバンクの競合を減らします。PはTMEMに留まるため、スウィズリングは適用されません。

スウィズリングの動機を説明するリソースはありますが、ここでは繰り返しません。私たちの目標はB200関連の概念を理解することなので、スウィズリングは各アトム内のSMEMアドレスを置換して、テンソルコアのアクセスがSMEMバンク全体でより良くマッピングされ、バンクの競合を回避することを意味すると言えば十分です。ここではB200固有の詳細をカバーしますが、一般的なスウィズリングの説明はしません。

メジャーモードとスウィズリングは別個のものです（チャプター1で議論しました。必要であれば復習してください）。メジャーモードは、16B要素内にパックされる行列乗算次元を示します。Kメジャーの場合はK、MNメジャーの場合はM/Nです。スウィズリングは、そのレイアウトの固定サイズアトム内のアドレス置換を追加します。したがって、まず行列乗算でのオペランドの役割に基づいてメジャーモードを選択します。スウィズリング（SW128）を使用することにした場合、必要なアトム形状は以下のtcgen05レイアウト表から導き出されます。

［画像］

この表を、入力テンソルのdtypeであるBF16要素に変換してみましょう。

［画像］

スウィズリングを追加したとき、カバーする価値のある2つのポイントがあります。（i）アトムが幅広くなり、（ii）内部でスウィズルされるようになりました。これら2つのポイントは順に以下で扱われます。

(i) アトムが幅広くなる

NVIDIAは、BF16値ではなく、128ビット要素（16バイト）の観点からtcgen05共有メモリレイアウトテーブルを説明しています。1つのBF16値は16ビットなので、1つの128ビット要素には128/16 = 8 BF16値が含まれます。

ドキュメントによると、アトムの先頭次元のみを拡張する必要があります。表では、アトム次元は(M/N, K)の順序で並んでいますが、これはすぐには明らかではありませんでした10。

したがって、Kメジャーの場合、Kは先頭次元であり、2番目の軸です。

(8,8) -> (8, 8*8) = (8,64) BF16

Q/Kの場合、ソースタイルを[M/N,K]座標で既に描画しているため、Kメジャーアトムは(8,64)のように見えます。

MNメジャーの場合、M/Nは先頭次元であり、1番目の軸です。

(8,8) -> (8*8, 8) = (64,8) BF16

Vの場合、MNメジャーアトムはNVIDIAの[N,K]座標では(8,8)です。ソースのV[K,N]論理ビューでは、それらの軸は入れ替わっており、(8,64)になります。そのため、それらは同じレイアウトではありませんが、私のソース座標図では両方のアトムタイプが(8,64)のように見えます。

今のところ、まずKメジャーレイアウト（Q@Kで使用）にズームインします。後でMNメジャーレイアウト（P@VでVに使用）に戻ります。

(ii) 内部でスウィズルされる

私の理解では、スウィズリングが追加された場合、16B要素内の要素は置換されず、16B要素自体が置換されます。

正確な行ごとの置換は気にする必要はありません。手動で計算しないからです。代わりに、Blackwellでは、a) TMAはスウィズル-128Bレイアウトで書き込み、b) tcgen05はスウィズル-128Bレイアウトを期待して読み取ります。

SMEMレイアウトと後続のMMAはレイアウトについて合意する必要があります。チャプター1から知っているように。具体的には、TMAテンソルマップでCU_TENSOR_MAP_SWIZZLE_128Bを設定するため、TMAはQ/K/Vをスウィズルして書き込みます。そして、対応するSW128モードをtcgen05オペランドディスクリプタに設定するため、tcgen05はスウィズルレイアウトを消費する方法を認識します。

［画像］

ここでKメジャーレイアウト（Q@Kの両オペランドに使用）にズームインしましょう。

SMEMにコピーする際、上記で説明したように、(8,64)アトムに従ってデータを配置したいと考えています。1つのKメジャーSW128アトムは幅64 BF16要素（つまり128B）です。しかし、私たちの完全なQ/Kタイルは幅128要素（つまり256B）です。

使用する単純な2D SW128 TMAマップでは、タイル全体を一度にコピーして、TMAに2つの別個のスウィズルされた128B領域として扱うように依頼することはできません。CUDAは、1つのSW128 TMAボックスの内部幅が最大128Bであることを要求します。私たちのQ/Kテンソルマップでは、その内部幅はHEAD_DIM（図の水平幅）に対応します。したがって、HEAD_DIMを2つの64幅のスライスに分割します。

しかし、それは(8,64)アトムごとに1回のTMA発行が必要であることを意味するわけではありません。128B制限は内部幅にのみ適用されます。幅が64になったら、同じTMAボックスはタイルのすべての行に拡張できます。

• Qコピー：(BLOCK_M, 64)
• Kコピー：(BLOCK_N, 64)

このカーネルでは、両方のタイルの高さは128です。したがって、1回のTMA発行は(128,64)スライスを1つコピーします。これには、垂直にスタックされた16個の(8,64)アトムが含まれます。

---

上記では、TMAにデータをKメジャー・スウィズルSMEMレイアウトにコピーする際にスウィズリングを使用するように指示する方法をカバーしました。しかし、2番目の部分があります。それは、MMAにKメジャー・スウィズルレイアウトを消費する方法を伝えることです。

［画像］

ここで、これらの64幅スライス1つの内部にズームインしましょう。

1つのKメジャーSW128アトムは(8,64) BF16（1024B）です。幅64のQスライス1つは、BLOCK_Mに沿ってスタックされたこれらのアトムで構成されます。

マイナーな注意点ですが、ここでは読みやすくするためにBLOCK_M=32のみを描いています。実際のカーネルでは、各スライスに16個のアトムが含まれます。

［画像］

1回のMMA発行は依然として幅MMA_K=32B（16 BF16要素）を操作します。しかし、私たちのBLOCK_KはMMA_Kよりも大きいため（ノースウィズルケースでチャプター1で説明したのと同様）、MMA_K=16要素ステップでBLOCK_Kを歩行するforループが必要になります。

ここでの違い（以前に説明したノンスウィズルケースと比較して）は、論理的には、アトム自体を分割することになるということです。各アトムの幅は64ですが、MMA_Kはわずか16です。128Bアトム幅 / 32B mma幅 = 4マイクロタイル。

したがって、ノンスウィズルケースではアトムの2列を取ってマイクロタイルと呼んでいましたが、ここではマイクロタイルはアトムの分割（結合ではなく）の結果です。

［画像］

KメジャーSW128スライス内では、SBOフィールドはハードウェアに次の8行アトム内の対応するフラグメントに到達する方法を伝えると理解しています。

LBOはここでは使用されていません。なぜなら、私の理解では次のとおりです。

1. SBOは既に同じマイクロタイル内の他のアトムチャンクに到達する方法を伝えており、別途、
2. ディスクリプタベースを進めること（マイクロタイルのステップを囲むソフトウェアループ）は、次のマイクロタイルに移動する方法を伝えています。したがって、これら2つを組み合わせることで、データの両方の軸がカバーされます。

したがって、この場合、別の独立した方向を指定してデータをステップ処理するために、追加のオフセット（LBOのような）は必要ありません。

---

要約すると：

メジャーモードが選択されると、SW128は1つのアトムの形状と、その16B要素が内部でどのように置換されるかを教えてくれます。しかし、SW128だけでは、完全なタイルのすべてのアトムの順序を教えてくれません。

TMAコピーの宛先が、各アトムがSMEM内のどこに着地するかを決定します。MMA側では、ディスクリプタベース（ソフトウェアループが各マイクロタイルステップで進めるもの）、SBO、および該当する場合はLBOが、同じ配置を記述する必要があります。そうしないと、tcgen05は間違ったデータを読み取ります。

---

これまでは、スウィズリングが適用されたKメジャーレイアウトをカバーしてきました。次に、PV行列乗算で使用される、スウィズリングが適用されたMNメジャーレイアウトについて議論しましょう。

［画像］

［画像］

各PV MMAマイクロタイルは、Vの16×128スライスを読み取ります。（削減Kに沿って16行、出力Nに沿って128列）。各アトムは8×64なので、このスライスには4つの完全なアトムが含まれます。

このV[K,N]ビューでは、LBOは右側の次アトムを見つけ、SBOは下の次の8行ストライプを見つけます。

図は、読みやすくするために8つのPVマイクロタイルのうち2つを示しています。

カーネルチェックポイント

3_swizzle.cuを参照し、カーネル2と比較してください。

それは私たちが議論したのと同じ最適化に従います。

今、コードに慣れてください。

これまでのパフォーマンス

Chapter 4 – Warp Specialization

［画像］

これまでのところ、私たちのカーネルは4つのワープを使用しており、それぞれが独立してスケジュール可能です。しかし、4つのワープすべて（128スレッド）はカーネル全体で複数回同期されるため、同期は実質的にCTA全体で行われます。以前のカーネルは既にmbarriersを使用しており、これはより細かく使用できますが、それでもすべてのワープをそれらに待機させ、いくつかの__syncthreads()呼び出しは依然としてCTA全体を同期します（上記図参照）。

ワーカー（ワープ）を独立させたいと考えています。これにより、後でそれらを並列化し、それらの間で作業をオーバーラップさせることができます。しかし、まずワープを独立させる必要があります（まだ独立した作業を与えていませんが、独立して実行する能力を与えています）。

例えば、合計128スレッド（4ワープ）があり、現在、各ステージでそれらすべてを同期しています。load K/V(i) -> QK(i) -> softmax(i) -> PV(i) -> load K/V(i+1)。しかし、カーネルを再構成して、すべてのスレッドが各矢印で同期するのではなく、特定のトランザクションに必要なサブセットのみを同期できるようにすることができます（例えば、4ワープのうち2つだけが同期でき、残りの2つのワープをブロックしない）。

例えば、ロードプロデューサーとMMAコンシューマーがバリアを通じてハンドシェイクします。つまり、1つのプロデューサーワープがデータ準備完了の信号を送るためにバリアに到着し20、別のコンシューマーワープがバリアで待機します。他の2つのワープはハンドシェイクに参加する必要がないため、将来、それらは自身の作業を続行でき、グローバルハンドシェイクでブロックされません。これにより、他のワープ（トランザクションに必要ないもの）はブロック/同期されず、不必要に作業を一時停止しないようになります。これは、カーネル4が行う変換です。以前のカーネルのように、すべての128スレッドを各ステージで同期するのではなく、各ワープにその作業を与え、トランザクションに実際に必要なワープのサブセットのみを同期します。

［画像］

特に、カーネル4はワープに独立した役割を与えますが、まだオーバーラップさせるのに十分な独立した作業を与えていません。つまり、与えられたハンドシェイクに参加しないワープは、まだ自身の独立した作業をオーバーラップさせていません。これは私の後続のチャプターのトピックです。これは、すべての最適化を一度に投入するのではなく、教示のために行ったため、フォローしやすくなっています。

この変更自体はパフォーマンスの向上を期待できません（まだより並列な作業を作成していません）。しかし、それは後続の最適化を可能にします。基本的に、カーネル4はカーネルを個々のワーカーに概念的に分割します（一部は同期しますが、他は独立して進行する能力を得ます）。そして後続のカーネルは、これらのワーカーに独立した作業を与え、その能力を実際の作業オーバーラップに変換します。これらの独立したワーカーを作成するために、カーネル4は以前のカーネルよりも多くのワープを必要とします。

まず、TMEM全体のS->Pの並列処理には4つのワープのみが必要になります（128x128 Sがあり、単一のワープは一度に32行しかロードできないため、128行すべてを並列にロードおよび処理したい。したがって、4つのワープが必要になります）。これらは「行ワープ」となり、ソフトマックスとOリ スケールを担当します。11

次に、ロードワープを独立したワーカーにしたい場合、そのためのワープが1つ必要です（SMのスケジューリング粒度はワープ単位であり、個々のスレッド単位ではないため、ロードワーカーはTMAを起動するために単一のワープしか必要としませんが、ワープ全体を割り当てます）。したがって、この時点で5つのワープがあります。

次に、MMAを別のワーカーにしたい場合、それにも別のワープが必要です（上記のTMAの場合と同じロジック）。結果として合計6つのワープになります。

したがって、カーネル4は合計6つのワープを使用します。4つの行ワープ、1つのロードワープ、1つのMMAワープ。

これまでのパフォーマンス

Part II — Coarse Scheduling

Chapter 5 – Two Q Tiles

［画像］

カーネル4は各CTAに1つのQタイルを割り当てます。したがって、2つの隣接するCTA（同じバッチ/ヘッドに対して連続するQタイルが割り当てられたCTA）は、同じK/Vタイルのシーケンスを独立してループし、各CTAは同じK/Vタイルをロードします。これらは各CTAがロードする同じKおよびVオペランドですが、2回ロードすることを回避できます（CTAあたり1回ずつ）。代わりに再利用できます（単一CTA内の各kvループイテレーションで）。私たちの次の最適化は、1つのCTAに2つのQタイルを処理させ、ロードされた各K/Vタイルを両方に再利用することです。これにより、各K/Vロードが2つのQタイル分の作業に償却されます。

単純に両方のQタイルを1つのCTAに移動させるだけでは（図の左側）十分ではありません。もし無造作に1つのCTAに2つのQ/Outputタイルを処理させると、直接的なSMEM K,Vの再利用は作成されません（CTAが最初のQタイルのKVループ全体を完了してから、2番目のQタイルのループを繰り返した場合、同じK/Vタイルを2回ロードすることになります）。

カーネル5は代わりにK/Vを1回ループし、ロードされた各タイルを両方のQタイルに使用します。

Q0とQ1を一度ロード

各KVタイルiについて：
    K(i)とV(i)を一度ロード

    Q0K(i)とQ1K(i)にK(i)を使用
    P0V(i)とP1V(i)にV(i)を使用

同じ2つのQタイルについて、これは以前のカーネルと比較して、K/V TMAロードの発行数を半分に削減します。SMEMから同じKVタイルを直接再利用します（L2再利用ではありません。別々のCTA間でデータがL2に残ることに依存しているわけではありません）。

以下の表記では、0/1はQステージを識別し、iはK/Vタイルを識別します。例えば、Q0K(i) = Q0 @ K(i)^T、そしてP0V(i) = P0(i) @ V(i)です。将来のカーネルでもこの命名規則を使用します。

［画像］

1つのCTAは、KVループ全体で2つの連続するQ/Oタイルの作業を所有します（KVループの期間中、2Qタイルを保持し、最終的に2つのOutputタイルを生成します）。

両方のQステージは、SMEM内の同じK/Vタイルを再利用しますが、QK、オンラインソフトマックス、およびPVは各Qタイルに対して別々に実行されます。各Qステージは、独自のS/Pタイル、Oアキュムレータ、行最大値、および行合計を保持します。

［画像］

2番目のQストリームも独自のTMEM状態を必要とするため、レイアウトは上記のようになります。

チャプター2で説明したように、各BF16 Pタイルは、32ビットTMEMセルの各ペアの値をパックした値で構成され、そのFP32 S割り当ての下半分を上書きします。したがって、P0はS0を超える追加ストレージを必要とせず、P1もS1を超える追加ストレージを必要としません。これはチャプター2の変更に依存しています。もしPをデッドSバッファに書き込まず（代わりに別のTMEM領域に格納した場合）、2つのQタイルのためのS/P/Oを単純に格納するのに十分なTMEMがなかったでしょう。12

別の見方をすると、カーネル5はTMEMを空間的に使用します（同じKVループイテレーション内）、しかし時間的には使用しません（KVループイテレーションをまたいで）。つまり、2つのQ/Outputストリームは同時にアクティブであり、2つのライブQタイル間でK/Vを再利用します（同じKVループイテレーションで）。しかし、時間的な計算パイプライン処理ではありません（異なるKVループイテレーションをまたいで）。すべての512個のTMEM列が使用されるため、将来のkvループイテレーションに対する追加のフルステージS(i+1)/P(i+1)スロットを保持する利用可能なTMEMはありません。13

これはまだロードパイプライン処理を実装していません。カーネル5は2Qの再利用を追加しますが、まだ現在のKスロットとVスロットは1つだけです。

［画像］

以前は6つのワープがありました（チャプター4を参照）。今、CTAは10個のワープを持っています。

ワープ0..3:     Q0の行/ソフトマックス作業
ワープ4..7:     Q1の行/ソフトマックス作業
ワープ8:         TMAロード
ワープ9:         QKとPV MMA発行

2つの行ワープグループは独立しています。Q0の行作業は、MMAワープがQ1Kを処理している間に開始できます。

命名に関するマイナーな注意点ですが、私はそれらをソフトマックスワープではなく行ワープと呼んでいます。なぜなら、これらのワープはソフトマックスだけでなく、Oリ スケールも行うからです（チャプター4でも同じ命名が使用されていたため、これは新しいものではありません）。

それでもMMAワープは1つだけなので、Q0K、Q1K、P0V、およびP1Vは、2つの並列テンソルコアストリームではなく、1つの順序付けられたMMAストリームによって発行されます。別のMMAワープを追加しても、同じテンソルコアハードウェアをターゲットにしているため、追加のハードウェアテンソルコアパイプラインは作成されません。

そしてtcgen05.mmaは非同期であり、1つの発行者（選択されたレーン）はMMA作業を送信するのに十分です。代わりに、ボトルネックは一般的にMMAストリームに発行するMMAがない場合に発生します。例えば、オペランドが準備できていない場合です。したがって、問題は別のMMAワープがあるかどうかではなく、スケジュールが既存のテンソルコアパイプラインを忙しく保つのに十分な作業を公開しているかどうかです（オペランドを十分に速く作成し、MMAが作業のために待機しないようにする。そしてMMA発行順序がMMAの発行を不必要にブロックしないようにする）。後続の最適化は、発行者の追加ではなく、オペランドの準備とスケジュールの改善によってこれを対処します。そのため、ここでは別のMMAワープを追加せず、残りのラインエイジではMMAワープを1つだけ保持します。

2Qステージを追加したため、新しいバリアが必要になりました。追加されたバリアの一部は両方のQステージで共有され、他のバリアは各ステージで個別に設定されます。

• q_ready, kv_ready: 両方のQステージで共有されます。q_readyはKVループ開始前に一度使用されます。kv_readyはKVイテレーションごとに一度再利用されます。
• qk_done[q], softmax_done[q], pv_done[q]: 各Qステージの独立したストリーム。

各softmax_done[q]は4つの到着を期待します。1つは各Qステージに割り当てられた行ワープから。これらの行ワープはソフトマックス/P生成とOリ スケールの両方を行います。4つすべてが到着すると、Pタイルが準備完了であり、古いOアキュムレータがPVに使用可能であることを示します。したがって、このバリアはP準備完了とO安全の両方を表します。MMAは、PVを発行する前にこのバリアで待機します。

図で説明したように、K/Vオペランドが2つのQ/O作業項目間で共有されるようになったため、次のK/Vタイルで再利用する前に、両方がこれらのバッファの使用を終了したことを確認する必要があります。そのため、P1Vが完了した後のみK/Vストレージを再利用します。14

カーネルチェックポイント

5_two_q_tiles.cuを参照し、カーネル4と比較してください。

それは私たちが議論したのと同じ最適化に従います。

今、コードに慣れてください。

これまでのパフォーマンス

Chapter 6 – Load Pipeline

［画像］

まだ上記の図を解析しようとしないでください。テキストに従ってください。必要になったときに図に注意を向けます。

これまでのカーネルでは、ループの各イテレーションでK(i),V(i)タイルをロードする際に、両方がロードされるまで停止し、MMAワープは待機します。その後、QKのみが進行できるため、QKの下流の残りの作業（S->P、そして最終的にはPV）も、KとVの両方のオペランドがロードされるまでブロックされます。

そこには2つの明白な非効率性があります。

1. QKはKタイルのみを必要とします（Q@KはVタイルを必要としません） - そのため、Kのみがロードされた後にQKを進行させることで、ロード待機時間を短縮できます（Vを待つ必要はありません）。これは、QKの下流の作業（S->P）も解除し、ソフトマックスをより早く開始できるようにします（アイドル時間が短縮されます）。

2. ロードをパイプライン化することで、個々のK/Vロードを待つアイドル時間を大幅に削減できます。

2番目の変更がこのチャプターの主な最適化です。

私はこれをプロデューサー（ロードワープ）とコンシューマー（MMAワープ）として考えるのが好きです。パイプライン化のアイデアは、プロデューサーがコンシューマーが必要とするものより先にオペランドをロードすることです。これにより、コンシューマーが必要とする前にロードが完了するのに十分な時間があるはずです。コンシューマーがついにそのオペランドに到達すると、それは既にロードされている可能性が高いため、MMAはそれを待機せずに使用できます。そして、将来のKまたはVオペランドが非同期にロードされている間、コンシューマーは既にSMEMにロードされているオペランドで忙しく作業しています。このようにして、ロードと計算をオーバーラップさせます。

最初の変更により、パイプラインがより細かくなります。

KまたはVを個別のオペランドとして扱い、パイプライン化します。K(i)とV(i)を1つのペアとして扱うのではなく、より大きな物理SMEMバッファが必要になります。15 3つの物理パイプラインスロットのそれぞれが、KまたはVタイルを1つ保持します。KV[3]は共有の3スロットパイプラインバッファであり、どのスロットもKまたはVを保持できます。したがって、KとVは論理的に分離されますが、別々のパイプラインには配置されません。これにより、ロードワープは完全なK/Vペアが再利用可能になるのを待つのではなく、先へ進むことができます。

これはチャプター2の最適化にも依存しています。PはもはやKのSMEMストレージを占有しません。したがって、QKの後でKをリサイクルでき、PはPVのためにTMEMに保持されます。

より詳細に：

Q0とQ1はKVループ全体で一度ロードされるため、パイプラインに参加するのは変化するK/Vオペランドのみです。

私たちのロードパイプラインの深さは3です（3つの物理バッファ）。

KVループが開始する前に、K0、V0、およびK1をプリロードし、すべての3つのパイプラインバッファを埋めます。K0とV0は最初のKVループイテレーションに属し、K1は次のイテレーションのために先行ロードされます。これは、最初のコンシューマーMMAが必要とするもの（Q0はループ前に一度ロードされているため、K0のみがオペランドです）よりも先にデータをプリロードし、プロデューサーにコンシューマーに対する先行リードを与えます。

KVタイルのforループでは、各QKまたはPVはプリロードされたKまたはVオペランドを使用します。2番目のQステージがそのオペランドの使用を終了すると、対応するコンシューマーQ1K(i)またはP1V(i)は、古いトークンはもはや必要ないことを信号で伝えます。独自のトークンループで待機しているロードワープは、そのスロットを再利用し、論理的に3つの位置先（今すぐ、上記の図を見てください。コードでの実装方法を示しています）の将来のオペランドのロードを発行します。

3つのトークンパイプラインにより、新しいロードされたトークンが必要になる前に2つのK/Vトークンがあります。各トークンは両方のQステージによって消費されるため、MMAストリームは新しいトークンを使用する前に4つの完全なQK/PV行列乗算を処理します。これにより、特定のロードはバックグラウンドで完了でき、一般的にそのオペランドのロードが完了するのに十分な時間があるはずです。これは効果的にロードと計算をオーバーラップさせます。

［画像］

もちろん、ハードウェアバッファは円形ではありませんが、インデックスがラップアラウンドし、パイプライン深度（3）の位置で分離されたトークンが同じ物理スロットにマッピングされるため、論理的には円形バッファとして扱うことができます。データは物理的に円形に移動しません。

今、上記の両方の図をより詳細にカバーします。

tokenは論理的なKまたはVオペランドを識別します。stageはそのオペランドを3つの物理SMEMスロットの1つにマッピングします。それはそのスロットが準備完了か上書き可能かを示しません。

トークンtとt+3は同じ物理スロットを使用します。したがって、トークンtをロードする前に、ローダーは古いトークンt-3が最後のコンシューマーによって使用されるまで待機する必要があります。

2つのQタイルがあり、両方のQが同じK/Vオペランドを再利用するため、最後のQタイルがそれを使用するのを終了した後のみ、特定のKまたはVトークンをリサイクルできます。

したがって：

• Kトークンの場合、最後のコンシューマーは最後のQのQKです。
• Vトークンの場合、最後のコンシューマーは最後のQのPVです。

具体的な例（以下の図を参照）で計算のロジックを見てみましょう。

［画像］

パイプラインの進行方法の視覚化（以下の図を参照）。

［画像］

カーネルチェックポイント

6_load_pipeline.cuを参照し、カーネル5と比較してください。

それは私たちが議論したのと同じ最適化に従います。

今、コードに慣れてください。

これまでのパフォーマンス

Chapter 7 – Compute Pipeline

これまでのスケジュールを要約しましょう（以前のカーネルで）。これは次の最適化を理解するために重要です。

［画像］

以前のカーネルでは、MMAワープはkvループイテレーションの粒度で動作します。つまり、Q0K(i)、Q1K(i)を発行し、次にP0V(i)、P1V(i)を発行します。そして、次のkvループイテレーションに進みます。

したがって、1番目のQタイルのPVが発行されると、その行ワープ（ソフトマックス+Oリ スケール）はアイドル状態になります（次のkvループイテレーションに進むことができません）。なぜなら、MMAワープは、2番目のQタイルのPVが発行されるまで、次のイテレーションのQ0K(i+1)（つまり、最初のQタイルの次のイテレーションのQK）に到達しないからです。

これは明らかに最適ではありません。なぜなら、1番目のQタイルのQK(i+1)、そしてさらに重要なことにそのソフトマックス(i+1)は、2番目のQタイルのPV(i)の後ろにブロックされているからです。

しかし、チャプター5で「2Qタイル」最適化を導入したときに議論したように、CTAが処理する2つのQ作業ストリームのそれぞれは、ほとんど独立したストリームであり、独自のバッファ（S/P/Oなど）を持っています。それらはK/Vを共有しますが、それはQ0K(i+1)がP1V(i)の後ろにブロックされる理由ではありません。16したがって、Qステージ間のデータ依存関係はなく、最初のQタイルの次のkvループイテレーション（つまり、次のKVタイルのQK）に進むのを妨げるものはありません。つまり、Q0K(i+1)がP1V(i)の後ろに留まることを要求する依存関係はありません。

代わりに、KVタイルのi番目のイテレーションの終盤近くでQ0K(i+1)を発行し、次にi+1イテレーションの開始時に再度発行するという、同じQ0K(i+1)を誤って2回発行することになります。そして、2番目のQステージについても同様の問題が発生します。

したがって、カーネルにはいくつかの変更が必要です。以下で説明します。

［画像］

したがって、各Qステージの最初のQKをループから剥がします。チャプター6で行ったロードパイプラインのプリフィルに似ていますが、今度は計算パイプラインをプリフィルしているようなものです。

これで、ループの定常状態では、同じ次のタイルQKを2回発行しなくなります。

しかし、ループ前のプリフィルの結果として、定常状態ループの後にも論理的に対応するドレインが必要になります（2つのQステージの作業ストリームを使い切るため。KVタイルのループの終わりに、後続のイテレーションがないため、ドレインを追加しない限り、最後の2つのPVタイルは実行されません）。

［画像］

そして、これがこのチャプターの最終設計に至った方法です。

スケジュールは、QKとPVを反転させたように見えますが、そうではありません。ここではqはQステージインデックス（0または1）です。定常状態では、各ループイテレーションは、現在のKVタイルをPqV(i)で終了し、次に次のKVタイルをQqK(i+1)で開始します。任意の個々のKVタイルについて、順序はQqK(i) -> softmax-q(i) -> PqV(i)のままです。したがって、各QタイルについてQKとPVを反転させたように見えるだけですが、ループ境界を移動させただけです。各KVタイル内では、何も反転されていません。

カーネルチェックポイント

7_compute_pipeline.cuを参照し、カーネル6と比較してください。

それは私たちが議論したのと同じ最適化に従います。

今、コードに慣れてください。

これまでのパフォーマンス

Chapter 8 – Hardware Approximations

このチャプターには、私が描くのに意味のあるものは何もありません。そのため、少し擬似コードで進めます。チャプター10で再び図に進みます。

これまでの最適化のほとんどはスケジュールを変更しましたが、行側のソフトマックス作業は依然として主なボトルネックです。このexp（S->Pロジックの一部）は、Qステージごと、KVタイルごと、行ごと、8スコアチャンクごとに実行され、PVが待機しているのはこれです。ここでの追加作業はテンソルコアを制限します。私たちの次のいくつかのカーネル（カーネル8を含む）は、そのホットパスをより安価にし始めます。つまり、ソフトマックスのホットパスでのexpをより速くしたいのです。

GPUは、近似的な2のべき乗指数関数命令ex2.approx.ftz.f32を提供します。「近似」という名前は、ハードウェアがexp2を近似的に計算することに由来します17。したがって、ホットパスを高速化しようとすると、この命令（ハードウェア近似）を、以前のカーネルで使用していた__expf組み込み関数ではなく使用できます。しかし、私たちのオンラインソフトマックスは依然としてexp（指数関数）を必要とし、ハードウェア命令はexp2を計算します。通常のexpと同じ値を生成するために、指数を2のべき乗で表現します。基本的に数学の公式exp(x) = exp2(x * log2(e))を使用します。したがって、exp2が期待する単位に（入力にlog2(e)を掛けることによって）入力を変換すると、その結果は既に必要なexp(x)値になります。この同一性自体は正確です。

実際には…コンパイラが__expf組み込み関数を次のようなものに展開することが判明しました。

// 元々__expfに渡された自然指数引数
// float x = s * softmax_scale - rowmax;

// 2のべき乗指数（exp2が期待するもの）に変換
// （上記の数学的同一性を参照）
float u = x * LOG2_E;

bool small = u < -126.0f;

// FTZ exp2命令内で
// アンダーフローする結果を保持
if (small) {
    u *= 0.5f;
}

float y;
asem("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(u));

if (small) {
    y *= y;
}

したがって、__expfはex2.approx.ftz.f32コアに加えて、その周りのいくつかのラッパー命令（後で議論します）に分解されます。私たちの使用例では、そのラッパーにはホットループで回避可能な2つの作業部分があります。

1. log2基底への繰り返し変換
2. アンダーフロー動作の保持（述語での比較）

そのラッパーコードが存在する理由は、コンパイラが私たちのより大きなセマンティクスを推測できないため、保守的にこの変換（exp2が期待するlog2基底への）をホットループ内のすべてのexpfに対してローカルに行い、EX2（ex2.approx.ftz.f32）コアの周りのアンダーフロー処理もホットループ内のすべての__expfに対して発行されるためです。そして、この追加作業はパフォーマンスに目立った影響を与えます。

カーネル8はMUFU指数関数の数は同じままですが、基本的に__expfラッパー全体を削除します。

1. その繰り返し変換（log2基底への）をホットループの外にホイストし、既存のSスケールに折り込み、
2. アンダーフローチェックを削除し、代わりにFTZに依存します。

まず、上記の2つのポイントを順番にカバーします。

---

1. ホットループの外への繰り返し変換（log2基底へ）のホイスト

以前のカーネルでは、ソフトマックスを自然指数関数で直接表現していました：p = __expf(s * softmax_scale - rowmax);。

コンパイラは各__expf呼び出しをローカルに処理します。生成されたコードでは、全体的なオンラインソフトマックスを書き換えて、その指数入力がexp2フレンドリーな単位で維持されるようにしません。そのため、ホットループ内のすべての値に対して、log2基底への変換を繰り返します。

カーネル8はこれを手動で行います。

// 通常の1/sqrt(HEAD_DIM)アテンションを折り込む
// log2(e)とともにスケールする
softmax_scale_log2 = softmax_scale * log2(e);
// 直接ex2.approx.ftzを呼び出す
// （ラッパーを追加する__expfではない）
p = fast_exp2(
        // Sをそのスケール係数でスケールする。これにはlog2が含まれるため、
        // 入力はexp2入力フレンドリーな単位になります。
        // rowmaxは通常通り減算する
        s * softmax_scale_log2 - rowmax
        );
rescale = fast_exp2(old_rowmax - new_rowmax);

Rowmaxもlog2単位である必要があります。チャプター1から思い出してください。私たちのカーネルは、SをPに変換するためにTMEMを2回走査します。最初のS走査は、生のSスコアの最大値を依然として見つけます。これは変更されません。削減後、カーネル8では、そのスカラー値（行あたりの最大値）にsoftmax_scale * log2(e)を掛けます。これにより、行最大値は直接exp2入力単位になります。したがって、上記の減算の両側はexp2が期待する単位になります。

score:  s * softmax_scale * log2(e)
rowmax: すでに同じ単位で格納されている

P、行合計、およびOはlog2空間に格納されません。行最大値とスケールされたS（exp2に渡されるもの）のみがこれらの単位を使用します。

---

2) アンダーフロー処理を削除し、FTZに依存する

第二に、__expfのアンダーフロー処理を削除します。したがって、比較および分岐ロジックを削除します。

この部分は、上記の最初の部分とは異なり、実際の数値近似です。結果がFP32の最小値（サブノーマル値）よりも小さい場合、__expfラッパーはexp2(u / 2)を計算してから結果を2乗することで、それを保持しようとします。

__expfラッパーは常にif/else分岐の述語を計算しますが、ほとんどの場合、S値は十分に大きいので、分岐ロジックは適用されませんが、述語での比較は常にコストがかかります。この比較を実行して、チェックの結果が適用されるかどうかを判断する必要がありました。したがって、この比較命令は、スコア値に関係なく、すべての__expf呼び出しで実行されます。しかし、これらの2つの予測された修正乗算は、exp2の結果が2^-126（最小のFP32値）を下回った場合にのみ重要です。

したがって、カーネル8は、サブノーマルなPとリ スケール値を意図的にゼロにフラッシュします。これは許容できると判断します。これはコンパイラが行うことができなかったことです。なぜなら、これはセマンティクスを保持する変更ではなく、コンパイラはこのような変更を行う際に保守的だからです。

サブノーマル結果の保持をやめると、選択する修正がなくなるため、述語とそれに対応する比較も削除できます。直接のex2.approx.ftz.f32は、その処理をスキップし、値をゼロに変換します（命令名でFTZで示されています）。

SASS比較：

MUFU.EX2:             129 -> 129
静的命令: 3296 -> 2776
述語命令:       368 -> 110
FMUL:                   644 -> 257
FSETP.GEU.AND:          129 -> 0

したがって、カーネル8の変更はMUFUの数を減らしませんでしたが、各MUFUの周りの作業を安価にしました。

これまでのパフォーマンス

Chapter 9 – Software Approximations

このチャプターには、私が描くのに意味のあるものは何もありません。そのため、少し擬似コードで進めます。次のチャプターで図に進みます。

カーネル8は各ハードウェア指数関数の周りの追加ラッパー作業を削除しましたが、すべてのP指数関数は同じMUFU.EX2パイプラインを通過しました。

カーネル9は、これらの指数関数のいくつかをMUFUから移動させ、通常のALU/FMA命令を使用してソフトウェアで近似します。

したがって、指数計算の数（Sの各値を指数化する必要があるという点では、以前のカーネルすべてと同じ）は減らしませんが、ワークロードを2つの異なるハードウェアに分散させ、MUFUパイプラインへの圧力を軽減します。

分割は次のとおりです。

各16個のP値について：
    最初の12個 -> ハードウェアexp2
    最後の4個  -> ソフトウェアexp2

したがって、P指数関数の25%がソフトウェアパスを使用します。この選択は列位置にのみ依存し、個々のP値には依存しません。

正確な分割（ALU用25%、MUFU用75%）は実験的に決定されており、ハードな要件ではありません。ソフトウェアexp2はより多くの命令、一時値、およびレジスタを必要とします。すべてをソフトウェアに移動すると、ボトルネックが他のリソースに移動します。

なぜいくつかのソフトウェア命令が1つのハードウェア命令よりも速いことがあるのでしょうか？ソフトウェアexp2は異なる実行リソースを使用するためです。ハードウェアパスは短いですが、それらの指数関数すべてが比較的限られたMUFUパイプラインを競合します。それらのいくつかを通常のALUパイプラインに移動させることで、GPUはより独立したリソースで作業を行うことができます。

ソフトウェアexp2ウォークスルー

上記の簡単な説明で、ブログの残りをフォローするのに十分です。

近似数学のオプションのウォークスルーについては、付録のソフトウェアexp2の仕組みを参照してください。

ヘルパーは次のようになります。

float software_exp2_scalar(float x) {
    // 多項式係数
    float C3 = 0.077f;
    float C2 = 0.227f;
    float C1 = 0.695f;

    // 入力は基底2単位であり、2**xを計算したい。
    // xが-127未満の場合、2**xは非常に小さくなることを意味します。
    // FTZのように、私たちが気にしないほど小さい。
    x = fmaxf(x, -127.0f);

    int int_part = floor(x);
    float frac_part = x - int_part;

    // 3つの連鎖FMAを使用して2^frac_partを近似します。
    // 最適化として、各結果を次のFMAにフィードし、
    // 別個のfrac_part^2およびfrac_part^3計算を回避します。
    float frac_exp2 = C3;
    frac_exp2 = fma(frac_exp2, frac_part, C2);
    frac_exp2 = fma(frac_exp2, frac_part, C1);
    frac_exp2 = fma(frac_exp2, frac_part, 1.0f);

    // 近似2^frac_partのビットを公開します。
    int frac_exp2_bits = __float_as_int(frac_exp2);

    // FP32は23個の小数ビットを持つため、1 << 23は1つの指数ステップです。
    int exponent_adjustment = int_part * (1 << 23);

    // 近似2^frac_partと2^int_partを結合します。
    int result_bits = frac_exp2_bits + exponent_adjustment;
    return __int_as_float(result_bits);
}

このアルゴリズムは、FA4のex2_emulation_2から適応されています。これは簡略化されたバージョンであり、明確さのために1つの値を処理します。実際のヘルパーは、パックされたf32x2命令を使用して、同じアルゴリズムを2つの値同時に実行します。

ハードウェアパス（チャプター8で追加）と、このチャプターのソフトウェアパスの両方が近似的です。この最適化は、選択されたP値の計算方法のみを変更し、指数計算の数を減らしません。

これまでのパフォーマンス

Chapter 10 – Register Caching

［画像］

以前のカーネルでは、各Qステージについて、SをTMEMで走査し、2回読み取りました（チャプター1で説明）。最初に、行最大値を計算するため（2番目のステップでの数値的安定性のために使用されます）。次に、行最大値を減算し、結果を指数化するため（したがって、まず行最大値を計算してから、このステップで減算する必要があります。そのため、同じSをTMEMで2回走査しました）。

私たちの次の最適化は、レジスタにSの一部を保持することにより、2回目のTMEMの再読み取りを削減します。

［画像］

したがって、レジスタのライフタイムを削減したいので、次のことを行います。

• (a) TMEMの最初の走査で、読み取ったSの最も最近の値の一部をキャッシュします（Sを低から高へ読み取るため、最新の値は高い96列です）。
• (b) P生成パス（以前は2回目のTMEM走査）を変更して、高から低へ走査するようにします（したがって、最近キャッシュされた値が最初に消費されます）。

［画像］

以前のカーネル（チャプター2から開始）では、PをTMEMのS領域の下半分に格納していました。このチャプターでは、高い96列のみがキャッシュされるため（前述）、低い32列は後で再読み取りする必要があります（合計128列）。したがって、PをSの下半分に格納することはできません（Pを下半分Sに格納すると、再読み取り前に低32 Sを上書きしてしまうため）。したがって、チャプター10から、PをSの上半分に格納します。

［画像］

もう1つの詳細：再読み取りされたSの低32列は、Pの1/4のみを生成するために必要です。残りの3/4のPは、レジスタにキャッシュされたS値を使用して生成されます。

したがって、Sに対する2回目の走査は残りますが、TMEMから128スコアのうち32のみを再読み取りし、再読み取りの3/4を削除します。

カーネルチェックポイント

10_cache96_reread32.cuを参照し、カーネル9と比較してください。

今、コードに慣れてください。

これまでのパフォーマンス

Chapter 11 – Skip O rescale

このチャプターには、私が描くのに意味のあるものは何もありません。そのため、少し擬似コードで進めます。次のチャプターで図に進みます。

通常のソフトマックスでは、行全体を見て、最終的な行最大値を一度見つけ、それを減算し、次に指数関数を計算します。

オンラインソフトマックスは、KVタイルを1つずつ処理するため、最終的な行最大値を事前に知りません。タイルごとに発見していくため、後続のタイルにさらに大きな最大値が含まれる可能性があります。したがって、オンラインソフトマックスが（通常の完全なソフトマックスとは異なり）行う必要があるユニークな決定は次のとおりです。

• 新しい最大値にすぐに切り替えますか？
• それとも古い行最大値を使用し続けますか？

以前のカーネルでは、より大きな最大値が見つかるたびに、すぐに使用を開始しました。

new_rowmax = max(old_rowmax, tile_rowmax)

次に、古い状態をすべてその新しい基底で表現しました。

old O      *= rescale
old rowsum *= rescale

行合計の更新は安価ですが、Oのリ スケールは、各行ごとにTMEMの128個の値全体を再度走査する必要があります。カーネル10は、最大値が変わらない場合（rescale == 1）でも、kvループの各イテレーション後にこのOパスを実行しました。したがって、カーネル10は128個のO値をロードし、1で乗算し、それらをストアしました。これは数値的には安全ですが、過度に保守的でした。

カーネル11では、新しいタイル最大値が現在保持している行最大値よりもわずかに大きい場合、古い行最大値を使用し続け、Oと行合計を変更しません。差が十分に大きい場合のみ、新しい行最大値に切り替えて、古いOと行合計を新しい基底にリ スケールします。

リ スケールをスキップしても情報は失われない

これはすぐに明白ではありませんでしたが18、リ スケールをスキップした場合に何が起こるかを小さな例で見て、情報が失われないことを自分で納得しましょう。

例えば、あるタイルをリ スケールすることをスキップしたとしましょう（具体的な閾値については後で議論します）。今のところ、現在のタイルの最大値が古い行最大値より7.78大きいとだけ言っておきましょう。

古い行最大値に基づいてPを計算し続けます。

P = exp2(scaled_S - kept_rowmax)

したがって、より新しい最大値は次を生成します。

exp2(7.78) ~= 219

したがって、Oリ スケールをスキップしても、より新しい最大値は消えず、情報も失われません。古い行最大値を使用してPを計算し続けるため、より大きなスコアは約219倍大きくなる正規化されていないP値に変わります。そのより大きなP値は、差が保持される場所です。

さらに良いことに、代わりにリ スケールを選択した場合と比較して、カーネル全体の最終結果はまったく変更されません。なぜなら、より大きなPはソフトマックスの両方の部分に入るため、より大きなP値はO（分子）と行合計（分母）の両方に入ります。

同じP値が両方の注意部分に入ります。

O      += P @ V
rowsum += sum(P)

したがって、私の以前の219の例では：

O      += 219 * V_new
rowsum += 219

その後、KVループ全体の後：

output = O / rowsum

したがって、より新しいスコアは、Oと行合計の両方で依然として約219倍大きいものとして扱われます。最終的な除算は、Oと行合計を共有する共通のスケールをキャンセルするだけで、その219倍の相対的な違いを消去しません。

では、なぜリベースするのですか？

最終的なO / rowsumが共通のスケールを処理する場合、最初の行最大値を永遠に保持できるようです。数学的には、中間値がオーバーフローしない限り、可能であるように思えます。例えば、古い行最大値が非常に古くなった場合：

gap = 20
largest P = exp2(20) = 1,048,576

最終的な答えは数学的には同等になりますが、P、行合計、およびOは巨大な中間値を使用するようになります。これらがオーバーフローした場合、明らかに最終的な正規化（行合計で割る）はそのオーバーフローを元に戻すことができません。

行ごとの決定がワープ全体TMEM Oパスにどのようにマッピングされるか

リベースの決定は各行ごとに別々に行われますが、TMEM Oのロードとストアはワープコレクティブです。したがって、実際の実際の実装は次のようになります。

// 現在のタイルの生の行最大値を基底2指数単位に変換
tile_rowmax *= softmax_scale_log2;

// kept_rowmaxはすでに同じ単位で格納されている
float candidate_rowmax = max(kept_rowmax, tile_rowmax);
float rowmax_gap = candidate_rowmax - kept_rowmax;
bool row_needs_rebase = rowmax_gap >= REBASE_THRESH_LOG2;

float rescale = 1.0f;

// 1) 必要であれば、この行の状態をリ スケールし、新しい行最大値の使用を開始する
if (row_needs_rebase) {
    // この係数は行合計とOの両方に適用される
    rescale = fast_exp2(kept_rowmax - candidate_rowmax);

    // この行のスカラーアキュムレータを新しい行最大値基底に移動させる
    rowsum *= rescale;

    // このタイルのPと将来のPを生成する際に新しい行最大値を使用する
    kept_rowmax = candidate_rowmax;
}

// 2) Oに同じリ スケール係数を適用する
//
// TMEMロード/ストアはワープコレクティブです。どの行もリ スケールが必要な場合、
// ワープ全体がその32x128 Oスラブを処理する必要があります。
// リベースしない行は、rescaleが1のままでも参加します。
if (__any_sync(FULL_WARP_MASK, row_needs_rebase)) {
    float o8[8];

    // 前のカーネルで使用されたのと同じO走査
    for (int col = 0; col < HEAD_DIM; col += 8) {
        int taddr = taddr_o_stage + trow + col;

        tcgen05::ld_32x32b_x8(taddr, o8);

        for (int i = 0; i < 8; ++i) {
            o8[i] *= rescale;
        }

        tcgen05::st_32x32b_x8(taddr, o8);
    }

    asm("tcgen05.wait::st...");
    asm("tcgen05.fence::before_thread_sync...");
}

// P生成と後続のPVは、このタイルに対して行最大値としてkept_rowmaxを使用して、
// 前のカーネルと同じように続行します。

実際の閾値

カーネル8以降、指数入力を基底2単位で表現しています。カーネル11では、8基底2指数単位の閾値を使用します。これは公式FA4と一致します。私の理解では、8には深い数学的な崖はありません。12や15のようなより大きな閾値は、BF16-P/FP32-アキュムレータパスで数値的に妥当です。8は、既にほとんどの有用なOリ スケール機会をスキップしているように見える保守的なポイントであり、オーバーフロー境界ではありません。

exp2(8) = 256

したがって、新しい行最大値に切り替える前に、正規化されていないP値の最大値が約256未満のままである間、古い行最大値を使用し続けます。そして、「正規化されていないP値」と言ったのは、新しい値の指数化の前に古い行最大値を減算したためです。それがそれを超えると、新しい行最大値に切り替えて、古い状態をリベースします。

直感のために、この基底2の閾値8が、カーネルで使用される他の単位で何に対応するかを見てみましょう。HEAD_DIM = 128の場合：

基底2指数ギャップ：  8
自然指数ギャップ：      8 / log2(e) ~= 5.545
生のQKスコアギャップ：     5.545 * sqrt(128) ~= 62.7

この62.7は、現在のタイルの行ごとの最大値と、kept_rowmaxに対応する生のQKスコアとの間のギャップです。カーネルでは62.7は使用しません（閾値は8で、log2基底で表現されています）。この変換は、リベースする前に生のQKスコア単位で閾値がどれだけの余裕を許容するかを感覚的に示すためだけです。

カーネルチェックポイント

11_skip_o_rescale.cuを参照し、カーネル10と比較してください。

今、コードに慣れてください。

これまでのパフォーマンス

Chapter 12 – Split Correction from Softmax

［画像］

以前のカーネルでは、同じ行ワープが次のことを行っていました。

1. 現在のタイルの行最大値とOリ スケール係数を計算します。
2. 新しいタイル行最大値が実行中の行最大値を閾値を超えた場合、Oをリ スケールします（カーネル11を参照）。
3. 行最大値をSから減算し、指数化してPを生成します。

したがって、（2）は行ワープの作業完了を遅延させます（1-3）。そのため、O補正とP生成は同じ行ワープで直列化されます。

［画像］

Oリ スケール（2）は正規化されていないPの計算（3）に依存しないため、Oリ スケールを、正規化されていないPを計算する行ワープとは別のワープに移動させることができます。これにより、並列に実行できます。したがって、メインのボトルネックであるP生成を削減します（P生成の前からO補正を削除し、2つのブランチをオーバーラップさせることができます）。行ワープは主にソフトマックス作業を行っているため、今後は「ソフトマックスワープ」と呼びます。

PVは依然として両方の条件が満たされる必要があります。P準備完了とO安全。各Qステージについて、これらはp_ready_o_safeバリアに結合されます。したがって、PV mmaは8つの到着を持つバリアで待機します。4つはソフトマックスワープから（P準備完了を示す）、4つは補正ワープから（O安全を示す）。したがって、遅い方のブランチがPVを遅延させます。

注意：Oが常にリ スケールされるわけではないため（カーネル11の最適化に従って）、o-safeと言い、o-rescaledとは言いません。

ソフトマックスワープは依然としてリ スケール係数を計算し、リ スケール係数をSMEMに公開し、stats_readyを信号で伝えます。その後、Pの生成を続行します。補正ワープは、その信号を待機し、それらの係数を読み取り、Oをリ スケールするかどうかを決定します（公開されたスケールが1でない場合）。

ソフトマックスワープは依然としてリ スケール係数を計算します（リ スケールワープにそれを行わせるのではなく）。なぜなら、ソフトマックスワープは既に古い実行中の行最大値と新しいタイル行最大値をレジスタに持っているからです。そうでなければ、その状態を補正ワープに送信し、そこでリベースロジックを繰り返す必要があります。したがって、ソフトマックスは行ごとに1つの最終スケール値のみを公開し、補正は単純にそれをOに適用します。

同じ4つの補正ワープは、Q0とQ1の両方で共有されます。各Qステージについて、それらはソフトマックスワープによって公開された行ごとのスケールを読み取り、そのステージのOタイルをリ スケールするために使用します。

これまでのパフォーマンス

Chapter 13 – Early PV

［画像］

Pは右から左へ生成されることを思い出してください（TMEM領域の高位側から低位側へ書き込まれる）。これは以前の最適化の結果です（カーネル10を参照）。したがって、Pの高い96列が生成され、古いOが安全になった時点で、PV MMAの一部を開始できます。アイデアは、P全体が準備完了になるまでPVをブロックするのを避け、代わりに次のことを行うことです。まずp_high96 @ V_slice1を実行し、次にP_low32 @ V_slice2を後で実行します（P_low32が生成されたとき）。

これは、Pチャンクの生成を開始する時点で、Sタイルの行最大値全体を既に計算しており、実行中の行最大値/参照が固定されているため、後でP_low32を生成しても、行最大値が変わらず、したがって既に生成されたP_high96を変更できないという事実に依存しています。