diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index de4264c95..2eec4cbaf 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -483,6 +483,7 @@ class Envs: SGLANG_DG_USE_NVRTC = EnvBool(False) SGLANG_USE_DEEPGEMM_BMM = EnvBool(False) SGLANG_DEEPGEMM_SANITY_CHECK = EnvBool(False) + SGLANG_DEEPGEMM_PDL = EnvBool(True) SGLANG_PP_PARALLEL_DEEPGEMM_WARMUP = EnvBool(False) # DeepSeek MHA Optimization diff --git a/python/sglang/srt/layers/deep_gemm_wrapper/entrypoint.py b/python/sglang/srt/layers/deep_gemm_wrapper/entrypoint.py index 08fa2159e..2cefd9461 100644 --- a/python/sglang/srt/layers/deep_gemm_wrapper/entrypoint.py +++ b/python/sglang/srt/layers/deep_gemm_wrapper/entrypoint.py @@ -20,6 +20,9 @@ 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()