在大语言模型推理场景中,FP8 量化已成为平衡计算吞吐与数值精度的关键技术路径。然而,FP8 格式的动态范围限制导致矩阵乘法累加过程中容易出现精度损失,这在深层 Transformer 堆叠中会逐级放大,最终影响模型输出质量。DeepGEMM 作为 DeepSeek 开源的 FP8 GEMM 库,通过细粒度缩放(Fine-grained Scaling)与两级累加机制,在 NVIDIA Hopper 张量核心上实现了高效且数值稳定的矩阵运算,成为支撑大模型推理落地的核心技术组件。
细粒度缩放的核心设计
传统 FP8 GEMM 通常采用 per-tensor 缩放策略,即对整个矩阵使用单一缩放因子。这种方式实现简单,但无法有效捕捉输入数据的局部分布特征,导致部分数值在 FP8 表示中溢出或精度严重下降。DeepGEMM 采用了更精细的缩放粒度:左侧矩阵(LHS)按 1×128 的瓦片进行缩放,右侧矩阵(RHS)按 128×128 的瓦片进行缩放。这种不对称的粒度设计源于矩阵乘法的计算特性 ——LHS 的每一行需要与 RHS 的不同列块进行乘法,缩放因子需要跟随数据流精确对齐到每个计算单元。
在实际 kernel 实现中,缩放因子作为独立的张量存储在全局内存中,通过 TMA(Tensor Memory Accelerator)异步加载到共享内存。NVIDIA Hopper 架构的 TMA 引擎支持多播传输(multicast),这意味着同一个缩放因子可以被多个 warp 共享,减少了重复访存开销。kernel 执行过程中,每个线程块首先将 FP8 输入数据与对应缩放因子从全局内存加载到共享内存,然后执行分块矩阵乘法。这种数据布局确保了缩放操作与矩阵乘法的高度流水线化,避免了额外的同步开销。
两级累加策略与精度保障
FP8 张量核心的执行效率极高,但其累加精度受限。当参与乘法的元素数量较大时,中间结果的舍入误差会显著累积。DeepGEMM 引入的两级累加策略有效缓解了这一问题:第一级在 FP8 张量核心上完成初始矩阵乘法,利用其高吞吐量产生初步结果;第二级通过 CUDA 核心执行 promotion 操作,将 FP8 结果提升至 FP32 进行精细累加,最后在必要时刻再 demote 回 FP8 格式输出。
这种设计的关键在于两级累加的切换时机与阈值控制。DeepGEMM 提供了可配置的 promotion 阈值参数,当累加器中的值超过设定阈值时自动触发精度提升。这一机制类似于 TransformerEngine 中的 amax 重分配策略,但 DeepGEMM 将其深度集成在 kernel 内部,减少了框架层的额外开销。对于典型的大模型推理负载,建议将 promotion 阈值设置为 1.0e3 至 1.0e4 之间的动态范围,并根据实际推理精度要求进行微调。
Hopper 架构的深度优化
DeepGEMM 的性能优势建立在对 NVIDIA Hopper 架构特性的充分利用之上。首先,库中所有 kernel 均针对 Hopper FP8 张量核心优化,支持 TF32 与 FP8 的混合精度计算模式。Tensor Core 的 MMA(Matrix Multiply-Accumulate)指令可以直接消费 FP8 输入并输出 FP8 结果,大幅减少了数据类型转换的寄存器压力。其次,TMA 引擎的异步数据传输能力被充分发挥 —— 数据加载与计算完全重叠,线程块可以在等待前一批数据加载完成的同时执行当前批次的乘法运算。
在内存访问模式方面,DeepGEMM 采用了 128×128 的瓦片尺寸,这与 Hopper 张量核心的计算粒度完美匹配。每个线程块负责处理一个独立的瓦片,瓦片内部进一步划分为 64×64 的子块供各个 warp 协作完成。共享内存的 bank 冲突是影响 kernel 效率的关键因素,DeepGEMM 通过调整数据布局和访问顺序,将 bank 冲突率控制在可接受范围内。实测表明,在 H100 GPU 上处理 4096×4096 的方阵乘法时,DeepGEMM 的吞吐量可达到同场景下 CUTLASS 实现的 1.15 至 1.25 倍。
推理部署的关键参数配置
将 DeepGEMM 集成到大模型推理服务中时,需要关注以下几个核心配置参数。第一个参数是 block_m 与 block_n,分别控制线程块在 M 与 N 维度上的瓦片大小,默认值为 128,可根据批量大小(batch size)动态调整 —— 当 batch size 较大时应适当增大瓦片尺寸以提高 occupancy。第二个参数是 num_stages,用于控制流水线阶段数,典型值为 3 至 5,较高值可以更好地隐藏内存访问延迟,但会增加共享内存消耗。第三个参数是 enable_tma_multicast,启用后可通过 TMA 多播机制减少缩放因子的重复加载,建议在单次推理调用多个专家矩阵时开启。
对于 MoE(Mixture of Experts)场景下的 grouped GEMM,DeepGEMM 提供了专用的 kernel 变体。这类 kernel 的输入包含多个独立的矩阵对(num_experts 个),输出也相应地为多个结果矩阵。关键配置参数 num_experts 决定了每次 kernel _launch 处理的专家数量,建议将其设置为 8 的倍数以充分利用 SIMT 调度。此外,grouped GEMM kernel 需要额外配置 expert_stride 参数,确保不同专家的数据在全局内存中正确隔离。
数值稳定性验证与监控
部署 DeepGEMM kernel 前,强烈建议在目标模型上进行数值稳定性验证。验证流程应包含以下几个步骤:准备一组具有代表性的输入数据,涵盖模型各层的典型激活范围;使用 FP32 基准实现计算参考结果;将 DeepGEMM 的 FP8 结果与参考结果进行逐元素对比,统计相对误差的均值、最大值与标准差;特别关注误差在 Softmax 前的残差连接与注意力计算中的传播情况。
线上部署时的监控指标应包括 kernel 执行时间、SM 利用率、内存带宽利用率以及 FP8 overflow 检测计数。NVIDIA Nsight Compute 可以采集详细的 hardware counter 数据,帮助识别潜在的瓶颈。当检测到 overflow 计数异常升高时,可通过降低缩放阈值或切换至更粗粒度的缩放策略来缓解。DeepGEMM 社区建议在生产环境保留一份 FP32 fallback 路径,用于处理边界 case 或模型微调后的首轮验证。
综合来看,DeepGEMM 通过细粒度缩放与两级累加机制,在 NVIDIA Hopper 架构上实现了 FP8 GEMM 的效率与精度平衡。理解和掌握上述 kernel 设计要点与配置参数,是将这一技术成功落地到实际大模型推理系统的关键前提。
资料来源
- DeepGEMM 官方 GitHub 仓库:https://github.com/deepseek-ai/DeepGEMM
- NVIDIA Hopper FP8 Tensor Core 技术文档