DeepGEMM 离线预编译：Cubin Bundle、Cold Compile 清零与 Wheel 交付

核心结果与交付边界

目标：让隔离网络环境中的 vLLM/DeepGEMM 直接加载预编译 cubin，跳过运行时 NVCC JIT，并以可校验 wheel 交付。

核心结果

工程链路

模型 shape 收集
  -> 运行时 JIT 产物归档
  -> 多模型 bundle union / 去重
  -> SHA-256 一致性检查
  -> 单卡与 TP=2 cold=0 验证
  -> baseline + DeepGEMM / combined wheel
  -> 离线安装与包内 bundle 自动解析

关键工作还包括修复 TP=2 环境的 NCCL/CUDA 包冲突，并区分“bundle 命中”“加载耗时”和“稳态 GEMM latency”三个指标。

文档

结论边界

• 7.7x 是 cold/load 启动阶段收益，不是 hot GEMM TFLOPS 提升。
• 8 个模型表示验证覆盖，不表示所有可能 shape 已预编译。
• 其中一个 BF16 dense 小模型不会触发 DeepGEMM，保留在清单中是为了验证“不误触发”。
• 超大模型和 H100 因当时硬件条件未覆盖，不能写成全平台支持。

简历建议

面向离线 H200 推理环境设计 DeepGEMM cubin 预编译与 wheel-only 交付，将 union bundle 从 109 扩展到 359 kernels，覆盖 8 个单卡模型及 Qwen3-8B-FP8 TP=2，验证范围内实现 cold compile 清零，并通过 SHA-256 保证 bundle 合并一致性。

离线安装与包内解析

日期：2026-05-21
目标：把离线 precompiled 交付收敛为只交付 deep_gemm wheel，并把默认 bundle 路径切到 Hopper union bundle。

1. 本次落地内容

• 恢复 deepgemm-packaged-precompiled/deep_gemm/precompiled/generate_precompiled.py
• 恢复 deepgemm-packaged-precompiled/deep_gemm/precompiled/seed_cache_root.py
• 新增 deepgemm-packaged-precompiled/scripts/build_hopper_union_bundle.py
• 新增包内元数据 deep_gemm/precompiled/release_coverage.json
• 将 package 默认 bundle 切换到：

deep_gemm/precompiled/sm90/nvcc12/fp8/qwen3_dpsk_v32_h100_h200_fp8/cache

2. 为什么这样改

本轮 mentor 已经明确最终更希望拿到的是 deep_gemm wheel，而不是 baseline vllm wheel 或 combined wheel。

所以交付形态收敛为：

• 目标环境保留自己的 vllm
• 单独安装 deep_gemm-*.whl
• deep_gemm wheel 自带 package 内 precompiled bundle
• 安装后如果未手动设置 DG_JIT_PRECOMPILED_DIR，默认自动解析 package 内 default.json

3. 统一生成 union bundle

在 deepgemm-packaged-precompiled 目录执行：

python3 scripts/build_hopper_union_bundle.py --force

默认行为：

• 从现有 qwen3_5_35b_a3b_fp8/cache 读取 320-kernel seed
• 生成新 workload：

sm90/nvcc12/fp8/qwen3_dpsk_v32_h100_h200_fp8/cache

• 自动写入新的 default.json
• 自动生成 deep_gemm/precompiled/release_delivery.json

如果后续补了 H100 / H200 TP / DeepSeek 的新 cache，可以继续 merge：

python3 scripts/build_hopper_union_bundle.py \
  --source-cache-dir [redacted-path]/qwen3_h200_tp2/cache \
  --source-cache-dir [redacted-path]/qwen3_h100/cache \
  --source-cache-dir [redacted-path]/deepseek_v32_h200/cache \
  --source-cache-dir [redacted-path]/deepseek_v32_h100/cache \
  --summary-json [redacted-path]/qwen3_h200_tp2_verify_summary.json \
  --summary-json [redacted-path]/deepseek_v32_h200_verify_summary.json

脚本会逐个 merge，并对同名 kernel 做 sha256 一致性检查；如果 hash 相同但内容不同，会直接报错停下，不会悄悄覆盖。

4. wheel-only 打包流程

在 deepgemm-packaged-precompiled 目录执行：

python3 -m pytest tests/test_precompiled_bundle.py
python3 deep_gemm/precompiled/generate_precompiled.py \
  --check-bundle-dir deep_gemm/precompiled/sm90/nvcc12/fp8/qwen3_dpsk_v32_h100_h200_fp8
python3 setup.py bdist_wheel
sha256sum dist/deep_gemm-*.whl > dist/SHA256SUMS

5. 安装验证

在干净环境中：

python3 -m pip install --force-reinstall dist/deep_gemm-*.whl --break-system-packages
python3 - <<'PY'
import os
import deep_gemm
print("deep_gemm =", deep_gemm.__file__)
print("DG_JIT_PRECOMPILED_DIR =", os.environ.get("DG_JIT_PRECOMPILED_DIR"))
PY

如果要把 package 内 bundle seed 到新的 DG_JIT_CACHE_DIR，可以执行：

python3 "$(python3 - <<'PY'
import pathlib
import deep_gemm
print(pathlib.Path(deep_gemm.__file__).resolve().parent / 'precompiled' / 'seed_cache_root.py')
PY
)" \
  --cache-root /tmp/dg_jit_cache \
  --force

6. 当前覆盖情况

当前 release 覆盖元数据放在：

deepgemm-packaged-precompiled/deep_gemm/precompiled/release_coverage.json

现阶段已经有明确证据的条目：

• H200 单卡 Qwen3-8B-FP8：verify JSON 可直接解析，cold=0
• H200 单卡 Qwen3-14B-FP8：two-wheel 手工记录，cold=0
• H200 单卡 Qwen3.5-35B-A3B-FP8：报告记录显示 cold=0

仍待补齐：

• H200 TP=2
• H100 单卡
• DPSK / DeepSeek-V3.2 cook workload

7. 现在的边界

本次已经把“只交付 deep_gemm wheel”的源码、默认 bundle 入口、bundle 生成脚本、包内覆盖元数据都落到仓库里了。

但下面这部分仍然需要后续实机补充，不能伪造：

• H200 TP=2 fresh verify
• H100 fresh verify
• DPSK / DeepSeek-V3.2 实际 cook 模型路径 collect/verify

换句话说，现在仓库已经具备了 wheel-only 交付骨架；剩下是继续把新 cache 和新 verify 结果灌进同一个 union bundle 里。

预编译目录与 cache 设计

1. 背景与目标

当前行为

调用 fp8_gemm_nt(m, n, k, ...) 
  → 构造 GemmDesc
  → get_best_config<ArchSpec>(desc)   ← 每次重新算，无缓存
      → 遍历所有 layout candidates
      → 对每个 candidate 调用 get_layout_info() 计算 wave/cycle 估算
      → 选出最优 layout
      → 推导 storage/pipeline/launch config
  → compiler->build(config)           ← 已有 cubin 磁盘缓存

问题

get_best_config 每次都重新遍历 layout candidates 计算，相同 shape 多次调用时完全重复劳动。

目标

为 get_best_config 加一层持久化缓存：

get_best_config(desc)
  → 先查缓存（内存 map）
    → 命中：直接返回 GemmConfig ✓
    → 未命中：走原有逻辑算出 config → 写入缓存 → 返回

2. 数据结构分析

缓存 Key：GemmDesc

GemmDesc 中决定 config 结果的字段（与 shape 和运行环境相关）：

gemm_type, kernel_type, m, n, k, num_groups
a_dtype, b_dtype, cd_dtype
major_a, major_b
with_accumulation
num_sms, tc_util
compiled_dims
expected_m, expected_n, expected_k, expected_num_groups

已有 operator<< 可直接序列化为唯一字符串，用作 key。

缓存 Value：GemmConfig

Layout:         swap_ab, block_m, block_n, block_k, cluster_m, cluster_n   (6 int)
StorageConfig:  load_block_m, load_block_n, store_block_m, store_block_n,
                swizzle_a_mode, swizzle_b_mode, swizzle_cd_mode             (7 int)
PipelineConfig: smem_size, num_stages                                        (2 int)
LaunchConfig:   num_sms, num_sms_per_cluster, num_threads,
                num_tma_threads, num_math_threads,
                num_non_epilogue_threads, num_epilogue_threads               (7 int)

总共 22 个 int，序列化为空格分隔的字符串即可，无需外部依赖。

3. 实现方案

3.1 新增文件：csrc/jit/config_cache.hpp

class ConfigCache {
    // 内存缓存
    std::unordered_map<std::string, GemmConfig> cache;
    // 磁盘缓存路径
    std::filesystem::path cache_file;
    std::mutex mtx;  // 多线程安全

public:
    ConfigCache();                                        // 读取磁盘缓存到内存
    std::optional<GemmConfig> get(const std::string& key);
    void put(const std::string& key, const GemmConfig& config);  // 同时写磁盘

private:
    static std::string serialize(const GemmConfig& config);
    static GemmConfig deserialize(const std::string& str);
};

磁盘缓存文件：$DG_JIT_CACHE_DIR/config_cache.txt（默认 ~/.deep_gemm/config_cache.txt）

文件格式（每行一条记录，| 分隔 key 和 value）：

GemmDesc(gemm_type=0,...,expected_k=0)|0 128 128 64 1 1 0 128 0 0 128 0 1 2 49152 4 ...

3.2 修改文件：csrc/jit_kernels/heuristics/common.hpp

在 get_best_config() 开头加缓存查询，末尾加缓存写入：

template <typename ArchSpec>
static GemmConfig get_best_config(const GemmDesc& desc) {
    desc.check_validity();

    // ── 新增：查缓存 ──────────────────────────────
    std::ostringstream oss;
    oss << desc;
    const auto key = oss.str();
    if (const auto cached = config_cache->get(key); cached.has_value())
        return cached.value();
    // ─────────────────────────────────────────────

    // 原有逻辑（不动）
    const auto layout_candidates = ArchSpec::get_layout_candidates(desc);
    ...
    const auto gemm_config = GemmConfig { ... };

    // ── 新增：存缓存 ──────────────────────────────
    config_cache->put(key, gemm_config);
    // ─────────────────────────────────────────────

    return gemm_config;
}

3.3 暴露给 Python（可选）：csrc/apis/runtime.hpp

m.def("clear_config_cache", []() { config_cache->clear(); });
m.def("get_config_cache_size", []() { return config_cache->size(); });
m.def("set_config_cache_dir", [](const std::string& path) {
    config_cache->set_cache_dir(path);
});

4. 涉及改动的文件

5. 工作流

仓库结构

origin   → personal fork（地址省略）
upstream → team repository（地址省略）

开发循环

CPU 机器（CC 辅助写代码）
  → git push origin feature/config-cache
GPU 机器
  → git pull
  → pip install -e .  / python setup.py develop
  → python tests/test_fp8.py  验证正确性
  → 验证缓存是否命中（DG_PRINT_CONFIGS=1）

6. 待验证

• [ ] GPU 机器能否正常编译安装 deepgemm（pip install -e .）
• [ ] 并发场景下缓存写入的线程安全
• [ ] 缓存文件在多机分布式下的共享策略（NFS / 各机独立）
• [ ] num_sms 是运行时动态值，是否应纳入 key（不同机器 SM 数不同）