deepseek-ai/DeepGEMM-Ascend

DeepGEMM-Ascend: clean and efficient matrix multiplication kernel library for Huawei Ascend NPUs

DeepGEMM‑Ascend とは？

DeepGEMM‑Ascend は、オープンソースの DeepGEMM 行列乗算ライブラリのハードウェア固有の移植版であり、Huawei の Ascend NPU ファミリー（例：Ascend 950）をターゲットにしています。DeepGEMM と同じ Python レベルの API を実装しているため、他のプラットフォーム（CUDA、ROCm など）用に書かれたコードは、このパッケージをインストールするだけで Ascend 上で実行できます。

このライブラリは、Ascend の MAD（行列加算乗算）プリミティブを薄い C++/Python 拡張でラップし、フラクタルレイアウト、アライメント、アドレス計算などの低レベルな詳細を処理します。また、スパースデータロードやコルーチンベースのパイプラインなどの Ascend 固有の最適化も追加し、さまざまな行列形状でハードウェアの理論上のピーク性能を目指しています。

誰が使うのか？

対象	重要性
AI 研究者／エンジニアが Ascend NPU 上で実行する大規模言語モデル（LLM）を構築する場合 – 高速な GEMM、MQA、MoE カーネルが必要です。
フレームワーク開発者が PyTorch‑NPU（`torch_npu`）をカスタム高性能演算で拡張する場合。
パフォーマンスエンジニアが極限まで最適化された Ascend カーネルのリファレンス実装を探している場合。
ハードウェアを意識した ML 実践者が Ascend 上で BF16、FP8、FP4、混合精度行列演算をベンチマークしたい場合。

コア機能（README に記載）

機能	詳細
API 互換性	アップストリームの DeepGEMM プロジェクトと同じパッケージ名と関数シグネチャを使用。
精度サポート	BF16、FP8、FP4 GEMM。MQA logits と MegaMoE 用の専用カーネルも提供。
軽量な抽象化	Ascend 固有のレイアウトとアライメントの複雑さを隠しつつ、カーネルコードを簡潔に保つ。
Ascend 固有の最適化	スパースデータロード、コルーチンベースのパイプライン、カスタムスケーリングファクターパッキング（UE8M0 ペア → `int16`）。
ユーティリティヘルパー	AI コア数の設定、NPU アーキテクチャの選択、決定論的モードの制御、スケーリングファクターレイアウトの変換、理論上のアライメント制限の照会。
JIT ランタイム	DeepJIT ランタイムに依存。カーネルは初回使用時にコンパイルされ、キャッシュされます（`DG_JIT_*` 環境変数で設定可能）。
環境変数チューニング	デバッグ、キャッシュ場所、起動タイムアウト、その他の設定を `DG_JIT_*` と標準の Ascend 変数で公開。
パフォーマンス	ベンチマークでは、BF16/FP8/FP4 の高密度 GEMM でハードウェア限界の 98〜99% を達成。グループ化 GEMM、MQA、MegaMoE、カスタム「HC Prenorm」カーネルの詳細な数値も提供。

始め方（クイックスタート概要）

ハードウェアとソフトウェアの前提条件
- Huawei Ascend NPU（950 シリーズで検証済み）
- CANN 9.20 ツールキット（bisheng、ld.lld）
- torch_npu Python パッケージ
- Python 3.10+、C++20 互換コンパイラ、および tilelang、tree‑sitter、tree‑sitter‑cpp

サブモジュール付きでクローン

git clone --recursive https://github.com/deepseek-ai/DeepGEMM-Ascend.git
cd DeepGEMM-Ascend

C++ 拡張をビルド（リポジトリには include をリンクしてコンパイラを実行する develop.sh スクリプトが用意されています）。
インストール
```
pip install . --no-build-isolation
```

Python コードで同じ DeepGEMM API を使用、例：

import deep_gemm as dg
dg.set_num_sms(8)          # 使用する AI コア数を選択
out = dg.aclnn_fp8_fp4_gemm_nt(A, B, scaling_factors)

いつ最高のパフォーマンスを発揮するか？

README には、ライブラリが Ascend ハードウェアの上限に近づくいくつかのベンチマークカテゴリがリストされています：

高密度 GEMM – BF16、FP8、FP4 で理論上の TFLOPS の最大 99.8%。
M‑Grouped GEMM – Mixture‑of‑Experts（MoE）エキスパートに有用で、現実的なエキスパート・トークンパターンで > 800 TFLOPS を達成。
MQA Logits – マルチクエリ注意の高速スコアリング、レイテンシは 100〜300 µs の範囲。
MegaMoE – エンドツーエンドの融合カーネル（dispatch + 2 つの GEMM + SwiGLU）で、計算とメモリ帯域幅の両方を高く維持。
HC Prenorm GEMM – 「Manifold‑Constrained Hyper‑Connections」モジュール用の専用カーネルで、HBM 書き込み帯域幅を飽和させます。

これらの数値は、ライブラリが Ascend ハードウェア上の本番規模の LLM 推論／トレーニングワークロードに ready であることを示しています。

制限事項／注意点

ハードウェア固有 – Huawei Ascend NPU でのみ動作。CPU や GPU のフォールバックはありません。
ツールチェーンが重い – CANN 9.20、bisheng コンパイラ、<format> をサポートする C++20 コンパイラが必要。
環境変数チューニング – パフォーマンスは DG_JIT_* 変数と Ascend の起動ブロッキング／デバッグフラグを正しく設定することに依存する場合があります。
ベータステータス – リポジトリは初期リリース（2026‑09‑30）と報告。コミュニティサポートとドキュメントはまだ限られている可能性があります。

AI エコシステムへの適合

DeepGEMM‑Ascend は、AI スタックの低レベル線形代数層に位置します。現代の LLM は、計算の大部分を大規模な行列乗算（GEMM）と関連するアテンションカーネルに費やしています。広く使用されている DeepGEMM API をミラーリングした、高度に最適化された Ascend ネイティブ実装を提供することで、このプロジェクトは以下を可能にします：

シームレスな移植 – 既存の PyTorch ベースのモデルをカーネルを書き直すことなく Ascend に移植。
正確なパフォーマンスベースライン – Ascend を NVIDIA/AMD GPU と比較する研究者向け。
リファレンスコード – 新しい Ascend 最適化技術（スパースロード、コルーチンパイプライン、カスタムスケーリングファクターパッキング）を探求するハードウェアエンジニア向け。

ライセンスと引用

ライセンス： MIT – 商用および研究利用に寛容。
引用： 学術論文用にコピー可能な BibTeX エントリが提供されています。

TL;DR

DeepGEMM‑Ascend は、Huawei Ascend NPU 向けのドロップインで高性能な GEMM ライブラリであり、元の DeepGEMM プロジェクトと同じ Python API を公開しつつ、Ascend 固有の最適化を追加しています。LLM ワークロード（高密度 GEMM、MoE、MQA）をターゲットとし、さまざまな行列形状でほぼピークのハードウェア利用率を主張しています。pip でインストール可能で、Ascend CANN ツールキットと C++20 コンパイラが必要です。