[AMD] remove the redundant projection (#15178)

This commit is contained in:
yctseng0211
2025-12-19 22:10:46 -08:00
committed by GitHub
parent e220da1723
commit 6468cb5823
+5 -4
View File
@@ -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,6 +1796,11 @@ class DeepseekV2AttentionMLA(nn.Module):
q.dtype, q.dtype,
forward_batch, forward_batch,
) )
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 = 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]