NVIDIAとジェーン・ストリートが、Mamba-2の入力から出力を生成する順伝播計算を高速化するGPUカーネルを3つ公開した。GPUカーネルはGPUで特定の演算を行うコードだ。モデルを最適化する開発者はPythonで書かれた公開コードを利用でき、同社の発表では既存のオープンソースカーネルより1.14~1.82倍高速な結果となった。

1. 再帰型モデルの計算をチャンク単位に変換

Mamba-2は、それまでに処理した情報を固定サイズの内部記憶である「隠れ状態」に格納して更新する再帰型モデルだ。学習やプロンプト処理では、トークンを1つずつ計算する代わりに、複数のトークンをまとめたチャンクの計算を行列積に変換する。GPUの行列演算専用装置であるTensor Coreを活用するためだ。

NVIDIAのcuTileは、GPUカーネルをPythonで記述するツールだ。開発者はデータのまとまりであるタイルの演算を記述し、詳細な実行配置はコンパイラに任せる。ジェーン・ストリートのエンジニアがこれを使ってMamba-2を実装し、NVIDIAとともにカーネルを最適化した。

2. 3つのカーネルで状態と出力を分担して計算

最初のカーネルchunk_cumsumは、過去の情報が減少する度合いを表す減衰率と時間間隔の積を、チャンク内で累積して保存する。後続の2つのカーネルはこの値を再利用する。

2番目のfwd_pass_stateは、チャンク開始時の隠れ状態と最後のチャンク後の状態を計算する。チャンクの順序は守りつつ、互いに独立した状態の部分は並列に処理する。

3番目のcompute_outは、チャンク開始時の状態からトークンの出力を生成する。中間行列を大きなグローバルメモリに別途書き出す代わりに、計算を融合し、演算装置に近い小さな記憶領域であるレジスタに保持する。

ブログの累積和カーネルは、タイルを読み込み、チャンク内で累積して保存する流れを次のように示している。

    dt_tile = ct.load(dt, index=(batch, chunk_idx, head_tile_idx),
                          shape=(1, chunk_size, BLOCK_H),
                          padding_mode=ct.PaddingMode.ZERO).astype(ct.float32)
    a_tile = ct.load(a, index=(head_tile_idx,),
                        shape=(BLOCK_H,),
                        padding_mode=ct.PaddingMode.ZERO).astype(ct.float32)

    d_tile = ct.cumsum(a_tile * dt_tile, axis=1) * LOG2E
    ct.store(output, index=(batch, chunk_idx, head_tile_idx), tile=d_tile.astype(output.dtype))

3. ハードウェアに合わせたメモリ・行列積の最適化

実装ではメモリの読み書きに予想遅延の情報を付け、コンパイラが命令順序を調整できるようにした。Blackwellの第5世代Tensor Coreの入力配置特性に合わせて乗算入力の順序を変え、複数の処理が記憶領域を取り合う共有メモリ競合も減らした。減衰の累積値と時間間隔のメモリ配置を変え、データを連続して読み込めるようにしたことも最適化の要素だ。

入力形状によってボトルネックが変わるため、ブロックサイズ、つまり一度に処理するデータのまとまりの大きさと、GPUリソースの活用設定を自動探索する。実行時間を短縮できる一方で、入力形状ごとに設定の組み合わせを探索する作業が必要となる仕組みだ。

4. 同社発表のカーネル測定では最大1.82倍

NVIDIAブログの測定では、単一のB300 SXM6でcuTile 1.5.0を使用した。GPUの最大演算クロックは1920 MHz、高帯域幅メモリの容量は268 GiBだ。

比較対象はTriton 3.7.1で書かれたオープンソースカーネルで、mamba_ssm 2.3.2.post1のNVIDIA Dynamo系のリファレンス実装だ。

測定にはPyTorch 2.13.0とCUDA 13.3を使用した。

同社の発表ではcuTileカーネルはTritonより1.14~1.82倍高速で、シーケンス長と状態サイズが大きいほど速度向上が大きくなる。グラフは1K~256Kトークンと状態サイズ64、128、256、512を扱うが、各棒に正確な数値ラベルはない。

この結果は順伝播カーネル単位の比較であり、学習全体の性能は、出力から入力側へさかのぼる逆伝播計算の実装にも左右されるとブログは説明している。

5. TileGymで実装を確認

コードとベンチマークはNVIDIAのTileGymリポジトリで確認できる。デフォルトの実行方式はcuTileで、リポジトリのインストール手順に従って演算例と性能測定を確認できる。

リポジトリにはcuTile 1.5.0以上と実行時に必要なコンパイラが明記されている。インストール案内には、Blackwell GPUではCUDA 13.1以上、Ampere GPUではCUDA 13.2以上をサポートすると記載されている。

今回の公開は、Mamba-2の順伝播計算の行列化とハードウェア別のチューニングが実測での効果につながり得ることを、コードとともに示している。