DeepGEMM C++ 核心架构与 JIT 编译引擎深度分析

graph TB
    subgraph "Python Layer"
        PY[deep_gemm/__init__.py]
    end

    subgraph "C++ API Layer (csrc/apis/)"
        GEMM[gemm.hpp]
        ATTN[attention.hpp]
        MEGA[mega.hpp]
        EINSUM[einsum.hpp]
        HC[hyperconnection.hpp]
        LAYOUT[layout.hpp]
        RUNTIME[runtime.hpp]
    end

    subgraph "JIT Kernel Implementations (csrc/jit_kernels/impls/)"
        K_SM90[sm90_fp8_gemm_1d1d.hpp]
        K_SM100[sm100_fp8_fp4_gemm_1d1d.hpp]
        K_SM90_BF16[sm90_bf16_gemm.hpp]
        K_SM100_BF16[sm100_bf16_gemm.hpp]
        K_OTHERS[...]
    end

    subgraph "JIT Engine (csrc/jit/)"
        COMP[compiler.hpp<br/>NVCC/NVRTC]
        CACHE[cache.hpp]
        KRUN[kernel_runtime.hpp]
        DEV[device_runtime.hpp]
        HANDLE[handle.hpp]
        INCPAR[include_parser.hpp]
    end

    subgraph "Heuristics (csrc/jit_kernels/heuristics/)"
        H_COMMON[common.hpp]
        H_SM90[sm90.hpp]
        H_SM100[sm100.hpp]
        H_CONFIG[config.hpp]
        H_RUNTIME[runtime.hpp]
    end

    subgraph "Utils (csrc/utils/)"
        UTILS[exception/format/hash/layout/<br/>lazy_init/math/system/compatibility]
    end

    subgraph "CUDA Infrastructure"
        NVCC[NVCC Compiler]
        NVRTC[NVRTC Runtime Compiler]
        CUDA_DRV[CUDA Driver API]
        CUDA_RT[CUDA Runtime API]
        CUBLASLT[cuBLASLt]
    end

    PY --> GEMM & ATTN & MEGA & EINSUM & HC & LAYOUT & RUNTIME
    GEMM & ATTN --> K_SM90 & K_SM100 & K_SM90_BF16 & K_SM100_BF16 & K_OTHERS
    K_SM90 & K_SM100 --> COMP
    K_SM90 & K_SM100 --> H_COMMON & H_SM90 & H_SM100
    COMP --> CACHE & KRUN
    KRUN --> DEV & HANDLE
    COMP --> INCPAR
    DEV --> HANDLE
    COMP --> NVCC & NVRTC
    KRUN --> CUDA_DRV & CUDA_RT
    DEV --> CUBLASLT
    COMP & KRUN & DEV --> UTILS

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.doc() = "DeepGEMM C++ library";
    deep_gemm::attention::register_apis(m);
    deep_gemm::einsum::register_apis(m);
    deep_gemm::hyperconnection::register_apis(m);
    deep_gemm::gemm::register_apis(m);
    deep_gemm::layout::register_apis(m);
    deep_gemm::mega::register_apis(m);
    deep_gemm::runtime::register_apis(m);
}

namespace deep_gemm::xxx {

static void register_apis(pybind11::module_& m) {
#if DG_FP8_COMPATIBLE and DG_TENSORMAP_COMPATIBLE
    m.def("function_name", &function_impl,
          py::arg("arg1"), py::arg("arg2") = default_value, ...);
#endif
}
}

flowchart TD
    A[fp8_fp4_gemm_nt] --> B{arch_major}
    B -->|9| C[sm90_fp8_gemm_1d1d<br/>sm90_fp8_gemm_1d2d]
    B -->|10| D[sm100_fp8_fp4_gemm_1d1d]

    A2[bf16_gemm_nt] --> E{arch_major}
    E -->|9| F[sm90_bf16_gemm]
    E -->|10| G[sm100_bf16_gemm]

    A3[m_grouped_fp8_fp4_gemm_nt_contiguous] --> H{arch_major}
    H -->|9| I[sm90_m_grouped_fp8_gemm_contiguous_1d2d]
    H -->|10| J[sm100_m_grouped_fp8_fp4_gemm_contiguous_1d1d]

Python Tensor → early_return 检查 → layout::transform_sf_pair_into_required_layout 
→ check_ab_fp8_fp4 (形状校验) → device_runtime->get_arch_major() 
→ 分派到具体 sm90/sm100 实现

static void fp8_fp4_mega_moe(
    const torch::Tensor& y,
    const std::tuple<torch::Tensor, torch::Tensor>& l1_weights_,
    const std::tuple<torch::Tensor, torch::Tensor>& l2_weights_,
    const torch::Tensor& sym_buffer,  // 对称通信 buffer
    const std::vector<int64_t>& sym_buffer_ptrs,
    ...
);

static torch::Tensor transform_sf_into_required_layout(
    const torch::Tensor& sf,
    const int& mn, const int& k,
    const std::variant<std::tuple<int,int,int>, std::tuple<int,int>>& recipe,
    ...
);

m.def("init", [&](const std::string& library_root_path, 
                 const std::string& cuda_home_path_by_python) {
    Compiler::prepare_init(library_root_path, cuda_home_path_by_python);
    KernelRuntime::prepare_init(cuda_home_path_by_python);
    IncludeParser::prepare_init(library_root_path);
});

sequenceDiagram
    participant PY as Python
    participant API as C++ API (apis/)
    participant KI as Kernel Impl (jit_kernels/impls/)
    participant COMP as Compiler
    participant CACHE as KernelRuntimeCache
    participant FS as File System
    participant CUDA as CUDA Driver

    PY->>API: fp8_fp4_gemm_nt(a, b, d, ...)
    API->>API: early_return 检查
    API->>API: layout::transform_sf_pair_into_required_layout
    API->>KI: sm90_fp8_gemm_1d1d(a, sfa, b, sfb, ...)

    KI->>KI: 构造 GemmDesc
    KI->>KI: get_best_config<SM90ArchSpec>(desc)
    KI->>KI: make_tma_a_desc / make_tma_b_desc / ...
    KI->>KI: LaunchRuntime::generate(args)
    Note over KI: 生成 C++ 代码字符串

    KI->>COMP: compiler->build("sm90_fp8_gemm_1d1d", code)

    COMP->>COMP: 计算 kernel_signature hash
    COMP->>CACHE: kernel_runtime_cache->get(dir_path)

    alt Cache Hit
        CACHE-->>COMP: 返回缓存的 KernelRuntime
    else Cache Miss
        COMP->>FS: 创建临时目录
        COMP->>FS: 写入 kernel.cu
        COMP->>CUDA: NVCC/NVRTC 编译
        CUDA-->>FS: 生成 kernel.cubin
        COMP->>FS: 可选：cuobjdump 反汇编
        COMP->>FS: fsync_dir + atomic rename
        COMP->>CACHE: 创建新 KernelRuntime
        CACHE->>CUDA: load_kernel(cubin_path)
        CUDA-->>CACHE: KernelHandle
    end

    COMP-->>KI: 返回 KernelRuntime
    KI->>KI: LaunchRuntime::launch(runtime, args)
    KI->>CUDA: launch_kernel(kernel, config, args...)
    CUDA-->>PY: 执行完成

class Compiler {
    std::string signature, flags;
    std::filesystem::path cache_dir_path;

    std::shared_ptr<KernelRuntime> build(const std::string& name, 
                                          const std::string& code) const;
    virtual void compile(const std::string& code, 
                         const std::filesystem::path& dir_path,
                         const std::filesystem::path &cubin_path, ...) const = 0;
};

class NVCCCompiler final: public Compiler {
    std::filesystem::path nvcc_path;

    void compile(...) const override {
        // 1. 写入 kernel.cu
        put(code_path, code);

        // 2. 在临时目录中编译（避免 cwd 文件污染）
        const auto command = fmt::format("cd {} && {} {} -cubin -o {} {}",
            compile_dir.c_str(), nvcc_path.c_str(), code_path.c_str(), 
            cubin_path.c_str(), flags);
        call_external_command(command);

        // 3. 可选：生成 PTX
        if (ptx_path.has_value()) { ... }
    }
};

-std=c++20 --diag-suppress=39,161,174,177,186,940
--ptxas-options=--register-usage-level=10
-I{library_include_path}
--gpu-architecture=sm_{arch}
--compiler-options=-fPIC,-O3,-fconcepts,-Wno-deprecated-declarations,-Wno-abi
-O3 --expt-relaxed-constexpr --expt-extended-lambda

class NVRTCCompiler final: public Compiler {
    void compile(...) const override {
        // 1. 创建 NVRTC program
        nvrtcCreateProgram(&program, code.c_str(), "kernel.cu", 0, nullptr, nullptr);

        // 2. 编译
        nvrtcCompileProgram(program, num_options, options);

        // 3. 获取 CUBIN
        nvrtcGetCUBINSize(program, &cubin_size);
        nvrtcGetCUBIN(program, cubin_data.data());

        // 4. 写入文件
        put(cubin_path, cubin_data);
    }
};

class IncludeParser {
    std::unordered_map<std::string, std::optional<std::string>> cache;

    std::string get_hash_value(const std::string& code, bool exclude_code = true) {
        for (const auto& i: get_includes(code))
            ss << get_hash_value_by_path(library_include_path / i) << "$";
        return get_hex_digest(ss.str());
    }
};

const auto kernel_signature = fmt::format("{}$${}$${}$${}", 
    name,           // kernel 名称，如 "sm90_fp8_gemm_1d1d"
    signature,      // 编译器签名，如 "NVCC12.9"
    flags,          // 编译 flags
    code            // 生成的 C++ 代码
);
const auto dir_path = cache_dir_path / "cache" / 
    fmt::format("kernel.{}.{}", name, get_hex_digest(kernel_signature));

graph LR
    A[Kernel Name] --> E[kernel_signature]
    B[Compiler Signature<br/>NVCC/NVRTC + version] --> E
    C[Compile Flags<br/>arch, include paths, options] --> E
    D[Generated Code<br/>template args + include hash] --> E
    E --> F[get_hex_digest<br/>FNV-1a + SplitMix64]
    F --> G[Cache Dir Path]

~/.deep_gemm/
├── tmp/                    # 临时目录（编译中）
│   └── {uuid}/
└── cache/                  # 最终缓存
    └── kernel.{name}.{hash}/
        ├── kernel.cu       # 源代码
        ├── kernel.cubin    # 编译产物
        ├── kernel.ptx      # 可选：PTX
        └── kernel.sass     # 可选：SASS

class KernelRuntimeCache {
    std::unordered_map<std::string, std::shared_ptr<KernelRuntime>> cache;

    std::shared_ptr<KernelRuntime> get(const std::filesystem::path& dir_path) {
        // 1. 内存缓存命中
        if (cache.contains(dir_path)) return cache[dir_path];

        // 2. 文件系统校验
        if (KernelRuntime::check_validity(dir_path))
            return cache[dir_path] = std::make_shared<KernelRuntime>(dir_path);

        return nullptr;  // 需要重新编译
    }
};

flowchart TD
    A[build(name, code)] --> B{内存缓存命中?}
    B -->|Yes| C[返回缓存的 KernelRuntime]
    B -->|No| D{文件系统缓存命中?}
    D -->|Yes| E[加载 CUBIN → 创建 KernelRuntime]
    D -->|No| F[NVCC/NVRTC 编译]
    F --> G[写入文件系统]
    G --> E
    E --> H[存入内存缓存]
    H --> C

class KernelRuntime {
    LibraryHandle library;
    KernelHandle kernel;

    explicit KernelRuntime(const std::filesystem::path& dir_path) {
        const auto cubin_path = dir_path / "kernel.cubin";

        // 方案1: CUDA 12.08+ Runtime API
        // cudaLibraryLoadFromFile → cudaLibraryGetKernel

        // 方案2: Driver API (默认)
        // cuModuleLoad → cuModuleGetFunction
        // 或 CUDA 12.4+: cuLibraryLoadFromFile → cuLibraryEnumerateKernels

        kernel = load_kernel(cubin_path, symbol_name, &library);
    }
};

template <typename Derived>
class LaunchRuntime {
    // 代码生成（静态多态）
    template <typename Args>
    static std::string generate(const Args& args) {
        auto code = Derived::generate_impl(args);
        code = fmt::format("// Includes' hash value: {}\n{}", 
                           include_hash, code);
        return code;
    }

    // 内核启动
    template <typename Args>
    static void launch(const std::shared_ptr<KernelRuntime>& runtime, 
                       const Args& args) {
        // 1. 获取当前 CUDA stream
        const auto stream = at::cuda::getCurrentCUDAStream();

        // 2. 构造 launch config
        auto config = construct_launch_config(kernel, stream, 
            smem_size, grid_dim, block_dim, cluster_dim, enable_pdl);

        // 3. 派生类实现具体 launch 逻辑
        Derived::launch_impl(kernel, config, args);
    }
};

class SM90FP8Gemm1D1DRuntime final: public LaunchRuntime<SM90FP8Gemm1D1DRuntime> {
    struct Args {
        GemmDesc gemm_desc;
        GemmConfig gemm_config;
        LaunchArgs launch_args;
        void *gmem_a_ptr, *gmem_b_ptr, *grouped_layout, *tensor_map_buffer;
        CUtensorMap tensor_map_a_base, tensor_map_b_base;
        CUtensorMap tensor_map_sfa, tensor_map_sfb, tensor_map_cd;
    };

    static std::string generate_impl(const Args& args) {
        return fmt::format(R"(
            #include <deep_gemm/impls/sm90_fp8_gemm_1d1d.cuh>
            static void __instantiate_kernel() {{
                auto ptr = reinterpret_cast<void*>(&sm90_fp8_gemm_1d1d_impl<
                    {}, {}, {}, ...  // 编译时常量参数
                >);
            }};
        )", ...);
    }

    static void launch_impl(const KernelHandle& kernel, 
                           const LaunchConfigHandle& config, Args args) {
        launch_kernel(kernel, config,
            args.gmem_a_ptr, args.gmem_b_ptr,
            args.grouped_layout, args.tensor_map_buffer,
            args.gemm_desc.m, args.gemm_desc.n, args.gemm_desc.k,
            args.tensor_map_a_base, args.tensor_map_b_base,
            args.tensor_map_sfa, args.tensor_map_sfb, args.tensor_map_cd);
    }
};

static CUtensorMap make_tma_2d_desc(
    const torch::Tensor& t,
    int gmem_inner_dim, int gmem_outer_dim,
    int smem_inner_dim, int smem_outer_dim,
    const int& gmem_outer_stride,
    const int& swizzle_mode, ...);

flowchart TD
    A[LaunchRuntime::launch] --> B[获取当前 CUDA Stream]
    A --> C[读取 device_runtime->get_pdl()]
    B --> D[construct_launch_config]
    C --> D
    D --> E{smem_size > 0?}
    E -->|Yes| F[cudaFuncSetAttribute<br/>MaxDynamicSharedMemorySize]
    E -->|No| G[构建 gridDim/blockDim]
    F --> G
    G --> H{cluster_dim > 1?}
    H -->|Yes| I[添加 ClusterDimension 属性]
    H -->|No| J{enable_pdl?}
    I --> J
    J -->|Yes| K[添加 PDL 属性]
    J -->|No| L[launch_kernel]
    K --> L

class DeviceRuntime {
    int num_sms = 0, tc_util = 0;
    bool enable_pdl = false;
    std::shared_ptr<cudaDeviceProp> cached_prop;

    // 架构检测
    std::pair<int, int> get_arch_pair();      // {major, minor}
    std::string get_arch(bool number_only, bool support_arch_family);
    int get_arch_major();                      // 9 (Hopper) or 10 (Blackwell)

    // SM 管理
    int get_num_sms();                         // 默认全部 SM
    void set_num_sms(int new_num_sms);

    // TC util
    int get_tc_util();                         // 默认 100%

    // PDL 控制
    bool get_pdl();
    void set_pdl(bool new_enable_pdl);
};

class DeviceRuntime {
    cublasLtHandle_t cublaslt_handle;
    torch::Tensor cublaslt_workspace;  // 32 MB

    cublasLtHandle_t get_cublaslt_handle() const;
    torch::Tensor get_cublaslt_workspace() const;
};

std::string get_arch(const bool& number_only, const bool& support_arch_family) {
    const auto [major, minor] = get_arch_pair();
    if (major == 10 and minor != 1) {
        if (number_only) return "100";
        return support_arch_family ? "100f" : "100a";
    }
    return std::to_string(major * 10 + minor) + (number_only ? "" : "a");
}

// 懒加载 CUDA driver
static void* get_driver_handle() {
    static void* handle = dlopen("libcuda.so.1", RTLD_LAZY | RTLD_LOCAL);
    return handle;
}

// 宏定义懒加载的 driver API
#define DECL_LAZY_CUDA_DRIVER_FUNCTION(name) \
    template <typename... Args> \
    static auto lazy_##name(Args&&... args) -> decltype(name(args...)) { \
        static FuncType func = nullptr; \
        if (func == nullptr) \
            func = reinterpret_cast<FuncType>(dlsym(get_driver_handle(), #name)); \
        return func(std::forward<Args>(args)...); \
    }

#if CUDART_VERSION >= 12080 and defined(DG_JIT_USE_RUNTIME_API)
    // CUDA Runtime Unified API
    using LibraryHandle = cudaLibrary_t;
    using KernelHandle = cudaKernel_t;
    cudaLibraryLoadFromFile(...);
    cudaLaunchKernelExC(...);
#else
    // CUDA Driver API
    using LibraryHandle = CUlibrary;  // or CUmodule
    using KernelHandle = CUfunction;
    cuLibraryLoadFromFile(...);
    cuLaunchKernelEx(...);
#endif

#define DG_FP8_COMPATIBLE (TORCH_VERSION_MAJOR > 2 or (TORCH_VERSION_MAJOR == 2 and TORCH_VERSION_MINOR >= 1))
#define DG_TENSORMAP_COMPATIBLE (CUDA_VERSION >= 12010)
#define DG_CUBLAS_GET_ERROR_STRING_COMPATIBLE (CUDART_VERSION >= 11042)
#define DG_CUBLASLT_ADVANCED_FEATURES_COMPATIBLE (CUDART_VERSION >= 11080)

class DGException final : public std::exception {
    std::string message;
public:
    explicit DGException(const char *name, const char* file, 
                         const int line, const std::string& error);
};

// 检查宏
#define DG_HOST_ASSERT(cond) ...
#define DG_HOST_UNREACHABLE(reason) ...
#define DG_NVRTC_CHECK(cmd) ...
#define DG_CUDA_DRIVER_CHECK(cmd) ...
#define DG_CUDA_RUNTIME_CHECK(cmd) ...
#define DG_CUBLASLT_CHECK(cmd) ...

C++ 异常 → pybind11 自动转换 → Python RuntimeError

#define FMT_HEADER_ONLY
#include <fmt/base.h>
#include <fmt/format.h>

static std::string get_hex_digest(const std::string& data) {
    const auto state_0 = fnv1a(data, 0xc6a4a7935bd1e995ull);
    const auto state_1 = fnv1a(data, 0x9e3779b97f4a7c15ull);
    // SplitMix64 混淆
    return hex(state_0) + hex(state_1);
}

// Major type 检测
cute::UMMA::Major get_major_type_ab(const torch::Tensor& t);

// SF layout 校验
torch::Tensor check_sf_layout(const torch::Tensor& sf, 
    const int& mn, const int& k, const int& gran_mn, const int& gran_k, ...);

// Recipe 默认值
std::tuple<int, int, int> get_default_recipe(
    const torch::ScalarType& sfa_dtype, const torch::ScalarType& sfb_dtype);

template <typename T>
class LazyInit {
    std::shared_ptr<T> ptr;
    std::function<std::shared_ptr<T>()> factory;
public:
    T* operator -> () {
        if (ptr == nullptr) ptr = factory();
        return ptr.get();
    }
};

constexpr auto kPackedFP4 = torch::kInt8;  // FP4 打包为 int8

template <typename T> static T ceil_div(const T& a, const T& b);
template <typename T> static constexpr T align(const T& a, const T& b);
static int get_tma_aligned_size(const int& x, const int& element_size);

// 环境变量读取
template <typename dtype_t>
static dtype_t get_env(const std::string& name, const dtype_t& default_value);

// 外部命令执行
static std::tuple<int, std::string> call_external_command(std::string command);

// 文件收集
static std::vector<std::filesystem::path> collect_files(const std::filesystem::path& root);

// 目录操作
static std::filesystem::path make_dirs(const std::filesystem::path& path);
static void safe_remove_all(const std::filesystem::path& path);

// UUID 生成
static std::string get_uuid();

template <typename ArchSpec>
static GemmConfig get_best_config(const GemmDesc& desc) {
    // 1. 生成 layout 候选
    const auto layout_candidates = ArchSpec::get_layout_candidates(desc);

    // 2. 评估每个候选
    for (auto& candidate : layout_candidates) {
        auto info = ArchSpec::get_layout_info(desc, candidate);
        if (ArchSpec::compare(info, best_info)) best = candidate;
    }

    // 3. 推导其他配置
    auto storage_config = ArchSpec::get_storage_config(desc, layout);
    auto pipeline_config = ArchSpec::get_pipeline_config(desc, layout, storage_config);
    auto launch_config = ArchSpec::get_launch_config(desc, layout);

    return {layout, storage_config, pipeline_config, launch_config};
}

graph TD
    A[错误发生] --> B{错误类型}
    B -->|主机断言失败| C[DGException<br/>Assertion]
    B -->|NVRTC 错误| D[DGException<br/>NVRTC]
    B -->|CUDA Driver 错误| E[DGException<br/>CUDA driver]
    B -->|CUDA Runtime 错误| F[DGException<br/>CUDA runtime]
    B -->|cuBLASLt 错误| G[DGException<br/>cuBLASLt]
    C --> H[what() 返回<br/>name + file:line + error]
    D --> H
    E --> H
    F --> H
    G --> H
    H --> I[pybind11 自动转换]
    I --> J[Python RuntimeError]

// 从 fused KV cache 中派生子张量
kv_cache = torch::from_blob(
    fused_kv_cache.data_ptr(),
    {num_kv_blocks, block_kv, head_dim / 2},
    {kv_cache_stride_bytes, head_dim / 2, 1},
    torch::TensorOptions().dtype(kPackedFP4)
);

static int get_tma_aligned_size(const int& x, const int& element_size) {
    constexpr int kNumTMAAlignmentBytes = 16;
    return align(x, kNumTMAAlignmentBytes / element_size);
}

classDiagram
    class Compiler {
        +signature: string
        +flags: string
        +cache_dir_path: path
        +build(name, code) shared_ptr~KernelRuntime~
        +compile(code, dir, cubin) void
    }

    class NVCCCompiler {
        +nvcc_path: path
        +compile(code, dir, cubin) void
    }

    class NVRTCCompiler {
        +compile(code, dir, cubin) void
    }

    class KernelRuntime {
        +library: LibraryHandle
        +kernel: KernelHandle
        +check_validity(dir) bool
    }

    class KernelRuntimeCache {
        +cache: unordered_map
        +get(dir) shared_ptr~KernelRuntime~
    }

    class DeviceRuntime {
        +num_sms: int
        +tc_util: int
        +enable_pdl: bool
        +get_arch() string
        +get_arch_major() int
    }

    class LaunchArgs {
        +grid_dim: pair~int,int~
        +num_threads: int
        +smem_size: int
        +cluster_dim: int
        +enable_pdl: bool
    }

    class GemmDesc {
        +gemm_type: GemmType
        +kernel_type: KernelType
        +m, n, k: int
        +num_sms, tc_util: int
        +compiled_dims: string
    }

    class GemmConfig {
        +layout: Layout
        +storage_config: StorageConfig
        +pipeline_config: PipelineConfig
        +launch_config: LaunchConfig
    }

    Compiler <|-- NVCCCompiler
    Compiler <|-- NVRTCCompiler
    Compiler --> KernelRuntime : creates
    KernelRuntimeCache --> KernelRuntime : caches
    KernelRuntime --> DeviceRuntime : uses
    LaunchArgs --> KernelRuntime : launches
    GemmDesc --> GemmConfig : generates
    GemmConfig --> LaunchArgs : derives