DOCSEARCHbyAlgolia
Submit Docs

DeepGEMM Documentation

repository·main·Indexed 27 days ago

https://github.com/deepseek-ai/deepgemm

A high-performance CUDA kernel library for modern LLMs providing optimized tensor core primitives including FP8, FP4, and BF16 GEMMs. It features a lightweight JIT compilation approach, support for NVIDIA SM90 and SM100 architectures, fused Mega MoE kernels with overlapped communication, and MQA scoring for lightning indexers.

Tokens
2.2K
Snippets
3
Records
11
Agent score
43%

What's inside DeepGEMM

  1. Overview of DeepGEMM

    main

    DeepGEMM is a unified, high-performance tensor core kernel library designed for modern large language models. It provides key computation primitives including GEMMs (FP8, FP4, BF16), fused MoE with overlapped communication (Mega MoE), MQA scoring for the lightning indexer, and HyperConnection (HC).

    Key features:

    • Kernels are compiled at runtime via a lightweight Just-In-Time (JIT) module, eliminating the need for CUDA compilation during installation.
    • Designed for simplicity with a limited number of core kernel functions.
    • Performance matches or exceeds expert-tuned libraries across various matrix shapes.
    • Supports NVIDIA SM90 and SM100 architectures.

  2. Requirements for DeepGEMM

    main

    To use DeepGEMM, ensure your environment meets the following requirements:

    • GPU Architecture: NVIDIA SM90 or SM100.
    • Python: version 3.8 or higher.
    • Compiler: Must support C++20.
    • CUDA Toolkit:
      • For SM90: CUDA 12.3 or higher (CUDA 12.9+ highly recommended for best performance).
      • For SM100: CUDA 12.9 or higher.
    • PyTorch: version 2.1 or higher.
    • Dependencies (can be provided via Git submodules):
      • CUTLASS 4.0 or higher.
      • {fmt} library.

  3. Set up DeepGEMM for development

    main

    To set up a development environment, you must clone the repository recursively to include submodules, then run the development script to build the CPP JIT module.

    # Clone the repository with submodules
git clone --recursive git@github.com:deepseek-ai/DeepGEMM.git
cd DeepGEMM

# Build the CPP JIT module
./develop.sh
    git clone --recursive git@github.com:deepseek-ai/DeepGEMM.git
cd DeepGEMM
./develop.sh

  4. Understand DeepGEMM GEMM Naming and Layouts

    main

    DeepGEMM provides optimized GEMM kernels for NVIDIA GPUs using the convention D = C + A @ B.

    Memory Layouts:

    • NT (Non-transposed A, Transposed B): The default input shape layout.
    • SM90 Support: Only supports NT layout (row-major A, col-major B).
    • SM100 Support: Supports all layouts: NT, TN, NN, and TT.

    Scaling Factor Requirements: LHS scaling factors must have a TMA-aligned and transposed layout. The required format depends on the architecture:

    • SM90: Requires scaling factors in FP32 format.
    • SM100: Requires scaling factors in packed [UE8M0] format (4 UE8M0 values packed into a single torch.int).

    Note: Users are responsible for input transposition and FP8 casting. While utility functions exist, they may be slower than fused implementations.

  6. Configure DeepGEMM via Environment Variables

    main

    DeepGEMM can be configured using several environment variables for JIT, compilation, and debugging:

    CategoryVariableDescription
    GeneralDG_JIT_DEBUGPrint JIT debugging info (0 or 1, default 0)
    DG_PRINT_CONFIGSPrint selected configs for each shape (0 or 1, default 0)
    JIT CacheDG_JIT_CACHE_DIRCache directory for compiled kernels (default $HOME/.deep_gemm)
    CompilerDG_JIT_USE_NVRTCUse NVRTC instead of NVCC (0 or 1, default 0)
    DG_JIT_NVCC_COMPILERPath to NVCC compiler (defaults to torch.utils.cpp_extension.CUDA_HOME)
    DG_JIT_CPP_STANDARDC++ standard version (default 20)
    Compiler OutputDG_JIT_PRINT_COMPILER_COMMANDPrint compilation commands (0 or 1, default 0)
    DG_JIT_PTXAS_VERBOSEShow detailed PTXAS output (0 or 1, default 0)
    DG_JIT_PTXAS_CHECKAssert no local memory usage in kernels (0 or 1, default 0)
    Debug/ProfileDG_JIT_WITH_LINEINFOEmbed source line info for profiling (0 or 1, default 0)
    DG_JIT_DUMP_ASMDump both PTX and SASS (0 or 1, default 0)
    DG_JIT_DUMP_PTXDump PTX output (0 or 1, default 0)
    DG_JIT_DUMP_SASSDump SASS output (0 or 1, default 0)
    DG_COMM_KERNEL_DEBUGZero symmetric buffer before Mega MoE calls (0 or 1, default 0)
    Build OptionsDG_SKIP_CUDA_BUILDSkip CUDA extension build during installation (0 or 1, default 0)
    DG_FORCE_BUILDForce local build instead of pre-built wheels (0 or 1, default 0)
    DG_JIT_USE_RUNTIME_APIUse CUDA Runtime API (requires CUDA >= 12.8) (0 or 1, default 0)

  7. Enable NVRTC for faster JIT compilation

    main

    DeepGEMM supports NVRTC, which can provide up to 10x compilation speedup. You can enable it by setting the DG_JIT_USE_NVRTC environment variable to 1.

    Note: Using NVRTC may result in performance loss in some specific cases.

  8. Run Mega MoE Fused Kernels

    main

    Mega MoE fuses EP dispatch, linear 1 (FP8xFP4), SwiGLU, linear 2 (FP8xFP4), and EP combine into a single kernel. This requires a multi-process launch with symmetric memory and PyTorch >= 2.9.

    Workflow:

    1. Allocate symmetric memory using get_symm_buffer_for_mega_moe.
    2. Transform weights using transform_weights_for_mega_moe.
    3. Copy inputs into the buffer.
    4. Execute fp8_fp4_mega_moe.
    # Allocate symmetric memory buffer
# NOTES: requires PyTorch >= 2.9
buffer = deep_gemm.get_symm_buffer_for_mega_moe(
    group, num_experts, num_max_tokens_per_rank, num_topk, hidden, intermediate_hidden
)

# Transform weights (FP4 with UE8M0 SF) into the required layout
transformed_l1, transformed_l2 = deep_gemm.transform_weights_for_mega_moe(l1_weights, l2_weights)

# Copy inputs into the buffer before each call
# You may fuse these into previous kernels
buffer.x[:num_tokens].copy_(x_fp8)
buffer.x_sf[:num_tokens].copy_(x_sf)
buffer.topk_idx[:num_tokens].copy_(topk_idx)
buffer.topk_weights[:num_tokens].copy_(topk_weights)

# Run the fused mega MoE kernel
y = torch.empty((num_tokens, hidden), dtype=torch.bfloat16, device='cuda')
deep_gemm.fp8_fp4_mega_moe(y, transformed_l1, transformed_l2, buffer)

  9. Use Normal and Grouped GEMM Kernels

    main

    DeepGEMM offers several GEMM kernel types depending on your workload:

    1. Normal Dense GEMMs: For basic non-grouped FP8 GEMM, use fp8_gemm_{nt, nn, tn, tt}.
    2. Grouped GEMMs (Contiguous Layout): Designed for MoE models where experts share the same shape. It groups only the M-axis (N and K remain fixed). Each expert segment must be aligned to the GEMM M block size using get_mk_alignment_for_contiguous_layout(). Use m_grouped_fp8_gemm_{nt, nn}_contiguous.
    3. K-axis-grouped GEMMs: For MoE weight backward passes (M and N remain fixed). Use k_grouped_fp8_gemm_tn_contiguous.
    4. Grouped GEMMs (Masked Layout): For inference decoding with CUDA graphs where the number of tokens per expert is unknown. Use m_grouped_fp8_gemm_nt_masked with a mask tensor.

  10. Use DeepGEMM Utility Functions

    main

    The library provides several utilities for managing SM counts, alignment, and tensor layouts:

    • SM/TC Control: set_num_sms/get_num_sms, set_tc_util/get_tc_util, set_pdl/get_pdl (Programmatic Dependent Launch).
    • Alignment: set_mk_alignment_for_contiguous_layout/get_mk_alignment_for_contiguous_layout, get_theoretical_mk_alignment_for_contiguous_layout.
    • JIT/Compilation: set_ignore_compile_dims, set_block_size_multiple_of.
    • Layout/Scaling: transform_sf_into_required_layout, get_tma_aligned_size, get_mn_major_tma_aligned_tensor, get_mn_major_tma_aligned_packed_ue8m0_tensor, get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor.

  11. Use V3.2 MQA Kernels for Indexers

    main

    DeepGEMM provides two versions of MQA kernels: non-paged (prefilling) and paged (decoding).

    Non-paged version (fp8_mqa_logits) inputs:

    • q: E4M3 tensor [seq_len, num_heads, head_dim]
    • kv: E4M3 tensor [seq_len_kv, head_dim] with float SF [seq_len_kv]
    • weights: float tensor [seq_len, num_heads]
    • cu_seq_len_k_start & cu_seq_len_k_end: int tensors [seq_len]
    • clean_logits: boolean (whether to clean unfilled logits to -inf)

    Output: A tensor of shape [seq_len, seq_len_kv] representing token-to-token logits. For the paged version, see tests/test_attention.py.