[NPU] Add MTP support for GLM-4.7-Flash (#28516)
This commit is contained in:
@@ -2004,13 +2004,46 @@ class AscendAttnBackend(AttentionBackend):
|
|||||||
self.speculative_num_draft_tokens,
|
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(
|
workspace = torch_npu._npu_fused_infer_attention_score_get_max_workspace(
|
||||||
q_nope,
|
q_nope,
|
||||||
c_kv_cache,
|
c_kv_cache,
|
||||||
c_kv_cache,
|
c_kv_cache,
|
||||||
query_rope=q_rope,
|
query_rope=q_rope,
|
||||||
key_rope=k_rope_cache,
|
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,
|
num_key_value_heads=layer.tp_k_head_num,
|
||||||
input_layout="TND",
|
input_layout="TND",
|
||||||
scale=layer.scaling,
|
scale=layer.scaling,
|
||||||
@@ -2031,7 +2064,7 @@ class AscendAttnBackend(AttentionBackend):
|
|||||||
c_kv_cache,
|
c_kv_cache,
|
||||||
query_rope=q_rope,
|
query_rope=q_rope,
|
||||||
key_rope=k_rope_cache,
|
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,
|
num_key_value_heads=layer.tp_k_head_num,
|
||||||
input_layout="TND",
|
input_layout="TND",
|
||||||
scale=layer.scaling,
|
scale=layer.scaling,
|
||||||
@@ -2046,6 +2079,7 @@ class AscendAttnBackend(AttentionBackend):
|
|||||||
workspace=workspace,
|
workspace=workspace,
|
||||||
out=[attn_output, softmax_lse],
|
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)
|
attn_output = attn_output.view(-1, layer.tp_q_head_num * layer.v_head_dim)
|
||||||
if (
|
if (
|
||||||
not self.graph_mode
|
not self.graph_mode
|
||||||
|
|||||||
@@ -547,6 +547,8 @@ class Glm4MoeLiteDecoderLayer(nn.Module):
|
|||||||
) -> None:
|
) -> None:
|
||||||
|
|
||||||
super().__init__()
|
super().__init__()
|
||||||
|
# Required for MTP: Glm4MoeLiteModelNextN bypasses Glm4MoeLiteForCausalLM.__init__
|
||||||
|
config.moe_layer_freq = 1
|
||||||
self.hidden_size = config.hidden_size
|
self.hidden_size = config.hidden_size
|
||||||
self.config = config
|
self.config = config
|
||||||
rope_theta, rope_scaling = get_rope_config(config)
|
rope_theta, rope_scaling = get_rope_config(config)
|
||||||
|
|||||||
Reference in New Issue
Block a user