diff --git a/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py b/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py index 561f7ef43..62bf02555 100644 --- a/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py +++ b/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py @@ -2004,13 +2004,46 @@ class AscendAttnBackend(AttentionBackend): self.speculative_num_draft_tokens, ) + if ( + self.q_head_num_padding is not None + and self.q_head_num_padding > self.tp_q_head_num + ): + nope_padding = torch.empty( + [ + q_nope.shape[0], + self.q_head_num_padding - self.tp_q_head_num, + self.kv_lora_rank, + ], + dtype=( + self.model_dtype + if self.model_dtype is not None + else torch.bfloat16 + ), + device=q_nope.device, + ) + rope_padding = torch.empty( + [ + q_rope.shape[0], + self.q_head_num_padding - self.tp_q_head_num, + self.qk_rope_head_dim, + ], + dtype=( + self.model_dtype + if self.model_dtype is not None + else torch.bfloat16 + ), + device=q_rope.device, + ) + q_nope = torch.cat([q_nope, nope_padding], dim=1).contiguous() + q_rope = torch.cat([q_rope, rope_padding], dim=1).contiguous() + workspace = torch_npu._npu_fused_infer_attention_score_get_max_workspace( q_nope, c_kv_cache, c_kv_cache, query_rope=q_rope, key_rope=k_rope_cache, - num_heads=layer.tp_q_head_num, + num_heads=self.q_head_num_padding, num_key_value_heads=layer.tp_k_head_num, input_layout="TND", scale=layer.scaling, @@ -2031,7 +2064,7 @@ class AscendAttnBackend(AttentionBackend): c_kv_cache, query_rope=q_rope, key_rope=k_rope_cache, - num_heads=layer.tp_q_head_num, + num_heads=self.q_head_num_padding, num_key_value_heads=layer.tp_k_head_num, input_layout="TND", scale=layer.scaling, @@ -2046,6 +2079,7 @@ class AscendAttnBackend(AttentionBackend): workspace=workspace, out=[attn_output, softmax_lse], ) + attn_output = attn_output[:, : layer.tp_q_head_num, :] attn_output = attn_output.view(-1, layer.tp_q_head_num * layer.v_head_dim) if ( not self.graph_mode diff --git a/python/sglang/srt/models/glm4_moe_lite.py b/python/sglang/srt/models/glm4_moe_lite.py index 7530b4b5d..e821fe89f 100644 --- a/python/sglang/srt/models/glm4_moe_lite.py +++ b/python/sglang/srt/models/glm4_moe_lite.py @@ -547,6 +547,8 @@ class Glm4MoeLiteDecoderLayer(nn.Module): ) -> None: super().__init__() + # Required for MTP: Glm4MoeLiteModelNextN bypasses Glm4MoeLiteForCausalLM.__init__ + config.moe_layer_freq = 1 self.hidden_size = config.hidden_size self.config = config rope_theta, rope_scaling = get_rope_config(config)