Hotdry.

Article

DeepSeek DeepGEMM 详解:面向 LLM 推理的 FP8 精细化缩放矩阵乘核

深入解析 DeepGEMM 的 FP8 矩阵乘实现,聚焦细粒度缩放机制与 LLM 推理落地的关键参数。

2026-04-18ai-systems

在大语言模型推理场景中,矩阵乘法(GEMM)占据着核心计算负载。随着模型规模突破千亿参数并引入 MoE(混合专家)架构,如何在保持数值精度的前提下充分挖掘 GPU Tensor Core 的峰值算力,成为推理优化的关键课题。DeepSeek 开源的 DeepGEMM 库正是针对这一需求设计的轻量级、高性能 FP8 GEMM 核实现,其核心理念在于通过细粒度缩放(fine-grained scaling)在计算效率与模型精度之间取得平衡。

DeepGEMM 的设计定位与架构选择

DeepGEMM 并非一个追求功能完备性的通用 BLAS 库,而是专为现代大语言模型定制的计算原语集合。从功能边界来看,它涵盖了三类核心 GEMM:传统稠密矩阵乘(FP8、FP4、BF16)、融合 MoE 的专家路由矩阵乘(Mega MoE),以及面向 Lightning Indexer 的 MQA 打分核。这种收敛的设计哲学使得库维护者能够将优化精力集中在推理场景最高频的 kernel 类型上,避免 CUTLASS 等通用库因过度抽象导致的性能损耗。

在硬件适配层面,DeepGEMM 明确面向 NVIDIA Hopper 架构(SM90)和 Blackwell 架构(SM100)两类 GPU。两者在 FP8 数据格式支持上存在重要差异:SM90 要求的缩放因子为 FP32 格式,而 SM100 则需使用打包后的 UE8M0 格式(每 4 个 UE8M0 缩放因子打包为单个 torch.int)。这一差异直接影响部署时的内存布局策略 —— 在 SM100 上使用正确的 UE8M0 打包格式可将缩放因子内存占用降低 75%。

值得注意的是,DeepGEMM 采用了运行时 JIT(Just-In-Time)编译架构。所有 kernel 在首次调用时动态编译为 SASS(机器码),无需在安装阶段进行 CUDA 编译。这一设计显著降低了部署复杂度,同时通过缓存机制(默认路径 $HOME/.deep_gemm)确保后续调用的编译开销可忽略。对于需要快速迭代的推理服务而言,首次冷启动延迟是唯一需要关注的编译成本。

细粒度缩放的技术本质与工程权衡

FP8 矩阵乘相比 BF16 的性能优势来源于两方面:更高的算力密度(FP8 Tensor Core 理论算力为 BF16 的 2 倍)和更低的内存带宽需求。然而,FP8 的数值表示范围仅为 FP16 的约 1/256,直接进行矩阵乘会导致显著的量化误差。缩放因子(Scaling Factor)的引入正是为了在计算过程中恢复动态范围。

所谓细粒度缩放,是指缩放因子的作用维度从整个矩阵(tensor-wise)细化到更小的数据块。常见的粒度包括:per-channel(按输出通道)、per-group(按自定义块大小)。粒度越细,对异常值的容忍度越高,模型精度损失越小;但同时缩放因子的存储开销增大,且在 kernel 内部需要更复杂的坐标计算来获取正确的缩放值。

DeepGEMM 在接口设计上将缩放因子的预处理责任交给调用方。对于 SM90,缩放因子需预先转换为 TMA 对齐的转置布局,并保持 FP32 格式;对于 SM100,则需通过 transform_sf_into_required_layoutget_mn_major_tma_aligned_packed_ue8m0_tensor 等工具函数完成 UE8M0 打包。这种设计保持了 kernel 本体的简洁性,但要求集成方在数据流水线中正确处理格式转换。

在 LLM 推理的典型场景中,attention 的 QK^T 计算对精度尤为敏感,建议采用更细粒度的缩放(如 per-head 或 per-group);而 FFN 层的 Gemm 操作可使用较粗粒度以换取更高的 occupancy。实践中可通过在离线测试集上对比不同缩放粒度下的 PPL(困惑度)变化来确定最优配置。

面向推理落地的关键配置参数

将 DeepGEMM 集成到生产推理服务时,以下参数需要重点关注和调优。

SM 调度策略:通过 deep_gemm.set_num_sms(n) 限制使用的流多处理器数量。默认行为是尝试全部占用,但在多租户共享 GPU 或需保留部分算力用于其他 kernel(如 attention)时,手动降低 SM 占用率可避免资源争用。建议将至少 20% 的 SM 预留用于 attention kernel,以维持显存访问与计算的平衡。

Tensor Core 利用率预估set_tc_util(ratio) 提供近似的 Tensor Core 利用率目标。在 JIT 编译阶段,DeepGEMM 会根据这一参数选择不同的 tile 策略。过高估计实际可达的 TC 利用率可能导致编译出的 kernel 在真实 workload 下表现不佳。建议通过 ncu(NVIDIA Nsight Compute)实际测量后再反馈调整。

连续布局的 M/K 对齐:对于 MoE 场景下的 grouped GEMM,调用 m_grouped_fp8_gemm_nt_contiguous 时需确保 M 轴按 get_mk_alignment_for_contiguous_layout() 返回值对齐。若未对齐,kernel 会自动截断至对齐边界,导致部分 token 被忽略。调试时可设置 DG_PRINT_CONFIGS=1 查看每种 shape 选中的具体配置。

掩码布局的支持:在 decode 阶段使用 CUDA Graph 时,token 到专家的分配在 CPU 侧不可预知,此时应使用 m_grouped_fp8_gemm_nt_masked 接口。掩码 tensor 的构建需与 DeepEP 等通信库的输出格式对齐,确保有效 token 位置的标记正确。

JIT 编译优化:对于追求首次请求低延迟的场景,可在服务启动阶段预热(warmup)所有可能遇到的 batch size 和 sequence length 组合。编译缓存默认持久化至本地磁盘,重启服务后仍可复用。若使用 CUDA 12.9 及以上版本,建议保持默认使用 NVCC 而非 NVRTC(设置 DG_JIT_USE_NVRTC=0),以获得更优的后编译 SASS 优化。

性能基准与监控要点

根据官方披露的数据,DeepGEMM 在 H800(SM90)上单精度 FP8 理论峰值可达 1550 TFLOPS。实际推理性能取决于 batch size、sequence length 和模型结构三个维度的交互。 batch size 较小时(≤8),kernel launch overhead 和缩放因子读取延迟占比上升,FP8 相比 BF16 的优势减弱;随着 batch size 增大,FP8 的内存带宽节省开始主导,推理吞吐量可提升 30% 至 60%。

部署监控建议关注以下指标:GPU SM Activity(目标 > 85%)、FP64/FP32/FP16/TF32 Tensor Core Activity(区分不同精度的时间占比)、DRAM(显存)带宽使用率(FP8 场景下通常低于 50% 的理论峰值,因计算而非带宽受限)、以及 JIT 编译缓存命中率(首次冷启动后的命中率应接近 100%)。若发现 SM Activity 持续低于 70%,需检查是否存在 GEMM 之外的计算瓶颈(如 attention 的 softmax 或 RoPE 位置编码)。

在精度验证层面,建议在上线前使用 hellaSwag、MMLU 等标准 benchmark 对量化后的模型进行端到端评估,确保 FP8 + 细粒度缩放的精度损失在可接受范围内(通常 < 0.5% 的绝对指标变化)。

小结

DeepGEMM 通过精简的 kernel 设计、运行时 JIT 编译和对细粒度缩放的原生支持,为 LLM 推理提供了一条直接利用 Hopper/Blackwell FP8 算力的工程化路径。在实际落地时,重点应放在缩放因子格式与硬件版本的匹配、batch size 与 kernel 粒度的协同调优,以及基于实际 workload 的性能监控上。FP8 矩阵乘的优化不是一个独立的算子调优问题,而是需要与 attention kernel、MoE 通信、显存管理等多方面协同的系统工程。

资料来源:DeepSeek DeepGEMM GitHub 仓库(https://github.com/deepseek-ai/DeepGEMM)

ai-systems