DeepGEMM 构建系统与基础设施架构分析

flowchart TD
    A[git clone --recursive] --> B[./install.sh]
    B --> C[python setup.py bdist_wheel]
    C --> D{CachedWheelsCommand}
    D -->|DG_FORCE_BUILD=1 或 local version| E[本地构建]
    D -->|默认发布模式| F[尝试下载预编译 wheel]
    F -->|下载成功| G[直接重命名放置到 dist/]
    F -->|下载失败| E
    E --> H[CustomBuildPy.run]
    H --> H1[prepare_includes: 复制 CuTe/cutlass 头文件到 build 目录]
    H1 --> H2[generate_default_envs: 生成 envs.py]
    H2 --> H3[generate_pyi_file: 从 C++ 生成 _C.pyi 类型桩]
    H3 --> H4[标准 build_py + CUDAExtension 编译]
    H4 --> I[生成 deep_gemm-<ver>+cu<cu>-torch<torch>-cxx11abi<abi>-<py>-<platform>.whl]
    I --> J[pip install dist/*.whl --force-reinstall]

deep_gemm-<version>+cu<cuda>-torch<torch>-cxx11abi<abi>-cp<py>-linux_x86_64.whl
└── deep_gemm/
    ├── __init__.py          # 入口，调用 _C.init()
    ├── _C.so                # 轻量 JIT 宿主（不含 kernel 代码）
    ├── _C.pyi               # 自动生成的类型桩
    ├── envs.py              # 构建时捕获的环境变量默认值
    ├── include/
    │   ├── deep_gemm/**/*   # 真实 kernel 头文件（JIT 编译时 include）
    │   ├── cute/**/*         # CUTLASS CuTe 子集
    │   └── cutlass/**/*     # CUTLASS 子集
    ├── utils/               # Python 工具函数
    ├── testing/             # 测试基础设施
    ├── legacy/              # Triton 遗留内核（A100）
    └── mega/                # Mega MoE 对称内存管理

cxx_flags = ['-std=c++17', '-O3', '-fPIC', '-Wno-psabi', '-Wno-deprecated-declarations',
             f'-D_GLIBCXX_USE_CXX11_ABI={int(torch.compiled_with_cxx11_abi())}']

sources = ['csrc/python_api.cpp']  # 单一入口文件
build_include_dirs = [
    f'{CUDA_HOME}/include',
    f'{CUDA_HOME}/include/cccl',       # CUDA Core Compute Libraries
    'deep_gemm/include',               # 项目自身头文件
    'third-party/cutlass/include',     # CUTLASS
    'third-party/fmt/include',         # fmt 格式化库
]
build_libraries = ['cudart', 'nvrtc']  # CUDA Runtime + NVRTC

def get_package_version():
    # 从 deep_gemm/__init__.py 读取公开版本号
    public_version = ast.literal_eval(version_match.group(1))  # 如 '2.4.2'

    revision = ''
    if DG_USE_LOCAL_VERSION:
        # 检查 git 工作区是否干净
        # 干净: +<short_sha>  不干净: assert 失败
        revision = '+' + subprocess.check_output(['git', 'rev-parse', '--short', 'HEAD'])
    return f'{public_version}{revision}'

class CustomBuildPy(build_py):
    def run(self):
        self.prepare_includes()      # 1. 准备头文件
        self.generate_default_envs() # 2. 生成 envs.py
        self.generate_pyi_file()     # 3. 生成 .pyi 类型桩
        build_py.run(self)           # 4. 标准构建

def prepare_includes(self):
    build_include_dir = os.path.join(self.build_lib, 'deep_gemm/include')
    for d in third_party_include_dirs:  # ['third-party/cutlass/include/cute', '.../cutlass']
        shutil.copytree(src_dir, dst_dir)  # 复制到 build 目录

def generate_default_envs(self):
    code = '# Pre-installed environment variables\n'
    code += 'persistent_envs = dict()\n'
    for name in ('DG_JIT_CACHE_DIR', 'DG_JIT_PRINT_COMPILER_COMMAND', 'DG_JIT_CPP_STANDARD'):
        code += f"persistent_envs['{name}'] = '{os.environ[name]}'\n" if name in os.environ else ''

class CachedWheelsCommand(_bdist_wheel):
    def run(self):
        if DG_FORCE_BUILD or DG_USE_LOCAL_VERSION:
            return super().run()  # 强制本地构建

        wheel_url, wheel_filename = get_wheel_url()
        try:
            # 尝试从 GitHub Release 下载（1 秒超时）
            urllib.request.urlopen(wheel_url, timeout=1)
            # 下载成功：直接重命名为 wheel 格式
        except (HTTPError, URLError):
            # 下载失败：回退到本地源码构建
            super().run()

def get_wheel_url():
    wheel_filename = (
        f'deep_gemm-{deep_gemm_version}+cu{cuda_version}'
        f'-torch{torch_version}-cxx11abi{cxx11_abi}'
        f'-{python_version}-{platform_name}.whl'
    )
    # 示例: deep_gemm-2.4.2+cu129-torch2.8-cxx11abi0-cp312-linux_x86_64.whl
    wheel_url = f'https://github.com/DeepSeek-AI/DeepGEMM/releases/download/v{version}/{wheel_filename}'

def get_ext_modules():
    if DG_SKIP_CUDA_BUILD:
        return []  # 跳过 CUDA 构建（用于 PyPI 纯 Python 发布）
    return [CUDAExtension(name='deep_gemm._C', ...)]

flowchart LR
    subgraph install.sh
        I1[rm -rf build dist] --> I2[python setup.py bdist_wheel]
        I2 --> I3[pip install dist/*.whl --force-reinstall]
    end

    subgraph develop.sh
        D1[ln -sf cutlass 头文件到 deep_gemm/include] --> D2[rm -rf build dist]
        D2 --> D3[python setup.py build]
        D3 --> D4[find .so 创建符号链接到 deep_gemm/]
    end

    subgraph build.sh
        B1[rm -rf build dist] --> B2[python setup.py bdist_wheel]
    end

# 创建符号链接，让 Python 能找到头文件
ln -sf $script_dir/third-party/cutlass/include/cutlass deep_gemm/include
ln -sf $script_dir/third-party/cutlass/include/cute deep_gemm/include

# 构建后创建 .so 符号链接
so_file=$(find build -name "*.so" -type f | head -n 1)
ln -sf "../$so_file" deep_gemm/

# NOTES: current just for CMake-based IDE (e.g. CLion) indexing, the real compilation is done via JIT
cmake_minimum_required(VERSION 3.10)
project(deep_gemm LANGUAGES CXX CUDA)

# C++20/CUDA20 标准（与 kernel JIT 编译一致）
set(CMAKE_CXX_STANDARD 20)
set(CMAKE_CUDA_STANDARD 20)

# 主 Python API 入口（IDE 索引用）
pybind11_add_module(_C csrc/python_api.cpp)

# 关键：所有 kernel 头文件都 include 进来供索引
cuda_add_library(deep_gemm_indexing_cuda STATIC csrc/indexing/main.cu)

// GEMM kernels
#include <deep_gemm/impls/sm90_bf16_gemm.cuh>
#include <deep_gemm/impls/sm90_fp8_gemm_1d1d.cuh>
// ... 所有 kernel 头文件

int main() { return 0; }  // 空 main，仅用于编译索引

传统 CUTLASS 流程：
  setup.py → nvcc 编译所有 SM arch → 链接巨大 .so → 分发 GB 级 wheel

DeepGEMM 流程：
  setup.py → 仅编译轻量 _C.so（< 1MB） → 分发 wheel
  运行时 → 首次调用 kernel → JIT 编译为 .cubin → 缓存到 ~/.deep_gemm/

flowchart TD
    A[Python 调用 deep_gemm.fp8_gemm_nt] --> B[_C.so 中的 pybind11 绑定]
    B --> C[JIT Compiler: 解析 #include <deep_gemm/...>]
    C --> D[递归展开所有头文件依赖]
    D --> E[计算依赖哈希值]
    E --> F{缓存命中?}
    F -->|命中| G[加载 ~/.deep_gemm/<hash>/kernel.cubin]
    F -->|未命中| H[nvcc 编译为 .cubin]
    H --> I[缓存到磁盘]
    I --> G
    G --> J[cuModuleLoadData 加载 cubin]
    J --> K[cuLaunchKernel 执行]

[submodule "third-party/cutlass"]
    path = third-party/cutlass
    url = https://github.com/NVIDIA/cutlass.git
[submodule "third-party/fmt"]
    path = third-party/fmt
    url = https://github.com/fmtlib/fmt.git

// 使用 fmt 进行编译时字符串格式化
flags = fmt::format("-std=c++{} --diag-suppress=39,161,...", 
                    get_env<int>("DG_JIT_CPP_STANDARD", 20));

# tests/test_mega_moe.py
spec = importlib.util.spec_from_file_location(
    'tilelang_ops',
    os.path.join('..', 'third-party', 'tilelang_ops', '__init__.py'))
tilelang_ops = importlib.util.module_from_spec(spec)

git submodule status:
  -f3fde58... third-party/cutlass  (- 表示未初始化)
  -553ec11...... third-party/fmt    (- 表示未初始化)

# deep_gemm/envs.py (自动生成)
persistent_envs = dict()
persistent_envs['DG_JIT_CACHE_DIR'] = '/home/user/.deep_gemm'
# ... 仅包含构建时存在的变量

# deep_gemm/__init__.py
from .envs import persistent_envs
for key, value in persistent_envs.items():
    if key not in os.environ:  # 不覆盖用户显式设置
        os.environ[key] = value

flowchart TD
    A[scripts/generate_pyi.py] --> B[build_cpp_function_index]
    B --> B1[扫描 csrc/ 下所有 .cpp/.h 文件]
    B1 --> B2[正则匹配函数签名]
    B2 --> B3[过滤非函数关键字 if/for/return]
    B3 --> B4[构建 func_name → signature 索引]

    A --> C[extract_m_def_statements]
    C --> C1[扫描 m.def(...)]
    C --> C2[BracketTracker 解析嵌套括号]
    C --> C3[提取 python_name, cpp_name, py::arg]

    A --> D[parse_mdef_and_attach_cpp_signatures]
    D --> D1[匹配 m.def 与 C++ 函数]
    D --> D2[提取返回值类型、参数列表]

    A --> E[cpp_type_to_python_type]
    E --> E1[std::vector<T> → list[T]]
    E --> E2[std::optional<T> → Optional[T]]
    E --> E3[torch::Tensor → torch.Tensor]
    E --> E4[bool/int/float/str 映射]

    A --> F[generate_pyi_file_content]
    F --> G[stubs/_C.pyi]

class BracketTracker:
    def __init__(self):
        self.paren = 0      # ()
        self.bracket = 0    # []
        self.brace = 0      # {}
        self.angle = 0      # <> (仅顶层)

def cpp_default_to_python_default(cpp_default: str):
    if s == 'false': return 'False'
    if s == 'true': return 'True'
    if s in ('nullptr', 'NULL'): return 'None'
    # std::tuple<int, int>({128, 128}) → (128, 128)
    # std::vector<int>({1,2,3}) → [1, 2, 3]

# generators.py 提供枚举器 + 生成器分离的设计

def enumerate_normal(dtype):
    """枚举所有测试参数组合"""
    for kernel_type in get_kernel_types(dtype):
        for quant_config in QuantConfig.get_list_from_dtype(dtype):
            for m in m_fwd_list:
                for n, k in nk_list:
                    yield kernel_type, quant_config, m, n, k, ...

def generate_normal(m, n, k, ...):
    """生成具体的测试张量"""
    a = torch.randn((m, k), dtype=torch.bfloat16)
    b = torch.randn((n, k), dtype=torch.bfloat16)
    # FP8/FP4 量化
    a = cast_fp8_fp4_with_major(a, ...)
    return a, b, c, d, ref_d

# deep_gemm/testing/numeric.py
def calc_diff(x: torch.Tensor, y: torch.Tensor):
    """余弦相似度: 1 - 2*x·y/(|x|²+|y|²)"""
    x, y = x.double(), y.double()
    denominator = (x * x + y * y).sum()
    sim = 2 * (x * y).sum() / denominator
    return 1 - sim

# 容差标准
quant_config.max_diff()  # FP4: 0.02, FP4×FP8: 0.01, FP8: 0.001

# deep_gemm/testing/bench.py
def bench_kineto(fn, kernel_names, num_tests=30, ...):
    """使用 torch.profiler + Kineto 进行 GPU 内核计时"""
    # 1. L2 cache flush (8 GB memset)
    # 2. warmup + active profiling
    # 3. 解析 profiler 表格提取内核时间

# tests/test_layout.py
@test_filter(lambda: get_arch_major() >= 10)
def test_fp8_bhd_hdr_bhr():
    """仅在 SM100+ 上运行"""

# generators.py
def get_kernel_types(dtype):
    if dtype == torch.bfloat16:
        return (KernelType.KernelNoSF, )
    return (KernelType.Kernel1D2D, ) if get_arch_major() == 9 else (KernelType.Kernel1D1D, )

flowchart TD
    A[触发: tag v* 创建] --> B[setup_release]
    B --> B1[提取版本号]
    B --> B2[创建 GitHub Release]

    B --> C[build_wheels 矩阵构建]
    C --> C1[_build.yml: ubuntu-22.04]
    C1 --> C2[Python 3.8-3.13]
    C2 --> C3[PyTorch 2.4.0-2.8.0]
    C3 --> C4[CUDA 12.9.1]
    C4 --> C5[cxx11_abi: FALSE/TRUE]

    C5 --> D[构建 wheel]
    D --> E[上传到 GitHub Release]

    E --> F[publish_package]
    F --> F1[安装 CPU-only PyTorch]
    F --> F2[DG_SKIP_CUDA_BUILD=1 构建 sdist]
    F --> F3[twine upload 到 PyPI]

matrix:
  os: [ubuntu-22.04]
  python-version: ["3.8", "3.9", "3.10", "3.11", "3.12", "3.13"]
  torch-version: ["2.4.0", "2.5.1", "2.6.0", "2.7.1", "2.8.0"]
  cuda-version: ["12.9.1"]
  cxx11_abi: ["FALSE", "TRUE"]
  exclude:
    - torch-version: "2.4.0"
      python-version: "3.13"  # PyTorch < 2.5 不支持 Python 3.13

- name: Free up disk space
  run: |
    sudo rm -rf /usr/share/dotnet
    sudo rm -rf /opt/ghc
    sudo rm -rf /opt/hostedtoolcache/CodeQL

- name: Set up swap space
  uses: pierotofy/set-swap-space@v1.0
  with:
    swap-size-gb: 10

- name: Restore build cache
  uses: actions/cache/restore@v4
  with:
    path: build.tar
    key: build-${{ inputs.release-version }}-${{ inputs.python-version }}-${{ inputs.cuda-version }}-${{ inputs.torch-version }}-${{ inputs.cxx11_abi }}-${{ github.run_number }}-${{ github.run_attempt }}

export MAX_JOBS=$([ "$MATRIX_CUDA_VERSION" == "129" ] && echo 1 || echo 2)
export NVCC_THREADS=2

export TORCH_CUDA_ARCH_LIST="7.0 7.2 7.5 8.0 8.6 8.7 9.0+PTX"

- name: Build core package
  env:
    DG_USE_LOCAL_VERSION: "0"
    DG_SKIP_CUDA_BUILD: "1"
  run: python setup.py sdist --dist-dir=dist

timeline
    title DeepGEMM 架构演进

    section 2025 Q1
        2025.04.18 : 1550 TFLOPS on H800<br/>初始版本发布
        2025.05.07 : NVRTC 支持<br/>DG_JIT_USE_NVRTC=1<br/>编译速度 10x 提升
        2025.05.14 : Weight gradient 内核<br/>支持 dense/MoE 反向传播

    section 2025 Q3
        2025.07.20 : 重大重构<br/>支持 SM90/SM100<br/>NVRTC/SASS 优化禁用<br/>低 CPU 开销 JIT
        2025.09.28 : MQA Scoring 内核<br/>V3.2 Lightning Indexer 支持<br/>加权 ReLU MQA logits

    section 2026 Q1-Q2
        2026.04.16 : Mega MoE 发布<br/>FP8xFP4 GEMM 支持<br/>FP4 Indexer<br/>PDL 支持<br/>更快 JIT 编译

DeepGEMM/
├── build/                    # 构建临时目录（gitignore）
│   └── lib.linux-x86_64-3.12/
│       └── deep_gemm/
│           ├── include/      # 复制的 CuTe/cutlass 头文件
│           └── _C.cpython-312-x86_64-linux-gnu.so
├── dist/                     # 最终分发目录
│   └── deep_gemm-2.4.2+cu129-torch2.8-cxx11abi0-cp312-linux_x86_64.whl
├── stubs/                    # 生成的 .pyi 文件（gitignore）
│   └── _C.pyi
├── deep_gemm/
│   ├── include/              # 符号链接（开发模式）
│   │   ├── cute -> third-party/...
│   │   └── cutlass -> third-party/...
│   └── *.so                  # 符号链接（开发模式）
└── *.egg-info/               # 构建元数据

# 开发调试
export DG_JIT_DEBUG=1                    # 打印 JIT 调试信息
export DG_JIT_PRINT_COMPILER_COMMAND=1   # 显示编译命令
export DG_JIT_DUMP_SASS=1                # 转储 SASS 代码

# 性能分析
export DG_USE_NVIDIA_TOOLS=1             # 跳过内部 profiling
export DG_JIT_WITH_LINEINFO=1            # 嵌入行信息

# 构建控制
export DG_FORCE_BUILD=1                  # 强制本地构建
export DG_SKIP_CUDA_BUILD=1              # 跳过 CUDA 构建
export DG_JIT_USE_NVRTC=1                # 使用 NVRTC（更快编译）

# 缓存配置
export DG_JIT_CACHE_DIR=/tmp/dg_cache    # 自定义缓存目录

# 运行单个测试
python tests/test_fp8_fp4.py

# 运行 sanitizer（内存检查）
python tests/test_sanitizer.py --funcs test_fp8_fp4.test_gemm --tools memcheck

# 运行 Mega MoE（多进程）
torchrun --nproc_per_node=8 tests/test_mega_moe.py --num-processes 8

# 条件过滤测试
python tests/test_bf16.py  # 自动跳过 SM90 不支持的测试