diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index 5406522c1..83d4f527f 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -61,6 +61,7 @@ def is_deepseek_nsa(config: PretrainedConfig) -> bool: "DeepseekV3ForCausalLMNextN", "MistralLarge3ForCausalLM", "PixtralForConditionalGeneration", + "GlmMoeDsaForCausalLM", ] and getattr(config, "index_topk", None) is not None ) @@ -271,10 +272,10 @@ class ModelConfig: def _config_draft_model(self): is_draft_model = self.is_draft_model - if ( - is_draft_model - and self.hf_config.architectures[0] == "DeepseekV3ForCausalLM" - ): + if is_draft_model and self.hf_config.architectures[0] in [ + "DeepseekV3ForCausalLM", + "GlmMoeDsaForCausalLM", + ]: self.hf_config.architectures[0] = "DeepseekV3ForCausalLMNextN" if is_draft_model and self.hf_config.architectures[0] in [ @@ -411,7 +412,6 @@ class ModelConfig: "swa_v_head_dim", self.v_head_dim, ) - # FIXME: temporary special judge for MLA architecture if ( "DeepseekV2ForCausalLM" in self.hf_config.architectures @@ -419,6 +419,7 @@ class ModelConfig: or "DeepseekV3ForCausalLM" in self.hf_config.architectures or "DeepseekV3ForCausalLMNextN" in self.hf_config.architectures or "Glm4MoeLiteForCausalLM" in self.hf_config.architectures + or "GlmMoeDsaForCausalLM" in self.hf_config.architectures or "LongcatFlashForCausalLM" in self.hf_config.architectures or "LongcatFlashForCausalLMNextN" in self.hf_config.architectures or "DotsVLMForCausalLM" in self.hf_config.architectures diff --git a/python/sglang/srt/models/glm4_moe.py b/python/sglang/srt/models/glm4_moe.py index 3954f501f..e3eeb7d10 100644 --- a/python/sglang/srt/models/glm4_moe.py +++ b/python/sglang/srt/models/glm4_moe.py @@ -79,6 +79,7 @@ from sglang.srt.layers.vocab_parallel_embedding import ( from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors from sglang.srt.model_loader.weight_utils import default_weight_loader +from sglang.srt.models.deepseek_v2 import DeepseekV2ForCausalLM from sglang.srt.models.utils import apply_qk_norm from sglang.srt.server_args import get_global_server_args from sglang.srt.utils import ( @@ -1279,4 +1280,8 @@ class Glm4MoeForCausalLM(nn.Module): self.model.layers_to_capture = [val + 1 for val in layer_ids] -EntryClass = [Glm4MoeForCausalLM] +class GlmMoeDsaForCausalLM(DeepseekV2ForCausalLM): + pass + + +EntryClass = [Glm4MoeForCausalLM, GlmMoeDsaForCausalLM] diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index ecee665ad..a688911c2 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -1194,9 +1194,15 @@ class ServerArgs: "KimiK25ForConditionalGeneration", "MistralLarge3ForCausalLM", "PixtralForConditionalGeneration", + "GlmMoeDsaForCausalLM", ]: # Set attention backend for DeepSeek - if is_deepseek_nsa(hf_config): # DeepSeek 3.2 + if is_deepseek_nsa(hf_config): # DeepSeek 3.2, GlmMoeDsaForCausalLM + if model_arch == "GlmMoeDsaForCausalLM" and is_blackwell_supported(): + envs.SGLANG_NSA_FORCE_MLA.set(True) + logger.warning( + "Force NSA prefill to use MLA (i.e. disable MHA_ONE_SHOT) for GlmMoeDsaForCausalLM on SM100." + ) if self.is_attention_backend_not_set(): self.attention_backend = "nsa" logger.info("Use nsa attention backend for DeepSeek with DSA.") @@ -2323,6 +2329,7 @@ class ServerArgs: "DeepseekV3ForCausalLM", "Glm4MoeForCausalLM", "Glm4MoeLiteForCausalLM", + "GlmMoeDsaForCausalLM", "BailingMoeForCausalLM", "BailingMoeV2ForCausalLM", "MistralLarge3ForCausalLM", @@ -2652,6 +2659,7 @@ class ServerArgs: "DeepseekV32ForCausalLM", "MistralLarge3ForCausalLM", "PixtralForConditionalGeneration", + "GlmMoeDsaForCausalLM", ] except Exception: pass @@ -5648,6 +5656,7 @@ def auto_choose_speculative_params(self: ServerArgs): "GptOssForCausalLM", "Glm4MoeForCausalLM", "Glm4MoeLiteForCausalLM", + "GlmMoeDsaForCausalLM", "BailingMoeForCausalLM", "BailingMoeV2ForCausalLM", "MistralLarge3ForCausalLM",