[DeepGEMM] Deduplicate JIT precompile across local ranks (#36231)

Co-authored-by: Nikita Shulga <nikita@thinkingmachines.ai>

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(
+        with _local_rank_compile_lock(kernel_type, n, k, num_groups):
-            kernel_type=kernel_type,
+            _compile_deep_gemm_one_type_all(
-            n=n,
+                kernel_type=kernel_type,
-            k=k,
+                n=n,
-            num_groups=num_groups,
+                k=k,
-            m_list=_BUILTIN_M_LIST,
+                num_groups=num_groups,
-        )
+                m_list=_BUILTIN_M_LIST,
+            )
 
 
 # NOTE(alcanderian): get_num_sms should be change when 2-batch-overlap is introduced