diff --git a/python/sglang/compile_deep_gemm.py b/python/sglang/compile_deep_gemm.py index aa6f0703c..77ddbadce 100644 --- a/python/sglang/compile_deep_gemm.py +++ b/python/sglang/compile_deep_gemm.py @@ -33,7 +33,7 @@ envs.SGLANG_IN_DEEPGEMM_PRECOMPILE_STAGE.set(True) # Force enable deep gemm envs.SGLANG_ENABLE_JIT_DEEPGEMM.set(True) # Force enable mha chunked kv for DeepSeek V3 to avoid missing kv_b_proj DeepGEMM case -os.environ["SGLANG_CHUNKED_PREFIX_CACHE_THRESHOLD"] = "0" +envs.SGLANG_CHUNKED_PREFIX_CACHE_THRESHOLD.set(0) @dataclasses.dataclass diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 47f570821..0a687e250 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -288,6 +288,7 @@ class Envs: SGLANG_DG_CACHE_DIR = EnvStr(os.path.expanduser("~/.cache/deep_gemm")) SGLANG_DG_USE_NVRTC = EnvBool(False) SGLANG_USE_DEEPGEMM_BMM = EnvBool(False) + SGLANG_CHUNKED_PREFIX_CACHE_THRESHOLD = EnvInt(8192) # DeepEP SGLANG_DEEPEP_BF16_DISPATCH = EnvBool(False) diff --git a/python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py b/python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py index fee23cc49..dca005dd1 100644 --- a/python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py +++ b/python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py @@ -14,7 +14,7 @@ from sglang.srt.distributed.device_communicators.pynccl_allocator import ( from sglang.srt.environ import envs from sglang.srt.layers.deep_gemm_wrapper.configurer import ENABLE_JIT_DEEPGEMM from sglang.srt.server_args import ServerArgs -from sglang.srt.utils import ceil_div, get_available_gpu_memory, get_bool_env_var +from sglang.srt.utils import ceil_div, get_available_gpu_memory logger = logging.getLogger(__name__) @@ -25,8 +25,8 @@ if ENABLE_JIT_DEEPGEMM: _BUILTIN_M_LIST = list(range(1, 1024 * 16 + 1)) _ENABLE_JIT_DEEPGEMM_PRECOMPILE = envs.SGLANG_JIT_DEEPGEMM_PRECOMPILE.get() _DO_COMPILE_ALL = True -_IS_FIRST_RANK_ON_NODE = get_bool_env_var("SGLANG_IS_FIRST_RANK_ON_NODE", "true") -_IN_PRECOMPILE_STAGE = get_bool_env_var("SGLANG_IN_DEEPGEMM_PRECOMPILE_STAGE", "false") +_IS_FIRST_RANK_ON_NODE = envs.SGLANG_IS_FIRST_RANK_ON_NODE.get() +_IN_PRECOMPILE_STAGE = envs.SGLANG_IN_DEEPGEMM_PRECOMPILE_STAGE.get() # Force redirect deep_gemm cache_dir os.environ["DG_JIT_CACHE_DIR"] = os.getenv( diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 4fd8376e4..b0e6ec1af 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -148,7 +148,6 @@ from sglang.srt.utils import ( cpu_has_amx_support, get_bool_env_var, get_device_sm, - get_int_env_var, is_cpu, is_cuda, is_gfx95_supported, @@ -1466,8 +1465,8 @@ class DeepseekV2AttentionMLA(nn.Module): ) # TODO: Design a finer way to determine the threshold - self.chunked_prefix_cache_threshold = get_int_env_var( - "SGLANG_CHUNKED_PREFIX_CACHE_THRESHOLD", 8192 + self.chunked_prefix_cache_threshold = ( + envs.SGLANG_CHUNKED_PREFIX_CACHE_THRESHOLD.get() ) # If we have self.fused_qkv_a_proj_with_mqa and we're running on CPU, we will choose the torch.ops.sgl_kernel.qkv_proj_with_rope_fused_weight kernel