minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Tiny env cleanup in deepgemm (#15706)

Browse Source 
Signed-off-by: vincentzed <207368749+vincentzed@users.noreply.github.com>
This commit is contained in:
vincentzed
2025-12-24 00:58:46 -08:00
committed by GitHub
parent f65fa04748
commit ddb3970eb4
4 changed files with 7 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/compile_deep_gemm.py
+1 -1
View File
@@ -33,7 +33,7 @@ envs.SGLANG_IN_DEEPGEMM_PRECOMPILE_STAGE.set(True)
# Force enable deep gemm # Force enable deep gemm
envs.SGLANG_ENABLE_JIT_DEEPGEMM.set(True) envs.SGLANG_ENABLE_JIT_DEEPGEMM.set(True)
# Force enable mha chunked kv for DeepSeek V3 to avoid missing kv_b_proj DeepGEMM case # 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 @dataclasses.dataclass
python/sglang/srt/environ.py
+1
View File
@@ -288,6 +288,7 @@ class Envs:
SGLANG_DG_CACHE_DIR = EnvStr(os.path.expanduser("~/.cache/deep_gemm")) SGLANG_DG_CACHE_DIR = EnvStr(os.path.expanduser("~/.cache/deep_gemm"))
SGLANG_DG_USE_NVRTC = EnvBool(False) SGLANG_DG_USE_NVRTC = EnvBool(False)
SGLANG_USE_DEEPGEMM_BMM = EnvBool(False) SGLANG_USE_DEEPGEMM_BMM = EnvBool(False)
SGLANG_CHUNKED_PREFIX_CACHE_THRESHOLD = EnvInt(8192)
# DeepEP # DeepEP
SGLANG_DEEPEP_BF16_DISPATCH = EnvBool(False) SGLANG_DEEPEP_BF16_DISPATCH = EnvBool(False)
python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py
+3 -3
View File
@@ -14,7 +14,7 @@ from sglang.srt.distributed.device_communicators.pynccl_allocator import (
from sglang.srt.environ import envs from sglang.srt.environ import envs
from sglang.srt.layers.deep_gemm_wrapper.configurer import ENABLE_JIT_DEEPGEMM from sglang.srt.layers.deep_gemm_wrapper.configurer import ENABLE_JIT_DEEPGEMM
from sglang.srt.server_args import ServerArgs 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__) logger = logging.getLogger(__name__)
@@ -25,8 +25,8 @@ if ENABLE_JIT_DEEPGEMM:
_BUILTIN_M_LIST = list(range(1, 1024 * 16 + 1)) _BUILTIN_M_LIST = list(range(1, 1024 * 16 + 1))
_ENABLE_JIT_DEEPGEMM_PRECOMPILE = envs.SGLANG_JIT_DEEPGEMM_PRECOMPILE.get() _ENABLE_JIT_DEEPGEMM_PRECOMPILE = envs.SGLANG_JIT_DEEPGEMM_PRECOMPILE.get()
_DO_COMPILE_ALL = True _DO_COMPILE_ALL = True
_IS_FIRST_RANK_ON_NODE = get_bool_env_var("SGLANG_IS_FIRST_RANK_ON_NODE", "true") _IS_FIRST_RANK_ON_NODE = envs.SGLANG_IS_FIRST_RANK_ON_NODE.get()
_IN_PRECOMPILE_STAGE = get_bool_env_var("SGLANG_IN_DEEPGEMM_PRECOMPILE_STAGE", "false") _IN_PRECOMPILE_STAGE = envs.SGLANG_IN_DEEPGEMM_PRECOMPILE_STAGE.get()
# Force redirect deep_gemm cache_dir # Force redirect deep_gemm cache_dir
os.environ["DG_JIT_CACHE_DIR"] = os.getenv( os.environ["DG_JIT_CACHE_DIR"] = os.getenv(
python/sglang/srt/models/deepseek_v2.py
+2 -3
View File
@@ -148,7 +148,6 @@ from sglang.srt.utils import (
cpu_has_amx_support, cpu_has_amx_support,
get_bool_env_var, get_bool_env_var,
get_device_sm, get_device_sm,
get_int_env_var,
is_cpu, is_cpu,
is_cuda, is_cuda,
is_gfx95_supported, is_gfx95_supported,
@@ -1466,8 +1465,8 @@ class DeepseekV2AttentionMLA(nn.Module):
) )
# TODO: Design a finer way to determine the threshold # TODO: Design a finer way to determine the threshold
self.chunked_prefix_cache_threshold = get_int_env_var( self.chunked_prefix_cache_threshold = (
"SGLANG_CHUNKED_PREFIX_CACHE_THRESHOLD", 8192 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 # 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