minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Defer DeepGEMM PDL setup to worker init (#27671)

Browse Source 
Co-authored-by: lmzheng <lmzheng@fb.com>
This commit is contained in:
Lianmin Zheng
2026-06-09 13:52:30 -07:00
committed by GitHub
co-authored by lmzheng
parent 186f1e300a
commit eb8dceda44
1 changed files with 5 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/deep_gemm_wrapper/entrypoint.py
+5 -3
View File
@@ -20,9 +20,6 @@ if ENABLE_JIT_DEEPGEMM:
import deep_gemm import deep_gemm
from deep_gemm.utils.layout import get_mn_major_tma_aligned_tensor # noqa: F401 from deep_gemm.utils.layout import get_mn_major_tma_aligned_tensor # noqa: F401
if envs.SGLANG_DEEPGEMM_PDL.get() and hasattr(deep_gemm, "set_pdl"):
deep_gemm.set_pdl(True)
_SANITY_CHECK = envs.SGLANG_DEEPGEMM_SANITY_CHECK.get() _SANITY_CHECK = envs.SGLANG_DEEPGEMM_SANITY_CHECK.get()
@@ -208,6 +205,11 @@ def tf32_hc_prenorm_gemm(
def update_deep_gemm_config(gpu_id: int, server_args: ServerArgs): def update_deep_gemm_config(gpu_id: int, server_args: ServerArgs):
# deep_gemm.set_pdl can initialize CUDA state, so run it only after the
# scheduler/TP worker has been forked and assigned a GPU.
if envs.SGLANG_DEEPGEMM_PDL.get() and hasattr(deep_gemm, "set_pdl"):
deep_gemm.set_pdl(True)
compile_utils.update_deep_gemm_config(gpu_id, server_args) compile_utils.update_deep_gemm_config(gpu_id, server_args)