[AMD] remove the redundant projection (#15178)
This commit is contained in:
@@ -1770,13 +1770,9 @@ class DeepseekV2AttentionMLA(nn.Module):
|
|||||||
res1=None,
|
res1=None,
|
||||||
output_unquantized_inp1=True, # return unqaunt kv_a
|
output_unquantized_inp1=True, # return unqaunt kv_a
|
||||||
)
|
)
|
||||||
kv = self.kv_b_proj(
|
|
||||||
kv_a_quanted,
|
|
||||||
)[0]
|
|
||||||
|
|
||||||
else:
|
else:
|
||||||
kv_a = self.kv_a_layernorm(kv_a)
|
kv_a = self.kv_a_layernorm(kv_a)
|
||||||
kv = self.kv_b_proj(kv_a)[0]
|
|
||||||
|
|
||||||
# kv_a = self.kv_a_layernorm(kv_a)
|
# kv_a = self.kv_a_layernorm(kv_a)
|
||||||
|
|
||||||
@@ -1800,7 +1796,12 @@ class DeepseekV2AttentionMLA(nn.Module):
|
|||||||
q.dtype,
|
q.dtype,
|
||||||
forward_batch,
|
forward_batch,
|
||||||
)
|
)
|
||||||
kv = self.kv_b_proj(kv_a)[0]
|
if _use_aiter_gfx95 and self.kv_b_proj.weight.dtype == torch.float8_e4m3fn:
|
||||||
|
kv = self.kv_b_proj(
|
||||||
|
kv_a_quanted,
|
||||||
|
)[0]
|
||||||
|
else:
|
||||||
|
kv = self.kv_b_proj(kv_a)[0]
|
||||||
kv = kv.view(-1, self.num_local_heads, self.qk_nope_head_dim + self.v_head_dim)
|
kv = kv.view(-1, self.num_local_heads, self.qk_nope_head_dim + self.v_head_dim)
|
||||||
k_nope = kv[..., : self.qk_nope_head_dim]
|
k_nope = kv[..., : self.qk_nope_head_dim]
|
||||||
v = kv[..., self.qk_nope_head_dim :]
|
v = kv[..., self.qk_nope_head_dim :]
|
||||||
|
|||||||
Reference in New Issue
Block a user