fix cohere

Author: shihaobai

lightllm/common/basemodel/layer_infer/template/transformer_layer_infer_cohere_template.py

@@ -44,11 +44,10 @@ def _bind_rotary_emb_fwd(self):
     def _get_qkv(
         self, input, infer_state: InferStateInfo, layer_weight
     ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        q = torch.mm(input.view(-1, self.embed_dim_), layer_weight.q_weight_)
-        cache_kv = torch.mm(
-            input.view(-1, self.embed_dim_),
-            layer_weight.kv_weight_,
-        ).view(-1, (self.tp_k_head_num_ + self.tp_v_head_num_), self.head_dim_)
+        q = layer_weight.q_proj.mm(input.view(-1, self.embed_dim_))
+        cache_kv = layer_weight.kv_proj.mm(input.view(-1, self.embed_dim_)).view(
+            -1, (self.tp_k_head_num_ + self.tp_v_head_num_), self.head_dim_
+        )
 
         if self.use_qk_norm_:
             q = q.view(-1, self.tp_q_head_num_, self.head_dim_)