[AMD] remove the redundant projection (#15178)
This commit is contained in:
@@ -1770,13 +1770,9 @@ class DeepseekV2AttentionMLA(nn.Module):
|
||||
res1=None,
|
||||
output_unquantized_inp1=True, # return unqaunt kv_a
|
||||
)
|
||||
kv = self.kv_b_proj(
|
||||
kv_a_quanted,
|
||||
)[0]
|
||||
|
||||
else:
|
||||
kv_a = self.kv_a_layernorm(kv_a)
|
||||
kv = self.kv_b_proj(kv_a)[0]
|
||||
|
||||
# kv_a = self.kv_a_layernorm(kv_a)
|
||||
|
||||
@@ -1800,7 +1796,12 @@ class DeepseekV2AttentionMLA(nn.Module):
|
||||
q.dtype,
|
||||
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)
|
||||
k_nope = kv[..., : self.qk_nope_head_dim]
|
||||
v = kv[..., self.qk_nope_head_dim :]
|
||||
|
||||
Reference in New Issue
Block a user