- DeepGEMMは、現代のLLMにおける中核計算プリミティブであるGEMM、fused MoE、MQA scoring、HyperConnectionなどを1つのCUDAコードベースにまとめた高性能tensor coreカーネルライブラリ
- すべてのカーネルは軽量なJITモジュールとして実行時にコンパイルされ、インストール時のCUDAコンパイルは不要で、C++20、CUDA Toolkit、PyTorch、CUTLASS 4.0以上を必要とする
- CUTLASSとCuTeの一部概念を活用しているが、テンプレートや代数構造に重く依存せず、限られた数の中核カーネル関数によりNVIDIA GPUカーネル最適化の学習に取り組みやすいよう設計されている
- 対応範囲にはFP8、FP4、BF16 GEMM、grouped GEMM、DeepSeek v3.2向けのMQA logitsカーネル、通信と計算を重ねるMega MoEまでが含まれ、SM90とSM100ではメモリレイアウトの制約が異なる
- 軽量設計でありながら、さまざまな行列shapeで専門家チューニング済みライブラリと同等以上の性能を目指しており、H800で最大1550 TFLOPSを達成した更新が含まれる
DeepGEMMの目的と設計
- DeepGEMMは、最新の大規模言語モデルで使われる主要な計算プリミティブを1つのCUDAコードベースに統合したtensor coreカーネルライブラリ
- GEMM: FP8、FP4、BF16
- 通信を重ねたfused MoE: Mega MoE
- lightning indexer向けMQA scoring
- HyperConnection(HC)
- すべてのカーネルは軽量な**Just-In-Time(JIT)**モジュールとして実行時にコンパイルされる
- CUTLASSとCuTeの一部概念を活用
- ただし重いテンプレートや代数構造には大きく依存しない
- 中核カーネル関数の数を限定し、コードベースをシンプルに保っている
- 軽量設計でありながら、複数の行列shapeで専門家チューニング済みライブラリと同等以上の性能を出すとしている
主な更新
- 2026年4月16日の更新では、Mega MoE、FP8xFP4 GEMM、FP4 Indexer、PDL、より高速なJITコンパイルなどを含む
- 2025年9月28日の更新では、DeepSeek v3.2 lightning indexer向けのweighted ReLU MQA logits scoringカーネルを追加
- 2025年7月20日の更新では、SM90とSM100の両方をサポートし、低CPUオーバーヘッドのJIT CPPモジュールへ全面リファクタリングされた
- NVRTCとコンパイル後のSASS最適化は無効化された
- NVRTCは今後対応予定とされている
- NVCC 12.9がFFMA interleavingを自動で行うため、コンパイル後最適化はもはやサポートされない
- 詳細は#112
- 2025年5月14日の更新では、denseおよびMoE backward向けのweight gradientカーネルを追加
- 2025年5月7日の更新では、NVRTC対応により最大10倍のコンパイル高速化を提供
DG_JIT_USE_NVRTC=1で有効化可能- 一部のケースでは性能低下の可能性がある
- 詳細は#94
- 2025年4月18日の更新では、H800で最大1550 TFLOPSを達成
要件とインストール手順
- 実行環境にはNVIDIA SM90またはSM100アーキテクチャGPUが必要
- ソフトウェア要件は以下の通り
- Python 3.8以上
- C++20対応コンパイラ
- CUDA Toolkit
- SM90: CUDA 12.3以上
- 最高性能のためにはCUDA 12.9以上を強く推奨
- SM100: CUDA 12.9以上
- PyTorch 2.1以上
- CUTLASS 4.0以上
{fmt}ライブラリ
- 開発環境では、submoduleを含めてリポジトリをcloneした後、
develop.shで必要なincludeの接続とCPP JITモジュールのビルドを行う
- インストールは
install.shを実行し、その後Pythonプロジェクトでdeep_gemmをimportする方式
GEMMインターフェースとレイアウト制約
- DeepGEMMのGEMMカーネル命名規則は
D = C + A @ B
- 入力shapeレイアウトはNTを基準とする
fp8_gemm_ntはD = C + A @ B.Tを実行する
- SM90実装はNTメモリレイアウトのみをサポート
- row-major、col-majorの組み合わせに相当する
- SM100実装はNT、TN、NN、TTすべてのメモリレイアウトをサポート
- 両アーキテクチャとも、LHS scaling factorはTMA整列済みかつ転置されたレイアウトである必要がある
- SM90はscaling factorをFP32形式で要求する
- SM100はpacked UE8M0形式を要求し、4つのUE8M0を1つの
torch.intにパックする
- 入力転置やFP8キャストのような処理はユーザーが別途行う必要がある
- ライブラリは簡単なPyTorchユーティリティ関数を提供するが、性能は低い場合がある
- 主な焦点はGEMMカーネル最適化にある
DenseとGrouped GEMM
- 基本のnon-grouped FP8 GEMMは
fp8_gemm_{nt, nn, tn, tt}関数を使う
- contiguous layoutのgrouped GEMMは、CUTLASSの従来のgrouped GEMMと異なりM軸のみをグループ化する
- NとKは固定である必要がある
- MoEモデルでexpertsが同じshapeを共有する状況に合わせた設計
- 学習のforward passや推論prefillingでは、expertごとに処理するtoken数が異なる場合がある
- こうしたtokenを1つのtensorに連結した形をcontiguous layoutと呼ぶ
- 各expert segmentはGEMMのMブロックサイズに整列している必要がある
- 整列条件は
get_mk_alignment_for_contiguous_layout()で確認できる
- MoE weight backward向けにK軸grouped APIも提供される
- MとNは固定である必要がある
- 関連関数は
k_grouped_fp8_gemm_tn_contiguous
- 推論decoding段階でCUDA graphが有効で、CPUがexpertごとのtoken数を把握できない場合には、masked grouped GEMMをサポート
- mask tensorを与えると、カーネルが有効領域のみを計算する
- 関数は
m_grouped_fp8_gemm_nt_masked
- DeepEPのlow-latencyカーネル出力を入力に使う例がある
DeepSeek v3.2 Indexer向けMQAカーネル
- V3.2 MQAカーネル群はnon-paged版とpaged版を提供
- non-pagedはprefilling用
- pagedはdecoding用
fp8_mqa_logitsは6つの入力を受け取る
q: E4M3 tensor、shape [seq_len, num_heads, head_dim]kv: E4M3 tensorとfloat scaling factor
- tensor shapeは
[seq_len_kv, head_dim]
- scaling factor shapeは
[seq_len_kv]
weights: float tensor、shape [seq_len, num_heads]cu_seq_len_k_start、cu_seq_len_k_end: int tensor、shape [seq_len]clean_logits: 埋まっていないlogitsを-infで整理するかどうか
- 出力tensor shapeは
[seq_len, seq_len_kv]で、token-to-token logitsを表す
- 各
q token iはcu_seq_len_k_start[i]からcu_seq_len_k_end[i]の手前までのkv token jを走査する
kv_jにscaling factorを掛けるq[i, :, :] @ kv_jでheadごとの値を計算する- ReLU適用後に
weights[i, :]を掛けて合計し、スカラーlogitを作る
- paged版の関数は
fp8_paged_mqa_logits
Mega MoE
- Mega MoEは複数のMoEステップを1つのmega-kernelにfuseする
- EP dispatch
- linear 1, FP8xFP4
- SwiGLU
- linear 2, FP8xFP4
- EP combine
- Mega MoEはNVLink通信とtensor core計算をoverlapする
- 実行にはsymmetric memoryを用いるmulti-process launchが必要
- 使用手順は以下の通り
deep_gemm.get_symm_buffer_for_mega_moeでsymmetric memory bufferを確保する
deep_gemm.transform_weights_for_mega_moeで、FP4とUE8M0 SFを含むweightを必要なレイアウトへ変換する- 呼び出し前に入力、scaling factor、top-k index、top-k weightをbufferへコピーする
deep_gemm.fp8_fp4_mega_moeでfused mega MoEカーネルを実行する
- 完全なmulti-process設定とベンチマーク例は
tests/test_mega_moe.pyにある
ユーティリティと環境変数
- 主なユーティリティ関数は、実行リソース、整列、JITコンパイル、scaling factor変換を制御する
deep_gemm.set_num_sms / get_num_sms: 使用する最大SM数の設定と取得deep_gemm.set_tc_util / get_tc_util: おおよそのtensor core utilization ratioの設定と取得deep_gemm.set_pdl / get_pdl: Programmatic Dependent Launch(PDL)の有効化と無効化deep_gemm.set_mk_alignment_for_contiguous_layout / get_mk_alignment_for_contiguous_layout: contiguous layoutにおけるgroup-level M/K alignmentの設定と取得deep_gemm.transform_sf_into_required_layout: scaling factorを必要なレイアウトへ変換deep_gemm.get_tma_aligned_size: 必要なTMA alignment sizeを取得
- JIT関連の環境変数は、デバッグ出力、キャッシュ位置、コンパイラ選択、プロファイリングオプションを制御する
DG_JIT_DEBUG: JITデバッグ情報を出力DG_PRINT_CONFIGS: shapeごとに選択されたconfigを出力DG_JIT_CACHE_DIR: コンパイル済みカーネルのキャッシュディレクトリ。デフォルトは$HOME/.deep_gemmDG_JIT_USE_NVRTC: NVCCの代わりにNVRTCを使用。高速コンパイルが可能だが、一部ケースでは性能が低下する可能性があるDG_JIT_NVCC_COMPILER: NVCCコンパイラのパスDG_JIT_CPP_STANDARD: C++標準バージョン。デフォルトは20
- デバッグおよびプロファイリング用の環境変数も提供される
DG_JIT_DUMP_ASM、DG_JIT_DUMP_PTX、DG_JIT_DUMP_SASS: PTXとSASSの出力ダンプDG_JIT_WITH_LINEINFO: プロファイリングツール向けにsource line情報を含めるDG_COMM_KERNEL_DEBUG: Mega MoE呼び出し前にsymmetric bufferを0で初期化DG_USE_NVIDIA_TOOLS: 外部NVIDIAツール実行時に内部profilingをスキップ
- ビルドオプションは、インストール方法とカーネル読み込み方式を制御する
DG_SKIP_CUDA_BUILD: インストール時のCUDA extension buildをスキップDG_FORCE_BUILD: pre-built wheelのダウンロードの代わりにローカルビルドを強制DG_JIT_USE_RUNTIME_API: カーネル読み込みにCUDA Runtime APIを使用。CUDA runtime 12.8以上が必要
ライセンスと引用
- DeepGEMMリポジトリはMIT Licenseで公開されている
- プロジェクトはCUTLASSから着想を得たとしている
- 引用項目のタイトルは
DeepGEMM: clean and efficient BLAS kernel library on GPU
1件のコメント
Hacker Newsのコメント
FFMA SASSのインターリーブは本当に驚くべきレベルに見える
NVCC 12.2と12.3の間でCUTLASS FP8カーネルの性能が向上したのを見て、コンパイル済みSASSを比較したところ、複数のFADD命令で1ビットがインターリーブパターンで反転しており、オープンソースのCUDAアセンブラ実装を参考に、そのビットが現在のワープを譲って別のワープを実行させるyieldビットだと突き止めたようだ
これを活用してコンパイル済みバイナリのFFMA命令を修正するスクリプトを作り、yieldビットだけでなく、ワープが譲られるとレジスタ再利用ができないためreuseビットもあわせて反転させ、細粒度スケーリングFP8 GEMMでMMA命令と昇格FFMA命令をよりうまく重ね、一部のケースで性能を10%以上引き上げたという点が印象的
ただ、この特定の問題については他のAI企業がまだ必要性を感じておらず、適用していなかったようで、結局みな似た地点に到達する可能性が高い
こうした事例は、高水準コードだけでハードウェア性能を引き出すには、既存のコンパイラがどれほど遠い位置にあるかをよく示している
伝統的なコンパイラ技術やAIベースの最適化エージェントがこうした結果を出すには、何が必要なのか気になる
報告されている速度向上の数値は、独自のCUTLASSベースのベースラインとの比較だ
cuBLASと直接性能比較をした人がいるのか気になる
これまで見たCUTLASS GEMMの結果はcuBLAS比でおおむね10%以内だったので、論文で述べられている2倍〜2.5倍の向上が維持されるなら本当に印象的だと思う
まずcuBLASは、FP8のような混合精度処理を扱うにはcuBLASLt拡張APIが必要になる
また、A x BでE5M2 x E5M2のような妥当に見える型の組み合わせはサポートしていない一方、E5M2 x E4M3はサポートしており、Ampere、Hopper、Blackwellでは行列Aが常に転置レイアウトでなければならないなど、制約が続く
FP8 cuBLASLtベンチマークを自分の「Less Slow C++」リポジトリ<https://github.com/ashvardanian/less_slow.cpp>に統合し、既存のcuBLASおよび自作のCUDA/PTXベンチマーク一覧に追加した
H100と同等の性能であるはずのH200 GPUで実行しており、正方入力ではスループットがおよそ1.35 Peta-opsでピークに達する
256は2.68T/s、512は20.49T/s、1024は144.23T/s、2048は665.68T/s、4096は1.26P/s、8192は1.34P/s、16384は1.23P/s程度で、これは密GEMMについてNVIDIAが宣伝している数値<https://resources.nvidia.com/en-us-data-center-overview-mc/e...>の約67%だ
ベースラインはcuBLASとCUTLASSのうち良い方を選んだものだと思っていた
こういうオープンソースは、業界の効率性達成という目標を本当によく示している
ただし、このソフトウェアの恩恵は、学習や実験をしたり、消費者向けハードウェアでモデルをサーブしたりする一般のオープンソースコミュニティよりも、大規模にモデルをサーブする大手企業、つまりDeepSeekの潜在的な競合により大きく及びそうだ
ますます低い精度へ最適化していく方向が長期的に良いのかは、よく分からない
これはモデルが実際にはかなり疎であることを意味するが、今はそうであり得るとしても、本質的にそれほど疎でなければならないからというより、学習方法に悪いアイデアが混ざっているためである可能性が高いと思う
より高い精度でしか非常に良いモデルを学習できないようにするのは研究課題で、低精度の学習と推論は工学の問題だ
CNNの時代から、少なくとも9年前からこういうことをやってきたし、今後もまだ数年は続くと思う
これはBlackwellのネイティブなマイクロスケーリング対応であるMXFPのせいで無意味になるかもしれない
Hopperではそれをより粗い単位で、代わりにFP32のスケーリング係数を使って手動実装しているようなものだ
こういう高品質な公開デモは、$NVDAの堀がどこにあるのかをよく示している
汎用GPUは非常に柔軟なので、ハードウェアベンダーが最初から思いつかなかったが十分に筋の通ったさまざまな処理を、プログラミングで実行できる
ただし、未来がますます専用ハードウェア対応へ収束し、このようなソフトウェア最適化の余地が消えていくと予測するなら、いわゆるCUDAの堀は崩れる
このゲームに残るために、NVIDIAは自らの堀を自分で壊しているようなものだ :p
わあ、MITライセンスだ
大手企業がこういうオープンソース協業のやり方を受け入れてくれるといい
なぜ文書化されていない命令が存在するのか、ずっと気になっている
完全に安定していなくても、ユーザーに提供するほうが正しいのではないかと思う
こうした内容は社内では文書化されているはずなのに、なぜ公開しないのか分からない
曖昧さに依存したセキュリティは通用しないし、競合他社はいずれにせよ全部リバースエンジニアリングする
時間不足だったり、不安定または実験的な機能についてサポートを示唆したくなかったりするからかもしれない
影響が隣のチーム程度にとどまるなら、変更するのもはるかに簡単になる
おそらくアーキテクチャ設計文書や仕様のような場所にだけ入っている可能性が高く、そういう文書は当然共有したくないはずだ
正直、自分の利用範囲と理解を超える内容だ
それでも、こうした発見と改善を共有し、誰もが恩恵を受けられるようにしてくれる点は本当にありがたく、新鮮に感じる
行列乗算やディープラーニングのワークロードで非常に重要である
NVIDIA の SASS では、FFMA 命令は64ビットまたは128ビット命令としてエンコードされ、正確な動作を決める複数の制御ビットを持つ
yield ビットが設定されると、ワープスケジューラに対して、この命令の後に現在のワープが実行を譲れることを知らせ、ハードウェアは別のワープを実行してレイテンシを隠蔽できる
GPU は大規模な並列性によって高いスループットを得ており、あるワープがメモリ待ちなどで停止すると、別のワープが進行できる
reuse ビットはソースレジスタを直後の演算で再利用できるかを示し、yield ビットが設定されている場合は必ずオフにしなければならない
ワープが譲ると、次に実行されるのがそのワープではない可能性があり、別のワープがレジスタファイルの状態を変える可能性もあるため、ハードウェアは yield をまたいでレジスタ値が維持されることを保証できないからだ
FFMA 命令群に yield ビットを交差パターンで設定すると、コンパイラは別のワープが進行できる明示的なスケジューリングポイントを作り、正確性を保つために該当命令の reuse ビットも同時にクリアする必要がある
この変更は、行列乗算の中核であるMMA命令と、FP8 をより高い精度で累算するために変換する昇格 FFMA 命令を重ね合わせるうえで特に役立つ
FP8 GEMM は通常、累算のためにより高い精度へ変換し、再び戻す過程が必要になるため追加の FFMA が発生し、メモリ帯域幅の要求は減らす一方で、昇格/降格演算が混在する複雑な計算パターンを生む
「細かなスケーリング」は、計算の複数の時点で精度を慎重に管理する作業を指しているようだ
yield ビットの操作は、計算演算と形式変換がよりうまくインターリーブされるようにし、GPU の実行ユニットをより効率的に使えるようにする。この最適化がないと、ワープスケジューラが自然な切り替え機会を見つけられず、計算リソースの活用度が下がる可能性がある