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)