NVIDIAのGPU言語「CUDA Tile」でMamba-2を高速化、Jane Streetと共同の順伝播カーネルを公開

要点

• 米NVIDIAは、Python向けGPUプログラミング環境「CUDA Tile」による「Mamba-2」の順伝播高速化手法を公開した。

• 金融取引大手Jane Streetの技術者と共同で、順伝播処理を担う3つの専用GPUカーネルを最適化した。

• 「chunk_cumsum」「fwd_pass_state」「compute_out」に処理を分割し、Tensor Coreの計算性能を引き出している。

• メモリ配置の転置やレジスタ保持、パラメータの自動チューニングなどを組み合わせ、ハードウェア性能を最大化した。

• 米NVIDIAは2026年10月19日、同社の公式技術ブログにおいて、Python向け並列プログラミング環境「NVIDIA CUDA Tile（cuTile）」を用いたシーケンスモデル「Mamba-2」の順伝播（フォワードパス）高速化手法を公開した。グローバルな金融取引企業であるJane Streetの機械学習エンジニアチームがcuTile上でMamba-2を実装し、NVIDIAと共同で3つのカスタムカーネルを最適化したという。低レベルなGPU並列制御を意識することなく、PythonからTensor Coreの計算性能を最大限に引き出す手法が紹介されている。

シーケンスモデルの計算課題とcuTileの役割

自然言語処理などのシーケンスモデルには、計算効率上の課題がある。代表的なSoftmax Attentionは、Tensor Core（行列演算専用プロセッサ）により高速に学習できるが、推論時に系列長に応じて計算量とメモリ消費量が線形に増加する。一方の再帰型ニューラルネットワーク（RNN）は推論負荷が一定だが、並列化が難しく学習が遅い。

この双方の利点を狙ったのが「Mamba-2」だ。RNNのような固定サイズの隠れ状態を持ちながら、系列をチャンク（塊）単位でTensor Coreを用いて効率よく再帰計算できる。隠れ状態はベクトルではなく行列で保持され、高い表現能力を持つという。

実装基盤の「NVIDIA CUDA Tile（cuTile）」は、Python向けの並列プログラミングDSL（特定領域向け言語）である。スレッドやワープ単位の複雑なスケジューリングをタイル単位の操作へと抽象化し、Tensor Coreなどのハードウェア機能を自動的に活用できるのが特徴だ。

順伝播を担う3つの最適化カーネル

Jane StreetとNVIDIAは、Mamba-2の順伝播処理を3つの専用カーネルに分割して最適化を進めた。

1つ目の「chunk_cumsum」は、トークンごと・ヘッドごとの減衰項について、対数空間での累積和を算出するカーネルだ。cuTileの「ct.cumsum」を用いて計算し、後続処理で再利用するために結果を保持する。

2つ目の「fwd_pass_state」は、系列のチャンクをまたいで隠れ状態を更新する。行列積プリミティブ「ct.mma」でTensor Core演算を実行し、指数関数的減衰を伴う「ジャンプ公式」と事前乗算されたオペランドを適用して状態を進める。

3つ目の「compute_out」は、引き継がれた状態項、減衰と因果マスクが適用された相互作用行列、スキップ接続をまとめて計算する。これらを単一のタイル化カーネルに融合（フュージョン）し、中間行列をレジスタに保持したまま「TF32（TensorFloat-32）」形式へキャストしてTensor Coreで高速実行する。

ハードウェア最適化と利用環境

本取り組みでは、GPUのハードウェア性能を引き出す多角的なチューニングが施された。

GPUメインメモリ（DRAM）との通信を効率化するため、レイテンシに関するヒント情報をコンパイラに与えるモデリングが導入された。また、融合されたオペランドをTensor Memoryへ適切にルーティングするオペランドスワップや、データ転送支援機構「TMA（Tensor Memory Accelerator）」に適した転置メモリレイアウトを採用している。

さらに、入力データの形状ごとに最適なブロックサイズやGPU占有率を探索する「オートチューナー」も備え、多様な条件下で安定した性能を発揮できるようにした。

なお、cuTileはPython環境で「cuda.tile（ct）」として利用でき、NVIDIAはGitHubなどでドキュメントや実装サンプルを公開している。