DeepSeek开源周第三天DeepGEMM

项目地址: https://github.com/deepseek-ai/DeepGEMM

项目介绍

DeepGEMM是一个为NVIDIA Hopper架构GPU设计的高性能FP8矩阵乘法库，具备细粒度缩放能力，支持DeepSeek-V3/R1模型所需的各种矩阵形状和布局。能在H800 GPU上实现高达1350+TFLOPS的计算性能。核心逻辑仅约300行代码，却能超越许多调优库的性能。

核心功能

FP8 GEMM支持:专为NVIDIA Hopper张量核心优化的FP8计算。e多种GEMM模式

细粒度缩放:支持精确控制缩放因子，提高数值稳定性

JIT编译优化:运行时编译确保最佳性能配置

两级累加机制:解决FP8张量核心累加精度问题

SM资源控制:可调控使用的流多处理器数量

DeepGEMM提供三种FP8GEMM实现，满足不同场景需求

1.普通密集GEMM

。标准矩阵乘法运算，适用于常规深度学习模型

。支持NT格式(非转置输入A，转置输入B)

。优化的TMA加载和存储操作

2.连续布局分组GEMM

。为MOE模型训练前向传播或推理预填充设计

支持在M轴分组，不同专家接收变长token0

。需要满足对齐要求，以获得最佳性能

3.掩码布局分组GEMM

。为启用CUDA图的推理解码阶段设计

支持CPU无感知专家接收token数量的场景0

通过掩码张量指定计算的有效部分0

与DeepEP的低延迟内核配合使用

系统需求

GPU架构:NVIDIA Hopper架构GPU(如H100/H800，需支持sm 90a)

Python:3.8或更高版本

PyTorch:2.1或更高版本

CUDA: 12.3或更高版本(建议12.8以获得最佳性能)

CUTLASS:3.6或更高版本(可通过Git子模块克隆)

目录概览

背景知识

Tensor Memory Accelerator(TMA): NVIDIA 在其最新 Hopper GPU 架构中引入的一种技术。TMA 的主要功能是优化矩阵乘法运算，通过异步数据传输和索引计算来提高性能。

Parallel Thread Execution(PTX): CUDA编程模型中的一种中间表示语言(IntermediateRepresentationLanguage)，主要用于桥接高级语言代码与GPU硬件的机器码。

Streaming Assembly(SASS): NVIDIA GPU的原生汇编语言，用于直接控制硬件资源。Float Fused Multiply Add(FFMA): SASS中的单精度浮点融合乘加指令。

NVIDIA CUDA Compiler(NVCC): 用于编译CUDA程序的专用编译器。

Shared Memory(SMEM): GPU架构中的关键组成部分，主要用于线程块内的线程协作与优化内存访问模式。

安装与使用

安装后，在Python项目中导入 deep_gemm 即可使用。

基本用法

常规GEMM操作示例:

import torch

import deep_gemm

# 准备输入数据

# LHS:Fp8 m x k矩阵 + FP32缩放因子

# RHS:FP8 n x k矩阵 + FP32缩放因子

m,n，k= 128,4096，7168

device = torch.device('cuda')

lhs = torch.randn(m,k,dtype=torch.float8_e4m3fn,device=device)lhs_scales =torch,randn(m,(k+127)//128,dtype=torch.float32, device=device)

rhs = torch.randn(n,k,dtype=torch.float8_e4m3fn,device=device)rhs_scales =torch.randn((n+127)//128,(k+ 127)// 128,dtype=torch.float32device=device)

# 确保RHS是转置格式的，LHS缩放因子需要TMA对齐和转置

rhs_scales = deep_gemm.get_col major_tma_aligned tensor(rhs_scales)

# 创建输出张量

out = torch.zeros(m,n,dtype=torch.bfloat16,device=device

#执行FP8 GEMM

deep_gemm.gem fp8_fp8_bf16_nt((lhs, lhs_scales),(rhs, rhs_scales), out)

环境变量控制

• DeepGEMM提供多个环境变量用于调整行为:
• DG CACHE DIR:存储编译内核的缓存目录，默认 $HOME/.deep_gemum
• DG NVCCCOMPILER:指定NVCC编译器路径
• DG DISABLE FFMAINTERLEAVE:禁用FFMA交织优化(0或1)
• DG PTXAS VERBOSE: 显示详细编译器输出(0或1)
• DG PRINT REG REUSE:打印FFMA交织详情(0或1)。
• DG JIT DEBUG: 输出更多调试信息(0或1)

密集模型 FP8 GEMM性能

小批量高效率: 小批量(M=64/128)下性能尤为出色，最高达2.7倍加速，能大幅提高推理速度。

计算峰值:大批量下达到1358 TFLOPS纯计算性能(H800 GEMM 理论峰值约 1979 TFLOPS)

带宽优化:小批量下内存带宽高达2.6 TB/S

测试环境:H800 GPU+NVCC 12.8，与基于CUTLASS 3.6优化的实现对比:

MoE模型连续布局分组GEMM性能

训练友好:专为MOE模型训练和预填充阶段优化0

专家数量可扩展:支持不同数量专家组(4/8组)配置

计算密集优化:计算密集场景高达 1297 TFLOPS(H800 GEMM 理论峰值约 1979 TFLOPS)

MoE模型掩码布局分组GEMM性能

CUDA图兼容:专为启用CUDA图的推理解码阶段设计

动态专家分配:CPU无需预知每个专家接收的token数

小M高效率:每组小批量(M=256)下带宽高达2064GB/S

延迟敏感场景:适合与DeepEP低延迟内核配合提升推理解码场景的响应速度

DeepGEMM 时序图

下面将结合AI深入分析3段关键代码片段。我们将可以看到DeepGEMM虽然仅有约300行核心代码，但通过精心设计的数据流、巧妙的内存管理以及对Hopper架构特性的充分利用，实现了卓越的性能。

特别是充分利用了TMA多播、两级精度累加等技术，使其在各种矩阵形状下都能高效运行。

源码剖析 fp8 gemm.cuh

1.共享内存分配与流水线设计

流水线和共享内存管理是高性能GEMM的关键，这段代码展示了DeepGEMM如何精心安排共享内存布局，确保各个流水线阶段的数据能够高效并行处理。1024字节对齐特别针对Hopper架构的TMA优化，可以大幅提升内存带宽利用率。

static constexpr uint32_t SEM_D_SIZE = BLOCK_M * BLOCK_N * sizeof(_nv_bfloat16);



static constexpr uint32_t SEM_A_SIZE_PER_STAGE = BLOCK_M * BLOCK_K * sizeof(nv_fp8_e4m3);



static constexpr uint32_t SEM_B_SIZE_PER_STAGE = BLOCK_N * BLOCK_K * sizeof(nv_fp8_e4m3);



static constexpr uint32_t SEM_SCALES_A_SIZE_PER_STAGE = BLOCK_M * sizeof(float);



extern shared align(1024) uint8_t sem_buffer[];



auto sem_d = reinterpret_cast<nv_bfloat16*>(sem_buffer);

nv_fp8_e4m3* sem_a[kNumStages];

nv_fp8_e4m3* sem_b[kNumStages];

float* sem_scales_a[kNumStages];

float* sem_scales_b;



#pragma unroll

for (int i = 0; i < kNumStages; ++i) {

    sem_a[i] = reinterpret_cast<nv_fp8_e4m3*>(sem_buffer + SEM_D_SIZE + i * SEM_A_SIZE_PER_STAGE);

    sem_b[i] = reinterpret_cast<nv_fp8_e4m3*>(sem_buffer + SEM_D_SIZE + kNumStages * SEM_A_SIZE_PER_STAGE + i * SEM_B_SIZE_PER_STAGE);

    sem_scales_a[i] = reinterpret_cast<float*>(sem_buffer + SEM_D_SIZE + kNumStages * (SEM_A_SIZE_PER_STAGE + SEM_B_SIZE_PER_STAGE) + i * SEM_SCALES_A_SIZE_PER_STAGE);

}

2.TMA和纵横切分优化

DeepGEMM的关键优化之一是高效的数据加载和块调度。这段代码实现了TMA加载线程与计算线程的分离，同时使用了异步屏障机制实现数据传输与计算的重叠。TMA多播技术(kNumrMAMulticast )是Hopper架构的特性，能在一次操作中复制数据到多个计算单元，显著提升带宽利用率。

//分离TMA线程和计算线程，优化资源利用

if(threadIdx.x>= kNumMathThreads )// IMA专用线程组，解除寄存器分配限制，cutlass::arch::warpgroup_reg_dealloc<kNumTMARegisters>();

//只使用一个线程执行TMA操作

if(threadIdx.x == kNumMathThreads){

// 持续调度块计算，实现全GPU占用

while(scheduler.get_next block(m block idx,n block idx))1launch k iterations([&](int k iter, auto type){//流水线内循环，加载数据

#pragma unroll

for(uint32_ts=0;s<kNumInnerstages;++ s){// 等待消费者释放

empty_barriers[s]->wait((scheduler.current_iter * kNumIterations

k iter + 1)& 1);

//发起TMA，使用多播加速

auto& full barrier = *full barriers[s];int k idx=k iter * kFullKOfAllStages + s * BLOCK K;

//多播复制A矩阵，一次传输到多个计算单元

tma_copy<kNumTMAMulticast>(&tensor_map_a,reinterpret_cast<uint64_t*>

(&full_barrier),

smem_a[s],k idx,scheduler.get_global_idx(shape_m,

BLOCK M,m_block_idx));

//加载A的缩放因子

tma_copy<kNumTMAMulticast>(&tensor_map_scales_a,reinterpret_cast<uint64_t*>(&full barrier),smem scales a[s],m block idx * BLOCK Mscheduler.get_global idx(SHAPE K SCALES,1,k_idx /

BLOCK_K));

//常规加载B矩阵，无需多播tma_copy(&tensor_map_b, reinterpret_cast<uint64_t*>(&full barrier)smem b[s],k idx,scheduler.get global idx<false>(SHAPE NBLOCK N,n block idx,m block idx));

//通知计算线程数据已准备就绪full barrier.arrive and expect tx(SMEM A SIZE PER STAGESMEM_B_SIZE_PER_STAGE+ SMEM_SCALES_A_SIZE_PER_STAGE);

});

}else 

// 计算线程逻辑

}

3.两级精度累加与细粒度缩放

FP8计算的关键挑战在于保持精度，DeepGEMM通过CUDA核心的两级累加机制解决。这段代码展示了两级累加的精妙设计:先在张量核心上执行FP8MMA操作，然后在CUDA核心上进行高精度累加和缩放。每个缩放医子由A和B两个矩阵的缩放因子乘积构成，实现了细粒度的数值控制，有效解决了FP8张量核心累加精度不足的问题。

//执行张量核心MMA操作

fpragma unrollfor(int i=0;i<WGMMA::kNumAccum; ++ i)

warpgroup_fence_operand(accum[i]); // 同步操作数访问//告知所有线程已就绪warpgroup_arrive();

#pragma unroll

for(int k=0;k<BLOCKK/WGMMA::K;++ k){auto desc_a = make_smem_desc(smem a[s]+ math_wg_idx * WGMMA::M * BLOCK_K + K * WGMMA::K, 1);auto desc b = make_smem desc(smem b[s]+ k * WGMMA::K, 1);WGMMA::wgmma(desc a,desc b,accum，k);// 执行矩阵乘法累加

warpgroup_commit batch(); // 提交批处理指令#pragma unrollfor(int i=0;i<WGMMA::kNumAccum; ++ i)warpgroup_fence_operand(accum[i]);// 同步访问结果warpgroup_wait<0>(); // 等待计算完成

//通知IMA可以复用内存

empty barrier arrive(s);

//应用细粒度缩放因子提升精度float scale_0 0=scale_a 0 *scale b 0,scale 1 0 = scale_a 1 * scale b 0;float scale_0 1,scale 1 1;if constexpr(not kMustUseUniformedscaleB)

scale_01=scale_a0*scaleb1,scale11=scale_a 1 * scale b 1;#pragma unroll

for(inti=0;i<WGMMA::kNumAccum/4;++ i){

bool predicate =kMustUseUniformedScaleB or i < num former iters;

//在CUDA核心上执行高精度累加，避免FP8精度损失

final_accum[i *4+0]+=(predicate ?scale_0_0 :scale_0 1)* accum[i * 4 + 0];

final_accum[i *4 + 1]+=(predicate ?scale_0_0:scale_0 1)* accum[i * 4 + 1];

final_accum accum[i*4+2]+=(predicate ?scale_1_0:scale_1_1)* accum[i * 4 + 2];

final_accum[i *4+3]+=(predicate ?scale_1_0 : scale_1_1)* accum[i * 4 + 3];

}

DeepSeek三个核心开源项目的相互关系

文章转载自：DeepSeek开源周第三天DeepGEMM