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

This commit is contained in:
Baizhou Zhang
2026-06-03 17:16:54 -07:00
committed by GitHub
parent 3790173b3b
commit d097cd2212
+10 -16
View File
@@ -1831,26 +1831,20 @@ class ServerArgs:
]:
# Set attention backend for DeepSeek
if is_deepseek_dsa(hf_config): # DeepSeek 3.2/GLM 5
if model_arch == "GlmMoeDsaForCausalLM" and is_blackwell_supported():
envs.SGLANG_DSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD.set(0)
if envs.SGLANG_DSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD.is_set():
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:
if envs.SGLANG_DSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD.is_set():
logger.warning(
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
# 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(
get_dsa_index_topk(hf_config)
)
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."
)
envs.SGLANG_DSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD.set(
get_dsa_index_topk(hf_config)
)
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."
)
if self.is_attention_backend_not_set():
self.attention_backend = "dsa"
logger.info("Use dsa attention backend for DeepSeek with DSA.")