From 7769ff8f1e9f70680ee6fb5cdf17672d4aa95666 Mon Sep 17 00:00:00 2001 From: Eric Zhang Date: Tue, 25 Aug 2026 01:58:03 -0400 Subject: [PATCH] [DeepGEMM] Deduplicate JIT precompile across local ranks (#36231) Co-authored-by: Nikita Shulga --- .../layers/deep_gemm_wrapper/compile_utils.py | 40 +++++++++++++++---- 1 file changed, 33 insertions(+), 7 deletions(-) diff --git a/python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py b/python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py index 8facfaca5..1e2db45ff 100644 --- a/python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py +++ b/python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py @@ -1,9 +1,11 @@ +import fcntl import logging import math import os import time from contextlib import contextmanager, nullcontext from enum import IntEnum, auto +from pathlib import Path from typing import Dict, List, Tuple import torch @@ -115,6 +117,29 @@ class DeepGemmKernelType(IntEnum): _INITIALIZATION_DICT: Dict[Tuple[DeepGemmKernelType, int, int, int], bool] = dict() +@contextmanager +def _local_rank_compile_lock( + kernel_type: DeepGemmKernelType, n: int, k: int, num_groups: int +): + """Serialize one pre-compile group across the ranks sharing DG_JIT_CACHE_DIR. + + Every rank lazily walks the same (kernel_type, n, k, num_groups) groups in + the same order, so without a lock N local ranks nvcc-compile N identical + copies of every kernel. The lock holder compiles into the shared cache; + waiters then find the cubins already present and their pass over the M + list is execution warmup only. + """ + lock_dir = Path(os.environ["DG_JIT_CACHE_DIR"]) / "locks" + lock_dir.mkdir(parents=True, exist_ok=True) + 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) + + # TODO improve code def _maybe_compile_deep_gemm_one_type_all( kernel_type: DeepGemmKernelType, @@ -151,13 +176,14 @@ def _maybe_compile_deep_gemm_one_type_all( f"{' It only takes a little time (typically 1 sec) if you have run `python3 -m sglang.compile_deep_gemm`. ' if not _IN_PRECOMPILE_STAGE else ''}" ) - _compile_deep_gemm_one_type_all( - kernel_type=kernel_type, - n=n, - k=k, - num_groups=num_groups, - m_list=_BUILTIN_M_LIST, - ) + with _local_rank_compile_lock(kernel_type, n, k, num_groups): + _compile_deep_gemm_one_type_all( + kernel_type=kernel_type, + n=n, + k=k, + num_groups=num_groups, + m_list=_BUILTIN_M_LIST, + ) # NOTE(alcanderian): get_num_sms should be change when 2-batch-overlap is introduced