[GLM-5] Apply trtllm MHA kernel for GLM-5 on Blackwell (#21332)
This commit is contained in:
@@ -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.")
|
||||
|
||||
Reference in New Issue
Block a user