deepseek-ai

DeepGEMM

AI#CUDA#LLM#GPU#Deep Learning#Optimization

7,104

// 概要

DeepGEMM は、最新の LLM に最適化された高性能な Tensor Core カーネルを提供する統合 CUDA ライブラリです。インストール時の CUDA コンパイルを不要にする軽量な Just-In-Time コンパイルモジュールを備えています。FP8、FP4、BF16 の GEMM や、fused MoE および MQA スコアリングなど、多様な行列演算において専門家が調整したパフォーマンスを実現します。

// 技術解説

DeepGEMM は、現代の LLM に不可欠な計算プリミティブを提供するために設計された、高性能で統合された CUDA カーネルライブラリです。これには FP8/FP4 GEMM、fused MoE、MQA スコアリングなどが含まれます。軽量な JIT コンパイルモジュールを活用することで、インストール時の複雑な CUDA コンパイルを不要にしつつ、専門家が調整したライブラリに匹敵するパフォーマンスを維持します。このプロジェクトは、複雑なテンプレートメタプログラミングへの過度な依存を避けることでシンプルさとアクセシビリティを優先しており、開発者が高度な NVIDIA GPU カーネル最適化を学習・実装するためのクリーンなコードベースを提供します。

// 主要ハイライト

FP8、FP4、BF16 を含む複数のデータフォーマットをサポートする高性能な GEMM カーネルを提供します。

NVLink 通信と Tensor Core の計算をオーバーラップさせ、スループットを最大化する fused カーネルである Mega MoE を実装しています。

実行時にカーネルをコンパイルする軽量な JIT モジュールを備えており、手動での CUDA コンパイルの負担を解消します。

lightning indexer 用に設計された特殊な MQA スコアリングカーネルが含まれており、非ページメモリとページメモリの両方のレイアウトをサポートします。

最適化された TMA アラインメントメモリ操作により、SM90 や SM100 アーキテクチャなどの高度なハードウェア機能をサポートします。

効率的な MoE エキスパート処理に特化した、連続およびマスクされたレイアウト向けの柔軟な grouped GEMM API を提供します。

// ユースケース

LLM 向けの高性能な FP8、FP4、および BF16 GEMM 演算

通信と Tensor Core 計算を融合した Mega MoE カーネル

大規模推論における lightning indexer 用の MQA スコアリングカーネル

// クイックスタート

開始するには、すべてのサブモジュールが含まれるように 'git clone --recursive' を使用してリポジトリをクローンしてください。提供されている 'develop.sh' スクリプトを実行して必要なインクルードをリンクし、C++ JIT モジュールをビルドした後、'install.sh' を実行してインストールを完了させます。インストールが完了したら、'deep_gemm' モジュールを Python 環境に直接インポートして、最適化されたカーネルにアクセスできます。