EngineX-Hygon/sglang
5
0
Fork 0
Code Issues Pull Requests Actions 7 Projects Releases Wiki Activity

Transpose mla weight offline (#1261)

Browse Source 
Co-authored-by: Yineng Zhang <me@zhyncs.com>
This commit is contained in:
Ke Bao
2024-08-30 14:45:40 +08:00
committed by GitHub
parent a362340b33
commit f414352ae6
1 changed files with 13 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

20
python/sglang/srt/models/deepseek_v2.py
View File

@@ -417,12 +417,8 @@ class DeepseekV2AttentionMLA(nn.Module):
v_head_dim=self.kv_lora_rank,
)
kv_b_proj = self.kv_b_proj
w_kc, w_vc = kv_b_proj.weight.unflatten(
0, (-1, qk_nope_head_dim + v_head_dim)
).split([qk_nope_head_dim, v_head_dim], dim=1)
self.w_kc = w_kc
self.w_vc = w_vc
self.w_kc = None
self.w_vc = None
def forward(
self,
@@ -464,7 +460,7 @@ class DeepseekV2AttentionMLA(nn.Module):
)
torch.bmm(
attn_output.transpose(0, 1),
self.w_vc.transpose(1, 2).contiguous(),
self.w_vc,
out=attn_bmm_output.transpose(0, 1),
)
@@ -715,5 +711,15 @@ class DeepseekV2ForCausalLM(nn.Module):
)
weight_loader(param, loaded_weight)
if global_server_args_dict["enable_mla"]:
for layer_id in range(self.config.num_hidden_layers):
self_attn = self.model.layers[layer_id].self_attn
w_kc, w_vc = self_attn.kv_b_proj.weight.unflatten(
0, (-1, self_attn.qk_nope_head_dim + self_attn.v_head_dim)
).split([self_attn.qk_nope_head_dim, self_attn.v_head_dim], dim=1)
self_attn.w_kc = w_kc.contiguous()
self_attn.w_vc = w_vc.transpose(1, 2).contiguous()
del self_attn.kv_b_proj
EntryClass = DeepseekV2ForCausalLM