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

fix(compile_utils, ep_moe): update environment variable and dtype check (#12034)

Browse Source
This commit is contained in:
ishandhanani
2025-10-24 11:00:12 -07:00
committed by GitHub
parent d7f0d88fa2
commit 14203432b4
3 changed files with 6 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py
View File

@@ -26,7 +26,7 @@ _IN_PRECOMPILE_STAGE = get_bool_env_var("SGL_IN_DEEPGEMM_PRECOMPILE_STAGE", "fal
# Force redirect deep_gemm cache_dir
os.environ["DG_JIT_CACHE_DIR"] = os.getenv(
"SGL_DG_CACHE_DIR", os.path.join(os.path.expanduser("~"), ".cache", "deep_gemm")
"SGLANG_DG_CACHE_DIR", os.path.join(os.path.expanduser("~"), ".cache", "deep_gemm")
)
# Refer to https://github.com/deepseek-ai/DeepGEMM/commit/d75b218b7b8f4a5dd5406ac87905039ead3ae42f

7
python/sglang/srt/layers/moe/ep_moe/layer.py
View File

@@ -440,9 +440,10 @@ class DeepEPMoE(FusedMoE):
hidden_states, hidden_states_scale, _, _, masked_m, expected_m = dispatch_output
assert self.quant_method is not None
assert self.moe_runner_config.activation == "silu"
assert (
hidden_states_scale.dtype == torch.float32
), f"hidden_states_scale.dtype: {hidden_states_scale.dtype}"
assert hidden_states_scale.dtype == torch.float32 or (
deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0
and hidden_states_scale.dtype == torch.int32
), f"hidden_states_scale.dtype: {hidden_states_scale.dtype}, DEEPGEMM_SCALE_UE8M0: {deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0}"
# GroupGemm-0
num_groups, m, k = hidden_states.size()