From eb8dceda44a50c89bb640a454fa0f937c59aa26b Mon Sep 17 00:00:00 2001 From: Lianmin Zheng Date: Tue, 9 Jun 2026 13:52:30 -0700 Subject: [PATCH] Defer DeepGEMM PDL setup to worker init (#27671) Co-authored-by: lmzheng --- python/sglang/srt/layers/deep_gemm_wrapper/entrypoint.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/python/sglang/srt/layers/deep_gemm_wrapper/entrypoint.py b/python/sglang/srt/layers/deep_gemm_wrapper/entrypoint.py index 2cefd9461..3d967f1e1 100644 --- a/python/sglang/srt/layers/deep_gemm_wrapper/entrypoint.py +++ b/python/sglang/srt/layers/deep_gemm_wrapper/entrypoint.py @@ -20,9 +20,6 @@ if ENABLE_JIT_DEEPGEMM: import deep_gemm 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() @@ -208,6 +205,11 @@ def tf32_hc_prenorm_gemm( 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)