开源项目 GitHub Trending 2026-10-06T12:01 正文 ok 1785 words

DeepSeek 开源 DeepGEMM:简洁高效的 GPU BLAS 算子库,全面支持 FP8/FP4 与 Mega MoE

英文标题:🔥 deepseek-ai/DeepGEMM

DeepSeek 正式开源统一高性能 Tensor Core 算子库 DeepGEMM,汇集了现代大语言模型的核心计算原语,包括 FP8/FP4/BF16 GEMM、通信重叠的融合 MoE(Mega MoE)、针对 Lightning Indexer 的 MQA 打分等。该库基于 C++ 与 CUDA 构建,所有算子均通过 DeepJIT 运行时编译,无需在安装时预编译 CUDA 代码,并已扩展支持 NVIDIA SM90(Hopper)、SM100(Blackwell)及华为昇腾(Ascend)。在实测基准中,DeepGEMM 在 NVIDIA H800 上实现了高达 1550 TFLOPS 的计算性能,在保持代码轻量易读的同时达到或超越专家手工调优水平,为开发者和底层工程优化提供了极具价值的参考实现。
打开原文相关 2 · 机会 2
media49.5 KB

中文

DeepGEMM 是一个统一的高性能 Tensor Core 算子库,将现代大语言模型的关键计算原语——GEMM(FP8、FP4、BF16)、包含通信重叠的融合 MoE(Mega MoE)、用于 Lightning Indexer 的 MQA 打分、HyperConnection (HC) 等——整合到了单一、整洁的 CUDA 代码库中。所有算子均通过 DeepJIT 在运行时进行即时编译,安装期间无需进行 CUDA 编译。

DeepGEMM 借鉴了 CUTLASS 和 CuTe 的部分设计理念,但避免了对其复杂模板和代数系统的过度依赖。该库以简洁为核心设计目标,仅包含有限数量的核心算子函数,使其成为学习 NVIDIA GPU 算子优化技术的清晰且易于上手的优质资源。

尽管设计轻量,DeepGEMM 在各种矩阵形状下的性能依然能够匹敌或超越专家手动调优的库。

更新动态

  • 2026.09.30
    • DeepGEMM Ascend 版本发布!更多详情请查看 DeepGEMM-Ascend
    • 增加了更多优化,包括局部性域(locality domain)特性,更多详情请查看 #462
  • 2026.09.10:Sparse Indexer、Mega Gate、Mega mHC、DeepJIT、MoE 和 Indexer 优化等。
    • 更多详情请查看 #432。
  • 2026.04.16:Mega MoE、FP8xFP4 GEMM、FP4 Indexer、PDL、更快的 JIT 编译等。
    • 更多详情请查看 #304。
    • 有关 Mega MoE 基准测试,请参考 #316。
  • 2025.09.28:DeepGEMM 现已支持用于 DeepSeek v3.2 Lightning Indexer 的打分算子(带权重 ReLU MQA logits)。
    • 更多详情请查看 #200。
  • 2025.07.20:DeepGEMM 现已同时支持 SM90/SM100,并重构为低 CPU 开销的 JIT C++ 模块。
    • 由于 NVCC 12.9 会自动进行 FFMA 交错重叠(interleaving),所有后置优化将不再受支持。
    • 更多详情请查看 #112。
  • 2025.05.14:DeepGEMM 现已支持用于 Dense 和 MoE 反向传播的权重梯度算子!详情请查看 #95。
  • 2025.04.18:DeepGEMM 在 H800 上现已实现高达 1550 TFLOPS 的性能!详情请查看 #74、#78、#81、#86 以及 340d988。

环境要求

  • NVIDIA SM90 或 SM100 架构 GPU
  • Python 3.8 或更高版本
  • 支持 C++20 <format> 的编译器和标准库
  • CUDA Toolkit 12.9 或更高版本
  • PyTorch 2.3 或更高版本
  • CUTLASS 4.0 或更高版本(可通过 Git submodule 克隆)

安装指南

bash

必须克隆子模块

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

链接必要的头文件并构建 C++ 扩展

cat develop.sh ./develop.sh

cat install.sh ./install.sh

随后,只需在 Python 项目中导入 deep_gemm 即可开始使用!

使用说明

该库为 NVIDIA GPU 提供了高度优化的 GEMM 算子,采用命名规范:D = C + A @ B。输入形状布局为 NT(A 不转置,B 转置)。SM90 实现仅支持 NT 内存布局(行主序、列主序),而 SM100 实现则支持所有内存布局(NT、TN、NN、TT)。例如,fp8_gemm_nt 将执行 D = C + A @ B.T。

对于这两种架构,左操作数(LHS)的缩放因子都必须满足 TMA 对齐且为转置布局。此外,SM90 与 SM100 的缩放因子数据格式有所不同:

  • SM90 要求缩放因子采用 FP32 格式。
  • SM100 要求缩放因子采用打包的 UE8M0 格式,将 4 个 UE8M0 打包为一个 torch.int。

请注意,输入转置或 FP8 类型转换等操作必须由用户单独处理,请独立实现或将其融合到前序算子中。虽然该库提供了一些简易的 PyTorch 工具函数,但直接使用它们可能会影响性能;我们的核心重点在于优化 GEMM 算子本身。

普通 Dense GEMM(非分组)

如需执行基础的非分组 FP8 GEMM,可调用 fp8_gemm_{nt, nn, tn, tt} 函数。更多详细信息请参考函数文档。

分组 GEMM(连续布局)

与 CUTLASS 中的传统分组 GEMM 不同,DeepGEMM 仅在 M 轴上进行分组,而 N 轴和 K 轴必须保持固定。这一设计专门针对 MoE 模型中专家具有相同形状的场景。在训练前向传播或推理 Prefill 阶段,每个专家处理的 Token 数量可能有所不同,我们将这些 Token 拼接成单个张量,称为“连续(contiguous)”布局。请注意,每个专家分段必须与 GEMM M 块大小对齐(get_mk_alignment_for_contiguous_layout())。更多信息请参阅 m_grouped_fp8_gemm_{nt, nn}_contiguous 函数文档。

我们还为 MoE 权重反向传播提供了 K 轴分组 API(此时 M 和 N 轴必须保持固定),更多信息请参阅 k_grouped_fp8_gemm_tn_contiguous。

分组 GEMM(掩码布局)

在推理 Decoding 阶段,当启用 CUDA Graph 且 CPU 无法获知每个专家分配到的具体 Token 数量时,我们支持带掩码的分组 GEMM(masked grouped GEMM)。通过传入一个掩码张量,算子仅计算有效部分。

该场景下可使用 m_grouped_fp8_gemm_nt_masked 并查阅相关文档。一个典型用例是将来自 DeepEP 的低延迟算子输出作为该算子的输入。

用于 Indexer 的 V3.2 MQA 算子

该算子家族有两个版本:非分页(non-paged,用于 Prefill)和分页(paged,用于 Decoding)。以非分页版本 fp8_fp4_mqa_logits 为例,其主要输入包括:

  • q:一个 (q_data, q_sf) 元组;SM100 接受 MXFP4/MXFP8 数据以及打包的 UE8M0 缩放因子
  • kv:一个 (kv_data, kv_sf) 元组,逻辑形状为 [seq_len_kv, head_dim]
  • weights:形状为 [seq_len, num_heads] 的张量(在 SM100 上为 BF16)
  • cu_seq_len_k_start 和 cu_seq_len_k_end:形状为 [seq_len] 的整型张量
  • max_seqlen_k:任意 Query 行的最大有效 KV 跨度

输出被压缩为 [seq_len, max_seqlen_k];第 i 行从第 0 列开始存储其有效 KV 跨度。对于 q 中的每个 Token i,它将遍历 [cu_seq_len_k_start[i], cu_seq_len_k_end[i]) 区间内的所有 Token j,并按如下公式计算对应的压缩 Logit:

python kv_j = kv[0][j, :] * kv[1][j].unsqueeze(1) # [head_dim] out_ij = q[i, :, :] @ kv_j # [num_heads] out_ij = out_ij.relu() * weights[i, :] # [num_heads] out_ij = out_ij.sum() # Scalar

如需了解更多详情及分页版本 fp8_fp4_paged_mqa_logits,请参阅 tests/test_attention.py。

Mega MoE

Mega MoE 将 EP Dispatch、Linear 1 与 Linear 2(FP8xFP4 或 FP8xFP8)、SwiGLU 以及 EP Combine 融合重叠为一个单一的巨型算子(mega-kernel),实现了 NVLink 通信与 Tensor Core 计算的高度重叠。它需要配合对称内存(symmetric memory)进行多进程启动。使用示例如下:

python

分配对称内存缓冲区

注意:需要 PyTorch >= 2.9

buffer = deep_gemm.get_symm_buffer_for_mega_moe( group, num_experts, num_max_tokens_per_rank, num_topk, hidden, intermediate_hidden, mma_type='fp8xfp4', # 对于 FP8 路由专家权重,请使用 'fp8xfp8' )

将权重(FP4 或带有 UE8M0 SF 的 FP8)转换为所需布局

transformed_l1, transformed_l2 = deep_gemm.transform_weights_for_mega_moe(l1_weights, l2_weights)

(可选)将权重局部化到局部性域中

transformed_l1 = (deep_gemm.localize(transformed_l1[0]), transformed_l1[1]) transformed_l2 = (deep_gemm.localize(transformed_l2[0]), transformed_l2[1]) deep_gemm.destroy_localizer()

在每次调用前将输入复制到缓冲区中

您也可以将这些操作融合到前序算子中

buffer.x[:num_tokens].copy_(x_fp8) buffer.x_sf[:num_tokens].copy_(x_sf) buffer.topk_idx[:num_tokens].copy_(topk_idx) buffer.topk_weights[:num_tokens].copy_(topk_weights)

运行融合 Mega MoE 算子

y = torch.empty((num_tokens, hidden), dtype=torch.bfloat16, device='cuda') deep_gemm.fp8_fp4_mega_moe(y, transformed_l1, transformed_l2, buffer)

完整的包含多进程设置和基准测试的示例,请参阅 tests/test_mega_moe.py。

工具函数

除上述算子外,该库还提供了一些实用工具函数:

  • deep_gemm.set_num_sms / get_num_sms:设置/获取使用的最大 SM 数量
  • deep_gemm.set_tc_util / get_tc_util:设置/获取近似的 Tensor Core 利用率比例
  • deep_gemm.set_pdl / get_pdl:启用/禁用编程依赖启动(Programmatic Dependent Launch, PDL)
  • deep_gemm.use_deterministic_algorithms:启用/禁用确定性算法
  • deep_gemm.set_mk_alignment_for_contiguous_layout / get_mk_alignment_for_contiguous_layout:设置/获取连续布局的分组级 M/K 对齐大小
  • deep_gemm.get_theoretical_mk_alignment_for_contiguous_layout:获取理论上的最小 M/K 对齐大小
  • deep_gemm.set_ignore_compile_dims:配置 JIT 编译期间忽略的维度
  • deep_gemm.set_block_size_multiple_of:约束 Block 大小为给定值的倍数
  • deep_gemm.transform_sf_into_required_layout:将缩放因子转换为所需布局
  • deep_gemm.get_tma_aligned_size:获取所需的 TMA 对齐大小
  • deep_gemm.get_mn_major_tma_aligned_tensor:获取 MN 主序的 TMA 对齐张量
  • deep_gemm.get_mn_major_tma_aligned_packed_ue8m0_tensor:获取 MN 主序的 TMA 对齐张量(将 FP32 打包为 UE8M0)
  • deep_gemm.get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor:K 分组 GEMM 打包算子

环境变量

该库还提供了若干环境变量,在特定场景下非常实用:

未设置时,各个 DG_JIT_* 变量将默认回退到对应的全局 DJ_JIT_* 变量。

  • 通用设置
    • DG_JIT_DEBUG:0 或 1,启用 JIT 调试功能,包括编译器命令和 PTXAS 输出、加载时间报告、行号信息以及 PTX/SASS 转储;默认为 0
    • DG_PRINT_CONFIGS:0 或 1,打印每个形状选定的配置参数,默认为 0
  • JIT 缓存
    • DG_JIT_CACHE_DIR:字符串,编译算子的缓存目录(或以 : 分隔的目录列表);查找时从前到后搜索所有路径(首次命中优先),未命中则编译至首个路径,默认为 $HOME/.dj
  • 编译器选择
    • DG_JIT_NVCC_COMPILER:字符串,NVCC 编译器路径;若未指定,CUDA 将依次通过 CUDA_HOME、CUDA_PATH、which nvcc 以及 /usr/local/cuda 自动定位
    • DG_JIT_CPP_STANDARD:整数,C++ 标准版本,默认为 20
  • 编译器输出
    • DG_JIT_PRINT_COMPILER_COMMAND:0 或 1,打印编译命令,默认为 0
    • DG_JIT_PTXAS_VERBOSE:0 或 1,显示详细的 PTXAS 输出,默认为 0
    • DG_JIT_CHECK_NO_SPILLS:0 或 1,断言编译后算子中无寄存器溢出(spill),默认为 0
    • DG_JIT_CHECK_NO_LOCAL_MEMORY:0 或 1,断言编译后算子中无局部内存使用,默认为 0
    • DG_JIT_PRINT_LOAD_TIME:0 或 1,打印算子加载时间,默认为 0
  • 调试与性能分析
    • DG_JIT_WITH_LINEINFO:0 或 1,为性能分析工具嵌入源码行信息,默认为 0
    • DG_JIT_DUMP_ASM:0 或 1,同时转储 PTX 和 SASS,默认为 0
    • DG_JIT_DUMP_PTX:0 或 1,转储 PTX 输出,默认为 0
    • DG_JIT_DUMP_SASS:0 或 1,转储 SASS 输出,默认为 0
    • DG_COMM_KERNEL_DEBUG:0 或 1,在每次 Mega MoE 调用前将对称内存缓冲区清零以供调试,默认为 0
    • DG_USE_NVIDIA_TOOLS:0 或 1,在外部 NVIDIA 工具下运行时跳过内部 Profiling,默认为 0
  • 构建选项
    • DG_SKIP_CUDA_BUILD:0 或 1,在安装期间跳过 CUDA 扩展构建,默认为 0
    • DG_FORCE_BUILD:0 或 1,强制本地构建而非下载预构建 Wheel 包,默认为 0

更多示例和详细信息,请参考测试代码或查阅对应的 Python 文档。

致谢

DeepGEMM 受 CUTLASS 项目启发而开发。向相关开发者致以诚挚的感谢与敬意!

许可证与引用

本代码仓库采用 MIT 许可证 开源。

bibtex @misc{deepgemm2025, title={DeepGEMM: clean and efficient BLAS kernel library on GPU}, author={Chenggang Zhao and Zhean Xu and Liang Zhao and Jiashi Li and Chenhao Xu and Anyi Xu and Shengyu Liu and Kexing Zhou and Kuai Yu}, year={2025}, publisher = {GitHub}, howpublished = {\url{https://github.com/deepseek-ai/DeepGEMM}}, }

English / 原文

DeepGEMM is a unified, high-performance tensor core kernel library that brings together the key computation primitives of modern large language models — GEMMs (FP8, FP4, BF16), fused MoE with overlapped communication (Mega MoE), MQA scoring for the lightning indexer, HyperConnection (HC), and more — into a single, cohesive CUDA codebase. All kernels are compiled at runtime through DeepJIT, requiring no CUDA compilation during installation.

DeepGEMM leverages some concepts from CUTLASS and CuTe, but avoids heavy reliance on their templates or algebras. The library is designed for simplicity, with only a limited number of core kernel functions, making it a clean and accessible resource for learning NVIDIA GPU kernel optimization techniques.

Despite its lightweight design, DeepGEMM's performance matches or exceeds expert-tuned libraries across various matrix shapes.

  • 2026.09.30

    • DeepGEMM Ascend is available! Check DeepGEMM-Ascend for more details
    • Add more optimizations, including locality domain features, check #462 for more details
  • 2026.09.10: Sparse Indexer, Mega Gate, Mega mHC, DeepJIT, MoE and Indexer optimizations and more.

    • Please see #432 for more details.
  • 2026.04.16: Mega MoE, FP8xFP4 GEMM, FP4 Indexer, PDL, faster JIT compilation and more.

    • Please see #304 for more details.
    • For Mega MoE benchmarks, refer to #316.
  • 2025.09.28: DeepGEMM now supports scoring kernels (weighted ReLU MQA logits) for the lightning indexer for DeepSeek v3.2.

    • Please see #200 for more details.
  • 2025.07.20: DeepGEMM now supports both SM90/SM100, and has a full refactor with a low-CPU-overhead JIT CPP module.

    • As NVCC 12.9 will automatically do the FFMA interleaving, all post optimizations will be no longer supported.
    • Please see #112 for more details.
  • 2025.05.14: DeepGEMM now offers weight gradient kernels for dense and MoE backward! See #95 for details.

  • 2025.04.18: DeepGEMM now achieves up to 1550 TFLOPS on H800! See #74, #78, #81, #86 and 340d988 for details.

  • NVIDIA SM90 or SM100 architecture GPU

  • Python 3.8 or higher

  • Compilers and standard libraries with C++20 <format> support

  • CUDA Toolkit 12.9 or higher

  • PyTorch 2.3 or higher

  • CUTLASS 4.0 or higher (could be cloned by Git submodule)

    Submodule must be cloned

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

    Link some essential includes and build the C++ extension

    cat develop.sh ./develop.sh

    cat install.sh ./install.sh

Then, import deep_gemm in your Python project, and enjoy!

This library provides optimized GEMM kernels for NVIDIA GPUs with a naming convention: D = C + A @ B. The input shape layout is NT (non-transposed A, transposed B). While the SM90 implementation supports only the NT memory layout (row-major, col-major), the SM100 implementation supports all memory layouts (NT, TN, NN, TT). For example, fp8_gemm_nt will do a D = C + A @ B.T

For both architectures, the LHS scaling factor is required to have a TMA-aligned and transposed layout. And the data format for the scaling factor of SM90 and SM100 is different:

  • SM90 requires scaling factors in FP32 format.
  • SM100 requires scaling factors in packed UE8M0 format, which packs 4 UE8M0 into a single torch.int.

Please note that operations like input transposition or FP8 casting must be handled separately by the user, please implement or fuse them into prior kernels independently. While the library provides some simple PyTorch utility functions, these may result in slower performance, but our primary focus is on optimizing the GEMM kernels themselves.

Normal dense GEMMs (non-grouped)

To perform a basic non-grouped FP8 GEMM, call the fp8_gemm_{nt, nn, tn, tt} function. For more details, please refer to the function documentation.

Grouped GEMMs (contiguous layout)

Unlike traditional grouped GEMMs in CUTLASS, DeepGEMM groups only the M-axis, while N and K must remain fixed. This design is tailored for scenarios where experts in an MoE model share the same shape. For training forward passes or inference prefilling, where each expert may process a varying number of tokens, we concatenate these tokens into a single tensor, referred to as the "contiguous" layout. Note that each expert segment must be aligned to the GEMM M block size (get_mk_alignment_for_contiguous_layout()). For more information, please refer to the m_grouped_fp8_gemm_{nt, nn}_contiguous function documentation.

We also provide a K-axis-grouped API for MoE weight backward (with M and N must remain fixed), please refer to k_grouped_fp8_gemm_tn_contiguous for more information.

Grouped GEMMs (masked layout)

During the inference decoding phase, when CUDA graph is enabled and the CPU is unaware of the number of tokens each expert receives, we support masked grouped GEMMs. By providing a mask tensor, the kernel computes only the valid portions.

Use m_grouped_fp8_gemm_nt_masked for this purpose and consult the relevant documentation. An example usage is to use the output of low-latency kernels from DeepEP as input.

V3.2 MQA kernels for the indexer

The kernel family has two versions, non-paged (for prefilling) and paged (for decoding). Take the non-paged version fp8_fp4_mqa_logits as an example. Its main inputs are:

  • q, a (q_data, q_sf) tuple; SM100 accepts MXFP4/MXFP8 data with packed UE8M0 scales
  • kv, a (kv_data, kv_sf) tuple with shape [seq_len_kv, head_dim] logically
  • weights, tensor with shape [seq_len, num_heads] (BF16 on SM100)
  • cu_seq_len_k_start and cu_seq_len_k_end, int tensor with shape [seq_len]
  • max_seqlen_k, the maximum valid KV span of any query row

The output is compressed to [seq_len, max_seqlen_k]; row i stores its valid KV span starting at column zero. For each token i in q, it will iterate all tokens j from [cu_seq_len_k_start[i], cu_seq_len_k_end[i]), and calculate the corresponding compressed logit as:

kv_j = kv[0][j, :] * kv[1][j].unsqueeze(1)  # [head_dim]
out_ij = q[i, :, :] @ kv_j  # [num_heads]
out_ij = out_ij.relu() * weights[i, :]  # [num_heads]
out_ij = out_ij.sum()  # Scalar

For more details and the paged version fp8_fp4_paged_mqa_logits, please refer to tests/test_attention.py.

Mega MoE fuses and overlaps EP dispatch, linear 1 and linear 2 (FP8xFP4 or FP8xFP8), SwiGLU, and EP combine into a single mega-kernel, overlapping NVLink communication and tensor core computation. It requires multi-process launch with symmetric memory. Usage:

# Allocate symmetric memory buffer
# NOTES: requires PyTorch >= 2.9
buffer = deep_gemm.get_symm_buffer_for_mega_moe(
    group, num_experts, num_max_tokens_per_rank, num_topk, hidden, intermediate_hidden,
    mma_type='fp8xfp4',  # Use 'fp8xfp8' for FP8 routed-expert weights
)

# Transform weights (FP4 or FP8 with UE8M0 SF) into the required layout
transformed_l1, transformed_l2 = deep_gemm.transform_weights_for_mega_moe(l1_weights, l2_weights)

# (Optional) Localize weights into locality domains
transformed_l1 = (deep_gemm.localize(transformed_l1[0]), transformed_l1[1])
transformed_l2 = (deep_gemm.localize(transformed_l2[0]), transformed_l2[1])
deep_gemm.destroy_localizer()

# Copy inputs into the buffer before each call
# You may fuse these into previous kernels
buffer.x[:num_tokens].copy_(x_fp8)
buffer.x_sf[:num_tokens].copy_(x_sf)
buffer.topk_idx[:num_tokens].copy_(topk_idx)
buffer.topk_weights[:num_tokens].copy_(topk_weights)

# Run the fused mega MoE kernel
y = torch.empty((num_tokens, hidden), dtype=torch.bfloat16, device='cuda')
deep_gemm.fp8_fp4_mega_moe(y, transformed_l1, transformed_l2, buffer)

For the full example with multi-process setup and benchmarking, please refer to tests/test_mega_moe.py.

The library provides some utility functions besides the above kernels:

  • deep_gemm.set_num_sms / get_num_sms: set/get the maximum SM count to use
  • deep_gemm.set_tc_util / get_tc_util: set/get an approximated tensor core utilization ratio
  • deep_gemm.set_pdl / get_pdl: enable/disable Programmatic Dependent Launch (PDL)
  • deep_gemm.use_deterministic_algorithms: enable/disable deterministic algorithms
  • 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.set_ignore_compile_dims: configure dimensions to ignore during JIT compilation
  • deep_gemm.set_block_size_multiple_of: constrain block sizes to be multiples of a given value
  • deep_gemm.transform_sf_into_required_layout: transform scaling factors into the required layout
  • deep_gemm.get_tma_aligned_size: get the required TMA alignment size
  • deep_gemm.get_mn_major_tma_aligned_tensor: get a MN-major TMA-aligned tensor
  • deep_gemm.get_mn_major_tma_aligned_packed_ue8m0_tensor: get a MN-major TMA-aligned tensor (with packing FP32 into UE8M0)
  • deep_gemm.get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor: K-grouped GEMM packing kernel

The library also provides some environment variables, which may be useful:

Each DG_JIT_* variable falls back to the corresponding global DJ_JIT_* variable when unset.

  • General
    • DG_JIT_DEBUG: 0 or 1, enable JIT debugging features, including compiler command and PTXAS output, load-time reporting, line info, and PTX/SASS dumps; 0 by default
    • DG_PRINT_CONFIGS: 0 or 1, print selected configs 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 selection
    • DG_JIT_NVCC_COMPILER: string, NVCC compiler path; otherwise CUDA is found through CUDA_HOME, CUDA_PATH, which nvcc, then /usr/local/cuda
    • DG_JIT_CPP_STANDARD: integer, C++ standard version, 20 by default
  • Compiler output
    • DG_JIT_PRINT_COMPILER_COMMAND: 0 or 1, print compilation commands, 0 by default
    • DG_JIT_PTXAS_VERBOSE: 0 or 1, show detailed PTXAS output, 0 by default
    • DG_JIT_CHECK_NO_SPILLS: 0 or 1, assert no register spills in compiled kernels, 0 by default
    • DG_JIT_CHECK_NO_LOCAL_MEMORY: 0 or 1, assert no local memory usage in compiled kernels, 0 by default
    • DG_JIT_PRINT_LOAD_TIME: 0 or 1, print kernel load time, 0 by default
  • Debug and profiling
    • DG_JIT_WITH_LINEINFO: 0 or 1, embed source line info for profiling tools, 0 by default
    • DG_JIT_DUMP_ASM: 0 or 1, dump both PTX and SASS, 0 by default
    • DG_JIT_DUMP_PTX: 0 or 1, dump PTX output, 0 by default
    • DG_JIT_DUMP_SASS: 0 or 1, dump SASS output, 0 by default
    • DG_COMM_KERNEL_DEBUG: 0 or 1, zero symmetric buffer before each Mega MoE call for debugging, 0 by default
    • DG_USE_NVIDIA_TOOLS: 0 or 1, skip internal profiling when running under external NVIDIA tools, 0 by default
  • Build options
    • DG_SKIP_CUDA_BUILD: 0 or 1, skip CUDA extension build during installation, 0 by default
    • DG_FORCE_BUILD: 0 or 1, force local build instead of downloading pre-built wheels, 0 by default

For additional examples and details, please refer to the test code or review the corresponding Python documentation.

DeepGEMM is inspired by the CUTLASS project. Thanks and respect to the developers!

This code repository is released under the MIT License.

@misc{deepgemm2025,
      title={DeepGEMM: clean and efficient BLAS kernel library on GPU}, 
      author={Chenggang Zhao and Zhean Xu and Liang Zhao and Jiashi Li and Chenhao Xu and Anyi Xu and Shengyu Liu and Kexing Zhou and Kuai Yu},
      year={2025},
      publisher = {GitHub},
      howpublished = {\url{https://github.com/deepseek-ai/DeepGEMM}},
}