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

[GLM-5] Apply trtllm MHA kernel for GLM-5 on Blackwell (#21332)

Browse Source
This commit is contained in:
Baizhou Zhang
2026-06-03 17:16:54 -07:00
committed by GitHub
parent 3790173b3b
commit d097cd2212
1 changed files with 10 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/server_args.py
+10 -16
View File
@@ -1831,26 +1831,20 @@ class ServerArgs:
]: ]:
# Set attention backend for DeepSeek # Set attention backend for DeepSeek
if is_deepseek_dsa(hf_config): # DeepSeek 3.2/GLM 5 if is_deepseek_dsa(hf_config): # DeepSeek 3.2/GLM 5
if model_arch == "GlmMoeDsaForCausalLM" and is_blackwell_supported(): if envs.SGLANG_DSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD.is_set():
envs.SGLANG_DSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD.set(0)
logger.warning( logger.warning(
"Force DSA prefill to use sparse MLA (i.e. disable MHA_ONE_SHOT) for GlmMoeDsaForCausalLM on Blackwell." f"Dense attention kv len threshold is manually set to {envs.SGLANG_DSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD.get()} for DSA. Caution: This may cause performance regression if the threshold is larger than the index topk of model."
) )
else: else:
if envs.SGLANG_DSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD.is_set(): # When threshold is not manually set, set it to the index topk of model
logger.warning( from sglang.srt.configs.model_config import get_dsa_index_topk
f"Dense attention kv len threshold is manually set to {envs.SGLANG_DSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD.get()} for DSA. Caution: This may cause performance regression if the threshold is larger than the index topk of model."
)
else:
# When threshold is not manually set, set it to the index topk of model
from sglang.srt.configs.model_config import get_dsa_index_topk
envs.SGLANG_DSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD.set( envs.SGLANG_DSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD.set(
get_dsa_index_topk(hf_config) get_dsa_index_topk(hf_config)
) )
logger.warning( logger.warning(
f"Set dense attention kv len threshold to model index_topk={envs.SGLANG_DSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD.get()} for DeepSeek with DSA." f"Set dense attention kv len threshold to model index_topk={envs.SGLANG_DSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD.get()} for DeepSeek with DSA."
) )
if self.is_attention_backend_not_set(): if self.is_attention_backend_not_set():
self.attention_backend = "dsa" self.attention_backend = "dsa"
logger.info("Use dsa attention backend for DeepSeek with DSA.") logger.info("Use dsa attention backend for DeepSeek with DSA.")