Overview

DeepGEMM is a cutting-edge, unified tensor core kernel library designed for high-performance computations in modern large language models. It provides a cohesive CUDA codebase for essential primitives like FP8, FP4, and BF16 GEMMs, fused Mixture-of-Experts (MoE) with overlapping communication (Mega MoE), and MQA scoring. Key to DeepGEMM's design is a lightweight Just-In-Time (JIT) compilation module, enabling runtime kernel compilation without requiring CUDA compilation during installation. Despite its clean and accessible design, DeepGEMM delivers performance that rivals or surpasses expert-tuned libraries across a wide range of matrix shapes.

Key Features

• Advanced Precision Support: Optimized kernels for FP8, FP4, and BF16 GEMM operations, crucial for efficient LLM training and inference.
• High-Performance GEMM Kernels: Delivers state-of-the-art performance for both standard and grouped GEMMs, including specialized support for MoE scenarios.
• Fused Mega MoE: Implements a unified Mega MoE kernel that fuses and overlaps EP dispatch, linear transformations, SwiGLU activation, and EP combine, while overlapping NVLink communication with tensor core computation.
• Low-Overhead JIT Compilation: Employs a lightweight JIT module for runtime kernel compilation, minimizing installation complexity and enabling dynamic kernel generation. Supports NVRTC for significantly faster compilation.
• MQA Scoring Kernels: Provides efficient kernels for weighted ReLU MQA logits, essential for lightning indexer implementations in advanced LLMs.
• Flexible Grouped GEMMs: Offers both contiguous and masked grouped GEMM APIs to efficiently handle various MoE configurations during training and inference.
• Extensive Utilities: Includes a rich set of utility functions for configuring SM count, tensor core utilization, TMA alignment, and weight transformations.
• Modern GPU Architecture Support: Optimized for NVIDIA SM90 and SM100 architectures, with broad CUDA Toolkit compatibility.

Typical Use Cases

• LLM Training & Inference: Accelerate core GEMM operations within transformer architectures, significantly boosting throughput and reducing latency.
• Mixture-of-Experts (MoE) Models: Efficiently handle the complexities of MoE models with optimized kernels for both forward and backward passes, including fused Mega MoE for overlapping communication and computation.
• High-Throughput Attention Mechanisms: Utilize the MQA scoring kernels for advanced indexing strategies in attention layers, enabling faster retrieval and processing of key-value pairs.
• Research & Development: Serve as a clean and accessible platform for researchers and developers to learn and experiment with NVIDIA GPU kernel optimization techniques, particularly for tensor core operations.
• Custom Kernel Integration: Integrate DeepGEMM's optimized primitives into existing deep learning frameworks or custom model architectures where fine-grained control over GEMM performance is critical.