DeepGEMM 简介
DeepGEMM 是一个用于 FP8 矩阵乘法（GEMM）的库，旨在实现简洁高效的计算，支持细粒度缩放以减轻特征异常值引起的量化误差，相关技术在 DeepSeek-V3 中被提出，如图7(a)所示。它支持普通矩阵乘法和MoE分组矩阵乘法。该库使用 CUDA 编写，安装时无需编译，所有Kernels在运行时通过轻量级的即时编译（JIT）模块动态编译。
目前，DeepGEMM 仅支持 NVIDIA Hopper 架构 tensor cores。为了应对 FP8 Tensor Cores 累加的不精确性，它采用了 CUDA Core 的两级累加技术来提高累加精度。虽然它借鉴了 CUTLASS 和 CuTe 的一些概念，但并未过度依赖它们的模板或代数结构。相反，该库的设计注重简洁性，其核心 Kernel 函数仅包含大约 300 行代码。这使得它成为学习 Hopper FP8 矩阵乘法和优化技术的清晰且易于理解的资源。尽管设计轻量，DeepGEMM 在各种矩阵形状上的性能均能达到或超过经过专家调优的库。
CUTLASS、CuTe、DeepGEMM 的区别
- CUTLASS：功能强大但代码复杂，依赖大量模板和预编译，适合通用场景。
- CuTe：专注于张量操作的抽象，灵活但需要较深理解。
- DeepGEMM：专注于 FP8 和 Hopper 架构，代码简洁，易于学习和修改，适合特定需求（如： DeepSeek-V3 的 MoE 模型）。
文章较长，建议先点赞收藏，再慢慢观看。另外，我撰写的大模型相关的博客及配套代码均整理放置在Github：llm-action，有需要的朋友自取。
补充：前置知识
GEMM
GEMM（General Matrix Multiplications，通用矩阵乘法）是将两个矩阵的进行相乘的计算。常用于深度学习中的GEMM有：
- 全连接层：直接使用GEMM。
- 卷积层：通过Im2col将卷积展开为矩阵乘法。
许多库都实现并优化了GEMM。
- CPU：OpenBLAS、Intel MKL、BLIS。
- GPU：cuBLAS（NVIDIA）、rocBLAS（AMD）
CPU需优化缓存和SIMD，GPU需最大化线程并行和共享内存使用。
在GPU中，GEMM 定义为运算 C=αAB+βC，其中 A、B、C 是矩阵，α 和 β 是标量。对于普通矩阵乘，α 等于 1 且 β 等于 0 。例如，在全连接层的前向传播中，权重矩阵为参数 A，传入激活为参数 B，α 和 β 通常分别为 1 和 0。在某些情况下，β 可以是 1。
GPU 通过将输出矩阵划分为图块来实现 GEMM，然后将其分配给线程块。图块大小（Tile Size）通常是指这些图块的尺寸。每个线程块通过单步执行图块中的 K 维度，从 A 和 B 矩阵加载所需的值，然后将它们相乘并累加到输出中来计算其输出图块。
系统内存架构
整个系统内存架构包括CPU/GPU等存储。
系统存储：
- L1/L2/L3：多级缓存，位置在 CPU 芯片内部；
- System DRAM：动态 RAM，CPU 芯片外部内存，如内存条。
- Disk/Buffer：外部存储，如磁盘或者固态硬盘。
GPU设备存储：
- L1/L2 cache：多级缓存，位置在 GPU 芯片内部；
- GPU DRAM：通常所指的显存；
传输通道：
- PCIE BUS：PCIE标准的数据通道，数据就是通过该通道从显卡到达主机；
- BUS：总线。计算机内部各个存储之间交互数据的通道；
- PCIE-to-PCIE：显卡之间通过PCIE直接传输数据；
- NVLINK：用于显卡之间的专用的超高速数据传输通道。
对于磁盘/硬盘（Disk/SSD）上面的数据传入到 GPU 的内存通常需要经过：硬盘 -> 系统内存 -> GPU 内存的过程。这个速度非常慢，要极力避免这种传输。
GPU 内部内存结构
GPU 内存可以分为局部内存（local memory）、全局内存（global memory）、常量内存（constant memory）、共享内存（shared memory）、寄存器（register）、L1/L2 缓存等。其中，全局内存、局部内存、常量内存都是片下内存，储存在 HBM 上。而所有的 SM 共享 L2 缓存。
CUDA Core
CUDA Core（Compute Unified Device Architecture Core）是英伟达 GPU 中的基础计算单元，专为并行计算任务设计。类似于 CPU 中的 Core，但CUDA Core针对大规模并行运算优化，能够同时处理成百上千个线程，适用于科学计算、深度学习、图形渲染等场景。
每个 CUDA Core 都能够执行计算，每个 CUDA Core 每个时钟周期可以执行一个操作。虽然其能力不如 CPU Core，但当多个 CUDA Core一起用于深度学习时，它们可以通过并行执行进程来加速计算。
CUDA Core 可以处理各种精度的运算，如 FP64、FP32 和 INT32 等，但它在执行深度学习中最常见的操作，如卷积（Conv）和矩阵乘法（GEMM）时仍然面临效率上的挑战。 首先，CUDA Core 在执行这些操作时，需要将数据在寄存器、算术逻辑单元（ALU）之间进行多次搬运，这种过程既耗时又低效。其次，由于每个 CUDA Core 在单个时钟周期只能进行一次计算，因此，受限于 CUDA Core 性能的 GPU 也受到可用 CUDA Core 数量和每个Core时钟速度的限制，这些因素共同限制了深度学习计算性能的提升。
Tensor Core
NVIDIA GPU 从 Volta 架构开始支持 Tensor Core，专为 AI 训练和推理设计的可编程矩阵乘法和累加单元。
相较于 CUDA Core，Tensor Core 能够在每个时钟周期内执行更多的运算，特别是它可以高效地完成矩阵乘法和累加操作，这两种操作是深度学习中最频繁且计算密集的任务。
通过利用 Tensor Core，V100 能够为 AI 训练和推理提供高达 125 Tensor TFLOPS 的算力。这种强大的性能，使得 V100 在处理深度学习任务时，相比于仅使用 CUDA Core 的早期架构，能够实现显著的加速。
此外，Tensor Core 和 CUDA Core 支持的数据类型也不同。下图展示了不同GPU架构 Tensor Core 和 CUDA Core 支持的数据类型。
可以看出，V100 是 Volta 架构，其 Tensor Core 只支持 FP16。
通常 Tensor Core 整体的算力一般是 CUDA Core 算力的几倍。而 GEMM 的实现效率与 Tensor Core 结构和数据格式密切相关，受数据的调度方式影响很大。比如：在 V100 上使用 INT8，只能运行在 CUDA Core 上，其性能可能反而比Tensor Core 使用 FP16 还差不少。同样，H100 的 Tensor Core 也不再支持 INT4 计算。 因此，基于Tensor Core的硬件架构进行计算优化就显得十分重要。好的优化往往能取得数倍的性能提升。
Tensor Core 采用融合乘法加法（FMA）的方式高效处理计算任务。每个 Tensor Core 每周期能执行 4x4x4 GEMM，即 64 个浮点乘法累加（FMA）运算。
在执行运算 D = A×B + C 时，其中 A、B、C 和 D 是 4×4 矩阵。矩阵乘法输入 A 和 B 是 FP16 精度，而累加矩阵 C 和 D 可以是 FP16 或 FP32 精度。通过将矩阵乘法的输入限定为 FP16 精度，减少计算资源和内存带宽需求，同时允许累加矩阵和输出矩阵使用 FP32 精度，保证运算结果的准确性和数值稳定性，从而实现了底层硬件上的混合精度计算。
Tensor Core 的高效性源于其专门的硬件架构，能够在一个周期内并行执行多个乘法和加法操作。这种架构使得 Tensor Core 能够快速处理大规模的矩阵运算，尤其适用于深度学习中的神经网络训练和推理过程。同时，Tensor Core 支持混合精度计算，即在矩阵乘法中使用不同的数据精度。例如，在执行 4x4x4 GEMM 时，输入矩阵 A 和 B 可以是 FP16（半精度浮点数）格式，而累加操作和结果矩阵 C 可以是 FP32（单精度浮点数）格式。这种方式既减少了计算资源和内存带宽的需求，又保证了运算结果的准确性和数值稳定性。
Hopper 架构主要特性
支持 Tensor Cores FP8
FP8 Tensor Cores 支持 FP32 和 FP16 累加器，以及两种新的 FP8 输入类型：
- E4M3 具有 4 个指数位、3 个尾数位和 1 个符号位
- E5M2，具有 5 个指数位、2 个尾数位和 1 个符号位
Thread Block Cluster
在Hopper架构之前，软件层级包含grid、block、thread三个层次。
为什么要划分Block？一是要实现大规模的并行处理，还有就是要利用“局部性”。一个Block中的所有线程会被调度到同一个SM中，这样，Block中的线程数据交换和通信就会高效得多。如下图，在图像处理kernel中，一整个图片被切分成相同大小的数据，分配到不同的SM中处理。
在Hopper架构中，新增了一个层次结构——Thread Block Cluster，可以把它理解为“Block of Blocks”。这给我们带来了一种新的编程视角： 我们仍然将每个Block视为相互独立的，但是我们有了一种方法可以定位grid中的一个子集，从而提供了更多的可编程性和加速的潜力。在引入Thread Block Cluster后，相当于新增了一层局部性的抽象，同样的一个kernel可以进行如下图所示的数据划分。
从硬件角度来看，同一个Thread Block Cluster中的Block被调度在同一个GPC（GPU processing cluster）上。GPC由一系列物理位置上相邻的SMs组成，运行在同一个GPC上的Blocks之间可以享有更快的local synchronization和更快的memory sharing。
一个Thread Block Cluster可以有最多16个Block，这些Block保证在同一时间被调度到同一个GPC的不同SM上运行。
Thread Block Cluster映射到硬件上的表示如下：
Cluster Distributed Shared Memory(DSMEM)
在同一个Thread Block Cluster中的Block可以直接读写其余Blocks的shared memory，这就是Cluster Distributed Shared Memory。需要说明的是，DSMEM并不是一个新增的内存空间，只是一个cluster中所有Block的shared memory的集合。
Asynchronous transaction barrier
异步屏障（Asynchronous barriers）是在Ampere架构中被引入的特性。如下图左侧，一组threads正在生产数据（async mem copies or data exchanges），这些数据将在barrier后面的代码中被使用（Consume Data）。异步屏障将同步过程分为两个步骤：
-
当某些线程完成自己所负责的共享数据的工作后，会将自己的状态标识为Arrive，表明自己完成生产数据。Arrive 是非阻塞的，因此这些线程可以自由地执行其他独立任务。
-
所有线程的后续执行需要共享数据全部完成，因此，当这些线程执行完自己的独立任务后，如果仍有线程未将自己的状态标识为Arrive，那么这些线程会触发Wait。Wait是阻塞的，直到所有线程都Arrive后，才可以继续执行Consume Data的动作。
异步屏障的优点就是提前Arrive的线程可以去做别的任务，这就带来了更高的并发性。 如果所有线程都有足够的独立任务去完成，那么所有线程的Wait动作都可以立即结束，因为所有线程都已经Arrive。
让线程不要白白等待，去做一些别的工作，这种想法其实是很自然的。Hopper架构保留了这一特性，并且增添了一个新的Feature，即 Asynchronous transaction barrier（如上图右）。
利用 Asynchronous transaction barrier 可以实现 ONE-SIDED MEMORY COPIE。在这种情况下，Barrier 知道自己期待收到多少数据，而每次的 asynchronous mem copy 都携带着自身拷贝的 transaction count。当数据到达后，更新 barrier 所接收到的总 transaction count 即可。这种方法比Amper架构下的{data+fance+flag}的handshake方式快7倍。
由此可见，Asynchronous transaction barrier不仅仅提供了一种线程同步的机制（等待其他线程到达arrive），还提供了“等待数据到达”的机制。利用这种特性，可以在Hopper GPU中实现纯粹的单向数据移动，等待数据到达的sleep线程只是纯粹的等待数据，不需要考虑数据从哪里来，被谁拷贝过来，也就是说，不需要一个类似”握手“的动作，从而提高拷贝的效率。这种机制是非常有价值的，考虑Hooper架构下的一个场景，如前文所说Hopper架构支持Cluster DSMEM，如果某个Block中的线程等待其他Blocks中的数据拷贝，那么这种ONE-SIDED MEMORY COPIES将会免于昂贵的跨SM的线程同步。
当然，Asynchronous transaction barrier并不止步于此。还可以利用这种ONE-SIDE的机制去传递一些messege或者operation。比如，transaction count为1，可以在某个Block中给同一个Cluster中的另一个Block发送一个Operation，内容是（add 100或者取max），而无需考虑实际操作的数据是什么。
看到这你可能会想，这种机制也挺好想出来的，为什么Amper架构不支持呢？其实，Hopper架构新增的TENSOR MEMORY ACCELERATOR UNIT (TMA) 硬件才使这种机制变成“更高效的机制”。而且，白皮书里说TMA的提出是为了更好的feed数据给更快速的第四代Tensor core，因此，猜想可能Ampere架构下的内存拷贝机制已经满足Ampere架构的算力需求了。
Tensor Memory Access（TMA）
异步拷贝的思想在cuda中其实很早就有对应的实现了。通过cudaMemcpyAsync，可以实现数据从CPU memory到 GPU global memory的搬运操作和cuda kernel执行操作的重叠。
在Amper架构之前，从global memory到shared memory之间的异步拷贝可以通过CudaDMA实现，但是，CudaDMA需要占用额外的thread warp来引导数据搬运。
在Amper架构之后，可以使用 cuda::memcpy_async 在 GPU global memory和shared memory之间异步拷贝数据，而无需绑定线程来引导数据移动。
这一点变化在数据流动的角度来说，在cuda::memcpy_async出现之前，搬运到shared memory的数据流动应该是GMEM→L2 Cache→L1 Cache→thread registers→shared memory。
而在cuda::memcpy_async出现之后则缩减为GMEM→L2 Cache→L1 Cache→shared memory，显而易见，数据搬运效率得到了提升。
但是，Amper架构和Hopper架构对于Asynchronous Copy的支持也是不同的。
Hopper架构新提供了一个名为Tensor Memory Accelerator (TMA)的硬件，专门用来处理global memory和shared memory之间的异步搬运（同时支持Cluster中DSMEMs之间的数据搬运）。TMA最核心的优点就是在执行异步拷贝时可以不给线程带来任何额外的负担。
在 A100上，异步内存复制是使用特殊的 LoadGlobalStoreShared 指令执行的，因此，线程负责在整个复制区域中循环生成所有拷贝地址。
在 Hopper GPU上，TMA 负责处理一切。单个线程在启动 TMA 之前创建一个copy descriptor（该descriptor使用张量维度和块坐标指定数据传输，而不是每个元素寻址。可以指定高达共享内存容量的大型数据块，并将其从全局内存加载到共享内存中，或从共享内存存储回全局内存），从那时起地址生成和数据移动由TMA硬件处理。
需要注意的是，TMA位于SM中，但是独立于SM中的线程运行，也就是说并不是所有的线程都要与TMA交互，一个 warp 中只需要选择单个thread发出异步 TMA 操作(cuda::memcpy_async)。
TMA 支持不同的张量布局（1D-5D 张量）、不同的内存访问模式、reductions和其他功能，从而显著降低了寻址开销并提高了效率。
DeepGEMM 主要特性及优化
DeepGEMM主要特性如下：
- 对于 FP8 的优化：DeepGEMM采用了 CUDA Core 两级累加。（换句话说这也是NV TensorCore的设计不足之处）FP8 是一种低比特浮点格式，能够在保持一定计算精度的同时大幅提升计算效率。FP8在大量累加时会累积出现随机误差。例如：FP8 GEMM在英伟达H800 GPU上的累加精度保留14位左右，明显低于FP32累加精度。以K= 4096的两个随机矩阵的GEMM 运算为例，Tensor Core中的有限累加精度可导致最大相对误差接近 2%。尽管存在这些问题，有限的累加精度在一些FP8框架中仍然是默认选项（TransformerEngine），这严重限制了训练的准确性。DeepSeek将中间结果的计算精度提升为FP32（32位浮点），实行高精度累加，然后再转换回 FP8，以降低大量微小误差累加带来的训练偏差。即在Tensor Core上执行矩阵乘累加（MMA）时，中间结果会使用有限的位宽进行累加。一旦达到 个元素的间隔，这些部分结果就会被复制到CUDA Core 上的FP32寄存器中，在那里进行全精度的FP32累加，如图7(b)所示。
- 支持分组GEMM：与 CUTLASS 中传统的分组 GEMM 不同，DeepGEMM 仅对 M 轴进行分组，而 N 和 K 可保持不变。（专门针对 MoE 模型中的专家量身定制）
- 即时编译（Just-In-Time, JIT）：通过 JIT 技术，代码可以在运行时动态生成和优化，进一步提升性能和灵活性。这也是跟Cutlass的最大区别。
- 支持Hopper架构中的TMA加速：
- TMA 加载：用于左侧矩阵（LHS）、左侧缩放因子和右侧矩阵（RHS）的异步加载
- TMA 存储：用于输出矩阵的异步写回
- TMA 多播：仅限于左侧矩阵，实现单次加载，多线程共享，减少重复传输
- TMA 描述符（张量映射）预取：提前加载数据布局元信息，减少指令延迟
- 使用 stmatrix PTX 指令进行性能优化。stmatrix 是一种 PTX 指令，用于高效地将计算好的结果写回共享内存（shared memory）中。它在矩阵乘法（GEMM）的计算过程中发挥着重要作用，特别是在将累加后的结果从寄存器写入共享内存时。
- FFMA SASS 交错：DeepSeek 深入分析了SASS编译结果，在FFMA/FADD中调整SASS指令的执行顺序，此调整创造更多机会将MMA指令与提升精度的FFMA指令重叠，提高了细粒度缩放 FP8 GEMM 的性能。
- 高性能：在 Hopper GPU（例如：H100）上，可达到 1350+ TFLOPS 的 FP8 计算性能，这表明DeepSeek针对Hopper进行了深度优化。特别是对Dense模型的加速比MoE更明显。
- 仍在发展：DeepGEMM 在某些形状上的表现并不是很好。
- DeepGEMM遵循 CUTLASS 设计， 其内核为warp-specialization，支持数据移动、Tensor Core MMA 指令和 CUDA Core 提升精度的重叠。Warp-specialization 是一种 GPU 编程模型，它允许在一个线程块（thread block）内的不同 warp（一组 32 个线程）被赋予不同的角色或任务。warp-specialization在NVIDIA Hopper架构的GPU上特别有效，因为它可以更好地利用TMA（Tensor Memory Accelerator）和WGMMA（warp-group-level matrix multiply-accumulate）指令。
- 支持未对齐的块大小使某些场景的SM得到充分利用
- 重叠 TMA 存储和非 TMA RHS 缩放因子加载
- 统一优化的块调度器：通过统一的调度器调度所有非分组和分组内核，光栅化（Rasterization ）以增强 L2 缓存的复用。
DeepGEMM 接口设计及代码结构
DeepGEMM 提供了多个GEMM接口函数，包括：
- 常规稠密 GEMM：deep_gemm.gemm_fp8_fp8_bf16_nt，执行基本的非分组 FP8 GEMM
- 分组 GEMM（连续布局）：m_grouped_gemm_fp8_fp8_bf16_nt_contiguous，与Cutlass中的传统分组GEMM不同，DeepGEMM组只有M轴，而N和K必须保持固定。该设计是针对MoE模型专家具有相同形状的场景量身定制的。对于训练前向传播或推理预填充，每个专家可能会处理不同数量的Token，将这些Token串联成一个张量，称为“连续”布局。请注意，每个专家段必须与 GEMM M 块大小对齐(get_m_alignment_for_contiguous_layout() ) 。
- 分组 GEMM（掩码布局）：m_grouped_gemm_fp8_fp8_bf16_nt_masked，适用于推理解码。在推理解码阶段，当启用CUDA graph 并且CPU不知道每个专家收到的Token数量时，支持掩码分组GEMM。通过提供掩码张量，内核仅计算有效的部分。
此外，DeepGEMM还提供了一些实用函数，比如：
- deep_gemm.set_num_sms ：设置要使用的最大SM数
- deep_gemm.get_num_sms ：获取当前的SM最大数量
- deep_gemm.get_m_alignment_for_contiguous_layout ：获取分组连续布局的组级对齐要求
- deep_gemm.get_tma_aligned_size ：获取所需的 TMA 对齐大小
- deep_gemm.get_col_major_tma_aligned_tensor ：获取一个按列主序的 TMA 对齐张量
环境变量：
- DG_CACHE_DIR：用于存储编译内核的高速缓存目录，默认 $HOME/.deep_gemm
- DG_NVCC_COMPILER：指定 NVCC 编译器路径，默认从 torch.utils.cpp_extension.CUDA_HOME 获取
- DG_DISABLE_FFMA_INTERLEAVE：0 或 1，禁用 FFMA 交错优化
- DG_PTXAS_VERBOSE：0 或 1，显示详细的PTXAS编译器输出
- DG_PRINT_REG_REUSE：0 或 1，打印FFMA插入细节
- DG_JIT_PRINT_NVCC_COMMAND：0 或 1，打印NVCC编译命令
- DG_JIT_DEBUG：0 或 1，打印更多调试信息
代码结构：
/include/deep_gemm：DeepGEMM CUDA计算库。
- fp8_gemm.cuh：主要实现了一个基于FP8（8位浮点）的分组通用矩阵乘法（GEMM）内核，用于在NVIDIA GPU上进行高效的矩阵乘法计算。
- mma_utils.cuh：定义了一系列用于CUDA设备的结构体，这些结构体封装了不同形状的矩阵乘法（GEMM）操作，使用了Warp Group Matrix Multiply Accumulate（WGMMA）指令。
- scheduler.cuh：定义了一个模板结构体Scheduler，用于调度不同类型的矩阵乘法（GEMM）操作。
- tma_utils.cuh：一系列用于CUDA TMA加速的工具函数和模板，主要用于在CUDA编程中处理不同数据类型的张量映射和数据复制。
jit_kernels/：JIT计算库。
- m_grouped_gemm.py：用于定义两个分组通用矩阵乘法（GEMM）函数，支持FP8输入和BF16输出，并且使用了即时编译（JIT）和自动调优功能。
- gemm.py：用于实现FP8（8位浮点）输入和BF16（16位脑浮点）输出的矩阵乘法（GEMM）操作，并且使用即时编译（JIT）和自动调优来提高性能。
- tuner.py：定义了一个名为JITTuner的类，用于即时编译（JIT）和自动调优CUDA内核代码。
jit/：定义JIT相关操作。
- interleave_ffma.py：对CUDA生成的目标文件（.so文件）中的汇编代码进行处理，特别是针对其中的FFMA（Fused Fused Multiply-Add，融合乘加）指令段进行寄存器复用的修改。
DeepGEMM 性能
在 H800 上使用 NVCC 12.8 测试了 DeepSeek-V3/R1 推理中可能使用的所有形状（包括预填充和解码，但不包括张量并行）。所有加速指标均与DeepSeek内部基于 CUTLASS 3.6 且经过精心优化的实现进行比较。大多数情况下 DeepGEMM 表现更优，在某些形状上的表现不佳。
用于稠密模型的普通 GEMM
| M | N | K | Computation | Memory bandwidth | Speedup |
|---|---|---|---|---|---|
| 64 | 2112 | 7168 | 206 TFLOPS | 1688 GB/s | 2.7x |
| 64 | 24576 | 1536 | 289 TFLOPS | 2455 GB/s | 1.7x |
| 64 | 32768 | 512 | 219 TFLOPS | 2143 GB/s | 1.8x |
| 64 | 7168 | 16384 | 336 TFLOPS | 2668 GB/s | 1.4x |
| 64 | 4096 | 7168 | 287 TFLOPS | 2320 GB/s | 1.4x |
| 64 | 7168 | 2048 | 295 TFLOPS | 2470 GB/s | 1.7x |
| 128 | 2112 | 7168 | 352 TFLOPS | 1509 GB/s | 2.4x |
| 128 | 24576 | 1536 | 535 TFLOPS | 2448 GB/s | 1.6x |
| 128 | 32768 | 512 | 358 TFLOPS | 2103 GB/s | 1.5x |
| 128 | 7168 | 16384 | 645 TFLOPS | 2604 GB/s | 1.4x |
| 128 | 4096 | 7168 | 533 TFLOPS | 2221 GB/s | 2.0x |
| 128 | 7168 | 2048 | 510 TFLOPS | 2277 GB/s | 1.7x |
| 4096 | 2112 | 7168 | 1058 TFLOPS | 527 GB/s | 1.1x |
| 4096 | 24576 | 1536 | 990 TFLOPS | 786 GB/s | 1.0x |
| 4096 | 32768 | 512 | 590 TFLOPS | 1232 GB/s | 1.0x |
| 4096 | 7168 | 16384 | 1358 TFLOPS | 343 GB/s | 1.2x |
| 4096 | 4096 | 7168 | 1304 TFLOPS | 500 GB/s | 1.1x |
| 4096 | 7168 | 2048 | 1025 TFLOPS | 697 GB/s | 1.1x |
用于MOE模型的分组 GEMM (连续布局)
| #Groups | M per group | N | K | Computation | Memory bandwidth | Speedup |
|---|---|---|---|---|---|---|
| 4 | 8192 | 4096 | 7168 | 1297 TFLOPS | 418 GB/s | 1.2x |
| 4 | 8192 | 7168 | 2048 | 1099 TFLOPS | 681 GB/s | 1.2x |
| 8 | 4096 | 4096 | 7168 | 1288 TFLOPS | 494 GB/s | 1.2x |
| 8 | 4096 | 7168 | 2048 | 1093 TFLOPS | 743 GB/s | 1.1x |
用于MOE模型的分组 GEMM (掩码布局)
| #Groups | M per group | N | K | Computation | Memory bandwidth | Speedup |
|---|---|---|---|---|---|---|
| 1 | 1024 | 4096 | 7168 | 1233 TFLOPS | 924 GB/s | 1.2x |
| 1 | 1024 | 7168 | 2048 | 925 TFLOPS | 968 GB/s | 1.2x |
| 2 | 512 | 4096 | 7168 | 1040 TFLOPS | 1288 GB/s | 1.2x |
| 2 | 512 | 7168 | 2048 | 916 TFLOPS | 1405 GB/s | 1.2x |
| 4 | 256 | 4096 | 7168 | 932 TFLOPS | 2064 GB/s | 1.1x |
| 4 | 256 | 7168 | 2048 | 815 TFLOPS | 2047 GB/s | 1.2x |
DeepGEMM 的不足
DeepGEMM 的现有不足如下:
- 硬件依赖：仅支持Hopper架构，无法在Volta、Ampere等其他 NVIDIA GPU 上运行。
- 功能有限：专注于FP8和MoE，未提供FP16/FP32或更广泛的矩阵运算支持。
- 文档不足：截至目前，README提供了基本用法，但缺乏详细的API 文档或性能测试数据。
