Defer DeepGEMM PDL setup to worker init (#27671)

Co-authored-by: lmzheng <lmzheng@fb.com>

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)