From d097cd2212d7ec0fef03d2679ddbb8008973b4a7 Mon Sep 17 00:00:00 2001 From: Baizhou Zhang Date: Wed, 3 Jun 2026 17:16:54 -0700 Subject: [PATCH] [GLM-5] Apply trtllm MHA kernel for GLM-5 on Blackwell (#21332) --- python/sglang/srt/server_args.py | 26 ++++++++++---------------- 1 file changed, 10 insertions(+), 16 deletions(-) diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index c0abb6fce..f27612cc2 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -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.")