Prhub

#36231 [DeepGEMM] Deduplicate JIT precompile across local ranks

原始 PR 作者 ekzhang 合并时间 2026-08-25 13:58 文件变更 1 提交数 2 评论 2 代码增减 +33 / -7

执行摘要

跨本地 rank 去重 DeepGEMM JIT 预编译,减少重复 NVCC 编译。

当多个本地 rank 共享 DG_JIT_CACHE_DIR 时,它们会惰性遍历相同的 (kernel_type, n, k, num_groups) 预编译组,可能并发编译完全相同的 DeepGEMM kernel,导致大量重复的 NVCC 编译工作。本 PR 旨在通过跨 rank 的编译锁消除这种重复,加快多 rank 启动速度。

值得精读,该 PR 展示了通过简单文件锁解决多进程并发编译去重的思路,设计简洁,对理解多 rank 启动优化有参考价值。建议关注锁粒度和异常处理,未来可考虑使用更细粒度的锁或异步编译。

讨论亮点

无 review 评论。Approver Fridge003 直接批准,未提出额外问题。

实现拆解

  1. 在 compile_utils.py 中新增 _local_rank_compile_lock 上下文管理器,利用 fcntl.flock 实现跨进程文件锁。
  2. 锁文件存放在 DG_JIT_CACHE_DIR/locks 目录下,文件名为 kernel 组唯一标识。
  3. 在 _maybe_compile_deep_gemm_one_type_all 中,将编译调用包裹在锁内,确保同一组 kernel 仅由一个 rank 编译,其他 rank 等待后复用缓存。
  4. 未修改任何测试文件,但该改动不改变编译结果,仅影响编译并发性。
文件 模块 状态 重要度
python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py 编译工具 modified 7.13

关键符号

_local_rank_compile_lock

关键源码片段

python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py core-logic

核心改动文件,新增文件锁机制控制 JIT 预编译并发。

# python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py
# 新增上下文管理器,用于序列化跨 rank 的编译组
@contextmanager
def local_rank_compile_lock(
    kernel_type: DeepGemmKernelType, n: int, k: int, num_groups: int
):
    """
    序列化一个预编译组,跨共享 DG_JIT_CACHE_DIR 的多个本地 rank。    每个 rank 惰性遍历相同的 (kernel_type, n, k, num_groups) 组,
    不锁时 N 个 rank 会对每个 kernel 编译 N 份副本。
    锁持有者将编译产物写入共享缓存;
    等待者发现 cubin 已存在后,仅执行 M 列表的 warmup 过程。
    """
    # 锁文件放在缓存目录下的 locks 子目录,确保各 rank 可见
    lock_dir = Path(os.environ["DG_JIT_CACHE_DIR"]) / "locks"
    lock_dir.mkdir(parents=True, exist_ok=True)
    # 锁文件名包含 kernel 类型和形状参数,保证不同组之间互不阻塞
    lock_path = lock_dir / f"{kernel_type.name}_n{n}_k{k}_g{num_groups}.lock"
    with open(lock_path, "w") as lock_file:
        fcntl.flock(lock_file, fcntl.LOCK_EX) # 获取排他锁
        try:
            yield
        finally:
            fcntl.flock(lock_file, fcntl.LOCK_UN) # 释放排他锁

评论区精华

没有提炼出高价值讨论线程

当前评论区没有形成足够清晰的争议点或结论,后续有更多讨论时会体现在这里。

风险与影响

  1. 锁机制依赖 DG_JIT_CACHE_DIR 环境变量,若未设置可能导致 KeyError;但该变量本应存在,风险较低。
  2. 锁粒度较粗,仅按 kernel 组加锁,多个 rank 仍可能同时执行 warmup,但编译部分已串行,整体争用有限。
  3. 如果编译过程中进程崩溃,锁文件可能残留,但 fcntl.flock 锁会在进程退出时自动释放,且下次启动会重新创建文件,风险可忽略。
  4. 无测试覆盖多 rank 并发场景,但该改动不改变编译结果,回归风险较低。

影响范围限于 DeepGEMM JIT 预编译启动路径,主要影响多 rank 场景下的启动时间,对单 rank 或多个不共享缓存目录的场景无影响。用户可感知到多卡部署时启动时间缩短,稳态推理性能不变。对团队而言,减少了并发编译的资源损耗,提升了集群利用率。

无测试覆盖 依赖环境变量 锁粒度粗

关联 Issue

未识别关联 Issue

当前没有检测到明确关联的 Issue 链接,后续同步到相关引用后会出现在这里。

完整报告

参与讨论