NVIDIA、cuTileでMamba-2フォワードパスを高速化

米NVIDIAは、Pythonベースの並列プログラミングDSL「NVIDIA CUDA Tile（cuTile）」を用い、次世代状態空間モデル（SSM（系列長に対して一定の計算量で推論できる状態空間モデル））「Mamba-2」の順伝播（フォワードパス）カーネルを大幅に高速化する実装手法を発表した。金融大手Jane Streetの機械学習エンジニアと共同開発を行い、Blackwell Ultra（B300）環境において既存のTriton実装比で最大1.82倍の高速化を実証した。成果はオープンソースの「TileGym」リポジトリにて公開されている。

cuTileによるMamba-2フォワードパスの3段階カーネル設計

系列モデルの代表格であるTransformerの自己注意機構（Attention）は、学習時には高速な並列行列積を実行できる一方、推論時には系列長に比例して計算量とメモリ消費が増加する。対照的に再帰型ニューラルネットワーク（RNN）は推論時のメモリ消費が一定だが、学習時の並列化が難しい。Mamba-2は固定サイズの隠れ状態（行列）を持ちながら、チャンク単位のジャンプ公式を用いることでTensor Coreに適合した行列積計算を可能にするアーキテクチャである。

cuTileは、スレッドやワープ単位の複雑なスケジューリングを隠蔽し、Python上でタイル単位の並列プリミティブを記述できるプログラミングモデルを提供する。今回のフォワードパス実装では、以下の3つのカスタムカーネルが設計された。

1. chunk_cumsum: チャンク内の減衰係数を計算するカーネル。ct.cumsumを用いてチャンク次元で累積和を算出し、コストの高いexpの代わりにexp2を用いてスケーリングすることで演算効率を高める。
2. fwd_pass_state: チャンク境界における隠れ状態を順次計算するカーネル。ジャンプ公式を行列積和演算（MMA）に落とし込み、Tensor Coreを実行するct.mmaによって各チャンクの初期隠れ状態を生成する。
3. compute_out: 各チャンクの初期隠れ状態とトークン入力を結合し、最終出力を算出するカーネル。通常のPyTorch実装では巨大な中間行列がグローバルメモリ（HBM）上に実体化されるが、本カーネルではタイルをレジスタ上に保持し融合（Fuse）することでメモリ帯域のボトルネックを解消する。

Blackwell Ultra環境でTriton比最大1.82倍の高速化を達成

性能検証は、NVIDIA Blackwell Ultra世代の単一GPU「B300 SXM6（コンピュート・ケイパビリティ 10.3、268 GiB HBM、1,100W制限）」上で実施された。比較対象には、オープンソースのTriton実装（mamba_ssm 2.3.2.post1、Triton 3.7.1）が用いられた。

ベンチマークの結果、cuTileによる最適化フォワードパスカーネルは、Triton実装に対して​1.14倍から1.82倍の高速化​​を達成した。特に系列長（Sequence Length）が長くなり、隠れ状態サイズ（State Size）が拡大するほど高速化の幅が広がる傾向が確認された。

Triton実装とcuTile最適化手法の技術的比較

cuTileによる大幅な高速化を支えるハードウェア最適化手法の差異は下表のとおりである。

最適化項目	従来の課題（Triton / PyTorch等）	cuTileによるアプローチ	得られた効果
DRAMアクセス​	コンパイラがメモリアクセスのレイテンシを推測	ct.load / ct.storeにLatency hint（遅延ヒント）を付与	命令スケジューリングとパイプライン深度の最適化
Tensor Core入力​	第5世代Tensor Coreで非対称なMMAオペランド配置により共有メモリ競合が発生	転置恒等式 $(AB)^\top = B^\top A^\top$ を用いてオペランドA（テンソルメモリ）へルーティング	共有メモリ（Shared Memory）の競合を解消
​メモリアライメント​	デフォルトのストライドアクセスがキャッシュ効率の悪いLDG命令へ変換	[batch, head, token]へ転置して連続メモリアクセスを確保	TMA（テンソルメモリへの非同期転送をハードウェア主導で実行するアクセラレーター）の直接活用
​カーネルチューニング​	入力形状ごとに手作業でタイルサイズや並列度を調整	cuTile-Python内蔵の探索空間オートチューナーを活用	モデル形状に応じた最適構成の自動選定

開発現場への影響と未確認事項

これまでGPU向けの極限的なカーネル最適化には、C++/CUDAやCUTLASSを用いた低レベルなコーディングが必要とされるケースが多かった。今回の発表は、PythonのDSLであるcuTileを用いることで、高度なハードウェア機能（第5世代Tensor Core、TMA、非同期MMA）を活用した最高性能カーネルを簡潔に記述可能であることを示した。

特に定量取引分野で超低レイテンシ推論を追求するJane Streetのような企業が本番環境レベルのモデル実装にcuTileを採用したことは、PythonベースのGPUプログラミングエコシステムの成熟を意味する。

ただし、今回の技術ブログで公表されたベンチマーク結果は「フォワードパス（順伝播）のカーネル単体」に限定されている。モデル全体の学習速度は逆伝播（バックワードパス）の実装や通信オーバーヘッドにも依存するため、エンドツーエンドの学習全体で同等の高速化が得られるかは未確認であり、今後の検証が待たれる。

引用元：NVIDIA Technical Blog