[AMD] remove the redundant projection (#15178)

This commit is contained in:
yctseng0211
2025-12-20 14:10:46 +08:00
committed by GitHub
parent e220da1723
commit 6468cb5823

View File

@@ -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 :]