deepseek-ai/DeepGEMM-Ascend

DeepGEMM Ascend

DeepGEMM Ascend is a port of DeepGEMM to the HUAWEI Ascend platform. It is fully API-compatible with DeepGEMM and supports BF16, FP8, FP4 GEMM, MQA logits, and MegaMoE. On Ascend platforms, users can simply install the package and use the same APIs and development workflow as DeepGEMM on other supported platforms.

DeepGEMM Ascend provides a lightweight abstraction over the Ascend MAD (matrix multiply-add) primitives, hiding much of the complexity associated with fractal layouts, alignment constraints, address calculations, and verbose low-level parameters. This enables GEMM kernels to remain both concise and efficient. DeepGEMM Ascend makes extensive use of Ascend-specific optimization techniques, such as sparse data loading and coroutine-based pipelining, to approach the performance limits of the Ascend hardware. These implementations can also serve as references for extreme performance optimization on the Ascend platform.

Despite its lightweight codebase, DeepGEMM Ascend can achieve peak hardware performance across a wide range of matrix shapes.

DeepGEMM Ascend 是华为昇腾平台上的 DeepGEMM 实现。它完全兼容 DeepGEMM 的 API，支持 BF16、FP8、FP4 GEMM、MQA logits 和 MegaMoE 算子。DeepGEMM Ascend 和 DeepGEMM 使用相同的包名，用户只需安装当前包，即可沿用其他平台上 DeepGEMM 的 API 和开发流程。

DeepGEMM Ascend 对昇腾平台的矩阵乘加原语（MAD）提供了一个轻量的抽象，能够隐藏分形矩阵布局、对齐约束、地址计算、参数转换等细节，让 GEMM kernel 的实现保持简洁高效。DeepGEMM Ascend 广泛采用昇腾平台特有的优化技术，包括稀疏数据加载、基于协程的流水线等，以接近硬件的性能极限。这些实现也可以作为昇腾平台极致性能优化的参考。

DeepGEMM Ascend 代码轻量，并且在多种矩阵形状上能达到硬件极限性能。

News

• 2026.09.30: Initial release of DeepGEMM Ascend, with support for Ascend 950 devices. The kernels are designed to achieve near-peak hardware performance.

Quick Start

Requirements

• HUAWEI Ascend NPU (developed and validated on the Ascend 950 series)
• CANN 9.20 toolkit providing bin/bisheng, bin/ld.lld
• Torch NPU package torch_npu
• Python 3.10 or higher
• Compilers and standard libraries with C++20 `` support
• tilelang, used by the HC prenorm kernel (declared as a package dependency)
• tree-sitter and tree-sitter-cpp, used to generate Python type stubs when building from source

Development

# Submodule must be cloned
git clone --recursive https://github.com/deepseek-ai/DeepGEMM-Ascend.git
cd DeepGEMM-Ascend
Link some essential includes and build the C++ extension
cat develop.sh
./develop.sh

Installation

pip install . --no-build-isolation

Interfaces

Kernel Interface

Please refer to DeepGEMM's interfaces for the kernel APIs.

[!NOTE]

The scaling factor format on Ascend differs from NVIDIA's: each pair of UE8M0 scaling factors along the K dimension is packed into an int16, and the packed values are stored in MN-major order for optimal hardware efficiency.

Utilities

The library provides some utility functions besides the above kernels:

• deep_gemm.set_num_sms / get_num_sms: set/get the number of AI cores the kernels may use
• deep_gemm.set_npu_arch / get_npu_arch: override/query the dav-* NPU architecture used for JIT compilation, 0 queries the device
• deep_gemm.set_mk_alignment_for_contiguous_layout / get_mk_alignment_for_contiguous_layout: set/get the group-level M/K alignment for contiguous layout
• deep_gemm.get_theoretical_mk_alignment_for_contiguous_layout: get the theoretical minimum M/K alignment
• deep_gemm.use_deterministic_algorithms / get_deterministic_algorithms: enable/disable deterministic algorithms
• deep_gemm.transform_sf_into_required_layout: transform scaling factors into the required layout
• deep_gemm.transform_k_grouped_sf_into_required_layout: transform K-grouped scaling factors into the required layout
• deep_gemm.get_paged_mqa_logits_metadata: build the scheduling metadata for the paged MQA kernels
• deep_gemm.aclnn_fp8_fp4_gemm_{nt, nn, tn, tt} and deep_gemm.aclnn_bf16_gemm_{nt, nn, tn, tt}: ACLNN reference GEMMs used to cross-check the kernels in tests

Environment Variables

Ascend home path is given by ASCEND_HOME_PATH or ASCEND_TOOLKIT_HOME.

Each DG_JIT_* variable falls back to the corresponding global DJ_JIT_* variable when unset, and is snapshotted when the JIT runtime is first created.

• General
• DG_JIT_DEBUG: 0 or 1, enable JIT debugging features, including compiler commands, load-time reporting, and the selected config per shape; 0 by default
• DG_PRINT_CONFIGS: 0 or 1, print the selected config for each shape, 0 by default
• JIT cache
• DG_JIT_CACHE_DIR: string, cache directory (or a :-separated list of directories) for compiled kernels; lookup searches all paths front-to-back (first hit wins) and a cache miss compiles into the first path, $HOME/.dj by default
• Compiler output and artifacts
• DG_JIT_PRINT_COMPILER_COMMAND: 0 or 1, print compiler and disassembler commands, 0 by default
• DG_JIT_PRINT_LOAD_TIME: 0 or 1, print kernel load time, 0 by default
• DG_JIT_KERNEL_DEBUG_INFO: 0 or 1, add Bisheng kernel debug information, 0 by default
• Debugging
• DG_JIT_LAUNCH_TIMEOUT: integer, Ascend kernel launch timeout in seconds; 300 by default, 0 disables it
• ASCEND_LAUNCH_BLOCKING: synchronize the device after every launch, turning asynchronous launch errors into in-line failures

Performance

Measured on Ascend 950DT (CANN 9.20) using bench_msprof with cold L2. Shapes and cases follow the DeepGEMM test suite and cover inference and training workloads of the DeepSeek model series; see tests/ for reproduction.

Dense GEMM

Dense GEMM reaches up to 99.8% of the hardware limit across types:

| Type | M | N | K | Latency (us) | Computation (TFLOPS) | Hardware limit (TFLOPS) | Utilization (%) | | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | | FP4×FP4 | 4096 | 7168 | 16384 | 565.6 | 1701 | 1730 | 98.3 | | FP8×FP4 | 4096 | 7168 | 16384 | 1117.8 | 861 | 865 | 99.5 | | FP8×FP8 | 4096 | 7168 | 16384 | 1117.6 | 861 | 865 | 99.5 | | BF16×BF16 | 4096 | 7168 | 16384 | 2229.9 | 431 | 432 | 99.8 |

FP8 GEMM for inference:

| M | N | K | Latency (us) | Computation (TFLOPS) | Memory bandwidth (GB/s) | | :---: | :---: | :---: | :---: | :---: | :---: | | 128 | 2112 | 7168 | 15.0 | 258 | 1205 | | 128 | 576 | 7168 | 10.1 | 105 | 564 | | 128 | 24576 | 1536 | 20.6 | 468 | 2068 | | 128 | 32768 | 512 | 10.4 | 415 | 1829 | | 128 | 7168 | 16384 | 54.8 | 549 | 2456 | | 128 | 4096 | 7168 | 19.0 | 397 | 1797 | | 128 | 7168 | 2048 | 10.1 | 373 | 1670 | | 4096 | 2112 | 7168 | 156.9 | 790 | 319 | | 4096 | 576 | 7168 | 54.6 | 619 | 689 | | 4096 | 24576 | 1536 | 362.2 | 854 | 137 | | 4096 | 32768 | 512 | 164.3 | 837 | 129 | | 4096 | 7168 | 16384 | 1117.6 | 861 | 186 | | 4096 | 4096 | 7168 | 282.9 | 850 | 234 | | 4096 | 7168 | 2048 | 143.4 | 838 | 181 |

M-Grouped GEMM

Grouped GEMM for DeepSeek MoE experts (FP8×FP4, BF16 output). #Groups is the number of experts, and M per group the average tokens each receives.

| #Groups | M per group | N | K | Latency (us) | Computation (TFLOPS) | Memory bandwidth (GB/s) | | :---: | :---: | :---: | :---: | :---: | :---: | :---: | | 4 | 8192 | 6144 | 7168 | 3722.4 | 818 | 104 | | 4 | 8192 | 7168 | 3072 | 1778.0 | 856 | 98 | | 4 | 8192 | 4096 | 4096 | 1356.1 | 855 | 148 | | 4 | 8192 | 4096 | 2048 | 680.6 | 852 | 148 | | 8 | 4096 | 6144 | 7168 | 3689.4 | 831 | 136 | | 8 | 4096 | 7168 | 3072 | 1778.9 | 862 | 130 | | 8 | 4096 | 4096 | 4096 | 1356.5 | 861 | 180 | | 8 | 4096 | 4096 | 2048 | 681.0 | 858 | 179 |

MQA Logits

MQA scoring for the DeepSeek Lightning Indexer. The kernel is FIX-pipe bound rather than compute bound, saturating the FIX pipe at 99% utilization.

| Type | Format | #Q Tokens | #K Tokens | #Heads | Dim | Latency (us) | Computation (TFLOPS) | Memory bandwidth (GB/s) | | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | | Prefill | FP8 | 4096 | 8192 | 32 | 128 | 302.6 | 681 | 283 | | Prefill | FP4 | 4096 | 8192 | 32 | 128 | 277.4 | 743 | 277 | | Decode | FP8 | 256 | 8192 | 32 | 128 | 150.9 | 565 | 2006 | | Decode | FP4 | 256 | 8192 | 32 | 128 | 124.2 | 708 | 1349 |

MegaMoE

Mega MoE fuses EP dispatch, two grouped GEMMs, SwiGLU, and combine. Benchmarked over EP8 with top-k=6 (each token routed to 6 experts) and one shared expert; all values are averaged across 8 ranks.

| #Experts | Hidden | Intermediate | Tokens | Latency (us) | Computation (TFLOPS) | Memory bandwidth (GB/s) | Communication bandwidth (GB/s) | | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | | 96 | 5120 | 2304 | 64 | 138.7 | 228.5 | 1850.0 | 37.5 | | 96 | 5120 | 2304 | 256 | 225.5 | 562.5 | 1257.4 | 91.6 | | 96 | 7168 | 3072 | 64 | 222.2 | 266.4 | 2136.9 | 32.8 | | 96 | 7168 | 3072 | 256 | 359.2 | 659.2 | 1425.4 | 80.5 | | 384 | 5120 | 2304 | 4096 | 2567.7 | 790.3 | 567.5 | 128.8 | | 384 | 5120 | 2304 | 16384 | 9768.6 | 831.0 | 325.0 | 135.2 | | 384 | 7168 | 3072 | 4096 | 4600.7 | 823.4 | 531.3 | 100.6 | | 384 | 7168 | 3072 | 16384 | 17904.2 | 846.3 | 269.3 | 103.3 |

HC Prenorm GEMM

Prenorm GEMM for the DeepSeek mHC (Manifold-Constrained Hyper-Connections) module, which nearly saturates the HBM write bandwidth.

| M | N | K | Latency (us) | Computation (TFLOPS) | Memory bandwidth (GB/s) | | :---: | :---: | :---: | :---: | :---: | :---: | | 13 | 24 | 28672 | 6.7 | 3 | 519 | | 137 | 24 | 28672 | 9.6 | 20 | 1110 | | 512 | 24 | 28672 | 15.4 | 46 | 2082 | | 4096 | 24 | 28672 | 72.7 | 78 | 3276 | | 8192 | 24 | 28672 | 136.7 | 82 | 3463 |

Contributors

• Project Leads: Kexing Zhou, Zhean Xu, Chenggang Zhao
• GEMM Kernels: Kexing Zhou, Zhean Xu, Yunfan Xiao, Yuhao Meng
• MQA Logits: Anyi Xu, Zhean Xu, Kaifeng Chen
• mHC Kernel: Yuxuan Zhou, Chenggang Zhao, Ruifan Xu, Huanqi Cao, Chenhao Xu
• MegaMoE: Zhean Xu
• SF Layout Kernels: Guanglin Li
• Infrastructure: Kexing Zhou, Kuai Yu

Acknowledgements

DeepGEMM-Ascend follows the design of the upstream DeepGEMM project, which is inspired by CUTLASS. The JIT runtime is provided by DeepJIT. The mHC kernel is backed by Tilelang. We sincerely thank the developers of these projects for their contributions, and gratefully acknowledge Huawei for its technical support and engineering expertise throughout the development of DeepGEMM-Ascend.

License

This code repository is released under the MIT License.

Citation

If you use DeepGEMM-Ascend in your work, please cite:

@misc{deepgemm_ascend2026,
  title     = {DeepGEMM-Ascend: Clean and Efficient BLAS Kernel Library on Ascend NPU},
  author    = {Kexing Zhou and Zhean Xu and Anyi Xu and Chenggang Zhao and
               Yuxuan Zhou and Yunfan Xiao and Guanglin Li and Kaifeng Chen and Yuhao Meng and
               Huanqi Cao and Ruifan Xu and Chenhao Xu and Kuai Yu},
  year      = {2026},
  publisher = {GitHub},
  url       = {https://github.com/deepseek-ai/DeepGEMM-Ascend}
}