AB
AiBoss
project

DeepGEMM - DeepSeek's open-source FP8 general-purpose matrix multiplication library

DeepGEMM is an open-source library from DeepSeek designed for efficient and concise FP8 matrix multiplication (GEMM), currently supporting only the NVIDIA Hopper tensor core. DeepGEMM supports GEMM operations in both normal and mixed expert (MoE) groups...

What is DeepGEMM?

DeepGEMM is an open-source library from DeepSeek designed for efficient and concise FP8 matrix multiplication (GEMM), currently supporting only NVIDIA Hopper tensor cores. DeepGEMM supports GEMM operations in both standard and mixed expert (MoE) groups, based on Just-In-Time (JIT) compilation technology, eliminating the need for installation-time compilation and supporting dynamic optimization at runtime. DeepGEMM addresses the precision limitations of FP8 by employing fine-grained scaling and CUDA core dual-level accumulation techniques, while leveraging Hopper's Tensor Memory Accelerator (TMA) feature to improve data transfer efficiency. With a core codebase of only about 300 lines, DeepGEMM is easy to learn and optimize. Its lightweight design achieves or surpasses expert-level optimized libraries across various matrix shapes.

Main functions of DeepGEMM

  • High-efficiency FP8 matrix multiplication (GEMM)DeepGEMM is a high-efficiency library designed specifically for FP8 (8-bit floating-point) matrix multiplication. It supports fine-grained scaling and significantly improves the performance and accuracy of matrix operations.
  • Supports both standard and grouped GEMM:
    • Ordinary GEMM: Applicable to regular matrix multiplication operations.
    • Grouping GEMMSupports grouped matrix multiplication in hybrid expert (MoE) models, including contiguous layout and masked layout, optimizing scenarios where multiple experts share shapes.
  • Just-in-Time (JIT) Compilation DesignBased on just-in-time (JIT) compilation technology, all kernels are dynamically compiled at runtime, eliminating the need for compilation during installation. Optimization is performed based on parameters such as matrix shape and block size, saving registers and improving performance.
  • Hopper architecture optimizationDesigned specifically for the NVIDIA Hopper architecture, it fully leverages the features of the Tensor Memory Accelerator (TMA), including TMA loading, storage, multicasting, and descriptor prefetching, significantly improving data transfer efficiency.
  • Fine-grained scaling and two-stage accumulationTo address the issue of insufficient precision in FP8, DeepGEMM introduces fine-grained scaling technology. Based on the CUDA core's two-level accumulation mechanism, it upscales FP8 calculation results to a higher-precision format (such as BF16) to ensure computational accuracy.
  • Lightweight designThe core code is concise, only about 300 lines, making it easy to understand and extend. It avoids dependencies on complex templates or algebraic structures, reducing the difficulty of learning and optimization.

DeepGEMM's project address

DeepGEMM performance

  • Standard GEMM (non-grouped) performance
    • Maximum acceleration ratioFor certain matrix shapes, DeepGEMM can achieve a speedup of 2.7 times, significantly improving the efficiency of matrix multiplication.
    • Computing performanceIn large-scale matrix operations, DeepGEMM can achieve computing performance of over 1000 TFLOPS, approaching the theoretical peak of the Hopper architecture GPU.
  • Grouped GEMM (MoE model) performance
    • acceleration ratioIn grouped GEMM, DeepGEMM can achieve a speedup of 1.1-1.2 times, significantly improving the training and inference efficiency of the MoE model.
    • Memory bandwidth optimizationBased on TMA features, DeepGEMM performs exceptionally well in memory bandwidth utilization, achieving performance close to the hardware limit.
      • Contiguous Layout
      • Masked Layout

DeepGEMM System Requirements

  • Hardware Requirements:
    • GPU architecture: Must support the NVIDIA Hopper architecture, specifically a GPU that supports sm_90a. H800 or H100 Hopper architecture GPUs are recommended, as they are optimized for FP8 computing and Tensor Cores.
    • CUDA compatibilityRequires CUDA 12.3 or later. CUDA 12.8 or later is recommended for optimal performance.
  • Software Requirements:
    • Recommended operating systemsUse a Linux operating system (such as Ubuntu, CentOS, etc.) because CUDA and PyTorch are better supported on Linux.
    • Python versionPython 3.8 or later.
    • CUDA ToolkitCUDA version 12.3 or later. The CUDA version must be compatible with the GPU architecture; version 12.8 or later is recommended to fully utilize the features of the Hopper architecture.
    • PyTorch: PyTorch 2.1 or later.
    • CUTLASS library: CUTLASS 3.6 or later.
  • Other requirements:
    • Conventional compilation tools (such as gcc, make, etc.).
    • The torch.utils.cpp_extension module is used for CUDA extensions.

Application scenarios of DeepGEMM

  • Large-scale AI model inferenceAccelerates high-dimensional matrix multiplication and improves reasoning speed.
  • Hybrid Expert (MoE) ModelOptimize grouped matrix multiplication to improve computational efficiency.
  • Low-precision calculation:The accuracy issues of FP8 are solved by fine-grained scaling, maintaining high-precision output.
  • High-performance computingBased on the characteristics of the Hopper architecture, improve the efficiency of matrix operations.
  • Deep learning framework optimizationAs a low-level optimization library, it accelerates model training and inference.